From 58c10b146ffd385538ebe3dbbb19a96a57f5c29f Mon Sep 17 00:00:00 2001 From: zhaozx-cn Date: Thu, 27 Aug 2026 07:34:39 +0000 Subject: [PATCH 01/11] add chunk kda fwd Signed-off-by: zhaozx-cn --- csrc/CMakeLists.txt | 30 + .../arch35/chunk_kda_fwd_tiling_impl.h | 53 + csrc/chunk_kda_fwd/op_host/chunk_kda_fwd.cpp | 686 +++ .../block_epilogue_kda_fwdh_update_arch22.hpp | 526 ++ .../block_epilogue_kda_fwdh_vnew_arch22.hpp | 396 ++ .../epilogue/kda_fwd_h_epilogue_policies.hpp | 28 + .../block_scheduler_kda_fwd_h_arch22.hpp | 445 ++ .../gemm/kernel/kda_fwd_h_kernel_arch22.hpp | 752 +++ .../arch35/chunk_kda_fwd_finalize_arch35.h | 1521 ++++++ .../op_kernel/arch35/chunk_kda_fwd_fwd_h.h | 1013 ++++ .../op_kernel/arch35/chunk_kda_fwd_impl.h | 61 + .../arch35/chunk_kda_fwd_post_wu_arch35.h | 2036 +++++++ .../arch35/chunk_kda_fwd_prepare_arch35.h | 4778 +++++++++++++++++ .../block/block_epilogue_kda_fwdh_regbase.hpp | 381 ++ .../block_epilogue_kda_fwdh_update_arch35.hpp | 434 ++ .../block_epilogue_kda_fwdh_vnew_arch35.hpp | 509 ++ .../epilogue/kda_fwd_h_epilogue_policies.hpp | 28 + .../block_scheduler_kda_fwd_h_arch35.hpp | 437 ++ .../gemm/kernel/kda_fwd_h_kernel_arch35.hpp | 1128 ++++ .../op_kernel/catlass/arch/kda_arch.hpp | 52 + .../catlass/arch/kda_cross_core_sync.hpp | 126 + .../catlass/arch/kda_local_tensor_buffer.hpp | 230 + .../op_kernel/catlass/arch/kda_resource.hpp | 43 + .../catlass/detail/kda_alignment.hpp | 103 + .../op_kernel/catlass/detail/kda_callback.hpp | 61 + .../catlass/detail/kda_dependent_false.hpp | 20 + .../op_kernel/catlass/detail/kda_macros.hpp | 26 + .../catlass/detail/kda_tag_to_layout.hpp | 86 + .../block_epilogue_elemwise_no_source.hpp | 185 + .../block/block_epilogue_fa_rescale_o.hpp | 337 ++ .../block_epilogue_fa_rescale_o_ascend950.hpp | 216 + .../block_epilogue_fa_softmax_ascend950.hpp | 405 ++ .../epilogue/block/block_epilogue_fixpipe.hpp | 83 + .../block/block_epilogue_mla_fd_rescale_o.hpp | 307 ++ .../block/block_epilogue_mla_rescale_o.hpp | 479 ++ .../block_epilogue_mla_tp1_rescale_o.hpp | 341 ++ .../block_epilogue_online_softmax_no_mask.hpp | 877 +++ .../block_epilogue_per_group_per_block.hpp | 863 +++ .../block_epilogue_per_token_dequant_tla.hpp | 784 +++ .../block_epilogue_rescale_o_no_split_row.hpp | 313 ++ ...gue_w4a4_per_token_per_channel_dequant.hpp | 252 + .../epilogue/block/kda_block_epilogue.hpp | 51 + ...kda_block_epilogue_elemwise_one_source.hpp | 185 + .../block/kda_block_epilogue_fa_softmax.hpp | 403 ++ .../block/kda_block_epilogue_gemm.hpp | 218 + .../block/kda_block_epilogue_gemv.hpp | 222 + .../block/kda_block_epilogue_mla_softmax.hpp | 441 ++ .../kda_block_epilogue_mla_tp1_softmax.hpp | 507 ++ .../kda_block_epilogue_per_token_dequant.hpp | 683 +++ .../catlass/epilogue/kda_dispatch_policy.hpp | 143 + .../epilogue/tile/copy_ub_to_l1_tla.hpp | 69 + .../epilogue/tile/kda_copy_gm_to_ub.hpp | 181 + .../epilogue/tile/kda_copy_gm_to_ub_tla.hpp | 143 + .../epilogue/tile/kda_copy_ub_to_gm.hpp | 136 + .../epilogue/tile/kda_copy_ub_to_gm_tla.hpp | 104 + .../epilogue/tile/kda_epilogue_tile_copy.hpp | 192 + .../catlass/epilogue/tile/kda_tile_cast.hpp | 46 + .../gemm/block/block_mmad_dynamic_aiv.hpp | 505 ++ .../gemm/block/block_mmad_dynamic_common.hpp | 353 ++ ...ad_dynamic_preload_async_with_callback.hpp | 441 ++ .../block_mmad_dynamic_single_core_splitk.hpp | 391 ++ .../gemm/block/block_mmad_dynamic_small.hpp | 284 + .../gemm/block/block_mmad_dynamic_streamk.hpp | 370 ++ .../block/block_mmad_fai_pv_head_tail.hpp | 315 ++ .../block/block_mmad_fai_pv_head_tail_tla.hpp | 350 ++ .../gemm/block/block_mmad_fai_pv_normal.hpp | 305 ++ .../block/block_mmad_fai_pv_normal_tla.hpp | 340 ++ .../gemm/block/block_mmad_fai_pv_tla.hpp | 306 ++ .../block/block_mmad_fai_qk_head_tail.hpp | 306 ++ .../block/block_mmad_fai_qk_head_tail_tla.hpp | 326 ++ .../gemm/block/block_mmad_fai_qk_normal.hpp | 337 ++ .../block/block_mmad_fai_qk_normal_tla.hpp | 347 ++ .../gemm/block/block_mmad_fai_qk_tla.hpp | 349 ++ .../block/block_mmad_pingpong_dequant_tla.hpp | 666 +++ .../block/block_mmad_pingpong_full_loadA.hpp | 347 ++ .../block_mmad_pingpong_full_loadA_tla.hpp | 582 ++ ..._mmad_pingpong_per_group_per_block_tla.hpp | 320 ++ ...ck_mmad_pingpong_slice_k_with_prologue.hpp | 806 +++ .../gemm/block/block_mmad_pingpong_tla.hpp | 630 +++ .../gemm/block/block_mmad_pingpong_tla_v2.hpp | 394 ++ .../block_mmad_pingpong_with_prologue.hpp | 535 ++ ...k_mmad_preload_async_with_callback_tla.hpp | 650 +++ .../gemm/block/block_mmad_preload_tla.hpp | 445 ++ .../block/block_mmad_single_core_splitk.hpp | 391 ++ .../catlass/gemm/block/block_mmad_small.hpp | 319 ++ .../gemm/block/block_mmad_sparse_tla.hpp | 397 ++ .../catlass/gemm/block/block_mmad_streamk.hpp | 406 ++ ...mad_w4a4_per_token_per_channel_dequant.hpp | 494 ++ .../catlass/gemm/block/kda_block_mmad.hpp | 143 + .../gemm/block/kda_block_mmad_fa_pv.hpp | 193 + .../gemm/block/kda_block_mmad_fa_qk.hpp | 199 + .../gemm/block/kda_block_mmad_gemm.hpp | 348 ++ .../gemm/block/kda_block_mmad_mla_pv.hpp | 207 + .../block/kda_block_mmad_mla_pv_tp1_spec.hpp | 221 + .../gemm/block/kda_block_mmad_mla_qk.hpp | 241 + .../block/kda_block_mmad_mla_qk_tp1_spec.hpp | 259 + .../gemm/block/kda_block_mmad_pingpong.hpp | 362 ++ .../block/kda_block_mmad_pingpong_bias.hpp | 421 ++ .../gemm/block/kda_block_mmad_preload.hpp | 402 ++ .../block/kda_block_mmad_preload_async.hpp | 448 ++ ...block_mmad_preload_async_with_callback.hpp | 449 ++ .../catlass/gemm/block/kda_block_swizzle.hpp | 1018 ++++ .../catlass/gemm/kda_gemm_dispatch_policy.hpp | 373 ++ .../catlass/gemm/kda_gemm_helper.hpp | 385 ++ .../op_kernel/catlass/gemm/kda_gemm_type.hpp | 29 + .../gemm/tile/ascend950/copy_l0c_to_ub.hpp | 166 + .../ascend950/kda_ascend950_copy_gm_to_l1.hpp | 727 +++ .../kda_ascend950_copy_l0c_to_dst.hpp | 274 + .../kda_ascend950_copy_l0c_to_gm.hpp | 533 ++ .../ascend950/kda_ascend950_copy_l1_to_bt.hpp | 107 + .../kda_ascend950_copy_l1_to_l0a.hpp | 467 ++ .../kda_ascend950_copy_l1_to_l0b.hpp | 437 ++ .../atlasa2/kda_atlasa2_cast_int4_to_int8.hpp | 180 + .../atlasa2/kda_atlasa2_copy_gm_to_l1.hpp | 2194 ++++++++ .../atlasa2/kda_atlasa2_copy_gm_to_ub.hpp | 99 + .../atlasa2/kda_atlasa2_copy_l0c_to_gm.hpp | 632 +++ .../atlasa2/kda_atlasa2_copy_l1_to_bt.hpp | 62 + .../atlasa2/kda_atlasa2_copy_l1_to_fp.hpp | 62 + .../atlasa2/kda_atlasa2_copy_l1_to_l0a.hpp | 719 +++ .../atlasa2/kda_atlasa2_copy_l1_to_l0b.hpp | 881 +++ .../atlasa2/kda_atlasa2_copy_ub_to_gm.hpp | 131 + .../tile/atlasa2/kda_cast_fp8_to_fp16.hpp | 415 ++ .../tile/atlasa2/kda_cast_int8_to_fp16.hpp | 188 + .../catlass/gemm/tile/copy_l0c_to_ub.hpp | 18 + .../gemm/tile/kda_cast_fp8_to_fp16.hpp | 18 + .../gemm/tile/kda_cast_int4_to_int8.hpp | 18 + .../gemm/tile/kda_cast_int8_to_fp16.hpp | 18 + .../catlass/gemm/tile/kda_copy_gm_to_l1.hpp | 21 + .../catlass/gemm/tile/kda_copy_gm_to_ub.hpp | 18 + .../catlass/gemm/tile/kda_copy_l0c_to_gm.hpp | 21 + .../catlass/gemm/tile/kda_copy_l1_to_bt.hpp | 21 + .../catlass/gemm/tile/kda_copy_l1_to_fp.hpp | 18 + .../catlass/gemm/tile/kda_copy_l1_to_l0a.hpp | 21 + .../catlass/gemm/tile/kda_copy_l1_to_l0b.hpp | 21 + .../catlass/gemm/tile/kda_copy_ub_to_gm.hpp | 18 + .../catlass/gemm/tile/kda_gemm_tile_copy.hpp | 520 ++ .../catlass/gemm/tile/kda_tile_copy_tla.hpp | 70 + .../catlass/gemm/tile/kda_tile_mmad.hpp | 285 + .../catlass/gemm/tile/kda_tile_muls.hpp | 49 + .../catlass/gemm/tile/kda_tile_traits.hpp | 48 + .../catlass/gemv/kda_gemv_helper.hpp | 122 + .../op_kernel/catlass/kda_catlass.hpp | 55 + .../op_kernel/catlass/kda_conv_coord.hpp | 1320 +++++ .../op_kernel/catlass/kda_coord.hpp | 336 ++ .../op_kernel/catlass/kda_debug.hpp | 73 + .../op_kernel/catlass/kda_gemm_coord.hpp | 161 + .../op_kernel/catlass/kda_gemv_coord.hpp | 108 + .../op_kernel/catlass/kda_matrix_coord.hpp | 108 + .../op_kernel/catlass/kda_numeric_size.hpp | 42 + .../op_kernel/catlass/layout/kda_layout.hpp | 19 + .../op_kernel/catlass/layout/kda_matrix.hpp | 1818 +++++++ .../op_kernel/catlass/layout/kda_tensor.hpp | 315 ++ .../op_kernel/catlass/layout/kda_vector.hpp | 139 + .../chunk_kda_fwd/op_kernel/chunk_kda_fwd.cpp | 188 + .../op_kernel/chunk_kda_fwd_common.h | 331 ++ .../op_kernel/chunk_kda_fwd_finalize.h | 810 +++ .../chunk_kda_fwd_h_runtime_tiling.h | 42 + .../op_kernel/chunk_kda_fwd_post_wu.h | 926 ++++ .../op_kernel/chunk_kda_fwd_prepare.h | 2518 +++++++++ .../op_kernel/chunk_kda_fwd_tiling_data.h | 99 + .../op_kernel/chunk_kda_fwd_varlen.h | 78 + .../op_kernel/kda_gate_cumsum_kernel.h | 749 +++ .../block_mmad_pingpong_tla_kernel_utils.hpp | 1014 ++++ .../block/block_mmad_pingpong_tla_multi.hpp | 677 +++ .../block_mmad_pingpong_tla_preloadA_l1B.hpp | 498 ++ .../kernel_utils/tile/copy_l0c_to_ub.hpp | 381 ++ .../op_kernel/kernel_utils/vector/regbase.hpp | 150 + .../chunk_kda_fwd/op_kernel/tla/int_tuple.hpp | 239 + .../op_kernel/tla/kda_tensor.hpp | 413 ++ .../op_kernel/tla/kda_tla_layout.hpp | 580 ++ .../tla/numeric/integer_sequence.hpp | 67 + .../tla/numeric/integral_constant.hpp | 196 + .../op_kernel/tla/numeric/math.hpp | 52 + csrc/chunk_kda_fwd/op_kernel/tla/tuple.hpp | 141 + .../op_kernel/tla/type_traits.hpp | 43 + csrc/pytorch_extensions.cpp | 13 + include/sgl_kenel_npu_ops.h | 12 + .../sgl_kernel_npu/fla/chunk_kda_fwd.py | 93 + 178 files changed, 68263 insertions(+) create mode 100644 csrc/chunk_kda_fwd/op_host/arch35/chunk_kda_fwd_tiling_impl.h create mode 100644 csrc/chunk_kda_fwd/op_host/chunk_kda_fwd.cpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/arch22/fwd_h/epilogue/block/block_epilogue_kda_fwdh_update_arch22.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/arch22/fwd_h/epilogue/block/block_epilogue_kda_fwdh_vnew_arch22.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/arch22/fwd_h/epilogue/kda_fwd_h_epilogue_policies.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/arch22/fwd_h/gemm/block/block_scheduler_kda_fwd_h_arch22.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/arch22/fwd_h/gemm/kernel/kda_fwd_h_kernel_arch22.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/arch35/chunk_kda_fwd_finalize_arch35.h create mode 100644 csrc/chunk_kda_fwd/op_kernel/arch35/chunk_kda_fwd_fwd_h.h create mode 100644 csrc/chunk_kda_fwd/op_kernel/arch35/chunk_kda_fwd_impl.h create mode 100644 csrc/chunk_kda_fwd/op_kernel/arch35/chunk_kda_fwd_post_wu_arch35.h create mode 100644 csrc/chunk_kda_fwd/op_kernel/arch35/chunk_kda_fwd_prepare_arch35.h create mode 100644 csrc/chunk_kda_fwd/op_kernel/arch35/fwd_h/epilogue/block/block_epilogue_kda_fwdh_regbase.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/arch35/fwd_h/epilogue/block/block_epilogue_kda_fwdh_update_arch35.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/arch35/fwd_h/epilogue/block/block_epilogue_kda_fwdh_vnew_arch35.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/arch35/fwd_h/epilogue/kda_fwd_h_epilogue_policies.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/arch35/fwd_h/gemm/block/block_scheduler_kda_fwd_h_arch35.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/arch35/fwd_h/gemm/kernel/kda_fwd_h_kernel_arch35.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/arch/kda_arch.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/arch/kda_cross_core_sync.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/arch/kda_local_tensor_buffer.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/arch/kda_resource.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/detail/kda_alignment.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/detail/kda_callback.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/detail/kda_dependent_false.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/detail/kda_macros.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/detail/kda_tag_to_layout.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_elemwise_no_source.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_fa_rescale_o.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_fa_rescale_o_ascend950.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_fa_softmax_ascend950.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_fixpipe.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_mla_fd_rescale_o.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_mla_rescale_o.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_mla_tp1_rescale_o.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_online_softmax_no_mask.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_per_group_per_block.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_per_token_dequant_tla.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_rescale_o_no_split_row.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_w4a4_per_token_per_channel_dequant.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/kda_block_epilogue.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/kda_block_epilogue_elemwise_one_source.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/kda_block_epilogue_fa_softmax.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/kda_block_epilogue_gemm.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/kda_block_epilogue_gemv.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/kda_block_epilogue_mla_softmax.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/kda_block_epilogue_mla_tp1_softmax.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/kda_block_epilogue_per_token_dequant.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/kda_dispatch_policy.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/tile/copy_ub_to_l1_tla.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/tile/kda_copy_gm_to_ub.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/tile/kda_copy_gm_to_ub_tla.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/tile/kda_copy_ub_to_gm.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/tile/kda_copy_ub_to_gm_tla.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/tile/kda_epilogue_tile_copy.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/tile/kda_tile_cast.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_dynamic_aiv.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_dynamic_common.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_dynamic_preload_async_with_callback.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_dynamic_single_core_splitk.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_dynamic_small.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_dynamic_streamk.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_pv_head_tail.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_pv_head_tail_tla.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_pv_normal.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_pv_normal_tla.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_pv_tla.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_qk_head_tail.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_qk_head_tail_tla.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_qk_normal.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_qk_normal_tla.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_qk_tla.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_pingpong_dequant_tla.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_pingpong_full_loadA.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_pingpong_full_loadA_tla.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_pingpong_per_group_per_block_tla.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_pingpong_slice_k_with_prologue.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_pingpong_tla.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_pingpong_tla_v2.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_pingpong_with_prologue.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_preload_async_with_callback_tla.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_preload_tla.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_single_core_splitk.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_small.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_sparse_tla.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_streamk.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_w4a4_per_token_per_channel_dequant.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_fa_pv.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_fa_qk.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_gemm.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_mla_pv.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_mla_pv_tp1_spec.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_mla_qk.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_mla_qk_tp1_spec.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_pingpong.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_pingpong_bias.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_preload.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_preload_async.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_preload_async_with_callback.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_swizzle.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/kda_gemm_dispatch_policy.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/kda_gemm_helper.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/kda_gemm_type.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/ascend950/copy_l0c_to_ub.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/ascend950/kda_ascend950_copy_gm_to_l1.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/ascend950/kda_ascend950_copy_l0c_to_dst.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/ascend950/kda_ascend950_copy_l0c_to_gm.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/ascend950/kda_ascend950_copy_l1_to_bt.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/ascend950/kda_ascend950_copy_l1_to_l0a.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/ascend950/kda_ascend950_copy_l1_to_l0b.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_atlasa2_cast_int4_to_int8.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_atlasa2_copy_gm_to_l1.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_atlasa2_copy_gm_to_ub.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_atlasa2_copy_l0c_to_gm.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_atlasa2_copy_l1_to_bt.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_atlasa2_copy_l1_to_fp.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_atlasa2_copy_l1_to_l0a.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_atlasa2_copy_l1_to_l0b.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_atlasa2_copy_ub_to_gm.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_cast_fp8_to_fp16.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_cast_int8_to_fp16.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/copy_l0c_to_ub.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_cast_fp8_to_fp16.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_cast_int4_to_int8.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_cast_int8_to_fp16.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_copy_gm_to_l1.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_copy_gm_to_ub.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_copy_l0c_to_gm.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_copy_l1_to_bt.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_copy_l1_to_fp.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_copy_l1_to_l0a.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_copy_l1_to_l0b.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_copy_ub_to_gm.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_gemm_tile_copy.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_tile_copy_tla.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_tile_mmad.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_tile_muls.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_tile_traits.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemv/kda_gemv_helper.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/kda_catlass.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/kda_conv_coord.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/kda_coord.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/kda_debug.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/kda_gemm_coord.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/kda_gemv_coord.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/kda_matrix_coord.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/kda_numeric_size.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/layout/kda_layout.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/layout/kda_matrix.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/layout/kda_tensor.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/layout/kda_vector.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd.cpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_common.h create mode 100644 csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_finalize.h create mode 100644 csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_h_runtime_tiling.h create mode 100644 csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_post_wu.h create mode 100644 csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_prepare.h create mode 100644 csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_tiling_data.h create mode 100644 csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_varlen.h create mode 100644 csrc/chunk_kda_fwd/op_kernel/kda_gate_cumsum_kernel.h create mode 100644 csrc/chunk_kda_fwd/op_kernel/kernel_utils/block/block_mmad_pingpong_tla_kernel_utils.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/kernel_utils/block/block_mmad_pingpong_tla_multi.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/kernel_utils/block/block_mmad_pingpong_tla_preloadA_l1B.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/kernel_utils/tile/copy_l0c_to_ub.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/kernel_utils/vector/regbase.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/tla/int_tuple.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/tla/kda_tensor.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/tla/kda_tla_layout.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/tla/numeric/integer_sequence.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/tla/numeric/integral_constant.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/tla/numeric/math.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/tla/tuple.hpp create mode 100644 csrc/chunk_kda_fwd/op_kernel/tla/type_traits.hpp create mode 100644 python/sgl_kernel_npu/sgl_kernel_npu/fla/chunk_kda_fwd.py diff --git a/csrc/CMakeLists.txt b/csrc/CMakeLists.txt index efd40c305b..5389f9d19e 100644 --- a/csrc/CMakeLists.txt +++ b/csrc/CMakeLists.txt @@ -45,6 +45,7 @@ set(OP_SRCS ${PROJECT_OP_SRC_BASE}/sparse_attn_sharedkv_metadata/op_host/sparse_attn_sharedkv_metadata.cpp ${PROJECT_OP_SRC_BASE}/sparse_attention_score/op_host/sparse_attention_score.cpp ${PROJECT_OP_SRC_BASE}/sparse_attention_score/op_host/tiling/sparse_attention_score_tiling.cpp + ${PROJECT_OP_SRC_BASE}/chunk_kda_fwd/op_host/chunk_kda_fwd.cpp ) # A3-only host sources if(SGL_KERNEL_ENABLE_A3_ONLY_OPS) @@ -147,6 +148,32 @@ ascendc_compile_definitions(workspace_kernel PRIVATE -DHAVE_TILING ) +# chunk_kda_fwd: dedicated kernel target with its own compile options. The kernel +# manages cross-pipe synchronization manually (SetFlag/WaitFlag/SyncAll), so the +# aclnn build compiles it with --cce-auto-sync=off; keep that here instead of +# applying it to the shared workspace_kernel target. +ascendc_library(chunk_kda_fwd_kernel STATIC + ${PROJECT_OP_SRC_BASE}/chunk_kda_fwd/op_kernel/chunk_kda_fwd.cpp +) +ascendc_include_directories(chunk_kda_fwd_kernel PRIVATE + ${PROJECT_OP_SRC_BASE}/chunk_kda_fwd/op_kernel + ${PROJECT_OP_SRC_BASE}/utils/kernel + ${ASCEND_INCLUDE_DIR} +) +ascendc_compile_definitions(chunk_kda_fwd_kernel PRIVATE + -DHAVE_WORKSPACE + -DHAVE_TILING +) +ascendc_compile_options(chunk_kda_fwd_kernel PRIVATE + --cce-auto-sync=off + -Wno-deprecated-declarations +) +if(_sgl_soc_lower MATCHES "^ascend950") + ascendc_compile_definitions(chunk_kda_fwd_kernel PRIVATE + -DENABLE_CV_COMM_VIA_SSBUF=true + ) +endif() + # Build catlass kernel in its own target (no sparse_attention_score include dir), # so catlass's #include "tla/layout.hpp" resolves to its own copy and avoids # symbol redefinition (RoundDown/CeilDiv/Max/Min/...) with sparse_attention_score. @@ -216,6 +243,7 @@ target_link_libraries(${OP_PLUGIN_NAME} PRIVATE ${A3_ONLY_KERNEL_TARGETS} workspace_kernel no_workspace_kernel + chunk_kda_fwd_kernel ${CATLASS_WORKSPACE_KERNEL_TARGETS} torch_npu ascendcl @@ -246,6 +274,8 @@ target_include_directories(${OP_PLUGIN_NAME} PRIVATE ${PROJECT_OP_SRC_BASE}/sparse_attn_sharedkv/op_host ${PROJECT_OP_SRC_BASE}/sparse_attention_score/op_host ${PROJECT_OP_SRC_BASE}/sparse_attention_score/op_host/tiling + ${PROJECT_OP_SRC_BASE}/chunk_kda_fwd/op_host + ${PROJECT_OP_SRC_BASE}/chunk_kda_fwd/op_kernel ${TORCH_DIR}/include ${TORCH_DIR}/include/torch/csrc/api/include ${TORCH_NPU_DIR}/include diff --git a/csrc/chunk_kda_fwd/op_host/arch35/chunk_kda_fwd_tiling_impl.h b/csrc/chunk_kda_fwd/op_host/arch35/chunk_kda_fwd_tiling_impl.h new file mode 100644 index 0000000000..4f72a7a012 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_host/arch35/chunk_kda_fwd_tiling_impl.h @@ -0,0 +1,53 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#pragma once + +namespace optiling::arch35 { + +struct ChunkKdaFwdArch35Options { + bool computeGateInPrepare = false; + bool fusePostWu = false; + bool fusePostWuIntoFwdH = false; + bool useDenseFwdH = false; +}; + +inline ChunkKdaFwdArch35Options ConfigureChunkKdaFwdArch35( + bool isAscend950, bool qIsBf16, bool rawGIsFp32, bool hasALog, + bool useGateInKernel, bool safeGate, bool isVarLen, int64_t seqlen, + int64_t vHeads, int64_t chunkSize, int64_t kDim, int64_t vDim, + bool storeQG, bool storeVNew, bool storeH) +{ + ChunkKdaFwdArch35Options options; + const bool shapeSupported = + isAscend950 && chunkSize == 64 && kDim == 128 && vDim == 128; + if (!shapeSupported) { + return options; + } + + // Tiling keys describe shape families independently of the SoC. These + // options only enable arch35 sub-pipelines within the selected family. + options.computeGateInPrepare = + qIsBf16 && rawGIsFp32 && hasALog && + useGateInKernel && safeGate; + const bool denseAligned = !isVarLen && seqlen % chunkSize == 0; + options.useDenseFwdH = denseAligned && qIsBf16; + const bool canFusePreparePostWu = + denseAligned && qIsBf16 && safeGate && vHeads % 2 == 0; + options.fusePostWuIntoFwdH = + options.useDenseFwdH && canFusePreparePostWu && + options.computeGateInPrepare && + !storeQG && !storeVNew && !storeH; + options.fusePostWu = + canFusePreparePostWu && !options.fusePostWuIntoFwdH; + return options; +} + +} // namespace optiling::arch35 diff --git a/csrc/chunk_kda_fwd/op_host/chunk_kda_fwd.cpp b/csrc/chunk_kda_fwd/op_host/chunk_kda_fwd.cpp new file mode 100644 index 0000000000..9b176669c4 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_host/chunk_kda_fwd.cpp @@ -0,0 +1,686 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file chunk_kda_fwd.cpp + * \brief host-side direct-launch implementation of chunk_kda_fwd + * + * Port of the aclnn ChunkKdaFwd front-end. Instead of building an aclnn executor + * graph (transpose/reshape + ADD_TO_LAUNCHER_LIST_AICORE), the layout conversions + * are done with plain torch ops and the kernel is launched directly through + * EXEC_KERNEL_CMD. The tiling computation mirrors Tiling4ChunkKdaFwd and the + * serialized tiling struct is copied GM->stack inside the kernel entry. + */ + +#include +#include +#include +#include +#include + +#include "acl/acl.h" + +#include +#include +#include + +#include "torch_npu/csrc/core/npu/NPUStream.h" +#include "torch_npu/csrc/core/npu/DeviceUtils.h" +#include "torch_npu/csrc/framework/OpCommand.h" + +#include "tiling/platform/platform_ascendc.h" + +#include "defines.h" +#include "torch_helper.h" +#include "aclrtlaunch_chunk_kda_fwd.h" + +#include "../op_kernel/chunk_kda_fwd_tiling_data.h" +#include "arch35/chunk_kda_fwd_tiling_impl.h" + +namespace sglang { +namespace npu_kernel { + +constexpr int64_t SYS_WORKSPACE_SIZE = 16777216; // 16 MB reserved system workspace +constexpr uint32_t PADDING_BYTE = 32U; + +constexpr uint64_t KDA_ALIGN = 512; +constexpr uint64_t KDA_SOLVE_SCRATCH_SLOTS = 5; +constexpr uint64_t KDA_SOLVE_PIPELINE_DEPTH = 4; +constexpr uint64_t KDA_SCORE_QUEUE_SLOTS = 4; +constexpr uint64_t KDA_SCORE_SCRATCH_PLANES = 3; +constexpr uint64_t KDA_GDN_PIPELINE_DEPTH = 2; + +constexpr int64_t MAX_KDA_K_DIM = 256; +constexpr int64_t MAX_KDA_HEAD_NUM = 128; +constexpr int64_t MAX_KDA_VARLEN_SEQUENCES = 1024; + +enum class KdaFwdLayout { + BSND, + BNSD, + TND, + NTD, +}; + +uint64_t AlignWorkspace(uint64_t bytes) +{ + return (bytes + KDA_ALIGN - 1) / KDA_ALIGN * KDA_ALIGN; +} + +uint64_t AllocateWorkspace(uint64_t &cursor, uint64_t bytes) +{ + const uint64_t offset = AlignWorkspace(cursor); + cursor = offset + bytes; + return offset; +} + +struct KdaShapeInfo { + bool isRank3 = false; + int64_t batch = 0; + int64_t seqlen = 0; + int64_t hNum = 0; + int64_t hvNum = 0; + int64_t kDim = 0; + int64_t vDim = 0; + int64_t seqNum = 0; + int64_t totalChunks = 0; +}; + +bool ParseLayout(const std::string &layout, KdaFwdLayout &parsed) +{ + if (layout == "BSND") { + parsed = KdaFwdLayout::BSND; + } else if (layout == "BNSD") { + parsed = KdaFwdLayout::BNSD; + } else if (layout == "TND") { + parsed = KdaFwdLayout::TND; + } else if (layout == "NTD") { + parsed = KdaFwdLayout::NTD; + } else { + return false; + } + return true; +} + +int64_t Dim(const at::Tensor &tensor, size_t idx) +{ + return tensor.size(static_cast(idx)); +} + +bool SameShape(const at::Tensor &lhs, const at::Tensor &rhs) +{ + return lhs.dim() == rhs.dim() && lhs.sizes() == rhs.sizes(); +} + +bool HasShape(const at::Tensor &tensor, std::initializer_list expected) +{ + if (tensor.dim() != static_cast(expected.size())) { + return false; + } + size_t idx = 0; + for (int64_t dim : expected) { + if (tensor.size(static_cast(idx++)) != dim) { + return false; + } + } + return true; +} + +int64_t CountChunks(const c10::optional &cuSeqlens, int64_t seqlen, int64_t chunkSize) +{ + if (!cuSeqlens.has_value()) { + return (seqlen + chunkSize - 1) / chunkSize; + } + const at::Tensor cu = cuSeqlens->cpu().contiguous(); + const int64_t *cuPtr = reinterpret_cast(cu.data_ptr()); + const int64_t n = cu.numel(); + int64_t chunks = 0; + for (int64_t idx = 0; idx + 1 < n; ++idx) { + chunks += (cuPtr[idx + 1] - cuPtr[idx] + chunkSize - 1) / chunkSize; + } + return chunks; +} + +void ResolveShapeInfo(const at::Tensor &q, const at::Tensor &v, const c10::optional &beta, + const c10::optional &g, KdaFwdLayout layout, int64_t chunkSize, + const c10::optional &cuSeqlens, KdaShapeInfo &info) +{ + info.isRank3 = layout == KdaFwdLayout::TND || layout == KdaFwdLayout::NTD; + const int64_t tensorRank = info.isRank3 ? 3 : 4; + const int64_t betaRank = info.isRank3 ? 2 : 3; + TORCH_CHECK(q.dim() == tensorRank && v.dim() == tensorRank && g->dim() == tensorRank && + beta->dim() == betaRank, + "chunk_kda_fwd: q/k/v/g and beta ranks must match layout: rank3/rank2 for TND/NTD, " + "rank4/rank3 for BSND/BNSD."); + + if (layout == KdaFwdLayout::TND) { + info.batch = 1; + info.seqlen = Dim(q, 0); + info.hNum = Dim(q, 1); + info.kDim = Dim(q, 2); + info.hvNum = Dim(v, 1); + info.vDim = Dim(v, 2); + TORCH_CHECK(HasShape(v, {info.seqlen, info.hvNum, info.vDim}) && + HasShape(*g, {info.seqlen, info.hvNum, info.kDim}) && + HasShape(*beta, {info.seqlen, info.hvNum}), + "chunk_kda_fwd: TND expects v/g/beta as [T,HV,V], [T,HV,K], [T,HV]."); + } else if (layout == KdaFwdLayout::NTD) { + info.batch = 1; + info.hNum = Dim(q, 0); + info.seqlen = Dim(q, 1); + info.kDim = Dim(q, 2); + info.hvNum = Dim(v, 0); + info.vDim = Dim(v, 2); + TORCH_CHECK(HasShape(v, {info.hvNum, info.seqlen, info.vDim}) && + HasShape(*g, {info.hvNum, info.seqlen, info.kDim}) && + HasShape(*beta, {info.hvNum, info.seqlen}), + "chunk_kda_fwd: NTD expects v/g/beta as [HV,T,V], [HV,T,K], [HV,T]."); + } else if (layout == KdaFwdLayout::BSND) { + info.batch = Dim(q, 0); + info.seqlen = Dim(q, 1); + info.hNum = Dim(q, 2); + info.kDim = Dim(q, 3); + info.hvNum = Dim(v, 2); + info.vDim = Dim(v, 3); + TORCH_CHECK(HasShape(v, {info.batch, info.seqlen, info.hvNum, info.vDim}) && + HasShape(*g, {info.batch, info.seqlen, info.hvNum, info.kDim}) && + HasShape(*beta, {info.batch, info.seqlen, info.hvNum}), + "chunk_kda_fwd: BSND expects v/g/beta as [B,T,HV,V], [B,T,HV,K], [B,T,HV]."); + } else { // BNSD + info.batch = Dim(q, 0); + info.hNum = Dim(q, 1); + info.seqlen = Dim(q, 2); + info.kDim = Dim(q, 3); + info.hvNum = Dim(v, 1); + info.vDim = Dim(v, 3); + TORCH_CHECK(HasShape(v, {info.batch, info.hvNum, info.seqlen, info.vDim}) && + HasShape(*g, {info.batch, info.hvNum, info.seqlen, info.kDim}) && + HasShape(*beta, {info.batch, info.hvNum, info.seqlen}), + "chunk_kda_fwd: BNSD expects v/g/beta as [B,HV,T,V], [B,HV,T,K], [B,HV,T]."); + } + + info.seqNum = cuSeqlens.has_value() ? cuSeqlens->numel() - 1 : info.batch; + info.totalChunks = CountChunks(cuSeqlens, info.seqlen, chunkSize); +} + +void ComputeTilingData(int64_t batch, int64_t seqlen, int64_t hNum, int64_t hvNum, int64_t kDim, + int64_t vDim, int64_t chunkSize, int64_t seqNum, int64_t totalChunks, + int64_t inputRank, double scale, double lowerBound, bool hasInitialState, + bool isVarLen, bool safeGate, bool sequenceMajor, bool useGateInKernel, + bool hasALog, bool hasDtBias, int64_t gateDataType, + bool storeFinalState, bool storeGk, bool storeW, + bool storeU, bool storeQG, bool storeKg, bool storeVNew, bool storeH, + uint32_t blockDim, const optiling::arch35::ChunkKdaFwdArch35Options &arch35Options, + int64_t betaDataType, ChunkKdaFwd::ChunkKdaFwdTilingData &td) +{ + std::memset(&td, 0, sizeof(td)); + + td.batch = batch; + td.seqNum = seqNum; + td.qHeadNum = hNum; + td.vHeadNum = hvNum; + td.seqlen = seqlen; + td.kHeadDim = kDim; + td.vHeadDim = vDim; + td.chunkSize = chunkSize; + td.totalChunks = totalChunks; + td.inputRank = inputRank; + td.scale = static_cast(scale); + td.lowerBound = static_cast(lowerBound); + td.hasInitialState = hasInitialState; + td.isVarLen = isVarLen; + td.safeGate = safeGate; + td.inputSequenceMajor = sequenceMajor; + td.useGateInKernel = useGateInKernel; + td.hasALog = hasALog; + td.hasDtBias = hasDtBias; + td.computeGateInPrepare = arch35Options.computeGateInPrepare; + td.fusePostWu = arch35Options.fusePostWu; + td.fusePostWuIntoFwdH = arch35Options.fusePostWuIntoFwdH; + td.useDenseFwdH = arch35Options.useDenseFwdH; + td.storeFinalState = storeFinalState; + td.storeGk = storeGk; + td.storeW = storeW; + td.storeU = storeU; + td.storeQG = storeQG; + td.storeKg = storeKg; + td.storeVNew = storeVNew; + td.storeH = storeH; + td.gateDataType = gateDataType; + td.gateUsedCoreNum = static_cast(blockDim) * 2; + td.prepareUsedCoreNum = blockDim; + td.postWuUsedCoreNum = blockDim; + td.outputUsedCoreNum = blockDim; + td.betaDataType = betaDataType; + + const uint64_t dataBytes = 2; // direct-launch currently only supports bf16 q/k/v + const uint64_t tokenHeads = static_cast(batch) * hvNum * seqlen; + const uint64_t kTensorBytes = tokenHeads * kDim * dataBytes; + const uint64_t vTensorBytes = tokenHeads * vDim * dataBytes; + const uint64_t gkBytes = tokenHeads * kDim * sizeof(float); + const uint64_t stateElements = static_cast(seqNum) * hvNum * kDim * vDim; + const uint64_t hChunkCount = + isVarLen ? static_cast(totalChunks) : static_cast(batch) * totalChunks; + const uint64_t hBytes = hChunkCount * hvNum * kDim * vDim * dataBytes; + + uint64_t cursor = 0; + td.gkStorageOffset = storeGk ? 0 : AllocateWorkspace(cursor, gkBytes); + td.finalStateStorageOffset = + storeFinalState ? 0 : AllocateWorkspace(cursor, stateElements * sizeof(float)); + td.wStorageOffset = storeW ? 0 : AllocateWorkspace(cursor, kTensorBytes); + td.uStorageOffset = storeU ? 0 : AllocateWorkspace(cursor, vTensorBytes); + td.qgStorageOffset = storeQG ? 0 : AllocateWorkspace(cursor, kTensorBytes); + td.kgStorageOffset = storeKg ? 0 : AllocateWorkspace(cursor, kTensorBytes); + const uint64_t vNewStorageBytes = + arch35Options.useDenseFwdH && !storeVNew + ? static_cast(batch) * hvNum * chunkSize * vDim * dataBytes + : vTensorBytes; + td.vNewStorageOffset = storeVNew ? 0 : AllocateWorkspace(cursor, vNewStorageBytes); + const uint64_t hStorageBytes = + arch35Options.useDenseFwdH && !storeH + ? static_cast(batch) * hvNum * kDim * vDim * dataBytes + : hBytes; + td.hStorageOffset = storeH ? 0 : AllocateWorkspace(cursor, hStorageBytes); + td.qgScaledOffset = AllocateWorkspace(cursor, kTensorBytes); + + const uint64_t matrixBytes = tokenHeads * chunkSize * sizeof(float); + td.prepareAqkFp32Offset = AllocateWorkspace(cursor, matrixBytes); + td.prepareAkkFp32Offset = AllocateWorkspace(cursor, matrixBytes); + td.prepareScratchOffset = AlignWorkspace(cursor); + const uint64_t solveDepth = safeGate ? KDA_SOLVE_PIPELINE_DEPTH : 1; + const uint64_t solveBytes = + static_cast(blockDim) * solveDepth * KDA_SOLVE_SCRATCH_SLOTS * chunkSize * chunkSize * + sizeof(float); + const uint64_t scoreBytes = static_cast(blockDim) * KDA_SCORE_QUEUE_SLOTS * + KDA_SCORE_SCRATCH_PLANES * chunkSize * kDim * dataBytes; + cursor = td.prepareScratchOffset + AlignWorkspace(solveBytes) + scoreBytes; + + td.postWuScratchOffset = AlignWorkspace(cursor); + if (!arch35Options.fusePostWu && !arch35Options.fusePostWuIntoFwdH) { + cursor = td.postWuScratchOffset + tokenHeads * kDim * sizeof(float); + } + + td.fwdHWorkspaceBaseOffset = AlignWorkspace(cursor); + uint64_t fwdHCursor = 0; + td.vWorkspaceOffset = AllocateWorkspace( + fwdHCursor, static_cast(blockDim) * chunkSize * vDim * sizeof(float) * + KDA_GDN_PIPELINE_DEPTH); + td.vUpdateWorkspaceOffset = AllocateWorkspace( + fwdHCursor, static_cast(blockDim) * chunkSize * vDim * sizeof(float) * + KDA_GDN_PIPELINE_DEPTH); + td.kDecayWorkspaceOffset = AllocateWorkspace( + fwdHCursor, static_cast(blockDim) * chunkSize * kDim * sizeof(float) * + KDA_GDN_PIPELINE_DEPTH); + td.hWorkspaceOffset = AllocateWorkspace( + fwdHCursor, static_cast(blockDim) * kDim * vDim * sizeof(float) * + KDA_GDN_PIPELINE_DEPTH); + const uint64_t tokenBatch = isVarLen ? static_cast(seqNum) : 1; + td.numSeqWorkspaceOffset = AllocateWorkspace(fwdHCursor, (tokenBatch + 1) * sizeof(int64_t)); + td.numChunksWorkspaceOffset = AllocateWorkspace(fwdHCursor, (tokenBatch + 1) * sizeof(int64_t)); + cursor = td.fwdHWorkspaceBaseOffset + AlignWorkspace(fwdHCursor); + + td.outputScratchOffset = AllocateWorkspace(cursor, 2 * tokenHeads * vDim * sizeof(float)); +} + +std::tuple, c10::optional, at::Tensor, at::Tensor, + c10::optional, c10::optional, c10::optional, + c10::optional, c10::optional, c10::optional> +chunk_kda_fwd(const at::Tensor &q, const at::Tensor &k, const at::Tensor &v, const at::Tensor &g, + const at::Tensor &beta, const c10::optional &aLog, + const c10::optional &dtBias, const c10::optional &initialState, + const c10::optional &cuSeqlens, + const c10::optional &chunkIndices, const std::string &layout, double scale, + int64_t chunkSize, bool safeGate, double lowerBound, bool useGateInKernel, + bool stateVFirst, bool outputFinalState, bool outputGk, bool outputW, bool outputU, + bool outputQG, bool outputKg, bool outputVNew, bool outputH) +{ + TORCH_CHECK(q.defined() && k.defined() && v.defined() && g.defined() && beta.defined(), + "chunk_kda_fwd: q, k, v, g and beta must be defined"); + + TORCH_CHECK(chunkSize == 64 || chunkSize == 128, "chunk_kda_fwd: chunkSize must be 64 or 128"); + + KdaFwdLayout parsedLayout; + TORCH_CHECK(ParseLayout(layout, parsedLayout), "chunk_kda_fwd: layout must be BSND, BNSD, TND or NTD"); + + const bool isRank3 = parsedLayout == KdaFwdLayout::TND || parsedLayout == KdaFwdLayout::NTD; + + // dtype checks: q/k/v must be bf16 (direct-launch restriction), g/beta fp32 or bf16. + TORCH_CHECK(q.scalar_type() == at::kBFloat16 && k.scalar_type() == at::kBFloat16 && + v.scalar_type() == at::kBFloat16, + "chunk_kda_fwd: q, k and v must be bfloat16 (direct-launch currently supports bf16)"); + const at::ScalarType gateType = g.scalar_type(); + TORCH_CHECK(gateType == at::kFloat || gateType == at::kBFloat16, + "chunk_kda_fwd: g must be float32 or bfloat16"); + const at::ScalarType betaType = beta.scalar_type(); + TORCH_CHECK(betaType == at::kFloat || betaType == at::kBFloat16, + "chunk_kda_fwd: beta must be float32 or bfloat16"); + const bool gIsFp32 = gateType == at::kFloat; + const int64_t gateDataType = gIsFp32 ? 2 : (gateType == at::kBFloat16 ? 1 : 0); + const int64_t betaDataType = betaType == at::kFloat ? 0 : 1; + + if (aLog.has_value() && aLog->defined()) { + TORCH_CHECK(aLog->scalar_type() == at::kFloat, "chunk_kda_fwd: a_log must be float32"); + } + if (dtBias.has_value() && dtBias->defined()) { + TORCH_CHECK(dtBias->scalar_type() == at::kFloat, "chunk_kda_fwd: dt_bias must be float32"); + } + if (initialState.has_value() && initialState->defined()) { + TORCH_CHECK(initialState->scalar_type() == at::kFloat, + "chunk_kda_fwd: initial_state must be float32"); + } + if (useGateInKernel) { + TORCH_CHECK(aLog.has_value() && aLog->defined(), + "chunk_kda_fwd: a_log is required when use_gate_in_kernel is true"); + } + + KdaShapeInfo info; + ResolveShapeInfo(q, v, beta, g, parsedLayout, chunkSize, cuSeqlens, info); + + TORCH_CHECK(info.hNum > 0 && info.hvNum >= info.hNum && info.hvNum % info.hNum == 0, + "chunk_kda_fwd: H and HV must be positive, HV must be a multiple of H"); + TORCH_CHECK(info.hNum <= MAX_KDA_HEAD_NUM && info.hvNum <= MAX_KDA_HEAD_NUM, + "chunk_kda_fwd: H and HV must be <= 128"); + TORCH_CHECK(info.kDim >= 16 && info.kDim <= MAX_KDA_K_DIM && info.kDim % 16 == 0 && + info.vDim >= 16 && info.vDim <= 256 && info.vDim % 16 == 0, + "chunk_kda_fwd: K/V must be multiples of 16, K <= 256 and V <= 256"); + TORCH_CHECK(SameShape(q, k) && q.dim() == k.dim(), + "chunk_kda_fwd: q and k must have identical shape"); + TORCH_CHECK(info.seqlen > 0 && info.totalChunks > 0, "chunk_kda_fwd: invalid sequence length"); + if (cuSeqlens.has_value() && cuSeqlens->defined()) { + TORCH_CHECK(cuSeqlens->scalar_type() == at::kLong, "chunk_kda_fwd: cu_seqlens must be int64"); + TORCH_CHECK(cuSeqlens->numel() >= 2, "chunk_kda_fwd: cu_seqlens must contain at least [0, T]"); + TORCH_CHECK(!isRank3 || info.batch == 1, "chunk_kda_fwd: rank4 varlen requires B=1"); + TORCH_CHECK(info.seqNum <= MAX_KDA_VARLEN_SEQUENCES, + "chunk_kda_fwd: varlen supports at most 1024 sequences"); + const at::Tensor cuCpu = cuSeqlens->cpu().contiguous(); + const int64_t *cuPtr = reinterpret_cast(cuCpu.data_ptr()); + TORCH_CHECK(cuPtr[0] == 0, "chunk_kda_fwd: cu_seqlens[0] must be 0"); + TORCH_CHECK(cuPtr[cuCpu.numel() - 1] == info.seqlen, + "chunk_kda_fwd: cu_seqlens last element must equal the sequence length"); + for (int64_t idx = 0; idx + 1 < cuCpu.numel(); ++idx) { + TORCH_CHECK(cuPtr[idx] <= cuPtr[idx + 1], + "chunk_kda_fwd: cu_seqlens must be nondecreasing"); + } + } + if (chunkIndices.has_value() && chunkIndices->defined()) { + TORCH_CHECK(cuSeqlens.has_value() && cuSeqlens->defined(), + "chunk_kda_fwd: chunk_indices requires cu_seqlens"); + TORCH_CHECK(chunkIndices->scalar_type() == at::kLong, "chunk_kda_fwd: chunk_indices must be int64"); + TORCH_CHECK(chunkIndices->numel() == info.totalChunks * 2, + "chunk_kda_fwd: chunk_indices must contain one (seq_id, chunk_id) pair per chunk"); + } + if (useGateInKernel && safeGate) { + TORCH_CHECK(lowerBound >= -5.0 && lowerBound < 0.0, + "chunk_kda_fwd: lower_bound must be in [-5, 0) when safe_gate is true"); + } + if (initialState.has_value() && initialState->defined()) { + const bool valid = stateVFirst + ? HasShape(*initialState, {info.seqNum, info.hvNum, info.vDim, info.kDim}) + : HasShape(*initialState, {info.seqNum, info.hvNum, info.kDim, info.vDim}); + TORCH_CHECK(valid, + "chunk_kda_fwd: initial_state must be [N,HV,K,V] (state_v_first=false) or " + "[N,HV,V,K] (state_v_first=true)"); + } + + const auto ascendcPlatform = platform_ascendc::PlatformAscendCManager::GetInstance(); + const uint32_t physicalCoreNum = std::max(ascendcPlatform->GetCoreNumAic(), 1); + const bool isAscend950 = + ascendcPlatform->GetSocVersion() == platform_ascendc::SocVersion::ASCEND950; + const bool isVarLen = cuSeqlens.has_value() && cuSeqlens->defined(); + const uint64_t fwdHTaskCount = static_cast(isVarLen ? info.seqNum : info.batch) * + info.hvNum; + const uint32_t blockDim = + isAscend950 ? static_cast( + std::min(physicalCoreNum, std::max(fwdHTaskCount, 1))) : + physicalCoreNum; + + const bool hasALog = aLog.has_value() && aLog->defined(); + const bool hasDtBias = dtBias.has_value() && dtBias->defined(); + const bool hasInitialState = initialState.has_value() && initialState->defined(); + + const auto arch35Options = optiling::arch35::ConfigureChunkKdaFwdArch35( + isAscend950, true /* qIsBf16 */, gIsFp32, hasALog, useGateInKernel, safeGate, isVarLen, + info.seqlen, info.hvNum, chunkSize, info.kDim, info.vDim, outputQG, outputVNew, outputH); + + const bool storeFinalState = outputFinalState; + const bool storeGk = outputGk; + const bool storeW = outputW; + const bool storeU = outputU; + const bool storeQG = outputQG; + const bool storeKg = outputKg; + const bool storeVNew = outputVNew; + const bool storeH = outputH; + + ChunkKdaFwd::ChunkKdaFwdTilingData tilingData; + ComputeTilingData(info.batch, info.seqlen, info.hNum, info.hvNum, info.kDim, info.vDim, chunkSize, + info.seqNum, info.totalChunks, info.isRank3 ? 3 : 4, scale, lowerBound, + hasInitialState, isVarLen, safeGate, parsedLayout == KdaFwdLayout::BSND, + useGateInKernel, hasALog, hasDtBias, gateDataType, storeFinalState, + storeGk, storeW, storeU, storeQG, storeKg, storeVNew, storeH, blockDim, + arch35Options, betaDataType, tilingData); + + // ---- layout conversion to the kernel's input convention ---- + at::Tensor qHead, kHead, vHead, gHead, betaHead; + if (parsedLayout == KdaFwdLayout::BSND) { + qHead = q.contiguous(); + kHead = k.contiguous(); + vHead = v.contiguous(); + gHead = g.contiguous(); + betaHead = beta.transpose(1, 2).contiguous(); // (B, S, HV) -> (B, HV, S) + } else if (parsedLayout == KdaFwdLayout::BNSD) { + qHead = q.contiguous(); + kHead = k.contiguous(); + vHead = v.contiguous(); + gHead = g.contiguous(); + betaHead = beta.contiguous(); + } else if (parsedLayout == KdaFwdLayout::TND) { + qHead = q.transpose(0, 1).contiguous().reshape({1, info.hNum, info.seqlen, info.kDim}); + kHead = k.transpose(0, 1).contiguous().reshape({1, info.hNum, info.seqlen, info.kDim}); + vHead = v.transpose(0, 1).contiguous().reshape({1, info.hvNum, info.seqlen, info.vDim}); + gHead = g.transpose(0, 1).contiguous().reshape({1, info.hvNum, info.seqlen, info.kDim}); + betaHead = beta.transpose(0, 1).contiguous().reshape({1, info.hvNum, info.seqlen}); + } else { // NTD + qHead = q.reshape({1, info.hNum, info.seqlen, info.kDim}); + kHead = k.reshape({1, info.hNum, info.seqlen, info.kDim}); + vHead = v.reshape({1, info.hvNum, info.seqlen, info.vDim}); + gHead = g.reshape({1, info.hvNum, info.seqlen, info.kDim}); + betaHead = beta.reshape({1, info.hvNum, info.seqlen}); + } + + at::Tensor initStateContig; + void *initStatePtr = nullptr; + if (hasInitialState) { + at::Tensor init = *initialState; + if (stateVFirst) { + init = init.transpose(-1, -2).contiguous(); // (N, HV, V, K) -> (N, HV, K, V) + } + initStateContig = init.contiguous(); + initStatePtr = initStateContig.data_ptr(); + } + + void *aLogPtr = nullptr; + if (hasALog) { + aLogPtr = aLog->contiguous().data_ptr(); + } + void *dtBiasPtr = nullptr; + if (hasDtBias) { + dtBiasPtr = dtBias->contiguous().data_ptr(); + } + + at::Tensor cuSeqlensContig, chunkIndicesContig; + void *cuSeqlensPtr = nullptr; + void *chunkIndicesPtr = nullptr; + if (cuSeqlens.has_value() && cuSeqlens->defined()) { + cuSeqlensContig = cuSeqlens->contiguous(); + cuSeqlensPtr = cuSeqlensContig.data_ptr(); + } + if (chunkIndices.has_value() && chunkIndices->defined()) { + chunkIndicesContig = chunkIndices->contiguous(); + chunkIndicesPtr = chunkIndicesContig.data_ptr(); + } + + // ---- outputs (kernel internal layouts, then reshaped/transposed for the user) ---- + const int64_t b = info.batch; + const int64_t s = info.seqlen; + const int64_t hv = info.hvNum; + const int64_t kd = info.kDim; + const int64_t vd = info.vDim; + const int64_t c = chunkSize; + + at::Tensor attnOut = at::empty({b, s, hv, vd}, q.options()); + at::Tensor aqkOut = at::empty({b, hv, s, c}, q.options()); + at::Tensor akkOut = at::empty({b, hv, s, c}, q.options()); + + at::Tensor finalStateOut; + void *finalStatePtr = nullptr; + if (storeFinalState) { + finalStateOut = at::empty({info.seqNum, hv, kd, vd}, + at::TensorOptions().dtype(at::kFloat).device(q.device())); + finalStatePtr = finalStateOut.data_ptr(); + } + + at::Tensor gkOut; + void *gkPtr = nullptr; + if (storeGk) { + gkOut = at::empty({b, hv, s, kd}, + at::TensorOptions().dtype(at::kFloat).device(q.device())); + gkPtr = gkOut.data_ptr(); + } + + at::Tensor wOut; + void *wPtr = nullptr; + if (storeW) { + wOut = at::empty({b, hv, s, kd}, q.options()); + wPtr = wOut.data_ptr(); + } + + at::Tensor uOut; + void *uPtr = nullptr; + if (storeU) { + uOut = at::empty({b, hv, s, vd}, q.options()); + uPtr = uOut.data_ptr(); + } + + at::Tensor qgOut; + void *qgPtr = nullptr; + if (storeQG) { + qgOut = at::empty({b, hv, s, kd}, q.options()); + qgPtr = qgOut.data_ptr(); + } + + at::Tensor kgOut; + void *kgPtr = nullptr; + if (storeKg) { + kgOut = at::empty({b, hv, s, kd}, q.options()); + kgPtr = kgOut.data_ptr(); + } + + at::Tensor vNewOut; + void *vNewPtr = nullptr; + if (storeVNew) { + vNewOut = at::empty({b, hv, s, vd}, q.options()); + vNewPtr = vNewOut.data_ptr(); + } + + at::Tensor hOutInternal; + void *hPtr = nullptr; + if (storeH) { + // kernel-internal h layout is (B, HV, NC, K, V); reshaped below for the user. + hOutInternal = at::empty({b, hv, info.totalChunks, kd, vd}, q.options()); + hPtr = hOutInternal.data_ptr(); + } + + // ---- workspace + tiling ---- + int32_t tilingSize = (static_cast(sizeof(ChunkKdaFwd::ChunkKdaFwdTilingData)) + + static_cast(PADDING_BYTE) - 1) / + PADDING_BYTE * PADDING_BYTE; + + auto cpuTiling = at::empty({tilingSize}, at::kByte); + std::memcpy(cpuTiling.data_ptr(), &tilingData, sizeof(ChunkKdaFwd::ChunkKdaFwdTilingData)); + at::Tensor tilingTensor = TorchNpuHelper::CopyTensorHostToDevice(cpuTiling); + + const uint64_t tokenHeads = static_cast(b) * hv * s; + const uint64_t totalUserWorkspace = tilingData.outputScratchOffset + 2 * tokenHeads * vd * sizeof(float); + const int64_t totalWorkspaceBytes = + SYS_WORKSPACE_SIZE + static_cast(AlignWorkspace(totalUserWorkspace)); + auto workspaceTensor = + at::empty({totalWorkspaceBytes}, at::TensorOptions().dtype(at::kByte).device(q.device())); + + EXEC_KERNEL_CMD(chunk_kda_fwd, blockDim, qHead, kHead, vHead, gHead, betaHead, aLogPtr, dtBiasPtr, + initStatePtr, cuSeqlensPtr, chunkIndicesPtr, attnOut, finalStatePtr, gkPtr, aqkOut, + akkOut, wPtr, uPtr, qgPtr, kgPtr, vNewPtr, hPtr, workspaceTensor, tilingTensor); + + // ---- post-process outputs to the aclnn-documented user layouts ---- + at::Tensor attn = attnOut; + at::Tensor aqk = aqkOut; + at::Tensor akk = akkOut; + if (isRank3) { + attn = attnOut.reshape({s, hv, vd}); + aqk = aqkOut.reshape({hv, s, c}); + akk = akkOut.reshape({hv, s, c}); + } + + c10::optional finalState = c10::nullopt; + if (storeFinalState) { + at::Tensor fs = finalStateOut; + if (stateVFirst) { + fs = fs.transpose(-1, -2).contiguous(); // (N, HV, K, V) -> (N, HV, V, K) + } + finalState = fs; + } + + c10::optional gk = c10::nullopt; + if (storeGk) { + gk = isRank3 ? gkOut.reshape({hv, s, kd}) : gkOut; + } + + c10::optional w = c10::nullopt; + if (storeW) { + w = isRank3 ? wOut.reshape({hv, s, kd}) : wOut; + } + + c10::optional u = c10::nullopt; + if (storeU) { + u = isRank3 ? uOut.reshape({hv, s, vd}) : uOut; + } + + c10::optional qg = c10::nullopt; + if (storeQG) { + qg = isRank3 ? qgOut.reshape({hv, s, kd}) : qgOut; + } + + c10::optional kg = c10::nullopt; + if (storeKg) { + kg = isRank3 ? kgOut.reshape({hv, s, kd}) : kgOut; + } + + c10::optional vNew = c10::nullopt; + if (storeVNew) { + vNew = isRank3 ? vNewOut.reshape({hv, s, vd}) : vNewOut; + } + + c10::optional h = c10::nullopt; + if (storeH) { + // (B, HV, NC, K, V) -> (B, NC, HV, K, V) [or state_v_first: (B, NC, HV, V, K)] + at::Tensor hView = hOutInternal.permute({0, 2, 1, 3, 4}); + if (stateVFirst) { + hView = hOutInternal.permute({0, 2, 1, 4, 3}); + } + at::Tensor hOut = hView.contiguous(); + if (isRank3) { + hOut = hOut.reshape({info.totalChunks, hv, stateVFirst ? vd : kd, stateVFirst ? kd : vd}); + } + h = hOut; + } + + return std::make_tuple(attn, finalState, gk, aqk, akk, w, u, qg, kg, vNew, h); +} + +} // namespace npu_kernel +} // namespace sglang diff --git a/csrc/chunk_kda_fwd/op_kernel/arch22/fwd_h/epilogue/block/block_epilogue_kda_fwdh_update_arch22.hpp b/csrc/chunk_kda_fwd/op_kernel/arch22/fwd_h/epilogue/block/block_epilogue_kda_fwdh_update_arch22.hpp new file mode 100644 index 0000000000..7d811b7cb8 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/arch22/fwd_h/epilogue/block/block_epilogue_kda_fwdh_update_arch22.hpp @@ -0,0 +1,526 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_KDA_FWDH_UPDATE_ARCH22_HPP +#define CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_KDA_FWDH_UPDATE_ARCH22_HPP +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "../kda_fwd_h_epilogue_policies.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/kda_matrix_coord.hpp" +#include "catlass/epilogue/tile/kda_epilogue_tile_copy.hpp" + +namespace Catlass::Epilogue::Block { + +template +class BlockEpilogue { + static constexpr bool kGated = KGatedTag::value; + static constexpr bool scalarGated = KGatedTag::scalarGated; + static constexpr bool useExp2 = KGatedTag::useExp2; + static constexpr float LN2 = 0.6931471805599453f; + +public: + using DispatchPolicy = EpilogueAtlasKdaFwdHUpdate; + using ArchTag = typename DispatchPolicy::ArchTag; + + using HElementOutput = typename HOutputType_::Element; + using GElementInput = typename GInputType_::Element; + using HElementInput = typename HInputType_::Element; + using HUpdateElementInput = typename HUpdateInputType_::Element; + using FinalStateElement = typename FinalStateType_::Element; + + CATLASS_DEVICE + BlockEpilogue(Arch::Resource &resource) + { + constexpr uint32_t CALC_BUF_OFFSET = 0; + constexpr uint32_t PING_BUF_0_OFFSET = 32 * 1024; + constexpr uint32_t PING_BUF_1_OFFSET = 48 * 1024; + constexpr uint32_t PING_BUF_2_OFFSET = 64 * 1024; + constexpr uint32_t PING_BUF_3_OFFSET = 80 * 1024; + constexpr uint32_t PONG_BUF_0_OFFSET = 96 * 1024; + constexpr uint32_t PONG_BUF_1_OFFSET = 112 * 1024; + constexpr uint32_t PONG_BUF_2_OFFSET = 128 * 1024; + constexpr uint32_t PONG_BUF_3_OFFSET = 144 * 1024; + constexpr uint32_t PING_G_BUF_OFFSET = 160 * 1024; + constexpr uint32_t PONG_G_BUF_OFFSET = 161 * 1024; + constexpr uint32_t PING_G_SUB_BUF_OFFSET = 162 * 1024; + constexpr uint32_t PONG_G_SUB_BUF_OFFSET = 163 * 1024; + constexpr uint32_t PING_G_INPUT_BUF_OFFSET = 164 * 1024; + constexpr uint32_t PONG_G_INPUT_BUF_OFFSET = 165 * 1024; + constexpr uint32_t SHARE_BUF_OFFSET = 166 * 1024; + + calcUbTensor = resource.ubBuf.template GetBufferByByte(CALC_BUF_OFFSET); + + hUpdateUbTensor_ping = resource.ubBuf.template GetBufferByByte(PING_BUF_0_OFFSET); + gkBroadcastUbTensor_ping = resource.ubBuf.template GetBufferByByte(PING_BUF_1_OFFSET); + hUbTensor_ping = resource.ubBuf.template GetBufferByByte(PING_BUF_3_OFFSET); + finalOutputUbTensor_ping = resource.ubBuf.template GetBufferByByte(PING_BUF_3_OFFSET); + glastUbTensor_ping = resource.ubBuf.template GetBufferByByte(PING_G_INPUT_BUF_OFFSET); + + hUpdateUbTensor_pong = resource.ubBuf.template GetBufferByByte(PONG_BUF_0_OFFSET); + gkBroadcastUbTensor_pong = resource.ubBuf.template GetBufferByByte(PONG_BUF_1_OFFSET); + hUbTensor_pong = resource.ubBuf.template GetBufferByByte(PONG_BUF_3_OFFSET); + finalOutputUbTensor_pong = resource.ubBuf.template GetBufferByByte(PONG_BUF_3_OFFSET); + glastUbTensor_pong = resource.ubBuf.template GetBufferByByte(PONG_G_INPUT_BUF_OFFSET); + + if constexpr (kGated) { + gkLastUbTensor_ping = resource.ubBuf.template GetBufferByByte(PING_G_SUB_BUF_OFFSET); + gkLastUbTensor_pong = resource.ubBuf.template GetBufferByByte(PONG_G_SUB_BUF_OFFSET); + gkInputUbTensor_ping = resource.ubBuf.template GetBufferByByte(PING_G_INPUT_BUF_OFFSET); + gkInputUbTensor_pong = resource.ubBuf.template GetBufferByByte(PONG_G_INPUT_BUF_OFFSET); + shareBufferGk_ = resource.ubBuf.template GetBufferByByte(SHARE_BUF_OFFSET); + } + } + + CATLASS_DEVICE + ~BlockEpilogue() {} + + template + CATLASS_DEVICE void CopyGmToUb(AscendC::LocalTensor dst, AscendC::GlobalTensor src, uint32_t rows, + uint32_t cols, uint32_t srcStride) + { + if (cols == srcStride) { + AscendC::DataCopy(dst, src, rows * cols); + return; + } + AscendC::DataCopyExtParams copyParams{static_cast(rows), + static_cast(cols * sizeof(Element)), + static_cast((srcStride - cols) * sizeof(Element)), 0, 0}; + AscendC::DataCopyPadExtParams padParams{false, 0, 0, 0}; + AscendC::DataCopyPad(dst, src, copyParams, padParams); + } + + template + CATLASS_DEVICE void CopyUbToGm(AscendC::GlobalTensor dst, AscendC::LocalTensor src, uint32_t rows, + uint32_t cols, uint32_t dstStride) + { + if (cols == dstStride) { + AscendC::DataCopy(dst, src, rows * cols); + return; + } + AscendC::DataCopyExtParams copyParams{static_cast(rows), + static_cast(cols * sizeof(Element)), 0, + static_cast((dstStride - cols) * sizeof(Element)), 0}; + AscendC::DataCopyPad(dst, src, copyParams); + } + + CATLASS_DEVICE + void operator()(AscendC::GlobalTensor hOutput, AscendC::GlobalTensor finalState, + AscendC::GlobalTensor gInput, AscendC::GlobalTensor hInput, + AscendC::GlobalTensor hUpdateInput, AscendC::GlobalTensor gkInput, + AscendC::GlobalTensor initialState, uint32_t chunkSize, uint32_t kHeadDim, + uint32_t vBlockDim, uint32_t vHeadDim, Arch::CrossCoreFlag cube2Done, bool isInitialState, + bool isFinalState, bool storeFinalState, bool useInitialState, bool isPing, bool cube2AlreadyWaited) + { + static constexpr uint32_t ROW_TILE = 16; + uint32_t mActual = kHeadDim; + uint32_t nActual = vBlockDim; + uint32_t outputStride = vHeadDim; + uint32_t subBlockIdx = AscendC::GetSubBlockIdx(); + uint32_t subBlockNum = AscendC::GetSubBlockNum(); + uint32_t rowsPerSubBlock = CeilDiv(mActual, subBlockNum); + uint32_t rowBegin = subBlockIdx * rowsPerSubBlock; + uint32_t rowEnd = rowBegin + rowsPerSubBlock; + if (rowEnd > mActual) { + rowEnd = mActual; + } + if (rowBegin >= mActual) { + if (!cube2AlreadyWaited) { + Arch::CrossCoreWaitFlag(cube2Done); + } + return; + } + + AscendC::ResetMask(); + + AscendC::GlobalTensor gInputThisSubBlock = gInput; + + uint32_t pingpongFlag = isPing ? 0 : pongBaseEvent; + AscendC::LocalTensor hUpdateUbTensor = isPing ? hUpdateUbTensor_ping : hUpdateUbTensor_pong; + AscendC::LocalTensor gkBroadcastUbTensor = isPing ? gkBroadcastUbTensor_ping : gkBroadcastUbTensor_pong; + AscendC::LocalTensor hUbTensor = isPing ? hUbTensor_ping : hUbTensor_pong; + AscendC::LocalTensor finalOutputUbTensor = + isPing ? finalOutputUbTensor_ping : finalOutputUbTensor_pong; + AscendC::LocalTensor glastUbTensor = isPing ? glastUbTensor_ping : glastUbTensor_pong; + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID2 + pingpongFlag); + bool useFp32StateUpdate = + storeFinalState && std::is_same::value && (!isInitialState || useInitialState); + + float muls = 1.0f; + if constexpr (scalarGated) { + GElementInput gLastVal = gInputThisSubBlock.GetValue(chunkSize - 1); + float gLastFloat = 0.0f; + if constexpr (std::is_same::value) { + gLastFloat = gLastVal; + } else if constexpr (std::is_same::value) { + gLastFloat = (float)gLastVal; + } else if constexpr (std::is_same::value) { + gLastFloat = AscendC::ToFloat(gLastVal); + } + glastUbTensor.SetValue(0, gLastFloat); + + AscendC::SetFlag(EVENT_ID3 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID3 + pingpongFlag); + if constexpr (useExp2) { + AscendC::Muls(glastUbTensor, glastUbTensor, LN2, 1); + AscendC::PipeBarrier(); + } + AscendC::Exp(glastUbTensor, glastUbTensor, 1); + AscendC::SetFlag(EVENT_ID3 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID3 + pingpongFlag); + muls = glastUbTensor.GetValue(0); + } + if constexpr (kGated) { + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + } + if constexpr (scalarGated) { + AscendC::SetFlag(EVENT_ID3 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID3 + pingpongFlag); + } + + if (nActual <= 128 && nActual == outputStride) { + uint32_t mActualThisSubBlock = rowEnd - rowBegin; + AscendC::GlobalTensor hOutputThisSubBlock = hOutput[rowBegin * outputStride]; + AscendC::GlobalTensor hInputThisSubBlock = hInput[rowBegin * outputStride]; + AscendC::GlobalTensor hUpdateInputThisSubBlock = hUpdateInput[rowBegin * nActual]; + AscendC::GlobalTensor finalStateThisSubBlock = finalState[rowBegin * outputStride]; + + if (storeFinalState && isInitialState && std::is_same::value) { + AscendC::WaitFlag(EVENT_ID2 + pingpongFlag); + } else { + AscendC::WaitFlag(EVENT_ID2 + pingpongFlag); + } + if constexpr (std::is_same::value) { + if (useFp32StateUpdate) { + if (isInitialState) { + AscendC::DataCopy(calcUbTensor, initialState[rowBegin * outputStride], + mActualThisSubBlock * nActual); + } else { + AscendC::DataCopy(calcUbTensor, finalStateThisSubBlock, mActualThisSubBlock * nActual); + } + } else { + AscendC::DataCopy(hUbTensor, hInputThisSubBlock, mActualThisSubBlock * nActual); + } + } else { + AscendC::DataCopy(hUbTensor, hInputThisSubBlock, mActualThisSubBlock * nActual); + } + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID2 + pingpongFlag); + + if (!useFp32StateUpdate) { + AscendC::Cast(calcUbTensor, hUbTensor, AscendC::RoundMode::CAST_NONE, mActualThisSubBlock * nActual); + AscendC::PipeBarrier(); + } + AscendC::Muls(calcUbTensor, calcUbTensor, muls, mActualThisSubBlock * nActual); + AscendC::PipeBarrier(); + + if constexpr (kGated) { + AscendC::GlobalTensor gkLastInput = gkInput[(chunkSize - 1) * kHeadDim + rowBegin]; + AscendC::LocalTensor gkLastUbTensor = isPing ? gkLastUbTensor_ping : gkLastUbTensor_pong; + AscendC::LocalTensor gkInputUbTensor = + isPing ? gkInputUbTensor_ping : gkInputUbTensor_pong; + + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + if constexpr (std::is_same::value) { + AscendC::DataCopy(gkLastUbTensor, gkLastInput, mActualThisSubBlock); + } else { + AscendC::DataCopy(gkInputUbTensor, gkLastInput, mActualThisSubBlock); + } + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID2 + pingpongFlag); + if constexpr (!std::is_same::value) { + AscendC::Cast(gkLastUbTensor, gkInputUbTensor, AscendC::RoundMode::CAST_NONE, mActualThisSubBlock); + } + AscendC::PipeBarrier(); + AscendC::Muls(gkLastUbTensor, gkLastUbTensor, LN2, mActualThisSubBlock); + AscendC::PipeBarrier(); + AscendC::Exp(gkLastUbTensor, gkLastUbTensor, mActualThisSubBlock); + AscendC::PipeBarrier(); + + uint32_t gkBrcReptime = (mActualThisSubBlock + 8 - 1) / 8; + uint32_t dstShapeGk[2] = {gkBrcReptime * 8, nActual}; + uint32_t srcShapeGk[2] = {gkBrcReptime * 8, 1}; + AscendC::Broadcast(hUpdateUbTensor, gkLastUbTensor, dstShapeGk, srcShapeGk, + shareBufferGk_); + AscendC::PipeBarrier(); + AscendC::Mul(calcUbTensor, calcUbTensor, hUpdateUbTensor, mActualThisSubBlock * nActual); + AscendC::PipeBarrier(); + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + } + + if (!cube2AlreadyWaited) { + Arch::CrossCoreWaitFlag(cube2Done); + } + + if constexpr (kGated) { + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + } + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + AscendC::DataCopy(hUpdateUbTensor, hUpdateInputThisSubBlock, mActualThisSubBlock * nActual); + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + AscendC::Add(hUpdateUbTensor, calcUbTensor, hUpdateUbTensor, mActualThisSubBlock * nActual); + AscendC::PipeBarrier(); + if (storeFinalState && isFinalState && std::is_same::value) { + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + } + + if constexpr (std::is_same::value) { + if (storeFinalState) { + if (!isFinalState) { + AscendC::Cast(hUbTensor, hUpdateUbTensor, AscendC::RoundMode::CAST_RINT, + mActualThisSubBlock * nActual); + AscendC::PipeBarrier(); + } + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + AscendC::DataCopy(finalStateThisSubBlock, hUpdateUbTensor, mActualThisSubBlock * nActual); + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + if (!isFinalState) { + AscendC::DataCopy(hOutputThisSubBlock, hUbTensor, mActualThisSubBlock * nActual); + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + } + } else { + AscendC::Cast(hUbTensor, hUpdateUbTensor, AscendC::RoundMode::CAST_RINT, + mActualThisSubBlock * nActual); + AscendC::PipeBarrier(); + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID2 + pingpongFlag); + AscendC::DataCopy(hOutputThisSubBlock, hUbTensor, mActualThisSubBlock * nActual); + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + } + } else { + if (storeFinalState && isFinalState) { + AscendC::Cast(finalOutputUbTensor, hUpdateUbTensor, AscendC::RoundMode::CAST_RINT, + mActualThisSubBlock * nActual); + AscendC::PipeBarrier(); + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID2 + pingpongFlag); + AscendC::DataCopy(finalStateThisSubBlock, finalOutputUbTensor, mActualThisSubBlock * nActual); + } else { + AscendC::Cast(hUbTensor, hUpdateUbTensor, AscendC::RoundMode::CAST_RINT, + mActualThisSubBlock * nActual); + AscendC::PipeBarrier(); + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID2 + pingpongFlag); + AscendC::DataCopy(hOutputThisSubBlock, hUbTensor, mActualThisSubBlock * nActual); + } + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + } + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + return; + } + + if (!cube2AlreadyWaited) { + Arch::CrossCoreWaitFlag(cube2Done); + } + + bool waitHFromV = storeFinalState && isInitialState && std::is_same::value; + bool waitUpdateFromMte3 = false; + uint32_t updateReadyEvent = EVENT_ID3 + pingpongFlag; + bool waitGkFromScalar = true; + for (uint32_t rowStart = rowBegin; rowStart < rowEnd; rowStart += ROW_TILE) { + uint32_t rowsThisTile = rowEnd - rowStart; + if (rowsThisTile > ROW_TILE) { + rowsThisTile = ROW_TILE; + } + + AscendC::GlobalTensor hOutputThisTile = hOutput[rowStart * outputStride]; + AscendC::GlobalTensor hInputThisTile = hInput[rowStart * outputStride]; + AscendC::GlobalTensor hUpdateInputThisTile = hUpdateInput[rowStart * nActual]; + AscendC::GlobalTensor finalStateThisTile = finalState[rowStart * outputStride]; + + if (waitHFromV) { + AscendC::WaitFlag(EVENT_ID2 + pingpongFlag); + } else { + AscendC::WaitFlag(EVENT_ID2 + pingpongFlag); + } + if constexpr (std::is_same::value) { + if (useFp32StateUpdate) { + if (isInitialState) { + CopyGmToUb(calcUbTensor, initialState[rowStart * outputStride], rowsThisTile, nActual, + outputStride); + } else { + CopyGmToUb(calcUbTensor, finalStateThisTile, rowsThisTile, nActual, outputStride); + } + } else { + CopyGmToUb(hUbTensor, hInputThisTile, rowsThisTile, nActual, outputStride); + } + } else { + CopyGmToUb(hUbTensor, hInputThisTile, rowsThisTile, nActual, outputStride); + } + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID2 + pingpongFlag); + + if (!useFp32StateUpdate) { + AscendC::Cast(calcUbTensor, hUbTensor, AscendC::RoundMode::CAST_NONE, rowsThisTile * nActual); + AscendC::PipeBarrier(); + } + AscendC::Muls(calcUbTensor, calcUbTensor, muls, rowsThisTile * nActual); + AscendC::PipeBarrier(); + + if constexpr (kGated) { + AscendC::GlobalTensor gkLastInput = gkInput[(chunkSize - 1) * kHeadDim + rowStart]; + AscendC::LocalTensor gkLastUbTensor = isPing ? gkLastUbTensor_ping : gkLastUbTensor_pong; + AscendC::LocalTensor gkInputUbTensor = + isPing ? gkInputUbTensor_ping : gkInputUbTensor_pong; + + if (waitGkFromScalar) { + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + waitGkFromScalar = false; + } else { + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + } + if constexpr (std::is_same::value) { + AscendC::DataCopy(gkLastUbTensor, gkLastInput, rowsThisTile); + } else { + AscendC::DataCopy(gkInputUbTensor, gkLastInput, rowsThisTile); + } + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID2 + pingpongFlag); + if constexpr (!std::is_same::value) { + AscendC::Cast(gkLastUbTensor, gkInputUbTensor, AscendC::RoundMode::CAST_NONE, rowsThisTile); + } + AscendC::PipeBarrier(); + AscendC::Muls(gkLastUbTensor, gkLastUbTensor, LN2, rowsThisTile); + AscendC::PipeBarrier(); + AscendC::Exp(gkLastUbTensor, gkLastUbTensor, rowsThisTile); + AscendC::PipeBarrier(); + + uint32_t gkBrcReptime = (rowsThisTile + 8 - 1) / 8; + uint32_t dstShapeGk[2] = {gkBrcReptime * 8, nActual}; + uint32_t srcShapeGk[2] = {gkBrcReptime * 8, 1}; + AscendC::Broadcast(gkBroadcastUbTensor, gkLastUbTensor, dstShapeGk, srcShapeGk, + shareBufferGk_); + AscendC::PipeBarrier(); + AscendC::Mul(calcUbTensor, calcUbTensor, gkBroadcastUbTensor, rowsThisTile * nActual); + AscendC::PipeBarrier(); + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + } + + if (waitUpdateFromMte3) { + AscendC::WaitFlag(updateReadyEvent); + } else { + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + } + AscendC::DataCopy(hUpdateUbTensor, hUpdateInputThisTile, rowsThisTile * nActual); + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + AscendC::Add(hUpdateUbTensor, calcUbTensor, hUpdateUbTensor, rowsThisTile * nActual); + AscendC::PipeBarrier(); + if (storeFinalState && isFinalState && std::is_same::value) { + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + waitHFromV = true; + } else { + waitHFromV = false; + } + + if constexpr (std::is_same::value) { + if (storeFinalState) { + if (!isFinalState) { + AscendC::Cast(hUbTensor, hUpdateUbTensor, AscendC::RoundMode::CAST_RINT, + rowsThisTile * nActual); + AscendC::PipeBarrier(); + } + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + CopyUbToGm(finalStateThisTile, hUpdateUbTensor, rowsThisTile, nActual, outputStride); + // A2 wide-V tiles reuse this UB in the next iteration; retire MTE3 before either writer advances. + AscendC::PipeBarrier(); + AscendC::SetFlag(updateReadyEvent); + AscendC::SetFlag(updateReadyEvent); + AscendC::WaitFlag(updateReadyEvent); + waitUpdateFromMte3 = true; + if (!isFinalState) { + CopyUbToGm(hOutputThisTile, hUbTensor, rowsThisTile, nActual, outputStride); + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + } + } else { + AscendC::Cast(hUbTensor, hUpdateUbTensor, AscendC::RoundMode::CAST_RINT, rowsThisTile * nActual); + AscendC::PipeBarrier(); + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID2 + pingpongFlag); + CopyUbToGm(hOutputThisTile, hUbTensor, rowsThisTile, nActual, outputStride); + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + waitUpdateFromMte3 = false; + } + } else { + if (storeFinalState && isFinalState) { + AscendC::Cast(finalOutputUbTensor, hUpdateUbTensor, AscendC::RoundMode::CAST_RINT, + rowsThisTile * nActual); + AscendC::PipeBarrier(); + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID2 + pingpongFlag); + CopyUbToGm(finalStateThisTile, finalOutputUbTensor, rowsThisTile, nActual, outputStride); + } else { + AscendC::Cast(hUbTensor, hUpdateUbTensor, AscendC::RoundMode::CAST_RINT, rowsThisTile * nActual); + AscendC::PipeBarrier(); + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID2 + pingpongFlag); + CopyUbToGm(hOutputThisTile, hUbTensor, rowsThisTile, nActual, outputStride); + } + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + waitUpdateFromMte3 = false; + } + } + if (storeFinalState && std::is_same::value) { + AscendC::WaitFlag(updateReadyEvent); + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + if (!isFinalState) { + AscendC::WaitFlag(EVENT_ID2 + pingpongFlag); + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + } + } else { + AscendC::WaitFlag(EVENT_ID2 + pingpongFlag); + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + } + if constexpr (kGated) { + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + } + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + } + +private: + uint32_t pongBaseEvent = 4; + + AscendC::LocalTensor calcUbTensor; + + AscendC::LocalTensor hUpdateUbTensor_ping; + AscendC::LocalTensor gkBroadcastUbTensor_ping; + AscendC::LocalTensor hUbTensor_ping; + AscendC::LocalTensor finalOutputUbTensor_ping; + AscendC::LocalTensor glastUbTensor_ping; + + AscendC::LocalTensor hUpdateUbTensor_pong; + AscendC::LocalTensor gkBroadcastUbTensor_pong; + AscendC::LocalTensor hUbTensor_pong; + AscendC::LocalTensor finalOutputUbTensor_pong; + AscendC::LocalTensor glastUbTensor_pong; + + AscendC::LocalTensor gkLastUbTensor_ping; + AscendC::LocalTensor gkLastUbTensor_pong; + AscendC::LocalTensor gkInputUbTensor_ping; + AscendC::LocalTensor gkInputUbTensor_pong; + AscendC::LocalTensor shareBufferGk_; +}; +} // namespace Catlass::Epilogue::Block + +#endif diff --git a/csrc/chunk_kda_fwd/op_kernel/arch22/fwd_h/epilogue/block/block_epilogue_kda_fwdh_vnew_arch22.hpp b/csrc/chunk_kda_fwd/op_kernel/arch22/fwd_h/epilogue/block/block_epilogue_kda_fwdh_vnew_arch22.hpp new file mode 100644 index 0000000000..f857cc780d --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/arch22/fwd_h/epilogue/block/block_epilogue_kda_fwdh_vnew_arch22.hpp @@ -0,0 +1,396 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_KDA_FWDH_VNEW_ARCH22_HPP +#define CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_KDA_FWDH_VNEW_ARCH22_HPP +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "../kda_fwd_h_epilogue_policies.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/kda_matrix_coord.hpp" +#include "catlass/epilogue/tile/kda_epilogue_tile_copy.hpp" + +namespace Catlass::Epilogue::Block { + +template +class BlockEpilogue { + static constexpr bool scalarGated = KGatedTag::scalarGated; + static constexpr bool useExp2 = KGatedTag::useExp2; + static constexpr float LN2 = 0.6931471805599453f; + +public: + // Type aliases + using DispatchPolicy = EpilogueAtlasKdaFwdHVnew; + using ArchTag = typename DispatchPolicy::ArchTag; + + using VElementOutput = typename VOutputType_::Element; + using GElementInput = typename GInputType_::Element; + using UElementInput = typename UInputType_::Element; + using WSElementInput = typename WSInputType_::Element; + using FinalStateElement = typename FinalStateType_::Element; + + CATLASS_DEVICE + BlockEpilogue(Arch::Resource &resource) + { + constexpr uint32_t CALC_BUF_OFFSET = 0; + constexpr uint32_t PING_BUF_0_OFFSET = 32 * 1024; + constexpr uint32_t PING_BUF_1_OFFSET = 48 * 1024; + constexpr uint32_t PING_BUF_2_OFFSET = 64 * 1024; + constexpr uint32_t PING_BUF_3_OFFSET = 80 * 1024; + constexpr uint32_t PONG_BUF_0_OFFSET = 96 * 1024; + constexpr uint32_t PONG_BUF_1_OFFSET = 112 * 1024; + constexpr uint32_t PONG_BUF_2_OFFSET = 128 * 1024; + constexpr uint32_t PONG_BUF_3_OFFSET = 144 * 1024; + constexpr uint32_t PING_G_BUF_OFFSET = 160 * 1024; + constexpr uint32_t PONG_G_BUF_OFFSET = 161 * 1024; + constexpr uint32_t PING_G_SUB_BUF_OFFSET = 162 * 1024; + constexpr uint32_t PONG_G_SUB_BUF_OFFSET = 163 * 1024; + constexpr uint32_t PING_G_INPUT_BUF_OFFSET = 164 * 1024; + constexpr uint32_t PONG_G_INPUT_BUF_OFFSET = 165 * 1024; + constexpr uint32_t SHARE_BUF_OFFSET = 166 * 1024; + + calcUbTensor = resource.ubBuf.template GetBufferByByte(CALC_BUF_OFFSET); + + uUbTensor_ping = resource.ubBuf.template GetBufferByByte(PING_BUF_2_OFFSET); + wsUbTensor_ping = resource.ubBuf.template GetBufferByByte(PING_BUF_0_OFFSET); + gUbTensor_ping = resource.ubBuf.template GetBufferByByte(PING_G_BUF_OFFSET); + gLastUbTensor_ping = resource.ubBuf.template GetBufferByByte(PING_G_SUB_BUF_OFFSET); + gInputUbTensor_ping = resource.ubBuf.template GetBufferByByte(PING_G_SUB_BUF_OFFSET); + vNewOutputUbTensor_ping = resource.ubBuf.template GetBufferByByte(PING_BUF_2_OFFSET); + vNewDecayUbTensor_ping = resource.ubBuf.template GetBufferByByte(PING_BUF_2_OFFSET); + + uUbTensor_pong = resource.ubBuf.template GetBufferByByte(PONG_BUF_2_OFFSET); + wsUbTensor_pong = resource.ubBuf.template GetBufferByByte(PONG_BUF_0_OFFSET); + gUbTensor_pong = resource.ubBuf.template GetBufferByByte(PONG_G_BUF_OFFSET); + gLastUbTensor_pong = resource.ubBuf.template GetBufferByByte(PONG_G_SUB_BUF_OFFSET); + gInputUbTensor_pong = resource.ubBuf.template GetBufferByByte(PONG_G_SUB_BUF_OFFSET); + vNewOutputUbTensor_pong = resource.ubBuf.template GetBufferByByte(PONG_BUF_2_OFFSET); + vNewDecayUbTensor_pong = resource.ubBuf.template GetBufferByByte(PONG_BUF_2_OFFSET); + + shareBuffer_ = resource.ubBuf.template GetBufferByByte(SHARE_BUF_OFFSET); + } + + CATLASS_DEVICE + ~BlockEpilogue() {} + + template + CATLASS_DEVICE void CopyGmToUb(AscendC::LocalTensor dst, AscendC::GlobalTensor src, uint32_t rows, + uint32_t cols, uint32_t srcStride) + { + if (cols == srcStride) { + AscendC::DataCopy(dst, src, rows * cols); + return; + } + AscendC::DataCopyExtParams copyParams{static_cast(rows), + static_cast(cols * sizeof(Element)), + static_cast((srcStride - cols) * sizeof(Element)), 0, 0}; + AscendC::DataCopyPadExtParams padParams{false, 0, 0, 0}; + AscendC::DataCopyPad(dst, src, copyParams, padParams); + } + + template + CATLASS_DEVICE void CopyUbToGm(AscendC::GlobalTensor dst, AscendC::LocalTensor src, uint32_t rows, + uint32_t cols, uint32_t dstStride) + { + if (cols == dstStride) { + AscendC::DataCopy(dst, src, rows * cols); + return; + } + AscendC::DataCopyExtParams copyParams{static_cast(rows), + static_cast(cols * sizeof(Element)), 0, + static_cast((dstStride - cols) * sizeof(Element)), 0}; + AscendC::DataCopyPad(dst, src, copyParams); + } + + CATLASS_DEVICE + void PrepareG(AscendC::LocalTensor gUbTensor, AscendC::LocalTensor gLastUbTensor, + AscendC::LocalTensor gInputUbTensor, + AscendC::GlobalTensor gInputThisSubBlock, uint32_t mActual, uint32_t pingpongFlag) + { + AscendC::WaitFlag(EVENT_ID3 + pingpongFlag); + if constexpr (!scalarGated) { + AscendC::Duplicate(gUbTensor, 1.0f, mActual); + AscendC::PipeBarrier(); + return; + } + if (mActual == 1) { + AscendC::Duplicate(gUbTensor, 1.0f, 1); + AscendC::PipeBarrier(); + return; + } + if constexpr (std::is_same::value) { + AscendC::DataCopyParams gUbParams{1, (uint16_t)(mActual * sizeof(float)), 0, 0}; + AscendC::DataCopyPadParams gUbPadParams{false, 0, 0, 0}; + AscendC::DataCopyPad(gUbTensor, gInputThisSubBlock, gUbParams, gUbPadParams); + } else { + AscendC::DataCopyParams gUbParams{1, (uint16_t)(mActual * sizeof(GElementInput)), 0, 0}; + AscendC::DataCopyPadParams gUbPadParams{false, 0, 0, 0}; + AscendC::DataCopyPad(gInputUbTensor, gInputThisSubBlock, gUbParams, gUbPadParams); + } + AscendC::SetFlag(EVENT_ID3 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID3 + pingpongFlag); + if constexpr (!std::is_same::value) { + AscendC::Cast(gUbTensor, gInputUbTensor, AscendC::RoundMode::CAST_NONE, mActual); + } + + AscendC::SetFlag(EVENT_ID3 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID3 + pingpongFlag); + float inputVal = gUbTensor.GetValue(mActual - 1); + AscendC::SetFlag(EVENT_ID3 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID3 + pingpongFlag); + + AscendC::PipeBarrier(); + AscendC::Duplicate(gLastUbTensor, inputVal, mActual); + AscendC::PipeBarrier(); + + AscendC::Sub(gUbTensor, gLastUbTensor, gUbTensor, mActual); + AscendC::PipeBarrier(); + if constexpr (useExp2) { + AscendC::Muls(gUbTensor, gUbTensor, LN2, mActual); + AscendC::PipeBarrier(); + } + AscendC::Exp(gUbTensor, gUbTensor, mActual); + AscendC::PipeBarrier(); + } + + CATLASS_DEVICE + void operator()(AscendC::GlobalTensor vnewOutput, + AscendC::GlobalTensor vnewdecayOutput, AscendC::GlobalTensor gInput, + AscendC::GlobalTensor uInput, AscendC::GlobalTensor wsInput, + uint32_t chunkSize, uint32_t vBlockDim, uint32_t vHeadDim, Arch::CrossCoreFlag cube1Done, + Arch::CrossCoreFlag vec1Done, bool isInitialState, bool isFinalState, bool storeFinalState, + bool waitWsFromMte3, bool isPing) + { + static constexpr uint32_t ROW_TILE = 16; + uint32_t mActual = chunkSize; + uint32_t nvActual = vBlockDim; + uint32_t inputStride = vHeadDim; + + uint32_t subBlockIdx = AscendC::GetSubBlockIdx(); + uint32_t subBlockNum = AscendC::GetSubBlockNum(); + uint32_t rowsPerSubBlock = CeilDiv(mActual, subBlockNum); + uint32_t rowBegin = subBlockIdx * rowsPerSubBlock; + uint32_t rowEnd = rowBegin + rowsPerSubBlock; + if (rowEnd > mActual) { + rowEnd = mActual; + } + uint32_t pingpongFlag = isPing ? 0 : pongBaseEvent; + if (rowBegin >= mActual) { + Arch::CrossCoreWaitFlag(cube1Done); + // A zero-row AIV lane still owns the EVENT0 hand-off consumed by V2. + if (waitWsFromMte3) { + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + } + Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(vec1Done); + return; + } + AscendC::ResetMask(); + + AscendC::GlobalTensor gInputThisSubBlock = gInput; + + AscendC::LocalTensor uUbTensor = isPing ? uUbTensor_ping : uUbTensor_pong; + AscendC::LocalTensor wsUbTensor = isPing ? wsUbTensor_ping : wsUbTensor_pong; + AscendC::LocalTensor gUbTensor = isPing ? gUbTensor_ping : gUbTensor_pong; + AscendC::LocalTensor gLastUbTensor = isPing ? gLastUbTensor_ping : gLastUbTensor_pong; + AscendC::LocalTensor gInputUbTensor = isPing ? gInputUbTensor_ping : gInputUbTensor_pong; + AscendC::LocalTensor vNewOutputUbTensor = + isPing ? vNewOutputUbTensor_ping : vNewOutputUbTensor_pong; + AscendC::LocalTensor vNewDecayUbTensor = + isPing ? vNewDecayUbTensor_ping : vNewDecayUbTensor_pong; + + if (nvActual <= 128 && nvActual == inputStride) { + uint32_t mActualThisSubBlock = rowEnd - rowBegin; + uint32_t gbrcRealStart = rowBegin & ~7; + uint32_t gbrcEffStart = rowBegin - gbrcRealStart; + uint32_t gbrcRealProcess = gbrcEffStart + mActualThisSubBlock; + uint32_t dstShape_[2] = {gbrcRealProcess, nvActual}; + uint32_t srcShape_[2] = {gbrcRealProcess, 1}; + + AscendC::GlobalTensor vnewOutputThisSubBlock = vnewOutput[rowBegin * inputStride]; + AscendC::GlobalTensor vnewdecayOutputThisSubBlock = vnewdecayOutput[rowBegin * nvActual]; + AscendC::GlobalTensor uInputThisSubBlock = uInput[rowBegin * inputStride]; + AscendC::GlobalTensor wsInputThisSubBlock = wsInput[rowBegin * nvActual]; + + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + AscendC::DataCopy(uUbTensor, uInputThisSubBlock, mActualThisSubBlock * nvActual); + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + AscendC::Cast(calcUbTensor, uUbTensor, AscendC::RoundMode::CAST_NONE, mActualThisSubBlock * nvActual); + AscendC::PipeBarrier(); + + if constexpr (scalarGated) { + PrepareG(gUbTensor, gLastUbTensor, gInputUbTensor, gInputThisSubBlock, mActual, pingpongFlag); + } else { + AscendC::WaitFlag(EVENT_ID3 + pingpongFlag); + } + Arch::CrossCoreWaitFlag(cube1Done); + + if (waitWsFromMte3) { + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + } else { + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + } + AscendC::DataCopy(wsUbTensor, wsInputThisSubBlock, mActualThisSubBlock * nvActual); + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + + AscendC::Sub(wsUbTensor, calcUbTensor, wsUbTensor, mActualThisSubBlock * nvActual); + AscendC::PipeBarrier(); + + uint32_t decayOffset = 0; + if constexpr (scalarGated) { + AscendC::Broadcast(calcUbTensor, gUbTensor[gbrcRealStart], dstShape_, srcShape_, + shareBuffer_); + AscendC::PipeBarrier(); + AscendC::Mul(calcUbTensor[gbrcEffStart * nvActual], wsUbTensor, calcUbTensor[gbrcEffStart * nvActual], + mActualThisSubBlock * nvActual); + decayOffset = gbrcEffStart * nvActual; + } else { + AscendC::Adds(calcUbTensor, wsUbTensor, 0.0f, mActualThisSubBlock * nvActual); + } + AscendC::PipeBarrier(); + AscendC::SetFlag(EVENT_ID3 + pingpongFlag); + + AscendC::Cast(vNewDecayUbTensor, calcUbTensor[decayOffset], AscendC::RoundMode::CAST_RINT, + mActualThisSubBlock * nvActual); + AscendC::PipeBarrier(); + + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + AscendC::DataCopy(vnewdecayOutputThisSubBlock, vNewDecayUbTensor, mActualThisSubBlock * nvActual); + + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + AscendC::Cast(vNewOutputUbTensor, wsUbTensor, AscendC::RoundMode::CAST_RINT, + mActualThisSubBlock * nvActual); + AscendC::PipeBarrier(); + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + AscendC::DataCopy(vnewOutputThisSubBlock, vNewOutputUbTensor, mActualThisSubBlock * nvActual); + Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(vec1Done); + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + + return; + } + + if constexpr (scalarGated) { + PrepareG(gUbTensor, gLastUbTensor, gInputUbTensor, gInputThisSubBlock, mActual, pingpongFlag); + } else { + AscendC::WaitFlag(EVENT_ID3 + pingpongFlag); + } + Arch::CrossCoreWaitFlag(cube1Done); + + bool waitWsThisTileFromMte3 = waitWsFromMte3; + for (uint32_t rowStart = rowBegin; rowStart < rowEnd;) { + uint32_t alignExtra = rowStart & 7; + uint32_t maxRowsThisTile = ROW_TILE - alignExtra; + uint32_t rowsThisTile = rowEnd - rowStart; + if (rowsThisTile > maxRowsThisTile) { + rowsThisTile = maxRowsThisTile; + } + uint32_t gbrcRealStart = rowStart & ~7; + uint32_t gbrcRealProcess = alignExtra + rowsThisTile; + uint32_t dstShape_[2] = {gbrcRealProcess, nvActual}; + uint32_t srcShape_[2] = {gbrcRealProcess, 1}; + + AscendC::GlobalTensor vnewOutputThisTile = vnewOutput[rowStart * inputStride]; + AscendC::GlobalTensor vnewdecayOutputThisTile = vnewdecayOutput[rowStart * nvActual]; + AscendC::GlobalTensor uInputThisTile = uInput[rowStart * inputStride]; + AscendC::GlobalTensor wsInputThisTile = wsInput[rowStart * nvActual]; + + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + CopyGmToUb(uUbTensor, uInputThisTile, rowsThisTile, nvActual, inputStride); + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + AscendC::Cast(calcUbTensor, uUbTensor, AscendC::RoundMode::CAST_NONE, rowsThisTile * nvActual); + AscendC::PipeBarrier(); + + if (waitWsThisTileFromMte3) { + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + } else { + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + } + AscendC::DataCopy(wsUbTensor, wsInputThisTile, rowsThisTile * nvActual); + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + waitWsThisTileFromMte3 = false; + + AscendC::Sub(wsUbTensor, calcUbTensor, wsUbTensor, rowsThisTile * nvActual); + AscendC::PipeBarrier(); + + uint32_t decayOffset = 0; + if constexpr (scalarGated) { + AscendC::Broadcast(calcUbTensor, gUbTensor[gbrcRealStart], dstShape_, srcShape_, + shareBuffer_); + AscendC::PipeBarrier(); + AscendC::Mul(calcUbTensor[alignExtra * nvActual], wsUbTensor, calcUbTensor[alignExtra * nvActual], + rowsThisTile * nvActual); + decayOffset = alignExtra * nvActual; + } else { + AscendC::Adds(calcUbTensor, wsUbTensor, 0.0f, rowsThisTile * nvActual); + } + AscendC::PipeBarrier(); + + AscendC::Cast(vNewDecayUbTensor, calcUbTensor[decayOffset], AscendC::RoundMode::CAST_RINT, + rowsThisTile * nvActual); + AscendC::PipeBarrier(); + + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + AscendC::DataCopy(vnewdecayOutputThisTile, vNewDecayUbTensor, rowsThisTile * nvActual); + + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + AscendC::Cast(vNewOutputUbTensor, wsUbTensor, AscendC::RoundMode::CAST_RINT, rowsThisTile * nvActual); + AscendC::PipeBarrier(); + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + CopyUbToGm(vnewOutputThisTile, vNewOutputUbTensor, rowsThisTile, nvActual, inputStride); + if (rowStart + rowsThisTile >= rowEnd) { + Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(vec1Done); + } + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + rowStart += rowsThisTile; + } + + AscendC::SetFlag(EVENT_ID3 + pingpongFlag); + } + +private: + uint32_t pongBaseEvent = 4; + + AscendC::LocalTensor calcUbTensor; + + AscendC::LocalTensor uUbTensor_ping; + AscendC::LocalTensor wsUbTensor_ping; + AscendC::LocalTensor gUbTensor_ping; + AscendC::LocalTensor gLastUbTensor_ping; + AscendC::LocalTensor gInputUbTensor_ping; + AscendC::LocalTensor vNewOutputUbTensor_ping; + AscendC::LocalTensor vNewDecayUbTensor_ping; + + AscendC::LocalTensor uUbTensor_pong; + AscendC::LocalTensor wsUbTensor_pong; + AscendC::LocalTensor gUbTensor_pong; + AscendC::LocalTensor gLastUbTensor_pong; + AscendC::LocalTensor gInputUbTensor_pong; + AscendC::LocalTensor vNewOutputUbTensor_pong; + AscendC::LocalTensor vNewDecayUbTensor_pong; + + AscendC::LocalTensor shareBuffer_; +}; +} // namespace Catlass::Epilogue::Block + +#endif diff --git a/csrc/chunk_kda_fwd/op_kernel/arch22/fwd_h/epilogue/kda_fwd_h_epilogue_policies.hpp b/csrc/chunk_kda_fwd/op_kernel/arch22/fwd_h/epilogue/kda_fwd_h_epilogue_policies.hpp new file mode 100644 index 0000000000..1dd37a521c --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/arch22/fwd_h/epilogue/kda_fwd_h_epilogue_policies.hpp @@ -0,0 +1,28 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_EPILOGUE_KDA_FWD_H_EPILOGUE_POLICIES_HPP +#define CATLASS_EPILOGUE_KDA_FWD_H_EPILOGUE_POLICIES_HPP + +#include "catlass/kda_catlass.hpp" + +namespace Catlass::Epilogue { + +struct EpilogueAtlasKdaFwdHVnew { + using ArchTag = Arch::AtlasA2; +}; + +struct EpilogueAtlasKdaFwdHUpdate { + using ArchTag = Arch::AtlasA2; +}; + +} // namespace Catlass::Epilogue + +#endif // CATLASS_EPILOGUE_KDA_FWD_H_EPILOGUE_POLICIES_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/arch22/fwd_h/gemm/block/block_scheduler_kda_fwd_h_arch22.hpp b/csrc/chunk_kda_fwd/op_kernel/arch22/fwd_h/gemm/block/block_scheduler_kda_fwd_h_arch22.hpp new file mode 100644 index 0000000000..961ce442bc --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/arch22/fwd_h/gemm/block/block_scheduler_kda_fwd_h_arch22.hpp @@ -0,0 +1,445 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#include "catlass/kda_gemm_coord.hpp" +#include "../../../../chunk_kda_fwd_h_runtime_tiling.h" +using namespace Catlass; + +#ifndef CATLASS_GEMM_SCHEDULER_KDA_FWD_H_HPP +#define CATLASS_GEMM_SCHEDULER_KDA_FWD_H_HPP + +// constexpr uint32_t PING_PONG_STAGES = 1; +constexpr uint32_t PING_PONG_STAGES = 2; +constexpr uint32_t BYTE_SIZE_16_BIT = 2; +constexpr uint32_t BYTES_PER_C0 = 32; +constexpr uint32_t BYTE_SIZE_PER_REPEAT = 256; +constexpr uint32_t SIZE_16_NUM_PER_C0 = BYTES_PER_C0 / BYTE_SIZE_16_BIT; +constexpr uint32_t FLOAT_NUM_PER_REPEAT = BYTE_SIZE_PER_REPEAT / sizeof(float); +constexpr uint32_t NZ_BLOCK_SIZE = 16; + +template +CATLASS_DEVICE T AlignUp(T a, T b) +{ + return (b == 0) ? 0 : (a + b - 1) / b * b; +} + +template +CATLASS_DEVICE T Min(T a, T b) +{ + return (a > b) ? b : a; +} + +template +CATLASS_DEVICE T Max(T a, T b) +{ + return (a > b) ? a : b; +} + +namespace Catlass::Gemm::Block { + +struct KdaFwdHOffsets { + uint32_t hSrcOffset; + uint32_t hDstOffset; + uint32_t uvOffset; + uint32_t wkOffset; + uint32_t wOffset; + uint32_t gOffset; + uint32_t gkOffset; + uint32_t hWorkOffset; + uint32_t vWorkOffset; + uint32_t kDecayWorkOffset; + uint32_t vBlockOffset; + uint32_t vBlockDim; + uint32_t initialStateOffset; + uint32_t finalStateOffset; + bool isInitialState; + bool isFinalState; + uint32_t blockTokens; + uint32_t streamId; + // for debug + uint32_t batchIdx; + uint32_t headIdx; + uint32_t chunkIdx; +}; + +struct KdaFwdHStream { + uint32_t batchIdx; + uint32_t chunkIdx{0}; + uint32_t vHeadIdx; + uint32_t kHeadIdx; + uint32_t shapeBatchIdx; + uint32_t tokenBatchIdx; + + uint32_t chunkOffset; + uint32_t tokenOffset; + uint32_t batchChunks{0}; + uint32_t batchTokens; + uint32_t nextTaskIdx{0}; + bool active{false}; + + KdaFwdHOffsets offset; +}; + +struct KdaFwdHRunningQ { + KdaFwdHStream streams[PING_PONG_STAGES]; +}; + +struct BlockSchedulerKdaFwdH { + uint32_t batch; + uint32_t seqlen; + uint32_t kNumHead; + uint32_t vNumHead; + uint32_t kHeadDim; + uint32_t vHeadDim; + uint32_t chunkSize; + uint32_t vBlockSize{128}; + uint32_t isVariedLen; + uint32_t shapeBatch; + uint32_t tokenBatch; + uint32_t inputTokenBatch; + bool useInitialState; + bool storeFinalState; + uint32_t numSeqWorkspaceOffset; + uint32_t numChunksWorkspaceOffset; + + uint32_t taskIdx; + uint32_t taskStride; + uint32_t cubeCoreIdx; + uint32_t cubeCoreNum; + uint32_t taskNum; + uint32_t headGroups; + uint32_t totalChunks; + uint32_t totalTokens; + + KdaFwdHRunningQ runningQ; + + bool isRunning; + + AscendC::GlobalTensor gmSeqlen; + AscendC::GlobalTensor gmNumSeq; + AscendC::GlobalTensor gmNumChunks; + + Arch::CrossCoreFlag cube1Done[PING_PONG_STAGES] = {0, 1}; + Arch::CrossCoreFlag vec1Done[PING_PONG_STAGES] = {2, 3}; + Arch::CrossCoreFlag cube2Done[PING_PONG_STAGES] = {4, 5}; + Arch::CrossCoreFlag vec2Done[PING_PONG_STAGES] = {6, 7}; + + CATLASS_DEVICE + BlockSchedulerKdaFwdH() {} + + CATLASS_DEVICE + void Init(GM_ADDR cu_seqlens, GM_ADDR chunk_indices, GM_ADDR tiling, GM_ADDR user, uint32_t coreIdx, + uint32_t coreNum) + { + __gm__ ChunkKdaFwdHRuntimeTiling *__restrict kdaFwdHTilingData = + reinterpret_cast<__gm__ ChunkKdaFwdHRuntimeTiling *__restrict>(tiling); + + batch = kdaFwdHTilingData->batch; + seqlen = kdaFwdHTilingData->seqlen; + kNumHead = kdaFwdHTilingData->kNumHead; + vNumHead = kdaFwdHTilingData->vNumHead; + kHeadDim = kdaFwdHTilingData->kHeadDim; + vHeadDim = kdaFwdHTilingData->vHeadDim; + chunkSize = kdaFwdHTilingData->chunkSize; + isVariedLen = kdaFwdHTilingData->isVariedLen; + shapeBatch = kdaFwdHTilingData->shapeBatch; + tokenBatch = kdaFwdHTilingData->tokenBatch; + useInitialState = kdaFwdHTilingData->useInitialState; + storeFinalState = kdaFwdHTilingData->storeFinalState; + numSeqWorkspaceOffset = kdaFwdHTilingData->numSeqWorkspaceOffset; + numChunksWorkspaceOffset = kdaFwdHTilingData->numChunksWorkspaceOffset; + + InitRuntime(cu_seqlens, chunk_indices, user, coreIdx, coreNum); + } + + template + CATLASS_DEVICE void InitFromData(GM_ADDR cu_seqlens, GM_ADDR chunk_indices, const TilingData &tilingData, + GM_ADDR user, uint32_t coreIdx, uint32_t coreNum) + { + batch = tilingData.batch; + seqlen = tilingData.seqlen; + kNumHead = tilingData.kNumHead; + vNumHead = tilingData.vNumHead; + kHeadDim = tilingData.kHeadDim; + vHeadDim = tilingData.vHeadDim; + chunkSize = tilingData.chunkSize; + isVariedLen = tilingData.isVariedLen; + shapeBatch = tilingData.shapeBatch; + tokenBatch = tilingData.tokenBatch; + useInitialState = tilingData.useInitialState; + storeFinalState = tilingData.storeFinalState; + numSeqWorkspaceOffset = tilingData.numSeqWorkspaceOffset; + numChunksWorkspaceOffset = tilingData.numChunksWorkspaceOffset; + + InitRuntime(cu_seqlens, chunk_indices, user, coreIdx, coreNum); + } + + CATLASS_DEVICE + void InitRuntime(GM_ADDR cu_seqlens, GM_ADDR chunk_indices, GM_ADDR user, uint32_t coreIdx, uint32_t coreNum) + { + gmSeqlen.SetGlobalBuffer((__gm__ int64_t *)cu_seqlens); + gmNumSeq.SetGlobalBuffer((__gm__ int64_t *)(user + numSeqWorkspaceOffset)); + gmNumChunks.SetGlobalBuffer((__gm__ int64_t *)(user + numChunksWorkspaceOffset)); + + if (isVariedLen) { + inputTokenBatch = tokenBatch; + uint32_t actualBatch = 0; + int64_t chunkPrefix = 0; + int64_t prevSeq = 0, currSeq; + for (uint32_t b = 1; b <= inputTokenBatch; b++) { + currSeq = gmSeqlen.GetValue(b); + int64_t batchSeqLen = currSeq - prevSeq; + if (batchSeqLen > 0) { + actualBatch++; + int64_t batchChunk = (batchSeqLen + chunkSize - 1) / chunkSize; + chunkPrefix += batchChunk; + } + prevSeq = currSeq; + } + tokenBatch = actualBatch; + batch = actualBatch; + totalChunks = chunkPrefix; + totalTokens = prevSeq; + } else { + totalChunks = (seqlen + chunkSize - 1) / chunkSize; + totalTokens = seqlen; + } + + cubeCoreIdx = coreIdx; + cubeCoreNum = coreNum; + vBlockSize = vHeadDim; + taskNum = batch * vNumHead; + headGroups = vNumHead / kNumHead; + InitTaskWave(0); + } + + CATLASS_DEVICE + void InitTaskWave(uint32_t waveIdx) + { + uint32_t firstTaskIdx = waveIdx * cubeCoreNum + cubeCoreIdx; + taskStride = taskNum; + for (uint32_t streamId = 0; streamId < PING_PONG_STAGES; ++streamId) { + auto &stream = runningQ.streams[streamId]; + stream.nextTaskIdx = streamId == 0 ? firstTaskIdx : taskNum; + stream.chunkIdx = 0; + stream.batchChunks = 0; + stream.active = false; + } + isRunning = firstTaskIdx < taskNum; + } + + CATLASS_DEVICE + uint32_t GetTaskWaveCount() const { return CeilDiv(taskNum, cubeCoreNum); } + + CATLASS_DEVICE + void ResolveVarlenSequence(uint32_t compactBatchIdx, KdaFwdHStream &stream) + { + uint32_t actualBatch = 0; + int64_t chunkPrefix = 0; + int64_t prevSeq = 0; + for (uint32_t b = 1; b <= inputTokenBatch; ++b) { + int64_t currSeq = gmSeqlen.GetValue(b); + int64_t batchTokens = currSeq - prevSeq; + if (batchTokens > 0) { + int64_t batchChunks = (batchTokens + chunkSize - 1) / chunkSize; + if (actualBatch == compactBatchIdx) { + stream.chunkOffset = static_cast(chunkPrefix); + stream.batchChunks = static_cast(batchChunks); + stream.tokenOffset = static_cast(prevSeq); + stream.batchTokens = static_cast(batchTokens); + return; + } + ++actualBatch; + chunkPrefix += batchChunks; + } + prevSeq = currSeq; + } + stream.chunkOffset = 0; + stream.batchChunks = 0; + stream.tokenOffset = 0; + stream.batchTokens = 0; + } + + CATLASS_DEVICE + uint32_t GetVarlenChunkOffset(uint32_t compactBatchIdx) + { + KdaFwdHStream stream; + ResolveVarlenSequence(compactBatchIdx, stream); + return stream.chunkOffset; + } + + CATLASS_DEVICE + void InitNewStream(KdaFwdHStream &newStream) + { + newStream.batchIdx = taskIdx / vNumHead; + newStream.vHeadIdx = taskIdx % vNumHead; + newStream.kHeadIdx = newStream.vHeadIdx / headGroups; + newStream.shapeBatchIdx = isVariedLen ? 0 : newStream.batchIdx; + newStream.tokenBatchIdx = isVariedLen ? newStream.batchIdx : 0; + if (isVariedLen) { + ResolveVarlenSequence(newStream.tokenBatchIdx, newStream); + } else { + newStream.chunkOffset = 0; + newStream.batchChunks = totalChunks; + newStream.tokenOffset = 0; + newStream.batchTokens = totalTokens; + } + newStream.chunkIdx = 0; + } + + CATLASS_DEVICE + void AssignNextStream(uint32_t streamId) + { + auto &stream = runningQ.streams[streamId]; + taskIdx = stream.nextTaskIdx; + if (taskIdx >= taskNum) { + stream.active = false; + stream.batchChunks = 0; + return; + } + + stream.nextTaskIdx += taskStride; + InitNewStream(stream); + stream.active = stream.batchChunks > 0; + if (stream.active) { + UpdateTask(streamId); + } + } + + CATLASS_DEVICE + void UpdateTask(uint32_t streamId) + { + auto &stream = runningQ.streams[streamId]; + auto &offset = stream.offset; + + offset.isInitialState = stream.chunkIdx == 0; + offset.isFinalState = stream.chunkIdx == (stream.batchChunks - 1); + uint32_t vBlockOffset = 0; + uint32_t vBlockDim = vBlockSize; + offset.initialStateOffset = (stream.batchIdx * vNumHead + stream.vHeadIdx) * kHeadDim * vHeadDim + vBlockOffset; + offset.finalStateOffset = (stream.batchIdx * vNumHead + stream.vHeadIdx) * kHeadDim * vHeadDim + vBlockOffset; + offset.hSrcOffset = (stream.shapeBatchIdx * vNumHead * totalChunks + stream.vHeadIdx * totalChunks + + stream.chunkOffset + stream.chunkIdx) * + kHeadDim * vHeadDim + + vBlockOffset; + offset.hDstOffset = offset.hSrcOffset + kHeadDim * vHeadDim; + if (storeFinalState && offset.isFinalState) { + offset.hDstOffset = offset.hSrcOffset; + } + offset.uvOffset = (stream.shapeBatchIdx * vNumHead * totalTokens + stream.vHeadIdx * totalTokens + + stream.tokenOffset + stream.chunkIdx * chunkSize) * + vHeadDim + + vBlockOffset; + offset.wkOffset = (stream.shapeBatchIdx * kNumHead * totalTokens + stream.kHeadIdx * totalTokens + + stream.tokenOffset + stream.chunkIdx * chunkSize) * + kHeadDim; + offset.wOffset = (stream.shapeBatchIdx * vNumHead * totalTokens + stream.vHeadIdx * totalTokens + + stream.tokenOffset + stream.chunkIdx * chunkSize) * + kHeadDim; + offset.gOffset = stream.shapeBatchIdx * vNumHead * totalTokens + stream.vHeadIdx * totalTokens + + stream.tokenOffset + stream.chunkIdx * chunkSize; + offset.gkOffset = (stream.shapeBatchIdx * vNumHead * totalTokens + stream.vHeadIdx * totalTokens + + stream.tokenOffset + stream.chunkIdx * chunkSize) * + kHeadDim; + offset.hWorkOffset = (cubeCoreIdx * PING_PONG_STAGES + streamId) * kHeadDim * vBlockSize; + offset.vWorkOffset = (cubeCoreIdx * PING_PONG_STAGES + streamId) * chunkSize * vBlockSize; + offset.kDecayWorkOffset = (cubeCoreIdx * PING_PONG_STAGES + streamId) * chunkSize * kHeadDim; + offset.vBlockOffset = vBlockOffset; + offset.vBlockDim = vBlockDim; + offset.blockTokens = offset.isFinalState ? (stream.batchTokens - stream.chunkIdx * chunkSize) : chunkSize; + offset.streamId = streamId; + offset.batchIdx = stream.batchIdx; + offset.headIdx = stream.vHeadIdx; + offset.chunkIdx = stream.chunkIdx; + } + + CATLASS_DEVICE + void InitTasks() + { + isRunning = false; + for (uint32_t streamId = 0; streamId < PING_PONG_STAGES; ++streamId) { + auto &stream = runningQ.streams[streamId]; + if (stream.active) { + stream.chunkIdx += 1; + if (stream.chunkIdx >= stream.batchChunks) { + stream.active = false; + stream.batchChunks = 0; + } + } + if (!stream.active) { + AssignNextStream(streamId); + } else { + UpdateTask(streamId); + } + if (stream.active) { + isRunning = true; + } + } + } + + CATLASS_DEVICE + const KdaFwdHStream &GetStream(uint32_t i) const { return runningQ.streams[i]; } + + CATLASS_DEVICE + uint32_t GetStreamId(uint32_t i) const { return i; } + + CATLASS_DEVICE + const KdaFwdHOffsets &GetCurTaskOffsets(const KdaFwdHStream &stream) const { return stream.offset; } + + CATLASS_DEVICE + bool StreamIsDone(const KdaFwdHStream &stream) const { return !stream.active; } + + CATLASS_DEVICE + bool NeedProcessStage2(const KdaFwdHStream &stream) { return storeFinalState || !stream.offset.isFinalState; } +}; + +struct BlockSchedulerKdaFwdHCube : public BlockSchedulerKdaFwdH { + CATLASS_DEVICE + BlockSchedulerKdaFwdHCube() {} + + CATLASS_DEVICE + void Init(GM_ADDR cu_seqlens, GM_ADDR chunk_indices, GM_ADDR tiling, GM_ADDR user) + { + BlockSchedulerKdaFwdH::Init(cu_seqlens, chunk_indices, tiling, user, AscendC::GetBlockIdx(), + AscendC::GetBlockNum()); + } + + template + CATLASS_DEVICE void InitFromData(GM_ADDR cu_seqlens, GM_ADDR chunk_indices, const TilingData &tilingData, + GM_ADDR user) + { + BlockSchedulerKdaFwdH::InitFromData(cu_seqlens, chunk_indices, tilingData, user, AscendC::GetBlockIdx(), + AscendC::GetBlockNum()); + } +}; + +struct BlockSchedulerKdaFwdHVec : public BlockSchedulerKdaFwdH { + CATLASS_DEVICE + BlockSchedulerKdaFwdHVec() {} + + CATLASS_DEVICE + void Init(GM_ADDR cu_seqlens, GM_ADDR chunk_indices, GM_ADDR tiling, GM_ADDR user) + { + BlockSchedulerKdaFwdH::Init(cu_seqlens, chunk_indices, tiling, user, + AscendC::GetBlockIdx() / AscendC::GetSubBlockNum(), AscendC::GetBlockNum()); + } + + template + CATLASS_DEVICE void InitFromData(GM_ADDR cu_seqlens, GM_ADDR chunk_indices, const TilingData &tilingData, + GM_ADDR user) + { + BlockSchedulerKdaFwdH::InitFromData(cu_seqlens, chunk_indices, tilingData, user, + AscendC::GetBlockIdx() / AscendC::GetSubBlockNum(), AscendC::GetBlockNum()); + } +}; + +} // namespace Catlass::Gemm::Block + +#endif // CATLASS_GEMM_SCHEDULER_KDA_FWD_H_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/arch22/fwd_h/gemm/kernel/kda_fwd_h_kernel_arch22.hpp b/csrc/chunk_kda_fwd/op_kernel/arch22/fwd_h/gemm/kernel/kda_fwd_h_kernel_arch22.hpp new file mode 100644 index 0000000000..bfd8cb9802 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/arch22/fwd_h/gemm/kernel/kda_fwd_h_kernel_arch22.hpp @@ -0,0 +1,752 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#define CATLASS_ARCH 2201 + +#include "catlass/arch/kda_arch.hpp" +#include "catlass/arch/kda_cross_core_sync.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/kda_catlass.hpp" +#include "catlass/kda_debug.hpp" +#include "catlass/epilogue/block/kda_block_epilogue.hpp" +#include "../../epilogue/block/block_epilogue_kda_fwdh_update_arch22.hpp" +#include "../../epilogue/block/block_epilogue_kda_fwdh_vnew_arch22.hpp" +#include "catlass/gemm/block/kda_block_mmad.hpp" +#include "../../../../kernel_utils/block/block_mmad_pingpong_tla_multi.hpp" +#include "catlass/gemm/block/kda_block_swizzle.hpp" +#include "../block/block_scheduler_kda_fwd_h_arch22.hpp" +#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" +#include "catlass/gemm/kda_gemm_type.hpp" +#include "catlass/layout/kda_layout.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "tla/kda_tensor.hpp" +#include "tla/kda_tla_layout.hpp" +#include "tla/kda_tensor.hpp" + +#include "kernel_operator.h" +using namespace Catlass; +using namespace tla; + +namespace Catlass::Gemm::Kernel { + +struct KdaFwdHTileShapes128 { + using L1TileShape = tla::Shape<_128, _128, _128>; + using L0TileShape = L1TileShape; +}; + +struct KdaFwdHTileShapes256 { + using L1TileShape = tla::Shape<_128, _256, _128>; + using L0TileShape = tla::Shape<_128, _256, _64>; +}; + +template +struct KdaFwdHGateTag { + static constexpr bool value = KGated; + static constexpr bool scalarGated = ScalarGated; + static constexpr bool useExp2 = UseExp2; +}; + +template +class KdaFwdHKernel { +public: + using ArchTag = Arch::AtlasA2; + using CubeScheduler = typename Catlass::Gemm::Block::BlockSchedulerKdaFwdHCube; + using VecScheduler = typename Catlass::Gemm::Block::BlockSchedulerKdaFwdHVec; + + using DispatchPolicyTla = Gemm::MmadPingpongTlaMulti; + using L1TileShapeVTla = typename TileShapes::L1TileShape; + using L0TileShapeVTla = typename TileShapes::L0TileShape; + + using WType = Gemm::GemmType; + using HType = Gemm::GemmType; + using VworkType = Gemm::GemmType; + using KType = Gemm::GemmType; + using HworkType = Gemm::GemmType; + using VType = Gemm::GemmType; + using GType = Gemm::GemmType; + using UType = Gemm::GemmType; + using FinalStateType = Gemm::GemmType; + + // cube 1 + using TileCopyWH = Catlass::Gemm::Tile::PackedTileCopyTla; + using BlockMmadWH = Gemm::Block::BlockMmadTla; + + // cube 2 + using TileCopyKV = Catlass::Gemm::Tile::PackedTileCopyTla; + using BlockMmadKV = Gemm::Block::BlockMmadTla; + + // vec 1 + using DispatchPolicyKdaFwdHVnew = Epilogue::EpilogueAtlasKdaFwdHVnew; + using GateTag = KdaFwdHGateTag; + using EpilogueKdaFwdHVnew = Epilogue::Block::BlockEpilogue; + + // vec 2 + using DispatchPolicyKdaFwdHUpdate = Epilogue::EpilogueAtlasKdaFwdHUpdate; + using EpilogueKdaFwdHUpdate = Epilogue::Block::BlockEpilogue; + + using KdaFwdHOffsets = Catlass::Gemm::Block::KdaFwdHOffsets; + + using ElementK = INPUT_TYPE; + using ElementW = INPUT_TYPE; + using ElementU = INPUT_TYPE; + using ElementG = G_TYPE; + using ElementH = INPUT_TYPE; + using ElementV = INPUT_TYPE; + using ElementVWork = WORKSPACE_TYPE; + using ElementHWork = WORKSPACE_TYPE; + using ElementInitialState = STATE_TYPE; + using ElementFinalState = STATE_TYPE; + + using LayoutW = Catlass::layout::RowMajor; + using LayoutH = Catlass::layout::RowMajor; + using LayoutV = Catlass::layout::RowMajor; + using LayoutK = Catlass::layout::ColumnMajor; + + uint32_t batch; + uint32_t seqlen; + uint32_t kNumHead; + uint32_t vNumHead; + uint32_t kHeadDim; + uint32_t vHeadDim; + uint32_t chunkSize; + bool useInitialState; + bool storeFinalState; + uint32_t isVariedLen; + uint32_t shapeBatch; + uint32_t tokenBatch; + uint32_t vWorkspaceOffset; + uint32_t vUpdateWorkspaceOffset; + uint32_t hWorkspaceOffset; + uint32_t numSeqWorkspaceOffset; + uint32_t numChunksWorkspaceOffset; + uint32_t kDecayWorkspaceOffset; + + AscendC::GlobalTensor gmK; + AscendC::GlobalTensor gmW; + AscendC::GlobalTensor gmU; + AscendC::GlobalTensor gmG; + AscendC::GlobalTensor gmInitialState; + AscendC::GlobalTensor gmH; + AscendC::GlobalTensor gmV; + AscendC::GlobalTensor gmFinalState; + AscendC::GlobalTensor gmVWorkspace; + AscendC::GlobalTensor gmVUpdateWorkspace; + AscendC::GlobalTensor gmHWorkspace; + + AscendC::GlobalTensor gmGk; + AscendC::GlobalTensor gmKDecayWorkspace; + + AscendC::GlobalTensor gmSeqlen; + AscendC::GlobalTensor gmNumSeq; + AscendC::GlobalTensor gmNumChunks; + + CubeScheduler cubeBlockScheduler; + VecScheduler vecBlockScheduler; + + Arch::Resource resource; + + __aicore__ inline KdaFwdHKernel() {} + + __aicore__ inline void Init(GM_ADDR k, GM_ADDR w, GM_ADDR u, GM_ADDR g, GM_ADDR gk, GM_ADDR inital_state, + GM_ADDR cu_seqlens, GM_ADDR chunk_indices, GM_ADDR h, GM_ADDR v_new, + GM_ADDR final_state, GM_ADDR tiling, GM_ADDR user) + { + __gm__ ChunkKdaFwdHRuntimeTiling *__restrict kdaFwdHTilingData = + reinterpret_cast<__gm__ ChunkKdaFwdHRuntimeTiling *__restrict>(tiling); + + batch = kdaFwdHTilingData->batch; + seqlen = kdaFwdHTilingData->seqlen; + kNumHead = kdaFwdHTilingData->kNumHead; + vNumHead = kdaFwdHTilingData->vNumHead; + kHeadDim = kdaFwdHTilingData->kHeadDim; + vHeadDim = kdaFwdHTilingData->vHeadDim; + chunkSize = kdaFwdHTilingData->chunkSize; + useInitialState = kdaFwdHTilingData->useInitialState; + storeFinalState = kdaFwdHTilingData->storeFinalState; + isVariedLen = kdaFwdHTilingData->isVariedLen; + shapeBatch = kdaFwdHTilingData->shapeBatch; + tokenBatch = kdaFwdHTilingData->tokenBatch; + vWorkspaceOffset = kdaFwdHTilingData->vWorkspaceOffset; + vUpdateWorkspaceOffset = kdaFwdHTilingData->vUpdateWorkspaceOffset; + hWorkspaceOffset = kdaFwdHTilingData->hWorkspaceOffset; + numSeqWorkspaceOffset = kdaFwdHTilingData->numSeqWorkspaceOffset; + numChunksWorkspaceOffset = kdaFwdHTilingData->numChunksWorkspaceOffset; + kDecayWorkspaceOffset = kdaFwdHTilingData->kDecayWorkspaceOffset; + + gmK.SetGlobalBuffer((__gm__ ElementK *)k); + gmW.SetGlobalBuffer((__gm__ ElementW *)w); + gmU.SetGlobalBuffer((__gm__ ElementU *)u); + gmG.SetGlobalBuffer((__gm__ ElementG *)(scalarGated ? g : gk)); + gmInitialState.SetGlobalBuffer((__gm__ ElementInitialState *)inital_state); + gmH.SetGlobalBuffer((__gm__ ElementH *)h); + gmV.SetGlobalBuffer((__gm__ ElementV *)v_new); + gmFinalState.SetGlobalBuffer((__gm__ ElementFinalState *)final_state); + gmVWorkspace.SetGlobalBuffer((__gm__ ElementVWork *)(user + vWorkspaceOffset)); + gmVUpdateWorkspace.SetGlobalBuffer((__gm__ ElementV *)(user + vUpdateWorkspaceOffset)); + gmHWorkspace.SetGlobalBuffer((__gm__ ElementHWork *)(user + hWorkspaceOffset)); + gmGk.SetGlobalBuffer((__gm__ ElementG *)(kGated ? gk : g)); + gmKDecayWorkspace.SetGlobalBuffer((__gm__ ElementK *)(user + kDecayWorkspaceOffset)); + + gmSeqlen.SetGlobalBuffer((__gm__ int64_t *)cu_seqlens); + gmNumSeq.SetGlobalBuffer((__gm__ int64_t *)(user + numSeqWorkspaceOffset)); + gmNumChunks.SetGlobalBuffer((__gm__ int64_t *)(user + numChunksWorkspaceOffset)); + + if ASCEND_IS_AIC { + cubeBlockScheduler.Init(cu_seqlens, chunk_indices, tiling, user); + } + + if ASCEND_IS_AIV { + vecBlockScheduler.Init(cu_seqlens, chunk_indices, tiling, user); + } + } + + template + __aicore__ inline void InitFromData(GM_ADDR k, GM_ADDR w, GM_ADDR u, GM_ADDR g, GM_ADDR gk, GM_ADDR inital_state, + GM_ADDR cu_seqlens, GM_ADDR chunk_indices, GM_ADDR h, GM_ADDR v_new, + GM_ADDR final_state, const TilingData &tilingData, GM_ADDR user) + { + batch = tilingData.batch; + seqlen = tilingData.seqlen; + kNumHead = tilingData.kNumHead; + vNumHead = tilingData.vNumHead; + kHeadDim = tilingData.kHeadDim; + vHeadDim = tilingData.vHeadDim; + chunkSize = tilingData.chunkSize; + useInitialState = tilingData.useInitialState; + storeFinalState = tilingData.storeFinalState; + isVariedLen = tilingData.isVariedLen; + shapeBatch = tilingData.shapeBatch; + tokenBatch = tilingData.tokenBatch; + vWorkspaceOffset = tilingData.vWorkspaceOffset; + vUpdateWorkspaceOffset = tilingData.vUpdateWorkspaceOffset; + hWorkspaceOffset = tilingData.hWorkspaceOffset; + numSeqWorkspaceOffset = tilingData.numSeqWorkspaceOffset; + numChunksWorkspaceOffset = tilingData.numChunksWorkspaceOffset; + kDecayWorkspaceOffset = tilingData.kDecayWorkspaceOffset; + + gmK.SetGlobalBuffer((__gm__ ElementK *)k); + gmW.SetGlobalBuffer((__gm__ ElementW *)w); + gmU.SetGlobalBuffer((__gm__ ElementU *)u); + gmG.SetGlobalBuffer((__gm__ ElementG *)(scalarGated ? g : gk)); + gmInitialState.SetGlobalBuffer((__gm__ ElementInitialState *)inital_state); + gmH.SetGlobalBuffer((__gm__ ElementH *)h); + gmV.SetGlobalBuffer((__gm__ ElementV *)v_new); + gmFinalState.SetGlobalBuffer((__gm__ ElementFinalState *)final_state); + gmVWorkspace.SetGlobalBuffer((__gm__ ElementVWork *)(user + vWorkspaceOffset)); + gmVUpdateWorkspace.SetGlobalBuffer((__gm__ ElementV *)(user + vUpdateWorkspaceOffset)); + gmHWorkspace.SetGlobalBuffer((__gm__ ElementHWork *)(user + hWorkspaceOffset)); + gmGk.SetGlobalBuffer((__gm__ ElementG *)(kGated ? gk : g)); + gmKDecayWorkspace.SetGlobalBuffer((__gm__ ElementK *)(user + kDecayWorkspaceOffset)); + gmSeqlen.SetGlobalBuffer((__gm__ int64_t *)cu_seqlens); + gmNumSeq.SetGlobalBuffer((__gm__ int64_t *)(user + numSeqWorkspaceOffset)); + gmNumChunks.SetGlobalBuffer((__gm__ int64_t *)(user + numChunksWorkspaceOffset)); + + if ASCEND_IS_AIC { + cubeBlockScheduler.InitFromData(cu_seqlens, chunk_indices, tilingData, user); + } + if ASCEND_IS_AIV { + vecBlockScheduler.InitFromData(cu_seqlens, chunk_indices, tilingData, user); + } + } + + template + __aicore__ inline float LoadScalarAsFloat(AscendC::GlobalTensor tensor, uint32_t offset) const + { + Element value = tensor.GetValue(offset); + if constexpr (std::is_same::value) { + return AscendC::ToFloat(value); + } + return static_cast(value); + } + + __aicore__ inline void ComputeTailVWorkspace(const KdaFwdHOffsets &offsets) + { + uint32_t subBlockIdx = AscendC::GetSubBlockIdx(); + uint32_t subBlockNum = AscendC::GetSubBlockNum(); + uint32_t rowsPerSubBlock = CeilDiv(offsets.blockTokens, subBlockNum); + uint32_t rowBegin = subBlockIdx * rowsPerSubBlock; + uint32_t rowEnd = Min(rowBegin + rowsPerSubBlock, offsets.blockTokens); + if (rowBegin >= rowEnd) { + return; + } + + constexpr uint32_t TAIL_INPUT_OFFSET = 166 * 1024; + constexpr uint32_t TAIL_FLOAT_OFFSET = 167 * 1024; + constexpr uint32_t TAIL_ACCUM_OFFSET = 168 * 1024; + AscendC::LocalTensor inputUb = resource.ubBuf.template GetBufferByByte(TAIL_INPUT_OFFSET); + AscendC::LocalTensor floatUb = resource.ubBuf.template GetBufferByByte(TAIL_FLOAT_OFFSET); + AscendC::LocalTensor accumUb = resource.ubBuf.template GetBufferByByte(TAIL_ACCUM_OFFSET); + + for (uint32_t tokenRow = rowBegin; tokenRow < rowEnd; ++tokenRow) { + AscendC::Duplicate(accumUb, 0.0f, offsets.vBlockDim); + AscendC::PipeBarrier(); + for (uint32_t kIdx = 0; kIdx < kHeadDim; ++kIdx) { + AscendC::DataCopy(inputUb, gmH[offsets.hSrcOffset + kIdx * vHeadDim], offsets.vBlockDim); + AscendC::SetFlag(EVENT_ID7); + AscendC::WaitFlag(EVENT_ID7); + AscendC::Cast(floatUb, inputUb, AscendC::RoundMode::CAST_NONE, offsets.vBlockDim); + AscendC::PipeBarrier(); + float weight = LoadScalarAsFloat(gmW, offsets.wOffset + tokenRow * kHeadDim + kIdx); + AscendC::Muls(floatUb, floatUb, weight, offsets.vBlockDim); + AscendC::PipeBarrier(); + AscendC::Add(accumUb, accumUb, floatUb, offsets.vBlockDim); + AscendC::PipeBarrier(); + } + AscendC::SetFlag(EVENT_ID7); + AscendC::WaitFlag(EVENT_ID7); + AscendC::DataCopy(gmVWorkspace[offsets.vWorkOffset + tokenRow * offsets.vBlockDim], accumUb, + offsets.vBlockDim); + AscendC::SetFlag(EVENT_ID7); + AscendC::WaitFlag(EVENT_ID7); + } + } + + __aicore__ inline void ComputeHWorkspaceVector(const KdaFwdHOffsets &offsets, uint32_t eventId) + { + uint32_t subBlockIdx = AscendC::GetSubBlockIdx(); + uint32_t subBlockNum = AscendC::GetSubBlockNum(); + uint32_t rowsPerSubBlock = CeilDiv(kHeadDim, subBlockNum); + uint32_t rowBegin = subBlockIdx * rowsPerSubBlock; + uint32_t rowEnd = Min(rowBegin + rowsPerSubBlock, kHeadDim); + + constexpr uint32_t MAX_K_ROWS = 64; + constexpr uint32_t ACCUM_OFFSET = 0; + constexpr uint32_t WEIGHT_BROADCAST_OFFSET = 32 * 1024; + constexpr uint32_t WEIGHT_INPUT_OFFSET = 64 * 1024; + constexpr uint32_t WEIGHT_FLOAT_OFFSET = 65 * 1024; + constexpr uint32_t V_INPUT_OFFSET = 66 * 1024; + constexpr uint32_t V_FLOAT_OFFSET = 67 * 1024; + constexpr uint32_t PRODUCT_OFFSET = 68 * 1024; + constexpr uint32_t BROADCAST_SCRATCH_OFFSET = 96 * 1024; + + AscendC::LocalTensor accumUb = resource.ubBuf.template GetBufferByByte(ACCUM_OFFSET); + AscendC::LocalTensor weightBroadcastUb = + resource.ubBuf.template GetBufferByByte(WEIGHT_BROADCAST_OFFSET); + AscendC::LocalTensor weightInputUb = + resource.ubBuf.template GetBufferByByte(WEIGHT_INPUT_OFFSET); + AscendC::LocalTensor weightFloatUb = resource.ubBuf.template GetBufferByByte(WEIGHT_FLOAT_OFFSET); + AscendC::LocalTensor vInputUb = resource.ubBuf.template GetBufferByByte(V_INPUT_OFFSET); + AscendC::LocalTensor vFloatUb = resource.ubBuf.template GetBufferByByte(V_FLOAT_OFFSET); + AscendC::LocalTensor productUb = resource.ubBuf.template GetBufferByByte(PRODUCT_OFFSET); + AscendC::LocalTensor broadcastScratchUb = + resource.ubBuf.template GetBufferByByte(BROADCAST_SCRATCH_OFFSET); + + AscendC::ResetMask(); + AscendC::WaitFlag(eventId); + + for (uint32_t kTileStart = rowBegin; kTileStart < rowEnd; kTileStart += MAX_K_ROWS) { + uint32_t rowCount = Min(MAX_K_ROWS, rowEnd - kTileStart); + uint32_t accumElements = rowCount * offsets.vBlockDim; + AscendC::Duplicate(accumUb, 0.0f, accumElements); + AscendC::PipeBarrier(); + for (uint32_t tokenRow = 0; tokenRow < offsets.blockTokens; ++tokenRow) { + AscendC::DataCopy(weightInputUb, gmK[offsets.wkOffset + tokenRow * kHeadDim], kHeadDim); + AscendC::DataCopy(vInputUb, gmV[offsets.uvOffset + tokenRow * vHeadDim], offsets.vBlockDim); + AscendC::SetFlag(eventId); + AscendC::WaitFlag(eventId); + AscendC::Cast(weightFloatUb, weightInputUb, AscendC::RoundMode::CAST_NONE, kHeadDim); + AscendC::Cast(vFloatUb, vInputUb, AscendC::RoundMode::CAST_NONE, offsets.vBlockDim); + AscendC::PipeBarrier(); + + uint32_t dstShape[2] = {rowCount, offsets.vBlockDim}; + uint32_t srcShape[2] = {rowCount, 1}; + AscendC::Broadcast(weightBroadcastUb, weightFloatUb[kTileStart], dstShape, srcShape, + broadcastScratchUb); + AscendC::PipeBarrier(); + + for (uint32_t localRow = 0; localRow < rowCount; ++localRow) { + uint32_t rowOffset = localRow * offsets.vBlockDim; + AscendC::Mul(productUb, vFloatUb, weightBroadcastUb[rowOffset], offsets.vBlockDim); + AscendC::PipeBarrier(); + AscendC::Add(accumUb[rowOffset], accumUb[rowOffset], productUb, offsets.vBlockDim); + AscendC::PipeBarrier(); + } + AscendC::SetFlag(eventId); + AscendC::WaitFlag(eventId); + } + AscendC::SetFlag(eventId); + AscendC::WaitFlag(eventId); + AscendC::DataCopy(gmHWorkspace[offsets.hWorkOffset + kTileStart * offsets.vBlockDim], accumUb, + accumElements); + AscendC::SetFlag(eventId); + AscendC::WaitFlag(eventId); + } + AscendC::SetFlag(eventId); + } + + __aicore__ inline void PresetVectorPipelineEvents() + { + constexpr uint32_t pongBaseEvent = 4; + if (storeFinalState && std::is_same::value) { + AscendC::SetFlag(EVENT_ID0); + AscendC::SetFlag(EVENT_ID0 + pongBaseEvent); + AscendC::SetFlag(EVENT_ID2); + AscendC::SetFlag(EVENT_ID2 + pongBaseEvent); + } else { + AscendC::SetFlag(EVENT_ID0); + AscendC::SetFlag(EVENT_ID0 + pongBaseEvent); + AscendC::SetFlag(EVENT_ID2); + AscendC::SetFlag(EVENT_ID2 + pongBaseEvent); + } + AscendC::SetFlag(EVENT_ID1); + AscendC::SetFlag(EVENT_ID1 + pongBaseEvent); + AscendC::SetFlag(EVENT_ID3); + AscendC::SetFlag(EVENT_ID3 + pongBaseEvent); + AscendC::SetFlag(EVENT_ID0); + AscendC::SetFlag(EVENT_ID0 + pongBaseEvent); + AscendC::SetFlag(EVENT_ID2); + AscendC::SetFlag(EVENT_ID2 + pongBaseEvent); + } + + __aicore__ inline void DrainVectorPipelineEvents(const bool (&event0FromMte3)[PING_PONG_STAGES], + const bool (&event2FromMte3)[PING_PONG_STAGES]) + { + constexpr uint32_t pongBaseEvent = 4; + if (storeFinalState && std::is_same::value) { + for (uint32_t streamId = 0; streamId < PING_PONG_STAGES; ++streamId) { + uint32_t eventOffset = streamId * pongBaseEvent; + if (event0FromMte3[streamId]) { + AscendC::WaitFlag(EVENT_ID0 + eventOffset); + } else { + AscendC::WaitFlag(EVENT_ID0 + eventOffset); + } + if (event2FromMte3[streamId]) { + AscendC::WaitFlag(EVENT_ID2 + eventOffset); + } else { + AscendC::WaitFlag(EVENT_ID2 + eventOffset); + } + } + } else { + AscendC::WaitFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0 + pongBaseEvent); + AscendC::WaitFlag(EVENT_ID2); + AscendC::WaitFlag(EVENT_ID2 + pongBaseEvent); + } + AscendC::WaitFlag(EVENT_ID1); + AscendC::WaitFlag(EVENT_ID1 + pongBaseEvent); + AscendC::WaitFlag(EVENT_ID3); + AscendC::WaitFlag(EVENT_ID3 + pongBaseEvent); + AscendC::WaitFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0 + pongBaseEvent); + AscendC::WaitFlag(EVENT_ID2); + AscendC::WaitFlag(EVENT_ID2 + pongBaseEvent); + } + + __aicore__ inline void Process() + { + if (isVariedLen) { + AscendC::SyncAll(); + } + + if ASCEND_IS_AIC { + uint32_t coreIdx = AscendC::GetBlockIdx(); + uint32_t coreNum = AscendC::GetBlockNum(); + + auto wLayout = + tla::MakeLayout(shapeBatch * kNumHead * cubeBlockScheduler.totalTokens, kHeadDim); + auto hLayout = tla::MakeLayout( + shapeBatch * vNumHead * cubeBlockScheduler.totalChunks * kHeadDim, vHeadDim); + auto vLayout = tla::MakeLayout(coreNum * chunkSize * PING_PONG_STAGES, + cubeBlockScheduler.vBlockSize); + auto kLayout = + tla::MakeLayout(kHeadDim, shapeBatch * kNumHead * cubeBlockScheduler.totalTokens); + auto vworkLayout = tla::MakeLayout(coreNum * chunkSize * PING_PONG_STAGES, + cubeBlockScheduler.vBlockSize); + auto hworkLayout = tla::MakeLayout(coreNum * kHeadDim * PING_PONG_STAGES, + cubeBlockScheduler.vBlockSize); + + uint32_t taskWaveCount = cubeBlockScheduler.GetTaskWaveCount(); + for (uint32_t waveIdx = 0; waveIdx < taskWaveCount; ++waveIdx) { + BlockMmadWH blockMmadWH(resource); + BlockMmadKV blockMmadKV(resource); + BlockMmadWH blockMmadWHTail(resource); + BlockMmadKV blockMmadKVTail(resource); + AscendC::SyncAll(); + cubeBlockScheduler.InitTaskWave(waveIdx); + uint32_t currStage = 0; // 0: C1, 1: C2 + while (cubeBlockScheduler.isRunning) { + if (currStage == 0) { + /* C1: v_work = w @ h[i] */ + cubeBlockScheduler.InitTasks(); + for (uint32_t i = 0; i < PING_PONG_STAGES; ++i) { + uint32_t streamId = cubeBlockScheduler.GetStreamId(i); + const auto &stream = cubeBlockScheduler.GetStream(i); + if (cubeBlockScheduler.StreamIsDone(stream)) { + continue; + } + + const KdaFwdHOffsets &cube1Offsets = cubeBlockScheduler.GetCurTaskOffsets(stream); + Arch::CrossCoreWaitFlag(cubeBlockScheduler.vec2Done[streamId]); + if (cube1Offsets.blockTokens < 16) { + Arch::CrossCoreSetFlag<0x2, PIPE_MTE2>(cubeBlockScheduler.cube1Done[streamId]); + continue; + } + int64_t cube1OffsetW = cube1Offsets.wOffset; + int64_t cube1OffsetH = cube1Offsets.hSrcOffset; + int64_t cube1OffsetVwork = cube1Offsets.vWorkOffset; + auto tensorW = tla::MakeTensor(gmW[cube1OffsetW], wLayout, Catlass::Arch::PositionGM{}); + auto tensorH = tla::MakeTensor(gmH[cube1OffsetH], hLayout, Catlass::Arch::PositionGM{}); + auto tensorV = + tla::MakeTensor(gmVWorkspace[cube1OffsetVwork], vLayout, Catlass::Arch::PositionGM{}); + GemmCoord cube1Shape{cube1Offsets.blockTokens, cube1Offsets.vBlockDim, kHeadDim}; + auto tensorBlockW = + GetTile(tensorW, tla::MakeCoord(0, 0), tla::MakeShape(cube1Shape.m(), cube1Shape.k())); + auto tensorBlockH = + GetTile(tensorH, tla::MakeCoord(0, 0), tla::MakeShape(cube1Shape.k(), cube1Shape.n())); + auto tensorBlockV = + GetTile(tensorV, tla::MakeCoord(0, 0), tla::MakeShape(cube1Shape.m(), cube1Shape.n())); + if (cube1Offsets.blockTokens < chunkSize) { + blockMmadWHTail.preSetFlags(); + blockMmadWHTail(tensorBlockW, tensorBlockH, tensorBlockV, cube1Shape, EmptyClass{}, + true); + blockMmadWHTail.finalWaitFlags(); + } else { + blockMmadWH.preSetFlags(); + blockMmadWH(tensorBlockW, tensorBlockH, tensorBlockV, cube1Shape); + blockMmadWH.finalWaitFlags(); + } + AscendC::PipeBarrier(); + Arch::CrossCoreSetFlag<0x2, PIPE_FIX>(cubeBlockScheduler.cube1Done[streamId]); + } + } else { + /* C2: h[i+1] = k.T @ v_work */ + for (uint32_t i = 0; i < PING_PONG_STAGES; ++i) { + uint32_t streamId = cubeBlockScheduler.GetStreamId(i); + const auto &stream = cubeBlockScheduler.GetStream(i); + if (cubeBlockScheduler.StreamIsDone(stream)) { + continue; + } + const KdaFwdHOffsets &cube2Offsets = cubeBlockScheduler.GetCurTaskOffsets(stream); + Arch::CrossCoreWaitFlag(cubeBlockScheduler.vec1Done[streamId]); + + if (cubeBlockScheduler.NeedProcessStage2(stream)) { + if (cube2Offsets.blockTokens < 16) { + Arch::CrossCoreSetFlag<0x2, PIPE_MTE2>(cubeBlockScheduler.cube2Done[streamId]); + continue; + } + auto tensorK = + tla::MakeTensor(gmK[cube2Offsets.wkOffset], kLayout, Catlass::Arch::PositionGM{}); + auto tensorVwork = tla::MakeTensor(gmVUpdateWorkspace[cube2Offsets.vWorkOffset], + vworkLayout, Catlass::Arch::PositionGM{}); + auto tensorHwork = tla::MakeTensor(gmHWorkspace[cube2Offsets.hWorkOffset], hworkLayout, + Catlass::Arch::PositionGM{}); + GemmCoord cube2Shape{kHeadDim, cube2Offsets.vBlockDim, cube2Offsets.blockTokens}; + auto tensorBlockK = GetTile(tensorK, tla::MakeCoord(0, 0), + tla::MakeShape(cube2Shape.m(), cube2Shape.k())); + auto tensorBlockVwork = GetTile(tensorVwork, tla::MakeCoord(0, 0), + tla::MakeShape(cube2Shape.k(), cube2Shape.n())); + auto tensorBlockHwork = GetTile(tensorHwork, tla::MakeCoord(0, 0), + tla::MakeShape(cube2Shape.m(), cube2Shape.n())); + if (cube2Offsets.blockTokens < chunkSize) { + blockMmadKVTail.preSetFlags(); + blockMmadKVTail(tensorBlockK, tensorBlockVwork, tensorBlockHwork, cube2Shape, + EmptyClass{}, true); + blockMmadKVTail.finalWaitFlags(); + } else { + blockMmadKV.preSetFlags(); + blockMmadKV(tensorBlockK, tensorBlockVwork, tensorBlockHwork, cube2Shape); + blockMmadKV.finalWaitFlags(); + } + AscendC::PipeBarrier(); + } + Arch::CrossCoreSetFlag<0x2, PIPE_FIX>(cubeBlockScheduler.cube2Done[streamId]); + } + } + currStage ^= 0x01; + } + Arch::CrossCoreWaitFlag(cubeBlockScheduler.vec2Done[0]); + Arch::CrossCoreWaitFlag(cubeBlockScheduler.vec2Done[1]); + } + } + + if ASCEND_IS_AIV { + uint32_t subBlockIdx = AscendC::GetSubBlockIdx(); + uint32_t subBlockNum = AscendC::GetSubBlockNum(); + uint32_t coreIdx = AscendC::GetBlockIdx() / subBlockNum; + uint32_t coreNum = AscendC::GetBlockNum(); + uint32_t taskCount = (isVariedLen ? vecBlockScheduler.tokenBatch : shapeBatch) * vNumHead; + uint32_t rowsPerSubBlock = (kHeadDim + subBlockNum - 1) / subBlockNum; + uint32_t rowBegin = subBlockIdx * rowsPerSubBlock; + uint32_t rowEnd = Min(rowBegin + rowsPerSubBlock, kHeadDim); + uint32_t hRowsPerTile = (32 * 1024) / (vHeadDim * sizeof(ElementH)); + uint32_t stateRowsPerTile = (64 * 1024) / (vHeadDim * sizeof(ElementInitialState)); + uint32_t rowsPerTile = Min(hRowsPerTile, stateRowsPerTile); + uint32_t totalChunks = isVariedLen ? vecBlockScheduler.totalChunks : ((seqlen + chunkSize - 1) / chunkSize); + uint32_t stateBlockSize = kHeadDim * vHeadDim; + AscendC::LocalTensor stateUbTensorPing = + resource.ubBuf.template GetBufferByByte(0); + AscendC::LocalTensor stateUbTensorPong = + resource.ubBuf.template GetBufferByByte(96 * 1024); + AscendC::LocalTensor hUbTensorPing = resource.ubBuf.template GetBufferByByte(64 * 1024); + AscendC::LocalTensor hUbTensorPong = + resource.ubBuf.template GetBufferByByte(160 * 1024); + uint32_t taskWaveCount = vecBlockScheduler.GetTaskWaveCount(); + for (uint32_t waveIdx = 0; waveIdx < taskWaveCount; ++waveIdx) { + EpilogueKdaFwdHVnew epilogueKdaFwdHVnew(resource); + EpilogueKdaFwdHUpdate epilogueKdaFwdHUpdate(resource); + uint32_t taskIdx = waveIdx * coreNum + coreIdx; + uint32_t pingpongFlag = 1; + AscendC::SetFlag(EVENT_ID0); + AscendC::SetFlag(EVENT_ID1); + if (taskIdx < taskCount) { + uint32_t batchIdx = taskIdx / vNumHead; + uint32_t vHeadIdx = taskIdx % vNumHead; + uint32_t chunkOffset = isVariedLen ? vecBlockScheduler.GetVarlenChunkOffset(batchIdx) : 0; + uint32_t shapeBatchIdx = isVariedLen ? 0 : batchIdx; + uint32_t hBaseOffset = + (shapeBatchIdx * vNumHead * totalChunks + vHeadIdx * totalChunks + chunkOffset) * + stateBlockSize; + uint32_t initialStateBaseOffset = taskIdx * stateBlockSize; + for (uint32_t rowOffset = rowBegin; rowOffset < rowEnd; rowOffset += rowsPerTile) { + uint32_t rowsThisTile = Min(rowsPerTile, rowEnd - rowOffset); + uint32_t stateTileElems = rowsThisTile * vHeadDim; + uint32_t hOffset = hBaseOffset + rowOffset * vHeadDim; + AscendC::LocalTensor stateUbTensor = + pingpongFlag ? stateUbTensorPing : stateUbTensorPong; + AscendC::LocalTensor hUbTensor = pingpongFlag ? hUbTensorPing : hUbTensorPong; + auto eventId = pingpongFlag ? EVENT_ID1 : EVENT_ID0; + AscendC::WaitFlag(eventId); + if (useInitialState) { + uint32_t initialStateOffset = initialStateBaseOffset + rowOffset * vHeadDim; + if constexpr (!std::is_same::value) { + AscendC::DataCopy(stateUbTensor, gmInitialState[initialStateOffset], stateTileElems); + AscendC::SetFlag(eventId); + AscendC::WaitFlag(eventId); + AscendC::Cast(hUbTensor, stateUbTensor, AscendC::RoundMode::CAST_RINT, stateTileElems); + AscendC::SetFlag(eventId); + AscendC::WaitFlag(eventId); + AscendC::DataCopy(gmH[hOffset], hUbTensor, stateTileElems); + } else { + AscendC::DataCopy(stateUbTensor, gmInitialState[initialStateOffset], stateTileElems); + AscendC::SetFlag(eventId); + AscendC::WaitFlag(eventId); + AscendC::DataCopy(gmH[hOffset], stateUbTensor, stateTileElems); + } + } else { + AscendC::Duplicate(hUbTensor, static_cast(0), stateTileElems); + AscendC::SetFlag(eventId); + AscendC::WaitFlag(eventId); + AscendC::DataCopy(gmH[hOffset], hUbTensor, stateTileElems); + } + AscendC::SetFlag(eventId); + pingpongFlag = 1 - pingpongFlag; + } + } + AscendC::WaitFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID1); + + AscendC::SyncAll(); + vecBlockScheduler.InitTaskWave(waveIdx); + Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(vecBlockScheduler.vec2Done[0]); + Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(vecBlockScheduler.vec2Done[1]); + PresetVectorPipelineEvents(); + uint32_t currStage = 0; // 0: V1, 1: V2 + bool waitStageFence = false; + bool event0FromMte3[PING_PONG_STAGES] = {false, false}; + bool event2FromMte3[PING_PONG_STAGES] = { + !(storeFinalState && std::is_same::value), + !(storeFinalState && std::is_same::value)}; + constexpr uint32_t pongEventOffset = 4; + while (vecBlockScheduler.isRunning) { + if (waitStageFence) { + AscendC::WaitFlag(EVENT_ID1); + } + if (currStage == 0) { + /* V1: + * gmV = gmU - gmVWorkspace + * g_buf = gmG[-1] - gmG + * g_buf = exp(g_buf) + * gmVWorkspace = g_buf * gmV + */ + vecBlockScheduler.InitTasks(); + for (uint32_t i = 0; i < PING_PONG_STAGES; ++i) { + uint32_t streamId = vecBlockScheduler.GetStreamId(i); + const auto &stream = vecBlockScheduler.GetStream(i); + if (vecBlockScheduler.StreamIsDone(stream)) { + continue; + } + const KdaFwdHOffsets &vec1Offsets = vecBlockScheduler.GetCurTaskOffsets(stream); + if (vec1Offsets.blockTokens < 16) { + ComputeTailVWorkspace(vec1Offsets); + } + bool waitWsFromMte3 = storeFinalState && std::is_same::value && + event0FromMte3[streamId]; + epilogueKdaFwdHVnew(gmV[vec1Offsets.uvOffset], gmVUpdateWorkspace[vec1Offsets.vWorkOffset], + gmG[vec1Offsets.gOffset], gmU[vec1Offsets.uvOffset], + gmVWorkspace[vec1Offsets.vWorkOffset], vec1Offsets.blockTokens, + vec1Offsets.vBlockDim, vHeadDim, vecBlockScheduler.cube1Done[streamId], + vecBlockScheduler.vec1Done[streamId], vec1Offsets.isInitialState, + vec1Offsets.isFinalState, storeFinalState, waitWsFromMte3, + (streamId == 0)); + AscendC::SetFlag(EVENT_ID1); + AscendC::WaitFlag(EVENT_ID1); + if (storeFinalState && std::is_same::value) { + event0FromMte3[streamId] = false; + } + } + } else { + /* V2: h[i+1] += h_work if i < num_chunks - 1 else None */ + for (uint32_t i = 0; i < PING_PONG_STAGES; ++i) { + uint32_t streamId = vecBlockScheduler.GetStreamId(i); + const auto &stream = vecBlockScheduler.GetStream(i); + if (vecBlockScheduler.StreamIsDone(stream)) { + continue; + } + const KdaFwdHOffsets &vec2Offsets = vecBlockScheduler.GetCurTaskOffsets(stream); + if (vecBlockScheduler.NeedProcessStage2(stream)) { + bool vectorC2Path = vec2Offsets.blockTokens < 16; + if (vectorC2Path) { + Arch::CrossCoreWaitFlag(vecBlockScheduler.cube2Done[streamId]); + ComputeHWorkspaceVector(vec2Offsets, + EVENT_ID3 + (streamId == 0 ? 0 : pongEventOffset)); + } + if (storeFinalState && std::is_same::value) { + event0FromMte3[streamId] = true; + event2FromMte3[streamId] = !vec2Offsets.isFinalState; + } + // step 4: h[i+1] += h_work if i < num_chunks - 1 else None + epilogueKdaFwdHUpdate( + gmH[vec2Offsets.hDstOffset], gmFinalState[vec2Offsets.finalStateOffset], + gmG[vec2Offsets.gOffset], gmH[vec2Offsets.hSrcOffset], + gmHWorkspace[vec2Offsets.hWorkOffset], gmGk[vec2Offsets.gkOffset], + gmInitialState[vec2Offsets.initialStateOffset], vec2Offsets.blockTokens, kHeadDim, + vec2Offsets.vBlockDim, vHeadDim, vecBlockScheduler.cube2Done[streamId], + vec2Offsets.isInitialState, vec2Offsets.isFinalState, storeFinalState, + useInitialState, (streamId == 0), vectorC2Path); + AscendC::SetFlag(EVENT_ID1); + AscendC::WaitFlag(EVENT_ID1); + } else { + Arch::CrossCoreWaitFlag(vecBlockScheduler.cube2Done[streamId]); + } + Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(vecBlockScheduler.vec2Done[streamId]); + } + } + waitStageFence = vecBlockScheduler.isRunning; + if (waitStageFence) { + AscendC::SetFlag(EVENT_ID1); + } + currStage ^= 0x01; + } + + DrainVectorPipelineEvents(event0FromMte3, event2FromMte3); + } + } + } +}; + +} // namespace Catlass::Gemm::Kernel diff --git a/csrc/chunk_kda_fwd/op_kernel/arch35/chunk_kda_fwd_finalize_arch35.h b/csrc/chunk_kda_fwd/op_kernel/arch35/chunk_kda_fwd_finalize_arch35.h new file mode 100644 index 0000000000..f77a4caa73 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/arch35/chunk_kda_fwd_finalize_arch35.h @@ -0,0 +1,1521 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#pragma once + +#ifndef CATLASS_ARCH +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 +#define CATLASS_ARCH 3510 +#else +#define CATLASS_ARCH 2201 +#endif +#endif + +#include "catlass/arch/kda_arch.hpp" +#include "catlass/arch/kda_cross_core_sync.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/kda_catlass.hpp" +#include "catlass/gemm/block/kda_block_mmad.hpp" +#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" +#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "../kernel_utils/block/block_mmad_pingpong_tla_multi.hpp" +#include "catlass/layout/kda_layout.hpp" +#include "kernel_operator.h" +#include "../chunk_kda_fwd_varlen.h" +#include "tla/kda_tla_layout.hpp" +#include "tla/kda_tensor.hpp" + +using namespace AscendC; + +namespace KdaFinalize { +namespace { +using KdaInt64 = tla::Int<64>; +using KdaInt128 = tla::Int<128>; +constexpr float LN2 = 0.69314718055994530942f; +constexpr float KDA_EXP2_CLAMP = 80.0f; +constexpr float KDA_EXP_INPUT_MAX = KDA_EXP2_CLAMP * LN2; +constexpr float KDA_EXP_INPUT_MIN = -KDA_EXP2_CLAMP * LN2; +constexpr float KDA_FP16_MAX = 65504.0f; +constexpr uint32_t EXP2_UB_ELEMENTS = 256; +constexpr uint32_t EXP2_UB_BYTES = EXP2_UB_ELEMENTS * (sizeof(float) + sizeof(uint16_t)); +constexpr uint32_t EXP2_EVENT_ID = 0; +constexpr uint32_t KDA_SOLVE_BT = 64; +constexpr uint32_t KDA_SOLVE_MATRIX_ELEMENTS = KDA_SOLVE_BT * KDA_SOLVE_BT; +constexpr uint32_t KDA_SOLVE_SCRATCH_X = 0; +constexpr uint32_t KDA_SOLVE_SCRATCH_Y0 = 1; +constexpr uint32_t KDA_SOLVE_SCRATCH_TMP = 2; +constexpr uint32_t KDA_SOLVE_SCRATCH_Y1 = 3; +constexpr uint32_t KDA_SOLVE_SCRATCH_IDENTITY = 4; +constexpr uint32_t KDA_SOLVE_SCRATCH_SLOTS = 5; +constexpr uint32_t KDA_SOLVE_DIAG_BT = 16; +constexpr uint32_t KDA_SOLVE_DIAG_BLOCKS = KDA_SOLVE_BT / KDA_SOLVE_DIAG_BT; +constexpr uint32_t KDA_SOLVE_DIAG_MCH_ITERS = 3; +constexpr uint32_t KDA_SCORE_REF_BC = 16; +constexpr uint32_t KDA_VEC_ARENA_ELEMENTS = 32768; +constexpr uint32_t KDA_BITS_PER_MASK_BYTE = 8; +constexpr uint32_t KDA_SELECT_COL_BLOCKS = 2; +constexpr uint32_t KDA_SELECT_COL_MASK_BYTES = KDA_SOLVE_MATRIX_ELEMENTS / KDA_BITS_PER_MASK_BYTE; +constexpr uint32_t KDA_SELECT_MASK_BYTES = KDA_SELECT_COL_BLOCKS * KDA_SELECT_COL_MASK_BYTES; +constexpr uint32_t KDA_SELECT_AQK_MASK_BYTE_OFFSET = 120 * 1024; +constexpr uint32_t KDA_SELECT_AKK_MASK_BYTE_OFFSET = KDA_SELECT_AQK_MASK_BYTE_OFFSET + KDA_SELECT_MASK_BYTES; +constexpr uint32_t KDA_SELECT_ZERO_BYTE_OFFSET = KDA_SELECT_AKK_MASK_BYTE_OFFSET + KDA_SELECT_MASK_BYTES; +constexpr uint32_t KDA_SELECT_ZERO_FLOAT_OFFSET = KDA_SELECT_ZERO_BYTE_OFFSET / sizeof(float); +constexpr uint8_t KDA_SCORE_DONE_FLAG0 = 2; +constexpr uint8_t KDA_SCORE_DONE_FLAG1 = 3; +constexpr uint8_t KDA_SCORE_READY_FLAG0 = 4; +constexpr uint8_t KDA_SCORE_READY_FLAG1 = 5; +constexpr uint32_t KDA_SCORE_QUEUE_DEPTH = 2; +constexpr uint32_t KDA_SYNC_REVERSE_DEPTH = 1; +constexpr uint32_t KDA_SCORE_SCRATCH_PLANES = 3; +constexpr uint32_t KDA_SCORE_SCRATCH_QG = 0; +constexpr uint32_t KDA_SCORE_SCRATCH_W = 1; +constexpr uint32_t KDA_SCORE_SCRATCH_KG = 2; +constexpr uint64_t KDA_WORKSPACE_ALIGN = 512; +constexpr uint32_t KDA_GATE_TILE_ROWS = 32; +constexpr uint32_t KDA_CUBE_MIN_REDUCTION = 16; + +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 +constexpr AscendC::MicroAPI::CastTrait KDA_TAIL_BF16_TO_FP32 = { + AscendC::MicroAPI::RegLayout::ZERO, + AscendC::MicroAPI::SatMode::SAT, + AscendC::MicroAPI::MaskMergeMode::ZEROING, + AscendC::RoundMode::CAST_NONE, +}; + +static __simd_vf__ inline void ComputeKdaTailOutputRegbase(__ubuf__ float *dst, __ubuf__ bfloat16_t *h, + __ubuf__ bfloat16_t *vNew, __ubuf__ bfloat16_t *qg, + __ubuf__ bfloat16_t *aqk, uint16_t rows, uint16_t curT, + float scale) +{ + using namespace AscendC::MicroAPI; + constexpr uint16_t DIM = 128; + constexpr uint16_t FP32_REG_ELEMENTS = AscendC::VECTOR_REG_WIDTH / sizeof(float); + constexpr uint16_t AQK_ROW_ELEMENTS = 16; + MaskReg fullMask = CreateMask(); + RegTensor scaleReg; + Duplicate(scaleReg, scale, fullMask); + + for (uint16_t row = 0; row < rows; ++row) { + RegTensor acc0; + RegTensor acc1; + Duplicate(acc0, 0.0f, fullMask); + Duplicate(acc1, 0.0f, fullMask); + + for (uint16_t d = 0; d < DIM; ++d) { + RegTensor matrixRaw0; + RegTensor matrixRaw1; + RegTensor weightRaw; + RegTensor matrix0; + RegTensor matrix1; + RegTensor weight; + RegTensor product0; + RegTensor product1; + DataCopy(matrixRaw0, h + static_cast(d) * DIM); + DataCopy(matrixRaw1, + h + static_cast(d) * DIM + FP32_REG_ELEMENTS); + LoadAlign(weightRaw, qg + static_cast(row) * DIM + d); + Cast(matrix0, matrixRaw0, fullMask); + Cast(matrix1, matrixRaw1, fullMask); + Cast(weight, weightRaw, fullMask); + Mul(product0, matrix0, weight, fullMask); + Mul(product1, matrix1, weight, fullMask); + Add(acc0, acc0, product0, fullMask); + Add(acc1, acc1, product1, fullMask); + } + Mul(acc0, acc0, scaleReg, fullMask); + Mul(acc1, acc1, scaleReg, fullMask); + + for (uint16_t token = 0; token < curT; ++token) { + RegTensor matrixRaw0; + RegTensor matrixRaw1; + RegTensor weightRaw; + RegTensor matrix0; + RegTensor matrix1; + RegTensor weight; + RegTensor product0; + RegTensor product1; + DataCopy(matrixRaw0, vNew + static_cast(token) * DIM); + DataCopy( + matrixRaw1, vNew + static_cast(token) * DIM + FP32_REG_ELEMENTS); + LoadAlign(weightRaw, + aqk + static_cast(row) * AQK_ROW_ELEMENTS + token); + Cast(matrix0, matrixRaw0, fullMask); + Cast(matrix1, matrixRaw1, fullMask); + Cast(weight, weightRaw, fullMask); + Mul(product0, matrix0, weight, fullMask); + Mul(product1, matrix1, weight, fullMask); + Add(acc0, acc0, product0, fullMask); + Add(acc1, acc1, product1, fullMask); + } + + DataCopy(dst + static_cast(row) * DIM, acc0, fullMask); + DataCopy(dst + static_cast(row) * DIM + FP32_REG_ELEMENTS, acc1, fullMask); + } +} +#endif + +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 +using KdaArchTag = Catlass::Arch::Ascend950; +#else +using KdaArchTag = Catlass::Arch::AtlasA2; +#endif +using KdaDispatchPolicy = Catlass::Gemm::MmadPingpong; +using KdaScoreDispatchPolicy = Catlass::Gemm::MmadPingpongTlaMulti; +static_assert(KdaScoreDispatchPolicy::ENABLE_L1_RESIDENT, + "KDA Aqk/Akk score MMAD must keep the shared right matrix resident in L1"); +static_assert(KdaScoreDispatchPolicy::L1B_STAGES == 1, + "KDA Aqk/Akk score MMAD needs one L1 B slot so the second MMAD reuses it"); +using KdaSolveDispatchPolicy = Catlass::Gemm::MmadPingpong; +static_assert(!KdaSolveDispatchPolicy::USE_HF32_MODE, "KDA triangular solve must use IEEE FP32 Cube mode"); +using KdaL1TileShape = tla::Shape; +using KdaL0TileShape = KdaL1TileShape; +using KdaSolveL1TileShape = tla::Shape; +using KdaSolveL0TileShape = KdaSolveL1TileShape; + +__aicore__ inline uint32_t FloatToBits(float value) +{ + union Bits { + __aicore__ Bits() {} + float f; + uint32_t u; + } bits; + bits.f = value; + return bits.u; +} + +__aicore__ inline float BitsToFloat(uint32_t value) +{ + union Bits { + __aicore__ Bits() {} + uint32_t u; + float f; + } bits; + bits.u = value; + return bits.f; +} + +__aicore__ inline uint16_t Bf16ToBits(bfloat16_t value) +{ + union Bits { + __aicore__ Bits() {} + bfloat16_t f; + uint16_t u; + } bits; + bits.f = value; + return bits.u; +} + +__aicore__ inline bfloat16_t BitsToBf16(uint16_t value) +{ + union Bits { + __aicore__ Bits() {} + uint16_t u; + bfloat16_t f; + } bits; + bits.u = value; + return bits.f; +} + +template +__aicore__ inline T FloatToType(float value) +{ + if constexpr (IsSameType::value) { + uint32_t bits = FloatToBits(value); + uint32_t bias = 0x7FFFu + ((bits >> 16) & 1u); + return BitsToBf16(static_cast((bits + bias) >> 16)); + } + return static_cast(value); +} + +template +class ChunkKdaFwdFinalizeKernel { +public: + using OUT_T = float; + using AKK_T = float; + template + __aicore__ inline void Init(GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR gk, GM_ADDR beta, GM_ADDR initialState, + GM_ADDR cuSeqlens, GM_ADDR chunkIndices, GM_ADDR preparedQG, GM_ADDR preparedAqk, + GM_ADDR propagatedVNew, GM_ADDR propagatedH, GM_ADDR o, GM_ADDR finalState, GM_ADDR aqk, + GM_ADDR akk, GM_ADDR w, GM_ADDR u, GM_ADDR qg, GM_ADDR kg, GM_ADDR vNew, GM_ADDR h, + GM_ADDR workspace, const TilingData &tiling, TPipe *pipe, bool initVecBuffers = true) + { + pipe_ = pipe; + q_.SetGlobalBuffer((__gm__ T *)q); + k_.SetGlobalBuffer((__gm__ T *)k); + v_.SetGlobalBuffer((__gm__ T *)v); + gk_.SetGlobalBuffer((__gm__ GK_T *)gk); + beta_.SetGlobalBuffer((__gm__ BETA_T *)beta); + if (initialState != nullptr) { + initialState_.SetGlobalBuffer((__gm__ float *)initialState); + } + cuSeqlensAddr_ = reinterpret_cast<__gm__ int64_t *>(cuSeqlens); + if (preparedQG != nullptr) { + preparedQG_.SetGlobalBuffer((__gm__ T *)preparedQG); + } + if (preparedAqk != nullptr) { + preparedAqk_.SetGlobalBuffer((__gm__ T *)preparedAqk); + } + if (propagatedVNew != nullptr) { + propagatedVNew_.SetGlobalBuffer((__gm__ T *)propagatedVNew); + } + if (propagatedH != nullptr) { + propagatedH_.SetGlobalBuffer((__gm__ T *)propagatedH); + } + chunkIndicesAddr_ = reinterpret_cast<__gm__ int64_t *>(chunkIndices); + o_.SetGlobalBuffer((__gm__ OUT_T *)o); + finalState_.SetGlobalBuffer((__gm__ float *)finalState); + aqk_.SetGlobalBuffer((__gm__ float *)aqk); + akk_.SetGlobalBuffer((__gm__ AKK_T *)akk); + w_.SetGlobalBuffer((__gm__ T *)w); + u_.SetGlobalBuffer((__gm__ OUT_T *)u); + qg_.SetGlobalBuffer((__gm__ T *)qg); + kg_.SetGlobalBuffer((__gm__ T *)kg); + vNew_.SetGlobalBuffer((__gm__ T *)vNew); + h_.SetGlobalBuffer((__gm__ float *)h); + solveWorkspace_.SetGlobalBuffer((__gm__ float *)workspace); + + B_ = tiling.batch; + N_ = tiling.seqNum; + H_ = tiling.qHeadNum; + HV_ = tiling.vHeadNum; + T_ = tiling.seqlen; + K_ = tiling.kHeadDim; + V_ = tiling.vHeadDim; + BT_ = tiling.chunkSize; + NT_ = tiling.totalChunks; + scale_ = tiling.scale; + hasInitial_ = tiling.hasInitialState; + isVarLen_ = tiling.isVarLen; + usedCoreNum_ = tiling.outputUsedCoreNum; + const uint64_t outputElements = B_ * HV_ * T_ * V_; + o_.SetGlobalBuffer((__gm__ OUT_T *)workspace); + u_.SetGlobalBuffer((__gm__ OUT_T *)workspace + outputElements); + if ASCEND_IS_AIV { + uint64_t subBlockNum = static_cast(GetSubBlockNum()); + solveCoreIdx_ = subBlockNum == 0 ? 0 : static_cast(GetBlockIdx()) / subBlockNum; + } else { + solveCoreIdx_ = static_cast(GetBlockIdx()); + } + if (pipe_ != nullptr && initVecBuffers) { + pipe_->InitBuffer(exp2Buf_, EXP2_UB_BYTES); + pipe_->InitBuffer(vecBuf_, KDA_VEC_ARENA_ELEMENTS * sizeof(float)); + const uint64_t gateWritebackRows = ScoreVectorMaxRows(5 * sizeof(float) + 2 * sizeof(T) + sizeof(GK_T)); + pipe_->InitBuffer(gateWritebackBuf_, + static_cast(gateWritebackRows * K_ * (3 * sizeof(T) + sizeof(GK_T)))); + AllocVectorEvents(); + } + } + __aicore__ inline void ProcessAiv() + { + ProcessOutAiv(); + ReleaseVectorEvents(); + } + + __aicore__ inline void ProcessAic() { ProcessOutAic(); } + +private: + __aicore__ inline void AllocVectorEvents() + { + mte2ToVEvent_ = pipe_->AllocEventID(); + vToMte2Event_ = pipe_->AllocEventID(); + vToMte3Event_ = pipe_->AllocEventID(); + mte3ToVEvent_ = pipe_->AllocEventID(); + mte2ToMte3Event_ = pipe_->AllocEventID(); + mte3ToMte2Event_ = pipe_->AllocEventID(); + vToSEvent_ = pipe_->AllocEventID(); + sToVEvent_ = pipe_->AllocEventID(); + sToMte2Event_ = pipe_->AllocEventID(); + vectorEventsAllocated_ = true; + } + + __aicore__ inline void ReleaseVectorEvents() + { + if (!vectorEventsAllocated_) { + return; + } + pipe_->ReleaseEventID(mte2ToVEvent_); + pipe_->ReleaseEventID(vToMte2Event_); + pipe_->ReleaseEventID(vToMte3Event_); + pipe_->ReleaseEventID(mte3ToVEvent_); + pipe_->ReleaseEventID(mte2ToMte3Event_); + pipe_->ReleaseEventID(mte3ToMte2Event_); + pipe_->ReleaseEventID(vToSEvent_); + pipe_->ReleaseEventID(sToVEvent_); + pipe_->ReleaseEventID(sToMte2Event_); + vectorEventsAllocated_ = false; + } + + __aicore__ inline uint64_t QOffset(uint64_t b, uint64_t h, uint64_t t, uint64_t d) const + { + return ((b * H_ + h) * T_ + t) * K_ + d; + } + + __aicore__ inline uint64_t KVOffset(uint64_t b, uint64_t hv, uint64_t t, uint64_t d, uint64_t dim) const + { + return ((b * HV_ + hv) * T_ + t) * dim + d; + } + + __aicore__ inline uint64_t OutputOffset(uint64_t b, uint64_t hv, uint64_t t, uint64_t d) const + { + return ((b * T_ + t) * HV_ + hv) * V_ + d; + } + + __aicore__ inline uint64_t BetaOffset(uint64_t b, uint64_t hv, uint64_t t) const { return (b * HV_ + hv) * T_ + t; } + + __aicore__ inline uint64_t AOffset(uint64_t b, uint64_t hv, uint64_t t, uint64_t j) const + { + return ((b * HV_ + hv) * T_ + t) * BT_ + j; + } + + __aicore__ inline uint64_t HOffset(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t d, uint64_t r) const + { + return (((b * HV_ + hv) * NT_ + chunkIdx) * K_ + d) * V_ + r; + } + + __aicore__ inline uint64_t WScratchOffset(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t t, uint64_t d) const + { + return (((b * HV_ + hv) * NT_ + chunkIdx) * BT_ + t) * K_ + d; + } + + __aicore__ inline uint64_t SolveScratchOffset(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t slot) const + { + (void)b; + (void)hv; + (void)chunkIdx; + uint64_t matrixElements = BT_ * BT_; + return solveCoreIdx_ * KDA_SOLVE_SCRATCH_SLOTS * matrixElements + slot * matrixElements; + } + + __aicore__ inline uint64_t ScoreScratchOffset(uint64_t slot, uint64_t plane, uint64_t t = 0, uint64_t d = 0) const + { + return (((solveCoreIdx_ * KDA_SCORE_QUEUE_DEPTH + slot) * KDA_SCORE_SCRATCH_PLANES + plane) * BT_ + t) * K_ + d; + } + + __aicore__ inline uint64_t ScoreRefBlockSize() const { return KDA_SCORE_REF_BC; } + + __aicore__ inline uint64_t ScoreRowBlockCount(uint64_t curT, uint64_t rowBegin) const + { + uint64_t blockSize = ScoreRefBlockSize(); + uint64_t rowCount = curT - rowBegin; + if (rowCount > blockSize) { + rowCount = blockSize; + } + return rowCount; + } + + __aicore__ inline uint64_t ScoreRefToken(uint64_t start, uint64_t curT, uint64_t rowBegin, uint64_t rowCount) const + { + uint64_t ref = rowBegin + rowCount / 2; + if (ref >= curT) { + ref = curT - 1; + } + return start + ref; + } + + __aicore__ inline void RunExp2(LocalTensor &tensor, uint32_t count) + { + SetFlag(EXP2_EVENT_ID); + WaitFlag(EXP2_EVENT_ID); + ClampExpInput(tensor, count); + Exp(tensor, tensor, count); + PipeBarrier(); + SetFlag(EXP2_EVENT_ID); + WaitFlag(EXP2_EVENT_ID); + } + + __aicore__ inline void ClampExpInput(LocalTensor &tensor, uint32_t count) + { + Mins(tensor, tensor, KDA_EXP_INPUT_MAX, count); + PipeBarrier(); + Maxs(tensor, tensor, KDA_EXP_INPUT_MIN, count); + PipeBarrier(); + } + + __aicore__ inline void ClampFp32ToOutputType(LocalTensor &tensor, uint32_t count) + { + if constexpr (IsSameType::value) { + Mins(tensor, tensor, KDA_FP16_MAX, count); + PipeBarrier(); + Maxs(tensor, tensor, -KDA_FP16_MAX, count); + PipeBarrier(); + } + } + + template + __aicore__ inline void CopyVectorIn(LocalTensor &dst, GlobalTensor &src, uint64_t offset, + uint64_t count) + { + uint64_t rowBytes = count * static_cast(sizeof(CopyT)); + if (rowBytes >= 32 && rowBytes % 32 == 0) { + DataCopy(dst, src[offset], static_cast(count)); + return; + } + DataCopyParams params{1, static_cast(rowBytes), 0, 0}; + DataCopyPadParams padParams{false, 0, 0, 0}; + DataCopyPad(dst, src[offset], params, padParams); + } + + template + __aicore__ inline void CopyVectorOut(GlobalTensor &dst, uint64_t offset, LocalTensor &src, + uint64_t count) + { + uint64_t rowBytes = count * static_cast(sizeof(CopyT)); + if (rowBytes >= 32 && rowBytes % 32 == 0) { + DataCopy(dst[offset], src, static_cast(count)); + return; + } + DataCopyParams params{1, static_cast(rowBytes), 0, 0}; + DataCopyPad(dst[offset], src, params); + } + + template + __aicore__ inline void CopyRowsOut(GlobalTensor &dst, uint64_t offset, LocalTensor &src, + uint64_t rows, uint64_t cols, uint64_t dstStride) + { + if (cols == dstStride) { + CopyVectorOut(dst, offset, src, rows * cols); + return; + } + constexpr uint64_t blockBytes = 32; + const uint64_t rowBytes = cols * sizeof(CopyT); + const uint64_t gapBytes = (dstStride - cols) * sizeof(CopyT); + DataCopyParams params{ +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + 1, +#else + static_cast(rows), +#endif + static_cast(rowBytes / blockBytes), 0, +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + 0 +#else + static_cast(gapBytes / blockBytes) +#endif + }; +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + const uint64_t dstRowBytes = dstStride * sizeof(CopyT); + LoopModeParams loopParams{static_cast(rows), 1, rowBytes, dstRowBytes, 0, 0}; + // Loop-mode registers are core-local state and must not leak across DMA calls. + ResetLoopModePara(DataCopyMVType::UB_TO_OUT); + SetLoopModePara(loopParams, DataCopyMVType::UB_TO_OUT); + DataCopy(dst[offset], src, params); + ResetLoopModePara(DataCopyMVType::UB_TO_OUT); +#else + DataCopy(dst[offset], src, params); +#endif + } + + template + __aicore__ inline void CopyRowIn(LocalTensor &dst, GlobalTensor &src, uint64_t offset) + { + CopyVectorIn(dst, src, offset, K_); + } + + template + __aicore__ inline void CopyRowOut(GlobalTensor &dst, uint64_t offset, LocalTensor &src) + { + CopyVectorOut(dst, offset, src, K_); + } + + __aicore__ inline LocalTensor VecScratch(uint64_t slot) + { + return vecBuf_.Get()[slot * EXP2_UB_ELEMENTS]; + } + + template + __aicore__ inline void LoadAsFloatRow(GlobalTensor &src, uint64_t srcOffset, LocalTensor &dst, + uint64_t count) + { + if constexpr (IsSameType::value) { + CopyVectorIn(dst, src, srcOffset, count); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + Adds(dst, dst, 0.0f, static_cast(count)); + PipeBarrier(); + SetFlag(vToMte2Event_); + WaitFlag(vToMte2Event_); + } else { + constexpr uint32_t typedOffset = EXP2_UB_ELEMENTS * sizeof(float) / sizeof(CopyT); + LocalTensor rowLocal = exp2Buf_.Get()[typedOffset]; + CopyVectorIn(rowLocal, src, srcOffset, count); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + Cast(dst, rowLocal, RoundMode::CAST_NONE, static_cast(count)); + PipeBarrier(); + SetFlag(vToMte2Event_); + WaitFlag(vToMte2Event_); + } + PipeBarrier(); + } + + __aicore__ inline void ComputeTailLocalRows(LocalTensor &dst, uint64_t b, uint64_t hv, uint64_t start, + uint64_t curT, uint64_t rowBegin, uint64_t rows) + { + LocalTensor vRow = exp2Buf_.Get(); + LocalTensor coefficientTyped = gateWritebackBuf_.Get(); + LocalTensor coefficients = gateWritebackBuf_.Get()[BT_]; + for (uint64_t localRow = 0; localRow < rows; ++localRow) { + LocalTensor dstRow = dst[localRow * V_]; + CopyVectorIn(coefficientTyped, preparedAqk_, AOffset(b, hv, start + rowBegin + localRow, 0), curT); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + Cast(coefficients, coefficientTyped, RoundMode::CAST_NONE, static_cast(curT)); + PipeBarrier(); + SetFlag(vToSEvent_); + WaitFlag(vToSEvent_); + Duplicate(dstRow, 0.0f, static_cast(V_)); + PipeBarrier(); + for (uint64_t j = 0; j < curT; ++j) { + LoadAsFloatRow(propagatedVNew_, KVOffset(b, hv, start + j, 0, V_), vRow, V_); + float weight = coefficients.GetValue(j); + SetFlag(sToVEvent_); + WaitFlag(sToVEvent_); + Muls(vRow, vRow, weight, static_cast(V_)); + PipeBarrier(); + Add(dstRow, dstRow, vRow, static_cast(V_)); + PipeBarrier(); + SetFlag(vToMte2Event_); + WaitFlag(vToMte2Event_); + } + SetFlag(sToMte2Event_); + WaitFlag(sToMte2Event_); + } + } + + __aicore__ inline void ComputeTailStateRows(LocalTensor &dst, uint64_t b, uint64_t hv, uint64_t chunkIdx, + uint64_t start, uint64_t rowBegin, uint64_t rows) + { + LocalTensor hRow = exp2Buf_.Get(); + LocalTensor coefficientTyped = gateWritebackBuf_.Get(); + LocalTensor coefficients = gateWritebackBuf_.Get()[BT_]; + for (uint64_t localRow = 0; localRow < rows; ++localRow) { + LocalTensor dstRow = dst[localRow * V_]; + CopyVectorIn(coefficientTyped, preparedQG_, KVOffset(b, hv, start + rowBegin + localRow, 0, K_), K_); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + Cast(coefficients, coefficientTyped, RoundMode::CAST_NONE, static_cast(K_)); + PipeBarrier(); + SetFlag(vToSEvent_); + WaitFlag(vToSEvent_); + Duplicate(dstRow, 0.0f, static_cast(V_)); + PipeBarrier(); + for (uint64_t d = 0; d < K_; ++d) { + LoadAsFloatRow(propagatedH_, HOffset(b, hv, chunkIdx, d, 0), hRow, V_); + float weight = coefficients.GetValue(d); + SetFlag(sToVEvent_); + WaitFlag(sToVEvent_); + Muls(hRow, hRow, weight, static_cast(V_)); + PipeBarrier(); + Add(dstRow, dstRow, hRow, static_cast(V_)); + PipeBarrier(); + SetFlag(vToMte2Event_); + WaitFlag(vToMte2Event_); + } + SetFlag(sToMte2Event_); + WaitFlag(sToMte2Event_); + } + } + +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + __aicore__ inline void ComputeTailOutputRegbaseRows(LocalTensor &dst, uint64_t b, uint64_t hv, + uint64_t chunkIdx, uint64_t start, uint64_t curT, + uint64_t rowBegin, uint64_t rows) + { + // Keep sub-16 tails on one MTE2-to-V handoff; the scalar event chain can retain stale state. + constexpr uint64_t dim = 128; + constexpr uint64_t maxTailRows = KDA_CUBE_MIN_REDUCTION - 1; + constexpr uint64_t aqkStagedCols = KDA_CUBE_MIN_REDUCTION; + constexpr uint64_t hOffset = 0; + constexpr uint64_t vNewOffset = hOffset + dim * dim; + constexpr uint64_t qgOffset = vNewOffset + maxTailRows * dim; + constexpr uint64_t aqkOffset = qgOffset + maxTailRows * dim; + LocalTensor staging = gateWritebackBuf_.Get(); + LocalTensor hLocal = staging[hOffset]; + LocalTensor vNewLocal = staging[vNewOffset]; + LocalTensor qgLocal = staging[qgOffset]; + LocalTensor aqkLocal = staging[aqkOffset]; + + CopyVectorIn(hLocal, propagatedH_, HOffset(b, hv, chunkIdx, 0, 0), dim * dim); + CopyVectorIn(vNewLocal, propagatedVNew_, KVOffset(b, hv, start, 0, V_), curT * dim); + CopyVectorIn(qgLocal, preparedQG_, KVOffset(b, hv, start + rowBegin, 0, K_), rows * dim); + for (uint64_t row = 0; row < rows; ++row) { + LocalTensor aqkRow = aqkLocal[row * aqkStagedCols]; + CopyVectorIn(aqkRow, preparedAqk_, AOffset(b, hv, start + rowBegin + row, 0), aqkStagedCols); + } + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + AscendC::VF_CALL(reinterpret_cast<__ubuf__ float *>(dst.GetPhyAddr()), + reinterpret_cast<__ubuf__ bfloat16_t *>(hLocal.GetPhyAddr()), + reinterpret_cast<__ubuf__ bfloat16_t *>(vNewLocal.GetPhyAddr()), + reinterpret_cast<__ubuf__ bfloat16_t *>(qgLocal.GetPhyAddr()), + reinterpret_cast<__ubuf__ bfloat16_t *>(aqkLocal.GetPhyAddr()), + static_cast(rows), static_cast(curT), scale_); + PipeBarrier(); + } +#endif + + template + __aicore__ inline void LoadAsFloatVector(GlobalTensor &src, uint64_t srcOffset, LocalTensor &dst, + LocalTensor &typedScratch, uint64_t count) + { + if constexpr (IsSameType::value) { + CopyVectorIn(dst, src, srcOffset, count); + } else { + CopyVectorIn(typedScratch, src, srcOffset, count); + } + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + if constexpr (!IsSameType::value) { + Cast(dst, typedScratch, RoundMode::CAST_NONE, static_cast(count)); + PipeBarrier(); + } + } + + template + __aicore__ inline void StoreFloatRow(GlobalTensor &dst, uint64_t dstOffset, LocalTensor &src, + uint64_t count) + { + if constexpr (IsSameType::value) { + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + CopyVectorOut(dst, dstOffset, src, count); + } else { + constexpr uint32_t typedOffset = EXP2_UB_ELEMENTS * sizeof(float) / sizeof(CopyT); + LocalTensor rowLocal = exp2Buf_.Get()[typedOffset]; + Cast(rowLocal, src, RoundMode::CAST_RINT, static_cast(count)); + PipeBarrier(); + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + CopyVectorOut(dst, dstOffset, rowLocal, count); + } + SetFlag(mte3ToMte2Event_); + WaitFlag(mte3ToMte2Event_); + SetFlag(mte3ToVEvent_); + WaitFlag(mte3ToVEvent_); + } + + __aicore__ inline LocalTensor Exp2NegG(uint64_t b, uint64_t hv, uint64_t t) + { + LocalTensor exp2Local = exp2Buf_.Get(); + LoadAsFloatRow(gk_, KVOffset(b, hv, t, 0, K_), exp2Local, K_); + Muls(exp2Local, exp2Local, -LN2, static_cast(K_)); + PipeBarrier(); + RunExp2(exp2Local, static_cast(K_)); + return exp2Local; + } + + __aicore__ inline uint64_t ScoreVectorMaxRows(uint64_t bytesPerElem) const + { + constexpr uint64_t arenaBytes = static_cast(KDA_VEC_ARENA_ELEMENTS) * sizeof(float); + uint64_t maxRows = (arenaBytes / bytesPerElem) / K_; + if (K_ >= 128 && maxRows > 32) { + maxRows = 32; + } + return maxRows; + } +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + __aicore__ inline void ComputeOutputCubeStagedArch35(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, + uint64_t curT) + { + SetMMLayoutTransform(true); + using ElementA = T; + using ElementB = T; + using ElementC = OUT_T; + using LayoutTagA = Catlass::layout::RowMajor; + using LayoutTagB = Catlass::layout::RowMajor; + using LayoutTagC = Catlass::layout::RowMajor; + using TileCopy = Catlass::Gemm::Tile::PackedTileCopyTla; + using LayoutTagL1A = typename TileCopy::LayoutTagL1A; + using LayoutTagL1B = typename TileCopy::LayoutTagL1B; + using LayoutTagL0A = typename TileCopy::LayoutTagL0A; + using LayoutTagL0B = typename TileCopy::LayoutTagL0B; + using CopyL1ToL0A = typename TileCopy::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopy::CopyL1ToL0B; + using TileMmad = Catlass::Gemm::Tile::TileMmadTla; + + constexpr uint16_t kMte2Event = 0; + constexpr uint16_t kMte1Event = 0; + constexpr uint16_t kMmadEvent = 0; + constexpr uint32_t kL1A0Offset = 0; + constexpr uint32_t kL1B0Offset = 64 * 128 * sizeof(ElementA); + constexpr uint32_t kL1A1Offset = kL1B0Offset + 128 * 128 * sizeof(ElementB); + constexpr uint32_t kL1B1Offset = kL1A1Offset + 64 * 64 * sizeof(ElementA); + + Catlass::Arch::Resource resource; + LocalTensor l1A0 = resource.l1Buf.template GetBufferByByte(kL1A0Offset); + LocalTensor l1B0 = resource.l1Buf.template GetBufferByByte(kL1B0Offset); + LocalTensor l1A1 = resource.l1Buf.template GetBufferByByte(kL1A1Offset); + LocalTensor l1B1 = resource.l1Buf.template GetBufferByByte(kL1B1Offset); + LocalTensor l0A = resource.l0ABuf.template GetBufferByByte(0); + LocalTensor l0B = resource.l0BBuf.template GetBufferByByte(0); + LocalTensor l0C = resource.l0CBuf.template GetBufferByByte(0); + + CopyL1ToL0A copyL1ToL0A; + CopyL1ToL0B copyL1ToL0B; + TileMmad tileMmad; + + auto layoutQ = tla::MakeLayout(BT_, K_); + auto layoutH = tla::MakeLayout(K_, V_); + auto layoutO = tla::MakeLayout(BT_, V_); + auto layoutAqk = tla::MakeLayout(BT_, BT_); + auto layoutV = tla::MakeLayout(BT_, V_); + + for (uint64_t nOffset = 0; nOffset < V_; nOffset += 128) { + const uint32_t curN = static_cast((V_ - nOffset) > 128 ? 128 : (V_ - nOffset)); + auto tensorH = tla::MakeTensor(propagatedH_[HOffset(b, hv, chunkIdx, 0, nOffset)], layoutH, + Catlass::Arch::PositionGM{}); + auto tensorVNew = tla::MakeTensor(propagatedVNew_[KVOffset(b, hv, start, nOffset, V_)], layoutV, + Catlass::Arch::PositionGM{}); + + for (uint64_t mOffset = 0; mOffset < curT; mOffset += 64) { + const uint32_t curM = static_cast((curT - mOffset) > 64 ? 64 : (curT - mOffset)); + auto tensorQ = tla::MakeTensor(preparedQG_[KVOffset(b, hv, start + mOffset, 0, K_)], layoutQ, + Catlass::Arch::PositionGM{}); + auto tensorAqk = tla::MakeTensor(preparedAqk_[AOffset(b, hv, start + mOffset, 0)], layoutAqk, + Catlass::Arch::PositionGM{}); + auto tensorO = tla::MakeTensor(o_[KVOffset(b, hv, start + mOffset, nOffset, V_)], layoutO, + Catlass::Arch::PositionGM{}); + auto tensorLocal = tla::MakeTensor(u_[KVOffset(b, hv, start + mOffset, nOffset, V_)], layoutO, + Catlass::Arch::PositionGM{}); + + auto blockQ = GetTile(tensorQ, tla::MakeCoord(0, 0), tla::MakeShape(curM, K_)); + auto blockH = GetTile(tensorH, tla::MakeCoord(0, 0), tla::MakeShape(K_, curN)); + auto blockAqk = GetTile(tensorAqk, tla::MakeCoord(0, 0), tla::MakeShape(curM, curT)); + auto blockVNew = GetTile(tensorVNew, tla::MakeCoord(0, 0), tla::MakeShape(curT, curN)); + auto blockO = GetTile(tensorO, tla::MakeCoord(0, 0), tla::MakeShape(curM, curN)); + auto blockLocal = GetTile(tensorLocal, tla::MakeCoord(0, 0), tla::MakeShape(curM, curN)); + + using CopyGmToL1A0 = typename TileCopy::template CopyGmToL1A; + using CopyGmToL1B0 = typename TileCopy::template CopyGmToL1B; + using CopyGmToL1A1 = typename TileCopy::template CopyGmToL1A; + using CopyGmToL1B1 = typename TileCopy::template CopyGmToL1B; + using CopyL0CToDst = typename TileCopy::template CopyL0CToDst; + CopyGmToL1A0 copyGmToL1A0; + CopyGmToL1B0 copyGmToL1B0; + CopyGmToL1A1 copyGmToL1A1; + CopyGmToL1B1 copyGmToL1B1; + CopyL0CToDst copyL0CToDst; + + auto layoutL1A0 = tla::MakeLayout(curM, K_); + auto layoutL1B0 = tla::MakeLayout(K_, curN); + auto layoutL1A1 = tla::MakeLayout(curM, curT); + auto layoutL1B1 = tla::MakeLayout(curT, curN); + auto layoutL0A0 = tla::MakeLayout(curM, K_); + auto layoutL0B0 = tla::MakeLayout(K_, curN); + auto layoutL0A1 = tla::MakeLayout(curM, curT); + auto layoutL0B1 = tla::MakeLayout(curT, curN); + auto layoutL0C = tla::MakeLayoutL0C(curM, curN); + + auto tensorL1A0 = tla::MakeTensor(l1A0, layoutL1A0, Catlass::Arch::PositionL1{}); + auto tensorL1B0 = tla::MakeTensor(l1B0, layoutL1B0, Catlass::Arch::PositionL1{}); + auto tensorL1A1 = tla::MakeTensor(l1A1, layoutL1A1, Catlass::Arch::PositionL1{}); + auto tensorL1B1 = tla::MakeTensor(l1B1, layoutL1B1, Catlass::Arch::PositionL1{}); + auto tensorL0A0 = tla::MakeTensor(l0A, layoutL0A0, Catlass::Arch::PositionL0A{}); + auto tensorL0B0 = tla::MakeTensor(l0B, layoutL0B0, Catlass::Arch::PositionL0B{}); + auto tensorL0A1 = tla::MakeTensor(l0A, layoutL0A1, Catlass::Arch::PositionL0A{}); + auto tensorL0B1 = tla::MakeTensor(l0B, layoutL0B1, Catlass::Arch::PositionL0B{}); + auto tensorL0C = tla::MakeTensor(l0C, layoutL0C, Catlass::Arch::PositionL0C{}); + uint32_t localRow = 0; + uint32_t localColumn = 0; + auto tileL1A0 = GetTile(tensorL1A0, tla::MakeCoord(localRow, localColumn), tla::MakeShape(curM, K_)); + auto tileL1B0 = GetTile(tensorL1B0, tla::MakeCoord(localRow, localColumn), tla::MakeShape(K_, curN)); + auto tileL1A1 = GetTile(tensorL1A1, tla::MakeCoord(localRow, localColumn), tla::MakeShape(curM, curT)); + auto tileL1B1 = GetTile(tensorL1B1, tla::MakeCoord(localRow, localColumn), tla::MakeShape(curT, curN)); + auto tileL0A0 = GetTile(tensorL0A0, tla::MakeCoord(localRow, localColumn), tla::MakeShape(curM, K_)); + auto tileL0B0 = GetTile(tensorL0B0, tla::MakeCoord(localRow, localColumn), tla::MakeShape(K_, curN)); + auto tileL0A1 = GetTile(tensorL0A1, tla::MakeCoord(localRow, localColumn), tla::MakeShape(curM, curT)); + auto tileL0B1 = GetTile(tensorL0B1, tla::MakeCoord(localRow, localColumn), tla::MakeShape(curT, curN)); + auto tileL0C = GetTile(tensorL0C, tla::MakeCoord(localRow, localColumn), tla::MakeShape(curM, curN)); + + copyGmToL1A0(tensorL1A0, blockQ); + copyGmToL1B0(tensorL1B0, blockH); + copyGmToL1A1(tensorL1A1, blockAqk); + copyGmToL1B1(tensorL1B1, blockVNew); + SetFlag(kMte2Event); + WaitFlag(kMte2Event); + + copyL1ToL0A(tileL0A0, tileL1A0); + copyL1ToL0B(tileL0B0, tileL1B0); + SetFlag(kMte1Event); + WaitFlag(kMte1Event); + tileMmad(tileL0C, tileL0A0, tileL0B0, curM, curN, static_cast(K_), true, 0b11); + SetFlag(kMmadEvent); + WaitFlag(kMmadEvent); + copyL0CToDst(blockO, tileL0C, 0b11); + PipeBarrier(); + + copyL1ToL0A(tileL0A1, tileL1A1); + copyL1ToL0B(tileL0B1, tileL1B1); + SetFlag(kMte1Event); + SetFlag(kMte2Event); + WaitFlag(kMte1Event); + WaitFlag(kMte2Event); + tileMmad(tileL0C, tileL0A1, tileL0B1, curM, curN, static_cast(curT), true, 0b11); + SetFlag(kMmadEvent); + WaitFlag(kMmadEvent); + copyL0CToDst(blockLocal, tileL0C, 0b11); + PipeBarrier(); + } + } + SetMMLayoutTransform(false); + } + + __aicore__ inline void PrefetchOutputTileArch35(Catlass::Arch::Resource &resource, uint32_t slot, + uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, + uint64_t curT, uint64_t nOffset, bool reuseSlot) + { + using ElementA = T; + using ElementB = T; + using ElementC = OUT_T; + using LayoutTagA = Catlass::layout::RowMajor; + using LayoutTagB = Catlass::layout::RowMajor; + using LayoutTagC = Catlass::layout::RowMajor; + using TileCopy = Catlass::Gemm::Tile::PackedTileCopyTla; + using LayoutTagL1A = typename TileCopy::LayoutTagL1A; + using LayoutTagL1B = typename TileCopy::LayoutTagL1B; + + constexpr uint32_t kL1SlotBytes = 96 * 1024; + constexpr uint32_t kL1A0Offset = 0; + constexpr uint32_t kL1B0Offset = 64 * 128 * sizeof(ElementA); + constexpr uint32_t kL1A1Offset = kL1B0Offset + 128 * 128 * sizeof(ElementB); + constexpr uint32_t kL1B1Offset = kL1A1Offset + 64 * 64 * sizeof(ElementA); + const uint32_t slotBase = slot * kL1SlotBytes; + const uint32_t curM = static_cast(curT); + const uint32_t curN = static_cast((V_ - nOffset) > 128 ? 128 : (V_ - nOffset)); + + auto layoutQ = tla::MakeLayout(BT_, K_); + auto layoutH = tla::MakeLayout(K_, V_); + auto layoutAqk = tla::MakeLayout(BT_, BT_); + auto layoutV = tla::MakeLayout(BT_, V_); + auto tensorQ = + tla::MakeTensor(preparedQG_[KVOffset(b, hv, start, 0, K_)], layoutQ, Catlass::Arch::PositionGM{}); + auto tensorH = + tla::MakeTensor(propagatedH_[HOffset(b, hv, chunkIdx, 0, nOffset)], layoutH, Catlass::Arch::PositionGM{}); + auto tensorAqk = + tla::MakeTensor(preparedAqk_[AOffset(b, hv, start, 0)], layoutAqk, Catlass::Arch::PositionGM{}); + auto tensorVNew = + tla::MakeTensor(propagatedVNew_[KVOffset(b, hv, start, nOffset, V_)], layoutV, Catlass::Arch::PositionGM{}); + auto blockQ = GetTile(tensorQ, tla::MakeCoord(0, 0), tla::MakeShape(curM, K_)); + auto blockH = GetTile(tensorH, tla::MakeCoord(0, 0), tla::MakeShape(K_, curN)); + auto blockAqk = GetTile(tensorAqk, tla::MakeCoord(0, 0), tla::MakeShape(curM, curT)); + auto blockVNew = GetTile(tensorVNew, tla::MakeCoord(0, 0), tla::MakeShape(curT, curN)); + + using CopyGmToL1A0 = typename TileCopy::template CopyGmToL1A; + using CopyGmToL1B0 = typename TileCopy::template CopyGmToL1B; + using CopyGmToL1A1 = typename TileCopy::template CopyGmToL1A; + using CopyGmToL1B1 = typename TileCopy::template CopyGmToL1B; + CopyGmToL1A0 copyGmToL1A0; + CopyGmToL1B0 copyGmToL1B0; + CopyGmToL1A1 copyGmToL1A1; + CopyGmToL1B1 copyGmToL1B1; + + LocalTensor l1A0 = resource.l1Buf.template GetBufferByByte(slotBase + kL1A0Offset); + LocalTensor l1B0 = resource.l1Buf.template GetBufferByByte(slotBase + kL1B0Offset); + LocalTensor l1A1 = resource.l1Buf.template GetBufferByByte(slotBase + kL1A1Offset); + LocalTensor l1B1 = resource.l1Buf.template GetBufferByByte(slotBase + kL1B1Offset); + auto tensorL1A0 = + tla::MakeTensor(l1A0, tla::MakeLayout(curM, K_), Catlass::Arch::PositionL1{}); + auto tensorL1B0 = + tla::MakeTensor(l1B0, tla::MakeLayout(K_, curN), Catlass::Arch::PositionL1{}); + auto tensorL1A1 = + tla::MakeTensor(l1A1, tla::MakeLayout(curM, curT), Catlass::Arch::PositionL1{}); + auto tensorL1B1 = + tla::MakeTensor(l1B1, tla::MakeLayout(curT, curN), Catlass::Arch::PositionL1{}); + + if (reuseSlot) { + WaitFlag(slot); + } + copyGmToL1A0(tensorL1A0, blockQ); + copyGmToL1B0(tensorL1B0, blockH); + copyGmToL1A1(tensorL1A1, blockAqk); + copyGmToL1B1(tensorL1B1, blockVNew); + SetFlag(slot); + } + + __aicore__ inline void ComputePrefetchedOutputTileArch35(Catlass::Arch::Resource &resource, + uint32_t slot, uint64_t b, uint64_t hv, uint64_t start, + uint64_t curT, uint64_t nOffset) + { + using ElementA = T; + using ElementB = T; + using ElementC = OUT_T; + using LayoutTagA = Catlass::layout::RowMajor; + using LayoutTagB = Catlass::layout::RowMajor; + using LayoutTagC = Catlass::layout::RowMajor; + using TileCopy = Catlass::Gemm::Tile::PackedTileCopyTla; + using LayoutTagL1A = typename TileCopy::LayoutTagL1A; + using LayoutTagL1B = typename TileCopy::LayoutTagL1B; + using LayoutTagL0A = typename TileCopy::LayoutTagL0A; + using LayoutTagL0B = typename TileCopy::LayoutTagL0B; + using CopyL1ToL0A = typename TileCopy::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopy::CopyL1ToL0B; + using TileMmad = Catlass::Gemm::Tile::TileMmadTla; + + constexpr uint16_t kMte1Event = 0; + constexpr uint16_t kMmadEvent = 0; + constexpr uint16_t kFixEvent = 0; + constexpr uint32_t kL1SlotBytes = 96 * 1024; + constexpr uint32_t kL1A0Offset = 0; + constexpr uint32_t kL1B0Offset = 64 * 128 * sizeof(ElementA); + constexpr uint32_t kL1A1Offset = kL1B0Offset + 128 * 128 * sizeof(ElementB); + constexpr uint32_t kL1B1Offset = kL1A1Offset + 64 * 64 * sizeof(ElementA); + const uint32_t slotBase = slot * kL1SlotBytes; + const uint32_t curM = static_cast(curT); + const uint32_t curN = static_cast((V_ - nOffset) > 128 ? 128 : (V_ - nOffset)); + + LocalTensor l1A0 = resource.l1Buf.template GetBufferByByte(slotBase + kL1A0Offset); + LocalTensor l1B0 = resource.l1Buf.template GetBufferByByte(slotBase + kL1B0Offset); + LocalTensor l1A1 = resource.l1Buf.template GetBufferByByte(slotBase + kL1A1Offset); + LocalTensor l1B1 = resource.l1Buf.template GetBufferByByte(slotBase + kL1B1Offset); + LocalTensor l0A = resource.l0ABuf.template GetBufferByByte(0); + LocalTensor l0B = resource.l0BBuf.template GetBufferByByte(0); + LocalTensor l0C = resource.l0CBuf.template GetBufferByByte(0); + + auto tensorL1A0 = + tla::MakeTensor(l1A0, tla::MakeLayout(curM, K_), Catlass::Arch::PositionL1{}); + auto tensorL1B0 = + tla::MakeTensor(l1B0, tla::MakeLayout(K_, curN), Catlass::Arch::PositionL1{}); + auto tensorL1A1 = + tla::MakeTensor(l1A1, tla::MakeLayout(curM, curT), Catlass::Arch::PositionL1{}); + auto tensorL1B1 = + tla::MakeTensor(l1B1, tla::MakeLayout(curT, curN), Catlass::Arch::PositionL1{}); + auto tensorL0A0 = + tla::MakeTensor(l0A, tla::MakeLayout(curM, K_), Catlass::Arch::PositionL0A{}); + auto tensorL0B0 = + tla::MakeTensor(l0B, tla::MakeLayout(K_, curN), Catlass::Arch::PositionL0B{}); + auto tensorL0A1 = + tla::MakeTensor(l0A, tla::MakeLayout(curM, curT), Catlass::Arch::PositionL0A{}); + auto tensorL0B1 = + tla::MakeTensor(l0B, tla::MakeLayout(curT, curN), Catlass::Arch::PositionL0B{}); + auto tensorL0C = tla::MakeTensor(l0C, tla::MakeLayoutL0C(curM, curN), Catlass::Arch::PositionL0C{}); + + uint32_t localRow = 0; + uint32_t localColumn = 0; + auto tileL1A0 = GetTile(tensorL1A0, tla::MakeCoord(localRow, localColumn), tla::MakeShape(curM, K_)); + auto tileL1B0 = GetTile(tensorL1B0, tla::MakeCoord(localRow, localColumn), tla::MakeShape(K_, curN)); + auto tileL1A1 = GetTile(tensorL1A1, tla::MakeCoord(localRow, localColumn), tla::MakeShape(curM, curT)); + auto tileL1B1 = GetTile(tensorL1B1, tla::MakeCoord(localRow, localColumn), tla::MakeShape(curT, curN)); + auto tileL0A0 = GetTile(tensorL0A0, tla::MakeCoord(localRow, localColumn), tla::MakeShape(curM, K_)); + auto tileL0B0 = GetTile(tensorL0B0, tla::MakeCoord(localRow, localColumn), tla::MakeShape(K_, curN)); + auto tileL0A1 = GetTile(tensorL0A1, tla::MakeCoord(localRow, localColumn), tla::MakeShape(curM, curT)); + auto tileL0B1 = GetTile(tensorL0B1, tla::MakeCoord(localRow, localColumn), tla::MakeShape(curT, curN)); + auto tileL0C = GetTile(tensorL0C, tla::MakeCoord(localRow, localColumn), tla::MakeShape(curM, curN)); + + auto layoutO = tla::MakeLayout(BT_, V_); + auto tensorO = tla::MakeTensor(o_[KVOffset(b, hv, start, nOffset, V_)], layoutO, Catlass::Arch::PositionGM{}); + auto tensorLocal = + tla::MakeTensor(u_[KVOffset(b, hv, start, nOffset, V_)], layoutO, Catlass::Arch::PositionGM{}); + auto blockO = GetTile(tensorO, tla::MakeCoord(0, 0), tla::MakeShape(curM, curN)); + auto blockLocal = GetTile(tensorLocal, tla::MakeCoord(0, 0), tla::MakeShape(curM, curN)); + using CopyL0CToDst = typename TileCopy::template CopyL0CToDst; + CopyL1ToL0A copyL1ToL0A; + CopyL1ToL0B copyL1ToL0B; + CopyL0CToDst copyL0CToDst; + TileMmad tileMmad; + + WaitFlag(slot); + if constexpr (IsSameType::value) { + WaitFlag(kFixEvent); + } + copyL1ToL0A(tileL0A0, tileL1A0); + copyL1ToL0B(tileL0B0, tileL1B0); + SetFlag(kMte1Event); + WaitFlag(kMte1Event); + tileMmad(tileL0C, tileL0A0, tileL0B0, curM, curN, static_cast(K_), true, 0b11); + SetFlag(kMmadEvent); + WaitFlag(kMmadEvent); + if constexpr (!IsSameType::value) { + copyL0CToDst(blockO, tileL0C, 0b11); + PipeBarrier(); + } + + copyL1ToL0A(tileL0A1, tileL1A1); + copyL1ToL0B(tileL0B1, tileL1B1); + SetFlag(kMte1Event); + SetFlag(slot); + WaitFlag(kMte1Event); + tileMmad(tileL0C, tileL0A1, tileL0B1, curM, curN, static_cast(curT), + !IsSameType::value, 0b11); + if constexpr (IsSameType::value) { + SetFlag(kFixEvent); + WaitFlag(kFixEvent); + auto fixParams = FixpipeParamsV220(curN, curM, curN, static_cast(HV_ * V_), false); + fixParams.quantPre = QuantMode_t::F322BF16; + Fixpipe(vNew_[OutputOffset(b, hv, start, nOffset)], l0C, fixParams); + SetFlag(kFixEvent); + } else { + SetFlag(kMmadEvent); + WaitFlag(kMmadEvent); + copyL0CToDst(blockLocal, tileL0C, 0b11); + PipeBarrier(); + } + } + + __aicore__ inline void ProcessOutAicPipelinedArch35() + { + SetLoadDataPaddingValue(static_cast(0)); + SetMMLayoutTransform(true); + uint64_t taskNum = static_cast((isVarLen_ ? NT_ : B_ * NT_) * HV_); + uint64_t coreNum = usedCoreNum_ == 0 ? 1 : usedCoreNum_; + uint64_t currentTask = static_cast(GetBlockIdx()); + uint64_t seq = 0; + uint64_t b = 0; + uint64_t h = 0; + uint64_t hv = 0; + uint64_t chunkIdx = 0; + uint64_t start = 0; + uint64_t end = 0; + while (currentTask < taskNum && !ResolveFlatChunk(currentTask, seq, b, h, hv, chunkIdx, start, end)) { + currentTask += coreNum; + } + if (currentTask >= taskNum) { + SetMMLayoutTransform(false); + return; + } + + Catlass::Arch::Resource resource; + uint64_t nOffset = 0; + uint32_t slot = 0; + if constexpr (IsSameType::value) { + SetFlag(0); + } + PrefetchOutputTileArch35(resource, slot, b, hv, chunkIdx, start, end - start, nOffset, false); + uint64_t outputTileIdx = 0; + + while (true) { + uint64_t nextTask = currentTask; + uint64_t nextSeq = seq; + uint64_t nextB = b; + uint64_t nextH = h; + uint64_t nextHv = hv; + uint64_t nextChunkIdx = chunkIdx; + uint64_t nextStart = start; + uint64_t nextEnd = end; + uint64_t nextNOffset = nOffset + 128; + bool hasNext = nextNOffset < V_; + if (!hasNext) { + nextTask += coreNum; + nextNOffset = 0; + while (nextTask < taskNum && + !ResolveFlatChunk(nextTask, nextSeq, nextB, nextH, nextHv, nextChunkIdx, nextStart, nextEnd)) { + nextTask += coreNum; + } + hasNext = nextTask < taskNum; + } + + const uint32_t nextSlot = slot ^ 1U; + if (hasNext) { + PrefetchOutputTileArch35(resource, nextSlot, nextB, nextHv, nextChunkIdx, nextStart, + nextEnd - nextStart, nextNOffset, outputTileIdx + 1 >= 2); + } + ComputePrefetchedOutputTileArch35(resource, slot, b, hv, start, end - start, nOffset); + if constexpr (!IsSameType::value) { + if (nOffset + 128 >= V_) { + Catlass::Arch::CrossCoreSetFlagWithReverse<0x2, PIPE_FIX>(syncDoneFlag_); + } + } + if (!hasNext) { + break; + } + ++outputTileIdx; + + currentTask = nextTask; + seq = nextSeq; + b = nextB; + h = nextH; + hv = nextHv; + chunkIdx = nextChunkIdx; + start = nextStart; + end = nextEnd; + nOffset = nextNOffset; + slot = nextSlot; + } + if constexpr (IsSameType::value) { + WaitFlag(0); + } + SetMMLayoutTransform(false); + } +#endif + + __aicore__ inline void ComputeOutputCube(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, uint64_t curT) + { +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if (curT < KDA_CUBE_MIN_REDUCTION) { + return; + } + SetLoadDataPaddingValue(static_cast(0)); + if (BT_ == 64 && curT == BT_) { + ComputeOutputCubeStagedArch35(b, hv, chunkIdx, start, curT); + return; + } +#endif + using ElementA = T; + using ElementB = T; + using ElementC = OUT_T; + using LayoutTagA = Catlass::layout::RowMajor; + using LayoutTagB = Catlass::layout::RowMajor; + using LayoutTagC = Catlass::layout::RowMajor; + using TileCopy = Catlass::Gemm::Tile::PackedTileCopyTla; + using BlockMmad = Catlass::Gemm::Block::BlockMmadTla; + + Catlass::Arch::Resource resource; + BlockMmad blockMmad(resource); + + auto layoutQ = tla::MakeLayout(BT_, K_); + auto layoutH = tla::MakeLayout(K_, V_); + auto layoutO = tla::MakeLayout(BT_, V_); + for (uint64_t nOffset = 0; nOffset < V_; nOffset += 128) { + uint32_t curN = static_cast((V_ - nOffset) > 128 ? 128 : (V_ - nOffset)); + auto tensorH = tla::MakeTensor(propagatedH_[HOffset(b, hv, chunkIdx, 0, nOffset)], layoutH, + Catlass::Arch::PositionGM{}); + for (uint64_t mOffset = 0; mOffset < curT; mOffset += 64) { + uint32_t curM = static_cast((curT - mOffset) > 64 ? 64 : (curT - mOffset)); + Catlass::GemmCoord shapeQH{curM, curN, static_cast(K_)}; + auto tensorQ = tla::MakeTensor(preparedQG_[KVOffset(b, hv, start + mOffset, 0, K_)], layoutQ, + Catlass::Arch::PositionGM{}); + auto tensorO = tla::MakeTensor(o_[KVOffset(b, hv, start + mOffset, nOffset, V_)], layoutO, + Catlass::Arch::PositionGM{}); + auto blockQ = GetTile(tensorQ, tla::MakeCoord(0, 0), tla::MakeShape(shapeQH.m(), shapeQH.k())); + auto blockH = GetTile(tensorH, tla::MakeCoord(0, 0), tla::MakeShape(shapeQH.k(), shapeQH.n())); + auto blockO = GetTile(tensorO, tla::MakeCoord(0, 0), tla::MakeShape(shapeQH.m(), shapeQH.n())); + blockMmad(blockQ, blockH, blockO, shapeQH); + PipeBarrier(); + } + } + + if (curT < KDA_CUBE_MIN_REDUCTION) { + return; + } + + auto layoutAqk = tla::MakeLayout(BT_, BT_); + auto layoutV = tla::MakeLayout(BT_, V_); + for (uint64_t nOffset = 0; nOffset < V_; nOffset += 128) { + uint32_t curN = static_cast((V_ - nOffset) > 128 ? 128 : (V_ - nOffset)); + auto tensorVNew = tla::MakeTensor(propagatedVNew_[KVOffset(b, hv, start, nOffset, V_)], layoutV, + Catlass::Arch::PositionGM{}); + for (uint64_t mOffset = 0; mOffset < curT; mOffset += 64) { + uint32_t curM = static_cast((curT - mOffset) > 64 ? 64 : (curT - mOffset)); + Catlass::GemmCoord shapeAV{curM, curN, static_cast(curT)}; + auto tensorAqk = tla::MakeTensor(preparedAqk_[AOffset(b, hv, start + mOffset, 0)], layoutAqk, + Catlass::Arch::PositionGM{}); + auto tensorLocal = tla::MakeTensor(u_[KVOffset(b, hv, start + mOffset, nOffset, V_)], layoutO, + Catlass::Arch::PositionGM{}); + auto blockAqk = GetTile(tensorAqk, tla::MakeCoord(0, 0), tla::MakeShape(shapeAV.m(), shapeAV.k())); + auto blockVNew = GetTile(tensorVNew, tla::MakeCoord(0, 0), tla::MakeShape(shapeAV.k(), shapeAV.n())); + auto blockLocal = GetTile(tensorLocal, tla::MakeCoord(0, 0), tla::MakeShape(shapeAV.m(), shapeAV.n())); + blockMmad(blockAqk, blockVNew, blockLocal, shapeAV); + PipeBarrier(); + } + } + } + + __aicore__ inline void FinalizeOutputRows(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, uint64_t curT, + uint64_t subBlockIdx, uint64_t subBlockNum) + { + if (subBlockNum == 0 || subBlockIdx >= subBlockNum || V_ == 0) { + return; + } + const uint64_t rowBegin = (curT * subBlockIdx) / subBlockNum; + const uint64_t rowEnd = (curT * (subBlockIdx + 1)) / subBlockNum; + const uint64_t gateWritebackRows = ScoreVectorMaxRows(5 * sizeof(float) + 2 * sizeof(T) + sizeof(GK_T)); + const uint64_t gateWritebackBytes = gateWritebackRows * K_ * (3 * sizeof(T) + sizeof(GK_T)); + uint64_t maxRows = KDA_VEC_ARENA_ELEMENTS / (3 * V_); + const uint64_t typedMaxRows = gateWritebackBytes / (V_ * sizeof(T)); + if (maxRows > typedMaxRows) { + maxRows = typedMaxRows; + } + if (maxRows == 0) { + return; + } + + for (uint64_t tileRow = rowBegin; tileRow < rowEnd; tileRow += maxRows) { + uint64_t tileRows = rowEnd - tileRow; + if (tileRows > maxRows) { + tileRows = maxRows; + } + const uint64_t elems = tileRows * V_; + const uint64_t ti = start + tileRow; + LocalTensor arena = vecBuf_.Get(); + LocalTensor stateLocal = arena; + LocalTensor localLocal = arena[elems]; + LocalTensor outLocal = arena[2 * elems]; + LocalTensor outTyped = gateWritebackBuf_.Get(); + + if (curT < KDA_CUBE_MIN_REDUCTION) { +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if constexpr (IsSameType::value) { + if (BT_ == 64 && K_ == 128 && V_ == 128) { + ComputeTailOutputRegbaseRows(outLocal, b, hv, chunkIdx, start, curT, tileRow, tileRows); + } else { + ComputeTailStateRows(stateLocal, b, hv, chunkIdx, start, tileRow, tileRows); + ComputeTailLocalRows(localLocal, b, hv, start, curT, tileRow, tileRows); + Muls(stateLocal, stateLocal, scale_, static_cast(elems)); + PipeBarrier(); + Add(outLocal, stateLocal, localLocal, static_cast(elems)); + PipeBarrier(); + } + } else { + ComputeTailStateRows(stateLocal, b, hv, chunkIdx, start, tileRow, tileRows); + ComputeTailLocalRows(localLocal, b, hv, start, curT, tileRow, tileRows); + Muls(stateLocal, stateLocal, scale_, static_cast(elems)); + PipeBarrier(); + Add(outLocal, stateLocal, localLocal, static_cast(elems)); + PipeBarrier(); + } +#else + ComputeTailStateRows(stateLocal, b, hv, chunkIdx, start, tileRow, tileRows); + ComputeTailLocalRows(localLocal, b, hv, start, curT, tileRow, tileRows); + Muls(stateLocal, stateLocal, scale_, static_cast(elems)); + PipeBarrier(); + Add(outLocal, stateLocal, localLocal, static_cast(elems)); + PipeBarrier(); +#endif + } else { + CopyVectorIn(stateLocal, o_, KVOffset(b, hv, ti, 0, V_), elems); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + CopyVectorIn(localLocal, u_, KVOffset(b, hv, ti, 0, V_), elems); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + Muls(stateLocal, stateLocal, scale_, static_cast(elems)); + PipeBarrier(); + Add(outLocal, stateLocal, localLocal, static_cast(elems)); + PipeBarrier(); + } + ClampFp32ToOutputType(outLocal, static_cast(elems)); + Cast(outTyped, outLocal, RoundMode::CAST_RINT, static_cast(elems)); + PipeBarrier(); + + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + CopyRowsOut(vNew_, OutputOffset(b, hv, ti, 0), outTyped, tileRows, V_, HV_ * V_); + SetFlag(mte3ToMte2Event_); + WaitFlag(mte3ToMte2Event_); + SetFlag(mte3ToVEvent_); + WaitFlag(mte3ToVEvent_); + } + } + __aicore__ inline bool ResolveFlatChunk(uint64_t task, uint64_t &seq, uint64_t &b, uint64_t &h, uint64_t &hv, + uint64_t &chunkIdx, uint64_t &start, uint64_t &end) + { + hv = task % HV_; + uint64_t flatChunk = task / HV_; + if (!isVarLen_) { + seq = flatChunk / NT_; + b = seq; + chunkIdx = flatChunk % NT_; + start = chunkIdx * BT_; + end = start + BT_; + if (end > T_) { + end = T_; + } + } else { + if (!KdaVarlen::ResolveChunkRange(cuSeqlensAddr_, chunkIndicesAddr_, N_, T_, BT_, flatChunk, seq, start, + end)) { + return false; + } + b = 0; + chunkIdx = flatChunk; + } + h = hv / (HV_ / H_); + return start < end; + } + + __aicore__ inline void ProcessChunkOutAiv(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, uint64_t end, + uint64_t subBlockIdx, uint64_t subBlockNum) + { + uint64_t curT = end - start; + if (curT == 0) { + return; + } + if constexpr (IsSameType::value) { + return; + } + Catlass::Arch::CrossCoreWaitFlagWithReverse<0x2, PIPE_MTE2>(syncDoneFlag_); + FinalizeOutputRows(b, hv, chunkIdx, start, curT, subBlockIdx, subBlockNum); + } + + __aicore__ inline void ProcessChunkOutAic(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, uint64_t end) + { + uint64_t curT = end - start; + if (curT == 0) { + return; + } + ComputeOutputCube(b, hv, chunkIdx, start, curT); + Catlass::Arch::CrossCoreSetFlagWithReverse<0x2, PIPE_FIX>(syncDoneFlag_); + } + + __aicore__ inline void ProcessOutAiv() + { + if constexpr (IsSameType::value) { + return; + } +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if constexpr (IsSameType::value) { + if (!isVarLen_ && T_ % BT_ == 0 && BT_ == 64 && K_ == 128 && V_ == 128) { + return; + } + } +#endif + uint64_t subBlockNum = static_cast(GetSubBlockNum()); + if (subBlockNum == 0) { + return; + } + uint64_t subBlockIdx = static_cast(GetSubBlockIdx()); + uint64_t coreNum = usedCoreNum_ == 0 ? 1 : usedCoreNum_; + uint64_t coreIdx = static_cast(GetBlockIdx()) / subBlockNum; + uint64_t taskNum = static_cast((isVarLen_ ? NT_ : B_ * NT_) * HV_); + for (uint64_t task = coreIdx; task < taskNum; task += coreNum) { + uint64_t seq = 0; + uint64_t b = 0; + uint64_t h = 0; + uint64_t hv = 0; + uint64_t chunkIdx = 0; + uint64_t start = 0; + uint64_t end = 0; + if (ResolveFlatChunk(task, seq, b, h, hv, chunkIdx, start, end)) { + (void)seq; + (void)h; + (void)chunkIdx; + ProcessChunkOutAiv(b, hv, chunkIdx, start, end, subBlockIdx, subBlockNum); + } + } + } + + __aicore__ inline void ProcessOutAic() + { + if constexpr (IsSameType::value) { + return; + } +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if (!isVarLen_ && T_ % BT_ == 0 && BT_ == 64 && K_ == 128 && V_ == 128) { + ProcessOutAicPipelinedArch35(); + return; + } +#endif + uint64_t taskNum = static_cast((isVarLen_ ? NT_ : B_ * NT_) * HV_); + uint64_t coreNum = usedCoreNum_ == 0 ? 1 : usedCoreNum_; + for (uint64_t task = GetBlockIdx(); task < taskNum; task += coreNum) { + uint64_t seq = 0; + uint64_t b = 0; + uint64_t h = 0; + uint64_t hv = 0; + uint64_t chunkIdx = 0; + uint64_t start = 0; + uint64_t end = 0; + if (ResolveFlatChunk(task, seq, b, h, hv, chunkIdx, start, end)) { + (void)seq; + (void)h; + ProcessChunkOutAic(b, hv, chunkIdx, start, end); + } + } + } + +private: + GlobalTensor q_; + GlobalTensor k_; + GlobalTensor v_; + GlobalTensor gk_; + GlobalTensor beta_; + GlobalTensor initialState_; + GlobalTensor o_; + GlobalTensor finalState_; + GlobalTensor aqk_; + GlobalTensor akk_; + GlobalTensor w_; + GlobalTensor u_; + GlobalTensor qg_; + GlobalTensor kg_; + GlobalTensor vNew_; + GlobalTensor h_; + GlobalTensor preparedQG_; + GlobalTensor preparedAqk_; + GlobalTensor propagatedVNew_; + GlobalTensor propagatedH_; + GlobalTensor solveWorkspace_; + GlobalTensor scoreWorkspace_; + TPipe *pipe_ = nullptr; + TBuf exp2Buf_; + TBuf vecBuf_; + TBuf gateWritebackBuf_; + TEventID mte2ToVEvent_ = 0; + TEventID vToMte2Event_ = 0; + TEventID vToMte3Event_ = 0; + TEventID mte3ToVEvent_ = 0; + TEventID mte2ToMte3Event_ = 0; + TEventID mte3ToMte2Event_ = 0; + TEventID vToSEvent_ = 0; + TEventID sToVEvent_ = 0; + TEventID sToMte2Event_ = 0; + bool vectorEventsAllocated_ = false; + Catlass::Arch::CrossCoreFlagWithReverse scoreReadyFlag_{KDA_SCORE_READY_FLAG0, + KDA_SCORE_READY_FLAG1}; + Catlass::Arch::CrossCoreFlagWithReverse scoreDoneFlag_{KDA_SCORE_DONE_FLAG0, + KDA_SCORE_DONE_FLAG1}; + // Score production is fully drained before solve starts, so the solve handshake can safely reuse + // the existing score flags without consuming additional hardware flag IDs. + Catlass::Arch::CrossCoreFlagWithReverse syncReadyFlag_{KDA_SCORE_READY_FLAG0, + KDA_SCORE_READY_FLAG1}; + Catlass::Arch::CrossCoreFlagWithReverse syncDoneFlag_{KDA_SCORE_DONE_FLAG0, + KDA_SCORE_DONE_FLAG1}; + uint64_t B_ = 0; + uint64_t N_ = 0; + uint64_t H_ = 0; + uint64_t HV_ = 0; + uint64_t T_ = 0; + uint64_t K_ = 0; + uint64_t V_ = 0; + uint64_t BT_ = 0; + uint64_t NT_ = 0; + float scale_ = 1.0f; + bool hasInitial_ = false; + bool isVarLen_ = false; + bool isAivOnly_ = false; + uint64_t usedCoreNum_ = 1; + uint64_t solveCoreIdx_ = 0; + __gm__ int64_t *chunkIndicesAddr_ = nullptr; + __gm__ int64_t *cuSeqlensAddr_ = nullptr; +}; +} // namespace + +template +__aicore__ inline void RunChunkKdaOutput(GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR gk, GM_ADDR beta, + GM_ADDR initialState, GM_ADDR cuSeqlens, GM_ADDR chunkIndices, + GM_ADDR qgScaled, GM_ADDR aqk, GM_ADDR propagatedVNew, GM_ADDR propagatedH, + GM_ADDR o, GM_ADDR userWorkspace, const TilingData &tiling, TPipe &pipe) +{ + GM_ADDR outputScratch = userWorkspace + tiling.outputScratchOffset; + uint64_t outputElements = static_cast(tiling.batch) * static_cast(tiling.vHeadNum) * + static_cast(tiling.seqlen) * static_cast(tiling.vHeadDim); + GM_ADDR stateScratch = outputScratch; + GM_ADDR localScratch = outputScratch + outputElements * sizeof(float); + if ASCEND_IS_AIC { + ChunkKdaFwdFinalizeKernel op; + op.Init(q, k, v, gk, beta, initialState, cuSeqlens, chunkIndices, qgScaled, aqk, propagatedVNew, propagatedH, + stateScratch, userWorkspace, aqk, userWorkspace, userWorkspace, localScratch, userWorkspace, + userWorkspace, o, propagatedH, outputScratch, tiling, &pipe, false); + op.ProcessAic(); + } + if ASCEND_IS_AIV { + ChunkKdaFwdFinalizeKernel op; + op.Init(q, k, v, gk, beta, initialState, cuSeqlens, chunkIndices, qgScaled, aqk, propagatedVNew, propagatedH, + stateScratch, userWorkspace, aqk, userWorkspace, userWorkspace, localScratch, userWorkspace, + userWorkspace, o, propagatedH, outputScratch, tiling, &pipe); + op.ProcessAiv(); + } +} + +} // namespace KdaFinalize diff --git a/csrc/chunk_kda_fwd/op_kernel/arch35/chunk_kda_fwd_fwd_h.h b/csrc/chunk_kda_fwd/op_kernel/arch35/chunk_kda_fwd_fwd_h.h new file mode 100644 index 0000000000..fcb979cc91 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/arch35/chunk_kda_fwd_fwd_h.h @@ -0,0 +1,1013 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CHUNK_KDA_FWD_ARCH35_FWD_H_H +#define CHUNK_KDA_FWD_ARCH35_FWD_H_H + +#include "kernel_operator.h" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" +#include "catlass/gemm/tile/kda_tile_mmad.hpp" +#include "../kernel_utils/tile/copy_l0c_to_ub.hpp" + +namespace KdaForward::arch35 { + +using namespace AscendC; + +// Prepare is fully drained before FwdH starts. Keep FwdH mode-2 +// traffic outside Matmul's possible 0..7 range and SyncAll's 11..14 range. +constexpr uint64_t KDA_FWD_H_DIRECT_FREE_FLAG = 8; +constexpr uint64_t KDA_FWD_H_DIRECT_READY_FLAG = 8; +constexpr uint64_t KDA_FWD_H_L1_FREE_FLAG = 9; +constexpr uint64_t KDA_FWD_H_L1_READY_FLAG = 9; +constexpr uint64_t KDA_FWD_H_STATE_FREE_FLAG = KDA_FWD_H_L1_FREE_FLAG; +constexpr uint64_t KDA_FWD_H_VNEW_FREE_FLAG = KDA_FWD_H_L1_FREE_FLAG; +constexpr uint64_t KDA_FWD_H_STATE_READY_FLAG = KDA_FWD_H_L1_READY_FLAG; +constexpr uint64_t KDA_FWD_H_VNEW_READY_FLAG = KDA_FWD_H_L1_READY_FLAG; + +constexpr TEventID KDA_FWD_H_MTE_W_EVENT = 0; +constexpr TEventID KDA_FWD_H_MTE_Q_EVENT = 1; +constexpr TEventID KDA_FWD_H_MTE_B_EVENT = 2; +constexpr TEventID KDA_FWD_H_MTE_A_EVENT = 3; +constexpr TEventID KDA_FWD_H_M_EVENT = 4; +constexpr TEventID KDA_FWD_H_FIX_EVENT = 5; +constexpr TEventID KDA_FWD_H_IO_REUSE_EVENT = 6; + +constexpr uint32_t KDA_FWD_H_CHUNK = 64; +constexpr uint32_t KDA_FWD_H_DIM = 128; +constexpr uint32_t KDA_FWD_H_SUB_CHUNK = KDA_FWD_H_CHUNK / 2; +constexpr uint32_t KDA_FWD_H_SUB_DIM = KDA_FWD_H_DIM / 2; +constexpr uint32_t KDA_FWD_H_STATE_SUB_ELEMS = KDA_FWD_H_SUB_DIM * KDA_FWD_H_DIM; +constexpr uint32_t KDA_FWD_H_TOKEN_SUB_ELEMS = KDA_FWD_H_SUB_CHUNK * KDA_FWD_H_DIM; + +constexpr uint32_t KDA_FWD_H_L1_W_OFFSET = 0; +constexpr uint32_t KDA_FWD_H_L1_Q_OFFSET = 16 * 1024; +constexpr uint32_t KDA_FWD_H_L1_H_OFFSET = 32 * 1024; +constexpr uint32_t KDA_FWD_H_L1_KG_OFFSET = 64 * 1024; +constexpr uint32_t KDA_FWD_H_L1_AQK_OFFSET = 80 * 1024; +constexpr uint32_t KDA_FWD_H_L1_V_OFFSET = 96 * 1024; +constexpr uint32_t KDA_FWD_H_L1_W1_OFFSET = 112 * 1024; +constexpr uint32_t KDA_FWD_H_L1_Q1_OFFSET = 128 * 1024; +constexpr uint32_t KDA_FWD_H_L1_KG1_OFFSET = 144 * 1024; +constexpr uint32_t KDA_FWD_H_L1_AQK1_OFFSET = 160 * 1024; +constexpr uint32_t KDA_FWD_H_L1_W2_OFFSET = 176 * 1024; +constexpr uint32_t KDA_FWD_H_L1_Q2_OFFSET = 192 * 1024; +constexpr uint32_t KDA_FWD_H_L1_KG2_OFFSET = 208 * 1024; +constexpr uint32_t KDA_FWD_H_L1_AQK2_OFFSET = 224 * 1024; +constexpr uint32_t KDA_FWD_H_L1_W3_OFFSET = 240 * 1024; +constexpr uint32_t KDA_FWD_H_L1_Q3_OFFSET = 256 * 1024; +constexpr uint32_t KDA_FWD_H_L1_KG3_OFFSET = 272 * 1024; +constexpr uint32_t KDA_FWD_H_L1_AQK3_OFFSET = 288 * 1024; +constexpr uint32_t KDA_FWD_H_L1_AKK_OFFSET = 304 * 1024; +constexpr uint32_t KDA_FWD_H_L1_U_OFFSET = 336 * 1024; +constexpr uint32_t KDA_FWD_H_L1_STAGING_DEPTH = 4; +constexpr uint32_t KDA_FWD_H_L1_AKK_SLOT_BYTES = 8 * 1024; +constexpr uint32_t KDA_FWD_H_L1_U_SLOT_BYTES = 16 * 1024; + +constexpr uint32_t KDA_FWD_H_L0A_STATE_OFFSET = 0; +constexpr uint32_t KDA_FWD_H_L0A_VNEW_OFFSET = 16 * 1024; +constexpr uint32_t KDA_FWD_H_L0A_POST_OFFSET = 32 * 1024; +constexpr uint32_t KDA_FWD_H_L0B_STATE_OFFSET = 0; +constexpr uint32_t KDA_FWD_H_L0B_VNEW_OFFSET = 32 * 1024; +constexpr uint32_t KDA_FWD_H_L0B_POST_OFFSET = 32 * 1024; + +constexpr uint32_t KDA_FWD_H_UB_STATE_OFFSET = 0; +constexpr uint32_t KDA_FWD_H_UB_STATE_TYPED_OFFSET = 32 * 1024; +constexpr uint32_t KDA_FWD_H_UB_DIRECT_OFFSET = 48 * 1024; +constexpr uint32_t KDA_FWD_H_UB_OUT1_OFFSET = 80 * 1024; +constexpr uint32_t KDA_FWD_H_UB_VNEW_OFFSET = 96 * 1024; +constexpr uint32_t KDA_FWD_H_UB_IO_OFFSET = 112 * 1024; +constexpr uint32_t KDA_FWD_H_UB_GATE_OFFSET = 128 * 1024; + +template +class ChunkKdaFwdFwdH { +public: + using ArchTag = Catlass::Arch::Ascend950; + using LayoutRM = Catlass::layout::RowMajor; + using LayoutCM = Catlass::layout::ColumnMajor; + using TileCopyRM = Catlass::Gemm::Tile::PackedTileCopyTla< + ArchTag, T, LayoutRM, T, LayoutRM, float, LayoutRM>; + using DirectTileCopyRM = Common::Tile::PackedTileCopyTlaToUB< + ArchTag, T, LayoutRM, T, LayoutRM, float, LayoutRM, void, + Catlass::Gemm::Tile::CopyL0CToUBMode::SPLIT_M>; + using TileCopyCM = Catlass::Gemm::Tile::PackedTileCopyTla< + ArchTag, T, LayoutCM, T, LayoutRM, float, LayoutRM>; + using DirectTileCopyCM = Common::Tile::PackedTileCopyTlaToUB< + ArchTag, T, LayoutCM, T, LayoutRM, float, LayoutRM, void, + Catlass::Gemm::Tile::CopyL0CToUBMode::SPLIT_M>; + + __aicore__ inline void Init( + GM_ADDR gk, GM_ADDR initialState, GM_ADDR attnOut, GM_ADDR finalState, + GM_ADDR aqk, GM_ADDR akk, GM_ADDR w, GM_ADDR u, GM_ADDR qg, GM_ADDR kg, + GM_ADDR vNew, GM_ADDR h, const TilingData &tiling) + { + gk_.SetGlobalBuffer(reinterpret_cast<__gm__ GK_T *>(gk)); + if (initialState != nullptr) { + initialState_.SetGlobalBuffer(reinterpret_cast<__gm__ float *>(initialState)); + } + attnOut_.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(attnOut)); + finalState_.SetGlobalBuffer(reinterpret_cast<__gm__ float *>(finalState)); + aqk_.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(aqk)); + akk_.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(akk)); + w_.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(w)); + u_.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(u)); + qg_.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(qg)); + kg_.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(kg)); + vNew_.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(vNew)); + h_.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(h)); + batch_ = tiling.batch; + heads_ = tiling.vHeadNum; + seqlen_ = tiling.seqlen; + totalChunks_ = tiling.totalChunks; + hasInitialState_ = tiling.hasInitialState; + storeFinalState_ = tiling.storeFinalState; + storeVNew_ = tiling.storeVNew; + storeH_ = tiling.storeH; + fusePostWuIntoFwdH_ = tiling.fusePostWuIntoFwdH; + scale_ = tiling.scale; + coreNum_ = tiling.prepareUsedCoreNum; + statePublishCount_[0] = 0; + statePublishCount_[1] = 0; + vnewPublishCount_[0] = 0; + vnewPublishCount_[1] = 0; + wPublishCount_[0] = 0; + wPublishCount_[1] = 0; + } + + __aicore__ inline void Process() + { + if ASCEND_IS_AIC { + ProcessAic(); + } + if ASCEND_IS_AIV { + ProcessAiv(); + } + } + +private: + __aicore__ inline void WaitL1SlotFreeMte3( + uint64_t freeFlag, uint32_t publishCount) + { + if (publishCount != 0) { + CrossCoreWaitFlag(freeFlag); + } + } + + template + __aicore__ inline void SetL1SlotFlagAicToAiv(uint64_t flag) + { + CrossCoreSetFlag<0x2, PIPE>(flag); + } + + template + __aicore__ inline void SetL1SlotFlagAivToAic(uint64_t flag) + { + CrossCoreSetFlag<0x2, PIPE>(flag); + } + + __aicore__ inline void WaitL1SlotReadyMte1(uint64_t readyFlag) + { + CrossCoreWaitFlag(readyFlag); + } + + __aicore__ inline uint64_t MatrixOffset(uint64_t b, uint64_t hv, uint64_t t) const + { + return ((b * heads_ + hv) * seqlen_ + t) * KDA_FWD_H_DIM; + } + + __aicore__ inline uint64_t ChunkMatrixOffset( + uint64_t b, uint64_t hv, uint64_t chunk, uint64_t row = 0) const + { + return ((b * heads_ + hv) * seqlen_ + + chunk * KDA_FWD_H_CHUNK + row) * + KDA_FWD_H_DIM; + } + + __aicore__ inline uint64_t ScoreOffset(uint64_t b, uint64_t hv, uint64_t t) const + { + return ((b * heads_ + hv) * seqlen_ + t) * KDA_FWD_H_CHUNK; + } + + __aicore__ inline uint64_t ChunkScoreOffset( + uint64_t b, uint64_t hv, uint64_t chunk) const + { + return ((b * heads_ + hv) * seqlen_ + + chunk * KDA_FWD_H_CHUNK) * + KDA_FWD_H_CHUNK; + } + + __aicore__ inline uint64_t StateOffset(uint64_t b, uint64_t hv) const + { + return (b * heads_ + hv) * KDA_FWD_H_DIM * KDA_FWD_H_DIM; + } + + __aicore__ inline uint64_t HOffset(uint64_t b, uint64_t hv, uint64_t chunk) const + { + if (!storeH_) { + return StateOffset(b, hv); + } + return ((b * heads_ + hv) * totalChunks_ + chunk) * + KDA_FWD_H_DIM * KDA_FWD_H_DIM; + } + + __aicore__ inline uint64_t VNewOffset( + uint64_t b, uint64_t hv, uint64_t chunk, uint64_t row = 0) const + { + if (!storeVNew_) { + return ((b * heads_ + hv) * KDA_FWD_H_CHUNK + row) * KDA_FWD_H_DIM; + } + return ChunkMatrixOffset(b, hv, chunk, row); + } + + __aicore__ inline uint64_t OutputOffset(uint64_t b, uint64_t hv, uint64_t t) const + { + return ((b * seqlen_ + t) * heads_ + hv) * KDA_FWD_H_DIM; + } + + __aicore__ inline uint32_t L1WOffset(uint32_t slot) const + { + if (slot == 0) { + return KDA_FWD_H_L1_W_OFFSET; + } + if (slot == 1) { + return KDA_FWD_H_L1_W1_OFFSET; + } + return slot == 2 ? KDA_FWD_H_L1_W2_OFFSET : KDA_FWD_H_L1_W3_OFFSET; + } + + __aicore__ inline uint32_t L1QOffset(uint32_t slot) const + { + if (slot == 0) { + return KDA_FWD_H_L1_Q_OFFSET; + } + if (slot == 1) { + return KDA_FWD_H_L1_Q1_OFFSET; + } + return slot == 2 ? KDA_FWD_H_L1_Q2_OFFSET : KDA_FWD_H_L1_Q3_OFFSET; + } + + __aicore__ inline uint32_t L1KgOffset(uint32_t slot) const + { + if (slot == 0) { + return KDA_FWD_H_L1_KG_OFFSET; + } + if (slot == 1) { + return KDA_FWD_H_L1_KG1_OFFSET; + } + return slot == 2 ? KDA_FWD_H_L1_KG2_OFFSET : KDA_FWD_H_L1_KG3_OFFSET; + } + + __aicore__ inline uint32_t L1AqkOffset(uint32_t slot) const + { + if (slot == 0) { + return KDA_FWD_H_L1_AQK_OFFSET; + } + if (slot == 1) { + return KDA_FWD_H_L1_AQK1_OFFSET; + } + return slot == 2 ? KDA_FWD_H_L1_AQK2_OFFSET : KDA_FWD_H_L1_AQK3_OFFSET; + } + + __aicore__ inline uint32_t L1AkkOffset(uint32_t slot) const + { + return KDA_FWD_H_L1_AKK_OFFSET + slot * KDA_FWD_H_L1_AKK_SLOT_BYTES; + } + + __aicore__ inline uint32_t L1UOffset(uint32_t slot) const + { + return KDA_FWD_H_L1_U_OFFSET + slot * KDA_FWD_H_L1_U_SLOT_BYTES; + } + + template + __aicore__ inline void PublishDirectTile( + TensorL0C tensorL0C, uint32_t m, uint32_t n, + TEventID mToFixEvent, TEventID fixToMEvent) + { + auto layoutUb = tla::MakeLayout(m, n); + auto tensorUb = tla::MakeTensor( + resource_.ubBuf.template GetBufferByByte(KDA_FWD_H_UB_DIRECT_OFFSET), + layoutUb, Catlass::Arch::PositionUB{}); + using CopyL0CToDst = + typename DirectTileCopy::template CopyL0CToDst; + CopyL0CToDst copyL0CToDst; + + CrossCoreWaitFlag(KDA_FWD_H_DIRECT_FREE_FLAG); + SetFlag(mToFixEvent); + WaitFlag(mToFixEvent); + copyL0CToDst(tensorUb, tensorL0C); + CrossCoreSetFlag<0x2, PIPE_FIX>(KDA_FWD_H_DIRECT_READY_FLAG); + SetFlag(fixToMEvent); + } + + template + __aicore__ inline void PublishDirect( + LocalTensor l0C, uint32_t m, uint32_t n, + TEventID mToFixEvent, TEventID fixToMEvent) + { + auto layoutL0C = tla::MakeLayoutL0C(m, n); + auto tensorL0C = tla::MakeTensor(l0C, layoutL0C, Catlass::Arch::PositionL0C{}); + PublishDirectTile( + tensorL0C, m, n, mToFixEvent, fixToMEvent); + } + + __aicore__ inline void PrefetchIndependentProductsAic( + uint64_t b, uint64_t hv, uint64_t chunk) + { + const uint32_t slot = static_cast(chunk & 3); + WaitFlag(aicL1ReuseEvents_[slot]); + using LayoutTagL1ARm = typename TileCopyRM::LayoutTagL1A; + using LayoutTagL1BRm = typename TileCopyRM::LayoutTagL1B; + using LayoutTagL1ACm = typename TileCopyCM::LayoutTagL1A; + + auto layoutToken = tla::MakeLayout(KDA_FWD_H_CHUNK, KDA_FWD_H_DIM); + auto layoutKg = tla::MakeLayout(KDA_FWD_H_DIM, KDA_FWD_H_CHUNK); + auto layoutAqk = tla::MakeLayout(KDA_FWD_H_CHUNK, KDA_FWD_H_CHUNK); + auto layoutAkk = tla::MakeLayout(KDA_FWD_H_CHUNK, KDA_FWD_H_CHUNK); + auto tensorW = tla::MakeTensor( + w_[ChunkMatrixOffset(b, hv, chunk)], layoutToken, Catlass::Arch::PositionGM{}); + auto tensorQ = tla::MakeTensor( + qg_[ChunkMatrixOffset(b, hv, chunk)], layoutToken, + Catlass::Arch::PositionGM{}); + auto tensorKg = tla::MakeTensor( + kg_[ChunkMatrixOffset(b, hv, chunk)], layoutKg, Catlass::Arch::PositionGM{}); + auto tensorAqk = tla::MakeTensor( + aqk_[ChunkScoreOffset(b, hv, chunk)], layoutAqk, Catlass::Arch::PositionGM{}); + auto tensorAkk = tla::MakeTensor( + akk_[ChunkScoreOffset(b, hv, chunk)], layoutAkk, Catlass::Arch::PositionGM{}); + auto tensorU = tla::MakeTensor( + u_[ChunkMatrixOffset(b, hv, chunk)], layoutToken, Catlass::Arch::PositionGM{}); + auto blockW = GetTile(tensorW, tla::MakeCoord(0, 0), + tla::MakeShape(KDA_FWD_H_CHUNK, KDA_FWD_H_DIM)); + auto blockQ = GetTile(tensorQ, tla::MakeCoord(0, 0), + tla::MakeShape(KDA_FWD_H_CHUNK, KDA_FWD_H_DIM)); + auto blockKg = GetTile(tensorKg, tla::MakeCoord(0, 0), + tla::MakeShape(KDA_FWD_H_DIM, KDA_FWD_H_CHUNK)); + auto blockAqk = GetTile(tensorAqk, tla::MakeCoord(0, 0), + tla::MakeShape(KDA_FWD_H_CHUNK, KDA_FWD_H_CHUNK)); + auto blockAkk = GetTile(tensorAkk, tla::MakeCoord(0, 0), + tla::MakeShape(KDA_FWD_H_CHUNK, KDA_FWD_H_CHUNK)); + auto blockU = GetTile(tensorU, tla::MakeCoord(0, 0), + tla::MakeShape(KDA_FWD_H_CHUNK, KDA_FWD_H_DIM)); + using CopyGmToL1ARmW = typename TileCopyRM::template CopyGmToL1A; + using CopyGmToL1ARmQ = typename TileCopyRM::template CopyGmToL1A; + using CopyGmToL1ACm = typename TileCopyCM::template CopyGmToL1A; + using CopyGmToL1ARmAqk = typename TileCopyRM::template CopyGmToL1A; + using CopyGmToL1ARmAkk = typename TileCopyRM::template CopyGmToL1A; + using CopyGmToL1BRmU = typename TileCopyRM::template CopyGmToL1B; + + LocalTensor l1W = resource_.l1Buf.template GetBufferByByte( + L1WOffset(slot)); + LocalTensor l1Q = resource_.l1Buf.template GetBufferByByte( + L1QOffset(slot)); + LocalTensor l1Kg = resource_.l1Buf.template GetBufferByByte( + L1KgOffset(slot)); + LocalTensor l1Aqk = resource_.l1Buf.template GetBufferByByte( + L1AqkOffset(slot)); + LocalTensor l1Akk = resource_.l1Buf.template GetBufferByByte( + L1AkkOffset(slot)); + LocalTensor l1U = resource_.l1Buf.template GetBufferByByte( + L1UOffset(slot)); + auto tensorL1W = tla::MakeTensor( + l1W, tla::MakeLayout(KDA_FWD_H_CHUNK, KDA_FWD_H_DIM), + Catlass::Arch::PositionL1{}); + auto tensorL1Q = tla::MakeTensor( + l1Q, tla::MakeLayout(KDA_FWD_H_CHUNK, KDA_FWD_H_DIM), + Catlass::Arch::PositionL1{}); + auto tensorL1Kg = tla::MakeTensor( + l1Kg, tla::MakeLayout(KDA_FWD_H_DIM, KDA_FWD_H_CHUNK), + Catlass::Arch::PositionL1{}); + auto tensorL1Aqk = tla::MakeTensor( + l1Aqk, tla::MakeLayout(KDA_FWD_H_CHUNK, KDA_FWD_H_CHUNK), + Catlass::Arch::PositionL1{}); + auto tensorL1Akk = tla::MakeTensor( + l1Akk, tla::MakeLayout(KDA_FWD_H_CHUNK, KDA_FWD_H_CHUNK), + Catlass::Arch::PositionL1{}); + auto tensorL1U = tla::MakeTensor( + l1U, tla::MakeLayout(KDA_FWD_H_CHUNK, KDA_FWD_H_DIM), + Catlass::Arch::PositionL1{}); + + CopyGmToL1ARmW{}(tensorL1W, blockW); + CopyGmToL1ARmQ{}(tensorL1Q, blockQ); + CopyGmToL1ACm{}(tensorL1Kg, blockKg); + CopyGmToL1ARmAqk{}(tensorL1Aqk, blockAqk); + if (fusePostWuIntoFwdH_) { + CopyGmToL1ARmAkk{}(tensorL1Akk, blockAkk); + CopyGmToL1BRmU{}(tensorL1U, blockU); + } + SetFlag(aicMte2ToMte1Event_); + } + + __aicore__ inline void ComputePostWuAic(uint64_t chunk) + { + const uint32_t slot = static_cast(chunk & 3); + using LayoutTagL1A = typename TileCopyRM::LayoutTagL1A; + using LayoutTagL1B = typename TileCopyRM::LayoutTagL1B; + using LayoutTagL0A = typename TileCopyRM::LayoutTagL0A; + using LayoutTagL0B = typename TileCopyRM::LayoutTagL0B; + using CopyL1ToL0A = typename TileCopyRM::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopyRM::CopyL1ToL0B; + using TileMmad = Catlass::Gemm::Tile::TileMmadTla; + + LocalTensor l1Akk = resource_.l1Buf.template GetBufferByByte( + L1AkkOffset(slot)); + LocalTensor l1W = resource_.l1Buf.template GetBufferByByte( + L1WOffset(slot)); + LocalTensor l1U = resource_.l1Buf.template GetBufferByByte( + L1UOffset(slot)); + LocalTensor l0A = resource_.l0ABuf.template GetBufferByByte( + KDA_FWD_H_L0A_POST_OFFSET); + LocalTensor l0B = resource_.l0BBuf.template GetBufferByByte( + KDA_FWD_H_L0B_POST_OFFSET); + LocalTensor l0C = resource_.l0CBuf.template GetBufferByByte(0); + + auto tensorL1Akk = tla::MakeTensor( + l1Akk, tla::MakeLayout(KDA_FWD_H_CHUNK, KDA_FWD_H_CHUNK), + Catlass::Arch::PositionL1{}); + auto tensorL1W = tla::MakeTensor( + l1W, tla::MakeLayout(KDA_FWD_H_CHUNK, KDA_FWD_H_DIM), + Catlass::Arch::PositionL1{}); + auto tensorL1U = tla::MakeTensor( + l1U, tla::MakeLayout(KDA_FWD_H_CHUNK, KDA_FWD_H_DIM), + Catlass::Arch::PositionL1{}); + auto tensorL0A = tla::MakeTensor( + l0A, tla::MakeLayout(KDA_FWD_H_CHUNK, KDA_FWD_H_CHUNK), + Catlass::Arch::PositionL0A{}); + auto tensorL0B = tla::MakeTensor( + l0B, tla::MakeLayout(KDA_FWD_H_CHUNK, 2 * KDA_FWD_H_DIM), + Catlass::Arch::PositionL0B{}); + auto tensorL0C = tla::MakeTensor( + l0C, tla::MakeLayoutL0C(KDA_FWD_H_CHUNK, 2 * KDA_FWD_H_DIM), + Catlass::Arch::PositionL0C{}); + auto tileL1Akk = GetTile(tensorL1Akk, tla::MakeCoord(0, 0), + tla::MakeShape(KDA_FWD_H_CHUNK, KDA_FWD_H_CHUNK)); + auto tileL1W = GetTile(tensorL1W, tla::MakeCoord(0, 0), + tla::MakeShape(KDA_FWD_H_CHUNK, KDA_FWD_H_DIM)); + auto tileL1U = GetTile(tensorL1U, tla::MakeCoord(0, 0), + tla::MakeShape(KDA_FWD_H_CHUNK, KDA_FWD_H_DIM)); + auto tileL0A = GetTile(tensorL0A, tla::MakeCoord(0, 0), + tla::MakeShape(KDA_FWD_H_CHUNK, KDA_FWD_H_CHUNK)); + auto tileL0BW = GetTile(tensorL0B, tla::MakeCoord(0, 0), + tla::MakeShape(KDA_FWD_H_CHUNK, KDA_FWD_H_DIM)); + auto tileL0BU = GetTile(tensorL0B, tla::MakeCoord(0, KDA_FWD_H_DIM), + tla::MakeShape(KDA_FWD_H_CHUNK, KDA_FWD_H_DIM)); + auto tileL0B = GetTile(tensorL0B, tla::MakeCoord(0, 0), + tla::MakeShape(KDA_FWD_H_CHUNK, 2 * KDA_FWD_H_DIM)); + auto tileL0C = GetTile(tensorL0C, tla::MakeCoord(0, 0), + tla::MakeShape(KDA_FWD_H_CHUNK, 2 * KDA_FWD_H_DIM)); + auto tileL0CW = GetTile(tensorL0C, tla::MakeCoord(0, 0), + tla::MakeShape(KDA_FWD_H_CHUNK, KDA_FWD_H_DIM)); + auto tileL0CU = GetTile(tensorL0C, tla::MakeCoord(0, KDA_FWD_H_DIM), + tla::MakeShape(KDA_FWD_H_CHUNK, KDA_FWD_H_DIM)); + + WaitFlag(aicMte2ToMte1Event_); + WaitFlag(stateL0FreeEvent_); + CopyL1ToL0A{}(tileL0A, tileL1Akk); + CopyL1ToL0B{}(tileL0BW, tileL1W); + CopyL1ToL0B{}(tileL0BU, tileL1U); + SetFlag(aicMte1ToMEvent_); + WaitFlag(aicMte1ToMEvent_); + TileMmad{}(tileL0C, tileL0A, tileL0B, KDA_FWD_H_CHUNK, + 2 * KDA_FWD_H_DIM, KDA_FWD_H_CHUNK, true, 0); + SetFlag(stateL0FreeEvent_); + PublishDirectTile( + tileL0CW, KDA_FWD_H_CHUNK, KDA_FWD_H_DIM, + aicMToFixEvent_, aicFixToMEvent_); + WaitFlag(aicFixToMEvent_); + PublishDirectTile( + tileL0CU, KDA_FWD_H_CHUNK, KDA_FWD_H_DIM, + aicMToFixEvent_, aicFixToMEvent_); + WaitFlag(aicFixToMEvent_); + } + + __aicore__ inline void ComputeStateProductsAic( + uint64_t b, uint64_t hv, uint64_t chunk, bool inputsReady = false) + { + const uint32_t slot = static_cast(chunk & 3); + using LayoutTagL1A = typename TileCopyRM::LayoutTagL1A; + using LayoutTagL1B = typename TileCopyRM::LayoutTagL1B; + using LayoutTagL0A = typename TileCopyRM::LayoutTagL0A; + using LayoutTagL0B = typename TileCopyRM::LayoutTagL0B; + using CopyL1ToL0A = typename TileCopyRM::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopyRM::CopyL1ToL0B; + using TileMmad = Catlass::Gemm::Tile::TileMmadTla; + + LocalTensor l1A0 = resource_.l1Buf.template GetBufferByByte( + L1WOffset(slot)); + LocalTensor l1A1 = resource_.l1Buf.template GetBufferByByte( + L1QOffset(slot)); + LocalTensor l1B = + resource_.l1Buf.template GetBufferByByte(KDA_FWD_H_L1_H_OFFSET); + LocalTensor l0A = resource_.l0ABuf.template GetBufferByByte( + KDA_FWD_H_L0A_STATE_OFFSET); + LocalTensor l0B = resource_.l0BBuf.template GetBufferByByte( + KDA_FWD_H_L0B_STATE_OFFSET); + LocalTensor l0C = resource_.l0CBuf.template GetBufferByByte(0); + + auto layoutL1A = tla::MakeLayout(KDA_FWD_H_CHUNK, KDA_FWD_H_DIM); + auto layoutL1B = tla::MakeLayout(KDA_FWD_H_DIM, KDA_FWD_H_DIM); + auto layoutL0A = tla::MakeLayout(KDA_FWD_H_CHUNK, KDA_FWD_H_DIM); + auto layoutL0B = tla::MakeLayout(KDA_FWD_H_DIM, KDA_FWD_H_DIM); + auto layoutL0C = tla::MakeLayoutL0C(KDA_FWD_H_CHUNK, KDA_FWD_H_DIM); + auto tensorL1A0 = tla::MakeTensor(l1A0, layoutL1A, Catlass::Arch::PositionL1{}); + auto tensorL1A1 = tla::MakeTensor(l1A1, layoutL1A, Catlass::Arch::PositionL1{}); + auto tensorL1B = tla::MakeTensor(l1B, layoutL1B, Catlass::Arch::PositionL1{}); + auto tensorL0A = tla::MakeTensor(l0A, layoutL0A, Catlass::Arch::PositionL0A{}); + auto tensorL0B = tla::MakeTensor(l0B, layoutL0B, Catlass::Arch::PositionL0B{}); + auto tensorL0C = tla::MakeTensor(l0C, layoutL0C, Catlass::Arch::PositionL0C{}); + auto tileL1B = GetTile(tensorL1B, tla::MakeCoord(0, 0), + tla::MakeShape(KDA_FWD_H_DIM, KDA_FWD_H_DIM)); + auto tileL0A = GetTile(tensorL0A, tla::MakeCoord(0, 0), + tla::MakeShape(KDA_FWD_H_CHUNK, KDA_FWD_H_DIM)); + auto tileL1W = GetTile(tensorL1A0, tla::MakeCoord(0, 0), + tla::MakeShape(KDA_FWD_H_CHUNK, KDA_FWD_H_DIM)); + auto tileL1Q = GetTile(tensorL1A1, tla::MakeCoord(0, 0), + tla::MakeShape(KDA_FWD_H_CHUNK, KDA_FWD_H_DIM)); + auto tileL0B = GetTile(tensorL0B, tla::MakeCoord(0, 0), + tla::MakeShape(KDA_FWD_H_DIM, KDA_FWD_H_DIM)); + auto tileL0C = GetTile(tensorL0C, tla::MakeCoord(0, 0), + tla::MakeShape(KDA_FWD_H_CHUNK, KDA_FWD_H_DIM)); + + CopyL1ToL0A copyL1ToL0A; + CopyL1ToL0B copyL1ToL0B; + TileMmad tileMmad; + + if (!inputsReady) { + WaitFlag(aicMte2ToMte1Event_); + } + WaitFlag(stateL0FreeEvent_); + copyL1ToL0A(tileL0A, tileL1W); + copyL1ToL0B(tileL0B, tileL1B); + SetFlag(aicMte1ToMEvent_); + WaitFlag(aicMte1ToMEvent_); + tileMmad(tileL0C, tileL0A, tileL0B, KDA_FWD_H_CHUNK, + KDA_FWD_H_DIM, KDA_FWD_H_DIM, true, 0); + SetFlag(stateL0FreeEvent_); + PublishDirect( + l0C, KDA_FWD_H_CHUNK, KDA_FWD_H_DIM, + aicMToFixEvent_, aicFixToMEvent_); + SetL1SlotFlagAicToAiv(KDA_FWD_H_STATE_FREE_FLAG); + if (fusePostWuIntoFwdH_) { + SetL1SlotFlagAicToAiv(KDA_FWD_H_STATE_FREE_FLAG); + } + + WaitFlag(stateL0FreeEvent_); + WaitFlag(aicFixToMEvent_); + copyL1ToL0A(tileL0A, tileL1Q); + SetFlag(aicMte1ToMEvent_); + WaitFlag(aicMte1ToMEvent_); + tileMmad(tileL0C, tileL0A, tileL0B, KDA_FWD_H_CHUNK, + KDA_FWD_H_DIM, KDA_FWD_H_DIM, true, 0); + SetFlag(stateL0FreeEvent_); + PublishDirect( + l0C, KDA_FWD_H_CHUNK, KDA_FWD_H_DIM, + aicMToFixEvent_, aicFixToMEvent_); + WaitFlag(aicFixToMEvent_); + } + + __aicore__ inline void ComputeVnewProductsAic( + uint64_t b, uint64_t hv, uint64_t chunk, bool prefetchNext) + { + const uint32_t slot = static_cast(chunk & 3); + using LayoutTagL1AK = typename TileCopyCM::LayoutTagL1A; + using LayoutTagL1AA = typename TileCopyRM::LayoutTagL1A; + using LayoutTagL1B = typename TileCopyRM::LayoutTagL1B; + using LayoutTagL0AK = typename TileCopyCM::LayoutTagL0A; + using LayoutTagL0AA = typename TileCopyRM::LayoutTagL0A; + using LayoutTagL0B = typename TileCopyRM::LayoutTagL0B; + using CopyL1ToL0AK = typename TileCopyCM::CopyL1ToL0A; + using CopyL1ToL0AA = typename TileCopyRM::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopyRM::CopyL1ToL0B; + using TileMmadK = Catlass::Gemm::Tile::TileMmadTla; + using TileMmadA = Catlass::Gemm::Tile::TileMmadTla; + + auto layoutKg = tla::MakeLayout(KDA_FWD_H_DIM, KDA_FWD_H_CHUNK); + auto layoutAqk = tla::MakeLayout(KDA_FWD_H_CHUNK, KDA_FWD_H_CHUNK); + LocalTensor l1Kg = resource_.l1Buf.template GetBufferByByte( + L1KgOffset(slot)); + LocalTensor l1Aqk = resource_.l1Buf.template GetBufferByByte( + L1AqkOffset(slot)); + LocalTensor l1V = + resource_.l1Buf.template GetBufferByByte(KDA_FWD_H_L1_V_OFFSET); + LocalTensor l0A = resource_.l0ABuf.template GetBufferByByte( + KDA_FWD_H_L0A_VNEW_OFFSET); + LocalTensor l0B = resource_.l0BBuf.template GetBufferByByte( + KDA_FWD_H_L0B_VNEW_OFFSET); + LocalTensor l0C = resource_.l0CBuf.template GetBufferByByte(0); + + auto layoutL1Kg = tla::MakeLayout(KDA_FWD_H_DIM, KDA_FWD_H_CHUNK); + auto layoutL1Aqk = tla::MakeLayout(KDA_FWD_H_CHUNK, KDA_FWD_H_CHUNK); + auto layoutL1V = tla::MakeLayout(KDA_FWD_H_CHUNK, KDA_FWD_H_DIM); + auto layoutL0Kg = tla::MakeLayout(KDA_FWD_H_DIM, KDA_FWD_H_CHUNK); + auto layoutL0Aqk = tla::MakeLayout(KDA_FWD_H_CHUNK, KDA_FWD_H_CHUNK); + auto layoutL0V = tla::MakeLayout(KDA_FWD_H_CHUNK, KDA_FWD_H_DIM); + auto baseL1Kg = tla::MakeTensor(l1Kg, layoutL1Kg, Catlass::Arch::PositionL1{}); + auto baseL1Aqk = tla::MakeTensor(l1Aqk, layoutL1Aqk, Catlass::Arch::PositionL1{}); + auto baseL1V = tla::MakeTensor(l1V, layoutL1V, Catlass::Arch::PositionL1{}); + auto baseL0Kg = tla::MakeTensor(l0A, layoutL0Kg, Catlass::Arch::PositionL0A{}); + auto baseL0Aqk = tla::MakeTensor(l0A, layoutL0Aqk, Catlass::Arch::PositionL0A{}); + auto baseL0V = tla::MakeTensor(l0B, layoutL0V, Catlass::Arch::PositionL0B{}); + auto baseL0Update = tla::MakeTensor( + l0C, tla::MakeLayoutL0C(KDA_FWD_H_DIM, KDA_FWD_H_DIM), + Catlass::Arch::PositionL0C{}); + auto baseL0Out = tla::MakeTensor( + l0C, tla::MakeLayoutL0C(KDA_FWD_H_CHUNK, KDA_FWD_H_DIM), + Catlass::Arch::PositionL0C{}); + auto tensorL1Kg = GetTile(baseL1Kg, tla::MakeCoord(0, 0), + tla::MakeShape(KDA_FWD_H_DIM, KDA_FWD_H_CHUNK)); + auto tensorL1Aqk = GetTile(baseL1Aqk, tla::MakeCoord(0, 0), + tla::MakeShape(KDA_FWD_H_CHUNK, KDA_FWD_H_CHUNK)); + auto tensorL1V = GetTile(baseL1V, tla::MakeCoord(0, 0), + tla::MakeShape(KDA_FWD_H_CHUNK, KDA_FWD_H_DIM)); + auto tensorL0Kg = GetTile(baseL0Kg, tla::MakeCoord(0, 0), + tla::MakeShape(KDA_FWD_H_DIM, KDA_FWD_H_CHUNK)); + auto tensorL0Aqk = GetTile(baseL0Aqk, tla::MakeCoord(0, 0), + tla::MakeShape(KDA_FWD_H_CHUNK, KDA_FWD_H_CHUNK)); + auto tensorL0V = GetTile(baseL0V, tla::MakeCoord(0, 0), + tla::MakeShape(KDA_FWD_H_CHUNK, KDA_FWD_H_DIM)); + auto tensorL0Update = GetTile(baseL0Update, tla::MakeCoord(0, 0), + tla::MakeShape(KDA_FWD_H_DIM, KDA_FWD_H_DIM)); + auto tensorL0Out = GetTile(baseL0Out, tla::MakeCoord(0, 0), + tla::MakeShape(KDA_FWD_H_CHUNK, KDA_FWD_H_DIM)); + + CopyL1ToL0AK copyL1ToL0AK; + CopyL1ToL0AA copyL1ToL0AA; + CopyL1ToL0B copyL1ToL0B; + TileMmadK tileMmadK; + TileMmadA tileMmadA; + + WaitFlag(vnewL0FreeEvent_); + copyL1ToL0AK(tensorL0Kg, tensorL1Kg); + copyL1ToL0B(tensorL0V, tensorL1V); + SetFlag(aicMte1ToMEvent_); + WaitFlag(aicMte1ToMEvent_); + tileMmadK(tensorL0Update, tensorL0Kg, tensorL0V, + KDA_FWD_H_DIM, KDA_FWD_H_DIM, KDA_FWD_H_CHUNK, true, 0); + SetFlag(vnewL0FreeEvent_); + PublishDirect( + l0C, KDA_FWD_H_DIM, KDA_FWD_H_DIM, + aicMToFixEvent_, aicFixToMEvent_); + SetL1SlotFlagAicToAiv(KDA_FWD_H_VNEW_FREE_FLAG); + + WaitFlag(vnewL0FreeEvent_); + WaitFlag(aicFixToMEvent_); + copyL1ToL0AA(tensorL0Aqk, tensorL1Aqk); + SetFlag(aicL1ReuseEvents_[slot]); + if (prefetchNext) { + PrefetchIndependentProductsAic(b, hv, chunk + 1); + } + SetFlag(aicMte1ToMEvent_); + WaitFlag(aicMte1ToMEvent_); + tileMmadA(tensorL0Out, tensorL0Aqk, tensorL0V, + KDA_FWD_H_CHUNK, KDA_FWD_H_DIM, KDA_FWD_H_CHUNK, true, 0); + SetFlag(vnewL0FreeEvent_); + PublishDirect( + l0C, KDA_FWD_H_CHUNK, KDA_FWD_H_DIM, + aicMToFixEvent_, aicFixToMEvent_); + WaitFlag(aicFixToMEvent_); + } + + __aicore__ inline void ProcessAic() + { + SetLoadDataPaddingValue(static_cast(0)); + SetFlag(stateL0FreeEvent_); + for (uint32_t slot = 0; slot < KDA_FWD_H_L1_STAGING_DEPTH; ++slot) { + SetFlag(aicL1ReuseEvents_[slot]); + } + const uint64_t coreIdx = static_cast(GetBlockIdx()); + const uint64_t coreNum = coreNum_ == 0 ? 1 : coreNum_; + for (uint64_t task = coreIdx; task < batch_ * heads_; task += coreNum) { + const uint64_t b = task / heads_; + const uint64_t hv = task % heads_; + PrefetchIndependentProductsAic(b, hv, 0); + for (uint64_t chunk = 0; chunk < totalChunks_; ++chunk) { + if (fusePostWuIntoFwdH_) { + ComputePostWuAic(chunk); + } + WaitL1SlotReadyMte1(KDA_FWD_H_STATE_READY_FLAG); + if (fusePostWuIntoFwdH_) { + WaitL1SlotReadyMte1(KDA_FWD_H_STATE_READY_FLAG); + } + ComputeStateProductsAic(b, hv, chunk, fusePostWuIntoFwdH_); + WaitL1SlotReadyMte1(KDA_FWD_H_VNEW_READY_FLAG); + ComputeVnewProductsAic( + b, hv, chunk, chunk + 1 < totalChunks_); + } + CrossCoreWaitFlag(KDA_FWD_H_DIRECT_FREE_FLAG); + } + WaitFlag(stateL0FreeEvent_); + for (uint32_t slot = 0; slot < KDA_FWD_H_L1_STAGING_DEPTH; ++slot) { + WaitFlag(aicL1ReuseEvents_[slot]); + } + } + + __aicore__ inline void CopyOutputRows( + uint64_t b, uint64_t hv, uint64_t tokenStart, + LocalTensor src, uint32_t rows) + { + DataCopyExtParams params{ + static_cast(rows), + static_cast(KDA_FWD_H_DIM * sizeof(T)), + 0, + static_cast((heads_ * KDA_FWD_H_DIM - KDA_FWD_H_DIM) * sizeof(T)), + 0}; + DataCopyPad(attnOut_[OutputOffset(b, hv, tokenStart)], src, params); + } + + __aicore__ inline void InitializeStateAiv( + uint64_t b, uint64_t hv, uint32_t rowBegin, + LocalTensor state) + { + if (hasInitialState_) { + DataCopy(state, initialState_[StateOffset(b, hv) + rowBegin * KDA_FWD_H_DIM], + KDA_FWD_H_STATE_SUB_ELEMS); + SetFlag(aivMte2ToVEvent_); + WaitFlag(aivMte2ToVEvent_); + } else { + Duplicate(state, 0.0f, KDA_FWD_H_STATE_SUB_ELEMS); + PipeBarrier(); + } + } + + __aicore__ inline void StoreCurrentStateAiv( + uint64_t b, uint64_t hv, uint64_t chunk, uint32_t rowBegin, + LocalTensor state, LocalTensor stateTyped) + { + if (storeH_) { + Cast(stateTyped, state, RoundMode::CAST_RINT, KDA_FWD_H_STATE_SUB_ELEMS); + PipeBarrier(); + SetFlag(aivVToMte3Event_); + WaitFlag(aivVToMte3Event_); + DataCopy(h_[HOffset(b, hv, chunk) + rowBegin * KDA_FWD_H_DIM], + stateTyped, KDA_FWD_H_STATE_SUB_ELEMS); + SetFlag(aivMte3ToVEvent_); + WaitFlag(aivMte3ToVEvent_); + } + + constexpr uint32_t columnGroup = 64; + constexpr uint32_t columnGroups = KDA_FWD_H_DIM / columnGroup; + for (uint32_t group = 0; group < columnGroups; ++group) { + Cast(stateTyped[group * KDA_FWD_H_SUB_DIM * columnGroup], + state[group * columnGroup], RoundMode::CAST_RINT, + columnGroup, KDA_FWD_H_SUB_DIM, + {static_cast(KDA_FWD_H_SUB_DIM), 1, 1, + static_cast(columnGroups * 8)}); + } + PipeBarrier(); + SetFlag(aivVToMte3Event_); + WaitFlag(aivVToMte3Event_); + LocalTensor l1State = + resource_.l1Buf.template GetBufferByByte(KDA_FWD_H_L1_H_OFFSET); + const uint32_t subBlockIdx = rowBegin / KDA_FWD_H_SUB_DIM; + WaitL1SlotFreeMte3( + KDA_FWD_H_STATE_FREE_FLAG, statePublishCount_[subBlockIdx]); + DataCopyParams stateCopyParams; + stateCopyParams.blockCount = KDA_FWD_H_DIM / 16; + stateCopyParams.blockLen = KDA_FWD_H_SUB_DIM; + stateCopyParams.srcGap = 0; + stateCopyParams.dstGap = KDA_FWD_H_DIM - KDA_FWD_H_SUB_DIM; + DataCopy(l1State[rowBegin * 16], stateTyped, stateCopyParams); + SetL1SlotFlagAivToAic(KDA_FWD_H_STATE_READY_FLAG); + ++statePublishCount_[subBlockIdx]; + SetFlag(aivMte3ToVEvent_); + WaitFlag(aivMte3ToVEvent_); + } + + __aicore__ inline void ProcessChunkAiv( + uint64_t b, uint64_t hv, uint32_t chunk, + uint32_t subBlockIdx, LocalTensor state, + LocalTensor stateTyped, LocalTensor direct, + LocalTensor out1, LocalTensor vnew, + LocalTensor ioTyped, LocalTensor gate) + { + const uint32_t tokenBegin = subBlockIdx * KDA_FWD_H_SUB_CHUNK; + const uint32_t stateRowBegin = subBlockIdx * KDA_FWD_H_SUB_DIM; + StoreCurrentStateAiv(b, hv, chunk, stateRowBegin, state, stateTyped); + + CrossCoreWaitFlag(KDA_FWD_H_DIRECT_READY_FLAG); + WaitFlag(aivMte3ToMte2Event_); + if (fusePostWuIntoFwdH_) { + Cast(ioTyped, direct, RoundMode::CAST_RINT, KDA_FWD_H_TOKEN_SUB_ELEMS); + PipeBarrier(); + CrossCoreSetFlag<0x2, PIPE_V>(KDA_FWD_H_DIRECT_FREE_FLAG); + + SetFlag(aivVToMte3Event_); + WaitFlag(aivVToMte3Event_); + LocalTensor l1W = resource_.l1Buf.template GetBufferByByte( + L1WOffset(chunk & 3)); + WaitL1SlotFreeMte3( + KDA_FWD_H_STATE_FREE_FLAG, wPublishCount_[subBlockIdx]); + DataCopyParams wCopyParams; + wCopyParams.blockCount = KDA_FWD_H_SUB_CHUNK; + wCopyParams.blockLen = 1; + wCopyParams.srcGap = KDA_FWD_H_DIM / 16 - 1; + wCopyParams.dstGap = 0; + for (uint32_t colBlock = 0; colBlock < KDA_FWD_H_DIM / 16; ++colBlock) { + const uint32_t dstOffset = + colBlock * KDA_FWD_H_CHUNK * 16 + tokenBegin * 16; + DataCopy(l1W[dstOffset], ioTyped[colBlock * 16], wCopyParams); + } + SetFlag(aivMte3ToVEvent_); + WaitFlag(aivMte3ToVEvent_); + SetL1SlotFlagAivToAic(KDA_FWD_H_STATE_READY_FLAG); + ++wPublishCount_[subBlockIdx]; + + CrossCoreWaitFlag(KDA_FWD_H_DIRECT_READY_FLAG); + Cast(ioTyped, direct, RoundMode::CAST_RINT, KDA_FWD_H_TOKEN_SUB_ELEMS); + PipeBarrier(); + Cast(vnew, ioTyped, RoundMode::CAST_NONE, KDA_FWD_H_TOKEN_SUB_ELEMS); + PipeBarrier(); + CrossCoreSetFlag<0x2, PIPE_V>(KDA_FWD_H_DIRECT_FREE_FLAG); + CrossCoreWaitFlag(KDA_FWD_H_DIRECT_READY_FLAG); + } else { + DataCopy(ioTyped, u_[ChunkMatrixOffset(b, hv, chunk, tokenBegin)], + KDA_FWD_H_TOKEN_SUB_ELEMS); + SetFlag(aivMte2ToVEvent_); + WaitFlag(aivMte2ToVEvent_); + Cast(vnew, ioTyped, RoundMode::CAST_NONE, KDA_FWD_H_TOKEN_SUB_ELEMS); + PipeBarrier(); + } + Sub(vnew, vnew, direct, KDA_FWD_H_TOKEN_SUB_ELEMS); + PipeBarrier(); + CrossCoreSetFlag<0x2, PIPE_V>(KDA_FWD_H_DIRECT_FREE_FLAG); + if (storeVNew_) { + Cast(ioTyped, vnew, RoundMode::CAST_RINT, KDA_FWD_H_TOKEN_SUB_ELEMS); + PipeBarrier(); + SetFlag(aivVToMte3Event_); + WaitFlag(aivVToMte3Event_); + DataCopy(vNew_[VNewOffset(b, hv, chunk, tokenBegin)], + ioTyped, KDA_FWD_H_TOKEN_SUB_ELEMS); + SetFlag(aivMte3ToVEvent_); + WaitFlag(aivMte3ToVEvent_); + } + + constexpr uint32_t columnGroup = 64; + constexpr uint32_t columnGroups = KDA_FWD_H_DIM / columnGroup; + for (uint32_t group = 0; group < columnGroups; ++group) { + Cast(ioTyped[group * KDA_FWD_H_SUB_CHUNK * columnGroup], + vnew[group * columnGroup], RoundMode::CAST_RINT, + columnGroup, KDA_FWD_H_SUB_CHUNK, + {static_cast(KDA_FWD_H_SUB_CHUNK), 1, 1, + static_cast(columnGroups * 8)}); + } + PipeBarrier(); + SetFlag(aivVToMte3Event_); + WaitFlag(aivVToMte3Event_); + LocalTensor l1Vnew = + resource_.l1Buf.template GetBufferByByte(KDA_FWD_H_L1_V_OFFSET); + WaitL1SlotFreeMte3( + KDA_FWD_H_VNEW_FREE_FLAG, vnewPublishCount_[subBlockIdx]); + DataCopyParams vnewL1CopyParams; + vnewL1CopyParams.blockCount = KDA_FWD_H_DIM / 16; + vnewL1CopyParams.blockLen = KDA_FWD_H_SUB_CHUNK; + vnewL1CopyParams.srcGap = 0; + vnewL1CopyParams.dstGap = KDA_FWD_H_CHUNK - KDA_FWD_H_SUB_CHUNK; + DataCopy(l1Vnew[tokenBegin * 16], ioTyped, vnewL1CopyParams); + SetFlag(aivMte3ToVEvent_); + WaitFlag(aivMte3ToVEvent_); + SetL1SlotFlagAivToAic(KDA_FWD_H_VNEW_READY_FLAG); + ++vnewPublishCount_[subBlockIdx]; + + CrossCoreWaitFlag(KDA_FWD_H_DIRECT_READY_FLAG); + Muls(out1, direct, scale_, KDA_FWD_H_TOKEN_SUB_ELEMS); + PipeBarrier(); + CrossCoreSetFlag<0x2, PIPE_V>(KDA_FWD_H_DIRECT_FREE_FLAG); + + CrossCoreWaitFlag(KDA_FWD_H_DIRECT_READY_FLAG); + DataCopy(gate, gk_[ChunkMatrixOffset(b, hv, chunk, KDA_FWD_H_CHUNK - 1) + stateRowBegin], KDA_FWD_H_SUB_DIM); + SetFlag(aivMte2ToVEvent_); + WaitFlag(aivMte2ToVEvent_); + Muls(gate, gate, 0.6931471805599453f, KDA_FWD_H_SUB_DIM); + PipeBarrier(); + Exp(gate, gate, KDA_FWD_H_SUB_DIM); + PipeBarrier(); + AscendC::VF_CALL>( + reinterpret_cast<__ubuf__ float *>(direct.GetPhyAddr()), + reinterpret_cast<__ubuf__ float *>(state.GetPhyAddr()), + reinterpret_cast<__ubuf__ float *>(gate.GetPhyAddr()), + static_cast(KDA_FWD_H_SUB_DIM), + static_cast(KDA_FWD_H_DIM)); + PipeBarrier(); + Adds(state, direct, 0.0f, KDA_FWD_H_STATE_SUB_ELEMS); + PipeBarrier(); + + CrossCoreSetFlag<0x2, PIPE_V>(KDA_FWD_H_DIRECT_FREE_FLAG); + CrossCoreWaitFlag(KDA_FWD_H_DIRECT_READY_FLAG); + Add(direct, direct, out1, KDA_FWD_H_TOKEN_SUB_ELEMS); + PipeBarrier(); + Cast(ioTyped, direct, RoundMode::CAST_RINT, KDA_FWD_H_TOKEN_SUB_ELEMS); + PipeBarrier(); + SetFlag(aivVToMte3Event_); + WaitFlag(aivVToMte3Event_); + CopyOutputRows( + b, hv, chunk * KDA_FWD_H_CHUNK + tokenBegin, + ioTyped, KDA_FWD_H_SUB_CHUNK); + CrossCoreSetFlag<0x2, PIPE_V>(KDA_FWD_H_DIRECT_FREE_FLAG); + SetFlag(aivMte3ToMte2Event_); + } + + __aicore__ inline void ProcessAiv() + { + const uint32_t subBlockIdx = static_cast(GetSubBlockIdx()); + const uint32_t subBlockNum = static_cast(GetSubBlockNum()); + const uint64_t coreIdx = static_cast(GetBlockIdx()) / subBlockNum; + const uint64_t coreNum = coreNum_ == 0 ? 1 : coreNum_; + LocalTensor state = + resource_.ubBuf.template GetBufferByByte(KDA_FWD_H_UB_STATE_OFFSET); + LocalTensor stateTyped = + resource_.ubBuf.template GetBufferByByte(KDA_FWD_H_UB_STATE_TYPED_OFFSET); + LocalTensor direct = + resource_.ubBuf.template GetBufferByByte(KDA_FWD_H_UB_DIRECT_OFFSET); + LocalTensor out1 = + resource_.ubBuf.template GetBufferByByte(KDA_FWD_H_UB_OUT1_OFFSET); + LocalTensor vnew = + resource_.ubBuf.template GetBufferByByte(KDA_FWD_H_UB_VNEW_OFFSET); + LocalTensor ioTyped = + resource_.ubBuf.template GetBufferByByte(KDA_FWD_H_UB_IO_OFFSET); + LocalTensor gate = + resource_.ubBuf.template GetBufferByByte(KDA_FWD_H_UB_GATE_OFFSET); + + SetFlag(aivMte3ToMte2Event_); + for (uint64_t task = coreIdx; task < batch_ * heads_; task += coreNum) { + const uint64_t b = task / heads_; + const uint64_t hv = task % heads_; + const uint32_t stateRowBegin = subBlockIdx * KDA_FWD_H_SUB_DIM; + InitializeStateAiv(b, hv, stateRowBegin, state); + CrossCoreSetFlag<0x2, PIPE_V>(KDA_FWD_H_DIRECT_FREE_FLAG); + for (uint32_t chunk = 0; + chunk < static_cast(totalChunks_); ++chunk) { + ProcessChunkAiv( + b, hv, chunk, subBlockIdx, + state, stateTyped, direct, out1, vnew, ioTyped, gate); + } + if (storeFinalState_) { + SetFlag(aivVToMte3Event_); + WaitFlag(aivVToMte3Event_); + DataCopy(finalState_[StateOffset(b, hv) + + stateRowBegin * KDA_FWD_H_DIM], + state, KDA_FWD_H_STATE_SUB_ELEMS); + SetFlag(aivMte3ToVEvent_); + WaitFlag(aivMte3ToVEvent_); + } + } + WaitL1SlotFreeMte3( + KDA_FWD_H_STATE_FREE_FLAG, statePublishCount_[subBlockIdx]); + if (fusePostWuIntoFwdH_) { + WaitL1SlotFreeMte3( + KDA_FWD_H_STATE_FREE_FLAG, wPublishCount_[subBlockIdx]); + } + WaitL1SlotFreeMte3( + KDA_FWD_H_VNEW_FREE_FLAG, vnewPublishCount_[subBlockIdx]); + WaitFlag(aivMte3ToMte2Event_); + } + +private: + GlobalTensor gk_; + GlobalTensor initialState_; + GlobalTensor attnOut_; + GlobalTensor finalState_; + GlobalTensor aqk_; + GlobalTensor akk_; + GlobalTensor w_; + GlobalTensor u_; + GlobalTensor qg_; + GlobalTensor kg_; + GlobalTensor vNew_; + GlobalTensor h_; + uint32_t statePublishCount_[2] = {0, 0}; + uint32_t vnewPublishCount_[2] = {0, 0}; + uint32_t wPublishCount_[2] = {0, 0}; + TEventID aicMte2ToMte1Event_ = KDA_FWD_H_MTE_W_EVENT; + TEventID aicL1ReuseEvents_[KDA_FWD_H_L1_STAGING_DEPTH] = {0, 1, 2, 3}; + TEventID stateL0FreeEvent_ = KDA_FWD_H_M_EVENT; + TEventID vnewL0FreeEvent_ = KDA_FWD_H_M_EVENT; + TEventID aicMte1ToMEvent_ = KDA_FWD_H_M_EVENT; + TEventID aicMToFixEvent_ = KDA_FWD_H_FIX_EVENT; + TEventID aicFixToMEvent_ = KDA_FWD_H_FIX_EVENT; + TEventID aivMte2ToVEvent_ = KDA_FWD_H_MTE_W_EVENT; + TEventID aivVToMte3Event_ = KDA_FWD_H_MTE_Q_EVENT; + TEventID aivMte3ToVEvent_ = KDA_FWD_H_MTE_B_EVENT; + TEventID aivMte3ToMte2Event_ = KDA_FWD_H_IO_REUSE_EVENT; + Catlass::Arch::Resource resource_; + uint64_t batch_ = 0; + uint64_t heads_ = 0; + uint64_t seqlen_ = 0; + uint64_t totalChunks_ = 0; + uint64_t coreNum_ = 1; + bool hasInitialState_ = false; + bool storeFinalState_ = false; + bool storeVNew_ = false; + bool storeH_ = false; + bool fusePostWuIntoFwdH_ = false; + float scale_ = 1.0f; +}; + +} // namespace KdaForward::arch35 + +#endif diff --git a/csrc/chunk_kda_fwd/op_kernel/arch35/chunk_kda_fwd_impl.h b/csrc/chunk_kda_fwd/op_kernel/arch35/chunk_kda_fwd_impl.h new file mode 100644 index 0000000000..dab78aadb6 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/arch35/chunk_kda_fwd_impl.h @@ -0,0 +1,61 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#pragma once + +#include "../chunk_kda_fwd_common.h" +#include "chunk_kda_fwd_fwd_h.h" + +namespace KdaForward::arch35 { + +template +__aicore__ inline void Run( + GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR g, GM_ADDR beta, + GM_ADDR aLog, GM_ADDR dtBias, GM_ADDR initialState, + GM_ADDR cuSeqlens, GM_ADDR chunkIndices, GM_ADDR attnOut, + GM_ADDR finalState, GM_ADDR gk, GM_ADDR aqk, GM_ADDR akk, + GM_ADDR w, GM_ADDR u, GM_ADDR qg, GM_ADDR kg, GM_ADDR vNew, GM_ADDR h, + GM_ADDR userWorkspace, const TilingData &tiling, AscendC::TPipe &pipe) +{ + const auto addresses = ResolveAddresses( + finalState, gk, w, u, qg, kg, vNew, h, userWorkspace, tiling); + RunFrontEnd( + q, k, v, g, beta, aLog, dtBias, initialState, cuSeqlens, + chunkIndices, aqk, akk, addresses, userWorkspace, tiling, pipe); + if (tiling.useDenseFwdH) { + ChunkKdaFwdFwdH fwdH; + fwdH.Init( + addresses.gk, initialState, attnOut, addresses.finalState, + aqk, akk, addresses.w, addresses.u, addresses.qgScaled, + addresses.kg, addresses.vNew, addresses.h, tiling); + fwdH.Process(); + return; + } + + const int64_t fwdHTaskCount = + (tiling.isVarLen ? tiling.seqNum : tiling.batch) * tiling.vHeadNum; + const bool isolateGenericBackEnd = + (!tiling.isVarLen && tiling.seqlen % tiling.chunkSize == 0) || + fwdHTaskCount > tiling.prepareUsedCoreNum; + if (isolateGenericBackEnd) { + pipe.Destroy(); + RunGenericBackEnd( + q, k, v, beta, initialState, cuSeqlens, chunkIndices, aqk, + attnOut, addresses, userWorkspace, tiling); + } else { + RunGenericBackEnd( + q, k, v, beta, initialState, cuSeqlens, chunkIndices, aqk, + attnOut, addresses, userWorkspace, tiling, pipe); + } +} + +} // namespace KdaForward::arch35 diff --git a/csrc/chunk_kda_fwd/op_kernel/arch35/chunk_kda_fwd_post_wu_arch35.h b/csrc/chunk_kda_fwd/op_kernel/arch35/chunk_kda_fwd_post_wu_arch35.h new file mode 100644 index 0000000000..a27984d2be --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/arch35/chunk_kda_fwd_post_wu_arch35.h @@ -0,0 +1,2036 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#pragma once + +#ifndef CATLASS_ARCH +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 +#define CATLASS_ARCH 3510 +#else +#define CATLASS_ARCH 2201 +#endif +#endif + +#include "catlass/arch/kda_arch.hpp" +#include "catlass/arch/kda_cross_core_sync.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/kda_catlass.hpp" +#include "catlass/gemm/block/kda_block_mmad.hpp" +#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" +#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" +#include "catlass/gemm/tile/kda_tile_mmad.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "../kernel_utils/block/block_mmad_pingpong_tla_multi.hpp" +#include "catlass/layout/kda_layout.hpp" +#include "kernel_operator.h" +#include "../chunk_kda_fwd_varlen.h" +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 +#ifndef FLA_NPU_REGBASE_HPP_INCLUDED +#define FLA_NPU_REGBASE_HPP_INCLUDED +#include "../kernel_utils/vector/regbase.hpp" +#endif +#endif +#include "tla/kda_tla_layout.hpp" +#include "tla/kda_tensor.hpp" + +using namespace AscendC; + +namespace KdaPostWu { +namespace { +using KdaInt64 = tla::Int<64>; +using KdaInt128 = tla::Int<128>; +constexpr float LN2 = 0.69314718055994530942f; +constexpr float KDA_EXP2_CLAMP = 80.0f; +constexpr float KDA_EXP_INPUT_MAX = KDA_EXP2_CLAMP * LN2; +constexpr float KDA_EXP_INPUT_MIN = -KDA_EXP2_CLAMP * LN2; +constexpr float KDA_FP16_MAX = 65504.0f; +constexpr uint32_t EXP2_UB_ELEMENTS = 256; +constexpr uint32_t EXP2_UB_BYTES = EXP2_UB_ELEMENTS * (sizeof(float) + sizeof(uint16_t)); +constexpr uint32_t EXP2_EVENT_ID = 0; +constexpr uint32_t KDA_SOLVE_BT = 64; +constexpr uint32_t KDA_SOLVE_MATRIX_ELEMENTS = KDA_SOLVE_BT * KDA_SOLVE_BT; +constexpr uint32_t KDA_SOLVE_SCRATCH_X = 0; +constexpr uint32_t KDA_SOLVE_SCRATCH_Y0 = 1; +constexpr uint32_t KDA_SOLVE_SCRATCH_TMP = 2; +constexpr uint32_t KDA_SOLVE_SCRATCH_Y1 = 3; +constexpr uint32_t KDA_SOLVE_SCRATCH_IDENTITY = 4; +constexpr uint32_t KDA_SOLVE_SCRATCH_SLOTS = 5; +constexpr uint32_t KDA_SOLVE_DIAG_BT = 16; +constexpr uint32_t KDA_SOLVE_DIAG_BLOCKS = KDA_SOLVE_BT / KDA_SOLVE_DIAG_BT; +constexpr uint32_t KDA_SOLVE_DIAG_MCH_ITERS = 3; +constexpr uint32_t KDA_SCORE_REF_BC = 16; +constexpr uint32_t KDA_VEC_ARENA_ELEMENTS = 32768; +constexpr uint32_t KDA_BITS_PER_MASK_BYTE = 8; +constexpr uint32_t KDA_SELECT_COL_BLOCKS = 2; +constexpr uint32_t KDA_SELECT_COL_MASK_BYTES = KDA_SOLVE_MATRIX_ELEMENTS / KDA_BITS_PER_MASK_BYTE; +constexpr uint32_t KDA_SELECT_MASK_BYTES = KDA_SELECT_COL_BLOCKS * KDA_SELECT_COL_MASK_BYTES; +constexpr uint32_t KDA_SELECT_AQK_MASK_BYTE_OFFSET = 120 * 1024; +constexpr uint32_t KDA_SELECT_AKK_MASK_BYTE_OFFSET = KDA_SELECT_AQK_MASK_BYTE_OFFSET + KDA_SELECT_MASK_BYTES; +constexpr uint32_t KDA_SELECT_ZERO_BYTE_OFFSET = KDA_SELECT_AKK_MASK_BYTE_OFFSET + KDA_SELECT_MASK_BYTES; +constexpr uint32_t KDA_SELECT_ZERO_FLOAT_OFFSET = KDA_SELECT_ZERO_BYTE_OFFSET / sizeof(float); +constexpr uint8_t KDA_SCORE_DONE_FLAG0 = 2; +constexpr uint8_t KDA_SCORE_DONE_FLAG1 = 3; +constexpr uint8_t KDA_SCORE_READY_FLAG0 = 4; +constexpr uint8_t KDA_SCORE_READY_FLAG1 = 5; +constexpr uint32_t KDA_SCORE_QUEUE_DEPTH = 2; +constexpr uint32_t KDA_SYNC_REVERSE_DEPTH = 1; +constexpr uint32_t KDA_SCORE_SCRATCH_PLANES = 3; +constexpr uint32_t KDA_SCORE_SCRATCH_QG = 0; +constexpr uint32_t KDA_SCORE_SCRATCH_W = 1; +constexpr uint32_t KDA_SCORE_SCRATCH_KG = 2; +constexpr uint64_t KDA_WORKSPACE_ALIGN = 512; +constexpr uint32_t KDA_GATE_TILE_ROWS = 32; +constexpr uint32_t KDA_TYPICAL_GATE_TILE_ROWS = 16; +constexpr uint32_t KDA_TYPICAL_GATE_PIPELINE_ROWS = 32; +constexpr uint16_t KDA_TYPICAL_GATE_PIPELINE_STAGES = 3; +constexpr uint32_t KDA_POST_EVENT = 3; +constexpr uint32_t KDA_POST_EVENT_NEXT = 4; +constexpr uint32_t KDA_POST_EVENT_FIX = 5; +constexpr uint32_t KDA_POST_PIPELINE_L1_SLOT_BYTES = 24 * 1024; +constexpr uint32_t KDA_POST_PIPELINE_L1_A_BYTES = 64 * 64 * sizeof(uint16_t); +constexpr uint32_t KDA_POST_PIPELINE_L1_B_BYTES = 64 * 128 * sizeof(uint16_t); +constexpr uint32_t KDA_POST_PIPELINE_L1_U_SLOT_BYTES = 64 * 128 * sizeof(uint16_t); +constexpr uint32_t KDA_POST_PIPELINE_L0_A_SLOT_BYTES = 64 * 64 * sizeof(uint16_t); +constexpr uint32_t KDA_POST_PIPELINE_L0_B_SLOT_BYTES = 64 * 256 * sizeof(uint16_t); +constexpr uint32_t KDA_POST_PIPELINE_L0_C_SLOT_BYTES = 64 * 256 * sizeof(float); +constexpr uint16_t KDA_POST_PIPELINE_STAGE_COUNT = 2; +constexpr uint16_t KDA_POST_FUSED_BATCH_TASKS = 4; +constexpr uint16_t KDA_POST_HEAD_PAIR_LANES = 2; +constexpr uint16_t KDA_POST_PIPELINE_U_EVENT = KDA_POST_EVENT_FIX; +constexpr uint16_t KDA_POST_REGBASE_TAIL_LIMIT = 16; +constexpr bool KDA_ENABLE_POST_AIC_PIPELINE = true; + +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 +static __simd_vf__ inline void ComputeKdaTailWuRegbase(__ubuf__ float *wOut, __ubuf__ float *uOut, + __ubuf__ bfloat16_t *akk, __ubuf__ bfloat16_t *wSeed, + __ubuf__ bfloat16_t *uSeed, uint16_t rows, uint16_t curT) +{ + using namespace AscendC::MicroAPI; + constexpr uint16_t DIM = 128; + constexpr uint16_t FP32_REG_ELEMENTS = AscendC::VECTOR_REG_WIDTH / sizeof(float); + constexpr uint16_t AKK_ROW_ELEMENTS = 64; + MaskReg fullMask = CreateMask(); + + for (uint16_t row = 0; row < rows; ++row) { + RegTensor wAcc0; + RegTensor wAcc1; + RegTensor uAcc0; + RegTensor uAcc1; + Duplicate(wAcc0, 0.0f, fullMask); + Duplicate(wAcc1, 0.0f, fullMask); + Duplicate(uAcc0, 0.0f, fullMask); + Duplicate(uAcc1, 0.0f, fullMask); + + for (uint16_t token = 0; token < curT; ++token) { + RegTensor wRaw0; + RegTensor wRaw1; + RegTensor uRaw0; + RegTensor uRaw1; + RegTensor weightRaw; + RegTensor w0; + RegTensor w1; + RegTensor u0; + RegTensor u1; + RegTensor weight; + RegTensor product; + DataCopy(wRaw0, wSeed + static_cast(token) * DIM); + DataCopy( + wRaw1, wSeed + static_cast(token) * DIM + FP32_REG_ELEMENTS); + DataCopy(uRaw0, uSeed + static_cast(token) * DIM); + DataCopy( + uRaw1, uSeed + static_cast(token) * DIM + FP32_REG_ELEMENTS); + LoadAlign(weightRaw, + akk + static_cast(row) * AKK_ROW_ELEMENTS + token); + Cast(w0, wRaw0, fullMask); + Cast(w1, wRaw1, fullMask); + Cast(u0, uRaw0, fullMask); + Cast(u1, uRaw1, fullMask); + Cast(weight, weightRaw, fullMask); + Mul(product, w0, weight, fullMask); + Add(wAcc0, wAcc0, product, fullMask); + Mul(product, w1, weight, fullMask); + Add(wAcc1, wAcc1, product, fullMask); + Mul(product, u0, weight, fullMask); + Add(uAcc0, uAcc0, product, fullMask); + Mul(product, u1, weight, fullMask); + Add(uAcc1, uAcc1, product, fullMask); + } + + uint32_t outputOffset = static_cast(row) * DIM; + DataCopy(wOut + outputOffset, wAcc0, fullMask); + DataCopy(wOut + outputOffset + FP32_REG_ELEMENTS, wAcc1, fullMask); + DataCopy(uOut + outputOffset, uAcc0, fullMask); + DataCopy(uOut + outputOffset + FP32_REG_ELEMENTS, uAcc1, fullMask); + } +} + +template +__simd_callee__ inline void LoadPostKdaGateRegbasePair(AscendC::MicroAPI::RegTensor &zeroReg, + AscendC::MicroAPI::RegTensor &oneReg, + __ubuf__ InputT *src, AscendC::MicroAPI::MaskReg &inputMask) +{ + using namespace AscendC::MicroAPI; + if constexpr (std::is_same()) { + LoadAlign(zeroReg, oneReg, src); + } else { + RegTensor inputReg; + LoadIn(inputReg, src); + CastHalf2Float(zeroReg, oneReg, inputReg, inputMask); + } +} + +template +__simd_callee__ inline void StorePostKdaGateRegbasePair(__ubuf__ OutputT *dst, + AscendC::MicroAPI::RegTensor &zeroReg, + AscendC::MicroAPI::RegTensor &oneReg, + AscendC::MicroAPI::MaskReg &inputMask, + AscendC::MicroAPI::MaskReg &floatMask) +{ + using namespace AscendC::MicroAPI; + if constexpr (std::is_same()) { + Mins(zeroReg, zeroReg, KDA_FP16_MAX, floatMask); + Mins(oneReg, oneReg, KDA_FP16_MAX, floatMask); + Maxs(zeroReg, zeroReg, -KDA_FP16_MAX, floatMask); + Maxs(oneReg, oneReg, -KDA_FP16_MAX, floatMask); + } + RegTensor outputReg; + CastFloat2Half(outputReg, zeroReg, oneReg, floatMask); + StoreAlign(dst, outputReg, inputMask); +} + +template +static __simd_vf__ inline void ComputePostKdaKgRegbase(__ubuf__ T *kAndKg, __ubuf__ GK_T *gate, __ubuf__ float *ref, + uint16_t rows, uint16_t cols) +{ + using namespace AscendC::MicroAPI; + constexpr uint16_t ELEMENTS_PER_REG = AscendC::VECTOR_REG_WIDTH / sizeof(T); + MaskReg floatMask = CreateMask(); + for (uint16_t row = 0; row < rows; ++row) { + uint32_t rowOffset = static_cast(row) * cols; + for (uint16_t col = 0; col < cols; col += ELEMENTS_PER_REG) { + uint32_t activeCount = static_cast(cols - col); + MaskReg inputMask = UpdateMask(activeCount); + uint32_t offset = rowOffset + col; + + RegTensor gateZeroReg; + RegTensor gateOneReg; + RegTensor refZeroReg; + RegTensor refOneReg; + RegTensor expZeroReg; + RegTensor expOneReg; + RegTensor inputZeroReg; + RegTensor inputOneReg; + RegTensor outputZeroReg; + RegTensor outputOneReg; + + LoadPostKdaGateRegbasePair(gateZeroReg, gateOneReg, gate + offset, inputMask); + LoadAlign(refZeroReg, refOneReg, ref + col); + SubFloatTwoReg(expZeroReg, expOneReg, refZeroReg, refOneReg, gateZeroReg, gateOneReg, floatMask); + Muls(expZeroReg, expZeroReg, LN2, floatMask); + Muls(expOneReg, expOneReg, LN2, floatMask); + MinsFloatTwoReg(expZeroReg, expOneReg, expZeroReg, expOneReg, KDA_EXP_INPUT_MAX, floatMask); + Maxs(expZeroReg, expZeroReg, KDA_EXP_INPUT_MIN, floatMask); + Maxs(expOneReg, expOneReg, KDA_EXP_INPUT_MIN, floatMask); + ExpFloatTwoReg(expZeroReg, expOneReg, expZeroReg, expOneReg, floatMask); + + LoadPostKdaGateRegbasePair(inputZeroReg, inputOneReg, kAndKg + offset, inputMask); + MulFloatTwoReg(outputZeroReg, outputOneReg, inputZeroReg, inputOneReg, expZeroReg, expOneReg, floatMask); + StorePostKdaGateRegbasePair(kAndKg + offset, outputZeroReg, outputOneReg, inputMask, floatMask); + } + } +} +#endif + +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 +using KdaArchTag = Catlass::Arch::Ascend950; +#else +using KdaArchTag = Catlass::Arch::AtlasA2; +#endif +using KdaDispatchPolicy = Catlass::Gemm::MmadPingpong; +using KdaScoreDispatchPolicy = Catlass::Gemm::MmadPingpongTlaMulti; +static_assert(KdaScoreDispatchPolicy::ENABLE_L1_RESIDENT, + "KDA Aqk/Akk score MMAD must keep the shared right matrix resident in L1"); +static_assert(KdaScoreDispatchPolicy::L1B_STAGES == 1, + "KDA Aqk/Akk score MMAD needs one L1 B slot so the second MMAD reuses it"); +using KdaSolveDispatchPolicy = Catlass::Gemm::MmadPingpong; +static_assert(!KdaSolveDispatchPolicy::USE_HF32_MODE, "KDA triangular solve must use IEEE FP32 Cube mode"); +using KdaL1TileShape = tla::Shape; +using KdaL0TileShape = KdaL1TileShape; +using KdaSolveL1TileShape = tla::Shape; +using KdaSolveL0TileShape = KdaSolveL1TileShape; + +__aicore__ inline uint32_t FloatToBits(float value) +{ + union Bits { + __aicore__ Bits() {} + float f; + uint32_t u; + } bits; + bits.f = value; + return bits.u; +} + +__aicore__ inline float BitsToFloat(uint32_t value) +{ + union Bits { + __aicore__ Bits() {} + uint32_t u; + float f; + } bits; + bits.u = value; + return bits.f; +} + +__aicore__ inline uint16_t Bf16ToBits(bfloat16_t value) +{ + union Bits { + __aicore__ Bits() {} + bfloat16_t f; + uint16_t u; + } bits; + bits.f = value; + return bits.u; +} + +__aicore__ inline bfloat16_t BitsToBf16(uint16_t value) +{ + union Bits { + __aicore__ Bits() {} + uint16_t u; + bfloat16_t f; + } bits; + bits.u = value; + return bits.f; +} + +template +__aicore__ inline T FloatToType(float value) +{ + if constexpr (IsSameType::value) { + uint32_t bits = FloatToBits(value); + uint32_t bias = 0x7FFFu + ((bits >> 16) & 1u); + return BitsToBf16(static_cast((bits + bias) >> 16)); + } + return static_cast(value); +} + +template +class ChunkKdaFwdPostWuKernel { +public: + using OUT_T = T; + using AKK_T = T; + template + __aicore__ inline void Init(GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR gk, GM_ADDR beta, GM_ADDR initialState, + GM_ADDR cuSeqlens, GM_ADDR chunkIndices, GM_ADDR preparedQG, GM_ADDR preparedAqk, + GM_ADDR propagatedVNew, GM_ADDR propagatedH, GM_ADDR o, GM_ADDR finalState, GM_ADDR aqk, + GM_ADDR akk, GM_ADDR w, GM_ADDR u, GM_ADDR qg, GM_ADDR kg, GM_ADDR vNew, GM_ADDR h, + GM_ADDR workspace, const TilingData &tiling, TPipe *pipe, bool initVecBuffers = true) + { + pipe_ = pipe; + q_.SetGlobalBuffer((__gm__ T *)q); + k_.SetGlobalBuffer((__gm__ T *)k); + v_.SetGlobalBuffer((__gm__ T *)v); + gk_.SetGlobalBuffer((__gm__ GK_T *)gk); + beta_.SetGlobalBuffer((__gm__ BETA_T *)beta); + if (initialState != nullptr) { + initialState_.SetGlobalBuffer((__gm__ float *)initialState); + } + cuSeqlensAddr_ = reinterpret_cast<__gm__ int64_t *>(cuSeqlens); + if (preparedQG != nullptr) { + preparedQG_.SetGlobalBuffer((__gm__ T *)preparedQG); + } + if (preparedAqk != nullptr) { + preparedAqk_.SetGlobalBuffer((__gm__ T *)preparedAqk); + } + if (propagatedVNew != nullptr) { + propagatedVNew_.SetGlobalBuffer((__gm__ T *)propagatedVNew); + } + if (propagatedH != nullptr) { + propagatedH_.SetGlobalBuffer((__gm__ T *)propagatedH); + } + chunkIndicesAddr_ = reinterpret_cast<__gm__ int64_t *>(chunkIndices); + o_.SetGlobalBuffer((__gm__ OUT_T *)o); + finalState_.SetGlobalBuffer((__gm__ float *)finalState); + aqk_.SetGlobalBuffer((__gm__ float *)aqk); + akk_.SetGlobalBuffer((__gm__ AKK_T *)akk); + w_.SetGlobalBuffer((__gm__ T *)w); + u_.SetGlobalBuffer((__gm__ OUT_T *)u); + qg_.SetGlobalBuffer((__gm__ T *)qg); + kg_.SetGlobalBuffer((__gm__ T *)kg); + vNew_.SetGlobalBuffer((__gm__ T *)vNew); + h_.SetGlobalBuffer((__gm__ float *)h); + solveWorkspace_.SetGlobalBuffer((__gm__ float *)workspace); + + B_ = tiling.batch; + N_ = tiling.seqNum; + H_ = tiling.qHeadNum; + HV_ = tiling.vHeadNum; + T_ = tiling.seqlen; + K_ = tiling.kHeadDim; + V_ = tiling.vHeadDim; + BT_ = tiling.chunkSize; + NT_ = tiling.totalChunks; + scale_ = tiling.scale; + hasInitial_ = tiling.hasInitialState; + isVarLen_ = tiling.isVarLen; + inputSequenceMajor_ = tiling.inputSequenceMajor; + usedCoreNum_ = tiling.postWuUsedCoreNum; + if ASCEND_IS_AIV { + uint64_t subBlockNum = static_cast(GetSubBlockNum()); + solveCoreIdx_ = subBlockNum == 0 ? 0 : static_cast(GetBlockIdx()) / subBlockNum; + } else { + solveCoreIdx_ = static_cast(GetBlockIdx()); + } + if (pipe_ != nullptr && initVecBuffers) { + pipe_->InitBuffer(exp2Buf_, EXP2_UB_BYTES); + pipe_->InitBuffer(vecBuf_, KDA_VEC_ARENA_ELEMENTS * sizeof(float)); + const uint64_t gateWritebackRows = ScoreVectorMaxRows(5 * sizeof(float) + 2 * sizeof(T) + sizeof(GK_T)); + pipe_->InitBuffer(gateWritebackBuf_, + static_cast(gateWritebackRows * K_ * (3 * sizeof(T) + sizeof(GK_T)))); + AllocVectorEvents(); + } + } + __aicore__ inline void ProcessAiv() + { + ProcessPostAiv(); + ReleaseVectorEvents(); + } + + __aicore__ inline void ProcessAic() { ProcessPostAic(); } + +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + __aicore__ inline void ProcessPreparedFullHeadPairBatchArch35(const uint64_t *batchB, const uint64_t *batchHvBase, + const uint64_t *batchStart, uint16_t taskCount) + { + if (taskCount == 0) { + return; + } + SetLoadDataPaddingValue(static_cast(0)); + Catlass::Arch::Resource resource; + const uint16_t itemCount = taskCount * KDA_POST_HEAD_PAIR_LANES; + uint16_t slot = 0; + uint16_t usedSlotCount = 1; + uint16_t taskIdx = 0; + uint16_t lane = 0; + uint64_t b = batchB[taskIdx]; + uint64_t hv = batchHvBase[taskIdx] + lane; + uint64_t start = batchStart[taskIdx]; + InitializePostWuPipelineEvents(); + PrefetchPostWuPipelineArch35(resource, slot, b, hv, start, BT_, false); + PrefetchPostWuPipelineU(resource, slot, b, hv, start, BT_, false); + + for (uint16_t item = 0; item < itemCount; ++item) { + const uint16_t nextItem = item + 1; + if (nextItem < itemCount) { + const uint16_t nextTaskIdx = nextItem / KDA_POST_HEAD_PAIR_LANES; + const uint16_t nextLane = nextItem % KDA_POST_HEAD_PAIR_LANES; + const uint16_t nextSlot = slot ^ 1; + const bool reuseSlot = nextItem >= KDA_POST_PIPELINE_STAGE_COUNT; + PrefetchPostWuPipelineArch35(resource, nextSlot, batchB[nextTaskIdx], + batchHvBase[nextTaskIdx] + nextLane, batchStart[nextTaskIdx], BT_, + reuseSlot); + PrefetchPostWuPipelineU(resource, nextSlot, batchB[nextTaskIdx], batchHvBase[nextTaskIdx] + nextLane, + batchStart[nextTaskIdx], BT_, reuseSlot); + if (!reuseSlot) { + ++usedSlotCount; + } + } + + ComputePrefetchedPostWuPipelineArch35(resource, slot, b, hv, start, BT_); + if (nextItem < itemCount) { + taskIdx = nextItem / KDA_POST_HEAD_PAIR_LANES; + lane = nextItem % KDA_POST_HEAD_PAIR_LANES; + b = batchB[taskIdx]; + hv = batchHvBase[taskIdx] + lane; + start = batchStart[taskIdx]; + slot ^= 1; + } + } + FinalizePostWuPipelineEvents(usedSlotCount); + } + + __aicore__ inline void ProcessPreparedTailHeadPairArch35(uint64_t b, uint64_t hvBase, uint64_t start, uint64_t curT) + { + SetLoadDataPaddingValue(static_cast(0)); + Catlass::Arch::Resource resource; + InitializePostWuPipelineEvents(); + for (uint16_t lane = 0; lane < KDA_POST_HEAD_PAIR_LANES; ++lane) { + PrefetchPostWuPipelineArch35(resource, lane, b, hvBase + lane, start, curT, false); + PrefetchPostWuPipelineU(resource, lane, b, hvBase + lane, start, curT, false); + } + for (uint16_t lane = 0; lane < KDA_POST_HEAD_PAIR_LANES; ++lane) { + ComputePrefetchedPostWuPipelineArch35(resource, lane, b, hvBase + lane, start, curT); + } + FinalizePostWuPipelineEvents(KDA_POST_HEAD_PAIR_LANES); + } + + __aicore__ inline void ProcessPreparedTailSingleArch35(uint64_t b, uint64_t hv, uint64_t start, uint64_t curT) + { + SetLoadDataPaddingValue(static_cast(0)); + Catlass::Arch::Resource resource; + InitializePostWuPipelineSlot(0); + PrefetchPostWuPipelineArch35(resource, 0, b, hv, start, curT, false); + PrefetchPostWuPipelineU(resource, 0, b, hv, start, curT, false); + ComputePrefetchedPostWuPipelineArch35(resource, 0, b, hv, start, curT); + FinalizePostWuPipelineEvents(1); + } + + __aicore__ inline void ProcessPreparedHeadPairBatchArch35(const uint64_t *batchB, const uint64_t *batchHvBase, + const uint64_t *batchStart, const uint64_t *batchEnd, + uint16_t taskCount) + { + uint16_t fullRunBegin = 0; + for (uint16_t task = 0; task < taskCount; ++task) { + if (batchEnd[task] - batchStart[task] == BT_) { + continue; + } + if (task > fullRunBegin) { + ProcessPreparedFullHeadPairBatchArch35(batchB + fullRunBegin, batchHvBase + fullRunBegin, + batchStart + fullRunBegin, task - fullRunBegin); + } + ProcessPreparedTailHeadPairArch35(batchB[task], batchHvBase[task], batchStart[task], + batchEnd[task] - batchStart[task]); + fullRunBegin = task + 1; + } + if (fullRunBegin < taskCount) { + ProcessPreparedFullHeadPairBatchArch35(batchB + fullRunBegin, batchHvBase + fullRunBegin, + batchStart + fullRunBegin, taskCount - fullRunBegin); + } + } +#endif + +private: + __aicore__ inline void AllocVectorEvents() + { + mte2ToVEvent_ = pipe_->AllocEventID(); + vToMte2Event_ = pipe_->AllocEventID(); + vToMte3Event_ = pipe_->AllocEventID(); + mte3ToVEvent_ = pipe_->AllocEventID(); + mte2ToMte3Event_ = pipe_->AllocEventID(); + mte3ToMte2Event_ = pipe_->AllocEventID(); + vectorEventsAllocated_ = true; + } + + __aicore__ inline void ReleaseVectorEvents() + { + if (!vectorEventsAllocated_) { + return; + } + pipe_->ReleaseEventID(mte2ToVEvent_); + pipe_->ReleaseEventID(vToMte2Event_); + pipe_->ReleaseEventID(vToMte3Event_); + pipe_->ReleaseEventID(mte3ToVEvent_); + pipe_->ReleaseEventID(mte2ToMte3Event_); + pipe_->ReleaseEventID(mte3ToMte2Event_); + vectorEventsAllocated_ = false; + } + + __aicore__ inline uint64_t QOffset(uint64_t b, uint64_t h, uint64_t t, uint64_t d) const + { + if (inputSequenceMajor_) { + return ((b * T_ + t) * H_ + h) * K_ + d; + } + return ((b * H_ + h) * T_ + t) * K_ + d; + } + + __aicore__ inline uint64_t KVOffset(uint64_t b, uint64_t hv, uint64_t t, uint64_t d, uint64_t dim) const + { + return ((b * HV_ + hv) * T_ + t) * dim + d; + } + + __aicore__ inline uint64_t BetaOffset(uint64_t b, uint64_t hv, uint64_t t) const { return (b * HV_ + hv) * T_ + t; } + + __aicore__ inline uint64_t AOffset(uint64_t b, uint64_t hv, uint64_t t, uint64_t j) const + { + return ((b * HV_ + hv) * T_ + t) * BT_ + j; + } + + __aicore__ inline uint64_t HOffset(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t d, uint64_t r) const + { + return (((b * HV_ + hv) * NT_ + chunkIdx) * K_ + d) * V_ + r; + } + + __aicore__ inline uint64_t WScratchOffset(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t t, uint64_t d) const + { + return (((b * HV_ + hv) * NT_ + chunkIdx) * BT_ + t) * K_ + d; + } + + __aicore__ inline uint64_t SolveScratchOffset(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t slot) const + { + (void)b; + (void)hv; + (void)chunkIdx; + uint64_t matrixElements = BT_ * BT_; + return solveCoreIdx_ * KDA_SOLVE_SCRATCH_SLOTS * matrixElements + slot * matrixElements; + } + + __aicore__ inline uint64_t ScoreScratchOffset(uint64_t slot, uint64_t plane, uint64_t t = 0, uint64_t d = 0) const + { + return (((solveCoreIdx_ * KDA_SCORE_QUEUE_DEPTH + slot) * KDA_SCORE_SCRATCH_PLANES + plane) * BT_ + t) * K_ + d; + } + + __aicore__ inline uint64_t ScoreRefBlockSize() const { return KDA_SCORE_REF_BC; } + + __aicore__ inline uint64_t ScoreRowBlockCount(uint64_t curT, uint64_t rowBegin) const + { + uint64_t blockSize = ScoreRefBlockSize(); + uint64_t rowCount = curT - rowBegin; + if (rowCount > blockSize) { + rowCount = blockSize; + } + return rowCount; + } + + __aicore__ inline uint64_t ScoreRefToken(uint64_t start, uint64_t curT, uint64_t rowBegin, uint64_t rowCount) const + { + uint64_t ref = rowBegin + rowCount / 2; + if (ref >= curT) { + ref = curT - 1; + } + return start + ref; + } + + __aicore__ inline void RunExp2(LocalTensor &tensor, uint32_t count) + { + SetFlag(EXP2_EVENT_ID); + WaitFlag(EXP2_EVENT_ID); + ClampExpInput(tensor, count); + Exp(tensor, tensor, count); + PipeBarrier(); + SetFlag(EXP2_EVENT_ID); + WaitFlag(EXP2_EVENT_ID); + } + + __aicore__ inline void ClampExpInput(LocalTensor &tensor, uint32_t count) + { + Mins(tensor, tensor, KDA_EXP_INPUT_MAX, count); + PipeBarrier(); + Maxs(tensor, tensor, KDA_EXP_INPUT_MIN, count); + PipeBarrier(); + } + + __aicore__ inline void ClampFp32ToOutputType(LocalTensor &tensor, uint32_t count) + { + if constexpr (IsSameType::value) { + Mins(tensor, tensor, KDA_FP16_MAX, count); + PipeBarrier(); + Maxs(tensor, tensor, -KDA_FP16_MAX, count); + PipeBarrier(); + } + } + + template + __aicore__ inline void CopyVectorIn(LocalTensor &dst, GlobalTensor &src, uint64_t offset, + uint64_t count) + { + uint64_t rowBytes = count * static_cast(sizeof(CopyT)); + if (rowBytes >= 32 && rowBytes % 32 == 0) { + DataCopy(dst, src[offset], static_cast(count)); + return; + } + DataCopyParams params{1, static_cast(rowBytes), 0, 0}; + DataCopyPadParams padParams{false, 0, 0, 0}; + DataCopyPad(dst, src[offset], params, padParams); + } + + template + __aicore__ inline void CopyRowsIn(LocalTensor &dst, GlobalTensor &src, uint64_t offset, uint64_t rows, + uint64_t cols, uint64_t rowStride) + { + if (rows == 0 || cols == 0) { + return; + } + if (rowStride == cols) { + CopyVectorIn(dst, src, offset, rows * cols); + return; + } + DataCopyExtParams params{static_cast(rows), static_cast(cols * sizeof(CopyT)), + static_cast((rowStride - cols) * sizeof(CopyT)), 0, 0}; + DataCopyPadExtParams padParams{false, 0, 0, 0}; + DataCopyPad(dst, src[offset], params, padParams); + } + + template + __aicore__ inline void CopyVectorOut(GlobalTensor &dst, uint64_t offset, LocalTensor &src, + uint64_t count) + { + uint64_t rowBytes = count * static_cast(sizeof(CopyT)); + if (rowBytes >= 32 && rowBytes % 32 == 0) { + DataCopy(dst[offset], src, static_cast(count)); + return; + } + DataCopyParams params{1, static_cast(rowBytes), 0, 0}; + DataCopyPad(dst[offset], src, params); + } + + template + __aicore__ inline void CopyRowIn(LocalTensor &dst, GlobalTensor &src, uint64_t offset) + { + CopyVectorIn(dst, src, offset, K_); + } + + template + __aicore__ inline void CopyRowOut(GlobalTensor &dst, uint64_t offset, LocalTensor &src) + { + CopyVectorOut(dst, offset, src, K_); + } + + __aicore__ inline LocalTensor VecScratch(uint64_t slot) + { + return vecBuf_.Get()[slot * EXP2_UB_ELEMENTS]; + } + + template + __aicore__ inline void LoadAsFloatRow(GlobalTensor &src, uint64_t srcOffset, LocalTensor &dst, + uint64_t count) + { + if constexpr (IsSameType::value) { + CopyVectorIn(dst, src, srcOffset, count); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + Adds(dst, dst, 0.0f, static_cast(count)); + PipeBarrier(); + SetFlag(vToMte2Event_); + WaitFlag(vToMte2Event_); + } else { + constexpr uint32_t typedOffset = EXP2_UB_ELEMENTS * sizeof(float) / sizeof(CopyT); + LocalTensor rowLocal = exp2Buf_.Get()[typedOffset]; + CopyVectorIn(rowLocal, src, srcOffset, count); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + Cast(dst, rowLocal, RoundMode::CAST_NONE, static_cast(count)); + PipeBarrier(); + SetFlag(vToMte2Event_); + WaitFlag(vToMte2Event_); + } + PipeBarrier(); + } + + template + __aicore__ inline void LoadAsFloatVector(GlobalTensor &src, uint64_t srcOffset, LocalTensor &dst, + LocalTensor &typedScratch, uint64_t count) + { + if constexpr (IsSameType::value) { + CopyVectorIn(dst, src, srcOffset, count); + } else { + CopyVectorIn(typedScratch, src, srcOffset, count); + } + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + if constexpr (!IsSameType::value) { + Cast(dst, typedScratch, RoundMode::CAST_NONE, static_cast(count)); + PipeBarrier(); + } + } + + template + __aicore__ inline void StoreFloatRow(GlobalTensor &dst, uint64_t dstOffset, LocalTensor &src, + uint64_t count) + { + if constexpr (IsSameType::value) { + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + CopyVectorOut(dst, dstOffset, src, count); + } else { + constexpr uint32_t typedOffset = EXP2_UB_ELEMENTS * sizeof(float) / sizeof(CopyT); + LocalTensor rowLocal = exp2Buf_.Get()[typedOffset]; + Cast(rowLocal, src, RoundMode::CAST_RINT, static_cast(count)); + PipeBarrier(); + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + CopyVectorOut(dst, dstOffset, rowLocal, count); + } + SetFlag(mte3ToMte2Event_); + WaitFlag(mte3ToMte2Event_); + SetFlag(mte3ToVEvent_); + WaitFlag(mte3ToVEvent_); + } + + __aicore__ inline LocalTensor Exp2NegG(uint64_t b, uint64_t hv, uint64_t t) + { + LocalTensor exp2Local = exp2Buf_.Get(); + LoadAsFloatRow(gk_, KVOffset(b, hv, t, 0, K_), exp2Local, K_); + Muls(exp2Local, exp2Local, -LN2, static_cast(K_)); + PipeBarrier(); + RunExp2(exp2Local, static_cast(K_)); + return exp2Local; + } + + __aicore__ inline uint64_t ScoreVectorMaxRows(uint64_t bytesPerElem) const + { + constexpr uint64_t arenaBytes = static_cast(KDA_VEC_ARENA_ELEMENTS) * sizeof(float); + uint64_t maxRows = (arenaBytes / bytesPerElem) / K_; + if (K_ >= 128 && maxRows > 32) { + maxRows = 32; + } + return maxRows; + } + __aicore__ inline bool UsePostWuCube(uint64_t curT) const + { + return curT > 0 && curT <= BT_ && (BT_ == 64 || BT_ == 128) && K_ >= 16 && V_ >= 16 && V_ <= 256 && + K_ % 16 == 0 && V_ % 16 == 0; + } + +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + __aicore__ inline bool UsePostWuCubeArch35(uint64_t curT) const + { + return curT == 64 && BT_ == 64 && K_ == 128 && V_ == 128; + } + + __aicore__ inline bool UseFullPostWuPipelineArch35(uint64_t curT) const + { + return curT == 64 && BT_ == 64 && K_ == 128 && V_ == 128; + } + + __aicore__ inline void ComputePostWuCubeFusedArch35(uint64_t b, uint64_t hv, uint64_t start, uint64_t curT) + { + using ElementA = T; + using ElementB = T; + using ElementC = T; + using LayoutTagA = Catlass::layout::RowMajor; + using LayoutTagB = Catlass::layout::RowMajor; + using LayoutTagC = Catlass::layout::RowMajor; + using TileCopy = Catlass::Gemm::Tile::PackedTileCopyTla; + + constexpr uint32_t capacityM = 64; + constexpr uint32_t n = 128; + constexpr uint32_t capacityK = 64; + const uint32_t m = static_cast(curT); + const uint32_t k = static_cast(curT); + SetLoadDataPaddingValue(static_cast(0)); + + auto layoutA = tla::MakeLayout(capacityM, capacityK); + auto layoutB = tla::MakeLayout(capacityK, n); + auto layoutC = tla::MakeLayout(capacityM, n); + auto tensorA = tla::MakeTensor(preparedAqk_[AOffset(b, hv, start, 0)], layoutA, Catlass::Arch::PositionGM{}); + auto tensorW = + tla::MakeTensor(preparedQG_[KVOffset(b, hv, start, 0, K_)], layoutB, Catlass::Arch::PositionGM{}); + auto tensorU = + tla::MakeTensor(propagatedVNew_[KVOffset(b, hv, start, 0, V_)], layoutB, Catlass::Arch::PositionGM{}); + auto tensorWOut = tla::MakeTensor(w_[KVOffset(b, hv, start, 0, K_)], layoutC, Catlass::Arch::PositionGM{}); + auto tensorUOut = tla::MakeTensor(u_[KVOffset(b, hv, start, 0, V_)], layoutC, Catlass::Arch::PositionGM{}); + auto blockA = GetTile(tensorA, tla::MakeCoord(0, 0), tla::MakeShape(m, k)); + auto blockW = GetTile(tensorW, tla::MakeCoord(0, 0), tla::MakeShape(k, n)); + auto blockU = GetTile(tensorU, tla::MakeCoord(0, 0), tla::MakeShape(k, n)); + auto blockWOut = GetTile(tensorWOut, tla::MakeCoord(0, 0), tla::MakeShape(m, n)); + auto blockUOut = GetTile(tensorUOut, tla::MakeCoord(0, 0), tla::MakeShape(m, n)); + + Catlass::Arch::Resource resource; + constexpr uint32_t aBytes = capacityM * capacityK * sizeof(ElementA); + constexpr uint32_t bBytes = capacityK * n * sizeof(ElementB); + LocalTensor l1A = resource.l1Buf.template GetBufferByByte(0); + LocalTensor l1B0 = resource.l1Buf.template GetBufferByByte(aBytes); + LocalTensor l1B1 = resource.l1Buf.template GetBufferByByte(aBytes + bBytes); + LocalTensor l0A = resource.l0ABuf.template GetBufferByByte(0); + LocalTensor l0B = resource.l0BBuf.template GetBufferByByte(0); + LocalTensor l0C = resource.l0CBuf.template GetBufferByByte(0); + + using LayoutTagL1A = typename TileCopy::LayoutTagL1A; + using LayoutTagL1B = typename TileCopy::LayoutTagL1B; + using LayoutTagL0A = typename TileCopy::LayoutTagL0A; + using LayoutTagL0B = typename TileCopy::LayoutTagL0B; + using CopyGmToL1A = typename TileCopy::template CopyGmToL1A; + using CopyGmToL1B = typename TileCopy::template CopyGmToL1B; + using CopyL1ToL0A = typename TileCopy::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopy::CopyL1ToL0B; + using CopyL0CToDst = typename TileCopy::template CopyL0CToDst; + using TileMmad = Catlass::Gemm::Tile::TileMmadTla; + + auto layoutL1A = tla::MakeLayout(capacityM, capacityK); + auto layoutL1B = tla::MakeLayout(capacityK, n); + auto layoutL0A = tla::MakeLayout(capacityM, capacityK); + auto layoutL0B = tla::MakeLayout(capacityK, n); + auto layoutL0C = tla::MakeLayoutL0C(capacityM, n); + auto tensorL1A = tla::MakeTensor(l1A, layoutL1A, Catlass::Arch::PositionL1{}); + auto tensorL1B0 = tla::MakeTensor(l1B0, layoutL1B, Catlass::Arch::PositionL1{}); + auto tensorL1B1 = tla::MakeTensor(l1B1, layoutL1B, Catlass::Arch::PositionL1{}); + auto tensorL0A = tla::MakeTensor(l0A, layoutL0A, Catlass::Arch::PositionL0A{}); + auto tensorL0B = tla::MakeTensor(l0B, layoutL0B, Catlass::Arch::PositionL0B{}); + auto tensorL0C = tla::MakeTensor(l0C, layoutL0C, Catlass::Arch::PositionL0C{}); + auto tileL1A = GetTile(tensorL1A, tla::MakeCoord(0, 0), tla::MakeShape(m, k)); + auto tileL1B0 = GetTile(tensorL1B0, tla::MakeCoord(0, 0), tla::MakeShape(k, n)); + auto tileL1B1 = GetTile(tensorL1B1, tla::MakeCoord(0, 0), tla::MakeShape(k, n)); + auto tileL0A = GetTile(tensorL0A, tla::MakeCoord(0, 0), tla::MakeShape(m, k)); + auto tileL0B = GetTile(tensorL0B, tla::MakeCoord(0, 0), tla::MakeShape(k, n)); + auto tileL0C = GetTile(tensorL0C, tla::MakeCoord(0, 0), tla::MakeShape(m, n)); + + CopyGmToL1A copyGmToL1A; + CopyGmToL1B copyGmToL1B; + CopyL1ToL0A copyL1ToL0A; + CopyL1ToL0B copyL1ToL0B; + CopyL0CToDst copyL0CToDst; + TileMmad tileMmad; + + copyGmToL1A(tensorL1A, blockA); + copyGmToL1B(tensorL1B0, blockW); + SetFlag(KDA_POST_EVENT); + WaitFlag(KDA_POST_EVENT); + copyL1ToL0A(tileL0A, tileL1A); + copyL1ToL0B(tileL0B, tileL1B0); + SetFlag(KDA_POST_EVENT); + copyGmToL1B(tensorL1B1, blockU); + SetFlag(KDA_POST_EVENT_NEXT); + WaitFlag(KDA_POST_EVENT); + tileMmad(tileL0C, tileL0A, tileL0B, m, n, k, true, 0); + SetFlag(KDA_POST_EVENT); + SetFlag(KDA_POST_EVENT); + WaitFlag(KDA_POST_EVENT_NEXT); + WaitFlag(KDA_POST_EVENT); + WaitFlag(KDA_POST_EVENT); + copyL0CToDst(blockWOut, tileL0C); + SetFlag(KDA_POST_EVENT_FIX); + + copyL1ToL0B(tileL0B, tileL1B1); + SetFlag(KDA_POST_EVENT_NEXT); + WaitFlag(KDA_POST_EVENT_FIX); + WaitFlag(KDA_POST_EVENT_NEXT); + tileMmad(tileL0C, tileL0A, tileL0B, m, n, k, true, 0); + SetFlag(KDA_POST_EVENT); + WaitFlag(KDA_POST_EVENT); + copyL0CToDst(blockUOut, tileL0C); + SetFlag(KDA_POST_EVENT_FIX); + WaitFlag(KDA_POST_EVENT_FIX); + } + + __aicore__ inline void FinalizePostWuPipelineEvents(uint16_t usedSlotCount) + { + for (uint16_t slot = 0; slot < usedSlotCount; ++slot) { + WaitFlag(KDA_POST_EVENT + slot); + WaitFlag(KDA_POST_PIPELINE_U_EVENT + slot); + WaitFlag(KDA_POST_EVENT + slot); + WaitFlag(KDA_POST_EVENT + slot); + } + } + + __aicore__ inline void InitializePostWuPipelineEvents() + { + for (uint16_t slot = 0; slot < KDA_POST_PIPELINE_STAGE_COUNT; ++slot) { + InitializePostWuPipelineSlot(slot); + } + } + + __aicore__ inline void InitializePostWuPipelineSlot(uint16_t slot) + { + SetFlag(KDA_POST_EVENT + slot); + SetFlag(KDA_POST_EVENT + slot); + } + + __aicore__ inline void PrefetchPostWuPipelineArch35(Catlass::Arch::Resource &resource, uint16_t slot, + uint64_t b, uint64_t hv, uint64_t start, uint64_t curT, + bool reuseSlot) + { + using LayoutTagA = Catlass::layout::RowMajor; + using LayoutTagB = Catlass::layout::RowMajor; + using LayoutTagC = Catlass::layout::RowMajor; + using TileCopy = + Catlass::Gemm::Tile::PackedTileCopyTla; + using LayoutTagL1A = typename TileCopy::LayoutTagL1A; + using LayoutTagL1B = typename TileCopy::LayoutTagL1B; + + constexpr uint32_t capacityM = 64; + constexpr uint32_t n = 128; + constexpr uint32_t capacityK = 64; + const uint32_t m = static_cast(curT); + const uint32_t k = static_cast(curT); + auto layoutA = tla::MakeLayout(capacityM, capacityK); + auto layoutB = tla::MakeLayout(capacityK, n); + auto tensorA = tla::MakeTensor(preparedAqk_[AOffset(b, hv, start, 0)], layoutA, Catlass::Arch::PositionGM{}); + auto tensorW = + tla::MakeTensor(preparedQG_[KVOffset(b, hv, start, 0, K_)], layoutB, Catlass::Arch::PositionGM{}); + auto blockA = GetTile(tensorA, tla::MakeCoord(0, 0), tla::MakeShape(m, k)); + auto blockW = GetTile(tensorW, tla::MakeCoord(0, 0), tla::MakeShape(k, n)); + using CopyGmToL1A = typename TileCopy::template CopyGmToL1A; + using CopyGmToL1B = typename TileCopy::template CopyGmToL1B; + CopyGmToL1A copyGmToL1A; + CopyGmToL1B copyGmToL1B; + + uint32_t slotBase = static_cast(slot) * KDA_POST_PIPELINE_L1_SLOT_BYTES; + LocalTensor l1A = resource.l1Buf.template GetBufferByByte(slotBase); + LocalTensor l1W = resource.l1Buf.template GetBufferByByte(slotBase + KDA_POST_PIPELINE_L1_A_BYTES); + auto layoutL1A = tla::MakeLayout(capacityM, capacityK); + auto layoutL1B = tla::MakeLayout(capacityK, n); + auto tensorL1A = tla::MakeTensor(l1A, layoutL1A, Catlass::Arch::PositionL1{}); + auto tensorL1W = tla::MakeTensor(l1W, layoutL1B, Catlass::Arch::PositionL1{}); + + uint16_t pipelineEvent = KDA_POST_EVENT + slot; + if (reuseSlot) { + WaitFlag(pipelineEvent); + } + copyGmToL1A(tensorL1A, blockA); + copyGmToL1B(tensorL1W, blockW); + SetFlag(pipelineEvent); + } + + __aicore__ inline void PrefetchPostWuPipelineU(Catlass::Arch::Resource &resource, uint16_t slot, + uint64_t b, uint64_t hv, uint64_t start, uint64_t curT, + bool reuseStage) + { + using LayoutTagB = Catlass::layout::RowMajor; + using TileCopy = Catlass::Gemm::Tile::PackedTileCopyTla; + using LayoutTagL1B = typename TileCopy::LayoutTagL1B; + + constexpr uint32_t capacityK = 64; + constexpr uint32_t n = 128; + const uint32_t k = static_cast(curT); + auto layoutB = tla::MakeLayout(capacityK, n); + auto tensorU = + tla::MakeTensor(propagatedVNew_[KVOffset(b, hv, start, 0, V_)], layoutB, Catlass::Arch::PositionGM{}); + auto blockU = GetTile(tensorU, tla::MakeCoord(0, 0), tla::MakeShape(k, n)); + using CopyGmToL1B = typename TileCopy::template CopyGmToL1B; + CopyGmToL1B copyGmToL1B; + + uint32_t uOffset = KDA_POST_PIPELINE_STAGE_COUNT * KDA_POST_PIPELINE_L1_SLOT_BYTES + + static_cast(slot) * KDA_POST_PIPELINE_L1_U_SLOT_BYTES; + LocalTensor l1U = resource.l1Buf.template GetBufferByByte(uOffset); + auto layoutL1B = tla::MakeLayout(capacityK, n); + auto tensorL1U = tla::MakeTensor(l1U, layoutL1B, Catlass::Arch::PositionL1{}); + + uint16_t pipelineEvent = KDA_POST_PIPELINE_U_EVENT + slot; + if (reuseStage) { + WaitFlag(pipelineEvent); + } + copyGmToL1B(tensorL1U, blockU); + SetFlag(pipelineEvent); + } + + __aicore__ inline void ComputePrefetchedPostWuPipelineArch35(Catlass::Arch::Resource &resource, + uint16_t slot, uint64_t b, uint64_t hv, uint64_t start, + uint64_t curT) + { + using LayoutTagA = Catlass::layout::RowMajor; + using LayoutTagB = Catlass::layout::RowMajor; + using LayoutTagC = Catlass::layout::RowMajor; + using TileCopy = + Catlass::Gemm::Tile::PackedTileCopyTla; + using LayoutTagL1A = typename TileCopy::LayoutTagL1A; + using LayoutTagL1B = typename TileCopy::LayoutTagL1B; + using LayoutTagL0A = typename TileCopy::LayoutTagL0A; + using LayoutTagL0B = typename TileCopy::LayoutTagL0B; + using CopyL1ToL0A = typename TileCopy::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopy::CopyL1ToL0B; + using TileMmad = Catlass::Gemm::Tile::TileMmadTla; + + constexpr uint32_t capacityM = 64; + constexpr uint32_t n = 128; + constexpr uint32_t packedN = 256; + constexpr uint32_t capacityK = 64; + const uint32_t m = static_cast(curT); + const uint32_t k = static_cast(curT); + auto layoutC = tla::MakeLayout(capacityM, n); + auto tensorWOut = tla::MakeTensor(w_[KVOffset(b, hv, start, 0, K_)], layoutC, Catlass::Arch::PositionGM{}); + auto tensorUOut = tla::MakeTensor(u_[KVOffset(b, hv, start, 0, V_)], layoutC, Catlass::Arch::PositionGM{}); + auto blockWOut = GetTile(tensorWOut, tla::MakeCoord(0, 0), tla::MakeShape(m, n)); + auto blockUOut = GetTile(tensorUOut, tla::MakeCoord(0, 0), tla::MakeShape(m, n)); + using CopyL0CToDst = typename TileCopy::template CopyL0CToDst; + + uint32_t l1Base = static_cast(slot) * KDA_POST_PIPELINE_L1_SLOT_BYTES; + LocalTensor l1A = resource.l1Buf.template GetBufferByByte(l1Base); + LocalTensor l1W = resource.l1Buf.template GetBufferByByte(l1Base + KDA_POST_PIPELINE_L1_A_BYTES); + uint32_t uOffset = KDA_POST_PIPELINE_STAGE_COUNT * KDA_POST_PIPELINE_L1_SLOT_BYTES + + static_cast(slot) * KDA_POST_PIPELINE_L1_U_SLOT_BYTES; + LocalTensor l1U = resource.l1Buf.template GetBufferByByte(uOffset); + LocalTensor l0A = resource.l0ABuf.template GetBufferByByte(static_cast(slot) * + KDA_POST_PIPELINE_L0_A_SLOT_BYTES); + LocalTensor l0B = resource.l0BBuf.template GetBufferByByte(static_cast(slot) * + KDA_POST_PIPELINE_L0_B_SLOT_BYTES); + LocalTensor l0C = resource.l0CBuf.template GetBufferByByte(static_cast(slot) * + KDA_POST_PIPELINE_L0_C_SLOT_BYTES); + + auto layoutL1A = tla::MakeLayout(capacityM, capacityK); + auto layoutL1B = tla::MakeLayout(capacityK, n); + auto layoutL0A = tla::MakeLayout(capacityM, capacityK); + auto layoutL0B = tla::MakeLayout(capacityK, packedN); + auto layoutL0C = tla::MakeLayoutL0C(capacityM, packedN); + auto tensorL1A = tla::MakeTensor(l1A, layoutL1A, Catlass::Arch::PositionL1{}); + auto tensorL1W = tla::MakeTensor(l1W, layoutL1B, Catlass::Arch::PositionL1{}); + auto tensorL1U = tla::MakeTensor(l1U, layoutL1B, Catlass::Arch::PositionL1{}); + auto tensorL0A = tla::MakeTensor(l0A, layoutL0A, Catlass::Arch::PositionL0A{}); + auto tensorL0B = tla::MakeTensor(l0B, layoutL0B, Catlass::Arch::PositionL0B{}); + auto tensorL0C = tla::MakeTensor(l0C, layoutL0C, Catlass::Arch::PositionL0C{}); + auto tileL1A = GetTile(tensorL1A, tla::MakeCoord(0, 0), tla::MakeShape(m, k)); + auto tileL1W = GetTile(tensorL1W, tla::MakeCoord(0, 0), tla::MakeShape(k, n)); + auto tileL1U = GetTile(tensorL1U, tla::MakeCoord(0, 0), tla::MakeShape(k, n)); + auto tileL0A = GetTile(tensorL0A, tla::MakeCoord(0, 0), tla::MakeShape(m, k)); + auto tileL0BW = GetTile(tensorL0B, tla::MakeCoord(0, 0), tla::MakeShape(k, n)); + auto tileL0BU = GetTile(tensorL0B, tla::MakeCoord(0, n), tla::MakeShape(k, n)); + auto tileL0B = GetTile(tensorL0B, tla::MakeCoord(0, 0), tla::MakeShape(k, packedN)); + auto tileL0C = GetTile(tensorL0C, tla::MakeCoord(0, 0), tla::MakeShape(m, packedN)); + auto tileL0CW = GetTile(tensorL0C, tla::MakeCoord(0, 0), tla::MakeShape(m, n)); + auto tileL0CU = GetTile(tensorL0C, tla::MakeCoord(0, n), tla::MakeShape(m, n)); + + CopyL1ToL0A copyL1ToL0A; + CopyL1ToL0B copyL1ToL0B; + CopyL0CToDst copyL0CToDst; + TileMmad tileMmad; + + uint16_t pipelineEvent = KDA_POST_EVENT + slot; + uint16_t uPipelineEvent = KDA_POST_PIPELINE_U_EVENT + slot; + WaitFlag(pipelineEvent); + WaitFlag(uPipelineEvent); + WaitFlag(pipelineEvent); + copyL1ToL0A(tileL0A, tileL1A); + copyL1ToL0B(tileL0BW, tileL1W); + copyL1ToL0B(tileL0BU, tileL1U); + SetFlag(pipelineEvent); + SetFlag(pipelineEvent); + SetFlag(uPipelineEvent); + WaitFlag(pipelineEvent); + WaitFlag(pipelineEvent); + tileMmad(tileL0C, tileL0A, tileL0B, m, packedN, k, true, 0); + SetFlag(pipelineEvent); + SetFlag(pipelineEvent); + WaitFlag(pipelineEvent); + copyL0CToDst(blockWOut, tileL0CW); + copyL0CToDst(blockUOut, tileL0CU); + SetFlag(pipelineEvent); + } +#endif + + __aicore__ inline void ComputePostWuCube(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, uint64_t curT) + { +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if (UsePostWuCubeArch35(curT)) { + ComputePostWuCubeFusedArch35(b, hv, start, curT); + return; + } +#endif + using ElementA = AKK_T; + using ElementB = T; + using LayoutTagA = Catlass::layout::RowMajor; + using LayoutTagB = Catlass::layout::RowMajor; + using LayoutTagC = Catlass::layout::RowMajor; +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + SetLoadDataPaddingValue(static_cast(0)); + using WTileCopy = Catlass::Gemm::Tile::PackedTileCopyTla; +#else + using WTileCopy = Catlass::Gemm::Tile::PackedTileCopyTla; +#endif + using UTileCopy = Catlass::Gemm::Tile::PackedTileCopyTla; + using PostL1TileShape128 = tla::Shape; + using PostL0TileShape128 = tla::Shape; + using PostL1TileShape256 = tla::Shape; + using PostL0TileShape256 = tla::Shape; +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + using WBlockMmad = Catlass::Gemm::Block::BlockMmadTla; +#else + using WBlockMmad = Catlass::Gemm::Block::BlockMmadTla; +#endif + using UBlockMmad128 = + Catlass::Gemm::Block::BlockMmadTla; + using UBlockMmad256 = + Catlass::Gemm::Block::BlockMmadTla; + LayoutTagA tagA = LayoutTagA::template MakeLayout(BT_, BT_); + auto layoutA = tla::MakeLayoutFromTag(tagA); + auto tensorA = tla::MakeTensor(preparedAqk_[AOffset(b, hv, start, 0)], layoutA, Catlass::Arch::PositionGM{}); + + { + LayoutTagB tagB = LayoutTagB::template MakeLayout(BT_, K_); +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + LayoutTagC tagC = LayoutTagC::template MakeLayout(BT_, K_); +#else + LayoutTagC tagC = LayoutTagC::template MakeLayout(BT_, K_); +#endif + auto layoutB = tla::MakeLayoutFromTag(tagB); + auto layoutC = tla::MakeLayoutFromTag(tagC); + Catlass::GemmCoord shape{static_cast(curT), static_cast(K_), + static_cast(curT)}; + auto tensorB = + tla::MakeTensor(preparedQG_[KVOffset(b, hv, start, 0, K_)], layoutB, Catlass::Arch::PositionGM{}); +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + auto tensorC = tla::MakeTensor(w_[KVOffset(b, hv, start, 0, K_)], layoutC, Catlass::Arch::PositionGM{}); +#else + auto tensorC = + tla::MakeTensor(h_[WScratchOffset(b, hv, chunkIdx, 0, 0)], layoutC, Catlass::Arch::PositionGM{}); +#endif + auto blockA = GetTile(tensorA, tla::MakeCoord(0, 0), tla::MakeShape(shape.m(), shape.k())); + auto blockB = GetTile(tensorB, tla::MakeCoord(0, 0), tla::MakeShape(shape.k(), shape.n())); + auto blockC = GetTile(tensorC, tla::MakeCoord(0, 0), tla::MakeShape(shape.m(), shape.n())); + Catlass::Arch::Resource wResource; + WBlockMmad wBlockMmad(wResource); + wBlockMmad(blockA, blockB, blockC, shape); + PipeBarrier(); + } + + { + LayoutTagB tagB = LayoutTagB::template MakeLayout(BT_, V_); + LayoutTagC tagC = LayoutTagC::template MakeLayout(BT_, V_); + auto layoutB = tla::MakeLayoutFromTag(tagB); + auto layoutC = tla::MakeLayoutFromTag(tagC); + Catlass::GemmCoord shape{static_cast(curT), static_cast(V_), + static_cast(curT)}; + auto tensorB = + tla::MakeTensor(propagatedVNew_[KVOffset(b, hv, start, 0, V_)], layoutB, Catlass::Arch::PositionGM{}); + auto tensorC = tla::MakeTensor(u_[KVOffset(b, hv, start, 0, V_)], layoutC, Catlass::Arch::PositionGM{}); + auto blockA = GetTile(tensorA, tla::MakeCoord(0, 0), tla::MakeShape(shape.m(), shape.k())); + auto blockB = GetTile(tensorB, tla::MakeCoord(0, 0), tla::MakeShape(shape.k(), shape.n())); + auto blockC = GetTile(tensorC, tla::MakeCoord(0, 0), tla::MakeShape(shape.m(), shape.n())); + Catlass::Arch::Resource uResource; + if (V_ <= 128) { + UBlockMmad128 uBlockMmad(uResource); + uBlockMmad(blockA, blockB, blockC, shape); + } else { + UBlockMmad256 uBlockMmad(uResource); + uBlockMmad(blockA, blockB, blockC, shape); + } + PipeBarrier(); + } + } + +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + __aicore__ inline bool UseTypicalPostWuGate(uint64_t curT) const + { + return curT == 64 && BT_ == 64 && K_ == 128 && V_ == 128; + } + + __aicore__ inline uint64_t TypicalGateStageElems() const + { + return static_cast(KDA_TYPICAL_GATE_TILE_ROWS) * 128; + } + + __aicore__ inline uint64_t TypicalGateStageBytes() const + { + return TypicalGateStageElems() * (sizeof(T) + sizeof(GK_T)); + } + + __aicore__ inline LocalTensor TypicalGateK(uint64_t slot) + { + return gateWritebackBuf_.Get()[slot * TypicalGateStageBytes() / sizeof(T)]; + } + + __aicore__ inline LocalTensor TypicalGateG(uint64_t slot) + { + uint64_t byteOffset = slot * TypicalGateStageBytes() + TypicalGateStageElems() * sizeof(T); + return gateWritebackBuf_.Get()[byteOffset / sizeof(GK_T)]; + } + + __aicore__ inline void PrefetchTypicalKg(uint64_t slot, uint64_t b, uint64_t h, uint64_t hv, uint64_t token, + uint64_t rows) + { + uint64_t elems = rows * K_; + LocalTensor kStage = TypicalGateK(slot); + LocalTensor gateStage = TypicalGateG(slot); + CopyRowsIn(kStage, k_, QOffset(b, h, token, 0), rows, K_, inputSequenceMajor_ ? H_ * K_ : K_); + DataCopy(gateStage, gk_[KVOffset(b, hv, token, 0, K_)], static_cast(elems)); + SetFlag(mte2ToVEvent_); + } + + __aicore__ inline void ComputeTypicalKg(uint64_t b, uint64_t h, uint64_t hv, uint64_t start, uint64_t curT, + uint64_t subBlockIdx, uint64_t subBlockNum) + { + uint64_t rowBegin = (curT * subBlockIdx) / subBlockNum; + uint64_t rowEnd = (curT * (subBlockIdx + 1)) / subBlockNum; + if (rowBegin >= rowEnd) { + return; + } + + LocalTensor gateLast = exp2Buf_.Get(); + LoadAsFloatRow(gk_, KVOffset(b, hv, start + curT - 1, 0, K_), gateLast, K_); + + uint64_t slot = 0; + uint64_t firstRows = rowEnd - rowBegin; + if (firstRows > KDA_TYPICAL_GATE_TILE_ROWS) { + firstRows = KDA_TYPICAL_GATE_TILE_ROWS; + } + PrefetchTypicalKg(slot, b, h, hv, start + rowBegin, firstRows); + + bool outputPending = false; + for (uint64_t tileRow = rowBegin; tileRow < rowEnd; tileRow += KDA_TYPICAL_GATE_TILE_ROWS) { + uint64_t tileRows = rowEnd - tileRow; + if (tileRows > KDA_TYPICAL_GATE_TILE_ROWS) { + tileRows = KDA_TYPICAL_GATE_TILE_ROWS; + } + uint64_t elems = tileRows * K_; + WaitFlag(mte2ToVEvent_); + + if (outputPending) { + WaitFlag(mte3ToMte2Event_); + } + uint64_t nextTileRow = tileRow + KDA_TYPICAL_GATE_TILE_ROWS; + if (nextTileRow < rowEnd) { + uint64_t nextRows = rowEnd - nextTileRow; + if (nextRows > KDA_TYPICAL_GATE_TILE_ROWS) { + nextRows = KDA_TYPICAL_GATE_TILE_ROWS; + } + PrefetchTypicalKg(slot ^ 1, b, h, hv, start + nextTileRow, nextRows); + } + + LocalTensor kAndKg = TypicalGateK(slot); + LocalTensor gateStage = TypicalGateG(slot); + ComputePostKdaKgRegbase((__ubuf__ T *)reinterpret_cast(kAndKg.GetPhyAddr()), + (__ubuf__ GK_T *)reinterpret_cast(gateStage.GetPhyAddr()), + (__ubuf__ float *)reinterpret_cast(gateLast.GetPhyAddr()), + static_cast(tileRows), static_cast(K_)); + + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + DataCopy(kg_[KVOffset(b, hv, start + tileRow, 0, K_)], kAndKg, static_cast(elems)); + SetFlag(mte3ToMte2Event_); + outputPending = true; + slot ^= 1; + } + if (outputPending) { + WaitFlag(mte3ToMte2Event_); + } + } + + __aicore__ inline uint64_t TypicalGatePipelineStageElems() const + { + return static_cast(KDA_TYPICAL_GATE_PIPELINE_ROWS) * 128; + } + + __aicore__ inline uint64_t TypicalGatePipelineStageBytes() const + { + return TypicalGatePipelineStageElems() * (sizeof(T) + sizeof(float)) + 128 * sizeof(float); + } + + __aicore__ inline LocalTensor TypicalGatePipelineK(uint64_t slot) + { + uint64_t byteOffset = slot * TypicalGatePipelineStageBytes(); + return vecBuf_.Get()[byteOffset / sizeof(T)]; + } + + __aicore__ inline LocalTensor TypicalGatePipelineG(uint64_t slot) + { + uint64_t byteOffset = slot * TypicalGatePipelineStageBytes() + TypicalGatePipelineStageElems() * sizeof(T); + return vecBuf_.Get()[byteOffset / sizeof(float)]; + } + + __aicore__ inline LocalTensor TypicalGatePipelineRef(uint64_t slot) + { + uint64_t byteOffset = + slot * TypicalGatePipelineStageBytes() + TypicalGatePipelineStageElems() * (sizeof(T) + sizeof(float)); + return vecBuf_.Get()[byteOffset / sizeof(float)]; + } + + __aicore__ inline bool CanPipelineTypicalKg(uint64_t curT, uint64_t subBlockIdx, uint64_t subBlockNum) const + { + if constexpr (!IsSameType::value) { + return false; + } + if (!UseTypicalPostWuGate(curT) || subBlockNum == 0) { + return false; + } + uint64_t rowBegin = (curT * subBlockIdx) / subBlockNum; + uint64_t rowEnd = (curT * (subBlockIdx + 1)) / subBlockNum; + return rowBegin < rowEnd && rowEnd - rowBegin <= KDA_TYPICAL_GATE_PIPELINE_ROWS; + } + + __aicore__ inline void PrefetchTypicalKgPipeline(uint64_t slot, uint64_t b, uint64_t h, uint64_t hv, uint64_t start, + uint64_t curT, uint64_t rowBegin, uint64_t rowEnd) + { + if constexpr (IsSameType::value) { + uint64_t elems = (rowEnd - rowBegin) * K_; + LocalTensor kStage = TypicalGatePipelineK(slot); + LocalTensor gateStage = TypicalGatePipelineG(slot); + LocalTensor refStage = TypicalGatePipelineRef(slot); + CopyRowsIn(kStage, k_, QOffset(b, h, start + rowBegin, 0), rowEnd - rowBegin, K_, + inputSequenceMajor_ ? H_ * K_ : K_); + DataCopy(gateStage, gk_[KVOffset(b, hv, start + rowBegin, 0, K_)], static_cast(elems)); + DataCopy(refStage, gk_[KVOffset(b, hv, start + curT - 1, 0, K_)], static_cast(K_)); + SetFlag(mte2ToVEvent_); + } + } + + __aicore__ inline void ComputeTypicalKgPipelineRegs(uint64_t slot, uint64_t rowBegin, uint64_t rowEnd) + { + uint64_t rows = rowEnd - rowBegin; + LocalTensor kAndKg = TypicalGatePipelineK(slot); + LocalTensor gateStage = TypicalGatePipelineG(slot); + LocalTensor refStage = TypicalGatePipelineRef(slot); + ComputePostKdaKgRegbase((__ubuf__ T *)reinterpret_cast(kAndKg.GetPhyAddr()), + (__ubuf__ float *)reinterpret_cast(gateStage.GetPhyAddr()), + (__ubuf__ float *)reinterpret_cast(refStage.GetPhyAddr()), + static_cast(rows), static_cast(K_)); + } + + __aicore__ inline void StoreTypicalKgPipeline(uint64_t slot, uint64_t b, uint64_t hv, uint64_t start, + uint64_t rowBegin, uint64_t rowEnd) + { + uint64_t elems = (rowEnd - rowBegin) * K_; + LocalTensor kAndKg = TypicalGatePipelineK(slot); + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + DataCopy(kg_[KVOffset(b, hv, start + rowBegin, 0, K_)], kAndKg, static_cast(elems)); + SetFlag(mte3ToMte2Event_); + } +#endif + + __aicore__ inline void CopyScratchWAndFinalizeKg(uint64_t b, uint64_t h, uint64_t hv, uint64_t chunkIdx, + uint64_t start, uint64_t curT, uint64_t subBlockIdx, + uint64_t subBlockNum) + { + constexpr uint64_t typedOffsetFloats = 20480; + constexpr uint64_t typedOffset = typedOffsetFloats * sizeof(float) / sizeof(T); + constexpr uint64_t kgFp32Planes = 4; + uint64_t rowBegin = (curT * subBlockIdx) / subBlockNum; + uint64_t rowEnd = (curT * (subBlockIdx + 1)) / subBlockNum; + if (rowBegin >= rowEnd) { + return; + } + uint64_t maxRows = (typedOffsetFloats / kgFp32Planes) / K_; + if (maxRows > 32) { + maxRows = 32; + } + if (maxRows == 0) { + return; + } + + uint64_t last = start + curT - 1; + LocalTensor arena = vecBuf_.Get(); + LocalTensor gateLast = exp2Buf_.Get(); + LocalTensor typedLocal = vecBuf_.Get()[typedOffset]; + LoadAsFloatRow(gk_, KVOffset(b, hv, last, 0, K_), gateLast, K_); + + for (uint64_t tileRow = rowBegin; tileRow < rowEnd; tileRow += maxRows) { + uint64_t tileRows = rowEnd - tileRow; + if (tileRows > maxRows) { + tileRows = maxRows; + } + uint64_t elemCount = tileRows * K_; +#if !defined(__CCE_AICORE__) || __CCE_AICORE__ != 310 + uint64_t scratchBase = WScratchOffset(b, hv, chunkIdx, tileRow, 0); +#else + (void)chunkIdx; +#endif + uint64_t token = start + tileRow; + +#if !defined(__CCE_AICORE__) || __CCE_AICORE__ != 310 + DataCopy(arena, h_[scratchBase], static_cast(elemCount)); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + Cast(typedLocal, arena, RoundMode::CAST_RINT, static_cast(elemCount)); + PipeBarrier(); + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + DataCopy(w_[KVOffset(b, hv, token, 0, K_)], typedLocal, static_cast(elemCount)); + SetFlag(mte3ToMte2Event_); + WaitFlag(mte3ToMte2Event_); +#endif + + LocalTensor kLocal = arena; + LocalTensor gLocal = arena[elemCount]; + LocalTensor expLocal = arena[2 * elemCount]; + LocalTensor outLocal = arena[3 * elemCount]; + const uint64_t gateOffsetBytes = (typedOffset + elemCount) * sizeof(T); + LocalTensor gateTyped = vecBuf_.Get()[(gateOffsetBytes + sizeof(GK_T) - 1) / sizeof(GK_T)]; + CopyRowsIn(typedLocal, k_, QOffset(b, h, token, 0), tileRows, K_, inputSequenceMajor_ ? H_ * K_ : K_); + LoadAsFloatVector(gk_, KVOffset(b, hv, token, 0, K_), gLocal, gateTyped, elemCount); + Cast(kLocal, typedLocal, RoundMode::CAST_NONE, static_cast(elemCount)); + PipeBarrier(); + + for (uint64_t row = 0; row < tileRows; ++row) { + Sub(expLocal[row * K_], gateLast, gLocal[row * K_], static_cast(K_)); + } + PipeBarrier(); + Muls(expLocal, expLocal, LN2, static_cast(elemCount)); + PipeBarrier(); + ClampExpInput(expLocal, static_cast(elemCount)); + Exp(expLocal, expLocal, static_cast(elemCount)); + PipeBarrier(); + Mul(outLocal, kLocal, expLocal, static_cast(elemCount)); + PipeBarrier(); + ClampFp32ToOutputType(outLocal, static_cast(elemCount)); + Cast(typedLocal, outLocal, RoundMode::CAST_RINT, static_cast(elemCount)); + PipeBarrier(); + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + CopyVectorOut(kg_, KVOffset(b, hv, token, 0, K_), typedLocal, elemCount); + SetFlag(mte3ToMte2Event_); + WaitFlag(mte3ToMte2Event_); + } + if (rowEnd == curT) { + CopyVectorIn(typedLocal, k_, QOffset(b, h, last, 0), K_); + SetFlag(mte2ToMte3Event_); + WaitFlag(mte2ToMte3Event_); + CopyVectorOut(kg_, KVOffset(b, hv, last, 0, K_), typedLocal, K_); + SetFlag(mte3ToMte2Event_); + WaitFlag(mte3ToMte2Event_); + } + SetFlag(mte3ToVEvent_); + WaitFlag(mte3ToVEvent_); + } + +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + template + __aicore__ inline void ComputeTailWuRow(GlobalTensor &src, GlobalTensor &dst, + uint64_t akkBase, uint64_t srcBase, uint64_t dstBase, uint64_t curT, + uint64_t dim) + { + LocalTensor acc = vecBuf_.Get(); + LocalTensor value = vecBuf_.Get()[512]; + LocalTensor typed = vecBuf_.Get()[4096]; + LocalTensor coefficientTyped = exp2Buf_.Get(); + LocalTensor coefficients = exp2Buf_.Get()[128]; + CopyVectorIn(coefficientTyped, preparedAqk_, akkBase, curT); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + Cast(coefficients, coefficientTyped, RoundMode::CAST_NONE, static_cast(curT)); + PipeBarrier(); + SetFlag(EXP2_EVENT_ID); + WaitFlag(EXP2_EVENT_ID); + for (uint64_t j = 0; j < curT; ++j) { + LoadAsFloatVector(src, srcBase + j * dim, value, typed, dim); + float coefficient = coefficients.GetValue(j); + SetFlag(EXP2_EVENT_ID); + WaitFlag(EXP2_EVENT_ID); + Muls(value, value, coefficient, static_cast(dim)); + PipeBarrier(); + if (j == 0) { + Adds(acc, value, 0.0f, static_cast(dim)); + } else { + Add(acc, acc, value, static_cast(dim)); + } + PipeBarrier(); + SetFlag(vToMte2Event_); + WaitFlag(vToMte2Event_); + } + SetFlag(EXP2_EVENT_ID); + WaitFlag(EXP2_EVENT_ID); + ClampFp32ToOutputType(acc, static_cast(dim)); + StoreFloatRow(dst, dstBase, acc, dim); + } + + __aicore__ inline void ComputeTailWuVector(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, + uint64_t curT, uint64_t subBlockIdx, uint64_t subBlockNum) + { + uint64_t rowBegin = (curT * subBlockIdx) / subBlockNum; + uint64_t rowEnd = (curT * (subBlockIdx + 1)) / subBlockNum; + for (uint64_t row = rowBegin; row < rowEnd; ++row) { + ComputeTailWuRow(preparedQG_, w_, AOffset(b, hv, start + row, 0), KVOffset(b, hv, start, 0, K_), + KVOffset(b, hv, start + row, 0, K_), curT, K_); + ComputeTailWuRow(propagatedVNew_, u_, AOffset(b, hv, start + row, 0), KVOffset(b, hv, start, 0, V_), + KVOffset(b, hv, start + row, 0, V_), curT, V_); + } + } + + __aicore__ inline void ComputeTailWuRegbaseArch35(uint64_t b, uint64_t hv, uint64_t start, uint64_t curT, + uint64_t subBlockIdx, uint64_t subBlockNum) + { + // W aliases W-seed in GM: one AIV must stage every seed row before either output is overwritten. + constexpr uint64_t dim = 128; + constexpr uint64_t maxTailRows = KDA_POST_REGBASE_TAIL_LIMIT - 1; + constexpr uint64_t tileRowsLimit = 32; + constexpr uint64_t akkStagedCols = 64; + constexpr uint64_t wSeedOffset = 0; + constexpr uint64_t uSeedOffset = wSeedOffset + maxTailRows * dim; + constexpr uint64_t akkOffset = uSeedOffset + maxTailRows * dim; + (void)subBlockNum; + if (subBlockIdx != 0 || curT == 0 || curT > maxTailRows) { + return; + } + + LocalTensor staging = gateWritebackBuf_.Get(); + LocalTensor wSeedLocal = staging[wSeedOffset]; + LocalTensor uSeedLocal = staging[uSeedOffset]; + LocalTensor akkLocal = staging[akkOffset]; + CopyVectorIn(wSeedLocal, preparedQG_, KVOffset(b, hv, start, 0, K_), curT * dim); + CopyVectorIn(uSeedLocal, propagatedVNew_, KVOffset(b, hv, start, 0, V_), curT * dim); + for (uint64_t row = 0; row < curT; ++row) { + LocalTensor akkRow = akkLocal[row * akkStagedCols]; + CopyVectorIn(akkRow, preparedAqk_, AOffset(b, hv, start + row, 0), akkStagedCols); + } + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + + for (uint64_t tileRow = 0; tileRow < curT; tileRow += tileRowsLimit) { + const uint64_t tileRows = (curT - tileRow) > tileRowsLimit ? tileRowsLimit : (curT - tileRow); + LocalTensor wFloat = vecBuf_.Get(); + LocalTensor uFloat = wFloat[tileRows * dim]; + LocalTensor akkTile = akkLocal[tileRow * akkStagedCols]; + AscendC::VF_CALL(reinterpret_cast<__ubuf__ float *>(wFloat.GetPhyAddr()), + reinterpret_cast<__ubuf__ float *>(uFloat.GetPhyAddr()), + reinterpret_cast<__ubuf__ bfloat16_t *>(akkTile.GetPhyAddr()), + reinterpret_cast<__ubuf__ bfloat16_t *>(wSeedLocal.GetPhyAddr()), + reinterpret_cast<__ubuf__ bfloat16_t *>(uSeedLocal.GetPhyAddr()), + static_cast(tileRows), static_cast(curT)); + PipeBarrier(); + + for (uint64_t row = 0; row < tileRows; ++row) { + LocalTensor wRow = wFloat[row * dim]; + LocalTensor uRow = uFloat[row * dim]; + StoreFloatRow(w_, KVOffset(b, hv, start + tileRow + row, 0, K_), wRow, dim); + StoreFloatRow(u_, KVOffset(b, hv, start + tileRow + row, 0, V_), uRow, dim); + } + } + } +#endif + __aicore__ inline bool ResolveFlatChunk(uint64_t task, uint64_t &seq, uint64_t &b, uint64_t &h, uint64_t &hv, + uint64_t &chunkIdx, uint64_t &start, uint64_t &end) + { + hv = task % HV_; + uint64_t flatChunk = task / HV_; + if (!isVarLen_) { + seq = flatChunk / NT_; + b = seq; + chunkIdx = flatChunk % NT_; + start = chunkIdx * BT_; + end = start + BT_; + if (end > T_) { + end = T_; + } + } else { + if (!KdaVarlen::ResolveChunkRange(cuSeqlensAddr_, chunkIndicesAddr_, N_, T_, BT_, flatChunk, seq, start, + end)) { + return false; + } + b = 0; + chunkIdx = flatChunk; + } + h = hv / (HV_ / H_); + return start < end; + } + +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + __aicore__ inline bool ResolveHeadMajorChunk(uint64_t task, uint64_t &seq, uint64_t &b, uint64_t &h, uint64_t &hv, + uint64_t &chunkIdx, uint64_t &start, uint64_t &end) + { + uint64_t flatTask = 0; + if (isVarLen_) { + hv = task / NT_; + uint64_t flatChunk = task % NT_; + flatTask = flatChunk * HV_ + hv; + } else { + uint64_t entity = task / NT_; + uint64_t localChunk = task % NT_; + b = entity / HV_; + hv = entity % HV_; + uint64_t flatChunk = b * NT_ + localChunk; + flatTask = flatChunk * HV_ + hv; + } + return ResolveFlatChunk(flatTask, seq, b, h, hv, chunkIdx, start, end); + } + + __aicore__ inline void GetHeadMajorTaskRange(uint64_t coreIdx, uint64_t coreNum, uint64_t taskNum, + uint64_t &taskBegin, uint64_t &taskEnd) const + { + uint64_t tasksPerCore = (taskNum + coreNum - 1) / coreNum; + taskBegin = coreIdx * tasksPerCore; + taskEnd = taskBegin + tasksPerCore; + if (taskBegin > taskNum) { + taskBegin = taskNum; + } + if (taskEnd > taskNum) { + taskEnd = taskNum; + } + } + + __aicore__ inline void ProcessPostAivPipelineArch35(uint64_t taskBegin, uint64_t taskEnd, uint64_t subBlockIdx, + uint64_t subBlockNum) + { + uint64_t task = taskBegin; + while (task < taskEnd) { + uint64_t seq = 0; + uint64_t b = 0; + uint64_t h = 0; + uint64_t hv = 0; + uint64_t chunkIdx = 0; + uint64_t start = 0; + uint64_t end = 0; + bool resolved = ResolveHeadMajorChunk(task, seq, b, h, hv, chunkIdx, start, end); + uint64_t curT = resolved ? end - start : 0; + if (!resolved || !CanPipelineTypicalKg(curT, subBlockIdx, subBlockNum)) { + if (resolved) { + (void)seq; + ProcessChunkPostAiv(b, h, hv, chunkIdx, start, end, subBlockIdx, subBlockNum); + } + ++task; + continue; + } + + uint64_t rowBegin = (curT * subBlockIdx) / subBlockNum; + uint64_t rowEnd = (curT * (subBlockIdx + 1)) / subBlockNum; + uint16_t slot = 0; + bool outputPending = false; + PrefetchTypicalKgPipeline(slot, b, h, hv, start, curT, rowBegin, rowEnd); + + while (true) { + WaitFlag(mte2ToVEvent_); + + uint64_t nextTask = task + 1; + uint64_t nextSeq = 0; + uint64_t nextB = 0; + uint64_t nextH = 0; + uint64_t nextHv = 0; + uint64_t nextChunkIdx = 0; + uint64_t nextStart = 0; + uint64_t nextEnd = 0; + bool nextResolved = nextTask < taskEnd && ResolveHeadMajorChunk(nextTask, nextSeq, nextB, nextH, nextHv, + nextChunkIdx, nextStart, nextEnd); + uint64_t nextCurT = nextResolved ? nextEnd - nextStart : 0; + bool nextIsTypical = nextResolved && CanPipelineTypicalKg(nextCurT, subBlockIdx, subBlockNum); + uint64_t nextRowBegin = 0; + uint64_t nextRowEnd = 0; + if (nextIsTypical) { + nextRowBegin = (nextCurT * subBlockIdx) / subBlockNum; + nextRowEnd = (nextCurT * (subBlockIdx + 1)) / subBlockNum; + uint16_t nextSlot = (slot + 1) % KDA_TYPICAL_GATE_PIPELINE_STAGES; + PrefetchTypicalKgPipeline(nextSlot, nextB, nextH, nextHv, nextStart, nextCurT, nextRowBegin, + nextRowEnd); + } + + ComputeTypicalKgPipelineRegs(slot, rowBegin, rowEnd); + if (outputPending) { + WaitFlag(mte3ToMte2Event_); + outputPending = false; + } + StoreTypicalKgPipeline(slot, b, hv, start, rowBegin, rowEnd); + outputPending = true; + + if (!nextIsTypical) { + WaitFlag(mte3ToMte2Event_); + task = nextTask; + break; + } + + task = nextTask; + seq = nextSeq; + b = nextB; + h = nextH; + hv = nextHv; + chunkIdx = nextChunkIdx; + start = nextStart; + end = nextEnd; + curT = nextCurT; + rowBegin = nextRowBegin; + rowEnd = nextRowEnd; + slot = (slot + 1) % KDA_TYPICAL_GATE_PIPELINE_STAGES; + (void)seq; + (void)chunkIdx; + (void)end; + (void)curT; + } + } + } +#endif + + __aicore__ inline void ProcessChunkPostAiv(uint64_t b, uint64_t h, uint64_t hv, uint64_t chunkIdx, uint64_t start, + uint64_t end, uint64_t subBlockIdx, uint64_t subBlockNum) + { + uint64_t curT = end - start; + if (curT == 0 || !UsePostWuCube(curT)) { + return; + } +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if (curT < BT_) { + if constexpr (IsSameType::value) { + if (BT_ == 64 && K_ == 128 && V_ == 128 && curT < KDA_POST_REGBASE_TAIL_LIMIT) { + ComputeTailWuRegbaseArch35(b, hv, start, curT, subBlockIdx, subBlockNum); + } else { + ComputeTailWuVector(b, hv, chunkIdx, start, curT, subBlockIdx, subBlockNum); + } + } else { + ComputeTailWuVector(b, hv, chunkIdx, start, curT, subBlockIdx, subBlockNum); + } + CopyScratchWAndFinalizeKg(b, h, hv, chunkIdx, start, curT, subBlockIdx, subBlockNum); + return; + } + if (UseTypicalPostWuGate(curT)) { + ComputeTypicalKg(b, h, hv, start, curT, subBlockIdx, subBlockNum); + return; + } else { + CopyScratchWAndFinalizeKg(b, h, hv, chunkIdx, start, curT, subBlockIdx, subBlockNum); + } +#else + Catlass::Arch::CrossCoreWaitFlagWithReverse<0x2, PIPE_MTE2>(syncDoneFlag_); + CopyScratchWAndFinalizeKg(b, h, hv, chunkIdx, start, curT, subBlockIdx, subBlockNum); +#endif + } + + __aicore__ inline void ProcessChunkPostAic(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, uint64_t end) + { + if constexpr (IsSameType::value) { + ProcessChunkPostAicTyped(b, hv, chunkIdx, start, end); + } + } + + __aicore__ inline void ProcessChunkPostAicTyped(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, + uint64_t end) + { + uint64_t curT = end - start; + if (curT == 0 || !UsePostWuCube(curT)) { + return; + } + ComputePostWuCube(b, hv, chunkIdx, start, curT); +#if !defined(__CCE_AICORE__) || __CCE_AICORE__ != 310 + Catlass::Arch::CrossCoreSetFlagWithReverse<0x2, PIPE_FIX>(syncDoneFlag_); +#endif + } + + __aicore__ inline void ProcessPostAiv() + { + if constexpr (IsSameType::value) { + return; + } + uint64_t subBlockNum = static_cast(GetSubBlockNum()); + if (subBlockNum == 0) { + return; + } + uint64_t subBlockIdx = static_cast(GetSubBlockIdx()); + uint64_t coreNum = usedCoreNum_ == 0 ? 1 : usedCoreNum_; + uint64_t coreIdx = static_cast(GetBlockIdx()) / subBlockNum; + uint64_t taskNum = static_cast((isVarLen_ ? NT_ : B_ * NT_) * HV_); +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if (BT_ == 64 && K_ == 128 && V_ == 128) { + uint64_t taskBegin = 0; + uint64_t taskEnd = 0; + GetHeadMajorTaskRange(coreIdx, coreNum, taskNum, taskBegin, taskEnd); + ProcessPostAivPipelineArch35(taskBegin, taskEnd, subBlockIdx, subBlockNum); + return; + } +#endif + for (uint64_t task = coreIdx; task < taskNum; task += coreNum) { + uint64_t seq = 0; + uint64_t b = 0; + uint64_t h = 0; + uint64_t hv = 0; + uint64_t chunkIdx = 0; + uint64_t start = 0; + uint64_t end = 0; + if (ResolveFlatChunk(task, seq, b, h, hv, chunkIdx, start, end)) { + (void)seq; + ProcessChunkPostAiv(b, h, hv, chunkIdx, start, end, subBlockIdx, subBlockNum); + } + } + } + + __aicore__ inline void ProcessPostAic() + { + if constexpr (IsSameType::value) { + return; + } +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if (KDA_ENABLE_POST_AIC_PIPELINE && BT_ == 64 && K_ == 128 && V_ == 128) { + ProcessPostAicPipelineArch35(); + return; + } +#endif + uint64_t taskNum = static_cast((isVarLen_ ? NT_ : B_ * NT_) * HV_); + uint64_t coreNum = usedCoreNum_ == 0 ? 1 : usedCoreNum_; + for (uint64_t task = GetBlockIdx(); task < taskNum; task += coreNum) { + uint64_t seq = 0; + uint64_t b = 0; + uint64_t h = 0; + uint64_t hv = 0; + uint64_t chunkIdx = 0; + uint64_t start = 0; + uint64_t end = 0; + if (ResolveFlatChunk(task, seq, b, h, hv, chunkIdx, start, end)) { + (void)seq; + (void)h; + ProcessChunkPostAic(b, hv, chunkIdx, start, end); + } + } + } + +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + __aicore__ inline void ProcessPostAicPipelineArch35() + { + static_assert(sizeof(T) == sizeof(uint16_t), "arch35 PostWU pipeline is specialized for fp16/bf16 inputs"); + SetLoadDataPaddingValue(static_cast(0)); + uint64_t taskNum = static_cast((isVarLen_ ? NT_ : B_ * NT_) * HV_); + uint64_t coreNum = usedCoreNum_ == 0 ? 1 : usedCoreNum_; + uint64_t coreIdx = static_cast(GetBlockIdx()); + uint64_t taskBegin = 0; + uint64_t taskEnd = 0; + GetHeadMajorTaskRange(coreIdx, coreNum, taskNum, taskBegin, taskEnd); + if (taskEnd - taskBegin < KDA_POST_PIPELINE_STAGE_COUNT) { + for (uint64_t task = taskBegin; task < taskEnd; ++task) { + uint64_t seq = 0; + uint64_t b = 0; + uint64_t h = 0; + uint64_t hv = 0; + uint64_t chunkIdx = 0; + uint64_t start = 0; + uint64_t end = 0; + if (ResolveHeadMajorChunk(task, seq, b, h, hv, chunkIdx, start, end)) { + (void)seq; + (void)h; + if (end - start == BT_) { + ProcessPreparedTailSingleArch35(b, hv, start, BT_); + } + } + } + return; + } + uint64_t task = taskBegin; + Catlass::Arch::Resource resource; + + while (task < taskEnd) { + uint64_t seq = 0; + uint64_t b = 0; + uint64_t h = 0; + uint64_t hv = 0; + uint64_t chunkIdx = 0; + uint64_t start = 0; + uint64_t end = 0; + bool resolved = ResolveHeadMajorChunk(task, seq, b, h, hv, chunkIdx, start, end); + uint64_t curT = resolved ? end - start : 0; + if (!resolved || !UseFullPostWuPipelineArch35(curT)) { + (void)seq; + (void)h; + ++task; + continue; + } + + uint16_t slot = 0; + uint16_t usedSlotCount = 1; + InitializePostWuPipelineSlot(slot); + PrefetchPostWuPipelineArch35(resource, slot, b, hv, start, curT, false); + PrefetchPostWuPipelineU(resource, slot, b, hv, start, curT, false); + while (true) { + uint64_t nextTask = task + 1; + uint64_t nextSeq = 0; + uint64_t nextB = 0; + uint64_t nextH = 0; + uint64_t nextHv = 0; + uint64_t nextChunkIdx = 0; + uint64_t nextStart = 0; + uint64_t nextEnd = 0; + bool nextResolved = nextTask < taskEnd && ResolveHeadMajorChunk(nextTask, nextSeq, nextB, nextH, nextHv, + nextChunkIdx, nextStart, nextEnd); + uint64_t nextCurT = nextResolved ? nextEnd - nextStart : 0; + bool nextIsTypical = nextResolved && UseFullPostWuPipelineArch35(nextCurT); + if (nextIsTypical) { + uint16_t nextSlot = slot ^ 1; + bool reuseSlot = usedSlotCount == KDA_POST_PIPELINE_STAGE_COUNT; + if (!reuseSlot) { + InitializePostWuPipelineSlot(nextSlot); + } + PrefetchPostWuPipelineArch35(resource, nextSlot, nextB, nextHv, nextStart, nextCurT, reuseSlot); + PrefetchPostWuPipelineU(resource, nextSlot, nextB, nextHv, nextStart, nextCurT, reuseSlot); + if (!reuseSlot) { + ++usedSlotCount; + } + } + + ComputePrefetchedPostWuPipelineArch35(resource, slot, b, hv, start, curT); + if (!nextIsTypical) { + FinalizePostWuPipelineEvents(usedSlotCount); + task = nextTask; + break; + } + task = nextTask; + seq = nextSeq; + b = nextB; + h = nextH; + hv = nextHv; + chunkIdx = nextChunkIdx; + start = nextStart; + end = nextEnd; + curT = nextCurT; + slot ^= 1; + (void)seq; + (void)h; + (void)chunkIdx; + (void)end; + (void)curT; + } + } + } +#endif + +private: + GlobalTensor q_; + GlobalTensor k_; + GlobalTensor v_; + GlobalTensor gk_; + GlobalTensor beta_; + GlobalTensor initialState_; + GlobalTensor o_; + GlobalTensor finalState_; + GlobalTensor aqk_; + GlobalTensor akk_; + GlobalTensor w_; + GlobalTensor u_; + GlobalTensor qg_; + GlobalTensor kg_; + GlobalTensor vNew_; + GlobalTensor h_; + GlobalTensor preparedQG_; + GlobalTensor preparedAqk_; + GlobalTensor propagatedVNew_; + GlobalTensor propagatedH_; + GlobalTensor solveWorkspace_; + GlobalTensor scoreWorkspace_; + TPipe *pipe_ = nullptr; + TBuf exp2Buf_; + TBuf vecBuf_; + TBuf gateWritebackBuf_; + TEventID mte2ToVEvent_ = 0; + TEventID vToMte2Event_ = 0; + TEventID vToMte3Event_ = 0; + TEventID mte3ToVEvent_ = 0; + TEventID mte2ToMte3Event_ = 0; + TEventID mte3ToMte2Event_ = 0; + bool vectorEventsAllocated_ = false; + Catlass::Arch::CrossCoreFlagWithReverse scoreReadyFlag_{KDA_SCORE_READY_FLAG0, + KDA_SCORE_READY_FLAG1}; + Catlass::Arch::CrossCoreFlagWithReverse scoreDoneFlag_{KDA_SCORE_DONE_FLAG0, + KDA_SCORE_DONE_FLAG1}; + // Score production is fully drained before solve starts, so the solve handshake can safely reuse + // the existing score flags without consuming additional hardware flag IDs. + Catlass::Arch::CrossCoreFlagWithReverse syncReadyFlag_{KDA_SCORE_READY_FLAG0, + KDA_SCORE_READY_FLAG1}; + Catlass::Arch::CrossCoreFlagWithReverse syncDoneFlag_{KDA_SCORE_DONE_FLAG0, + KDA_SCORE_DONE_FLAG1}; + uint64_t B_ = 0; + uint64_t N_ = 0; + uint64_t H_ = 0; + uint64_t HV_ = 0; + uint64_t T_ = 0; + uint64_t K_ = 0; + uint64_t V_ = 0; + uint64_t BT_ = 0; + uint64_t NT_ = 0; + float scale_ = 1.0f; + bool hasInitial_ = false; + bool isVarLen_ = false; + bool isAivOnly_ = false; + bool inputSequenceMajor_ = false; + uint64_t usedCoreNum_ = 1; + uint64_t solveCoreIdx_ = 0; + __gm__ int64_t *chunkIndicesAddr_ = nullptr; + __gm__ int64_t *cuSeqlensAddr_ = nullptr; +}; +} // namespace + +template +__aicore__ inline void RunChunkKdaPostWu(GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR gk, GM_ADDR beta, + GM_ADDR initialState, GM_ADDR cuSeqlens, GM_ADDR chunkIndices, GM_ADDR wSeed, + GM_ADDR akk, GM_ADDR uSeed, GM_ADDR w, GM_ADDR u, GM_ADDR kg, GM_ADDR vNew, + GM_ADDR userWorkspace, const TilingData &tiling, TPipe &pipe) +{ + GM_ADDR postScratch = userWorkspace + tiling.postWuScratchOffset; + if ASCEND_IS_AIC { + ChunkKdaFwdPostWuKernel op; + op.Init(q, k, v, gk, beta, initialState, cuSeqlens, chunkIndices, wSeed, akk, uSeed, nullptr, userWorkspace, + userWorkspace, userWorkspace, akk, w, u, userWorkspace, kg, vNew, postScratch, postScratch, tiling, + &pipe, false); + op.ProcessAic(); + } + if ASCEND_IS_AIV { + ChunkKdaFwdPostWuKernel op; + op.Init(q, k, v, gk, beta, initialState, cuSeqlens, chunkIndices, wSeed, akk, uSeed, nullptr, userWorkspace, + userWorkspace, userWorkspace, akk, w, u, userWorkspace, kg, vNew, postScratch, postScratch, tiling, + &pipe); + op.ProcessAiv(); + } +} + +} // namespace KdaPostWu diff --git a/csrc/chunk_kda_fwd/op_kernel/arch35/chunk_kda_fwd_prepare_arch35.h b/csrc/chunk_kda_fwd/op_kernel/arch35/chunk_kda_fwd_prepare_arch35.h new file mode 100644 index 0000000000..e9a34c1b1f --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/arch35/chunk_kda_fwd_prepare_arch35.h @@ -0,0 +1,4778 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#pragma once + +#ifndef CATLASS_ARCH +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 +#define CATLASS_ARCH 3510 +#else +#define CATLASS_ARCH 2201 +#endif +#endif + +#include "catlass/arch/kda_arch.hpp" +#include "catlass/arch/kda_cross_core_sync.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/kda_catlass.hpp" +#include "catlass/gemm/block/kda_block_mmad.hpp" +#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" +#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "../kernel_utils/block/block_mmad_pingpong_tla_multi.hpp" +#include "../kernel_utils/tile/copy_l0c_to_ub.hpp" +#include "catlass/layout/kda_layout.hpp" +#include "kernel_operator.h" +#include "../chunk_kda_fwd_varlen.h" +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 +#ifndef FLA_NPU_REGBASE_HPP_INCLUDED +#define FLA_NPU_REGBASE_HPP_INCLUDED +#include "../kernel_utils/vector/regbase.hpp" +#endif +#endif +#include "tla/kda_tla_layout.hpp" +#include "tla/kda_tensor.hpp" +#include "chunk_kda_fwd_post_wu_arch35.h" + +using namespace AscendC; + +namespace KdaPrepare { +namespace { +using KdaInt64 = tla::Int<64>; +using KdaInt128 = tla::Int<128>; +constexpr float LN2 = 0.69314718055994530942f; +constexpr float RCP_LN2 = 1.44269504088896340736f; +constexpr float KDA_EXP2_CLAMP = 80.0f; +constexpr float KDA_EXP_INPUT_MAX = KDA_EXP2_CLAMP * LN2; +constexpr float KDA_EXP_INPUT_MIN = -KDA_EXP2_CLAMP * LN2; +constexpr float KDA_SCORE_EXP2_CLAMP = 120.0f; +constexpr float KDA_SCORE_EXP2_MIN_CLAMP = 126.0f; +constexpr float KDA_SCORE_EXP_INPUT_MAX = KDA_SCORE_EXP2_CLAMP * LN2; +constexpr float KDA_SCORE_EXP_INPUT_MIN = -KDA_SCORE_EXP2_MIN_CLAMP * LN2; +constexpr float KDA_FP16_MAX = 65504.0f; +constexpr uint32_t EXP2_UB_ELEMENTS = 256; +constexpr uint32_t EXP2_UB_BYTES = EXP2_UB_ELEMENTS * (sizeof(float) + sizeof(uint16_t)); +constexpr uint32_t EXP2_EVENT_ID = 0; +constexpr uint32_t KDA_SOLVE_BT = 64; +constexpr uint32_t KDA_SOLVE_MATRIX_ELEMENTS = KDA_SOLVE_BT * KDA_SOLVE_BT; +constexpr uint32_t KDA_SOLVE_SCRATCH_X = 0; +constexpr uint32_t KDA_SOLVE_SCRATCH_Y0 = 1; +constexpr uint32_t KDA_SOLVE_SCRATCH_TMP = 2; +constexpr uint32_t KDA_SOLVE_SCRATCH_Y1 = 3; +constexpr uint32_t KDA_SOLVE_SCRATCH_IDENTITY = 4; +constexpr uint32_t KDA_SOLVE_SCRATCH_RAW_AKK = KDA_SOLVE_SCRATCH_Y1; +constexpr uint32_t KDA_SOLVE_SCRATCH_RAW_AQK = KDA_SOLVE_SCRATCH_IDENTITY; +constexpr uint32_t KDA_SOLVE_SCRATCH_SLOTS = 5; +constexpr uint32_t KDA_SOLVE_PIPELINE_DEPTH = 4; +constexpr uint32_t KDA_SOLVE_DIAG_BT = 16; +constexpr uint32_t KDA_SOLVE_DIAG_BLOCKS = KDA_SOLVE_BT / KDA_SOLVE_DIAG_BT; +constexpr uint32_t KDA_SOLVE_DIAG_MCH_ITERS = 3; +// Keep the local safe-gate exponent span within the BF16 score range while +// reducing repeated gate-factor work and AIV/AIC handshakes. +constexpr uint32_t KDA_SCORE_REF_BC = 32; +constexpr uint32_t KDA_SAFE_SCORE_REF_BC = 32; +constexpr uint32_t KDA_VEC_ARENA_ELEMENTS = 32768; +constexpr uint32_t KDA_BITS_PER_MASK_BYTE = 8; +constexpr uint32_t KDA_SELECT_COL_BLOCKS = 2; +constexpr uint32_t KDA_SELECT_COL_MASK_BYTES = KDA_SOLVE_MATRIX_ELEMENTS / KDA_BITS_PER_MASK_BYTE; +constexpr uint32_t KDA_SELECT_MASK_BYTES = KDA_SELECT_COL_BLOCKS * KDA_SELECT_COL_MASK_BYTES; +constexpr uint32_t KDA_SELECT_AQK_MASK_BYTE_OFFSET = 120 * 1024; +constexpr uint32_t KDA_SELECT_AKK_MASK_BYTE_OFFSET = KDA_SELECT_AQK_MASK_BYTE_OFFSET + KDA_SELECT_MASK_BYTES; +constexpr uint32_t KDA_SELECT_ZERO_BYTE_OFFSET = KDA_SELECT_AKK_MASK_BYTE_OFFSET + KDA_SELECT_MASK_BYTES; +constexpr uint32_t KDA_SELECT_ZERO_FLOAT_OFFSET = KDA_SELECT_ZERO_BYTE_OFFSET / sizeof(float); +constexpr uint8_t KDA_SCORE_DONE_FLAG0 = 2; +constexpr uint8_t KDA_SCORE_DONE_FLAG1 = 3; +constexpr uint8_t KDA_SCORE_READY_FLAG0 = 4; +constexpr uint8_t KDA_SCORE_READY_FLAG1 = 5; +constexpr uint8_t KDA_SOLVE_DONE_FLAG = 6; +constexpr uint8_t KDA_SOLVE_READY_FLAG = 7; +constexpr uint8_t KDA_POST_READY_FLAG = 0; +constexpr uint8_t KDA_POST_FREE_FLAG = 1; +constexpr uint32_t KDA_SCORE_QUEUE_DEPTH = 2; +constexpr uint32_t KDA_SCORE_LANES = 2; +constexpr uint32_t KDA_POST_QUEUE_DEPTH = 4; +constexpr uint32_t KDA_POST_QUEUE_STORAGE = KDA_POST_QUEUE_DEPTH + 1; +constexpr uint32_t KDA_DIRECT_SCORE_QUEUE_DEPTH = 2; +constexpr uint32_t KDA_DIRECT_SCORE_ROWS = 32; +constexpr uint32_t KDA_DIRECT_SCORE_MATRIX_ELEMENTS = KDA_DIRECT_SCORE_ROWS * 64; +constexpr uint32_t KDA_DIRECT_SCORE_SLOT_ELEMENTS = 3 * KDA_DIRECT_SCORE_MATRIX_ELEMENTS; +constexpr uint32_t KDA_DIRECT_SCORE_FLOAT_OFFSET = 20 * 1024; +constexpr uint32_t KDA_DIRECT_SCORE_UB_BYTE_OFFSET = EXP2_UB_BYTES + KDA_DIRECT_SCORE_FLOAT_OFFSET * sizeof(float); +constexpr uint32_t KDA_DIRECT_SCORE_L1_SLOT_ELEMENTS = 2 * KDA_DIRECT_SCORE_ROWS * 128; +constexpr uint32_t KDA_DIRECT_SCORE_L1_SLOT_BYTES = KDA_DIRECT_SCORE_L1_SLOT_ELEMENTS * sizeof(uint16_t); +constexpr uint32_t KDA_DIRECT_SCORE_L1_B_OFFSET = + KDA_SCORE_QUEUE_DEPTH * KDA_SCORE_LANES * KDA_DIRECT_SCORE_L1_SLOT_BYTES; +constexpr uint64_t KDA_DIRECT_SCORE_FREE_FLAG = 8; +constexpr uint64_t KDA_DIRECT_SCORE_READY_FLAG = 10; +constexpr uint64_t KDA_DIRECT_SCORE_SUBBLOCK_FLAG_STRIDE = 16; +constexpr uint32_t KDA_SCORE_SCRATCH_SLOTS = KDA_SCORE_QUEUE_DEPTH * KDA_SCORE_LANES; +constexpr uint32_t KDA_SYNC_REVERSE_DEPTH = 1; +constexpr uint32_t KDA_SCORE_SCRATCH_PLANES = 3; +constexpr uint32_t KDA_SCORE_SCRATCH_QG = 0; +constexpr uint32_t KDA_SCORE_SCRATCH_W = 1; +constexpr uint32_t KDA_SCORE_SCRATCH_KG = 2; +constexpr uint64_t KDA_WORKSPACE_ALIGN = 512; +constexpr uint32_t KDA_GATE_TILE_ROWS = 16; +constexpr uint32_t KDA_GATE_PIPELINE_DEPTH = 3; +constexpr uint32_t KDA_AIV_UB_BUDGET_BYTES = 192 * 1024; +constexpr uint32_t KDA_LOCAL_GK_FLOAT_OFFSET = 10 * 1024; +constexpr uint32_t KDA_SCALED_QG_FLOAT_OFFSET = 18 * 1024; +constexpr bool KDA_ARCH35_ENABLE_HEAD_PAIR = true; +constexpr bool KDA_ARCH35_ENABLE_MANUAL_SCORE_PIPELINE = false; +constexpr bool KDA_ARCH35_ENABLE_DIRECT_SCORE_UB = true; +constexpr bool KDA_ARCH35_ENABLE_DIRECT_SCORE_L1 = false; +constexpr uint16_t KDA_ARCH35_SCORE_EVENT = 3; +constexpr uint16_t KDA_ARCH35_SCORE_W_EVENT = 4; +constexpr uint16_t KDA_ARCH35_SOLVE_FIX_EVENT = 7; + +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 +template +static __simd_vf__ inline void AccumulateRawSafeGateChunk128Regbase(__ubuf__ float *input, __ubuf__ float *bias, + __ubuf__ float *acc, uint16_t rows, float expA, + float lowerBound) +{ + using namespace AscendC::MicroAPI; + constexpr uint16_t FLOAT_ELEMENTS_PER_REG = AscendC::VECTOR_REG_WIDTH / sizeof(float); + constexpr uint16_t ROW_ELEMENTS = 2 * FLOAT_ELEMENTS_PER_REG; + + MaskReg floatMask = CreateMask(); + RegTensor accZeroReg; + RegTensor accOneReg; + RegTensor oneZeroReg; + RegTensor oneOneReg; + RegTensor biasZeroReg; + RegTensor biasOneReg; + LoadAlign(accZeroReg, acc); + LoadAlign(accOneReg, acc + FLOAT_ELEMENTS_PER_REG); + Duplicate(oneZeroReg, 1.0f, floatMask); + Duplicate(oneOneReg, 1.0f, floatMask); + if constexpr (HAS_BIAS) { + LoadAlign(biasZeroReg, bias); + LoadAlign(biasOneReg, bias + FLOAT_ELEMENTS_PER_REG); + } + + const float gateScale = lowerBound * RCP_LN2; + for (uint16_t row = 0; row < rows; ++row) { + const uint32_t rowOffset = static_cast(row) * ROW_ELEMENTS; + RegTensor gateZeroReg; + RegTensor gateOneReg; + RegTensor sigmoidZeroReg; + RegTensor sigmoidOneReg; + LoadAlign(gateZeroReg, input + rowOffset); + LoadAlign(gateOneReg, input + rowOffset + FLOAT_ELEMENTS_PER_REG); + if constexpr (HAS_BIAS) { + Add(gateZeroReg, gateZeroReg, biasZeroReg, floatMask); + Add(gateOneReg, gateOneReg, biasOneReg, floatMask); + } + Muls(gateZeroReg, gateZeroReg, -expA, floatMask); + Muls(gateOneReg, gateOneReg, -expA, floatMask); + Exp(gateZeroReg, gateZeroReg, floatMask); + Exp(gateOneReg, gateOneReg, floatMask); + Adds(gateZeroReg, gateZeroReg, 1.0f, floatMask); + Adds(gateOneReg, gateOneReg, 1.0f, floatMask); + Div(sigmoidZeroReg, oneZeroReg, gateZeroReg, floatMask); + Div(sigmoidOneReg, oneOneReg, gateOneReg, floatMask); + Muls(sigmoidZeroReg, sigmoidZeroReg, gateScale, floatMask); + Muls(sigmoidOneReg, sigmoidOneReg, gateScale, floatMask); + Add(accZeroReg, accZeroReg, sigmoidZeroReg, floatMask); + Add(accOneReg, accOneReg, sigmoidOneReg, floatMask); + StoreAlign(input + rowOffset, accZeroReg, floatMask); + StoreAlign(input + rowOffset + FLOAT_ELEMENTS_PER_REG, accOneReg, floatMask); + } + StoreAlign(acc, accZeroReg, floatMask); + StoreAlign(acc + FLOAT_ELEMENTS_PER_REG, accOneReg, floatMask); +} + +template +__simd_callee__ inline void LoadKdaGateRegbasePair(AscendC::MicroAPI::RegTensor &zeroReg, + AscendC::MicroAPI::RegTensor &oneReg, __ubuf__ InputT *src, + AscendC::MicroAPI::MaskReg &inputMask) +{ + using namespace AscendC::MicroAPI; + if constexpr (std::is_same()) { + LoadAlign(zeroReg, oneReg, src); + } else { + RegTensor inputReg; + LoadIn(inputReg, src); + CastHalf2Float(zeroReg, oneReg, inputReg, inputMask); + } +} + +template +__simd_callee__ inline void ClampKdaGateRegbaseOutput(AscendC::MicroAPI::RegTensor &zeroReg, + AscendC::MicroAPI::RegTensor &oneReg, + AscendC::MicroAPI::MaskReg &floatMask) +{ + using namespace AscendC::MicroAPI; + if constexpr (std::is_same()) { + Mins(zeroReg, zeroReg, KDA_FP16_MAX, floatMask); + Mins(oneReg, oneReg, KDA_FP16_MAX, floatMask); + Maxs(zeroReg, zeroReg, -KDA_FP16_MAX, floatMask); + Maxs(oneReg, oneReg, -KDA_FP16_MAX, floatMask); + } +} + +template +__simd_callee__ inline void BuildKdaGateRegbaseExp(AscendC::MicroAPI::RegTensor &expZeroReg, + AscendC::MicroAPI::RegTensor &expOneReg, + AscendC::MicroAPI::RegTensor &gateZeroReg, + AscendC::MicroAPI::RegTensor &gateOneReg, __ubuf__ float *ref, + AscendC::MicroAPI::MaskReg &floatMask) +{ + using namespace AscendC::MicroAPI; + constexpr float expInputMax = std::is_same() ? KDA_SCORE_EXP_INPUT_MAX : KDA_EXP_INPUT_MAX; + constexpr float expInputMin = std::is_same() ? KDA_SCORE_EXP_INPUT_MIN : KDA_EXP_INPUT_MIN; + if constexpr (USE_REF) { + RegTensor refZeroReg; + RegTensor refOneReg; + LoadAlign(refZeroReg, refOneReg, ref); + if constexpr (NEGATIVE) { + SubFloatTwoReg(expZeroReg, expOneReg, refZeroReg, refOneReg, gateZeroReg, gateOneReg, floatMask); + } else { + SubFloatTwoReg(expZeroReg, expOneReg, gateZeroReg, gateOneReg, refZeroReg, refOneReg, floatMask); + } + } else if constexpr (NEGATIVE) { + Muls(expZeroReg, gateZeroReg, -1.0f, floatMask); + Muls(expOneReg, gateOneReg, -1.0f, floatMask); + } else { + Adds(expZeroReg, gateZeroReg, 0.0f, floatMask); + Adds(expOneReg, gateOneReg, 0.0f, floatMask); + } + Muls(expZeroReg, expZeroReg, LN2, floatMask); + Muls(expOneReg, expOneReg, LN2, floatMask); + MinsFloatTwoReg(expZeroReg, expOneReg, expZeroReg, expOneReg, expInputMax, floatMask); + Maxs(expZeroReg, expZeroReg, expInputMin, floatMask); + Maxs(expOneReg, expOneReg, expInputMin, floatMask); + ExpFloatTwoReg(expZeroReg, expOneReg, expZeroReg, expOneReg, floatMask); +} + +template +__simd_callee__ inline void StoreKdaGateRegbasePair(__ubuf__ OutputT *dst, AscendC::MicroAPI::RegTensor &zeroReg, + AscendC::MicroAPI::RegTensor &oneReg, + AscendC::MicroAPI::MaskReg &inputMask, + AscendC::MicroAPI::MaskReg &floatMask) +{ + using namespace AscendC::MicroAPI; + RegTensor outputReg; + ClampKdaGateRegbaseOutput(zeroReg, oneReg, floatMask); + CastFloat2Half(outputReg, zeroReg, oneReg, floatMask); + StoreAlign(dst, outputReg, inputMask); +} + +template +static __simd_vf__ inline void PrepareKdaGateQwRegbase(__ubuf__ InputT *q, __ubuf__ InputT *k, __ubuf__ OutputT *qOut, + __ubuf__ OutputT *kOut, __ubuf__ InputT *qDirect, + __ubuf__ InputT *kDirect, __ubuf__ GK_T *gate, + __ubuf__ float *ref, uint16_t rows, uint16_t cols) +{ + using namespace AscendC::MicroAPI; + constexpr uint16_t ELEMENTS_PER_REG = AscendC::VECTOR_REG_WIDTH / sizeof(InputT); + + MaskReg floatMask = CreateMask(); + for (uint16_t row = 0; row < rows; ++row) { + uint32_t rowOffset = static_cast(row) * cols; + for (uint16_t col = 0; col < cols; col += ELEMENTS_PER_REG) { + uint32_t activeCount = static_cast(cols - col); + MaskReg inputMask = UpdateMask(activeCount); + uint32_t offset = rowOffset + col; + + RegTensor gateZeroReg; + RegTensor gateOneReg; + RegTensor expZeroReg; + RegTensor expOneReg; + RegTensor directZeroReg; + RegTensor directOneReg; + RegTensor inputZeroReg; + RegTensor inputOneReg; + RegTensor outputZeroReg; + RegTensor outputOneReg; + + LoadKdaGateRegbasePair(gateZeroReg, gateOneReg, gate + offset, inputMask); + BuildKdaGateRegbaseExp(expZeroReg, expOneReg, gateZeroReg, gateOneReg, ref + col, + floatMask); + BuildKdaGateRegbaseExp(directZeroReg, directOneReg, gateZeroReg, gateOneReg, + ref + col, floatMask); + + LoadKdaGateRegbasePair(inputZeroReg, inputOneReg, q + offset, inputMask); + MulFloatTwoReg(outputZeroReg, outputOneReg, inputZeroReg, inputOneReg, directZeroReg, directOneReg, + floatMask); + StoreKdaGateRegbasePair(qDirect + offset, outputZeroReg, outputOneReg, inputMask, floatMask); + MulFloatTwoReg(outputZeroReg, outputOneReg, inputZeroReg, inputOneReg, expZeroReg, expOneReg, floatMask); + StoreKdaGateRegbasePair(qOut + offset, outputZeroReg, outputOneReg, inputMask, floatMask); + + LoadKdaGateRegbasePair(inputZeroReg, inputOneReg, k + offset, inputMask); + MulFloatTwoReg(outputZeroReg, outputOneReg, inputZeroReg, inputOneReg, directZeroReg, directOneReg, + floatMask); + StoreKdaGateRegbasePair(kDirect + offset, outputZeroReg, outputOneReg, inputMask, floatMask); + MulFloatTwoReg(outputZeroReg, outputOneReg, inputZeroReg, inputOneReg, expZeroReg, expOneReg, floatMask); + StoreKdaGateRegbasePair(kOut + offset, outputZeroReg, outputOneReg, inputMask, floatMask); + } + } +} + +template +static __simd_vf__ inline void PrepareKdaGateKgRegbase(__ubuf__ OutputT *kg, __ubuf__ InputT *k, __ubuf__ GK_T *gate, + __ubuf__ float *ref, uint16_t rows, uint16_t cols, + uint16_t validRows) +{ + using namespace AscendC::MicroAPI; + constexpr uint16_t ELEMENTS_PER_REG = AscendC::VECTOR_REG_WIDTH / sizeof(InputT); + + MaskReg floatMask = CreateMask(); + for (uint16_t row = 0; row < rows; ++row) { + uint32_t rowOffset = static_cast(row) * cols; + for (uint16_t col = 0; col < cols; col += ELEMENTS_PER_REG) { + uint32_t activeCount = static_cast(cols - col); + MaskReg inputMask = UpdateMask(activeCount); + uint32_t offset = rowOffset + col; + + RegTensor gateZeroReg; + RegTensor gateOneReg; + RegTensor expZeroReg; + RegTensor expOneReg; + RegTensor inputZeroReg; + RegTensor inputOneReg; + RegTensor outputZeroReg; + RegTensor outputOneReg; + + LoadKdaGateRegbasePair(gateZeroReg, gateOneReg, gate + offset, inputMask); + BuildKdaGateRegbaseExp(expZeroReg, expOneReg, gateZeroReg, gateOneReg, ref + col, + floatMask); + LoadKdaGateRegbasePair(inputZeroReg, inputOneReg, k + offset, inputMask); + MulFloatTwoReg(outputZeroReg, outputOneReg, inputZeroReg, inputOneReg, expZeroReg, expOneReg, floatMask); + if constexpr (USE_REF) { + if (row >= validRows) { + Duplicate(outputZeroReg, 0.0f, floatMask); + Duplicate(outputOneReg, 0.0f, floatMask); + } + } + StoreKdaGateRegbasePair(kg + offset, outputZeroReg, outputOneReg, inputMask, floatMask); + } + } +} + +template +static __simd_vf__ inline void PrepareKdaGateQwKgRegbase( + __ubuf__ InputT *q, __ubuf__ InputT *k, __ubuf__ OutputT *qOut, __ubuf__ OutputT *wOut, __ubuf__ OutputT *kgOut, + __ubuf__ InputT *qDirect, __ubuf__ InputT *wDirect, __ubuf__ InputT *v, __ubuf__ InputT *vDirect, + __ubuf__ InputT *finalKgOut, __ubuf__ float *beta, __ubuf__ GK_T *gate, __ubuf__ float *ref, + __ubuf__ float *finalRef, uint16_t rows, uint16_t cols, uint16_t validRows) +{ + using namespace AscendC::MicroAPI; + constexpr uint16_t ELEMENTS_PER_REG = AscendC::VECTOR_REG_WIDTH / sizeof(InputT); + constexpr float scoreExpInputMax = + std::is_same() ? KDA_SCORE_EXP_INPUT_MAX : KDA_EXP_INPUT_MAX; + constexpr float scoreExpInputMin = + std::is_same() ? KDA_SCORE_EXP_INPUT_MIN : KDA_EXP_INPUT_MIN; + + MaskReg floatMask = CreateMask(); + for (uint16_t row = 0; row < rows; ++row) { + uint32_t rowOffset = static_cast(row) * cols; + for (uint16_t col = 0; col < cols; col += ELEMENTS_PER_REG) { + uint32_t activeCount = static_cast(cols - col); + MaskReg inputMask = UpdateMask(activeCount); + uint32_t offset = rowOffset + col; + + RegTensor gateZeroReg; + RegTensor gateOneReg; + RegTensor posZeroReg; + RegTensor posOneReg; + RegTensor negZeroReg; + RegTensor negOneReg; + RegTensor directZeroReg; + RegTensor directOneReg; + RegTensor finalNegZeroReg; + RegTensor finalNegOneReg; + RegTensor inputZeroReg; + RegTensor inputOneReg; + RegTensor outputZeroReg; + RegTensor outputOneReg; + RegTensor betaReg; + + LoadKdaGateRegbasePair(gateZeroReg, gateOneReg, gate + offset, inputMask); + if constexpr (USE_REF) { + RegTensor refZeroReg; + RegTensor refOneReg; + LoadAlign(refZeroReg, refOneReg, ref + col); + SubFloatTwoReg(posZeroReg, posOneReg, gateZeroReg, gateOneReg, refZeroReg, refOneReg, floatMask); + SubFloatTwoReg(negZeroReg, negOneReg, refZeroReg, refOneReg, gateZeroReg, gateOneReg, floatMask); + } else { + Adds(posZeroReg, gateZeroReg, 0.0f, floatMask); + Adds(posOneReg, gateOneReg, 0.0f, floatMask); + Muls(negZeroReg, gateZeroReg, -1.0f, floatMask); + Muls(negOneReg, gateOneReg, -1.0f, floatMask); + } + if constexpr (STORE_DIRECT) { + Adds(directZeroReg, gateZeroReg, 0.0f, floatMask); + Adds(directOneReg, gateOneReg, 0.0f, floatMask); + Muls(directZeroReg, directZeroReg, LN2, floatMask); + Muls(directOneReg, directOneReg, LN2, floatMask); + MinsFloatTwoReg(directZeroReg, directOneReg, directZeroReg, directOneReg, KDA_EXP_INPUT_MAX, floatMask); + Maxs(directZeroReg, directZeroReg, KDA_EXP_INPUT_MIN, floatMask); + Maxs(directOneReg, directOneReg, KDA_EXP_INPUT_MIN, floatMask); + ExpFloatTwoReg(directZeroReg, directOneReg, directZeroReg, directOneReg, floatMask); + } + Muls(posZeroReg, posZeroReg, LN2, floatMask); + Muls(posOneReg, posOneReg, LN2, floatMask); + Muls(negZeroReg, negZeroReg, LN2, floatMask); + Muls(negOneReg, negOneReg, LN2, floatMask); + MinsFloatTwoReg(posZeroReg, posOneReg, posZeroReg, posOneReg, scoreExpInputMax, floatMask); + MinsFloatTwoReg(negZeroReg, negOneReg, negZeroReg, negOneReg, scoreExpInputMax, floatMask); + Maxs(posZeroReg, posZeroReg, scoreExpInputMin, floatMask); + Maxs(posOneReg, posOneReg, scoreExpInputMin, floatMask); + Maxs(negZeroReg, negZeroReg, scoreExpInputMin, floatMask); + Maxs(negOneReg, negOneReg, scoreExpInputMin, floatMask); + ExpFloatTwoReg(posZeroReg, posOneReg, posZeroReg, posOneReg, floatMask); + ExpFloatTwoReg(negZeroReg, negOneReg, negZeroReg, negOneReg, floatMask); + + LoadKdaGateRegbasePair(inputZeroReg, inputOneReg, q + offset, inputMask); + if constexpr (STORE_DIRECT) { + MulFloatTwoReg(outputZeroReg, outputOneReg, inputZeroReg, inputOneReg, directZeroReg, directOneReg, + floatMask); + StoreKdaGateRegbasePair(qDirect + offset, outputZeroReg, outputOneReg, inputMask, floatMask); + } + MulFloatTwoReg(outputZeroReg, outputOneReg, inputZeroReg, inputOneReg, posZeroReg, posOneReg, floatMask); + StoreKdaGateRegbasePair(qOut + offset, outputZeroReg, outputOneReg, inputMask, floatMask); + + LoadKdaGateRegbasePair(inputZeroReg, inputOneReg, k + offset, inputMask); + if constexpr (STORE_DIRECT || SCALE_SCORE_W) { + LoadAlign(betaReg, beta + row); + } + if constexpr (STORE_DIRECT) { + MulFloatTwoReg(outputZeroReg, outputOneReg, inputZeroReg, inputOneReg, directZeroReg, directOneReg, + floatMask); + RegTensor roundedReg; + ClampKdaGateRegbaseOutput(outputZeroReg, outputOneReg, floatMask); + CastFloat2Half(roundedReg, outputZeroReg, outputOneReg, floatMask); + CastHalf2Float(outputZeroReg, outputOneReg, roundedReg, inputMask); + Mul(outputZeroReg, outputZeroReg, betaReg, floatMask); + Mul(outputOneReg, outputOneReg, betaReg, floatMask); + StoreKdaGateRegbasePair(wDirect + offset, outputZeroReg, outputOneReg, inputMask, floatMask); + } + MulFloatTwoReg(outputZeroReg, outputOneReg, inputZeroReg, inputOneReg, posZeroReg, posOneReg, floatMask); + if constexpr (SCALE_SCORE_W) { + Mul(outputZeroReg, outputZeroReg, betaReg, floatMask); + Mul(outputOneReg, outputOneReg, betaReg, floatMask); + } + StoreKdaGateRegbasePair(wOut + offset, outputZeroReg, outputOneReg, inputMask, floatMask); + MulFloatTwoReg(outputZeroReg, outputOneReg, inputZeroReg, inputOneReg, negZeroReg, negOneReg, floatMask); + if constexpr (USE_REF) { + if (row >= validRows) { + Duplicate(outputZeroReg, 0.0f, floatMask); + Duplicate(outputOneReg, 0.0f, floatMask); + } + } + StoreKdaGateRegbasePair(kgOut + offset, outputZeroReg, outputOneReg, inputMask, floatMask); + + if constexpr (STORE_DIRECT) { + LoadKdaGateRegbasePair(inputZeroReg, inputOneReg, v + offset, inputMask); + Mul(outputZeroReg, inputZeroReg, betaReg, floatMask); + Mul(outputOneReg, inputOneReg, betaReg, floatMask); + StoreKdaGateRegbasePair(vDirect + offset, outputZeroReg, outputOneReg, inputMask, floatMask); + } + if constexpr (EXPORT_FINAL_KG) { + RegTensor finalRefZeroReg; + RegTensor finalRefOneReg; + LoadAlign(finalRefZeroReg, finalRefOneReg, finalRef + col); + SubFloatTwoReg(finalNegZeroReg, finalNegOneReg, finalRefZeroReg, finalRefOneReg, gateZeroReg, + gateOneReg, floatMask); + Muls(finalNegZeroReg, finalNegZeroReg, LN2, floatMask); + Muls(finalNegOneReg, finalNegOneReg, LN2, floatMask); + MinsFloatTwoReg(finalNegZeroReg, finalNegOneReg, finalNegZeroReg, finalNegOneReg, KDA_EXP_INPUT_MAX, + floatMask); + Maxs(finalNegZeroReg, finalNegZeroReg, KDA_EXP_INPUT_MIN, floatMask); + Maxs(finalNegOneReg, finalNegOneReg, KDA_EXP_INPUT_MIN, floatMask); + ExpFloatTwoReg(finalNegZeroReg, finalNegOneReg, finalNegZeroReg, finalNegOneReg, floatMask); + LoadKdaGateRegbasePair(inputZeroReg, inputOneReg, k + offset, inputMask); + MulFloatTwoReg(outputZeroReg, outputOneReg, inputZeroReg, inputOneReg, finalNegZeroReg, finalNegOneReg, + floatMask); + StoreKdaGateRegbasePair(finalKgOut + offset, outputZeroReg, outputOneReg, inputMask, floatMask); + } + } + } +} + +static __simd_vf__ inline void ForwardSubDiag16Regbase(__ubuf__ float *diag, uint16_t valid) +{ + using namespace AscendC::MicroAPI; + constexpr uint16_t DIAG_SIZE = KDA_SOLVE_DIAG_BT; + uint32_t activeCount = DIAG_SIZE; + MaskReg rowMask = UpdateMask(activeCount); + + for (uint16_t row = 2; row < valid; ++row) { + RegTensor currentReg; + RegTensor scaleReg; + RegTensor matrixReg; + RegTensor productReg; + RegTensor sumReg; + LoadAlign(currentReg, diag + static_cast(row) * DIAG_SIZE); + Duplicate(sumReg, 0.0f, rowMask); + + for (uint16_t sourceRow = 0; sourceRow < row; ++sourceRow) { + LoadAlign(scaleReg, + diag + static_cast(row) * DIAG_SIZE + sourceRow); + LoadAlign(matrixReg, diag + static_cast(sourceRow) * DIAG_SIZE); + Mul(productReg, matrixReg, scaleReg, rowMask); + Add(sumReg, sumReg, productReg, rowMask); + } + Add(currentReg, currentReg, sumReg, rowMask); + StoreAlign(diag + static_cast(row) * DIAG_SIZE, currentReg, rowMask); + LocalMemBar(); + } +} + +static __simd_vf__ inline void SelectCausalRows64Regbase(__ubuf__ float *aqk, __ubuf__ float *akk, uint16_t rowBegin, + uint16_t rowCount) +{ + using namespace AscendC::MicroAPI; + constexpr uint16_t ROW_ELEMENTS = 64; + + MaskReg fullMask = CreateMask(); + for (uint16_t localRow = 0; localRow < rowCount; ++localRow) { + const uint16_t row = rowBegin + localRow; + const uint32_t rowOffset = static_cast(localRow) * ROW_ELEMENTS; + RegTensor zeroReg; + RegTensor aqkInputReg; + RegTensor akkInputReg; + RegTensor aqkReg; + RegTensor akkReg; + uint32_t aqkCount = static_cast(row) + 1; + uint32_t akkCount = static_cast(row); + MaskReg aqkMask = UpdateMask(aqkCount); + MaskReg akkMask = UpdateMask(akkCount); + Duplicate(zeroReg, 0.0f, fullMask); + LoadAlign(aqkInputReg, aqk + rowOffset); + LoadAlign(akkInputReg, akk + rowOffset); + Select(aqkReg, aqkInputReg, zeroReg, aqkMask); + Select(akkReg, akkInputReg, zeroReg, akkMask); + StoreAlign(aqk + rowOffset, aqkReg, fullMask); + StoreAlign(akk + rowOffset, akkReg, fullMask); + } +} + +static __simd_vf__ inline void ForwardSubDiag16StridedRegbase(__ubuf__ float *matrix, uint16_t rowStride, + uint16_t rowBegin, uint16_t colBegin, uint16_t valid) +{ + using namespace AscendC::MicroAPI; + constexpr uint16_t DIAG_SIZE = KDA_SOLVE_DIAG_BT; + uint32_t activeCount = DIAG_SIZE; + MaskReg rowMask = UpdateMask(activeCount); + + for (uint16_t row = 2; row < valid; ++row) { + uint32_t currentOffset = static_cast(rowBegin + row) * rowStride + colBegin; + RegTensor currentReg; + RegTensor scaleReg; + RegTensor matrixReg; + RegTensor productReg; + RegTensor sumReg; + LoadAlign(currentReg, matrix + currentOffset); + Duplicate(sumReg, 0.0f, rowMask); + + for (uint16_t sourceRow = 0; sourceRow < row; ++sourceRow) { + LoadAlign(scaleReg, matrix + currentOffset + sourceRow); + uint32_t sourceOffset = static_cast(rowBegin + sourceRow) * rowStride + colBegin; + LoadAlign(matrixReg, matrix + sourceOffset); + Mul(productReg, matrixReg, scaleReg, rowMask); + Add(sumReg, sumReg, productReg, rowMask); + } + Add(currentReg, currentReg, sumReg, rowMask); + StoreAlign(matrix + currentOffset, currentReg, rowMask); + LocalMemBar(); + } +} + +static __simd_vf__ inline void ForwardSubDiag16PairStridedRegbase(__ubuf__ float *matrix, uint16_t rowStride, + uint16_t colBegin, uint16_t firstValid, + uint16_t secondValid) +{ + using namespace AscendC::MicroAPI; + constexpr uint16_t DIAG_SIZE = KDA_SOLVE_DIAG_BT; + constexpr uint16_t SECOND_LOCAL_ROW = DIAG_SIZE; + uint32_t activeCount = DIAG_SIZE; + MaskReg rowMask = UpdateMask(activeCount); + + for (uint16_t row = 2; row < DIAG_SIZE; ++row) { + uint32_t firstCurrentOffset = static_cast(row) * rowStride + colBegin; + uint32_t secondCurrentOffset = static_cast(SECOND_LOCAL_ROW + row) * rowStride + colBegin + DIAG_SIZE; + RegTensor firstCurrentReg; + RegTensor secondCurrentReg; + RegTensor firstSumReg; + RegTensor secondSumReg; + LoadAlign(firstCurrentReg, matrix + firstCurrentOffset); + LoadAlign(secondCurrentReg, matrix + secondCurrentOffset); + Duplicate(firstSumReg, 0.0f, rowMask); + Duplicate(secondSumReg, 0.0f, rowMask); + + if (row < firstValid || row < secondValid) { + for (uint16_t sourceRow = 0; sourceRow < row; ++sourceRow) { + RegTensor firstScaleReg; + RegTensor secondScaleReg; + RegTensor firstMatrixReg; + RegTensor secondMatrixReg; + RegTensor firstProductReg; + RegTensor secondProductReg; + LoadAlign(firstScaleReg, matrix + firstCurrentOffset + sourceRow); + LoadAlign(secondScaleReg, matrix + secondCurrentOffset + sourceRow); + uint32_t firstSourceOffset = static_cast(sourceRow) * rowStride + colBegin; + uint32_t secondSourceOffset = + static_cast(SECOND_LOCAL_ROW + sourceRow) * rowStride + colBegin + DIAG_SIZE; + LoadAlign(firstMatrixReg, matrix + firstSourceOffset); + LoadAlign(secondMatrixReg, matrix + secondSourceOffset); + Mul(firstProductReg, firstMatrixReg, firstScaleReg, rowMask); + Mul(secondProductReg, secondMatrixReg, secondScaleReg, rowMask); + Add(firstSumReg, firstSumReg, firstProductReg, rowMask); + Add(secondSumReg, secondSumReg, secondProductReg, rowMask); + } + } + if (row < firstValid) { + Add(firstCurrentReg, firstCurrentReg, firstSumReg, rowMask); + StoreAlign(matrix + firstCurrentOffset, firstCurrentReg, rowMask); + } + if (row < secondValid) { + Add(secondCurrentReg, secondCurrentReg, secondSumReg, rowMask); + StoreAlign(matrix + secondCurrentOffset, secondCurrentReg, rowMask); + } + LocalMemBar(); + } + + RegTensor indexReg; + Arange(indexReg, 0); + for (uint16_t row = 0; row < DIAG_SIZE; ++row) { + MaskReg diagMask; + CompareScalar(diagMask, indexReg, static_cast(row), rowMask); + uint32_t firstOffset = static_cast(row) * rowStride + colBegin; + uint32_t secondOffset = static_cast(SECOND_LOCAL_ROW + row) * rowStride + colBegin + DIAG_SIZE; + RegTensor firstReg; + RegTensor secondReg; + LoadAlign(firstReg, matrix + firstOffset); + LoadAlign(secondReg, matrix + secondOffset); + Adds(firstReg, firstReg, 1.0f, diagMask); + Adds(secondReg, secondReg, 1.0f, diagMask); + StoreAlign(matrix + firstOffset, firstReg, rowMask); + StoreAlign(matrix + secondOffset, secondReg, rowMask); + } +} + +static __simd_vf__ inline void ApplyKdaRowScaleRegbase(__ubuf__ float *matrix, __ubuf__ float *rowScale, uint16_t rows, + uint16_t cols) +{ + using namespace AscendC::MicroAPI; + constexpr uint16_t FP32_PER_REG = AscendC::VECTOR_REG_WIDTH / sizeof(float); + RegTensor matrixReg0; + RegTensor matrixReg1; + RegTensor scaleReg0; + RegTensor scaleReg1; + + uint16_t row = 0; + for (; row + 1 < rows; row += 2) { + LoadAlign(scaleReg0, rowScale + row); + LoadAlign(scaleReg1, rowScale + row + 1); + for (uint16_t col = 0; col < cols; col += FP32_PER_REG) { + uint32_t activeCount0 = static_cast(cols - col); + uint32_t activeCount1 = activeCount0; + MaskReg mask0 = UpdateMask(activeCount0); + MaskReg mask1 = UpdateMask(activeCount1); + uint32_t offset0 = static_cast(row) * cols + col; + uint32_t offset1 = static_cast(row + 1) * cols + col; + LoadAlign(matrixReg0, matrix + offset0); + LoadAlign(matrixReg1, matrix + offset1); + Mul(matrixReg0, matrixReg0, scaleReg0, mask0); + Mul(matrixReg1, matrixReg1, scaleReg1, mask1); + StoreAlign(matrix + offset0, matrixReg0, mask0); + StoreAlign(matrix + offset1, matrixReg1, mask1); + } + } + if (row < rows) { + LoadAlign(scaleReg0, rowScale + row); + for (uint16_t col = 0; col < cols; col += FP32_PER_REG) { + uint32_t activeCount = static_cast(cols - col); + MaskReg mask = UpdateMask(activeCount); + uint32_t offset = static_cast(row) * cols + col; + LoadAlign(matrixReg0, matrix + offset); + Mul(matrixReg0, matrixReg0, scaleReg0, mask); + StoreAlign(matrix + offset, matrixReg0, mask); + } + } +} +#endif + +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 +using KdaArchTag = Catlass::Arch::Ascend950; +#else +using KdaArchTag = Catlass::Arch::AtlasA2; +#endif +using KdaDispatchPolicy = Catlass::Gemm::MmadPingpong; +using KdaScoreDispatchPolicy = Catlass::Gemm::MmadPingpongTlaMulti; +static_assert(KdaScoreDispatchPolicy::ENABLE_L1_RESIDENT, + "KDA Aqk/Akk score MMAD must keep the shared right matrix resident in L1"); +static_assert(KdaScoreDispatchPolicy::L1B_STAGES == 1, + "KDA Aqk/Akk score MMAD needs one L1 B slot so the second MMAD reuses it"); +using KdaSolveDispatchPolicy = Catlass::Gemm::MmadPingpong; +static_assert(!KdaSolveDispatchPolicy::USE_HF32_MODE, "KDA triangular solve must use IEEE FP32 Cube mode"); +using KdaL1TileShape = tla::Shape; +using KdaL0TileShape = KdaL1TileShape; +using KdaSolveL1TileShape = tla::Shape; +using KdaSolveL0TileShape = KdaSolveL1TileShape; + +__aicore__ inline uint32_t FloatToBits(float value) +{ + union Bits { + __aicore__ Bits() {} + float f; + uint32_t u; + } bits; + bits.f = value; + return bits.u; +} + +__aicore__ inline float BitsToFloat(uint32_t value) +{ + union Bits { + __aicore__ Bits() {} + uint32_t u; + float f; + } bits; + bits.u = value; + return bits.f; +} + +__aicore__ inline uint16_t Bf16ToBits(bfloat16_t value) +{ + union Bits { + __aicore__ Bits() {} + bfloat16_t f; + uint16_t u; + } bits; + bits.f = value; + return bits.u; +} + +__aicore__ inline bfloat16_t BitsToBf16(uint16_t value) +{ + union Bits { + __aicore__ Bits() {} + uint16_t u; + bfloat16_t f; + } bits; + bits.u = value; + return bits.f; +} + +template +__aicore__ inline T FloatToType(float value) +{ + if constexpr (IsSameType::value) { + uint32_t bits = FloatToBits(value); + uint32_t bias = 0x7FFFu + ((bits >> 16) & 1u); + return BitsToBf16(static_cast((bits + bias) >> 16)); + } + return static_cast(value); +} + +template +class ChunkKdaFwdPrepareKernel { +public: + using OUT_T = T; + using AKK_T = float; + using SCORE_T = std::conditional_t::value, bfloat16_t, T>; + template + __aicore__ inline void Init(GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR gk, GM_ADDR rawG, GM_ADDR aLog, GM_ADDR dtBias, + GM_ADDR beta, GM_ADDR initialState, GM_ADDR cuSeqlens, GM_ADDR chunkIndices, + GM_ADDR preparedQG, GM_ADDR preparedAqk, GM_ADDR propagatedVNew, GM_ADDR propagatedH, + GM_ADDR o, GM_ADDR finalState, GM_ADDR aqk, GM_ADDR akk, GM_ADDR w, GM_ADDR u, + GM_ADDR qg, GM_ADDR kg, GM_ADDR vNew, GM_ADDR h, GM_ADDR finalKg, GM_ADDR workspace, + const TilingData &tiling, TPipe *pipe, bool initVecBuffers = true, bool storeQG = true) + { + pipe_ = pipe; + q_.SetGlobalBuffer((__gm__ T *)q); + k_.SetGlobalBuffer((__gm__ T *)k); + v_.SetGlobalBuffer((__gm__ T *)v); + gk_.SetGlobalBuffer((__gm__ GK_T *)gk); + rawG_.SetGlobalBuffer((__gm__ float *)rawG); + aLog_.SetGlobalBuffer((__gm__ float *)aLog); + dtBias_.SetGlobalBuffer((__gm__ float *)dtBias); + beta_.SetGlobalBuffer((__gm__ BETA_T *)beta); + if (initialState != nullptr) { + initialState_.SetGlobalBuffer((__gm__ float *)initialState); + } + cuSeqlensAddr_ = reinterpret_cast<__gm__ int64_t *>(cuSeqlens); + if (preparedQG != nullptr) { + preparedQG_.SetGlobalBuffer((__gm__ T *)preparedQG); + } + if (preparedAqk != nullptr) { + preparedAqk_.SetGlobalBuffer((__gm__ T *)preparedAqk); + } + if (propagatedVNew != nullptr) { + propagatedVNew_.SetGlobalBuffer((__gm__ T *)propagatedVNew); + } + if (propagatedH != nullptr) { + propagatedH_.SetGlobalBuffer((__gm__ T *)propagatedH); + } + chunkIndicesAddr_ = reinterpret_cast<__gm__ int64_t *>(chunkIndices); + o_.SetGlobalBuffer((__gm__ OUT_T *)o); + finalState_.SetGlobalBuffer((__gm__ float *)finalState); + aqk_.SetGlobalBuffer((__gm__ float *)aqk); + akk_.SetGlobalBuffer((__gm__ AKK_T *)akk); + w_.SetGlobalBuffer((__gm__ T *)w); + u_.SetGlobalBuffer((__gm__ OUT_T *)u); + qg_.SetGlobalBuffer((__gm__ T *)qg); + kg_.SetGlobalBuffer((__gm__ T *)kg); + vNew_.SetGlobalBuffer((__gm__ T *)vNew); + h_.SetGlobalBuffer((__gm__ float *)h); + finalKg_.SetGlobalBuffer((__gm__ T *)finalKg); + solveWorkspace_.SetGlobalBuffer((__gm__ float *)workspace); + + B_ = tiling.batch; + N_ = tiling.seqNum; + H_ = tiling.qHeadNum; + HV_ = tiling.vHeadNum; + T_ = tiling.seqlen; + K_ = COMPILE_K == 0 ? tiling.kHeadDim : COMPILE_K; + V_ = COMPILE_V == 0 ? tiling.vHeadDim : COMPILE_V; + BT_ = COMPILE_BT == 0 ? tiling.chunkSize : COMPILE_BT; + NT_ = tiling.totalChunks; + scale_ = tiling.scale; + hasInitial_ = tiling.hasInitialState; + isVarLen_ = tiling.isVarLen; + inputSequenceMajor_ = tiling.inputSequenceMajor; + fusePostWu_ = tiling.fusePostWu; + computeGateInPrepare_ = tiling.computeGateInPrepare; + hasALog_ = tiling.hasALog; + hasDtBias_ = tiling.hasDtBias; + lowerBound_ = tiling.lowerBound; + storeQG_ = storeQG; + usedCoreNum_ = tiling.prepareUsedCoreNum; +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if constexpr (SAFE_GATE && !IsSameType::value) { + headPairMode_ = KDA_ARCH35_ENABLE_HEAD_PAIR && HV_ % KDA_SCORE_LANES == 0; + } +#endif + constexpr uint64_t solvePipelineDepth = SAFE_GATE ? KDA_SOLVE_PIPELINE_DEPTH : 1; + const uint64_t solveBytes = + usedCoreNum_ * solvePipelineDepth * KDA_SOLVE_SCRATCH_SLOTS * BT_ * BT_ * sizeof(float); + const uint64_t alignedSolveBytes = + (solveBytes + KDA_WORKSPACE_ALIGN - 1) / KDA_WORKSPACE_ALIGN * KDA_WORKSPACE_ALIGN; + scoreWorkspace_.SetGlobalBuffer((__gm__ SCORE_T *)(workspace + alignedSolveBytes)); + if ASCEND_IS_AIV { + uint64_t subBlockNum = static_cast(GetSubBlockNum()); + solveCoreIdx_ = subBlockNum == 0 ? 0 : static_cast(GetBlockIdx()) / subBlockNum; + } else { + solveCoreIdx_ = static_cast(GetBlockIdx()); + } + if (pipe_ != nullptr && initVecBuffers) { + pipe_->InitBuffer(exp2Buf_, EXP2_UB_BYTES); + pipe_->InitBuffer(vecBuf_, KDA_VEC_ARENA_ELEMENTS * sizeof(float)); + const uint64_t gateStageElems = GatePipelineRows() * K_; + const uint64_t gateInputSlotBytes = GateInputSlotBytes(); + const uint64_t gatePipelineBytes = GateBufferDepth() * (gateInputSlotBytes + gateStageElems * sizeof(T)); + pipe_->InitBuffer(gateWritebackBuf_, static_cast(gatePipelineBytes)); + AllocVectorEvents(); + } + } + __aicore__ inline void ProcessAivOnly() + { + isAivOnly_ = true; + ProcessPreAiv(); + ReleaseVectorEvents(); + } + + __aicore__ inline void ProcessAiv() + { + ProcessPreAiv(); + ReleaseVectorEvents(); + } + + __aicore__ inline void ProcessAic() { ProcessPreAic(); } + + template + __aicore__ inline void ProcessAicFused(PostWuOp &postWu) + { + ProcessPreAicHeadPairFused(postWu); + } + +private: + __aicore__ inline void AllocVectorEvents() + { + mte2ToVEvent_ = pipe_->AllocEventID(); + vToMte2Event_ = pipe_->AllocEventID(); + vToMte3Event_ = pipe_->AllocEventID(); + mte3ToVEvent_ = pipe_->AllocEventID(); + mte2ToMte3Event_ = pipe_->AllocEventID(); + vToSEvent_ = pipe_->AllocEventID(); + for (uint32_t slot = 0; slot < KDA_GATE_PIPELINE_DEPTH; ++slot) { + mte3ToMte2Events_[slot] = pipe_->AllocEventID(); + } + vectorEventsAllocated_ = true; + } + + __aicore__ inline void ReleaseVectorEvents() + { + if (!vectorEventsAllocated_) { + return; + } + pipe_->ReleaseEventID(mte2ToVEvent_); + pipe_->ReleaseEventID(vToMte2Event_); + pipe_->ReleaseEventID(vToMte3Event_); + pipe_->ReleaseEventID(mte3ToVEvent_); + pipe_->ReleaseEventID(mte2ToMte3Event_); + pipe_->ReleaseEventID(vToSEvent_); + for (uint32_t slot = 0; slot < KDA_GATE_PIPELINE_DEPTH; ++slot) { + pipe_->ReleaseEventID(mte3ToMte2Events_[slot]); + } + vectorEventsAllocated_ = false; + } + + __aicore__ inline uint64_t QOffset(uint64_t b, uint64_t h, uint64_t t, uint64_t d) const + { + if (inputSequenceMajor_) { + return ((b * T_ + t) * H_ + h) * K_ + d; + } + return ((b * H_ + h) * T_ + t) * K_ + d; + } + + __aicore__ inline uint64_t VInputOffset(uint64_t b, uint64_t hv, uint64_t t, uint64_t d) const + { + if (inputSequenceMajor_) { + return ((b * T_ + t) * HV_ + hv) * V_ + d; + } + return ((b * HV_ + hv) * T_ + t) * V_ + d; + } + + __aicore__ inline uint64_t RawGateOffset(uint64_t b, uint64_t hv, uint64_t t, uint64_t d) const + { + if (inputSequenceMajor_) { + return ((b * T_ + t) * HV_ + hv) * K_ + d; + } + return ((b * HV_ + hv) * T_ + t) * K_ + d; + } + + __aicore__ inline uint64_t KVOffset(uint64_t b, uint64_t hv, uint64_t t, uint64_t d, uint64_t dim) const + { + return ((b * HV_ + hv) * T_ + t) * dim + d; + } + + __aicore__ inline uint64_t BetaOffset(uint64_t b, uint64_t hv, uint64_t t) const { return (b * HV_ + hv) * T_ + t; } + + __aicore__ inline uint64_t AOffset(uint64_t b, uint64_t hv, uint64_t t, uint64_t j) const + { + return ((b * HV_ + hv) * T_ + t) * BT_ + j; + } + + __aicore__ inline uint64_t HOffset(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t d, uint64_t r) const + { + return (((b * HV_ + hv) * NT_ + chunkIdx) * K_ + d) * V_ + r; + } + + __aicore__ inline uint64_t WScratchOffset(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t t, uint64_t d) const + { + return (((b * HV_ + hv) * NT_ + chunkIdx) * BT_ + t) * K_ + d; + } + + __aicore__ inline uint64_t SolveScratchOffset(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t slot) const + { + (void)b; + (void)hv; + (void)chunkIdx; + constexpr uint64_t solvePipelineDepth = SAFE_GATE ? KDA_SOLVE_PIPELINE_DEPTH : 1; + uint64_t matrixElements = BT_ * BT_; + return ((solveCoreIdx_ * solvePipelineDepth + activeSolveSlot_) * KDA_SOLVE_SCRATCH_SLOTS + slot) * + matrixElements; + } + + __aicore__ inline uint64_t ScoreScratchOffset(uint64_t slot, uint64_t plane, uint64_t t = 0, uint64_t d = 0) const + { + return (((solveCoreIdx_ * KDA_SCORE_SCRATCH_SLOTS + slot) * KDA_SCORE_SCRATCH_PLANES + plane) * BT_ + t) * K_ + + d; + } + + __aicore__ inline uint64_t ScoreScratchSlot(uint64_t queueSlot, uint64_t lane, bool pairHeads) const + { + return pairHeads ? queueSlot * KDA_SCORE_LANES + lane : queueSlot; + } + + __aicore__ inline uint64_t ScoreRefBlockSize() const + { + if constexpr (SAFE_GATE) { + return KDA_SAFE_SCORE_REF_BC; + } + return KDA_SCORE_REF_BC; + } + + __aicore__ inline uint64_t ScoreRowBlockCount(uint64_t curT, uint64_t rowBegin) const + { + uint64_t blockSize = ScoreRefBlockSize(); + uint64_t rowCount = curT - rowBegin; + if (rowCount > blockSize) { + rowCount = blockSize; + } + return rowCount; + } + + __aicore__ inline uint64_t ScoreRefToken(uint64_t start, uint64_t curT, uint64_t rowBegin, uint64_t rowCount) const + { + uint64_t ref = rowBegin + rowCount / 2; + if (ref >= curT) { + ref = curT - 1; + } + return start + ref; + } + + __aicore__ inline void RunExp2(LocalTensor &tensor, uint32_t count) + { + SetFlag(EXP2_EVENT_ID); + WaitFlag(EXP2_EVENT_ID); + ClampExpInput(tensor, count); + Exp(tensor, tensor, count); + PipeBarrier(); + SetFlag(EXP2_EVENT_ID); + WaitFlag(EXP2_EVENT_ID); + } + + __aicore__ inline void ClampExpInput(LocalTensor &tensor, uint32_t count) + { + Mins(tensor, tensor, KDA_EXP_INPUT_MAX, count); + PipeBarrier(); + Maxs(tensor, tensor, KDA_EXP_INPUT_MIN, count); + PipeBarrier(); + } + + __aicore__ inline void ClampScoreExpInput(LocalTensor &tensor, uint32_t count) + { + constexpr float expInputMax = + IsSameType::value ? KDA_SCORE_EXP_INPUT_MAX : KDA_EXP_INPUT_MAX; + constexpr float expInputMin = + IsSameType::value ? KDA_SCORE_EXP_INPUT_MIN : KDA_EXP_INPUT_MIN; + Mins(tensor, tensor, expInputMax, count); + PipeBarrier(); + Maxs(tensor, tensor, expInputMin, count); + PipeBarrier(); + } + + template + __aicore__ inline void ClampFp32ForCast(LocalTensor &tensor, uint32_t count) + { + if constexpr (IsSameType::value) { + Mins(tensor, tensor, KDA_FP16_MAX, count); + PipeBarrier(); + Maxs(tensor, tensor, -KDA_FP16_MAX, count); + PipeBarrier(); + } + } + + __aicore__ inline void ClampFp32ToOutputType(LocalTensor &tensor, uint32_t count) + { + ClampFp32ForCast(tensor, count); + } + + template + __aicore__ inline void CopyVectorIn(LocalTensor &dst, GlobalTensor &src, uint64_t offset, + uint64_t count) + { + uint64_t rowBytes = count * static_cast(sizeof(CopyT)); + if (rowBytes >= 32 && rowBytes % 32 == 0) { + DataCopy(dst, src[offset], static_cast(count)); + return; + } + DataCopyParams params{1, static_cast(rowBytes), 0, 0}; + DataCopyPadParams padParams{false, 0, 0, 0}; + DataCopyPad(dst, src[offset], params, padParams); + } + + template + __aicore__ inline void CopyRowsIn(LocalTensor &dst, GlobalTensor &src, uint64_t offset, uint64_t rows, + uint64_t cols, uint64_t rowStride) + { + if (rows == 0 || cols == 0) { + return; + } + if (rowStride == cols) { + CopyVectorIn(dst, src, offset, rows * cols); + return; + } + DataCopyExtParams params{static_cast(rows), static_cast(cols * sizeof(CopyT)), + static_cast((rowStride - cols) * sizeof(CopyT)), 0, 0}; + DataCopyPadExtParams padParams{false, 0, 0, 0}; + DataCopyPad(dst, src[offset], params, padParams); + } + + template + __aicore__ inline void CopyVectorOut(GlobalTensor &dst, uint64_t offset, LocalTensor &src, + uint64_t count) + { + uint64_t rowBytes = count * static_cast(sizeof(CopyT)); + if (rowBytes >= 32 && rowBytes % 32 == 0) { + DataCopy(dst[offset], src, static_cast(count)); + return; + } + DataCopyParams params{1, static_cast(rowBytes), 0, 0}; + DataCopyPad(dst[offset], src, params); + } + + template + __aicore__ inline void CopyRowIn(LocalTensor &dst, GlobalTensor &src, uint64_t offset) + { + CopyVectorIn(dst, src, offset, K_); + } + + template + __aicore__ inline void CopyRowOut(GlobalTensor &dst, uint64_t offset, LocalTensor &src) + { + CopyVectorOut(dst, offset, src, K_); + } + + __aicore__ inline LocalTensor VecScratch(uint64_t slot) + { + return vecBuf_.Get()[slot * EXP2_UB_ELEMENTS]; + } + + __aicore__ inline uint64_t GateStageElems() const { return GatePipelineRows() * K_; } + + __aicore__ inline uint64_t GatePipelineRows() const + { + constexpr uint64_t fixedBytes = static_cast(KDA_VEC_ARENA_ELEMENTS) * sizeof(float) + EXP2_UB_BYTES; + constexpr uint64_t availableBytes = KDA_AIV_UB_BUDGET_BYTES - fixedBytes; + uint64_t bytesPerRow = K_ * KDA_GATE_PIPELINE_DEPTH * (3 * sizeof(T) + sizeof(GK_T)); +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if constexpr (SAFE_GATE && COMPILE_BT == 64 && COMPILE_K == 128 && COMPILE_V == 128) { + bytesPerRow = GateBufferDepth() * (K_ * (4 * sizeof(T) + sizeof(GK_T)) + sizeof(BETA_T)); + } +#endif + uint64_t rows = bytesPerRow == 0 ? 0 : availableBytes / bytesPerRow; + return rows < KDA_GATE_TILE_ROWS ? rows : KDA_GATE_TILE_ROWS; + } + + __aicore__ inline constexpr uint64_t GateBufferDepth() const + { +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if constexpr (SAFE_GATE && COMPILE_BT == 64 && COMPILE_K == 128 && COMPILE_V == 128) { + return 2; + } +#endif + return KDA_GATE_PIPELINE_DEPTH; + } + + __aicore__ inline uint64_t GateInputSlotBytes() const + { +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if constexpr (SAFE_GATE && COMPILE_BT == 64 && COMPILE_K == 128 && COMPILE_V == 128) { + return GateStageElems() * (3 * sizeof(T) + sizeof(GK_T)) + GatePipelineRows() * sizeof(BETA_T); + } +#endif + return GateStageElems() * (2 * sizeof(T) + sizeof(GK_T)); + } + + __aicore__ inline LocalTensor GateQTyped(uint64_t slot) + { + uint64_t byteOffset = slot * GateInputSlotBytes(); + return gateWritebackBuf_.Get()[byteOffset / sizeof(T)]; + } + + __aicore__ inline LocalTensor GateKTyped(uint64_t slot) + { + uint64_t byteOffset = slot * GateInputSlotBytes() + GateStageElems() * sizeof(T); + return gateWritebackBuf_.Get()[byteOffset / sizeof(T)]; + } + + __aicore__ inline LocalTensor GateGTyped(uint64_t slot) + { + uint64_t byteOffset = slot * GateInputSlotBytes() + 2 * GateStageElems() * sizeof(T); + return gateWritebackBuf_.Get()[byteOffset / sizeof(GK_T)]; + } + + __aicore__ inline LocalTensor GateVTyped(uint64_t slot) + { + uint64_t byteOffset = slot * GateInputSlotBytes() + GateStageElems() * (2 * sizeof(T) + sizeof(GK_T)); + return gateWritebackBuf_.Get()[byteOffset / sizeof(T)]; + } + + __aicore__ inline LocalTensor GateBetaTyped(uint64_t slot) + { + uint64_t byteOffset = slot * GateInputSlotBytes() + GateStageElems() * (3 * sizeof(T) + sizeof(GK_T)); + return gateWritebackBuf_.Get()[byteOffset / sizeof(BETA_T)]; + } + + __aicore__ inline LocalTensor GateKgTyped(uint64_t slot) + { + uint64_t byteOffset = GateBufferDepth() * GateInputSlotBytes() + slot * GateStageElems() * sizeof(T); + return gateWritebackBuf_.Get()[byteOffset / sizeof(T)]; + } + + __aicore__ inline LocalTensor LocalGateChunk() + { + constexpr uint64_t byteOffset = static_cast(KDA_LOCAL_GK_FLOAT_OFFSET) * sizeof(float); + return vecBuf_.Get()[byteOffset / sizeof(GK_T)]; + } + + __aicore__ inline LocalTensor GateScoreTyped(uint64_t slot, uint64_t tileRow) + { + if constexpr (IsSameType::value) { + if (computeGateInPrepare_) { + return LocalGateChunk()[tileRow * K_]; + } + } + return GateGTyped(slot); + } + + __aicore__ inline void LoadGateScoreRef(LocalTensor dst, uint64_t b, uint64_t hv, uint64_t token) + { + if constexpr (IsSameType::value) { + if (computeGateInPrepare_) { + const uint64_t tileRow = token - activeGateChunkStart_; + Adds(dst, LocalGateChunk()[tileRow * K_], 0.0f, static_cast(K_)); + PipeBarrier(); + return; + } + } + LoadAsFloatRow(gk_, KVOffset(b, hv, token, 0, K_), dst, K_); + } + + __aicore__ inline void PrefetchQKGate(uint64_t slot, uint64_t b, uint64_t h, uint64_t hv, uint64_t token, + uint64_t elems) + { + const uint64_t rows = elems / K_; + LocalTensor qTyped = GateQTyped(slot); + LocalTensor kTyped = GateKTyped(slot); + LocalTensor gateTyped = GateGTyped(slot); + CopyRowsIn(qTyped, q_, QOffset(b, h, token, 0), rows, K_, inputSequenceMajor_ ? H_ * K_ : K_); + CopyRowsIn(kTyped, k_, QOffset(b, h, token, 0), rows, K_, inputSequenceMajor_ ? H_ * K_ : K_); + if (!computeGateInPrepare_) { + CopyVectorIn(gateTyped, gk_, KVOffset(b, hv, token, 0, K_), elems); + } +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if constexpr (SAFE_GATE && COMPILE_BT == 64 && COMPILE_K == 128 && COMPILE_V == 128) { + LocalTensor vTyped = GateVTyped(slot); + LocalTensor betaTyped = GateBetaTyped(slot); + CopyRowsIn(vTyped, v_, VInputOffset(b, hv, token, 0), rows, V_, inputSequenceMajor_ ? HV_ * V_ : V_); + CopyVectorIn(betaTyped, beta_, BetaOffset(b, hv, token), rows); + } +#endif + SetFlag(mte2ToVEvent_); + } + + __aicore__ inline float LoadGateExpA(uint64_t hv) + { + if (!hasALog_) { + return 1.0f; + } + LocalTensor scalar = exp2Buf_.Get(); + CopyVectorIn(scalar, aLog_, hv, 1); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + Exp(scalar, scalar, 1); + PipeBarrier(); + SetFlag(vToSEvent_); + WaitFlag(vToSEvent_); + __ubuf__ float *ptr = (__ubuf__ float *)scalar.GetPhyAddr(); + return ptr[0]; + } + + __aicore__ inline void PrefetchRawGateTile(uint64_t slot, uint64_t b, uint64_t hv, uint64_t token, uint64_t rows) + { + (void)slot; + const uint64_t tileRow = token - activeGateChunkStart_; + LocalTensor gate = LocalGateChunk().template ReinterpretCast()[tileRow * K_]; + CopyRowsIn(gate, rawG_, RawGateOffset(b, hv, token, 0), rows, K_, inputSequenceMajor_ ? HV_ * K_ : K_); + SetFlag(mte2ToVEvent_); + } + + __aicore__ inline void MaterializeRawGateChunkArch35(uint64_t b, uint64_t hv, uint64_t start, uint64_t rows) + { +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if (!computeGateInPrepare_) { + return; + } + if constexpr (!IsSameType::value) { + return; + } else { + activeGateChunkStart_ = start; + const float expA = LoadGateExpA(hv); + LocalTensor acc = exp2Buf_.Get(); + LocalTensor bias = exp2Buf_.Get()[K_]; + if (hasDtBias_) { + CopyVectorIn(bias, dtBias_, hv * K_, K_); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + } + Duplicate(acc, 0.0f, static_cast(K_)); + PipeBarrier(); + + const uint64_t tileRows = GatePipelineRows(); + const uint64_t tileCount = (rows + tileRows - 1) / tileRows; + uint64_t currentRows = rows < tileRows ? rows : tileRows; + PrefetchRawGateTile(0, b, hv, start, currentRows); + for (uint64_t tile = 0; tile < tileCount; ++tile) { + const uint64_t slot = tile & 1; + const uint64_t tileRow = tile * tileRows; + currentRows = rows - tileRow; + if (currentRows > tileRows) { + currentRows = tileRows; + } + WaitFlag(mte2ToVEvent_); + + const uint64_t nextTile = tile + 1; + if (nextTile < tileCount) { + const uint64_t nextSlot = nextTile & 1; + uint64_t nextRows = rows - nextTile * tileRows; + if (nextRows > tileRows) { + nextRows = tileRows; + } + PrefetchRawGateTile(nextSlot, b, hv, start + nextTile * tileRows, nextRows); + } + + LocalTensor gate = LocalGateChunk().template ReinterpretCast()[tileRow * K_]; + if (hasDtBias_) { + AccumulateRawSafeGateChunk128Regbase( + (__ubuf__ float *)gate.GetPhyAddr(), (__ubuf__ float *)bias.GetPhyAddr(), + (__ubuf__ float *)acc.GetPhyAddr(), static_cast(currentRows), expA, lowerBound_); + } else { + AccumulateRawSafeGateChunk128Regbase( + (__ubuf__ float *)gate.GetPhyAddr(), (__ubuf__ float *)bias.GetPhyAddr(), + (__ubuf__ float *)acc.GetPhyAddr(), static_cast(currentRows), expA, lowerBound_); + } + PipeBarrier(); + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + CopyVectorOut(gk_, KVOffset(b, hv, start + tileRow, 0, K_), gate, currentRows * K_); + } + SetFlag(mte3ToVEvent_); + WaitFlag(mte3ToVEvent_); + } +#else + (void)b; + (void)hv; + (void)start; + (void)rows; +#endif + } + + __aicore__ inline LocalTensor GateDirectQ(uint64_t slot) + { + return vecBuf_.Get()[slot * 3 * GateStageElems()]; + } + + __aicore__ inline LocalTensor GateDirectW(uint64_t slot) { return GateDirectQ(slot)[GateStageElems()]; } + + __aicore__ inline LocalTensor GateDirectV(uint64_t slot) { return GateDirectQ(slot)[2 * GateStageElems()]; } + + __aicore__ inline LocalTensor GateBetaFloat(uint64_t slot) + { + constexpr uint64_t directBytes = KDA_GATE_PIPELINE_DEPTH * 3 * KDA_GATE_TILE_ROWS * COMPILE_K * sizeof(T); + return vecBuf_.Get()[directBytes / sizeof(float) + slot * KDA_GATE_TILE_ROWS]; + } + + __aicore__ inline void StorePreparedQG(uint64_t b, uint64_t hv, uint64_t token, LocalTensor directQ, + uint64_t elems) + { + const uint64_t offset = KVOffset(b, hv, token, 0, K_); + if (storeQG_) { + CopyVectorOut(qg_, offset, directQ, elems); + SetFlag(mte3ToVEvent_); + WaitFlag(mte3ToVEvent_); + } + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + CopyVectorOut(kg_, offset, directQ, elems); + } + + __aicore__ inline void PrefetchKGate(uint64_t slot, uint64_t b, uint64_t h, uint64_t hv, uint64_t token, + uint64_t elems) + { + const uint64_t rows = elems / K_; + LocalTensor kTyped = GateQTyped(slot); + LocalTensor gateTyped = GateGTyped(slot); + CopyRowsIn(kTyped, k_, QOffset(b, h, token, 0), rows, K_, inputSequenceMajor_ ? H_ * K_ : K_); + if (!computeGateInPrepare_) { + CopyVectorIn(gateTyped, gk_, KVOffset(b, hv, token, 0, K_), elems); + } + SetFlag(mte2ToVEvent_); + } + + __aicore__ inline void WaitGateInputReady() { WaitFlag(mte2ToVEvent_); } + + __aicore__ inline void WaitGateOutputForMte2(uint64_t slot = 0) + { + WaitFlag(mte3ToMte2Events_[slot]); + } + + __aicore__ inline void WaitGateOutputForVector() { WaitFlag(mte3ToVEvent_); } + + __aicore__ inline void SignalGateOutputDone() + { + SetFlag(mte3ToMte2Events_[0]); + SetFlag(mte3ToVEvent_); + } + + __aicore__ inline void SignalGateOutputDoneForMte2(uint64_t slot) + { + SetFlag(mte3ToMte2Events_[slot]); + } + + template + __aicore__ inline void LoadAsFloatRow(GlobalTensor &src, uint64_t srcOffset, LocalTensor &dst, + uint64_t count) + { + if constexpr (IsSameType::value) { + CopyVectorIn(dst, src, srcOffset, count); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + Adds(dst, dst, 0.0f, static_cast(count)); + PipeBarrier(); + SetFlag(vToMte2Event_); + WaitFlag(vToMte2Event_); + } else { + constexpr uint32_t typedOffset = EXP2_UB_ELEMENTS * sizeof(float) / sizeof(CopyT); + LocalTensor rowLocal = exp2Buf_.Get()[typedOffset]; + CopyVectorIn(rowLocal, src, srcOffset, count); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + Cast(dst, rowLocal, RoundMode::CAST_NONE, static_cast(count)); + PipeBarrier(); + SetFlag(vToMte2Event_); + WaitFlag(vToMte2Event_); + } + PipeBarrier(); + } + + template + __aicore__ inline void LoadAsFloatVector(GlobalTensor &src, uint64_t srcOffset, LocalTensor &dst, + LocalTensor &typedScratch, uint64_t count) + { + if constexpr (IsSameType::value) { + CopyVectorIn(dst, src, srcOffset, count); + } else { + CopyVectorIn(typedScratch, src, srcOffset, count); + } + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + if constexpr (!IsSameType::value) { + Cast(dst, typedScratch, RoundMode::CAST_NONE, static_cast(count)); + PipeBarrier(); + } + } + + template + __aicore__ inline void StoreFloatRow(GlobalTensor &dst, uint64_t dstOffset, LocalTensor &src, + uint64_t count) + { + if constexpr (IsSameType::value) { + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + CopyVectorOut(dst, dstOffset, src, count); + } else { + constexpr uint32_t typedOffset = EXP2_UB_ELEMENTS * sizeof(float) / sizeof(CopyT); + LocalTensor rowLocal = exp2Buf_.Get()[typedOffset]; + Cast(rowLocal, src, RoundMode::CAST_RINT, static_cast(count)); + PipeBarrier(); + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + CopyVectorOut(dst, dstOffset, rowLocal, count); + } + SetFlag(mte3ToMte2Events_[0]); + WaitFlag(mte3ToMte2Events_[0]); + SetFlag(mte3ToVEvent_); + WaitFlag(mte3ToVEvent_); + } + + __aicore__ inline LocalTensor Exp2NegG(uint64_t b, uint64_t hv, uint64_t t) + { + LocalTensor exp2Local = exp2Buf_.Get(); + LoadAsFloatRow(gk_, KVOffset(b, hv, t, 0, K_), exp2Local, K_); + Muls(exp2Local, exp2Local, -LN2, static_cast(K_)); + PipeBarrier(); + RunExp2(exp2Local, static_cast(K_)); + return exp2Local; + } + + __aicore__ inline void PrepareScoreFactorsBulk(uint64_t b, uint64_t h, uint64_t hv, uint64_t start, + uint64_t subBlockIdx, uint64_t subBlockNum, uint64_t refToken, + uint64_t scoreRowBegin, uint64_t scoreRowCount, uint64_t validColEnd, + uint64_t finalRefToken, uint64_t scoreSlot) + { + const bool useDirectScoreL1 = +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + KDA_ARCH35_ENABLE_DIRECT_SCORE_L1 && KDA_ARCH35_ENABLE_DIRECT_SCORE_UB && SAFE_GATE && BT_ == 64 && + K_ == 128 && V_ == 128 && subBlockNum == 1 && scoreRowCount == KDA_DIRECT_SCORE_ROWS && + finalRefToken == start + BT_ - 1; +#else + false; +#endif + LocalTensor refFp32 = exp2Buf_.Get(); + LoadGateScoreRef(refFp32, b, hv, refToken); +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + constexpr bool exportFinalKg = SAFE_GATE && COMPILE_BT == 64 && COMPILE_K == 128 && COMPILE_V == 128; + LocalTensor finalRefFp32 = exp2Buf_.Get()[K_]; + if constexpr (exportFinalKg) { + LoadGateScoreRef(finalRefFp32, b, hv, finalRefToken); + } +#endif + + uint64_t qwBegin = scoreRowBegin + (scoreRowCount * subBlockIdx) / subBlockNum; + uint64_t qwEnd = scoreRowBegin + (scoreRowCount * (subBlockIdx + 1)) / subBlockNum; + uint64_t qwMaxRows = GatePipelineRows(); + bool qwOutputPending = false; + uint64_t qwSlot = 0; + if (qwBegin < qwEnd && qwMaxRows > 0) { + uint64_t firstRows = qwEnd - qwBegin; + if (firstRows > qwMaxRows) { + firstRows = qwMaxRows; + } + PrefetchQKGate(qwSlot, b, h, hv, start + qwBegin, firstRows * K_); + } + for (uint64_t tileRow = qwBegin; tileRow < qwEnd && qwMaxRows > 0; tileRow += qwMaxRows) { + uint64_t tileRows = qwEnd - tileRow; + if (tileRows > qwMaxRows) { + tileRows = qwMaxRows; + } + uint64_t elems = tileRows * K_; + LocalTensor qTyped = GateQTyped(qwSlot); + LocalTensor kTyped = GateKTyped(qwSlot); + LocalTensor qScore = qTyped.template ReinterpretCast(); + LocalTensor kScore = kTyped.template ReinterpretCast(); + LocalTensor gateTyped = GateScoreTyped(qwSlot, tileRow); +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + LocalTensor kgScore = GateKgTyped(qwSlot).template ReinterpretCast(); + LocalTensor vTyped = GateVTyped(qwSlot); + LocalTensor betaTyped = GateBetaTyped(qwSlot); + LocalTensor directQ = GateDirectQ(qwSlot); + LocalTensor directW = GateDirectW(qwSlot); + LocalTensor directV = GateDirectV(qwSlot); + LocalTensor betaFp32 = GateBetaFloat(qwSlot); +#endif +#if !defined(__CCE_AICORE__) || __CCE_AICORE__ != 310 + LocalTensor arena = vecBuf_.Get(); + LocalTensor qFp32 = arena; + LocalTensor kFp32 = arena[elems]; + LocalTensor gFp32 = arena[2 * elems]; + LocalTensor expFp32 = arena[3 * elems]; + LocalTensor outFp32 = arena[4 * elems]; +#endif + + WaitGateInputReady(); +#if !defined(__CCE_AICORE__) || __CCE_AICORE__ != 310 + Cast(qFp32, qTyped, RoundMode::CAST_NONE, static_cast(elems)); + Cast(kFp32, kTyped, RoundMode::CAST_NONE, static_cast(elems)); + if constexpr (IsSameType::value) { + gFp32 = gateTyped; + } else { + Cast(gFp32, gateTyped, RoundMode::CAST_NONE, static_cast(elems)); + } +#endif +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if constexpr (IsSameType::value) { + Adds(betaFp32, betaTyped, 0.0f, static_cast(tileRows)); + } else { + Cast(betaFp32, betaTyped, RoundMode::CAST_NONE, static_cast(tileRows)); + } + PipeBarrier(); +#endif + if (qwOutputPending) { + WaitGateOutputForMte2(); + } + uint64_t nextTileRow = tileRow + qwMaxRows; + if (nextTileRow < qwEnd) { + uint64_t nextRows = qwEnd - nextTileRow; + if (nextRows > qwMaxRows) { + nextRows = qwMaxRows; + } + PrefetchQKGate(qwSlot ^ 1, b, h, hv, start + nextTileRow, nextRows * K_); + } +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + bool fuseQwKg = SAFE_GATE && BT_ == 64 && K_ == 128 && V_ == 128 && subBlockNum == 1; + if (fuseQwKg) { + PrepareKdaGateQwKgRegbase( + (__ubuf__ T *)reinterpret_cast(qTyped.GetPhyAddr()), + (__ubuf__ T *)reinterpret_cast(kTyped.GetPhyAddr()), + (__ubuf__ SCORE_T *)reinterpret_cast(qScore.GetPhyAddr()), + (__ubuf__ SCORE_T *)reinterpret_cast(kScore.GetPhyAddr()), + (__ubuf__ SCORE_T *)reinterpret_cast(kgScore.GetPhyAddr()), + (__ubuf__ T *)reinterpret_cast(directQ.GetPhyAddr()), + (__ubuf__ T *)reinterpret_cast(directW.GetPhyAddr()), + (__ubuf__ T *)reinterpret_cast(vTyped.GetPhyAddr()), + (__ubuf__ T *)reinterpret_cast(directV.GetPhyAddr()), + (__ubuf__ T *)reinterpret_cast(vTyped.GetPhyAddr()), + (__ubuf__ float *)reinterpret_cast(betaFp32.GetPhyAddr()), + (__ubuf__ GK_T *)reinterpret_cast(gateTyped.GetPhyAddr()), + (__ubuf__ float *)reinterpret_cast(refFp32.GetPhyAddr()), + (__ubuf__ float *)reinterpret_cast(finalRefFp32.GetPhyAddr()), + static_cast(tileRows), static_cast(K_), static_cast(tileRows)); + } else { + PrepareKdaGateQwRegbase( + (__ubuf__ T *)reinterpret_cast(qTyped.GetPhyAddr()), + (__ubuf__ T *)reinterpret_cast(kTyped.GetPhyAddr()), + (__ubuf__ SCORE_T *)reinterpret_cast(qScore.GetPhyAddr()), + (__ubuf__ SCORE_T *)reinterpret_cast(kScore.GetPhyAddr()), + (__ubuf__ T *)reinterpret_cast(directQ.GetPhyAddr()), + (__ubuf__ T *)reinterpret_cast(directW.GetPhyAddr()), + (__ubuf__ GK_T *)reinterpret_cast(gateTyped.GetPhyAddr()), + (__ubuf__ float *)reinterpret_cast(refFp32.GetPhyAddr()), static_cast(tileRows), + static_cast(K_)); + } +#else + PipeBarrier(); + for (uint64_t row = 0; row < tileRows; ++row) { + Sub(expFp32[row * K_], gFp32[row * K_], refFp32, static_cast(K_)); + } + PipeBarrier(); + Muls(expFp32, expFp32, LN2, static_cast(elems)); + PipeBarrier(); + ClampScoreExpInput(expFp32, static_cast(elems)); + Exp(expFp32, expFp32, static_cast(elems)); + PipeBarrier(); + + Mul(outFp32, qFp32, expFp32, static_cast(elems)); + PipeBarrier(); + ClampFp32ForCast(outFp32, static_cast(elems)); + Cast(qScore, outFp32, RoundMode::CAST_RINT, static_cast(elems)); + PipeBarrier(); + Mul(outFp32, kFp32, expFp32, static_cast(elems)); + PipeBarrier(); + ClampFp32ForCast(outFp32, static_cast(elems)); + Cast(kScore, outFp32, RoundMode::CAST_RINT, static_cast(elems)); + PipeBarrier(); +#endif + + if (qwOutputPending) { + WaitGateOutputForVector(); + } + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + if (useDirectScoreL1) { + Catlass::Arch::Resource resource; + LocalTensor scoreL1 = + resource.l1Buf.template GetBufferByByte(scoreSlot * KDA_DIRECT_SCORE_L1_SLOT_BYTES); + const uint64_t localRow = tileRow - scoreRowBegin; + constexpr uint32_t c0Elements = 16; + constexpr uint32_t c0Blocks = 128 / c0Elements; + constexpr uint32_t rowBlockElements = 16 * c0Elements; + LocalTensor nzScratch = resource.ubBuf.template GetBufferByByte( + KDA_DIRECT_SCORE_UB_BYTE_OFFSET + + (scoreSlot / KDA_SCORE_LANES) * KDA_DIRECT_SCORE_SLOT_ELEMENTS * sizeof(float) + + 2 * KDA_DIRECT_SCORE_MATRIX_ELEMENTS * sizeof(float)); + LocalTensor qNz = nzScratch; + LocalTensor kNz = qNz[tileRows * K_]; + constexpr uint8_t srcRepeatStride = 128 * sizeof(SCORE_T) / 32; + for (uint32_t colBlock = 0; colBlock < c0Blocks; ++colBlock) { + const uint64_t srcOffset = colBlock * c0Elements; + const uint64_t dstOffset = colBlock * tileRows * c0Elements; + Copy(qNz[dstOffset], qScore[srcOffset], c0Elements, static_cast(tileRows), + {1, 1, 1, srcRepeatStride}); + Copy(kNz[dstOffset], kScore[srcOffset], c0Elements, static_cast(tileRows), + {1, 1, 1, srcRepeatStride}); + } + PipeBarrier(); + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + const uint64_t qL1Offset = (localRow / 16) * rowBlockElements; + const uint64_t kL1Offset = ((KDA_DIRECT_SCORE_ROWS + localRow) / 16) * rowBlockElements; + DataCopyParams nzCopyParams{c0Blocks, static_cast(tileRows), 0, + static_cast(64 - tileRows)}; + DataCopy(scoreL1[qL1Offset], qNz, nzCopyParams); + DataCopy(scoreL1[kL1Offset], kNz, nzCopyParams); + } else { + CopyVectorOut(scoreWorkspace_, ScoreScratchOffset(scoreSlot, KDA_SCORE_SCRATCH_QG, tileRow), qScore, + elems); + CopyVectorOut(scoreWorkspace_, ScoreScratchOffset(scoreSlot, KDA_SCORE_SCRATCH_W, tileRow), kScore, + elems); + } +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if (fuseQwKg) { + CopyVectorOut(scoreWorkspace_, ScoreScratchOffset(scoreSlot, KDA_SCORE_SCRATCH_KG, tileRow), kgScore, + elems); + } + StorePreparedQG(b, hv, start + tileRow, directQ, elems); + CopyVectorOut(w_, KVOffset(b, hv, start + tileRow, 0, K_), directW, elems); + CopyVectorOut(vNew_, KVOffset(b, hv, start + tileRow, 0, V_), directV, tileRows * V_); + if constexpr (exportFinalKg) { + CopyVectorOut(finalKg_, KVOffset(b, hv, start + tileRow, 0, K_), vTyped, elems); + } +#endif + SignalGateOutputDone(); + qwOutputPending = true; + qwSlot ^= 1; + } + if (qwOutputPending) { + WaitGateOutputForMte2(); + WaitGateOutputForVector(); + } + + bool fuseQwKg = false; +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + fuseQwKg = SAFE_GATE && BT_ == 64 && K_ == 128 && V_ == 128 && subBlockNum == 1; +#endif + uint64_t kgRows = fuseQwKg ? scoreRowBegin : validColEnd; + uint64_t kgBegin = (kgRows * subBlockIdx) / subBlockNum; + uint64_t kgEnd = (kgRows * (subBlockIdx + 1)) / subBlockNum; + uint64_t kgMaxRows = GatePipelineRows(); + bool kgOutputPending = false; + uint64_t kgSlot = 0; + if (kgBegin < kgEnd && kgMaxRows > 0) { + uint64_t firstRows = kgEnd - kgBegin; + if (firstRows > kgMaxRows) { + firstRows = kgMaxRows; + } + PrefetchKGate(kgSlot, b, h, hv, start + kgBegin, firstRows * K_); + } + for (uint64_t tileRow = kgBegin; tileRow < kgEnd && kgMaxRows > 0; tileRow += kgMaxRows) { + uint64_t tileRows = kgEnd - tileRow; + if (tileRows > kgMaxRows) { + tileRows = kgMaxRows; + } + uint64_t elems = tileRows * K_; + LocalTensor kTyped = GateQTyped(kgSlot); + LocalTensor kgScore = kTyped.template ReinterpretCast(); + LocalTensor gateTyped = GateScoreTyped(kgSlot, tileRow); +#if !defined(__CCE_AICORE__) || __CCE_AICORE__ != 310 + LocalTensor arena = vecBuf_.Get(); + LocalTensor kFp32 = arena; + LocalTensor gFp32 = arena[elems]; + LocalTensor expFp32 = arena[2 * elems]; + LocalTensor outFp32 = arena[3 * elems]; +#endif + + WaitGateInputReady(); +#if !defined(__CCE_AICORE__) || __CCE_AICORE__ != 310 + Cast(kFp32, kTyped, RoundMode::CAST_NONE, static_cast(elems)); + if constexpr (IsSameType::value) { + gFp32 = gateTyped; + } else { + Cast(gFp32, gateTyped, RoundMode::CAST_NONE, static_cast(elems)); + } +#endif + if (kgOutputPending) { + WaitGateOutputForMte2(); + } + uint64_t nextTileRow = tileRow + kgMaxRows; + if (nextTileRow < kgEnd) { + uint64_t nextRows = kgEnd - nextTileRow; + if (nextRows > kgMaxRows) { + nextRows = kgMaxRows; + } + PrefetchKGate(kgSlot ^ 1, b, h, hv, start + nextTileRow, nextRows * K_); + } +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + PrepareKdaGateKgRegbase( + (__ubuf__ SCORE_T *)reinterpret_cast(kgScore.GetPhyAddr()), + (__ubuf__ T *)reinterpret_cast(kTyped.GetPhyAddr()), + (__ubuf__ GK_T *)reinterpret_cast(gateTyped.GetPhyAddr()), + (__ubuf__ float *)reinterpret_cast(refFp32.GetPhyAddr()), static_cast(tileRows), + static_cast(K_), static_cast(tileRows)); +#else + PipeBarrier(); + for (uint64_t row = 0; row < tileRows; ++row) { + Sub(expFp32[row * K_], refFp32, gFp32[row * K_], static_cast(K_)); + } + PipeBarrier(); + Muls(expFp32, expFp32, LN2, static_cast(elems)); + PipeBarrier(); + ClampScoreExpInput(expFp32, static_cast(elems)); + Exp(expFp32, expFp32, static_cast(elems)); + PipeBarrier(); + Mul(outFp32, kFp32, expFp32, static_cast(elems)); + PipeBarrier(); + ClampFp32ForCast(outFp32, static_cast(elems)); + Cast(kgScore, outFp32, RoundMode::CAST_RINT, static_cast(elems)); + PipeBarrier(); +#endif + + if (kgOutputPending) { + WaitGateOutputForVector(); + } + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + CopyVectorOut(scoreWorkspace_, ScoreScratchOffset(scoreSlot, KDA_SCORE_SCRATCH_KG, tileRow), kgScore, + elems); + SignalGateOutputDone(); + kgOutputPending = true; + kgSlot ^= 1; + } + if (kgOutputPending) { + WaitGateOutputForMte2(); + WaitGateOutputForVector(); + } + } + + __aicore__ inline void PrepareGateProductsBulk(uint64_t b, uint64_t h, uint64_t hv, uint64_t start, uint64_t curT, + uint64_t subBlockIdx, uint64_t subBlockNum, bool useRef, + uint64_t refToken, uint64_t validColEnd, bool writeScoreScratch, + uint64_t scoreSlot) + { + if constexpr (IsSameType::value) { + return; + } + if (subBlockNum == 0 || subBlockIdx >= subBlockNum || K_ == 0) { + return; + } + uint64_t rowBegin = (curT * subBlockIdx) / subBlockNum; + uint64_t rowEnd = (curT * (subBlockIdx + 1)) / subBlockNum; + if (rowBegin >= rowEnd) { + return; + } + + uint64_t maxRows = GatePipelineRows(); + if (maxRows == 0) { + return; + } + LocalTensor refFp32 = exp2Buf_.Get(); + if (useRef) { + LoadGateScoreRef(refFp32, b, hv, refToken); + } +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + const bool fuseScoreWriteback = writeScoreScratch && useRef && K_ * 2 <= EXP2_UB_ELEMENTS; +#endif + + bool outputPending = false; + uint64_t gateSlot = 0; + uint64_t firstRows = rowEnd - rowBegin; + if (firstRows > maxRows) { + firstRows = maxRows; + } + PrefetchQKGate(gateSlot, b, h, hv, start + rowBegin, firstRows * K_); + for (uint64_t tileRow = rowBegin; tileRow < rowEnd; tileRow += maxRows) { + uint64_t tileRows = rowEnd - tileRow; + if (tileRows > maxRows) { + tileRows = maxRows; + } + uint64_t elems = tileRows * K_; + LocalTensor qTyped = GateQTyped(gateSlot); + LocalTensor kTyped = GateKTyped(gateSlot); + LocalTensor kgTyped = GateKgTyped(gateSlot); + LocalTensor qScore = qTyped.template ReinterpretCast(); + LocalTensor wScore = kTyped.template ReinterpretCast(); + LocalTensor kgScore = kgTyped.template ReinterpretCast(); + LocalTensor gateTyped = GateScoreTyped(gateSlot, tileRow); +#if !defined(__CCE_AICORE__) || __CCE_AICORE__ != 310 + LocalTensor arena = vecBuf_.Get(); + LocalTensor qFp32 = arena; + LocalTensor kFp32 = arena[elems]; + LocalTensor gFp32 = arena[2 * elems]; + LocalTensor expFp32 = arena[3 * elems]; + LocalTensor outFp32 = arena[4 * elems]; +#else + LocalTensor vTyped = GateVTyped(gateSlot); + LocalTensor betaTyped = GateBetaTyped(gateSlot); + LocalTensor directQ = GateDirectQ(gateSlot); + LocalTensor directW = GateDirectW(gateSlot); + LocalTensor directV = GateDirectV(gateSlot); + LocalTensor betaFp32 = GateBetaFloat(gateSlot); +#endif + + uint64_t token = start + tileRow; + WaitGateInputReady(); +#if !defined(__CCE_AICORE__) || __CCE_AICORE__ != 310 + Cast(qFp32, qTyped, RoundMode::CAST_NONE, static_cast(elems)); + Cast(kFp32, kTyped, RoundMode::CAST_NONE, static_cast(elems)); + if constexpr (IsSameType::value) { + gFp32 = gateTyped; + } else { + Cast(gFp32, gateTyped, RoundMode::CAST_NONE, static_cast(elems)); + } +#endif +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if constexpr (IsSameType::value) { + Adds(betaFp32, betaTyped, 0.0f, static_cast(tileRows)); + } else { + Cast(betaFp32, betaTyped, RoundMode::CAST_NONE, static_cast(tileRows)); + } + PipeBarrier(); +#endif + uint64_t nextTileRow = tileRow + maxRows; +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + uint64_t nextGateSlot = (gateSlot + 1) % KDA_GATE_PIPELINE_DEPTH; + if (nextTileRow < rowEnd) { + uint64_t tileIndex = (tileRow - rowBegin) / maxRows; + if (tileIndex + 1 >= KDA_GATE_PIPELINE_DEPTH) { + WaitGateOutputForMte2(nextGateSlot); + } + uint64_t nextRows = rowEnd - nextTileRow; + if (nextRows > maxRows) { + nextRows = maxRows; + } + PrefetchQKGate(nextGateSlot, b, h, hv, start + nextTileRow, nextRows * K_); + } +#else + if (outputPending) { + WaitGateOutputForMte2(); + } + if (nextTileRow < rowEnd) { + uint64_t nextRows = rowEnd - nextTileRow; + if (nextRows > maxRows) { + nextRows = maxRows; + } + PrefetchQKGate(gateSlot ^ 1, b, h, hv, start + nextTileRow, nextRows * K_); + } +#endif +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + uint16_t validRows = static_cast(tileRows); + if (useRef && tileRow >= validColEnd) { + validRows = 0; + } else if (useRef && tileRow + tileRows > validColEnd) { + validRows = static_cast(validColEnd - tileRow); + } + if (writeScoreScratch) { + if (fuseScoreWriteback) { + PrepareKdaGateQwKgRegbase( + (__ubuf__ T *)reinterpret_cast(qTyped.GetPhyAddr()), + (__ubuf__ T *)reinterpret_cast(kTyped.GetPhyAddr()), + (__ubuf__ SCORE_T *)reinterpret_cast(qScore.GetPhyAddr()), + (__ubuf__ SCORE_T *)reinterpret_cast(wScore.GetPhyAddr()), + (__ubuf__ SCORE_T *)reinterpret_cast(kgScore.GetPhyAddr()), + (__ubuf__ T *)reinterpret_cast(directQ.GetPhyAddr()), + (__ubuf__ T *)reinterpret_cast(directW.GetPhyAddr()), + (__ubuf__ T *)reinterpret_cast(vTyped.GetPhyAddr()), + (__ubuf__ T *)reinterpret_cast(directV.GetPhyAddr()), nullptr, + (__ubuf__ float *)reinterpret_cast(betaFp32.GetPhyAddr()), + (__ubuf__ GK_T *)reinterpret_cast(gateTyped.GetPhyAddr()), + (__ubuf__ float *)reinterpret_cast(refFp32.GetPhyAddr()), nullptr, + static_cast(tileRows), static_cast(K_), validRows); + } else { + PrepareKdaGateQwKgRegbase( + (__ubuf__ T *)reinterpret_cast(qTyped.GetPhyAddr()), + (__ubuf__ T *)reinterpret_cast(kTyped.GetPhyAddr()), + (__ubuf__ SCORE_T *)reinterpret_cast(qScore.GetPhyAddr()), + (__ubuf__ SCORE_T *)reinterpret_cast(wScore.GetPhyAddr()), + (__ubuf__ SCORE_T *)reinterpret_cast(kgScore.GetPhyAddr()), nullptr, nullptr, nullptr, + nullptr, nullptr, nullptr, (__ubuf__ GK_T *)reinterpret_cast(gateTyped.GetPhyAddr()), + (__ubuf__ float *)reinterpret_cast(refFp32.GetPhyAddr()), nullptr, + static_cast(tileRows), static_cast(K_), validRows); + } + } else if (useRef) { + PrepareKdaGateQwKgRegbase( + (__ubuf__ T *)reinterpret_cast(qTyped.GetPhyAddr()), + (__ubuf__ T *)reinterpret_cast(kTyped.GetPhyAddr()), + (__ubuf__ T *)reinterpret_cast(qTyped.GetPhyAddr()), + (__ubuf__ T *)reinterpret_cast(kTyped.GetPhyAddr()), + (__ubuf__ T *)reinterpret_cast(kgTyped.GetPhyAddr()), nullptr, nullptr, nullptr, nullptr, + nullptr, nullptr, (__ubuf__ GK_T *)reinterpret_cast(gateTyped.GetPhyAddr()), + (__ubuf__ float *)reinterpret_cast(refFp32.GetPhyAddr()), nullptr, + static_cast(tileRows), static_cast(K_), validRows); + } else { + PrepareKdaGateQwKgRegbase( + (__ubuf__ T *)reinterpret_cast(qTyped.GetPhyAddr()), + (__ubuf__ T *)reinterpret_cast(kTyped.GetPhyAddr()), + (__ubuf__ T *)reinterpret_cast(qTyped.GetPhyAddr()), + (__ubuf__ T *)reinterpret_cast(kTyped.GetPhyAddr()), + (__ubuf__ T *)reinterpret_cast(kgTyped.GetPhyAddr()), nullptr, nullptr, nullptr, nullptr, + nullptr, nullptr, (__ubuf__ GK_T *)reinterpret_cast(gateTyped.GetPhyAddr()), + (__ubuf__ float *)reinterpret_cast(refFp32.GetPhyAddr()), nullptr, + static_cast(tileRows), static_cast(K_), validRows); + } +#else + PipeBarrier(); + + if (useRef) { + for (uint64_t row = 0; row < tileRows; ++row) { + Sub(expFp32[row * K_], gFp32[row * K_], refFp32, static_cast(K_)); + } + } else { + Adds(expFp32, gFp32, 0.0f, static_cast(elems)); + } + PipeBarrier(); + Muls(expFp32, expFp32, LN2, static_cast(elems)); + PipeBarrier(); + if (writeScoreScratch) { + ClampScoreExpInput(expFp32, static_cast(elems)); + } else { + ClampExpInput(expFp32, static_cast(elems)); + } + Exp(expFp32, expFp32, static_cast(elems)); + PipeBarrier(); + + Mul(outFp32, qFp32, expFp32, static_cast(elems)); + PipeBarrier(); + if (writeScoreScratch) { + ClampFp32ForCast(outFp32, static_cast(elems)); + Cast(qScore, outFp32, RoundMode::CAST_RINT, static_cast(elems)); + } else { + ClampFp32ToOutputType(outFp32, static_cast(elems)); + Cast(qTyped, outFp32, RoundMode::CAST_RINT, static_cast(elems)); + } + PipeBarrier(); + + Mul(outFp32, kFp32, expFp32, static_cast(elems)); + PipeBarrier(); + if (writeScoreScratch) { + ClampFp32ForCast(outFp32, static_cast(elems)); + Cast(wScore, outFp32, RoundMode::CAST_RINT, static_cast(elems)); + } else { + ClampFp32ToOutputType(outFp32, static_cast(elems)); + Cast(kTyped, outFp32, RoundMode::CAST_RINT, static_cast(elems)); + } + PipeBarrier(); + + if (useRef) { + for (uint64_t row = 0; row < tileRows; ++row) { + Sub(expFp32[row * K_], refFp32, gFp32[row * K_], static_cast(K_)); + } + } else { + Muls(expFp32, gFp32, -1.0f, static_cast(elems)); + } + PipeBarrier(); + Muls(expFp32, expFp32, LN2, static_cast(elems)); + PipeBarrier(); + if (writeScoreScratch) { + ClampScoreExpInput(expFp32, static_cast(elems)); + } else { + ClampExpInput(expFp32, static_cast(elems)); + } + Exp(expFp32, expFp32, static_cast(elems)); + PipeBarrier(); + Mul(outFp32, kFp32, expFp32, static_cast(elems)); + PipeBarrier(); + if (useRef && tileRow + tileRows > validColEnd) { + for (uint64_t row = 0; row < tileRows; ++row) { + if (tileRow + row >= validColEnd) { + Duplicate(outFp32[row * K_], 0.0f, static_cast(K_)); + } + } + PipeBarrier(); + } + if (writeScoreScratch) { + ClampFp32ForCast(outFp32, static_cast(elems)); + } else { + ClampFp32ToOutputType(outFp32, static_cast(elems)); + } + if (outputPending) { + WaitGateOutputForVector(); + } + if (writeScoreScratch) { + Cast(kgScore, outFp32, RoundMode::CAST_RINT, static_cast(elems)); + } else { + Cast(kgTyped, outFp32, RoundMode::CAST_RINT, static_cast(elems)); + } + PipeBarrier(); +#endif + + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + if (writeScoreScratch) { + CopyVectorOut(scoreWorkspace_, ScoreScratchOffset(scoreSlot, KDA_SCORE_SCRATCH_QG, tileRow), qScore, + elems); + CopyVectorOut(scoreWorkspace_, ScoreScratchOffset(scoreSlot, KDA_SCORE_SCRATCH_W, tileRow), wScore, + elems); + CopyVectorOut(scoreWorkspace_, ScoreScratchOffset(scoreSlot, KDA_SCORE_SCRATCH_KG, tileRow), kgScore, + elems); +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if (fuseScoreWriteback) { + StorePreparedQG(b, hv, token, directQ, elems); + CopyVectorOut(w_, KVOffset(b, hv, token, 0, K_), directW, elems); + CopyVectorOut(vNew_, KVOffset(b, hv, token, 0, V_), directV, tileRows * V_); + } +#endif + } else { + CopyVectorOut(qg_, KVOffset(b, hv, token, 0, K_), qTyped, elems); + CopyVectorOut(w_, KVOffset(b, hv, token, 0, K_), kTyped, elems); + CopyVectorOut(kg_, KVOffset(b, hv, token, 0, K_), kgTyped, elems); + } +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + SignalGateOutputDoneForMte2(gateSlot); +#else + SignalGateOutputDone(); +#endif + outputPending = true; +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + gateSlot = (gateSlot + 1) % KDA_GATE_PIPELINE_DEPTH; +#else + gateSlot ^= 1; +#endif + } +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + uint64_t tileCount = (rowEnd - rowBegin + maxRows - 1) / maxRows; + uint64_t firstPending = tileCount > KDA_GATE_PIPELINE_DEPTH ? tileCount - KDA_GATE_PIPELINE_DEPTH : 0; + for (uint64_t tile = firstPending; tile < tileCount; ++tile) { + WaitGateOutputForMte2(tile % KDA_GATE_PIPELINE_DEPTH); + } +#else + if (outputPending) { + WaitGateOutputForMte2(); + WaitGateOutputForVector(); + } +#endif + return; + } + + __aicore__ inline void ZeroScoreScratchRange(uint64_t scoreSlot, uint64_t planeBegin, uint64_t planeEnd, + uint64_t firstRow, uint64_t rowEnd) + { + if (firstRow >= rowEnd) { + return; + } + const uint64_t maxRows = GatePipelineRows(); + LocalTensor zeroLocal = GateQTyped(0).template ReinterpretCast(); + for (uint64_t row = firstRow; row < rowEnd; row += maxRows) { + uint64_t rows = rowEnd - row; + if (rows > maxRows) { + rows = maxRows; + } + const uint64_t elems = rows * K_; + Duplicate(zeroLocal, static_cast(0), static_cast(elems)); + PipeBarrier(); + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + for (uint64_t plane = planeBegin; plane < planeEnd; ++plane) { + CopyVectorOut(scoreWorkspace_, ScoreScratchOffset(scoreSlot, plane, row), zeroLocal, elems); + } + SignalGateOutputDone(); + WaitGateOutputForMte2(); + WaitGateOutputForVector(); + } + } + + __aicore__ inline void ZeroScoreScratchPadding(uint64_t scoreSlot, uint64_t scoreRowBegin, uint64_t scoreRowCount, + uint64_t validColEnd, uint64_t subBlockIdx, uint64_t subBlockNum) + { + if (subBlockNum == 0 || subBlockIdx + 1 != subBlockNum) { + return; + } + const uint64_t validRowEnd = scoreRowBegin + scoreRowCount; + const uint64_t paddedRowEnd = (validRowEnd + 15) / 16 * 16; + const uint64_t paddedColEnd = BT_; + ZeroScoreScratchRange(scoreSlot, KDA_SCORE_SCRATCH_QG, KDA_SCORE_SCRATCH_KG, validRowEnd, paddedRowEnd); + ZeroScoreScratchRange(scoreSlot, KDA_SCORE_SCRATCH_KG, KDA_SCORE_SCRATCH_PLANES, validColEnd, paddedColEnd); + } + + __aicore__ inline void PrepareGateProducts(uint64_t b, uint64_t h, uint64_t hv, uint64_t start, uint64_t curT, + uint64_t subBlockIdx, uint64_t subBlockNum, bool useRef = false, + uint64_t refToken = 0, uint64_t validColEnd = 0, + bool writeScoreScratch = false, uint64_t scoreSlot = 0, + uint64_t scoreRowBegin = 0, uint64_t scoreRowCount = 0) + { + if (subBlockNum == 0 || subBlockIdx >= subBlockNum) { + return; + } + if (validColEnd == 0 || validColEnd > curT) { + validColEnd = curT; + } +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if (writeScoreScratch && curT == BT_ && scoreRowBegin == 0 && scoreRowCount == curT && validColEnd == curT) { + PrepareGateProductsBulk(b, h, hv, start, curT, subBlockIdx, subBlockNum, useRef, refToken, validColEnd, + writeScoreScratch, scoreSlot); + ZeroScoreScratchPadding(scoreSlot, scoreRowBegin, scoreRowCount, validColEnd, subBlockIdx, subBlockNum); + return; + } +#endif + if (writeScoreScratch) { + PrepareScoreFactorsBulk(b, h, hv, start, subBlockIdx, subBlockNum, refToken, scoreRowBegin, scoreRowCount, + validColEnd, start + curT - 1, scoreSlot); + ZeroScoreScratchPadding(scoreSlot, scoreRowBegin, scoreRowCount, validColEnd, subBlockIdx, subBlockNum); + return; + } + PrepareGateProductsBulk(b, h, hv, start, curT, subBlockIdx, subBlockNum, useRef, refToken, validColEnd, + writeScoreScratch, scoreSlot); + } + + __aicore__ inline void ComputeRawAqkAkkCube(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, + uint64_t curT) + { + ComputeRawAqkAkkCubeBlock(b, hv, chunkIdx, start, curT, 0, curT); + } + +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + template + __aicore__ inline void ComputeRawAqkAkkCubeStableBlockDirectUbArch35(uint64_t b, uint64_t hv, uint64_t start, + uint64_t rowBegin, uint64_t scoreSlot, + uint8_t subBlockIdx, uint32_t directSlot) + { + using ElementA = SCORE_T; + using ElementB = SCORE_T; + using ElementC = float; + using LayoutTagA = Catlass::layout::RowMajor; + using LayoutTagB = Catlass::layout::ColumnMajor; + using LayoutTagC = Catlass::layout::RowMajor; + using TileCopy = Catlass::Gemm::Tile::PackedTileCopyTla; + using DirectTileCopy = Common::Tile::PackedTileCopyTlaToUB; + using LayoutTagL1A = typename TileCopy::LayoutTagL1A; + using LayoutTagL1B = typename TileCopy::LayoutTagL1B; + using LayoutTagL0A = typename TileCopy::LayoutTagL0A; + using LayoutTagL0B = typename TileCopy::LayoutTagL0B; + using CopyL1ToL0A = typename TileCopy::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopy::CopyL1ToL0B; + using TileMmad = Catlass::Gemm::Tile::TileMmadTla; + + constexpr uint32_t scoreRows = KDA_SAFE_SCORE_REF_BC; + constexpr uint32_t packedRows = scoreRows * 2; + constexpr uint32_t k = 128; + static_assert(scoreRows == 32 && packedRows == 64); + static_assert(N == 32 || N == 64); + + Catlass::Arch::Resource resource; + auto layoutA = tla::MakeLayout(BT_, K_); + auto layoutB = tla::MakeLayout(K_, BT_); + auto tensorQPos = tla::MakeTensor(scoreWorkspace_[ScoreScratchOffset(scoreSlot, KDA_SCORE_SCRATCH_QG)], layoutA, + Catlass::Arch::PositionGM{}); + auto tensorKPos = tla::MakeTensor(scoreWorkspace_[ScoreScratchOffset(scoreSlot, KDA_SCORE_SCRATCH_W)], layoutA, + Catlass::Arch::PositionGM{}); + auto tensorKNeg = tla::MakeTensor(scoreWorkspace_[ScoreScratchOffset(scoreSlot, KDA_SCORE_SCRATCH_KG)], layoutB, + Catlass::Arch::PositionGM{}); + + auto blockQPos = GetTile(tensorQPos, tla::MakeCoord(rowBegin, 0), tla::MakeShape(scoreRows, k)); + auto blockKPos = GetTile(tensorKPos, tla::MakeCoord(rowBegin, 0), tla::MakeShape(scoreRows, k)); + auto blockKNeg = GetTile(tensorKNeg, tla::MakeCoord(0, 0), tla::MakeShape(k, N)); + + using CopyGmToL1A = typename TileCopy::template CopyGmToL1A; + using CopyGmToL1B = typename TileCopy::template CopyGmToL1B; + + static_assert(sizeof(ElementA) == sizeof(uint16_t)); + LocalTensor l1A = + resource.l1Buf.template GetBufferByByte(scoreSlot * KDA_DIRECT_SCORE_L1_SLOT_BYTES); + LocalTensor l1B = resource.l1Buf.template GetBufferByByte(KDA_DIRECT_SCORE_L1_B_OFFSET); + LocalTensor l0A = resource.l0ABuf.template GetBufferByByte(0); + LocalTensor l0B = resource.l0BBuf.template GetBufferByByte(0); + LocalTensor l0C = resource.l0CBuf.template GetBufferByByte(0); + auto layoutL1A = tla::MakeLayout(packedRows, k); + auto layoutL1B = tla::MakeLayout(k, N); + auto layoutL0A = tla::MakeLayout(packedRows, k); + auto layoutL0B = tla::MakeLayout(k, N); + auto layoutL0C = tla::MakeLayoutL0C(packedRows, N); + auto tensorL1A = tla::MakeTensor(l1A, layoutL1A, Catlass::Arch::PositionL1{}); + auto tensorL1B = tla::MakeTensor(l1B, layoutL1B, Catlass::Arch::PositionL1{}); + auto tensorL0A = tla::MakeTensor(l0A, layoutL0A, Catlass::Arch::PositionL0A{}); + auto tensorL0B = tla::MakeTensor(l0B, layoutL0B, Catlass::Arch::PositionL0B{}); + auto tensorL0C = tla::MakeTensor(l0C, layoutL0C, Catlass::Arch::PositionL0C{}); + auto tileL1A = GetTile(tensorL1A, tla::MakeCoord(0, 0), tla::MakeShape(packedRows, k)); + auto tileL1AQ = GetTile(tensorL1A, tla::MakeCoord(0, 0), tla::MakeShape(scoreRows, k)); + auto tileL1AK = GetTile(tensorL1A, tla::MakeCoord(scoreRows, 0), tla::MakeShape(scoreRows, k)); + auto tileL1B = GetTile(tensorL1B, tla::MakeCoord(0, 0), tla::MakeShape(k, N)); + auto tileL0A = GetTile(tensorL0A, tla::MakeCoord(0, 0), tla::MakeShape(packedRows, k)); + auto tileL0B = GetTile(tensorL0B, tla::MakeCoord(0, 0), tla::MakeShape(k, N)); + auto tileL0C = GetTile(tensorL0C, tla::MakeCoord(0, 0), tla::MakeShape(packedRows, N)); + auto tileL0CTop = GetTile(tensorL0C, tla::MakeCoord(0, 0), tla::MakeShape(scoreRows, N)); + auto tileL0CBottom = GetTile(tensorL0C, tla::MakeCoord(scoreRows, 0), tla::MakeShape(scoreRows, N)); + LocalTensor directBase = resource.ubBuf.template GetBufferByByte( + KDA_DIRECT_SCORE_UB_BYTE_OFFSET + directSlot * KDA_DIRECT_SCORE_SLOT_ELEMENTS * sizeof(ElementC)); + LocalTensor directAqk = directBase; + LocalTensor directAkk = directBase[KDA_DIRECT_SCORE_MATRIX_ELEMENTS]; + auto layoutDirect = tla::MakeLayout(scoreRows, BT_); + auto tensorDirectAqk = tla::MakeTensor(directAqk, layoutDirect, Catlass::Arch::PositionUB{}); + auto tensorDirectAkk = tla::MakeTensor(directAkk, layoutDirect, Catlass::Arch::PositionUB{}); + auto blockDirectAqk = GetTile(tensorDirectAqk, tla::MakeCoord(0, 0), tla::MakeShape(scoreRows, N)); + auto blockDirectAkk = GetTile(tensorDirectAkk, tla::MakeCoord(0, 0), tla::MakeShape(scoreRows, N)); + using CopyL0CToDirectUb = typename DirectTileCopy::template CopyL0CToDst; + + CopyGmToL1A copyGmToL1A; + CopyGmToL1B copyGmToL1B; + CopyL1ToL0A copyL1ToL0A; + CopyL1ToL0B copyL1ToL0B; + CopyL0CToDirectUb copyL0CToDirectUb; + TileMmad tileMmad; + + copyGmToL1A(tileL1AQ, blockQPos); + copyGmToL1A(tileL1AK, blockKPos); + copyGmToL1B(tensorL1B, blockKNeg); + SetFlag(KDA_ARCH35_SCORE_EVENT); + WaitFlag(KDA_ARCH35_SCORE_EVENT); + copyL1ToL0A(tileL0A, tileL1A); + copyL1ToL0B(tileL0B, tileL1B); + SetFlag(KDA_ARCH35_SCORE_EVENT); + WaitFlag(KDA_ARCH35_SCORE_EVENT); + tileMmad(tileL0C, tileL0A, tileL0B, packedRows, N, k, true, 0); + SetFlag(KDA_ARCH35_SCORE_EVENT); + WaitFlag(KDA_ARCH35_SCORE_EVENT); + const uint64_t flagOffset = directSlot + subBlockIdx * KDA_DIRECT_SCORE_SUBBLOCK_FLAG_STRIDE; + CrossCoreWaitFlag<0x4, PIPE_FIX>(KDA_DIRECT_SCORE_FREE_FLAG + flagOffset); + copyL0CToDirectUb(blockDirectAqk, tileL0CTop, KDA_DIRECT_SCORE_ROWS, subBlockIdx, 1, 0); + copyL0CToDirectUb(blockDirectAkk, tileL0CBottom, KDA_DIRECT_SCORE_ROWS, subBlockIdx, 1, 0); + CrossCoreSetFlag<0x4, PIPE_FIX>(KDA_DIRECT_SCORE_READY_FLAG + flagOffset); + SetFlag(KDA_ARCH35_SCORE_EVENT); + WaitFlag(KDA_ARCH35_SCORE_EVENT); + } + + template + __aicore__ inline void PrefetchRawAqkAkkHeadPairArch35(uint64_t rowBegin, uint64_t scoreSlotBase, + uint32_t l1BaseOffset, TEventID readyEvent) + { + using ElementA = SCORE_T; + using ElementB = SCORE_T; + using ElementC = float; + using LayoutTagA = Catlass::layout::RowMajor; + using LayoutTagB = Catlass::layout::ColumnMajor; + using LayoutTagC = Catlass::layout::RowMajor; + using TileCopy = Catlass::Gemm::Tile::PackedTileCopyTla; + using LayoutTagL1A = typename TileCopy::LayoutTagL1A; + using LayoutTagL1B = typename TileCopy::LayoutTagL1B; + + constexpr uint32_t scoreRows = KDA_SAFE_SCORE_REF_BC; + constexpr uint32_t packedRows = scoreRows * 2; + constexpr uint32_t n = N; + constexpr uint32_t k = 128; + constexpr uint32_t l1ABytes = packedRows * k * sizeof(ElementA); + constexpr uint32_t l1BBytes = k * n * sizeof(ElementB); + constexpr uint32_t l1LaneBytes = l1ABytes + l1BBytes; + static_assert(scoreRows == 32 && packedRows == 64); + static_assert(N == 32 || N == 64); + + Catlass::Arch::Resource resource; + auto layoutA = tla::MakeLayout(BT_, K_); + auto layoutB = tla::MakeLayout(K_, BT_); + auto layoutL1A = tla::MakeLayout(packedRows, k); + auto layoutL1B = tla::MakeLayout(k, n); + + for (uint64_t lane = 0; lane < KDA_SCORE_LANES; ++lane) { + const uint64_t scoreSlot = scoreSlotBase + lane; + LocalTensor l1A = resource.l1Buf.template GetBufferByByte( + l1BaseOffset + static_cast(lane) * l1LaneBytes); + LocalTensor l1B = resource.l1Buf.template GetBufferByByte( + l1BaseOffset + static_cast(lane) * l1LaneBytes + l1ABytes); + auto tensorL1A = tla::MakeTensor(l1A, layoutL1A, Catlass::Arch::PositionL1{}); + auto tensorL1B = tla::MakeTensor(l1B, layoutL1B, Catlass::Arch::PositionL1{}); + auto tileL1AQ = GetTile(tensorL1A, tla::MakeCoord(0, 0), tla::MakeShape(scoreRows, k)); + auto tileL1AK = GetTile(tensorL1A, tla::MakeCoord(scoreRows, 0), tla::MakeShape(scoreRows, k)); + + auto tensorQPos = tla::MakeTensor(scoreWorkspace_[ScoreScratchOffset(scoreSlot, KDA_SCORE_SCRATCH_QG)], + layoutA, Catlass::Arch::PositionGM{}); + auto tensorKPos = tla::MakeTensor(scoreWorkspace_[ScoreScratchOffset(scoreSlot, KDA_SCORE_SCRATCH_W)], + layoutA, Catlass::Arch::PositionGM{}); + auto tensorKNeg = tla::MakeTensor(scoreWorkspace_[ScoreScratchOffset(scoreSlot, KDA_SCORE_SCRATCH_KG)], + layoutB, Catlass::Arch::PositionGM{}); + auto blockQPos = GetTile(tensorQPos, tla::MakeCoord(rowBegin, 0), tla::MakeShape(scoreRows, k)); + auto blockKPos = GetTile(tensorKPos, tla::MakeCoord(rowBegin, 0), tla::MakeShape(scoreRows, k)); + auto blockKNeg = GetTile(tensorKNeg, tla::MakeCoord(0, 0), tla::MakeShape(k, n)); + using CopyGmToL1A = typename TileCopy::template CopyGmToL1A; + using CopyGmToL1B = typename TileCopy::template CopyGmToL1B; + CopyGmToL1A copyGmToL1A; + CopyGmToL1B copyGmToL1B; + copyGmToL1A(tileL1AQ, blockQPos); + copyGmToL1A(tileL1AK, blockKPos); + copyGmToL1B(tensorL1B, blockKNeg); + } + SetFlag(readyEvent); + } + + template + __aicore__ inline void ComputeRawAqkAkkCubeStableHeadPairDirectUbArch35(uint64_t rowBegin, uint64_t scoreSlotBase, + uint32_t directSlot) + { + using ElementA = SCORE_T; + using ElementB = SCORE_T; + using ElementC = float; + using LayoutTagA = Catlass::layout::RowMajor; + using LayoutTagB = Catlass::layout::ColumnMajor; + using LayoutTagC = Catlass::layout::RowMajor; + using TileCopy = Catlass::Gemm::Tile::PackedTileCopyTla; + using DirectTileCopy = Common::Tile::PackedTileCopyTlaToUB; + using LayoutTagL1A = typename TileCopy::LayoutTagL1A; + using LayoutTagL1B = typename TileCopy::LayoutTagL1B; + using LayoutTagL0A = typename TileCopy::LayoutTagL0A; + using LayoutTagL0B = typename TileCopy::LayoutTagL0B; + using CopyL1ToL0A = typename TileCopy::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopy::CopyL1ToL0B; + using TileMmad = Catlass::Gemm::Tile::TileMmadTla; + + constexpr uint32_t scoreRows = KDA_SAFE_SCORE_REF_BC; + constexpr uint32_t packedRows = scoreRows * 2; + constexpr uint32_t n = N; + constexpr uint32_t k = 128; + constexpr uint32_t l1ABytes = packedRows * k * sizeof(ElementA); + constexpr uint32_t l1BBytes = k * n * sizeof(ElementB); + constexpr uint32_t l1LaneBytes = l1ABytes + l1BBytes; + static_assert(scoreRows == 32 && packedRows == 64); + static_assert(N == 32 || N == 64); + + PrefetchRawAqkAkkHeadPairArch35(rowBegin, scoreSlotBase, 0, KDA_ARCH35_SCORE_EVENT); + WaitFlag(KDA_ARCH35_SCORE_EVENT); + + Catlass::Arch::Resource resource; + auto layoutL1A = tla::MakeLayout(packedRows, k); + auto layoutL1B = tla::MakeLayout(k, n); + + constexpr uint32_t l0ABytes = packedRows * k * sizeof(ElementA); + constexpr uint32_t l0BBytes = k * n * sizeof(ElementB); + constexpr uint32_t l0CBytes = packedRows * n * sizeof(ElementC); + LocalTensor l0A0 = resource.l0ABuf.template GetBufferByByte(0); + LocalTensor l0A1 = resource.l0ABuf.template GetBufferByByte(l0ABytes); + LocalTensor l0B0 = resource.l0BBuf.template GetBufferByByte(0); + LocalTensor l0B1 = resource.l0BBuf.template GetBufferByByte(l0BBytes); + LocalTensor l0C0 = resource.l0CBuf.template GetBufferByByte(0); + LocalTensor l0C1 = resource.l0CBuf.template GetBufferByByte(l0CBytes); + auto layoutL0A = tla::MakeLayout(packedRows, k); + auto layoutL0B = tla::MakeLayout(k, n); + auto layoutL0C = tla::MakeLayoutL0C(packedRows, n); + auto tensorL0A0 = tla::MakeTensor(l0A0, layoutL0A, Catlass::Arch::PositionL0A{}); + auto tensorL0A1 = tla::MakeTensor(l0A1, layoutL0A, Catlass::Arch::PositionL0A{}); + auto tensorL0B0 = tla::MakeTensor(l0B0, layoutL0B, Catlass::Arch::PositionL0B{}); + auto tensorL0B1 = tla::MakeTensor(l0B1, layoutL0B, Catlass::Arch::PositionL0B{}); + auto tensorL0C0 = tla::MakeTensor(l0C0, layoutL0C, Catlass::Arch::PositionL0C{}); + auto tensorL0C1 = tla::MakeTensor(l0C1, layoutL0C, Catlass::Arch::PositionL0C{}); + auto tileL0A0 = GetTile(tensorL0A0, tla::MakeCoord(0, 0), tla::MakeShape(packedRows, k)); + auto tileL0A1 = GetTile(tensorL0A1, tla::MakeCoord(0, 0), tla::MakeShape(packedRows, k)); + auto tileL0B0 = GetTile(tensorL0B0, tla::MakeCoord(0, 0), tla::MakeShape(k, n)); + auto tileL0B1 = GetTile(tensorL0B1, tla::MakeCoord(0, 0), tla::MakeShape(k, n)); + auto tileL0C0 = GetTile(tensorL0C0, tla::MakeCoord(0, 0), tla::MakeShape(packedRows, n)); + auto tileL0C1 = GetTile(tensorL0C1, tla::MakeCoord(0, 0), tla::MakeShape(packedRows, n)); + auto tileL0C0Top = GetTile(tensorL0C0, tla::MakeCoord(0, 0), tla::MakeShape(scoreRows, n)); + auto tileL0C0Bottom = GetTile(tensorL0C0, tla::MakeCoord(scoreRows, 0), tla::MakeShape(scoreRows, n)); + auto tileL0C1Top = GetTile(tensorL0C1, tla::MakeCoord(0, 0), tla::MakeShape(scoreRows, n)); + auto tileL0C1Bottom = GetTile(tensorL0C1, tla::MakeCoord(scoreRows, 0), tla::MakeShape(scoreRows, n)); + CopyL1ToL0A copyL1ToL0A; + CopyL1ToL0B copyL1ToL0B; + TileMmad tileMmad; + + LocalTensor l1A0 = resource.l1Buf.template GetBufferByByte(0); + LocalTensor l1A1 = resource.l1Buf.template GetBufferByByte(l1LaneBytes); + LocalTensor l1B0 = resource.l1Buf.template GetBufferByByte(l1ABytes); + LocalTensor l1B1 = resource.l1Buf.template GetBufferByByte(l1LaneBytes + l1ABytes); + auto tensorL1A0 = tla::MakeTensor(l1A0, layoutL1A, Catlass::Arch::PositionL1{}); + auto tensorL1A1 = tla::MakeTensor(l1A1, layoutL1A, Catlass::Arch::PositionL1{}); + auto tensorL1B0 = tla::MakeTensor(l1B0, layoutL1B, Catlass::Arch::PositionL1{}); + auto tensorL1B1 = tla::MakeTensor(l1B1, layoutL1B, Catlass::Arch::PositionL1{}); + auto tileL1A0 = GetTile(tensorL1A0, tla::MakeCoord(0, 0), tla::MakeShape(packedRows, k)); + auto tileL1A1 = GetTile(tensorL1A1, tla::MakeCoord(0, 0), tla::MakeShape(packedRows, k)); + auto tileL1B0 = GetTile(tensorL1B0, tla::MakeCoord(0, 0), tla::MakeShape(k, n)); + auto tileL1B1 = GetTile(tensorL1B1, tla::MakeCoord(0, 0), tla::MakeShape(k, n)); + + LocalTensor directBase = resource.ubBuf.template GetBufferByByte( + KDA_DIRECT_SCORE_UB_BYTE_OFFSET + directSlot * KDA_DIRECT_SCORE_SLOT_ELEMENTS * sizeof(ElementC)); + LocalTensor directAqk = directBase; + LocalTensor directAkk = directBase[KDA_DIRECT_SCORE_MATRIX_ELEMENTS]; + auto layoutDirect = tla::MakeLayout(scoreRows, BT_); + auto tensorDirectAqk = tla::MakeTensor(directAqk, layoutDirect, Catlass::Arch::PositionUB{}); + auto tensorDirectAkk = tla::MakeTensor(directAkk, layoutDirect, Catlass::Arch::PositionUB{}); + auto blockDirectAqk = GetTile(tensorDirectAqk, tla::MakeCoord(0, 0), tla::MakeShape(scoreRows, n)); + auto blockDirectAkk = GetTile(tensorDirectAkk, tla::MakeCoord(0, 0), tla::MakeShape(scoreRows, n)); + using CopyL0CToDirectUb = typename DirectTileCopy::template CopyL0CToDst; + CopyL0CToDirectUb copyL0CToDirectUb; + + copyL1ToL0A(tileL0A0, tileL1A0); + copyL1ToL0B(tileL0B0, tileL1B0); + SetFlag(KDA_ARCH35_SCORE_EVENT); + copyL1ToL0A(tileL0A1, tileL1A1); + copyL1ToL0B(tileL0B1, tileL1B1); + SetFlag(KDA_ARCH35_SCORE_W_EVENT); + + WaitFlag(KDA_ARCH35_SCORE_EVENT); + tileMmad(tileL0C0, tileL0A0, tileL0B0, packedRows, n, k, true, 0); + SetFlag(KDA_ARCH35_SCORE_EVENT); + WaitFlag(KDA_ARCH35_SCORE_EVENT); + const uint64_t flagOffset0 = directSlot; + CrossCoreWaitFlag<0x4, PIPE_FIX>(KDA_DIRECT_SCORE_FREE_FLAG + flagOffset0); + copyL0CToDirectUb(blockDirectAqk, tileL0C0Top, KDA_DIRECT_SCORE_ROWS, 0, 1, 0); + copyL0CToDirectUb(blockDirectAkk, tileL0C0Bottom, KDA_DIRECT_SCORE_ROWS, 0, 1, 0); + CrossCoreSetFlag<0x4, PIPE_FIX>(KDA_DIRECT_SCORE_READY_FLAG + flagOffset0); + SetFlag(KDA_ARCH35_SCORE_EVENT); + + WaitFlag(KDA_ARCH35_SCORE_W_EVENT); + tileMmad(tileL0C1, tileL0A1, tileL0B1, packedRows, n, k, true, 0); + SetFlag(KDA_ARCH35_SCORE_W_EVENT); + WaitFlag(KDA_ARCH35_SCORE_W_EVENT); + const uint64_t flagOffset1 = directSlot + KDA_DIRECT_SCORE_SUBBLOCK_FLAG_STRIDE; + CrossCoreWaitFlag<0x4, PIPE_FIX>(KDA_DIRECT_SCORE_FREE_FLAG + flagOffset1); + copyL0CToDirectUb(blockDirectAqk, tileL0C1Top, KDA_DIRECT_SCORE_ROWS, 1, 1, 0); + copyL0CToDirectUb(blockDirectAkk, tileL0C1Bottom, KDA_DIRECT_SCORE_ROWS, 1, 1, 0); + CrossCoreSetFlag<0x4, PIPE_FIX>(KDA_DIRECT_SCORE_READY_FLAG + flagOffset1); + SetFlag(KDA_ARCH35_SCORE_W_EVENT); + + WaitFlag(KDA_ARCH35_SCORE_EVENT); + WaitFlag(KDA_ARCH35_SCORE_W_EVENT); + } + + __aicore__ inline void ComputeRawAqkAkkCubeFullArch35(uint64_t b, uint64_t hv, uint64_t start, uint64_t scoreSlot) + { + using ElementA = SCORE_T; + using ElementB = SCORE_T; + using ElementC = float; + using LayoutTagA = Catlass::layout::RowMajor; + using LayoutTagB = Catlass::layout::ColumnMajor; + using LayoutTagC = Catlass::layout::RowMajor; + using TileCopy = Catlass::Gemm::Tile::PackedTileCopyTla; + using LayoutTagL1A = typename TileCopy::LayoutTagL1A; + using LayoutTagL1B = typename TileCopy::LayoutTagL1B; + using LayoutTagL0A = typename TileCopy::LayoutTagL0A; + using LayoutTagL0B = typename TileCopy::LayoutTagL0B; + using CopyL1ToL0A = typename TileCopy::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopy::CopyL1ToL0B; + using TileMmad = Catlass::Gemm::Tile::TileMmadTla; + + constexpr uint32_t m = 64; + constexpr uint32_t n = 64; + constexpr uint32_t k = 128; + Catlass::Arch::Resource resource; + auto layoutA = tla::MakeLayout(m, k); + auto layoutB = tla::MakeLayout(k, n); + auto layoutC = tla::MakeLayout(m, n); + auto tensorQPos = tla::MakeTensor(scoreWorkspace_[ScoreScratchOffset(scoreSlot, KDA_SCORE_SCRATCH_QG)], layoutA, + Catlass::Arch::PositionGM{}); + auto tensorKPos = tla::MakeTensor(scoreWorkspace_[ScoreScratchOffset(scoreSlot, KDA_SCORE_SCRATCH_W)], layoutA, + Catlass::Arch::PositionGM{}); + auto tensorKNeg = tla::MakeTensor(scoreWorkspace_[ScoreScratchOffset(scoreSlot, KDA_SCORE_SCRATCH_KG)], layoutB, + Catlass::Arch::PositionGM{}); + auto tensorAqk = tla::MakeTensor(aqk_[AOffset(b, hv, start, 0)], layoutC, Catlass::Arch::PositionGM{}); + auto tensorAkk = tla::MakeTensor(akk_[AOffset(b, hv, start, 0)], layoutC, Catlass::Arch::PositionGM{}); + auto blockQPos = GetTile(tensorQPos, tla::MakeCoord(0, 0), tla::MakeShape(m, k)); + auto blockKPos = GetTile(tensorKPos, tla::MakeCoord(0, 0), tla::MakeShape(m, k)); + auto blockKNeg = GetTile(tensorKNeg, tla::MakeCoord(0, 0), tla::MakeShape(k, n)); + auto blockAqk = GetTile(tensorAqk, tla::MakeCoord(0, 0), tla::MakeShape(m, n)); + auto blockAkk = GetTile(tensorAkk, tla::MakeCoord(0, 0), tla::MakeShape(m, n)); + + using CopyGmToL1A = typename TileCopy::template CopyGmToL1A; + using CopyGmToL1B = typename TileCopy::template CopyGmToL1B; + using CopyL0CToDst = typename TileCopy::template CopyL0CToDst; + + constexpr uint32_t l1ABytes = m * k * sizeof(ElementA); + LocalTensor l1A0 = resource.l1Buf.template GetBufferByByte(0); + LocalTensor l1A1 = resource.l1Buf.template GetBufferByByte(l1ABytes); + LocalTensor l1B = resource.l1Buf.template GetBufferByByte(2 * l1ABytes); + LocalTensor l0A = resource.l0ABuf.template GetBufferByByte(0); + LocalTensor l0B = resource.l0BBuf.template GetBufferByByte(0); + LocalTensor l0C = resource.l0CBuf.template GetBufferByByte(0); + auto layoutL1A = tla::MakeLayout(m, k); + auto layoutL1B = tla::MakeLayout(k, n); + auto layoutL0A = tla::MakeLayout(m, k); + auto layoutL0B = tla::MakeLayout(k, n); + auto layoutL0C = tla::MakeLayoutL0C(m, n); + auto tensorL1A0 = tla::MakeTensor(l1A0, layoutL1A, Catlass::Arch::PositionL1{}); + auto tensorL1A1 = tla::MakeTensor(l1A1, layoutL1A, Catlass::Arch::PositionL1{}); + auto tensorL1B = tla::MakeTensor(l1B, layoutL1B, Catlass::Arch::PositionL1{}); + auto tensorL0A = tla::MakeTensor(l0A, layoutL0A, Catlass::Arch::PositionL0A{}); + auto tensorL0B = tla::MakeTensor(l0B, layoutL0B, Catlass::Arch::PositionL0B{}); + auto tensorL0C = tla::MakeTensor(l0C, layoutL0C, Catlass::Arch::PositionL0C{}); + auto tileL1A0 = GetTile(tensorL1A0, tla::MakeCoord(0, 0), tla::MakeShape(m, k)); + auto tileL1A1 = GetTile(tensorL1A1, tla::MakeCoord(0, 0), tla::MakeShape(m, k)); + auto tileL1B = GetTile(tensorL1B, tla::MakeCoord(0, 0), tla::MakeShape(k, n)); + auto tileL0A = GetTile(tensorL0A, tla::MakeCoord(0, 0), tla::MakeShape(m, k)); + auto tileL0B = GetTile(tensorL0B, tla::MakeCoord(0, 0), tla::MakeShape(k, n)); + auto tileL0C = GetTile(tensorL0C, tla::MakeCoord(0, 0), tla::MakeShape(m, n)); + + CopyGmToL1A copyGmToL1A; + CopyGmToL1B copyGmToL1B; + CopyL1ToL0A copyL1ToL0A; + CopyL1ToL0B copyL1ToL0B; + CopyL0CToDst copyL0CToDst; + TileMmad tileMmad; + + copyGmToL1B(tensorL1B, blockKNeg); + copyGmToL1A(tensorL1A0, blockQPos); + SetFlag(KDA_ARCH35_SCORE_EVENT); + WaitFlag(KDA_ARCH35_SCORE_EVENT); + copyL1ToL0B(tileL0B, tileL1B); + copyL1ToL0A(tileL0A, tileL1A0); + SetFlag(KDA_ARCH35_SCORE_EVENT); + copyGmToL1A(tensorL1A1, blockKPos); + SetFlag(KDA_ARCH35_SCORE_W_EVENT); + WaitFlag(KDA_ARCH35_SCORE_EVENT); + tileMmad(tileL0C, tileL0A, tileL0B, m, n, k, true, 0); + SetFlag(KDA_ARCH35_SCORE_EVENT); + SetFlag(KDA_ARCH35_SCORE_EVENT); + WaitFlag(KDA_ARCH35_SCORE_W_EVENT); + WaitFlag(KDA_ARCH35_SCORE_EVENT); + WaitFlag(KDA_ARCH35_SCORE_EVENT); + copyL0CToDst(blockAqk, tensorL0C); + SetFlag(KDA_ARCH35_SCORE_EVENT); + copyL1ToL0B(tileL0B, tileL1B); + copyL1ToL0A(tileL0A, tileL1A1); + SetFlag(KDA_ARCH35_SCORE_EVENT); + WaitFlag(KDA_ARCH35_SCORE_EVENT); + WaitFlag(KDA_ARCH35_SCORE_EVENT); + tileMmad(tileL0C, tileL0A, tileL0B, m, n, k, true, 0); + SetFlag(KDA_ARCH35_SCORE_EVENT); + WaitFlag(KDA_ARCH35_SCORE_EVENT); + copyL0CToDst(blockAkk, tensorL0C); + SetFlag(KDA_ARCH35_SCORE_EVENT); + WaitFlag(KDA_ARCH35_SCORE_EVENT); + } +#endif + + __aicore__ inline void ComputeRawAqkAkkCubeBlock(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, + uint64_t curT, uint64_t rowBegin, uint64_t rowCount, + bool readScoreScratch = false, uint64_t scoreSlot = 0, + uint64_t colCount = 0) + { + if (colCount == 0 || colCount > curT) { + colCount = curT; + } +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if constexpr (COMPILE_BT == 64 && COMPILE_K == 128 && COMPILE_V == 128) { + if (KDA_ARCH35_ENABLE_MANUAL_SCORE_PIPELINE && fusePostWu_ && rowBegin == 0 && rowCount == 64 && + colCount == 64) { + ComputeRawAqkAkkCubeFullArch35(b, hv, start, scoreSlot); + return; + } + } +#endif + using ElementA = SCORE_T; + using ElementB = SCORE_T; + using ElementC = float; + using LayoutTagA = Catlass::layout::RowMajor; + using LayoutTagB = Catlass::layout::ColumnMajor; + using LayoutTagC = Catlass::layout::RowMajor; + using TileCopy = Catlass::Gemm::Tile::PackedTileCopyTla; + using BlockMmad = Catlass::Gemm::Block::BlockMmadTla; + + Catlass::Arch::Resource resource; + BlockMmad blockMmad(resource); + auto layoutA = tla::MakeLayout(BT_, K_); + auto layoutB = tla::MakeLayout(K_, BT_); + auto layoutC = tla::MakeLayout(BT_, BT_); + const bool paddedTail = curT < BT_; + const uint64_t mmRowCount = paddedTail ? (rowCount + 15) / 16 * 16 : rowCount; + const uint64_t mmColCount = paddedTail ? BT_ : colCount; + Catlass::GemmCoord shape{static_cast(mmRowCount), static_cast(mmColCount), + static_cast(K_)}; + + (void)readScoreScratch; + auto tensorQPos = tla::MakeTensor(scoreWorkspace_[ScoreScratchOffset(scoreSlot, KDA_SCORE_SCRATCH_QG)], layoutA, + Catlass::Arch::PositionGM{}); + auto tensorKPos = tla::MakeTensor(scoreWorkspace_[ScoreScratchOffset(scoreSlot, KDA_SCORE_SCRATCH_W)], layoutA, + Catlass::Arch::PositionGM{}); + auto tensorKNeg = tla::MakeTensor(scoreWorkspace_[ScoreScratchOffset(scoreSlot, KDA_SCORE_SCRATCH_KG)], layoutB, + Catlass::Arch::PositionGM{}); + auto aqkBase = paddedTail ? solveWorkspace_[SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_RAW_AQK)] : + aqk_[AOffset(b, hv, start, 0)]; + auto akkBase = paddedTail ? solveWorkspace_[SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_RAW_AKK)] : + akk_[AOffset(b, hv, start, 0)]; + auto tensorAqk = tla::MakeTensor(aqkBase, layoutC, Catlass::Arch::PositionGM{}); + auto tensorAkk = tla::MakeTensor(akkBase, layoutC, Catlass::Arch::PositionGM{}); + + auto blockQPos = GetTile(tensorQPos, tla::MakeCoord(rowBegin, 0), tla::MakeShape(shape.m(), shape.k())); + auto blockKPos = GetTile(tensorKPos, tla::MakeCoord(rowBegin, 0), tla::MakeShape(shape.m(), shape.k())); + auto blockKNeg = GetTile(tensorKNeg, tla::MakeCoord(0, 0), tla::MakeShape(shape.k(), shape.n())); + auto blockAqk = GetTile(tensorAqk, tla::MakeCoord(rowBegin, 0), tla::MakeShape(shape.m(), shape.n())); + auto blockAkk = GetTile(tensorAkk, tla::MakeCoord(rowBegin, 0), tla::MakeShape(shape.m(), shape.n())); + + blockMmad.preSetFlags(); + blockMmad(blockQPos, blockKNeg, blockAqk, shape); + blockMmad(blockKPos, blockKNeg, blockAkk, shape); + blockMmad.finalWaitFlags(); + } + + __aicore__ inline bool UseAkkCubeSolve(uint64_t curT) const + { + return curT > 0 && curT <= BT_ && (BT_ == 64 || BT_ == 128) && K_ >= 16 && V_ >= 16 && V_ <= 256 && + K_ % 16 == 0 && V_ % 16 == 0; + } + + __aicore__ inline bool UsePostWuCube(uint64_t curT) const + { + return curT > 0 && curT <= BT_ && (BT_ == 64 || BT_ == 128) && K_ >= 16 && V_ >= 16 && V_ <= 256 && + K_ % 16 == 0 && V_ % 16 == 0; + } + + __aicore__ inline void CopyLocalFloat(LocalTensor dst, LocalTensor src, uint64_t count) + { + if (count == 0) { + return; + } + Adds(dst, src, 0.0f, static_cast(count)); + PipeBarrier(); + } + + __aicore__ inline void FillLocalFloat(LocalTensor dst, float value, uint64_t count) + { + if (count == 0) { + return; + } + Duplicate(dst, value, static_cast(count)); + PipeBarrier(); + } + + __aicore__ inline void ForwardSubDiag16(LocalTensor diag, LocalTensor row, LocalTensor prod, + LocalTensor rowBrcb, LocalTensor reduced, uint64_t valid) + { + constexpr uint32_t brcbStride = 8; + constexpr uint32_t diagSize = KDA_SOLVE_DIAG_BT; + constexpr uint8_t rowBlk = diagSize * sizeof(float) / 32; + +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + ForwardSubDiag16Regbase((__ubuf__ float *)reinterpret_cast(diag.GetPhyAddr()), + static_cast(valid)); +#else + for (uint64_t i = 2; i < valid; ++i) { + uint32_t rowOffset = static_cast(i * diagSize); + DataCopy(row, diag[rowOffset], diagSize); + PipeBarrier(); + + Brcb(rowBrcb, row, diagSize / brcbStride, {1, 8}); + PipeBarrier(); + for (uint32_t col = 0; col < diagSize; col += brcbStride) { + Mul(prod[col], diag[col], rowBrcb, brcbStride, static_cast(diagSize), + {1, 1, 0, rowBlk, rowBlk, 1}); + } + PipeBarrier(); + + uint32_t remain = diagSize; + while (remain > 1) { + uint32_t calcCount = (remain / 2) * diagSize; + remain = (remain + 1) / 2; + Add(prod, prod, prod[remain * diagSize], calcCount); + PipeBarrier(); + } + DataCopy(reduced, prod, diagSize); + PipeBarrier(); + Add(row, row, reduced, diagSize); + PipeBarrier(); + DataCopy(diag[rowOffset], row, diagSize); + PipeBarrier(); + } +#endif + + SetFlag(EXP2_EVENT_ID); + WaitFlag(EXP2_EVENT_ID); + for (uint32_t i = 0; i < diagSize; ++i) { + uint32_t diagOffset = i * diagSize + i; + if (i < valid) { + diag.SetValue(diagOffset, diag.GetValue(diagOffset) + 1.0f); + } else { + diag.SetValue(diagOffset, 1.0f); + } + } + SetFlag(EXP2_EVENT_ID); + WaitFlag(EXP2_EVENT_ID); + } + + __aicore__ inline void SolveDiagonalBlocksInRows(LocalTensor akkMat, LocalTensor xMat, + LocalTensor arena, uint64_t scratchBase, uint64_t curT, + uint64_t rowBegin, uint64_t rowCount) + { + constexpr uint32_t diagSize = KDA_SOLVE_DIAG_BT; + constexpr uint32_t diagElements = diagSize * diagSize; + constexpr uint32_t brcbElements = diagSize * 8; + +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + (void)akkMat; + (void)arena; + (void)scratchBase; + uint64_t rowEnd = rowBegin + rowCount; + for (uint64_t blockBegin = 0; blockBegin < BT_; blockBegin += diagSize) { + if (blockBegin < rowBegin || blockBegin + diagSize > rowEnd) { + continue; + } + uint64_t localBlockRow = blockBegin - rowBegin; + uint64_t valid = blockBegin < curT ? curT - blockBegin : 0; + if (valid > diagSize) { + valid = diagSize; + } + ForwardSubDiag16StridedRegbase((__ubuf__ float *)reinterpret_cast(xMat.GetPhyAddr()), + static_cast(BT_), static_cast(localBlockRow), + static_cast(blockBegin), static_cast(valid)); + SetFlag(EXP2_EVENT_ID); + WaitFlag(EXP2_EVENT_ID); + for (uint32_t rowIdx = 0; rowIdx < diagSize; ++rowIdx) { + uint32_t diagOffset = static_cast((localBlockRow + rowIdx) * BT_ + blockBegin + rowIdx); + if (rowIdx < valid) { + xMat.SetValue(diagOffset, xMat.GetValue(diagOffset) + 1.0f); + } else { + xMat.SetValue(diagOffset, 1.0f); + } + } + SetFlag(EXP2_EVENT_ID); + WaitFlag(EXP2_EVENT_ID); + } +#else + LocalTensor diag = arena[scratchBase]; + LocalTensor row = diag[diagElements]; + LocalTensor prod = row[diagSize]; + LocalTensor rowBrcb = prod[diagElements]; + LocalTensor reduced = rowBrcb[brcbElements]; + + uint64_t rowEnd = rowBegin + rowCount; + for (uint64_t blockBegin = 0; blockBegin < BT_; blockBegin += diagSize) { + if (blockBegin < rowBegin || blockBegin + diagSize > rowEnd) { + continue; + } + Duplicate(diag, 0.0f, diagElements); + PipeBarrier(); + + uint64_t localBlockRow = blockBegin - rowBegin; + uint64_t valid = blockBegin < curT ? curT - blockBegin : 0; + if (valid > diagSize) { + valid = diagSize; + } + for (uint32_t rowIdx = 0; rowIdx < diagSize; ++rowIdx) { + uint64_t srcOffset = (localBlockRow + rowIdx) * BT_ + blockBegin; + Muls(diag[rowIdx * diagSize], akkMat[srcOffset], -1.0f, diagSize); + } + PipeBarrier(); + + ForwardSubDiag16(diag, row, prod, rowBrcb, reduced, valid); + for (uint32_t rowIdx = 0; rowIdx < diagSize; ++rowIdx) { + uint64_t dstOffset = (localBlockRow + rowIdx) * BT_ + blockBegin; + Adds(xMat[dstOffset], diag[rowIdx * diagSize], 0.0f, diagSize); + } + PipeBarrier(); + } +#endif + } + + __aicore__ inline void BuildPrefixMask(LocalTensor dst, uint64_t prefix, uint64_t count) + { + if (prefix > count) { + prefix = count; + } + Duplicate(dst, 0.0f, static_cast(count)); + if (prefix > 0) { + Duplicate(dst, 1.0f, static_cast(prefix)); + } + PipeBarrier(); + } + + __aicore__ inline uint64_t BuildCausalMask(uint64_t threshold, uint64_t colBegin) const + { + if (threshold <= colBegin) { + return ~0ULL; + } + if (threshold >= colBegin + KDA_SOLVE_BT) { + return 0ULL; + } + return ~0ULL << (threshold - colBegin); + } + + __aicore__ inline void BuildCausalSelectMasks(LocalTensor aqkMask, LocalTensor akkMask, + uint64_t rowBegin, uint64_t rowCount, uint64_t colBegin) + { + __ubuf__ uint64_t *aqkMaskPtr = reinterpret_cast<__ubuf__ uint64_t *>(aqkMask.GetPhyAddr()); + __ubuf__ uint64_t *akkMaskPtr = reinterpret_cast<__ubuf__ uint64_t *>(akkMask.GetPhyAddr()); + for (uint32_t localRow = 0; localRow < rowCount; ++localRow) { + uint32_t row = static_cast(rowBegin + localRow); + aqkMaskPtr[localRow] = BuildCausalMask(static_cast(row) + 1, colBegin); + akkMaskPtr[localRow] = BuildCausalMask(static_cast(row), colBegin); + } + } + + __aicore__ inline void SelectCausalRows(LocalTensor aqkMat, LocalTensor akkMat, uint64_t rowBegin, + uint64_t rowCount) + { +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if constexpr (COMPILE_BT == 64 && COMPILE_K == 128 && COMPILE_V == 128) { + SelectCausalRows64Regbase((__ubuf__ float *)reinterpret_cast(aqkMat.GetPhyAddr()), + (__ubuf__ float *)reinterpret_cast(akkMat.GetPhyAddr()), + static_cast(rowBegin), static_cast(rowCount)); + PipeBarrier(); + return; + } +#endif + LocalTensor aqkMask = vecBuf_.Get()[KDA_SELECT_AQK_MASK_BYTE_OFFSET]; + LocalTensor akkMask = vecBuf_.Get()[KDA_SELECT_AKK_MASK_BYTE_OFFSET]; + LocalTensor zeroLocal = vecBuf_.Get()[KDA_SELECT_ZERO_FLOAT_OFFSET]; + Duplicate(zeroLocal, 0.0f, 8); + PipeBarrier(); + + uint64_t colBlockCount = (BT_ + KDA_SOLVE_BT - 1) / KDA_SOLVE_BT; + for (uint64_t colBlock = 0; colBlock < colBlockCount; ++colBlock) { + uint64_t maskOffset = colBlock * KDA_SELECT_COL_MASK_BYTES; + uint64_t colBegin = colBlock * KDA_SOLVE_BT; + BuildCausalSelectMasks(aqkMask[maskOffset], akkMask[maskOffset], rowBegin, rowCount, colBegin); + } + SetFlag(EXP2_EVENT_ID); + WaitFlag(EXP2_EVENT_ID); + + uint8_t rowStride = static_cast(BT_ * sizeof(float) / 32); + BinaryRepeatParams repeatParams = {1, 0, 1, rowStride, 0, rowStride}; + for (uint64_t colBlock = 0; colBlock < colBlockCount; ++colBlock) { + uint64_t maskOffset = colBlock * KDA_SELECT_COL_MASK_BYTES; + uint64_t colBegin = colBlock * KDA_SOLVE_BT; + Select(aqkMat[colBegin], aqkMask[maskOffset], zeroLocal, aqkMat[colBegin], SELMODE::VSEL_TENSOR_TENSOR_MODE, + KDA_SOLVE_BT, static_cast(rowCount), repeatParams); + Select(akkMat[colBegin], akkMask[maskOffset], zeroLocal, akkMat[colBegin], SELMODE::VSEL_TENSOR_TENSOR_MODE, + KDA_SOLVE_BT, static_cast(rowCount), repeatParams); + } + PipeBarrier(); + SetFlag(EXP2_EVENT_ID); + WaitFlag(EXP2_EVENT_ID); + } + + __aicore__ inline void PrepareAqkAkkSolveInput64(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start) + { + LocalTensor arena = vecBuf_.Get(); + LocalTensor aqkMat = arena; + LocalTensor akkMat = arena[KDA_SOLVE_MATRIX_ELEMENTS]; + LocalTensor xMat = arena[2 * KDA_SOLVE_MATRIX_ELEMENTS]; + LocalTensor betaLocal = arena[3 * KDA_SOLVE_MATRIX_ELEMENTS]; + LocalTensor betaBrcb = arena[3 * KDA_SOLVE_MATRIX_ELEMENTS + KDA_SOLVE_BT]; + LocalTensor maskLocal = arena[3 * KDA_SOLVE_MATRIX_ELEMENTS + KDA_SOLVE_BT + 512]; + LocalTensor oneHotLocal = arena[3 * KDA_SOLVE_MATRIX_ELEMENTS + KDA_SOLVE_BT + 512 + KDA_SOLVE_BT]; + + LoadAsFloatRow(beta_, BetaOffset(b, hv, start), betaLocal, KDA_SOLVE_BT); + Brcb(betaBrcb, betaLocal, 8, {1, 8}); + PipeBarrier(); + + DataCopy(aqkMat, aqk_[AOffset(b, hv, start, 0)], KDA_SOLVE_MATRIX_ELEMENTS); + DataCopy(akkMat, akk_[AOffset(b, hv, start, 0)], KDA_SOLVE_MATRIX_ELEMENTS); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + + for (uint64_t col = 0; col < KDA_SOLVE_BT; col += 8) { + Mul(akkMat[col], akkMat[col], betaBrcb, 8, KDA_SOLVE_BT, {1, 1, 1, 8, 8, 1}); + PipeBarrier(); + } + SelectCausalRows(aqkMat, akkMat, 0, KDA_SOLVE_BT); + + Muls(xMat, akkMat, -1.0f, KDA_SOLVE_MATRIX_ELEMENTS); + PipeBarrier(); + for (uint64_t row = 0; row < KDA_SOLVE_BT; ++row) { + BuildPrefixMask(maskLocal, row + 1, KDA_SOLVE_BT); + BuildPrefixMask(oneHotLocal, row, KDA_SOLVE_BT); + Sub(maskLocal, maskLocal, oneHotLocal, KDA_SOLVE_BT); + PipeBarrier(); + Add(xMat[row * KDA_SOLVE_BT], xMat[row * KDA_SOLVE_BT], maskLocal, KDA_SOLVE_BT); + PipeBarrier(); + } + + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + DataCopy(aqk_[AOffset(b, hv, start, 0)], aqkMat, KDA_SOLVE_MATRIX_ELEMENTS); + DataCopy(akk_[AOffset(b, hv, start, 0)], akkMat, KDA_SOLVE_MATRIX_ELEMENTS); + DataCopy(h_[SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_X)], xMat, KDA_SOLVE_MATRIX_ELEMENTS); + SetFlag(mte3ToMte2Events_[0]); + WaitFlag(mte3ToMte2Events_[0]); + SetFlag(mte3ToVEvent_); + WaitFlag(mte3ToVEvent_); + } + + __aicore__ inline void PrepareAqkAkkSolveInputTail(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, + uint64_t curT) + { + uint64_t elemCount = curT * KDA_SOLVE_BT; + LocalTensor arena = vecBuf_.Get(); + LocalTensor aqkMat = arena; + LocalTensor akkMat = arena[KDA_SOLVE_MATRIX_ELEMENTS]; + LocalTensor xMat = arena[2 * KDA_SOLVE_MATRIX_ELEMENTS]; + LocalTensor betaLocal = arena[3 * KDA_SOLVE_MATRIX_ELEMENTS]; + LocalTensor betaBrcb = arena[3 * KDA_SOLVE_MATRIX_ELEMENTS + KDA_SOLVE_BT]; + LocalTensor maskLocal = arena[3 * KDA_SOLVE_MATRIX_ELEMENTS + KDA_SOLVE_BT + 512]; + LocalTensor oneHotLocal = arena[3 * KDA_SOLVE_MATRIX_ELEMENTS + KDA_SOLVE_BT + 512 + KDA_SOLVE_BT]; + + FillLocalFloat(betaLocal, 0.0f, KDA_SOLVE_BT); + SetFlag(vToMte2Event_); + WaitFlag(vToMte2Event_); + LoadAsFloatRow(beta_, BetaOffset(b, hv, start), betaLocal, curT); + Brcb(betaBrcb, betaLocal, 8, {1, 8}); + PipeBarrier(); + + DataCopy(aqkMat, aqk_[AOffset(b, hv, start, 0)], static_cast(elemCount)); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + if (elemCount < KDA_SOLVE_MATRIX_ELEMENTS) { + FillLocalFloat(aqkMat[elemCount], 0.0f, KDA_SOLVE_MATRIX_ELEMENTS - elemCount); + } + DataCopy(akkMat, akk_[AOffset(b, hv, start, 0)], static_cast(elemCount)); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + if (elemCount < KDA_SOLVE_MATRIX_ELEMENTS) { + FillLocalFloat(akkMat[elemCount], 0.0f, KDA_SOLVE_MATRIX_ELEMENTS - elemCount); + } + + for (uint64_t col = 0; col < KDA_SOLVE_BT; col += 8) { + Mul(akkMat[col], akkMat[col], betaBrcb, 8, KDA_SOLVE_BT, {1, 1, 1, 8, 8, 1}); + PipeBarrier(); + } + SelectCausalRows(aqkMat, akkMat, 0, KDA_SOLVE_BT); + + Muls(xMat, akkMat, -1.0f, KDA_SOLVE_MATRIX_ELEMENTS); + PipeBarrier(); + for (uint64_t row = 0; row < KDA_SOLVE_BT; ++row) { + BuildPrefixMask(maskLocal, row + 1, KDA_SOLVE_BT); + BuildPrefixMask(oneHotLocal, row, KDA_SOLVE_BT); + Sub(maskLocal, maskLocal, oneHotLocal, KDA_SOLVE_BT); + PipeBarrier(); + Add(xMat[row * KDA_SOLVE_BT], xMat[row * KDA_SOLVE_BT], maskLocal, KDA_SOLVE_BT); + PipeBarrier(); + } + + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + DataCopy(aqk_[AOffset(b, hv, start, 0)], aqkMat, static_cast(elemCount)); + DataCopy(h_[SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_X)], xMat, KDA_SOLVE_MATRIX_ELEMENTS); + DataCopy(h_[SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_Y0)], akkMat, KDA_SOLVE_MATRIX_ELEMENTS); + SetFlag(mte3ToMte2Events_[0]); + WaitFlag(mte3ToMte2Events_[0]); + SetFlag(mte3ToVEvent_); + WaitFlag(mte3ToVEvent_); + } + + __aicore__ inline void GetSolveRowRange(uint64_t curT, uint64_t subBlockIdx, uint64_t subBlockNum, + uint64_t &rowBegin, uint64_t &rowEnd) const + { + if (subBlockNum == 0 || subBlockIdx >= subBlockNum) { + rowBegin = 0; + rowEnd = 0; + return; + } + rowBegin = (curT * subBlockIdx) / subBlockNum; + rowEnd = (curT * (subBlockIdx + 1)) / subBlockNum; + } + +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + __aicore__ inline void InitializeDirectScoreUbArch35() + { + for (uint32_t slot = 0; slot < KDA_DIRECT_SCORE_QUEUE_DEPTH; ++slot) { + CrossCoreSetFlag<0x4, PIPE_V>(KDA_DIRECT_SCORE_FREE_FLAG + slot); + } + } + + __aicore__ inline void ProcessDirectScoreSolveRowsArch35(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, + uint64_t rowBegin, uint32_t directSlot) + { + CrossCoreWaitFlag<0x4, PIPE_V>(KDA_DIRECT_SCORE_READY_FLAG + directSlot); + + Catlass::Arch::Resource resource; + LocalTensor arena = vecBuf_.Get(); + LocalTensor directBase = resource.ubBuf.template GetBufferByByte( + KDA_DIRECT_SCORE_UB_BYTE_OFFSET + directSlot * KDA_DIRECT_SCORE_SLOT_ELEMENTS * sizeof(float)); + LocalTensor aqkMat = directBase; + LocalTensor akkMat = directBase[KDA_DIRECT_SCORE_MATRIX_ELEMENTS]; + LocalTensor xMat = directBase[2 * KDA_DIRECT_SCORE_MATRIX_ELEMENTS]; + + SelectCausalRows64Regbase((__ubuf__ float *)reinterpret_cast(aqkMat.GetPhyAddr()), + (__ubuf__ float *)reinterpret_cast(akkMat.GetPhyAddr()), + static_cast(rowBegin), KDA_DIRECT_SCORE_ROWS); + PipeBarrier(); + Muls(xMat, akkMat, -1.0f, KDA_DIRECT_SCORE_MATRIX_ELEMENTS); + PipeBarrier(); + SolveDiagonalBlocksInRows(akkMat, xMat, arena, 0, BT_, rowBegin, KDA_DIRECT_SCORE_ROWS); + + LocalTensor aqkTyped = GateQTyped(0); + Muls(aqkMat, aqkMat, scale_, KDA_DIRECT_SCORE_MATRIX_ELEMENTS); + PipeBarrier(); + ClampFp32ToOutputType(aqkMat, KDA_DIRECT_SCORE_MATRIX_ELEMENTS); + Cast(aqkTyped, aqkMat, RoundMode::CAST_RINT, KDA_DIRECT_SCORE_MATRIX_ELEMENTS); + PipeBarrier(); + + const uint64_t token = start + rowBegin; + const uint64_t xBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_X) + rowBegin * BT_; + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + CopyVectorOut(o_, AOffset(b, hv, token, 0), aqkTyped, KDA_DIRECT_SCORE_MATRIX_ELEMENTS); + DataCopy(solveWorkspace_[xBase], xMat, KDA_DIRECT_SCORE_MATRIX_ELEMENTS); + SetFlag(mte3ToMte2Events_[0]); + WaitFlag(mte3ToMte2Events_[0]); + SetFlag(mte3ToVEvent_); + WaitFlag(mte3ToVEvent_); + CrossCoreSetFlag<0x4, PIPE_V>(KDA_DIRECT_SCORE_FREE_FLAG + directSlot); + } +#endif + + __aicore__ inline void PrepareAqkAkkSolveInputRows(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, + uint64_t curT, uint64_t rowBegin, uint64_t rowEnd, + bool storeLToAkk, bool storeLToScratch) + { + uint64_t rowCount = rowEnd - rowBegin; + if (rowCount == 0) { + return; + } + uint64_t validRowCount = rowBegin < curT ? curT - rowBegin : 0; + if (validRowCount > rowCount) { + validRowCount = rowCount; + } + uint64_t elemCount = rowCount * BT_; + uint64_t validElemCount = validRowCount * BT_; + LocalTensor arena = vecBuf_.Get(); + LocalTensor aqkMat = arena; + LocalTensor akkMat = arena[elemCount]; + LocalTensor xMat = arena[2 * elemCount]; + LocalTensor betaLocal = arena[3 * elemCount]; + LocalTensor betaBrcb = arena[3 * elemCount + BT_]; + LocalTensor maskLocal = arena[3 * elemCount + BT_ + 512]; + LocalTensor oneHotLocal = arena[3 * elemCount + BT_ + 512 + BT_]; + + uint64_t token = start + rowBegin; + + constexpr bool scoreCanIncludeBeta = +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + SAFE_GATE && COMPILE_BT == 64 && COMPILE_K == 128 && COMPILE_V == 128; +#else + false; +#endif + const bool scoreIncludesBeta = scoreCanIncludeBeta && HV_ % KDA_SCORE_LANES == 0; + if (validRowCount < rowCount) { + FillLocalFloat(aqkMat, 0.0f, elemCount); + FillLocalFloat(akkMat, 0.0f, elemCount); + if (!scoreIncludesBeta) { + FillLocalFloat(betaLocal, 0.0f, rowCount); + } + } + SetFlag(vToMte2Event_); + WaitFlag(vToMte2Event_); + if (validRowCount > 0) { + if (!scoreIncludesBeta) { + LoadAsFloatRow(beta_, BetaOffset(b, hv, token), betaLocal, validRowCount); + } + if (curT < BT_) { + DataCopy( + aqkMat, + solveWorkspace_[SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_RAW_AQK) + rowBegin * BT_], + static_cast(validElemCount)); + DataCopy( + akkMat, + solveWorkspace_[SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_RAW_AKK) + rowBegin * BT_], + static_cast(validElemCount)); + } else { + DataCopy(aqkMat, aqk_[AOffset(b, hv, token, 0)], static_cast(validElemCount)); + DataCopy(akkMat, akk_[AOffset(b, hv, token, 0)], static_cast(validElemCount)); + } + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + } +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if (!scoreIncludesBeta) { + ApplyKdaRowScaleRegbase((__ubuf__ float *)reinterpret_cast(akkMat.GetPhyAddr()), + (__ubuf__ float *)reinterpret_cast(betaLocal.GetPhyAddr()), + static_cast(rowCount), static_cast(BT_)); + PipeBarrier(); + } +#else + Brcb(betaBrcb, betaLocal, static_cast((rowCount + 7) / 8), {1, 8}); + PipeBarrier(); + uint8_t rowStride = static_cast(BT_ * sizeof(float) / 32); + for (uint64_t col = 0; col < BT_; col += 8) { + Mul(akkMat[col], akkMat[col], betaBrcb, 8, static_cast(rowCount), + {1, 1, 0, rowStride, rowStride, 1}); + } + PipeBarrier(); +#endif + if (validRowCount > 0) { + SelectCausalRows(aqkMat, akkMat, rowBegin, validRowCount); + } + + Muls(xMat, akkMat, -1.0f, static_cast(elemCount)); + PipeBarrier(); + if constexpr (SAFE_GATE) { + uint64_t scratchBase = 3 * elemCount + BT_ + 512 + 2 * BT_; + SolveDiagonalBlocksInRows(akkMat, xMat, arena, scratchBase, curT, rowBegin, rowCount); + } else if (curT < BT_) { + uint64_t scratchBase = 3 * elemCount + BT_ + 512 + 2 * BT_; + SolveDiagonalBlocksInRows(akkMat, xMat, arena, scratchBase, curT, rowBegin, rowCount); + } else { + for (uint64_t localRow = 0; localRow < rowCount; ++localRow) { + uint64_t row = rowBegin + localRow; + BuildPrefixMask(maskLocal, row + 1, BT_); + BuildPrefixMask(oneHotLocal, row, BT_); + Sub(maskLocal, maskLocal, oneHotLocal, static_cast(BT_)); + PipeBarrier(); + Add(xMat[localRow * BT_], xMat[localRow * BT_], maskLocal, static_cast(BT_)); + PipeBarrier(); + } + } + + uint64_t xBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_X) + rowBegin * BT_; + uint64_t lBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_Y0) + rowBegin * BT_; +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if constexpr (SAFE_GATE && COMPILE_BT == 64 && COMPILE_K == 128 && COMPILE_V == 128) { + LocalTensor aqkTyped = GateQTyped(0); + if (validElemCount > 0) { + Muls(aqkMat, aqkMat, scale_, static_cast(validElemCount)); + PipeBarrier(); + ClampFp32ToOutputType(aqkMat, static_cast(validElemCount)); + Cast(aqkTyped, aqkMat, RoundMode::CAST_RINT, static_cast(validElemCount)); + PipeBarrier(); + } + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + if (validElemCount > 0) { + CopyVectorOut(o_, AOffset(b, hv, token, 0), aqkTyped, validElemCount); + } + DataCopy(solveWorkspace_[xBase], xMat, static_cast(elemCount)); + if (storeLToScratch) { + DataCopy(solveWorkspace_[lBase], akkMat, static_cast(elemCount)); + } + SetFlag(mte3ToMte2Events_[0]); + WaitFlag(mte3ToMte2Events_[0]); + SetFlag(mte3ToVEvent_); + WaitFlag(mte3ToVEvent_); + return; + } +#endif + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + if (validRowCount > 0) { + DataCopy(aqk_[AOffset(b, hv, token, 0)], aqkMat, static_cast(validElemCount)); + if (storeLToAkk) { + DataCopy(akk_[AOffset(b, hv, token, 0)], akkMat, static_cast(validElemCount)); + } + } + DataCopy(solveWorkspace_[xBase], xMat, static_cast(elemCount)); + if (storeLToScratch) { + DataCopy(solveWorkspace_[lBase], akkMat, static_cast(elemCount)); + } + SetFlag(mte3ToMte2Events_[0]); + WaitFlag(mte3ToMte2Events_[0]); + SetFlag(mte3ToVEvent_); + WaitFlag(mte3ToVEvent_); + } + + __aicore__ inline void CubeGemmSolveSub(GlobalTensor &tensorA, uint64_t baseA, uint64_t rowA, uint64_t colA, + GlobalTensor &tensorB, uint64_t baseB, uint64_t rowB, uint64_t colB, + GlobalTensor &tensorC, uint64_t baseC, uint64_t rowC, uint64_t colC, + uint32_t m, uint32_t n, uint32_t k) + { + using ElementA = float; + using ElementB = float; + using ElementC = float; + using LayoutTagA = Catlass::layout::RowMajor; + using LayoutTagB = Catlass::layout::RowMajor; + using LayoutTagC = Catlass::layout::RowMajor; + using TileCopy = Catlass::Gemm::Tile::PackedTileCopyTla; + using BlockMmad = + Catlass::Gemm::Block::BlockMmadTla; + Catlass::Arch::Resource resource; + auto layoutA = tla::MakeLayout(BT_, BT_); + auto layoutB = tla::MakeLayout(BT_, BT_); + auto layoutC = tla::MakeLayout(BT_, BT_); + auto tensorLayoutA = tla::MakeTensor(tensorA[baseA], layoutA, Catlass::Arch::PositionGM{}); + auto tensorLayoutB = tla::MakeTensor(tensorB[baseB], layoutB, Catlass::Arch::PositionGM{}); + auto tensorLayoutC = tla::MakeTensor(tensorC[baseC], layoutC, Catlass::Arch::PositionGM{}); + Catlass::GemmCoord shape{m, n, k}; + auto blockA = GetTile(tensorLayoutA, tla::MakeCoord(rowA, colA), tla::MakeShape(shape.m(), shape.k())); + auto blockB = GetTile(tensorLayoutB, tla::MakeCoord(rowB, colB), tla::MakeShape(shape.k(), shape.n())); + auto blockC = GetTile(tensorLayoutC, tla::MakeCoord(rowC, colC), tla::MakeShape(shape.m(), shape.n())); + BlockMmad blockMmad(resource); + blockMmad(blockA, blockB, blockC, shape); +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + SetFlag(KDA_ARCH35_SOLVE_FIX_EVENT); + WaitFlag(KDA_ARCH35_SOLVE_FIX_EVENT); +#else + PipeBarrier(); +#endif + } + + __aicore__ inline void AddSolveTmpToX(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, bool storeAkk) + { + LocalTensor arena = vecBuf_.Get(); + LocalTensor xLocal = arena; + LocalTensor tmpLocal = arena[KDA_SOLVE_MATRIX_ELEMENTS]; + uint64_t xBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_X); + uint64_t tmpBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_TMP); + + DataCopy(xLocal, h_[xBase], KDA_SOLVE_MATRIX_ELEMENTS); + DataCopy(tmpLocal, h_[tmpBase], KDA_SOLVE_MATRIX_ELEMENTS); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + + Add(xLocal, xLocal, tmpLocal, KDA_SOLVE_MATRIX_ELEMENTS); + PipeBarrier(); + + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + DataCopy(h_[xBase], xLocal, KDA_SOLVE_MATRIX_ELEMENTS); + if (storeAkk) { + DataCopy(akk_[AOffset(b, hv, start, 0)], xLocal, KDA_SOLVE_MATRIX_ELEMENTS); + } + SetFlag(mte3ToMte2Events_[0]); + WaitFlag(mte3ToMte2Events_[0]); + SetFlag(mte3ToVEvent_); + WaitFlag(mte3ToVEvent_); + } + + __aicore__ inline void AddSolveTmpToXTail(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, uint64_t curT, + bool storeAkk) + { + uint64_t elemCount = curT * KDA_SOLVE_BT; + LocalTensor arena = vecBuf_.Get(); + LocalTensor xLocal = arena; + LocalTensor tmpLocal = arena[KDA_SOLVE_MATRIX_ELEMENTS]; + uint64_t xBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_X); + uint64_t tmpBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_TMP); + + DataCopy(xLocal, h_[xBase], KDA_SOLVE_MATRIX_ELEMENTS); + DataCopy(tmpLocal, h_[tmpBase], KDA_SOLVE_MATRIX_ELEMENTS); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + + Add(xLocal, xLocal, tmpLocal, KDA_SOLVE_MATRIX_ELEMENTS); + PipeBarrier(); + + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + DataCopy(h_[xBase], xLocal, KDA_SOLVE_MATRIX_ELEMENTS); + if (storeAkk) { + DataCopy(akk_[AOffset(b, hv, start, 0)], xLocal, static_cast(elemCount)); + } + SetFlag(mte3ToMte2Events_[0]); + WaitFlag(mte3ToMte2Events_[0]); + SetFlag(mte3ToVEvent_); + WaitFlag(mte3ToVEvent_); + } + + __aicore__ inline void AddSolveTmpToXRows(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, uint64_t curT, + uint64_t rowBegin, uint64_t rowEnd, bool storeAkk) + { + uint64_t rowCount = rowEnd - rowBegin; + if (rowCount == 0) { + return; + } + uint64_t validRowCount = rowBegin < curT ? curT - rowBegin : 0; + if (validRowCount > rowCount) { + validRowCount = rowCount; + } + uint64_t elemCount = rowCount * BT_; + uint64_t validElemCount = validRowCount * BT_; + LocalTensor arena = vecBuf_.Get(); + LocalTensor xLocal = arena; + LocalTensor tmpLocal = arena[elemCount]; + uint64_t xBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_X) + rowBegin * BT_; + uint64_t tmpBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_TMP) + rowBegin * BT_; + uint64_t token = start + rowBegin; + + DataCopy(xLocal, solveWorkspace_[xBase], static_cast(elemCount)); + DataCopy(tmpLocal, solveWorkspace_[tmpBase], static_cast(elemCount)); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + + Add(xLocal, xLocal, tmpLocal, static_cast(elemCount)); + PipeBarrier(); + + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + DataCopy(solveWorkspace_[xBase], xLocal, static_cast(elemCount)); + if (storeAkk && validRowCount > 0) { + DataCopy(akk_[AOffset(b, hv, token, 0)], xLocal, static_cast(validElemCount)); + } + SetFlag(mte3ToMte2Events_[0]); + WaitFlag(mte3ToMte2Events_[0]); + SetFlag(mte3ToVEvent_); + WaitFlag(mte3ToVEvent_); + } + + __aicore__ inline void AddSolveTmpToXDiagRows(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, + uint64_t rowBegin, uint64_t rowEnd, bool storeAkk) + { + uint64_t rowCount = rowEnd - rowBegin; + if (rowCount == 0) { + return; + } + uint64_t elemCount = rowCount * BT_; + LocalTensor arena = vecBuf_.Get(); + LocalTensor xLocal = arena; + LocalTensor tmpLocal = arena[elemCount]; + uint64_t xBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_X) + rowBegin * BT_; + uint64_t tmpBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_TMP) + rowBegin * BT_; + uint64_t token = start + rowBegin; + + DataCopy(xLocal, solveWorkspace_[xBase], static_cast(elemCount)); + DataCopy(tmpLocal, solveWorkspace_[tmpBase], static_cast(elemCount)); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + + for (uint64_t localRow = 0; localRow < rowCount; ++localRow) { + uint64_t row = rowBegin + localRow; + uint64_t col = (row / KDA_SOLVE_DIAG_BT) * KDA_SOLVE_DIAG_BT; + uint64_t offset = localRow * BT_ + col; + Add(xLocal[offset], xLocal[offset], tmpLocal[offset], KDA_SOLVE_DIAG_BT); + PipeBarrier(); + } + + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + DataCopy(solveWorkspace_[xBase], xLocal, static_cast(elemCount)); + if (storeAkk) { + DataCopy(akk_[AOffset(b, hv, token, 0)], xLocal, static_cast(elemCount)); + } + SetFlag(mte3ToMte2Events_[0]); + WaitFlag(mte3ToMte2Events_[0]); + SetFlag(mte3ToVEvent_); + WaitFlag(mte3ToVEvent_); + } + + __aicore__ inline void StoreSolveXRowsToAkk(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, + uint64_t curT, uint64_t rowBegin, uint64_t rowEnd) + { + uint64_t validRowCount = rowBegin < curT ? curT - rowBegin : 0; + uint64_t rowCount = rowEnd - rowBegin; + if (validRowCount > rowCount) { + validRowCount = rowCount; + } + if (validRowCount == 0) { + return; + } + uint64_t elemCount = validRowCount * BT_; + LocalTensor xLocal = vecBuf_.Get(); + uint64_t xBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_X) + rowBegin * BT_; + + DataCopy(xLocal, solveWorkspace_[xBase], static_cast(elemCount)); + SetFlag(mte2ToMte3Event_); + WaitFlag(mte2ToMte3Event_); + DataCopy(akk_[AOffset(b, hv, start + rowBegin, 0)], xLocal, static_cast(elemCount)); + SetFlag(mte3ToMte2Events_[0]); + WaitFlag(mte3ToMte2Events_[0]); + } + + __aicore__ inline void ComputeAkkMergeCube(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start) + { + uint64_t aiBase = AOffset(b, hv, start, 0); + uint64_t negABase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_X); + uint64_t tmpBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_TMP); + + for (uint32_t mergeSize = 2 * KDA_SOLVE_DIAG_BT; mergeSize <= BT_; mergeSize *= 2) { + uint32_t half = mergeSize / 2; + for (uint32_t block = 0; block < BT_; block += mergeSize) { + uint32_t lower = block + half; + CubeGemmSolveSub(akk_, aiBase, lower, lower, solveWorkspace_, negABase, lower, block, solveWorkspace_, + tmpBase, 0, 0, half, half, half); + CubeGemmSolveSub(solveWorkspace_, tmpBase, 0, 0, akk_, aiBase, block, block, akk_, aiBase, lower, block, + half, half, half); + } + } + } + + __aicore__ inline void ComputeAkkMergeCubeWorkspace(uint64_t b, uint64_t hv, uint64_t chunkIdx) + { + uint64_t xBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_X); + uint64_t tmpBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_TMP); + + for (uint32_t mergeSize = 2 * KDA_SOLVE_DIAG_BT; mergeSize <= BT_; mergeSize *= 2) { + uint32_t half = mergeSize / 2; + for (uint32_t block = 0; block < BT_; block += mergeSize) { + uint32_t lower = block + half; + CubeGemmSolveSub(solveWorkspace_, xBase, lower, lower, solveWorkspace_, xBase, lower, block, + solveWorkspace_, tmpBase, 0, 0, half, half, half); + CubeGemmSolveSub(solveWorkspace_, tmpBase, 0, 0, solveWorkspace_, xBase, block, block, solveWorkspace_, + xBase, lower, block, half, half, half); + } + } + } + +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + template + __aicore__ inline void LoadSolveTile(LocalTensor dst, LocalTensor src) + { + static_assert(TILE_SIZE == 16 || TILE_SIZE == 32, "arch35 solve tile must be 16 or 32 rows"); + constexpr uint32_t rowFractals = TILE_SIZE / 16; + constexpr uint32_t columnFractals = TILE_SIZE / 8; + LoadData2DParamsV2 loadParams; + loadParams.mStartPosition = 0; + loadParams.kStartPosition = 0; + loadParams.mStep = rowFractals; + loadParams.kStep = columnFractals; + loadParams.srcStride = rowFractals; + loadParams.dstStride = rowFractals; + loadParams.ifTranspose = TRANSPOSE; + LoadData(dst, src, loadParams); + } + + __aicore__ inline void ComputeAkkMergeCubeWorkspaceArch35(uint64_t b, uint64_t hv, uint64_t chunkIdx) + { + (void)b; + (void)hv; + SetMMLayoutTransform(true); + + using Element = float; + using LayoutTag = Catlass::layout::RowMajor; + using TileCopy = Catlass::Gemm::Tile::PackedTileCopyTla; + using LayoutTagL1A = typename TileCopy::LayoutTagL1A; + using LayoutTagL1B = typename TileCopy::LayoutTagL1B; + using LayoutTagL0A = typename TileCopy::LayoutTagL0A; + using LayoutTagL0B = typename TileCopy::LayoutTagL0B; + using TileMmad = Catlass::Gemm::Tile::TileMmadTla; + + constexpr uint32_t maxTile = 32; + constexpr uint32_t tileSlotBytes = maxTile * maxTile * sizeof(Element); + constexpr uint32_t a0Slot = 0; + constexpr uint32_t b0Slot = 1; + constexpr uint32_t a1Slot = 2; + constexpr uint32_t b1Slot = 3; + constexpr uint32_t diag0Slot = 4; + constexpr uint32_t diag1Slot = 5; + constexpr uint32_t tmp0Slot = 6; + constexpr uint32_t tmp1Slot = 7; + + Catlass::Arch::Resource resource; + LocalTensor l1A0 = resource.l1Buf.template GetBufferByByte(a0Slot * tileSlotBytes); + LocalTensor l1B0 = resource.l1Buf.template GetBufferByByte(b0Slot * tileSlotBytes); + LocalTensor l1A1 = resource.l1Buf.template GetBufferByByte(a1Slot * tileSlotBytes); + LocalTensor l1B1 = resource.l1Buf.template GetBufferByByte(b1Slot * tileSlotBytes); + LocalTensor l1Diag0 = resource.l1Buf.template GetBufferByByte(diag0Slot * tileSlotBytes); + LocalTensor l1Diag1 = resource.l1Buf.template GetBufferByByte(diag1Slot * tileSlotBytes); + LocalTensor l1Tmp0 = resource.l1Buf.template GetBufferByByte(tmp0Slot * tileSlotBytes); + LocalTensor l1Tmp1 = resource.l1Buf.template GetBufferByByte(tmp1Slot * tileSlotBytes); + + LocalTensor l0A0 = resource.l0ABuf.template GetBufferByByte(a0Slot * tileSlotBytes); + LocalTensor l0B0 = resource.l0BBuf.template GetBufferByByte(b0Slot * tileSlotBytes); + LocalTensor l0A1 = resource.l0ABuf.template GetBufferByByte(a1Slot * tileSlotBytes); + LocalTensor l0B1 = resource.l0BBuf.template GetBufferByByte(b1Slot * tileSlotBytes); + LocalTensor l0A2 = resource.l0ABuf.template GetBufferByByte(diag0Slot * tileSlotBytes); + LocalTensor l0B2 = resource.l0BBuf.template GetBufferByByte(diag0Slot * tileSlotBytes); + LocalTensor l0A3 = resource.l0ABuf.template GetBufferByByte(diag1Slot * tileSlotBytes); + LocalTensor l0B3 = resource.l0BBuf.template GetBufferByByte(diag1Slot * tileSlotBytes); + LocalTensor l0C0 = resource.l0CBuf.template GetBufferByByte(a0Slot * tileSlotBytes); + LocalTensor l0C1 = resource.l0CBuf.template GetBufferByByte(a1Slot * tileSlotBytes); + LocalTensor l0C2 = resource.l0CBuf.template GetBufferByByte(diag0Slot * tileSlotBytes); + LocalTensor l0C3 = resource.l0CBuf.template GetBufferByByte(diag1Slot * tileSlotBytes); + + TileMmad tileMmad; + const uint64_t xBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_X); + const uint64_t tmpBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_TMP); + auto gmLayout = tla::MakeLayout(BT_, BT_); + auto tensorX = tla::MakeTensor(solveWorkspace_[xBase], gmLayout, Catlass::Arch::PositionGM{}); + auto tensorTmp = tla::MakeTensor(solveWorkspace_[tmpBase], gmLayout, Catlass::Arch::PositionGM{}); + + constexpr uint32_t tile16 = 16; + auto shape16 = tla::MakeShape(tile16, tile16); + auto blockA0 = GetTile(tensorX, tla::MakeCoord(16, 16), shape16); + auto blockB0 = GetTile(tensorX, tla::MakeCoord(16, 0), shape16); + auto blockA1 = GetTile(tensorX, tla::MakeCoord(48, 48), shape16); + auto blockB1 = GetTile(tensorX, tla::MakeCoord(48, 32), shape16); + auto blockDiag0 = GetTile(tensorX, tla::MakeCoord(0, 0), shape16); + auto blockDiag1 = GetTile(tensorX, tla::MakeCoord(32, 32), shape16); + auto blockTmp0 = GetTile(tensorTmp, tla::MakeCoord(0, 0), shape16); + auto blockTmp1 = GetTile(tensorTmp, tla::MakeCoord(16, 0), shape16); + auto blockOut0 = GetTile(tensorX, tla::MakeCoord(16, 0), shape16); + auto blockOut1 = GetTile(tensorX, tla::MakeCoord(48, 32), shape16); + + using CopyGmToL1A16 = typename TileCopy::template CopyGmToL1A; + using CopyGmToL1B16 = typename TileCopy::template CopyGmToL1B; + using CopyL0CToDst16 = typename TileCopy::template CopyL0CToDst; + CopyGmToL1A16 copyGmToL1A16; + CopyGmToL1B16 copyGmToL1B16; + CopyL0CToDst16 copyL0CToDst16; + + auto tensorL1A0 = + tla::MakeTensor(l1A0, tla::MakeLayout(tile16, tile16), Catlass::Arch::PositionL1{}); + auto tensorL1B0 = + tla::MakeTensor(l1B0, tla::MakeLayout(tile16, tile16), Catlass::Arch::PositionL1{}); + auto tensorL1A1 = + tla::MakeTensor(l1A1, tla::MakeLayout(tile16, tile16), Catlass::Arch::PositionL1{}); + auto tensorL1B1 = + tla::MakeTensor(l1B1, tla::MakeLayout(tile16, tile16), Catlass::Arch::PositionL1{}); + auto tensorL1Diag0 = tla::MakeTensor(l1Diag0, tla::MakeLayout(tile16, tile16), + Catlass::Arch::PositionL1{}); + auto tensorL1Diag1 = tla::MakeTensor(l1Diag1, tla::MakeLayout(tile16, tile16), + Catlass::Arch::PositionL1{}); + auto tensorL1Tmp0 = tla::MakeTensor(l1Tmp0, tla::MakeLayout(tile16, tile16), + Catlass::Arch::PositionL1{}); + auto tensorL1Tmp1 = tla::MakeTensor(l1Tmp1, tla::MakeLayout(tile16, tile16), + Catlass::Arch::PositionL1{}); + + auto tensorL0A0 = + tla::MakeTensor(l0A0, tla::MakeLayout(tile16, tile16), Catlass::Arch::PositionL0A{}); + auto tensorL0B0 = + tla::MakeTensor(l0B0, tla::MakeLayout(tile16, tile16), Catlass::Arch::PositionL0B{}); + auto tensorL0A1 = + tla::MakeTensor(l0A1, tla::MakeLayout(tile16, tile16), Catlass::Arch::PositionL0A{}); + auto tensorL0B1 = + tla::MakeTensor(l0B1, tla::MakeLayout(tile16, tile16), Catlass::Arch::PositionL0B{}); + auto tensorL0A2 = + tla::MakeTensor(l0A2, tla::MakeLayout(tile16, tile16), Catlass::Arch::PositionL0A{}); + auto tensorL0B2 = + tla::MakeTensor(l0B2, tla::MakeLayout(tile16, tile16), Catlass::Arch::PositionL0B{}); + auto tensorL0A3 = + tla::MakeTensor(l0A3, tla::MakeLayout(tile16, tile16), Catlass::Arch::PositionL0A{}); + auto tensorL0B3 = + tla::MakeTensor(l0B3, tla::MakeLayout(tile16, tile16), Catlass::Arch::PositionL0B{}); + auto tensorL0C0 = tla::MakeTensor(l0C0, tla::MakeLayoutL0C(tile16, tile16), Catlass::Arch::PositionL0C{}); + auto tensorL0C1 = tla::MakeTensor(l0C1, tla::MakeLayoutL0C(tile16, tile16), Catlass::Arch::PositionL0C{}); + auto tensorL0C2 = tla::MakeTensor(l0C2, tla::MakeLayoutL0C(tile16, tile16), Catlass::Arch::PositionL0C{}); + auto tensorL0C3 = tla::MakeTensor(l0C3, tla::MakeLayoutL0C(tile16, tile16), Catlass::Arch::PositionL0C{}); + uint32_t localRow16 = 0; + uint32_t localColumn16 = 0; + auto tileL0A0 = GetTile(tensorL0A0, tla::MakeCoord(localRow16, localColumn16), shape16); + auto tileL0B0 = GetTile(tensorL0B0, tla::MakeCoord(localRow16, localColumn16), shape16); + auto tileL0A1 = GetTile(tensorL0A1, tla::MakeCoord(localRow16, localColumn16), shape16); + auto tileL0B1 = GetTile(tensorL0B1, tla::MakeCoord(localRow16, localColumn16), shape16); + auto tileL0A2 = GetTile(tensorL0A2, tla::MakeCoord(localRow16, localColumn16), shape16); + auto tileL0B2 = GetTile(tensorL0B2, tla::MakeCoord(localRow16, localColumn16), shape16); + auto tileL0A3 = GetTile(tensorL0A3, tla::MakeCoord(localRow16, localColumn16), shape16); + auto tileL0B3 = GetTile(tensorL0B3, tla::MakeCoord(localRow16, localColumn16), shape16); + auto tileL0C0 = GetTile(tensorL0C0, tla::MakeCoord(localRow16, localColumn16), shape16); + auto tileL0C1 = GetTile(tensorL0C1, tla::MakeCoord(localRow16, localColumn16), shape16); + auto tileL0C2 = GetTile(tensorL0C2, tla::MakeCoord(localRow16, localColumn16), shape16); + auto tileL0C3 = GetTile(tensorL0C3, tla::MakeCoord(localRow16, localColumn16), shape16); + + copyGmToL1A16(tensorL1A0, blockA0); + copyGmToL1B16(tensorL1B0, blockB0); + copyGmToL1A16(tensorL1A1, blockA1); + copyGmToL1B16(tensorL1B1, blockB1); + copyGmToL1B16(tensorL1Diag0, blockDiag0); + copyGmToL1B16(tensorL1Diag1, blockDiag1); + SetFlag(0); + WaitFlag(0); + + LoadSolveTile(l0A0, l1A0); + LoadSolveTile(l0B0, l1B0); + LoadSolveTile(l0A1, l1A1); + LoadSolveTile(l0B1, l1B1); + SetFlag(0); + WaitFlag(0); + tileMmad(tileL0C0, tileL0A0, tileL0B0, tile16, tile16, tile16, true, 0b11); + SetFlag(0); + tileMmad(tileL0C1, tileL0A1, tileL0B1, tile16, tile16, tile16, true, 0b11); + SetFlag(1); + WaitFlag(0); + copyL0CToDst16(blockTmp0, tileL0C0, 0b11); + SetFlag(0); + WaitFlag(1); + copyL0CToDst16(blockTmp1, tileL0C1, 0b11); + SetFlag(1); + WaitFlag(0); + copyGmToL1A16(tensorL1Tmp0, blockTmp0); + WaitFlag(1); + copyGmToL1A16(tensorL1Tmp1, blockTmp1); + SetFlag(0); + WaitFlag(0); + + LoadSolveTile(l0A2, l1Tmp0); + LoadSolveTile(l0B2, l1Diag0); + LoadSolveTile(l0A3, l1Tmp1); + LoadSolveTile(l0B3, l1Diag1); + SetFlag(0); + WaitFlag(0); + tileMmad(tileL0C2, tileL0A2, tileL0B2, tile16, tile16, tile16, true, 0b11); + SetFlag(2); + tileMmad(tileL0C3, tileL0A3, tileL0B3, tile16, tile16, tile16, true, 0b11); + SetFlag(3); + WaitFlag(2); + copyL0CToDst16(blockOut0, tileL0C2, 0b11); + SetFlag(2); + WaitFlag(3); + copyL0CToDst16(blockOut1, tileL0C3, 0b11); + SetFlag(3); + + constexpr uint32_t tile32 = 32; + auto shape32 = tla::MakeShape(tile32, tile32); + auto blockA32 = GetTile(tensorX, tla::MakeCoord(32, 32), shape32); + auto blockB32 = GetTile(tensorX, tla::MakeCoord(32, 0), shape32); + auto blockDiag32 = GetTile(tensorX, tla::MakeCoord(0, 0), shape32); + auto blockTmp32 = GetTile(tensorTmp, tla::MakeCoord(0, 0), shape32); + auto blockOut32 = GetTile(tensorX, tla::MakeCoord(32, 0), shape32); + using CopyGmToL1A32 = typename TileCopy::template CopyGmToL1A; + using CopyGmToL1B32 = typename TileCopy::template CopyGmToL1B; + using CopyL0CToDst32 = typename TileCopy::template CopyL0CToDst; + CopyGmToL1A32 copyGmToL1A32; + CopyGmToL1B32 copyGmToL1B32; + CopyL0CToDst32 copyL0CToDst32; + + auto tensorL1A32 = + tla::MakeTensor(l1A0, tla::MakeLayout(tile32, tile32), Catlass::Arch::PositionL1{}); + auto tensorL1B32 = + tla::MakeTensor(l1B0, tla::MakeLayout(tile32, tile32), Catlass::Arch::PositionL1{}); + auto tensorL1Diag32 = tla::MakeTensor(l1Diag0, tla::MakeLayout(tile32, tile32), + Catlass::Arch::PositionL1{}); + auto tensorL1Tmp32 = tla::MakeTensor(l1Tmp0, tla::MakeLayout(tile32, tile32), + Catlass::Arch::PositionL1{}); + auto tensorL0A32 = + tla::MakeTensor(l0A0, tla::MakeLayout(tile32, tile32), Catlass::Arch::PositionL0A{}); + auto tensorL0B32 = + tla::MakeTensor(l0B0, tla::MakeLayout(tile32, tile32), Catlass::Arch::PositionL0B{}); + auto tensorL0C32 = tla::MakeTensor(l0C0, tla::MakeLayoutL0C(tile32, tile32), Catlass::Arch::PositionL0C{}); + uint32_t localRow32 = 0; + uint32_t localColumn32 = 0; + auto tileL0A32 = GetTile(tensorL0A32, tla::MakeCoord(localRow32, localColumn32), shape32); + auto tileL0B32 = GetTile(tensorL0B32, tla::MakeCoord(localRow32, localColumn32), shape32); + auto tileL0C32 = GetTile(tensorL0C32, tla::MakeCoord(localRow32, localColumn32), shape32); + + WaitFlag(2); + WaitFlag(3); + copyGmToL1A32(tensorL1A32, blockA32); + copyGmToL1B32(tensorL1B32, blockB32); + copyGmToL1B32(tensorL1Diag32, blockDiag32); + SetFlag(0); + WaitFlag(0); + LoadSolveTile(l0A0, l1A0); + LoadSolveTile(l0B0, l1B0); + SetFlag(0); + WaitFlag(0); + tileMmad(tileL0C32, tileL0A32, tileL0B32, tile32, tile32, tile32, true, 0b11); + SetFlag(0); + WaitFlag(0); + copyL0CToDst32(blockTmp32, tileL0C32, 0b11); + SetFlag(0); + WaitFlag(0); + copyGmToL1A32(tensorL1Tmp32, blockTmp32); + SetFlag(0); + WaitFlag(0); + LoadSolveTile(l0A0, l1Tmp0); + LoadSolveTile(l0B0, l1Diag0); + SetFlag(0); + WaitFlag(0); + tileMmad(tileL0C32, tileL0A32, tileL0B32, tile32, tile32, tile32, true, 0b11); + SetFlag(0); + WaitFlag(0); + copyL0CToDst32(blockOut32, tileL0C32, 0b11); + SetFlag(0); + WaitFlag(0); + SetMMLayoutTransform(false); + } +#endif + + __aicore__ inline void ComputeAkkMergeCubeWorkspaceDispatch(uint64_t b, uint64_t hv, uint64_t chunkIdx) + { +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if constexpr (SAFE_GATE && COMPILE_BT == 64 && COMPILE_K == 128 && COMPILE_V == 128) { + ComputeAkkMergeCubeWorkspace(b, hv, chunkIdx); + return; + } +#endif + ComputeAkkMergeCubeWorkspace(b, hv, chunkIdx); + } + + __aicore__ inline void ComputeAkkInverseMchFull(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start) + { + uint64_t aBase = AOffset(b, hv, start, 0); + uint64_t xBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_X); + uint64_t yBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_Y0); + uint64_t yNextBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_Y1); + uint64_t tmpBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_TMP); + + uint32_t diagBlocks = static_cast(BT_ / KDA_SOLVE_DIAG_BT); + for (uint32_t block = 0; block < diagBlocks; ++block) { + uint32_t off = block * KDA_SOLVE_DIAG_BT; + CubeGemmSolveSub(akk_, aBase, off, off, akk_, aBase, off, off, solveWorkspace_, yBase, off, off, + KDA_SOLVE_DIAG_BT, KDA_SOLVE_DIAG_BT, KDA_SOLVE_DIAG_BT); + } + for (uint32_t iter = 0; iter < KDA_SOLVE_DIAG_MCH_ITERS; ++iter) { + for (uint32_t block = 0; block < diagBlocks; ++block) { + uint32_t off = block * KDA_SOLVE_DIAG_BT; + CubeGemmSolveSub(solveWorkspace_, xBase, off, off, solveWorkspace_, yBase, off, off, solveWorkspace_, + tmpBase, off, off, KDA_SOLVE_DIAG_BT, KDA_SOLVE_DIAG_BT, KDA_SOLVE_DIAG_BT); + } + Catlass::Arch::CrossCoreSetFlagWithReverse<0x2, PIPE_FIX>(mchSyncDoneFlag_); + if (iter + 1 < KDA_SOLVE_DIAG_MCH_ITERS) { + for (uint32_t block = 0; block < diagBlocks; ++block) { + uint32_t off = block * KDA_SOLVE_DIAG_BT; + CubeGemmSolveSub(solveWorkspace_, yBase, off, off, solveWorkspace_, yBase, off, off, + solveWorkspace_, yNextBase, off, off, KDA_SOLVE_DIAG_BT, KDA_SOLVE_DIAG_BT, + KDA_SOLVE_DIAG_BT); + } + } + Catlass::Arch::CrossCoreWaitFlagWithReverse<0x2, PIPE_FIX>(mchSyncReadyFlag_); + if (iter + 1 < KDA_SOLVE_DIAG_MCH_ITERS) { + uint64_t oldYBase = yBase; + yBase = yNextBase; + yNextBase = oldYBase; + } + } + ComputeAkkMergeCube(b, hv, chunkIdx, start); + Catlass::Arch::CrossCoreSetFlagWithReverse<0x2, PIPE_FIX>(mchSyncDoneFlag_); + } + + __aicore__ inline void ScaleRowsByBeta(GlobalTensor &src, GlobalTensor &dst, uint64_t b, uint64_t hv, + uint64_t start, uint64_t rowBegin, uint64_t rowCount, uint64_t dim, + LocalTensor &betaLocal, LocalTensor &betaBrcb, + LocalTensor &matrixLocal, bool sourceSequenceMajor = false) + { + constexpr uint64_t vecElemsPerRepeat = 64; + constexpr uint64_t typedOffsetFloats = 20480; + constexpr uint64_t typedOffset = typedOffsetFloats * sizeof(float) / sizeof(T); + uint64_t elemCount = rowCount * dim; + uint64_t baseOffset = KVOffset(b, hv, start + rowBegin, 0, dim); + uint64_t sourceOffset = sourceSequenceMajor ? VInputOffset(b, hv, start + rowBegin, 0) : baseOffset; + uint64_t sourceStride = sourceSequenceMajor ? HV_ * dim : dim; + + if constexpr (IsSameType::value) { + CopyRowsIn(matrixLocal, src, sourceOffset, rowCount, dim, sourceStride); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + } else { + LocalTensor matrixTyped = vecBuf_.Get()[typedOffset]; + CopyRowsIn(matrixTyped, src, sourceOffset, rowCount, dim, sourceStride); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + Cast(matrixLocal, matrixTyped, RoundMode::CAST_NONE, static_cast(elemCount)); + PipeBarrier(); + } + +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + ApplyKdaRowScaleRegbase((__ubuf__ float *)reinterpret_cast(matrixLocal.GetPhyAddr()), + (__ubuf__ float *)reinterpret_cast(betaLocal.GetPhyAddr()), + static_cast(rowCount), static_cast(dim)); +#else + uint8_t repeatStride = static_cast(dim * sizeof(float) / 32); + for (uint64_t col = 0; col < dim; col += vecElemsPerRepeat) { + uint64_t mask = dim - col; + if (mask > vecElemsPerRepeat) { + mask = vecElemsPerRepeat; + } + Mul(matrixLocal[col], matrixLocal[col], betaBrcb, mask, static_cast(rowCount), + {1, 1, 0, repeatStride, repeatStride, 1}); + } + PipeBarrier(); +#endif + + if constexpr (IsSameType::value) { + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + DataCopy(dst[baseOffset], matrixLocal, static_cast(elemCount)); + } else { + LocalTensor matrixTyped = vecBuf_.Get()[typedOffset]; + Cast(matrixTyped, matrixLocal, RoundMode::CAST_RINT, static_cast(elemCount)); + PipeBarrier(); + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + DataCopy(dst[baseOffset], matrixTyped, static_cast(elemCount)); + } + SetFlag(mte3ToMte2Events_[0]); + WaitFlag(mte3ToMte2Events_[0]); + SetFlag(mte3ToVEvent_); + WaitFlag(mte3ToVEvent_); + } + + __aicore__ inline void PrepareWuCubeInputs(uint64_t b, uint64_t hv, uint64_t start, uint64_t curT, + uint64_t subBlockIdx, uint64_t subBlockNum) + { + uint64_t rowsPerSubBlock = (curT + subBlockNum - 1) / subBlockNum; + uint64_t rowBegin = subBlockIdx * rowsPerSubBlock; + if (rowBegin >= curT) { + return; + } + uint64_t rowCount = curT - rowBegin; + if (rowCount > rowsPerSubBlock) { + rowCount = rowsPerSubBlock; + } + LocalTensor arena = vecBuf_.Get(); + LocalTensor betaLocal = arena; + LocalTensor betaBrcb = arena[KDA_SOLVE_BT]; + LocalTensor matrixLocal = arena[KDA_SOLVE_BT + 512]; +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + // Keep each arch35 Cast/regbase panel within an 8K-element UB instruction span. + constexpr uint64_t maxScaleElements = 8192; + if (rowCount * K_ > maxScaleElements || rowCount * V_ > maxScaleElements) { + constexpr uint64_t tileRows = 16; + for (uint64_t tileRow = 0; tileRow < rowCount; tileRow += tileRows) { + uint64_t tileCount = rowCount - tileRow; + if (tileCount > tileRows) { + tileCount = tileRows; + } + LoadAsFloatRow(beta_, BetaOffset(b, hv, start + rowBegin + tileRow), betaLocal, tileCount); + ScaleRowsByBeta(w_, w_, b, hv, start, rowBegin + tileRow, tileCount, K_, betaLocal, betaBrcb, + matrixLocal); + ScaleRowsByBeta(v_, vNew_, b, hv, start, rowBegin + tileRow, tileCount, V_, betaLocal, betaBrcb, + matrixLocal, inputSequenceMajor_); + } + return; + } +#endif + LoadAsFloatRow(beta_, BetaOffset(b, hv, start + rowBegin), betaLocal, rowCount); +#if !defined(__CCE_AICORE__) || __CCE_AICORE__ != 310 + Brcb(betaBrcb, betaLocal, static_cast((rowCount + 7) / 8), {1, 8}); + PipeBarrier(); +#endif + ScaleRowsByBeta(w_, w_, b, hv, start, rowBegin, rowCount, K_, betaLocal, betaBrcb, matrixLocal); + ScaleRowsByBeta(v_, vNew_, b, hv, start, rowBegin, rowCount, V_, betaLocal, betaBrcb, matrixLocal, + inputSequenceMajor_); + } + + __aicore__ inline void FinalizePrepareIntermediates(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, + uint64_t curT, uint64_t subBlockIdx, uint64_t subBlockNum) + { +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + constexpr bool qgAlreadyStored = SAFE_GATE && COMPILE_BT == 64 && COMPILE_K == 128 && COMPILE_V == 128; +#else + constexpr bool qgAlreadyStored = false; +#endif + constexpr uint64_t tileRows = 32; + // Keep tail rows on the same AIV that owns their padded solve rows. Splitting by curT would + // move short-tail export to AIV1 while AIV0 is still writing the solved matrix. + const uint64_t rowBegin = (BT_ * subBlockIdx) / subBlockNum; + uint64_t rowEnd = (BT_ * (subBlockIdx + 1)) / subBlockNum; + if (rowEnd > curT) { + rowEnd = curT; + } + if (rowBegin >= rowEnd) { + return; + } + for (uint64_t tileRow = rowBegin; tileRow < rowEnd; tileRow += tileRows) { + const uint64_t rows = (rowEnd - tileRow) > tileRows ? tileRows : (rowEnd - tileRow); + const uint64_t matrixElems = rows * BT_; + const uint64_t qgElems = rows * K_; + LocalTensor arena = vecBuf_.Get(); + LocalTensor aqkLocal = arena; + LocalTensor akkLocal = arena[matrixElems]; + const uint64_t typedOffset = (2 * matrixElems + qgElems) * sizeof(float) / sizeof(T); + LocalTensor typedBase = vecBuf_.Get()[typedOffset]; + LocalTensor aqkTyped = typedBase; + LocalTensor akkTyped = typedBase[matrixElems]; + LocalTensor qgTyped = typedBase[2 * matrixElems]; + +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if constexpr (SAFE_GATE && COMPILE_BT == 64 && COMPILE_K == 128 && COMPILE_V == 128) { + const uint64_t xBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_X) + tileRow * BT_; + CopyVectorIn(akkLocal, solveWorkspace_, xBase, matrixElems); + } else { + CopyVectorIn(aqkLocal, aqk_, AOffset(b, hv, start + tileRow, 0), matrixElems); + CopyVectorIn(akkLocal, akk_, AOffset(b, hv, start + tileRow, 0), matrixElems); + } +#else + CopyVectorIn(aqkLocal, aqk_, AOffset(b, hv, start + tileRow, 0), matrixElems); + CopyVectorIn(akkLocal, akk_, AOffset(b, hv, start + tileRow, 0), matrixElems); +#endif + if constexpr (!qgAlreadyStored) { + CopyVectorIn(qgTyped, qg_, KVOffset(b, hv, start + tileRow, 0, K_), qgElems); + } + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if constexpr (!(SAFE_GATE && COMPILE_BT == 64 && COMPILE_K == 128 && COMPILE_V == 128)) { + Muls(aqkLocal, aqkLocal, scale_, static_cast(matrixElems)); + } +#else + Muls(aqkLocal, aqkLocal, scale_, static_cast(matrixElems)); +#endif +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if constexpr (!(SAFE_GATE && COMPILE_BT == 64 && COMPILE_K == 128 && COMPILE_V == 128)) { + ClampFp32ToOutputType(aqkLocal, static_cast(matrixElems)); + } +#else + ClampFp32ToOutputType(aqkLocal, static_cast(matrixElems)); +#endif + ClampFp32ToOutputType(akkLocal, static_cast(matrixElems)); +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if constexpr (!(SAFE_GATE && COMPILE_BT == 64 && COMPILE_K == 128 && COMPILE_V == 128)) { + Cast(aqkTyped, aqkLocal, RoundMode::CAST_RINT, static_cast(matrixElems)); + } +#else + Cast(aqkTyped, aqkLocal, RoundMode::CAST_RINT, static_cast(matrixElems)); +#endif + Cast(akkTyped, akkLocal, RoundMode::CAST_RINT, static_cast(matrixElems)); + PipeBarrier(); + + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if constexpr (!(SAFE_GATE && COMPILE_BT == 64 && COMPILE_K == 128 && COMPILE_V == 128)) { + CopyVectorOut(o_, AOffset(b, hv, start + tileRow, 0), aqkTyped, matrixElems); + } +#else + CopyVectorOut(o_, AOffset(b, hv, start + tileRow, 0), aqkTyped, matrixElems); +#endif + CopyVectorOut(u_, AOffset(b, hv, start + tileRow, 0), akkTyped, matrixElems); + if constexpr (!qgAlreadyStored) { + CopyVectorOut(kg_, KVOffset(b, hv, start + tileRow, 0, K_), qgTyped, qgElems); + } + SetFlag(mte3ToMte2Events_[0]); + WaitFlag(mte3ToMte2Events_[0]); + SetFlag(mte3ToVEvent_); + WaitFlag(mte3ToVEvent_); + } + } + + __aicore__ inline bool ResolveFlatChunk(uint64_t task, uint64_t &seq, uint64_t &b, uint64_t &h, uint64_t &hv, + uint64_t &chunkIdx, uint64_t &start, uint64_t &end) + { + hv = task % HV_; + uint64_t flatChunk = task / HV_; + if (!isVarLen_) { + seq = flatChunk / NT_; + b = seq; + chunkIdx = flatChunk % NT_; + start = chunkIdx * BT_; + end = start + BT_; + if (end > T_) { + end = T_; + } + } else { + if (!KdaVarlen::ResolveChunkRange(cuSeqlensAddr_, chunkIndicesAddr_, N_, T_, BT_, flatChunk, seq, start, + end)) { + return false; + } + b = 0; + chunkIdx = flatChunk; + } + h = hv / (HV_ / H_); + return start < end; + } + + __aicore__ inline void ProcessChunkPreAiv(uint64_t b, uint64_t h, uint64_t hv, uint64_t chunkIdx, uint64_t start, + uint64_t end, uint64_t subBlockIdx, uint64_t subBlockNum) + { + if constexpr (IsSameType::value) { + ProcessChunkPreAivFp32(b, h, hv, chunkIdx, start, end, subBlockIdx, subBlockNum); + } + } + + template + __aicore__ inline void JoinAivMte3() + { + if constexpr (CORE_TYPE == AscendC::AIV) { +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if (!isAivOnly_) { + if (!headPairMode_) { + Catlass::Arch::CrossCoreBarrier<0x1, PIPE_MTE3>(); + } + PipeBarrier(); + } +#endif + } + } + + template + __aicore__ inline void RunAicAfterBothAivReady(uint64_t subBlockIdx, uint64_t subBlockNum) + { + if constexpr (CORE_TYPE == AscendC::AIV) { + (void)subBlockIdx; + (void)subBlockNum; + JoinAivMte3(); + if constexpr (SAFE_GATE) { + Catlass::Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(syncReadyFlag_); + Catlass::Arch::CrossCoreWaitFlag(syncDoneFlag_); + } else { + Catlass::Arch::CrossCoreSetFlagWithReverse<0x2, PIPE_MTE3>(mchSyncReadyFlag_); + Catlass::Arch::CrossCoreWaitFlagWithReverse<0x2, PIPE_MTE2>(mchSyncDoneFlag_); + } + } + } + + template + __aicore__ inline void SignalAicSolveReady() + { + if constexpr (CORE_TYPE == AscendC::AIV) { + JoinAivMte3(); + Catlass::Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(syncReadyFlag_); + } + } + + template + __aicore__ inline void WaitAicSolveDone() + { + if constexpr (CORE_TYPE == AscendC::AIV) { + Catlass::Arch::CrossCoreWaitFlag(syncDoneFlag_); + } + } + + template + __aicore__ inline void SignalPostWuReady() + { + if constexpr (CORE_TYPE == AscendC::AIV) { + Catlass::Arch::CrossCoreSetFlagWithReverse<0x2, PIPE_MTE3>(postWuReadyFlag_); + } + } + + __aicore__ inline void ProcessChunkPreAivFp32(uint64_t b, uint64_t h, uint64_t hv, uint64_t chunkIdx, + uint64_t start, uint64_t end, uint64_t subBlockIdx, + uint64_t subBlockNum, bool deferSafeSolve = false, + bool waitPendingSafeSolve = false, uint64_t scoreLane = 0, + bool pairHeads = false) + { + uint64_t curT = end - start; + if (curT == 0) { + return; + } + if constexpr (IsSameType::value) { + return; + } + + if (K_ < 16) { + return; + } + MaterializeRawGateChunkArch35(b, hv, start, curT); + bool usePostWuCube = UsePostWuCube(curT); + bool useAkkCubeSolve = UseAkkCubeSolve(curT); + uint64_t solveRowBegin = 0; + uint64_t solveRowEnd = 0; + GetSolveRowRange(BT_, subBlockIdx, subBlockNum, solveRowBegin, solveRowEnd); + // Safe-gate score factors need the bounded 16-row reference span. + // A single 64-row reference loses BF16 dynamic range for valid gates. + const bool useFullChunkScore = false; + uint64_t scoreBlockSize = useFullChunkScore ? curT : ScoreRefBlockSize(); + uint64_t scoreBlockCount = (curT + scoreBlockSize - 1) / scoreBlockSize; + uint64_t pipelineBlockCount = useFullChunkScore ? scoreBlockCount : + (scoreBlockCount + KDA_SCORE_QUEUE_DEPTH - 1) / + KDA_SCORE_QUEUE_DEPTH * KDA_SCORE_QUEUE_DEPTH; + const bool useDirectScoreUb = +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + KDA_ARCH35_ENABLE_DIRECT_SCORE_UB && pairHeads && curT == 64 && scoreBlockCount == 2; +#else + false; +#endif + bool firstSolveRowsPrepared = false; + for (uint64_t block = 0; block < pipelineBlockCount; ++block) { + if (block < scoreBlockCount) { + uint64_t rowBegin = block * scoreBlockSize; + uint64_t rowCount = useFullChunkScore ? curT - rowBegin : ScoreRowBlockCount(curT, rowBegin); + uint64_t refToken = ScoreRefToken(start, curT, rowBegin, rowCount); + uint64_t queueSlot = + useFullChunkScore ? activeSolveSlot_ / KDA_SCORE_LANES : block % KDA_SCORE_QUEUE_DEPTH; + uint64_t scoreSlot = ScoreScratchSlot(queueSlot, scoreLane, pairHeads); + PrepareGateProducts(b, h, hv, start, curT, subBlockIdx, subBlockNum, true, refToken, + rowBegin + rowCount, true, scoreSlot, rowBegin, rowCount); + } + JoinAivMte3(); + Catlass::Arch::CrossCoreSetFlagWithReverse<0x2, PIPE_MTE3>(scoreReadyFlag_); + if (block > 0) { + if constexpr (SAFE_GATE) { + if (waitPendingSafeSolve) { + WaitAicSolveDone(); + waitPendingSafeSolve = false; + } + } + Catlass::Arch::CrossCoreWaitFlagWithReverse<0x2, PIPE_MTE2>(scoreDoneFlag_); +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if constexpr (SAFE_GATE && COMPILE_BT == 64 && COMPILE_K == 128 && COMPILE_V == 128) { + if (useDirectScoreUb && useAkkCubeSolve && block == 1) { + ProcessDirectScoreSolveRowsArch35(b, hv, chunkIdx, start, 0, 0); + firstSolveRowsPrepared = true; + } else if (pairHeads && useAkkCubeSolve && block == 1) { + uint64_t firstRowBegin = 0; + uint64_t firstRowEnd = 0; + GetSolveRowRange(BT_, 0, KDA_SCORE_LANES, firstRowBegin, firstRowEnd); + PrepareAqkAkkSolveInputRows(b, hv, chunkIdx, start, curT, firstRowBegin, firstRowEnd, false, + false); + firstSolveRowsPrepared = true; + } + } +#endif + } + } + bool fusedScoreWriteback = false; +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + fusedScoreWriteback = SAFE_GATE && BT_ == 64 && K_ == 128 && V_ == 128; +#endif + if (!fusedScoreWriteback) { + // The final score MMAD only consumes scoreWorkspace_. Run the + // independent gate writeback while AIC drains its MMAD/Fixpipe path. + PrepareGateProducts(b, h, hv, start, curT, subBlockIdx, subBlockNum); + } + if (pipelineBlockCount > 0) { + Catlass::Arch::CrossCoreWaitFlagWithReverse<0x2, PIPE_MTE2>(scoreDoneFlag_); + } +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if (useDirectScoreUb && useAkkCubeSolve) { + ProcessDirectScoreSolveRowsArch35(b, hv, chunkIdx, start, KDA_DIRECT_SCORE_ROWS, 1); + } +#endif + if constexpr (SAFE_GATE) { + if (waitPendingSafeSolve) { + WaitAicSolveDone(); + waitPendingSafeSolve = false; + } + } + if (useAkkCubeSolve) { + bool fullChunk = curT == BT_; + if constexpr (SAFE_GATE) { + if (pairHeads) { + if (!useDirectScoreUb) { + uint64_t firstRowPart = firstSolveRowsPrepared ? 1 : 0; + for (uint64_t rowPart = firstRowPart; rowPart < KDA_SCORE_LANES; ++rowPart) { + uint64_t pairRowBegin = 0; + uint64_t pairRowEnd = 0; + GetSolveRowRange(BT_, rowPart, KDA_SCORE_LANES, pairRowBegin, pairRowEnd); + PrepareAqkAkkSolveInputRows(b, hv, chunkIdx, start, curT, pairRowBegin, pairRowEnd, false, + false); + } + } + } else { + PrepareAqkAkkSolveInputRows(b, hv, chunkIdx, start, curT, solveRowBegin, solveRowEnd, false, false); + } + if (deferSafeSolve) { + SignalAicSolveReady(); + return; + } + RunAicAfterBothAivReady(subBlockIdx, subBlockNum); +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if constexpr (!(SAFE_GATE && COMPILE_BT == 64 && COMPILE_K == 128 && COMPILE_V == 128)) { + StoreSolveXRowsToAkk(b, hv, chunkIdx, start, curT, solveRowBegin, solveRowEnd); + } +#else + StoreSolveXRowsToAkk(b, hv, chunkIdx, start, curT, solveRowBegin, solveRowEnd); +#endif + } else { + PrepareAqkAkkSolveInputRows(b, hv, chunkIdx, start, curT, solveRowBegin, solveRowEnd, fullChunk, false); + if (!fullChunk) { + RunAicAfterBothAivReady(subBlockIdx, subBlockNum); + StoreSolveXRowsToAkk(b, hv, chunkIdx, start, curT, solveRowBegin, solveRowEnd); + } else { + uint32_t solveIters = KDA_SOLVE_DIAG_MCH_ITERS; + RunAicAfterBothAivReady(subBlockIdx, subBlockNum); + for (uint32_t iter = 0; iter < solveIters; ++iter) { + AddSolveTmpToXDiagRows(b, hv, chunkIdx, start, solveRowBegin, solveRowEnd, + iter + 1 == solveIters); + RunAicAfterBothAivReady(subBlockIdx, subBlockNum); + } + } + } + } + // Host validation guarantees every accepted shape has enough workspace for this cube path. +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if constexpr (SAFE_GATE && COMPILE_BT == 64 && COMPILE_K == 128 && COMPILE_V == 128) { + if (HV_ % KDA_SCORE_LANES != 0) { + PrepareWuCubeInputs(b, hv, start, curT, subBlockIdx, subBlockNum); + } + } else { + PrepareWuCubeInputs(b, hv, start, curT, subBlockIdx, subBlockNum); + } +#else + PrepareWuCubeInputs(b, hv, start, curT, subBlockIdx, subBlockNum); +#endif + FinalizePrepareIntermediates(b, hv, chunkIdx, start, curT, subBlockIdx, subBlockNum); + } + + __aicore__ inline void FinishDeferredSafeChunk(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, + uint64_t end, uint64_t subBlockIdx, uint64_t subBlockNum) + { + uint64_t curT = end - start; + uint64_t solveRowBegin = 0; + uint64_t solveRowEnd = 0; + GetSolveRowRange(BT_, subBlockIdx, subBlockNum, solveRowBegin, solveRowEnd); +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if constexpr (!(SAFE_GATE && COMPILE_BT == 64 && COMPILE_K == 128 && COMPILE_V == 128)) { + StoreSolveXRowsToAkk(b, hv, chunkIdx, start, curT, solveRowBegin, solveRowEnd); + } +#else + StoreSolveXRowsToAkk(b, hv, chunkIdx, start, curT, solveRowBegin, solveRowEnd); +#endif +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if constexpr (SAFE_GATE && COMPILE_BT == 64 && COMPILE_K == 128 && COMPILE_V == 128) { + if (HV_ % KDA_SCORE_LANES != 0) { + PrepareWuCubeInputs(b, hv, start, curT, subBlockIdx, subBlockNum); + } + } else { + PrepareWuCubeInputs(b, hv, start, curT, subBlockIdx, subBlockNum); + } +#else + PrepareWuCubeInputs(b, hv, start, curT, subBlockIdx, subBlockNum); +#endif + FinalizePrepareIntermediates(b, hv, chunkIdx, start, curT, subBlockIdx, subBlockNum); + } + + __aicore__ inline void FinishDeferredSafeChunkPair(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, + uint64_t end) + { + FinishDeferredSafeChunk(b, hv, chunkIdx, start, end, 0, 1); + } + + __aicore__ inline void ProcessChunkPreAic(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, uint64_t end) + { + if constexpr (IsSameType::value) { + ProcessChunkPreAicFp32(b, hv, chunkIdx, start, end); + } + } + + __aicore__ inline void ProcessChunkPreAicFp32(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, + uint64_t end) + { + uint64_t curT = end - start; + if (curT == 0 || K_ < 16) { + return; + } + uint64_t scoreBlockSize = ScoreRefBlockSize(); + uint64_t scoreBlockCount = (curT + scoreBlockSize - 1) / scoreBlockSize; + uint64_t pipelineBlockCount = + (scoreBlockCount + KDA_SCORE_QUEUE_DEPTH - 1) / KDA_SCORE_QUEUE_DEPTH * KDA_SCORE_QUEUE_DEPTH; + for (uint64_t block = 0; block < pipelineBlockCount; ++block) { + Catlass::Arch::CrossCoreWaitFlagWithReverse<0x2, PIPE_FIX>(scoreReadyFlag_); + if (block < scoreBlockCount) { + uint64_t rowBegin = block * scoreBlockSize; + uint64_t rowCount = ScoreRowBlockCount(curT, rowBegin); + ComputeRawAqkAkkCubeBlock(b, hv, chunkIdx, start, curT, rowBegin, rowCount, true, + block % KDA_SCORE_QUEUE_DEPTH, rowBegin + rowCount); + } + Catlass::Arch::CrossCoreSetFlagWithReverse<0x2, PIPE_FIX>(scoreDoneFlag_); + } + bool usePostWuCube = UsePostWuCube(curT); + bool useAkkCubeSolve = UseAkkCubeSolve(curT); + if (useAkkCubeSolve) { + if constexpr (SAFE_GATE) { + Catlass::Arch::CrossCoreWaitFlag(syncReadyFlag_); + ComputeAkkMergeCubeWorkspaceDispatch(b, hv, chunkIdx); + Catlass::Arch::CrossCoreSetFlag<0x2, PIPE_FIX>(syncDoneFlag_); + } else { + Catlass::Arch::CrossCoreWaitFlagWithReverse<0x2, PIPE_FIX>(mchSyncReadyFlag_); + if (curT == BT_) { + ComputeAkkInverseMchFull(b, hv, chunkIdx, start); + } else { + ComputeAkkMergeCubeWorkspace(b, hv, chunkIdx); + Catlass::Arch::CrossCoreSetFlagWithReverse<0x2, PIPE_FIX>(mchSyncDoneFlag_); + } + } + } + (void)usePostWuCube; + (void)chunkIdx; + } + + __aicore__ inline void ProcessChunkPreAicHeadPairFp32(uint64_t b, uint64_t hvBase, uint64_t chunkIdx, + uint64_t start, uint64_t end, uint64_t localTaskIdx) + { + uint64_t curT = end - start; + if (curT == 0 || K_ < 16) { + return; + } + const bool useFullChunkScore = false; + uint64_t scoreBlockSize = useFullChunkScore ? curT : ScoreRefBlockSize(); + uint64_t scoreBlockCount = (curT + scoreBlockSize - 1) / scoreBlockSize; + uint64_t pipelineBlockCount = useFullChunkScore ? scoreBlockCount : + (scoreBlockCount + KDA_SCORE_QUEUE_DEPTH - 1) / + KDA_SCORE_QUEUE_DEPTH * KDA_SCORE_QUEUE_DEPTH; + for (uint64_t block = 0; block < pipelineBlockCount; ++block) { + Catlass::Arch::CrossCoreWaitFlagWithReverse<0x2, PIPE_FIX>(scoreReadyFlag_); + if (block < scoreBlockCount) { + uint64_t rowBegin = block * scoreBlockSize; + uint64_t rowCount = useFullChunkScore ? curT - rowBegin : ScoreRowBlockCount(curT, rowBegin); + uint64_t queueSlot = + useFullChunkScore ? localTaskIdx % KDA_SCORE_QUEUE_DEPTH : block % KDA_SCORE_QUEUE_DEPTH; +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + bool directScoreDispatched = false; + if constexpr (SAFE_GATE) { + if (KDA_ARCH35_ENABLE_DIRECT_SCORE_UB && curT == 64 && rowCount == 32) { + uint64_t scoreSlotBase = ScoreScratchSlot(queueSlot, 0, true); + if (rowBegin == 0) { + ComputeRawAqkAkkCubeStableHeadPairDirectUbArch35<32>(rowBegin, scoreSlotBase, + static_cast(block)); + } else { + ComputeRawAqkAkkCubeStableHeadPairDirectUbArch35<64>(rowBegin, scoreSlotBase, + static_cast(block)); + } + directScoreDispatched = true; + } + } + if (!directScoreDispatched) +#endif + { + for (uint64_t lane = 0; lane < KDA_SCORE_LANES; ++lane) { + uint64_t hv = hvBase + lane; + uint64_t scoreSlot = ScoreScratchSlot(queueSlot, lane, true); + activeSolveSlot_ = + (localTaskIdx % (KDA_SOLVE_PIPELINE_DEPTH / KDA_SCORE_LANES)) * KDA_SCORE_LANES + lane; + ComputeRawAqkAkkCubeBlock(b, hv, chunkIdx, start, curT, rowBegin, rowCount, true, scoreSlot, + rowBegin + rowCount); + } + } + } + Catlass::Arch::CrossCoreSetFlagWithReverse<0x2, PIPE_FIX>(scoreDoneFlag_); + } + + if (UseAkkCubeSolve(curT)) { + Catlass::Arch::CrossCoreWaitFlag(syncReadyFlag_); + for (uint64_t lane = 0; lane < KDA_SCORE_LANES; ++lane) { + activeSolveSlot_ = + (localTaskIdx % (KDA_SOLVE_PIPELINE_DEPTH / KDA_SCORE_LANES)) * KDA_SCORE_LANES + lane; + ComputeAkkMergeCubeWorkspaceDispatch(b, hvBase + lane, chunkIdx); + } + Catlass::Arch::CrossCoreSetFlag<0x2, PIPE_FIX>(syncDoneFlag_); + } + } + + __aicore__ inline bool ResolveFlatChunkForHv(uint64_t flatChunk, uint64_t hv, uint64_t &seq, uint64_t &b, + uint64_t &h, uint64_t &chunkIdx, uint64_t &start, uint64_t &end) + { + if (!isVarLen_) { + seq = flatChunk / NT_; + b = seq; + chunkIdx = flatChunk % NT_; + start = chunkIdx * BT_; + end = start + BT_; + if (end > T_) { + end = T_; + } + } else { + if (!KdaVarlen::ResolveChunkRange(cuSeqlensAddr_, chunkIndicesAddr_, N_, T_, BT_, flatChunk, seq, start, + end)) { + return false; + } + b = 0; + chunkIdx = flatChunk; + } + h = hv / (HV_ / H_); + return start < end; + } + + __aicore__ inline void ProcessPreAivHeadPair() + { + const uint64_t subBlockIdx = static_cast(GetSubBlockIdx()); + const uint64_t subBlockNum = static_cast(GetSubBlockNum()); + const uint64_t coreNum = usedCoreNum_; + const uint64_t coreIdx = static_cast(GetBlockIdx()) / subBlockNum; + const uint64_t chunkCount = isVarLen_ ? NT_ : B_ * NT_; + const uint64_t headWindows = HV_ / KDA_SCORE_LANES; + const uint64_t taskNum = chunkCount * headWindows; + bool pendingValid = false; + uint64_t pendingB = 0; + uint64_t pendingHv = 0; + uint64_t pendingChunkIdx = 0; + uint64_t pendingStart = 0; + uint64_t pendingEnd = 0; + uint64_t pendingSlot = 0; + uint64_t localTaskIdx = 0; + +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if constexpr (SAFE_GATE && COMPILE_BT == 64 && COMPILE_K == 128 && COMPILE_V == 128) { + if (KDA_ARCH35_ENABLE_DIRECT_SCORE_UB) { + InitializeDirectScoreUbArch35(); + } + } +#endif + + for (uint64_t task = coreIdx; task < taskNum; task += coreNum, ++localTaskIdx) { + uint64_t flatChunk = task / headWindows; + uint64_t hv = (task % headWindows) * KDA_SCORE_LANES + subBlockIdx; + uint64_t seq = 0; + uint64_t b = 0; + uint64_t h = 0; + uint64_t chunkIdx = 0; + uint64_t start = 0; + uint64_t end = 0; + if (!ResolveFlatChunkForHv(flatChunk, hv, seq, b, h, chunkIdx, start, end)) { + continue; + } + (void)seq; + uint64_t currentSlot = + (localTaskIdx % (KDA_SOLVE_PIPELINE_DEPTH / KDA_SCORE_LANES)) * KDA_SCORE_LANES + subBlockIdx; + activeSolveSlot_ = currentSlot; + bool deferSolve = UseAkkCubeSolve(end - start); + ProcessChunkPreAivFp32(b, h, hv, chunkIdx, start, end, 0, 1, deferSolve, pendingValid, subBlockIdx, true); + if (pendingValid) { + activeSolveSlot_ = pendingSlot; + FinishDeferredSafeChunkPair(pendingB, pendingHv, pendingChunkIdx, pendingStart, pendingEnd); + if (fusePostWu_) { + SignalPostWuReady(); + } + } + pendingValid = deferSolve; + if (pendingValid) { + pendingB = b; + pendingHv = hv; + pendingChunkIdx = chunkIdx; + pendingStart = start; + pendingEnd = end; + pendingSlot = currentSlot; + } + } + if (pendingValid) { + WaitAicSolveDone(); + activeSolveSlot_ = pendingSlot; + FinishDeferredSafeChunkPair(pendingB, pendingHv, pendingChunkIdx, pendingStart, pendingEnd); + if (fusePostWu_) { + SignalPostWuReady(); + } + } + } + + __aicore__ inline void ProcessPreAicHeadPair() + { + const uint64_t chunkCount = isVarLen_ ? NT_ : B_ * NT_; + const uint64_t headWindows = HV_ / KDA_SCORE_LANES; + const uint64_t taskNum = chunkCount * headWindows; + const uint64_t coreNum = usedCoreNum_ == 0 ? 1 : usedCoreNum_; + uint64_t localTaskIdx = 0; + for (uint64_t task = GetBlockIdx(); task < taskNum; task += coreNum, ++localTaskIdx) { + uint64_t flatChunk = task / headWindows; + uint64_t hvBase = (task % headWindows) * KDA_SCORE_LANES; + uint64_t seq = 0; + uint64_t b = 0; + uint64_t h = 0; + uint64_t chunkIdx = 0; + uint64_t start = 0; + uint64_t end = 0; + if (ResolveFlatChunkForHv(flatChunk, hvBase, seq, b, h, chunkIdx, start, end)) { + (void)seq; + (void)h; + ProcessChunkPreAicHeadPairFp32(b, hvBase, chunkIdx, start, end, localTaskIdx); + } + } + } + + template + __aicore__ inline void ProcessPreAicHeadPairFused(PostWuOp &postWu) + { + const uint64_t chunkCount = isVarLen_ ? NT_ : B_ * NT_; + const uint64_t headWindows = HV_ / KDA_SCORE_LANES; + const uint64_t taskNum = chunkCount * headWindows; + const uint64_t coreNum = usedCoreNum_ == 0 ? 1 : usedCoreNum_; + uint64_t batchB[KDA_POST_QUEUE_STORAGE]; + uint64_t batchHvBase[KDA_POST_QUEUE_STORAGE]; + uint64_t batchStart[KDA_POST_QUEUE_STORAGE]; + uint64_t batchEnd[KDA_POST_QUEUE_STORAGE]; + uint16_t batchCount = 0; + uint64_t localTaskIdx = 0; + + for (uint64_t task = GetBlockIdx(); task < taskNum; task += coreNum, ++localTaskIdx) { + uint64_t flatChunk = task / headWindows; + uint64_t hvBase = (task % headWindows) * KDA_SCORE_LANES; + uint64_t seq = 0; + uint64_t b = 0; + uint64_t h = 0; + uint64_t chunkIdx = 0; + uint64_t start = 0; + uint64_t end = 0; + if (!ResolveFlatChunkForHv(flatChunk, hvBase, seq, b, h, chunkIdx, start, end)) { + continue; + } + (void)seq; + (void)h; + ProcessChunkPreAicHeadPairFp32(b, hvBase, chunkIdx, start, end, localTaskIdx); + batchB[batchCount] = b; + batchHvBase[batchCount] = hvBase; + batchStart[batchCount] = start; + batchEnd[batchCount] = end; + ++batchCount; + + if (batchCount == KDA_POST_QUEUE_STORAGE) { + for (uint16_t i = 0; i < KDA_POST_QUEUE_DEPTH; ++i) { + Catlass::Arch::CrossCoreWaitFlagWithReverse<0x2, PIPE_MTE2>(postWuReadyFlag_); + } + postWu.ProcessPreparedHeadPairBatchArch35(batchB, batchHvBase, batchStart, batchEnd, + KDA_POST_QUEUE_DEPTH); + batchB[0] = batchB[KDA_POST_QUEUE_DEPTH]; + batchHvBase[0] = batchHvBase[KDA_POST_QUEUE_DEPTH]; + batchStart[0] = batchStart[KDA_POST_QUEUE_DEPTH]; + batchEnd[0] = batchEnd[KDA_POST_QUEUE_DEPTH]; + batchCount = 1; + } + } + + if (batchCount > 0) { + for (uint16_t i = 0; i < batchCount; ++i) { + Catlass::Arch::CrossCoreWaitFlagWithReverse<0x2, PIPE_MTE2>(postWuReadyFlag_); + } + postWu.ProcessPreparedHeadPairBatchArch35(batchB, batchHvBase, batchStart, batchEnd, batchCount); + } + } + + __aicore__ inline void ProcessPreAiv() + { + if constexpr (IsSameType::value) { + isAivOnly_ = true; + } +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if constexpr (SAFE_GATE && !IsSameType::value) { + if (headPairMode_ && !isAivOnly_) { + ProcessPreAivHeadPair(); + return; + } + } +#endif + uint64_t subBlockNum = isAivOnly_ ? 1 : static_cast(GetSubBlockNum()); + if (subBlockNum == 0) { + return; + } + uint64_t subBlockIdx = isAivOnly_ ? 0 : static_cast(GetSubBlockIdx()); + uint64_t coreNum = isAivOnly_ ? static_cast(GetBlockNum()) : usedCoreNum_; + uint64_t coreIdx = + isAivOnly_ ? static_cast(GetBlockIdx()) : static_cast(GetBlockIdx()) / subBlockNum; + uint64_t taskNum = static_cast((isVarLen_ ? NT_ : B_ * NT_) * HV_); + if constexpr (SAFE_GATE && !IsSameType::value) { + bool pendingValid = false; + uint64_t pendingB = 0; + uint64_t pendingHv = 0; + uint64_t pendingChunkIdx = 0; + uint64_t pendingStart = 0; + uint64_t pendingEnd = 0; + uint64_t pendingSlot = 0; + uint64_t localTaskIdx = 0; + for (uint64_t task = coreIdx; task < taskNum; task += coreNum, ++localTaskIdx) { + uint64_t seq = 0; + uint64_t b = 0; + uint64_t h = 0; + uint64_t hv = 0; + uint64_t chunkIdx = 0; + uint64_t start = 0; + uint64_t end = 0; + if (!ResolveFlatChunk(task, seq, b, h, hv, chunkIdx, start, end)) { + continue; + } + (void)seq; + uint64_t currentSlot = localTaskIdx % KDA_SOLVE_PIPELINE_DEPTH; + activeSolveSlot_ = currentSlot; + bool deferSolve = UseAkkCubeSolve(end - start); + if (!deferSolve && pendingValid) { + WaitAicSolveDone(); + activeSolveSlot_ = pendingSlot; + FinishDeferredSafeChunk(pendingB, pendingHv, pendingChunkIdx, pendingStart, pendingEnd, subBlockIdx, + subBlockNum); + pendingValid = false; + activeSolveSlot_ = currentSlot; + } + ProcessChunkPreAivFp32(b, h, hv, chunkIdx, start, end, subBlockIdx, subBlockNum, deferSolve, + pendingValid); + if (pendingValid) { + activeSolveSlot_ = pendingSlot; + FinishDeferredSafeChunk(pendingB, pendingHv, pendingChunkIdx, pendingStart, pendingEnd, subBlockIdx, + subBlockNum); + } + pendingValid = deferSolve; + if (pendingValid) { + pendingB = b; + pendingHv = hv; + pendingChunkIdx = chunkIdx; + pendingStart = start; + pendingEnd = end; + pendingSlot = currentSlot; + } + } + if (pendingValid) { + WaitAicSolveDone(); + activeSolveSlot_ = pendingSlot; + FinishDeferredSafeChunk(pendingB, pendingHv, pendingChunkIdx, pendingStart, pendingEnd, subBlockIdx, + subBlockNum); + } + return; + } + for (uint64_t task = coreIdx; task < taskNum; task += coreNum) { + uint64_t seq = 0; + uint64_t b = 0; + uint64_t h = 0; + uint64_t hv = 0; + uint64_t chunkIdx = 0; + uint64_t start = 0; + uint64_t end = 0; + if (ResolveFlatChunk(task, seq, b, h, hv, chunkIdx, start, end)) { + (void)seq; + ProcessChunkPreAiv(b, h, hv, chunkIdx, start, end, subBlockIdx, subBlockNum); + } + } + } + + __aicore__ inline void ProcessPreAic() + { + if constexpr (IsSameType::value) { + return; + } +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if constexpr (SAFE_GATE) { + if (headPairMode_) { + ProcessPreAicHeadPair(); + return; + } + } +#endif + uint64_t taskNum = static_cast((isVarLen_ ? NT_ : B_ * NT_) * HV_); + uint64_t coreNum = usedCoreNum_ == 0 ? 1 : usedCoreNum_; + uint64_t localTaskIdx = 0; + for (uint64_t task = GetBlockIdx(); task < taskNum; task += coreNum, ++localTaskIdx) { + uint64_t seq = 0; + uint64_t b = 0; + uint64_t h = 0; + uint64_t hv = 0; + uint64_t chunkIdx = 0; + uint64_t start = 0; + uint64_t end = 0; + if (ResolveFlatChunk(task, seq, b, h, hv, chunkIdx, start, end)) { + if constexpr (SAFE_GATE) { + activeSolveSlot_ = localTaskIdx % KDA_SOLVE_PIPELINE_DEPTH; + } + (void)seq; + (void)h; + ProcessChunkPreAic(b, hv, chunkIdx, start, end); + } + } + } + +private: + GlobalTensor q_; + GlobalTensor k_; + GlobalTensor v_; + GlobalTensor gk_; + GlobalTensor rawG_; + GlobalTensor aLog_; + GlobalTensor dtBias_; + GlobalTensor beta_; + GlobalTensor initialState_; + GlobalTensor o_; + GlobalTensor finalState_; + GlobalTensor aqk_; + GlobalTensor akk_; + GlobalTensor w_; + GlobalTensor u_; + GlobalTensor qg_; + GlobalTensor kg_; + GlobalTensor vNew_; + GlobalTensor h_; + GlobalTensor finalKg_; + GlobalTensor preparedQG_; + GlobalTensor preparedAqk_; + GlobalTensor propagatedVNew_; + GlobalTensor propagatedH_; + GlobalTensor solveWorkspace_; + GlobalTensor scoreWorkspace_; + TPipe *pipe_ = nullptr; + TBuf exp2Buf_; + TBuf vecBuf_; + TBuf gateWritebackBuf_; + TEventID mte2ToVEvent_ = 0; + TEventID vToMte2Event_ = 0; + TEventID vToMte3Event_ = 0; + TEventID mte3ToVEvent_ = 0; + TEventID mte2ToMte3Event_ = 0; + TEventID vToSEvent_ = 0; + TEventID mte3ToMte2Events_[KDA_GATE_PIPELINE_DEPTH] = {0, 0, 0}; + bool vectorEventsAllocated_ = false; + Catlass::Arch::CrossCoreFlagWithReverse scoreReadyFlag_{KDA_SCORE_READY_FLAG0, + KDA_SCORE_READY_FLAG1}; + Catlass::Arch::CrossCoreFlagWithReverse scoreDoneFlag_{KDA_SCORE_DONE_FLAG0, + KDA_SCORE_DONE_FLAG1}; + // Solve has one outstanding task per core. Reuse the primary score IDs as an ordered token stream; + // score credits remain on the reverse IDs, so no additional hardware flag IDs are consumed. + Catlass::Arch::CrossCoreFlag syncReadyFlag_{KDA_SOLVE_READY_FLAG}; + Catlass::Arch::CrossCoreFlag syncDoneFlag_{KDA_SOLVE_DONE_FLAG}; + Catlass::Arch::CrossCoreFlagWithReverse postWuReadyFlag_{KDA_POST_READY_FLAG, + KDA_POST_FREE_FLAG}; + Catlass::Arch::CrossCoreFlagWithReverse mchSyncReadyFlag_{KDA_SCORE_READY_FLAG0, + KDA_SCORE_READY_FLAG1}; + Catlass::Arch::CrossCoreFlagWithReverse mchSyncDoneFlag_{KDA_SCORE_DONE_FLAG0, + KDA_SCORE_DONE_FLAG1}; + uint64_t B_ = 0; + uint64_t N_ = 0; + uint64_t H_ = 0; + uint64_t HV_ = 0; + uint64_t T_ = 0; + uint64_t K_ = 0; + uint64_t V_ = 0; + uint64_t BT_ = 0; + uint64_t NT_ = 0; + float scale_ = 1.0f; + bool hasInitial_ = false; + bool isVarLen_ = false; + bool isAivOnly_ = false; + bool headPairMode_ = false; + bool inputSequenceMajor_ = false; + bool fusePostWu_ = false; + bool computeGateInPrepare_ = false; + bool hasALog_ = false; + bool hasDtBias_ = false; + bool storeQG_ = true; + float lowerBound_ = -5.0f; + uint64_t usedCoreNum_ = 1; + uint64_t solveCoreIdx_ = 0; + uint64_t activeSolveSlot_ = 0; + uint64_t activeGateChunkStart_ = 0; + __gm__ int64_t *chunkIndicesAddr_ = nullptr; + __gm__ int64_t *cuSeqlensAddr_ = nullptr; +}; +} // namespace + +template +__aicore__ inline void RunChunkKdaPrepare(GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR gk, GM_ADDR rawG, GM_ADDR aLog, + GM_ADDR dtBias, GM_ADDR beta, GM_ADDR initialState, GM_ADDR cuSeqlens, + GM_ADDR chunkIndices, GM_ADDR aqk, GM_ADDR akk, GM_ADDR qg, GM_ADDR qgScaled, + GM_ADDR wSeed, GM_ADDR uSeed, GM_ADDR finalKg, GM_ADDR userWorkspace, + const TilingData &tiling, TPipe &pipe, bool storeQG = true) +{ + GM_ADDR aqkFp32 = userWorkspace + tiling.prepareAqkFp32Offset; + GM_ADDR akkFp32 = userWorkspace + tiling.prepareAkkFp32Offset; + GM_ADDR prepareScratch = userWorkspace + tiling.prepareScratchOffset; + + if ASCEND_IS_AIC { + ChunkKdaFwdPrepareKernel op; + op.Init(q, k, v, gk, rawG, aLog, dtBias, beta, initialState, cuSeqlens, chunkIndices, nullptr, nullptr, nullptr, + nullptr, aqk, userWorkspace, aqkFp32, akkFp32, wSeed, akk, qg, qgScaled, uSeed, userWorkspace, finalKg, + prepareScratch, tiling, &pipe, false, storeQG); + if (tiling.fusePostWu) { +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + KdaPostWu::ChunkKdaFwdPostWuKernel postWu; + postWu.Init(nullptr, k, nullptr, gk, beta, initialState, cuSeqlens, chunkIndices, wSeed, akk, uSeed, + nullptr, userWorkspace, userWorkspace, userWorkspace, akk, wSeed, uSeed, userWorkspace, finalKg, + userWorkspace, prepareScratch, prepareScratch, tiling, &pipe, false); + op.ProcessAicFused(postWu); +#else + op.ProcessAic(); +#endif + } else { + op.ProcessAic(); + } + } + if ASCEND_IS_AIV { + ChunkKdaFwdPrepareKernel op; + op.Init(q, k, v, gk, rawG, aLog, dtBias, beta, initialState, cuSeqlens, chunkIndices, nullptr, nullptr, nullptr, + nullptr, aqk, userWorkspace, aqkFp32, akkFp32, wSeed, akk, qg, qgScaled, uSeed, userWorkspace, finalKg, + prepareScratch, tiling, &pipe, true, storeQG); + op.ProcessAiv(); + } +} + +} // namespace KdaPrepare diff --git a/csrc/chunk_kda_fwd/op_kernel/arch35/fwd_h/epilogue/block/block_epilogue_kda_fwdh_regbase.hpp b/csrc/chunk_kda_fwd/op_kernel/arch35/fwd_h/epilogue/block/block_epilogue_kda_fwdh_regbase.hpp new file mode 100644 index 0000000000..be8af460d3 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/arch35/fwd_h/epilogue/block/block_epilogue_kda_fwdh_regbase.hpp @@ -0,0 +1,381 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef BLOCK_EPILOGUE_KDA_FWDH_REGBASE_HPP +#define BLOCK_EPILOGUE_KDA_FWDH_REGBASE_HPP + +#include "kernel_operator.h" +#ifndef FLA_NPU_REGBASE_HPP_INCLUDED +#define FLA_NPU_REGBASE_HPP_INCLUDED +#include "../../../../kernel_utils/vector/regbase.hpp" +#endif + +namespace Catlass::Epilogue::Block::detail { + +using namespace AscendC::MicroAPI; + +constexpr CastTrait KDA_B16_TO_F32_ZERO = { + RegLayout::ZERO, + SatMode::UNKNOWN, + MaskMergeMode::ZEROING, + AscendC::RoundMode::UNKNOWN, +}; + +template +__simd_callee__ inline void LoadKdaAsFloat( + RegTensor &dst, __ubuf__ T *src, MaskReg &mask) +{ + if constexpr (std::is_same()) { + DataCopy(dst, src); + } else if constexpr (std::is_same() || std::is_same()) { + RegTensor raw; + DataCopy(raw, src); + Cast(dst, raw, mask); + } else { + static_assert(!std::is_same::value, "KDA regbase only supports float/half/bfloat16_t"); + } +} + +template +__simd_callee__ inline void LoadKdaBroadcastAsFloat( + RegTensor &dst, __ubuf__ T *src, MaskReg &mask) +{ + if constexpr (std::is_same()) { + LoadAlign(dst, src); + } else if constexpr (std::is_same() || std::is_same()) { + RegTensor raw; + LoadAlign(raw, src); + Cast(dst, raw, mask); + } else { + static_assert(!std::is_same::value, "KDA regbase only supports float/half/bfloat16_t"); + } +} + +__simd_callee__ inline void LoadKdaFloat(RegTensor &dst, __ubuf__ float *src) +{ + DataCopy(dst, src); +} + +__simd_callee__ inline void StoreKdaFloat( + __ubuf__ float *dst, RegTensor &src, MaskReg &mask) +{ + DataCopy(dst, src, mask); +} + +template +static __simd_vf__ inline void PrepareKGateRegbase( + __ubuf__ float *gateOutput, __ubuf__ T *gateInput, uint16_t count) +{ + constexpr uint16_t FP32_PER_REG = AscendC::VECTOR_REG_WIDTH / sizeof(float); + constexpr float LN2 = 0.6931471805599453f; + RegTensor gateReg; + MaskReg mask; + uint32_t remaining = count; + uint32_t offset = 0; + while (remaining > 0) { + mask = UpdateMask(remaining); + LoadKdaAsFloat(gateReg, gateInput + offset, mask); + if constexpr (USE_EXP2) { + Muls(gateReg, gateReg, LN2, mask); + } + Exp(gateReg, gateReg, mask); + StoreKdaFloat(gateOutput + offset, gateReg, mask); + offset += FP32_PER_REG; + } +} + +template +static __simd_vf__ inline void ComputeVNewRegbaseDualIssue( + __ubuf__ float *workspace, __ubuf__ T *uInput, uint32_t count) +{ + constexpr uint32_t FP32_PER_REG = AscendC::VECTOR_REG_WIDTH / sizeof(float); + RegTensor uReg0; + RegTensor uReg1; + RegTensor wsReg0; + RegTensor wsReg1; + MaskReg mask0; + MaskReg mask1; + + uint32_t remaining = count; + uint32_t offset = 0; + while (remaining > FP32_PER_REG) { + mask0 = UpdateMask(remaining); + mask1 = UpdateMask(remaining); + LoadKdaAsFloat(uReg0, uInput + offset, mask0); + LoadKdaAsFloat(uReg1, uInput + offset + FP32_PER_REG, mask1); + LoadKdaFloat(wsReg0, workspace + offset); + LoadKdaFloat(wsReg1, workspace + offset + FP32_PER_REG); + Sub(wsReg0, uReg0, wsReg0, mask0); + Sub(wsReg1, uReg1, wsReg1, mask1); + StoreKdaFloat(workspace + offset, wsReg0, mask0); + StoreKdaFloat(workspace + offset + FP32_PER_REG, wsReg1, mask1); + offset += 2 * FP32_PER_REG; + } + if (remaining > 0) { + mask0 = UpdateMask(remaining); + LoadKdaAsFloat(uReg0, uInput + offset, mask0); + LoadKdaFloat(wsReg0, workspace + offset); + Sub(wsReg0, uReg0, wsReg0, mask0); + StoreKdaFloat(workspace + offset, wsReg0, mask0); + } +} + +template +static __simd_vf__ inline void ApplyKGateUpdateRegbaseDualIssue( + __ubuf__ float *update, __ubuf__ T *state, __ubuf__ float *rowScale, + uint16_t rows, uint16_t cols) +{ + constexpr uint16_t FP32_PER_REG = AscendC::VECTOR_REG_WIDTH / sizeof(float); + RegTensor stateReg0; + RegTensor stateReg1; + RegTensor updateReg0; + RegTensor updateReg1; + RegTensor scaleReg0; + RegTensor scaleReg1; + MaskReg mask0; + MaskReg mask1; + + uint16_t row = 0; + for (; row + 1 < rows; row += 2) { + LoadAlign(scaleReg0, rowScale + row); + LoadAlign(scaleReg1, rowScale + row + 1); + uint32_t remaining0 = cols; + uint32_t remaining1 = cols; + uint16_t colLoops = static_cast((cols + FP32_PER_REG - 1) / FP32_PER_REG); + for (uint16_t colLoop = 0; colLoop < colLoops; ++colLoop) { + uint32_t col = static_cast(colLoop) * FP32_PER_REG; + mask0 = UpdateMask(remaining0); + mask1 = UpdateMask(remaining1); + LoadKdaAsFloat(stateReg0, state + static_cast(row) * cols + col, mask0); + LoadKdaAsFloat(stateReg1, state + static_cast(row + 1) * cols + col, mask1); + LoadKdaFloat(updateReg0, update + static_cast(row) * cols + col); + LoadKdaFloat(updateReg1, update + static_cast(row + 1) * cols + col); + Mul(stateReg0, stateReg0, scaleReg0, mask0); + Mul(stateReg1, stateReg1, scaleReg1, mask1); + Add(updateReg0, stateReg0, updateReg0, mask0); + Add(updateReg1, stateReg1, updateReg1, mask1); + StoreKdaFloat(update + static_cast(row) * cols + col, updateReg0, mask0); + StoreKdaFloat(update + static_cast(row + 1) * cols + col, updateReg1, mask1); + } + } + + if (row < rows) { + LoadAlign(scaleReg0, rowScale + row); + uint32_t remaining = cols; + uint16_t colLoops = static_cast((cols + FP32_PER_REG - 1) / FP32_PER_REG); + for (uint16_t colLoop = 0; colLoop < colLoops; ++colLoop) { + uint32_t col = static_cast(colLoop) * FP32_PER_REG; + mask0 = UpdateMask(remaining); + LoadKdaAsFloat(stateReg0, state + static_cast(row) * cols + col, mask0); + LoadKdaFloat(updateReg0, update + static_cast(row) * cols + col); + Mul(stateReg0, stateReg0, scaleReg0, mask0); + Add(updateReg0, stateReg0, updateReg0, mask0); + StoreKdaFloat(update + static_cast(row) * cols + col, updateReg0, mask0); + } + } +} + +template +static __simd_vf__ inline void ComputeKdaTailFinalStateRegbaseDualIssue( + __ubuf__ float *state, __ubuf__ KType *kInput, __ubuf__ VType *vInput, + __ubuf__ float *rowScale, uint16_t rows, uint16_t cols, + uint16_t tokens, uint16_t kStride, uint16_t kRowOffset) +{ + constexpr uint16_t FP32_PER_REG = AscendC::VECTOR_REG_WIDTH / sizeof(float); + RegTensor stateReg0; + RegTensor stateReg1; + RegTensor valueReg; + RegTensor productReg0; + RegTensor productReg1; + RegTensor weightReg0; + RegTensor weightReg1; + RegTensor scaleReg0; + RegTensor scaleReg1; + MaskReg mask0; + MaskReg mask1; + + uint16_t row = 0; + for (; row + 1 < rows; row += 2) { + LoadAlign(scaleReg0, rowScale + row); + LoadAlign(scaleReg1, rowScale + row + 1); + uint32_t remaining0 = cols; + uint32_t remaining1 = cols; + uint16_t colLoops = static_cast((cols + FP32_PER_REG - 1) / FP32_PER_REG); + for (uint16_t colLoop = 0; colLoop < colLoops; ++colLoop) { + uint32_t col = static_cast(colLoop) * FP32_PER_REG; + mask0 = UpdateMask(remaining0); + mask1 = UpdateMask(remaining1); + LoadKdaFloat(stateReg0, state + static_cast(row) * cols + col); + LoadKdaFloat(stateReg1, state + static_cast(row + 1) * cols + col); + Mul(stateReg0, stateReg0, scaleReg0, mask0); + Mul(stateReg1, stateReg1, scaleReg1, mask1); + for (uint16_t token = 0; token < tokens; ++token) { + uint32_t kOffset = static_cast(token) * kStride + kRowOffset + row; + LoadKdaBroadcastAsFloat(weightReg0, kInput + kOffset, mask0); + LoadKdaBroadcastAsFloat(weightReg1, kInput + kOffset + 1, mask1); + LoadKdaAsFloat( + valueReg, vInput + static_cast(token) * cols + col, mask0); + Mul(productReg0, valueReg, weightReg0, mask0); + Mul(productReg1, valueReg, weightReg1, mask1); + Add(stateReg0, stateReg0, productReg0, mask0); + Add(stateReg1, stateReg1, productReg1, mask1); + } + StoreKdaFloat(state + static_cast(row) * cols + col, stateReg0, mask0); + StoreKdaFloat(state + static_cast(row + 1) * cols + col, stateReg1, mask1); + } + } + + if (row < rows) { + LoadAlign(scaleReg0, rowScale + row); + uint32_t remaining = cols; + uint16_t colLoops = static_cast((cols + FP32_PER_REG - 1) / FP32_PER_REG); + for (uint16_t colLoop = 0; colLoop < colLoops; ++colLoop) { + uint32_t col = static_cast(colLoop) * FP32_PER_REG; + mask0 = UpdateMask(remaining); + LoadKdaFloat(stateReg0, state + static_cast(row) * cols + col); + Mul(stateReg0, stateReg0, scaleReg0, mask0); + for (uint16_t token = 0; token < tokens; ++token) { + uint32_t kOffset = static_cast(token) * kStride + kRowOffset + row; + LoadKdaBroadcastAsFloat(weightReg0, kInput + kOffset, mask0); + LoadKdaAsFloat( + valueReg, vInput + static_cast(token) * cols + col, mask0); + Mul(productReg0, valueReg, weightReg0, mask0); + Add(stateReg0, stateReg0, productReg0, mask0); + } + StoreKdaFloat(state + static_cast(row) * cols + col, stateReg0, mask0); + } + } +} + +template +static __simd_vf__ inline void ComputeKdaTailVWorkspaceRegbaseDualIssue( + __ubuf__ float *output, __ubuf__ WType *wInput, __ubuf__ HType *hInput, + uint16_t rows, uint16_t cols, uint16_t reduction) +{ + constexpr uint16_t FP32_PER_REG = AscendC::VECTOR_REG_WIDTH / sizeof(float); + RegTensor accumReg0; + RegTensor accumReg1; + RegTensor valueReg; + RegTensor productReg0; + RegTensor productReg1; + RegTensor weightReg0; + RegTensor weightReg1; + MaskReg mask0; + MaskReg mask1; + + uint16_t row = 0; + for (; row + 1 < rows; row += 2) { + uint32_t remaining0 = cols; + uint32_t remaining1 = cols; + uint16_t colLoops = static_cast((cols + FP32_PER_REG - 1) / FP32_PER_REG); + for (uint16_t colLoop = 0; colLoop < colLoops; ++colLoop) { + uint32_t col = static_cast(colLoop) * FP32_PER_REG; + mask0 = UpdateMask(remaining0); + mask1 = UpdateMask(remaining1); + Duplicate(accumReg0, 0.0f, mask0); + Duplicate(accumReg1, 0.0f, mask1); + for (uint16_t reductionIdx = 0; reductionIdx < reduction; ++reductionIdx) { + LoadKdaBroadcastAsFloat( + weightReg0, + wInput + static_cast(row) * reduction + reductionIdx, + mask0); + LoadKdaBroadcastAsFloat( + weightReg1, + wInput + static_cast(row + 1) * reduction + reductionIdx, + mask1); + LoadKdaAsFloat( + valueReg, + hInput + static_cast(reductionIdx) * cols + col, + mask0); + Mul(productReg0, valueReg, weightReg0, mask0); + Mul(productReg1, valueReg, weightReg1, mask1); + Add(accumReg0, accumReg0, productReg0, mask0); + Add(accumReg1, accumReg1, productReg1, mask1); + } + StoreKdaFloat(output + static_cast(row) * cols + col, accumReg0, mask0); + StoreKdaFloat(output + static_cast(row + 1) * cols + col, accumReg1, mask1); + } + } + + if (row < rows) { + uint32_t remaining = cols; + uint16_t colLoops = static_cast((cols + FP32_PER_REG - 1) / FP32_PER_REG); + for (uint16_t colLoop = 0; colLoop < colLoops; ++colLoop) { + uint32_t col = static_cast(colLoop) * FP32_PER_REG; + mask0 = UpdateMask(remaining); + Duplicate(accumReg0, 0.0f, mask0); + for (uint16_t reductionIdx = 0; reductionIdx < reduction; ++reductionIdx) { + LoadKdaBroadcastAsFloat( + weightReg0, + wInput + static_cast(row) * reduction + reductionIdx, + mask0); + LoadKdaAsFloat( + valueReg, + hInput + static_cast(reductionIdx) * cols + col, + mask0); + Mul(productReg0, valueReg, weightReg0, mask0); + Add(accumReg0, accumReg0, productReg0, mask0); + } + StoreKdaFloat(output + static_cast(row) * cols + col, accumReg0, mask0); + } + } +} + +static __simd_vf__ inline void ApplyRowScaleDualIssue( + __ubuf__ float *matrix, __ubuf__ float *rowScale, uint32_t rowScaleOffset, + uint16_t rows, uint16_t cols) +{ + using namespace AscendC::MicroAPI; + constexpr uint16_t FP32_PER_REG = AscendC::VECTOR_REG_WIDTH / sizeof(float); + + RegTensor matrixReg0; + RegTensor matrixReg1; + RegTensor scaleReg0; + RegTensor scaleReg1; + MaskReg mask0; + MaskReg mask1; + + uint16_t row = 0; + for (; row + 1 < rows; row += 2) { + LoadAlign(scaleReg0, rowScale + rowScaleOffset + row); + LoadAlign(scaleReg1, rowScale + rowScaleOffset + row + 1); + uint32_t remaining0 = cols; + uint32_t remaining1 = cols; + uint16_t colLoops = static_cast((cols + FP32_PER_REG - 1) / FP32_PER_REG); + for (uint16_t colLoop = 0; colLoop < colLoops; ++colLoop) { + uint32_t col = static_cast(colLoop) * FP32_PER_REG; + mask0 = UpdateMask(remaining0); + mask1 = UpdateMask(remaining1); + LoadAlign(matrixReg0, matrix + static_cast(row) * cols + col); + LoadAlign(matrixReg1, matrix + static_cast(row + 1) * cols + col); + Mul(matrixReg0, matrixReg0, scaleReg0, mask0); + Mul(matrixReg1, matrixReg1, scaleReg1, mask1); + StoreAlign(matrix + static_cast(row) * cols + col, matrixReg0, mask0); + StoreAlign(matrix + static_cast(row + 1) * cols + col, matrixReg1, mask1); + } + } + + if (row < rows) { + LoadAlign(scaleReg0, rowScale + rowScaleOffset + row); + uint32_t remaining = cols; + uint16_t colLoops = static_cast((cols + FP32_PER_REG - 1) / FP32_PER_REG); + for (uint16_t colLoop = 0; colLoop < colLoops; ++colLoop) { + uint32_t col = static_cast(colLoop) * FP32_PER_REG; + mask0 = UpdateMask(remaining); + LoadAlign(matrixReg0, matrix + static_cast(row) * cols + col); + Mul(matrixReg0, matrixReg0, scaleReg0, mask0); + StoreAlign(matrix + static_cast(row) * cols + col, matrixReg0, mask0); + } + } +} + +} // namespace Catlass::Epilogue::Block::detail + +#endif diff --git a/csrc/chunk_kda_fwd/op_kernel/arch35/fwd_h/epilogue/block/block_epilogue_kda_fwdh_update_arch35.hpp b/csrc/chunk_kda_fwd/op_kernel/arch35/fwd_h/epilogue/block/block_epilogue_kda_fwdh_update_arch35.hpp new file mode 100644 index 0000000000..0346024e38 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/arch35/fwd_h/epilogue/block/block_epilogue_kda_fwdh_update_arch35.hpp @@ -0,0 +1,434 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_KDA_FWDH_UPDATE_ARCH35_HPP +#define CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_KDA_FWDH_UPDATE_ARCH35_HPP +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "../kda_fwd_h_epilogue_policies.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/kda_matrix_coord.hpp" +#include "catlass/epilogue/tile/kda_epilogue_tile_copy.hpp" +#include "block_epilogue_kda_fwdh_regbase.hpp" + +namespace Catlass::Epilogue::Block { + +template +class BlockEpilogue { + static constexpr bool kGated = KGatedTag::value; + static constexpr bool scalarGated = KGatedTag::scalarGated; + static constexpr bool useExp2 = KGatedTag::useExp2; + static constexpr float LN2 = 0.6931471805599453f; + +public: + // Type aliases + using DispatchPolicy = EpilogueAtlasKdaFwdHUpdate; + using ArchTag = typename DispatchPolicy::ArchTag; + + using HElementOutput = typename HOutputType_::Element; + using GElementInput = typename GInputType_::Element; + using HElementInput = typename HInputType_::Element; + using HUpdateElementInput = typename HUpdateInputType_::Element; + using FinalStateElement = typename FinalStateType_::Element; + + CATLASS_DEVICE + BlockEpilogue(Arch::Resource &resource) + { + constexpr uint32_t CALC_BUF_OFFSET = 0; + constexpr uint32_t PING_BUF_0_OFFSET = 32 * 1024; + constexpr uint32_t PING_BUF_1_OFFSET = 48 * 1024; + constexpr uint32_t PING_BUF_2_OFFSET = 64 * 1024; + constexpr uint32_t PING_BUF_3_OFFSET = 80 * 1024; + constexpr uint32_t PONG_BUF_0_OFFSET = 96 * 1024; + constexpr uint32_t PONG_BUF_1_OFFSET = 112 * 1024; + constexpr uint32_t PONG_BUF_2_OFFSET = 128 * 1024; + constexpr uint32_t PONG_BUF_3_OFFSET = 144 * 1024; + constexpr uint32_t PING_G_BUF_OFFSET = 168 * 1024; + constexpr uint32_t PONG_G_BUF_OFFSET = 169 * 1024; + constexpr uint32_t PING_G_SUB_BUF_OFFSET = 170 * 1024; + constexpr uint32_t PONG_G_SUB_BUF_OFFSET = 171 * 1024; + constexpr uint32_t PING_G_INPUT_BUF_OFFSET = 172 * 1024; + constexpr uint32_t PONG_G_INPUT_BUF_OFFSET = 173 * 1024; + constexpr uint32_t UPDATE_SCRATCH_BUF_OFFSET = 160 * 1024; + constexpr uint32_t UPDATE_G_BUF_OFFSET = 176 * 1024; + + calcUbTensor = resource.ubBuf.template GetBufferByByte(CALC_BUF_OFFSET); + + hUpdateUbTensor_ping = resource.ubBuf.template GetBufferByByte(PING_BUF_0_OFFSET); + hUbTensor_ping = resource.ubBuf.template GetBufferByByte(UPDATE_SCRATCH_BUF_OFFSET); + finalOutputUbTensor_ping = + resource.ubBuf.template GetBufferByByte(UPDATE_SCRATCH_BUF_OFFSET); + glastUbTensor_ping = resource.ubBuf.template GetBufferByByte(UPDATE_G_BUF_OFFSET); + + hUpdateUbTensor_pong = resource.ubBuf.template GetBufferByByte(PONG_BUF_0_OFFSET); + hUbTensor_pong = resource.ubBuf.template GetBufferByByte(UPDATE_SCRATCH_BUF_OFFSET); + finalOutputUbTensor_pong = + resource.ubBuf.template GetBufferByByte(UPDATE_SCRATCH_BUF_OFFSET); + glastUbTensor_pong = resource.ubBuf.template GetBufferByByte(UPDATE_G_BUF_OFFSET); + + if constexpr (kGated) { + gkLastUbTensor_ping = resource.ubBuf.template GetBufferByByte(PING_G_SUB_BUF_OFFSET); + gkLastUbTensor_pong = resource.ubBuf.template GetBufferByByte(PONG_G_SUB_BUF_OFFSET); + gkInputUbTensor_ping = resource.ubBuf.template GetBufferByByte(PING_G_INPUT_BUF_OFFSET); + gkInputUbTensor_pong = resource.ubBuf.template GetBufferByByte(PONG_G_INPUT_BUF_OFFSET); + gkBrcbUbTensor_ping = resource.ubBuf.template GetBufferByByte(PING_G_INPUT_BUF_OFFSET); + gkBrcbUbTensor_pong = resource.ubBuf.template GetBufferByByte(PONG_G_INPUT_BUF_OFFSET); + } + } + + CATLASS_DEVICE + ~BlockEpilogue() {} + + template + CATLASS_DEVICE void CopyGmToUb(AscendC::LocalTensor dst, AscendC::GlobalTensor src, uint32_t rows, + uint32_t cols, uint32_t srcStride) + { + if (cols == srcStride) { + AscendC::DataCopy(dst, src, rows * cols); + return; + } + AscendC::DataCopyExtParams copyParams{static_cast(rows), + static_cast(cols * sizeof(Element)), + static_cast((srcStride - cols) * sizeof(Element)), 0, 0}; + AscendC::DataCopyPadExtParams padParams{false, 0, 0, 0}; + AscendC::DataCopyPad(dst, src, copyParams, padParams); + } + + template + CATLASS_DEVICE void CopyUbToGm(AscendC::GlobalTensor dst, AscendC::LocalTensor src, uint32_t rows, + uint32_t cols, uint32_t dstStride) + { + if (cols == dstStride) { + AscendC::DataCopy(dst, src, rows * cols); + return; + } + AscendC::DataCopyExtParams copyParams{static_cast(rows), + static_cast(cols * sizeof(Element)), 0, + static_cast((dstStride - cols) * sizeof(Element)), 0}; + AscendC::DataCopyPad(dst, src, copyParams); + } + + CATLASS_DEVICE + void ApplyRowScale(AscendC::LocalTensor matrix, AscendC::LocalTensor rowScale, uint32_t rows, + uint32_t cols) + { + __ubuf__ float *matrixAddr = reinterpret_cast<__ubuf__ float *>(matrix.GetPhyAddr()); + __ubuf__ float *rowScaleAddr = reinterpret_cast<__ubuf__ float *>(rowScale.GetPhyAddr()); + AscendC::VF_CALL(matrixAddr, rowScaleAddr, 0, static_cast(rows), + static_cast(cols)); + AscendC::PipeBarrier(); + } + + CATLASS_DEVICE + void PrepareKGate(AscendC::LocalTensor gateOutput, AscendC::LocalTensor gateInput, + uint32_t count) + { + __ubuf__ float *gateOutputAddr = reinterpret_cast<__ubuf__ float *>(gateOutput.GetPhyAddr()); + __ubuf__ GElementInput *gateInputAddr = reinterpret_cast<__ubuf__ GElementInput *>(gateInput.GetPhyAddr()); + AscendC::VF_CALL>(gateOutputAddr, gateInputAddr, + static_cast(count)); + AscendC::PipeBarrier(); + } + + template + CATLASS_DEVICE void ApplyKGateUpdate(AscendC::LocalTensor update, AscendC::LocalTensor state, + AscendC::LocalTensor rowScale, uint32_t rows, uint32_t cols) + { + __ubuf__ float *updateAddr = reinterpret_cast<__ubuf__ float *>(update.GetPhyAddr()); + __ubuf__ StateElement *stateAddr = reinterpret_cast<__ubuf__ StateElement *>(state.GetPhyAddr()); + __ubuf__ float *rowScaleAddr = reinterpret_cast<__ubuf__ float *>(rowScale.GetPhyAddr()); + AscendC::VF_CALL>( + updateAddr, stateAddr, rowScaleAddr, static_cast(rows), static_cast(cols)); + AscendC::PipeBarrier(); + } + + CATLASS_DEVICE + void operator()(AscendC::GlobalTensor hOutput, AscendC::GlobalTensor finalState, + AscendC::GlobalTensor gInput, AscendC::GlobalTensor hInput, + AscendC::GlobalTensor hUpdateInput, AscendC::GlobalTensor gkInput, + AscendC::GlobalTensor initialState, uint32_t chunkSize, uint32_t kHeadDim, + uint32_t vBlockDim, uint32_t vHeadDim, Arch::CrossCoreFlag cube2Done, bool isInitialState, + bool isFinalState, bool storeFinalState, bool useInitialState, bool isPing, bool cube2AlreadyWaited, + bool useDirectFp32Ub, uint64_t directUbFreeFlagBegin, uint64_t directUbReadyFlagBegin, + uint64_t directUbFlagStride) + { + static constexpr uint32_t ROW_TILE = 16; + uint32_t mActual = kHeadDim; + uint32_t nActual = vBlockDim; + uint32_t outputStride = vHeadDim; + uint32_t subBlockIdx = AscendC::GetSubBlockIdx(); + uint32_t subBlockNum = AscendC::GetSubBlockNum(); + uint32_t rowsPerSubBlock = CeilDiv(mActual, subBlockNum); + uint32_t rowBegin = subBlockIdx * rowsPerSubBlock; + uint32_t rowEnd = rowBegin + rowsPerSubBlock; + if (rowEnd > mActual) { + rowEnd = mActual; + } + if (rowBegin >= mActual) { + if (useDirectFp32Ub) { + uint32_t directUbSlot = isPing ? 0 : 1; + uint64_t directUbFlagOffset = directUbSlot + AscendC::GetSubBlockIdx() * directUbFlagStride; + AscendC::CrossCoreWaitFlag<0x4, PIPE_V>(directUbReadyFlagBegin + directUbFlagOffset); + AscendC::CrossCoreSetFlag<0x4, PIPE_V>(directUbFreeFlagBegin + directUbFlagOffset); + } else if (!cube2AlreadyWaited) { + Arch::CrossCoreWaitFlag(cube2Done); + } + return; + } + + AscendC::ResetMask(); + + AscendC::GlobalTensor gInputThisSubBlock = gInput; + + uint32_t pingpongFlag = isPing ? 0 : pongBaseEvent; + AscendC::LocalTensor hUpdateUbTensor = isPing ? hUpdateUbTensor_ping : hUpdateUbTensor_pong; + AscendC::LocalTensor hUbTensor = isPing ? hUbTensor_ping : hUbTensor_pong; + AscendC::LocalTensor finalOutputUbTensor = + isPing ? finalOutputUbTensor_ping : finalOutputUbTensor_pong; + AscendC::LocalTensor glastUbTensor = isPing ? glastUbTensor_ping : glastUbTensor_pong; + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID2 + pingpongFlag); + bool useFp32StateUpdate = + storeFinalState && std::is_same::value && (!isInitialState || useInitialState); + float muls = 1.0f; + if constexpr (scalarGated) { + GElementInput gLastVal = gInputThisSubBlock.GetValue(chunkSize - 1); + float gLastFloat = 0.0f; + if constexpr (std::is_same::value) { + gLastFloat = gLastVal; + } else if constexpr (std::is_same::value) { + gLastFloat = (float)gLastVal; + } else if constexpr (std::is_same::value) { + gLastFloat = AscendC::ToFloat(gLastVal); + } + glastUbTensor.SetValue(0, gLastFloat); + + AscendC::SetFlag(EVENT_ID3 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID3 + pingpongFlag); + if constexpr (useExp2) { + AscendC::Muls(glastUbTensor, glastUbTensor, LN2, 1); + AscendC::PipeBarrier(); + } + AscendC::Exp(glastUbTensor, glastUbTensor, 1); + AscendC::SetFlag(EVENT_ID3 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID3 + pingpongFlag); + muls = glastUbTensor.GetValue(0); + } + if constexpr (kGated) { + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + } + if constexpr (scalarGated) { + AscendC::SetFlag(EVENT_ID3 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID3 + pingpongFlag); + } + + if (useDirectFp32Ub) { + uint32_t directUbSlot = isPing ? 0 : 1; + uint64_t directUbFlagOffset = directUbSlot + AscendC::GetSubBlockIdx() * directUbFlagStride; + AscendC::CrossCoreWaitFlag<0x4, PIPE_V>(directUbReadyFlagBegin + directUbFlagOffset); + } else if (!cube2AlreadyWaited) { + Arch::CrossCoreWaitFlag(cube2Done); + } + // fix: need to adapt kGated. issue: A5 do not have vdim128 branch. + bool waitHFromV = storeFinalState && isInitialState && std::is_same::value; + bool waitUpdateFromMte3 = false; + uint32_t updateReadyEvent = EVENT_ID3 + pingpongFlag; + for (uint32_t rowStart = rowBegin; rowStart < rowEnd; rowStart += ROW_TILE) { + uint32_t rowsThisTile = rowEnd - rowStart; + if (rowsThisTile > ROW_TILE) { + rowsThisTile = ROW_TILE; + } + + AscendC::GlobalTensor hOutputThisTile = hOutput[rowStart * outputStride]; + AscendC::GlobalTensor hInputThisTile = hInput[rowStart * outputStride]; + AscendC::GlobalTensor hUpdateInputThisTile = hUpdateInput[rowStart * nActual]; + AscendC::GlobalTensor finalStateThisTile = finalState[rowStart * outputStride]; + AscendC::LocalTensor hUpdateUbTensorThisTile = + useDirectFp32Ub ? hUpdateUbTensor[(rowStart - rowBegin) * nActual] : hUpdateUbTensor; + + if (waitHFromV) { + AscendC::WaitFlag(EVENT_ID2 + pingpongFlag); + } else { + AscendC::WaitFlag(EVENT_ID2 + pingpongFlag); + } + if constexpr (std::is_same::value) { + if (useFp32StateUpdate) { + if (isInitialState) { + CopyGmToUb(calcUbTensor, initialState[rowStart * outputStride], rowsThisTile, nActual, + outputStride); + } else { + CopyGmToUb(calcUbTensor, finalStateThisTile, rowsThisTile, nActual, outputStride); + } + } else { + CopyGmToUb(hUbTensor, hInputThisTile, rowsThisTile, nActual, outputStride); + } + } else { + CopyGmToUb(hUbTensor, hInputThisTile, rowsThisTile, nActual, outputStride); + } + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID2 + pingpongFlag); + + if (!useFp32StateUpdate) { + AscendC::Cast(calcUbTensor, hUbTensor, AscendC::RoundMode::CAST_NONE, rowsThisTile * nActual); + AscendC::PipeBarrier(); + } + if constexpr (scalarGated) { + AscendC::Muls(calcUbTensor, calcUbTensor, muls, rowsThisTile * nActual); + AscendC::PipeBarrier(); + } + + if constexpr (kGated) { + AscendC::GlobalTensor gkLastInput = gkInput[(chunkSize - 1) * kHeadDim + rowStart]; + AscendC::LocalTensor gkLastUbTensor = isPing ? gkLastUbTensor_ping : gkLastUbTensor_pong; + AscendC::LocalTensor gkInputUbTensor = + isPing ? gkInputUbTensor_ping : gkInputUbTensor_pong; + if (rowStart == rowBegin) { + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + } else { + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + } + if constexpr (std::is_same::value) { + AscendC::DataCopy(gkLastUbTensor, gkLastInput, rowsThisTile); + } else { + AscendC::DataCopy(gkInputUbTensor, gkLastInput, rowsThisTile); + } + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID2 + pingpongFlag); + if constexpr (std::is_same::value) { + PrepareKGate(gkLastUbTensor, gkLastUbTensor, rowsThisTile); + } else { + PrepareKGate(gkLastUbTensor, gkInputUbTensor, rowsThisTile); + } + ApplyRowScale(calcUbTensor, gkLastUbTensor, rowsThisTile, nActual); + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + } + + if (waitUpdateFromMte3) { + AscendC::WaitFlag(updateReadyEvent); + } else { + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + } + if (!useDirectFp32Ub) { + CopyGmToUb(hUpdateUbTensorThisTile, hUpdateInputThisTile, rowsThisTile, nActual, nActual); + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + } + AscendC::Add(hUpdateUbTensorThisTile, calcUbTensor, hUpdateUbTensorThisTile, rowsThisTile * nActual); + AscendC::PipeBarrier(); + if (storeFinalState && isFinalState && std::is_same::value) { + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + waitHFromV = true; + } else { + waitHFromV = false; + } + + if constexpr (std::is_same::value) { + if (storeFinalState) { + if (!isFinalState) { + AscendC::Cast(hUbTensor, hUpdateUbTensorThisTile, AscendC::RoundMode::CAST_RINT, + rowsThisTile * nActual); + AscendC::PipeBarrier(); + } + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + CopyUbToGm(finalStateThisTile, hUpdateUbTensorThisTile, rowsThisTile, nActual, outputStride); + AscendC::PipeBarrier(); + AscendC::SetFlag(updateReadyEvent); + AscendC::SetFlag(updateReadyEvent); + AscendC::WaitFlag(updateReadyEvent); + waitUpdateFromMte3 = true; + if (!isFinalState) { + CopyUbToGm(hOutputThisTile, hUbTensor, rowsThisTile, nActual, outputStride); + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + } + } else { + AscendC::Cast(hUbTensor, hUpdateUbTensorThisTile, AscendC::RoundMode::CAST_RINT, + rowsThisTile * nActual); + AscendC::PipeBarrier(); + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID2 + pingpongFlag); + CopyUbToGm(hOutputThisTile, hUbTensor, rowsThisTile, nActual, outputStride); + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + waitUpdateFromMte3 = false; + } + } else { + if (storeFinalState && isFinalState) { + AscendC::Cast(finalOutputUbTensor, hUpdateUbTensorThisTile, AscendC::RoundMode::CAST_RINT, + rowsThisTile * nActual); + AscendC::PipeBarrier(); + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID2 + pingpongFlag); + CopyUbToGm(finalStateThisTile, finalOutputUbTensor, rowsThisTile, nActual, outputStride); + } else { + AscendC::Cast(hUbTensor, hUpdateUbTensorThisTile, AscendC::RoundMode::CAST_RINT, + rowsThisTile * nActual); + AscendC::PipeBarrier(); + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID2 + pingpongFlag); + CopyUbToGm(hOutputThisTile, hUbTensor, rowsThisTile, nActual, outputStride); + } + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + waitUpdateFromMte3 = false; + } + } + + if (storeFinalState && std::is_same::value) { + AscendC::WaitFlag(updateReadyEvent); + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + if (!isFinalState) { + AscendC::WaitFlag(EVENT_ID2 + pingpongFlag); + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + } + } else { + AscendC::WaitFlag(EVENT_ID2 + pingpongFlag); + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + } + if constexpr (kGated) { + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + } + if (useDirectFp32Ub) { + uint32_t directUbSlot = isPing ? 0 : 1; + uint64_t directUbFlagOffset = directUbSlot + AscendC::GetSubBlockIdx() * directUbFlagStride; + AscendC::CrossCoreSetFlag<0x4, PIPE_MTE3>(directUbFreeFlagBegin + directUbFlagOffset); + } + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + } + +private: + uint32_t pongBaseEvent = 4; + + AscendC::LocalTensor calcUbTensor; + + AscendC::LocalTensor hUpdateUbTensor_ping; + AscendC::LocalTensor hUbTensor_ping; + AscendC::LocalTensor finalOutputUbTensor_ping; + AscendC::LocalTensor glastUbTensor_ping; + + AscendC::LocalTensor hUpdateUbTensor_pong; + AscendC::LocalTensor hUbTensor_pong; + AscendC::LocalTensor finalOutputUbTensor_pong; + AscendC::LocalTensor glastUbTensor_pong; + + AscendC::LocalTensor gkLastUbTensor_ping; + AscendC::LocalTensor gkLastUbTensor_pong; + AscendC::LocalTensor gkInputUbTensor_ping; + AscendC::LocalTensor gkInputUbTensor_pong; + AscendC::LocalTensor gkBrcbUbTensor_ping; + AscendC::LocalTensor gkBrcbUbTensor_pong; +}; +} // namespace Catlass::Epilogue::Block + +#endif diff --git a/csrc/chunk_kda_fwd/op_kernel/arch35/fwd_h/epilogue/block/block_epilogue_kda_fwdh_vnew_arch35.hpp b/csrc/chunk_kda_fwd/op_kernel/arch35/fwd_h/epilogue/block/block_epilogue_kda_fwdh_vnew_arch35.hpp new file mode 100644 index 0000000000..cd908f5077 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/arch35/fwd_h/epilogue/block/block_epilogue_kda_fwdh_vnew_arch35.hpp @@ -0,0 +1,509 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_KDA_FWDH_VNEW_ARCH35_HPP +#define CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_KDA_FWDH_VNEW_ARCH35_HPP +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "../kda_fwd_h_epilogue_policies.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/kda_matrix_coord.hpp" +#include "catlass/epilogue/tile/kda_epilogue_tile_copy.hpp" +#include "block_epilogue_kda_fwdh_regbase.hpp" + +namespace Catlass::Epilogue::Block { + +template +class BlockEpilogue { + static constexpr bool kGated = KGatedTag::value; + static constexpr bool scalarGated = KGatedTag::scalarGated; + static constexpr bool useExp2 = KGatedTag::useExp2; + static constexpr float LN2 = 0.6931471805599453f; + +public: + using DispatchPolicy = EpilogueAtlasKdaFwdHVnew; + using ArchTag = typename DispatchPolicy::ArchTag; + + using VElementOutput = typename VOutputType_::Element; + using GElementInput = typename GInputType_::Element; + using UElementInput = typename UInputType_::Element; + using WSElementInput = typename WSInputType_::Element; + using VElementUpdate = typename VUpdateType_::Element; + using VLayoutUpdate = typename VUpdateType_::Layout; + using FinalStateElement = typename FinalStateType_::Element; + + CATLASS_DEVICE + BlockEpilogue(Arch::Resource &resource) + { + constexpr uint32_t CALC_BUF_OFFSET = 0; + constexpr uint32_t PING_BUF_0_OFFSET = 32 * 1024; + constexpr uint32_t PING_BUF_1_OFFSET = 48 * 1024; + constexpr uint32_t PING_BUF_2_OFFSET = 64 * 1024; + constexpr uint32_t PING_BUF_3_OFFSET = 80 * 1024; + constexpr uint32_t PONG_BUF_0_OFFSET = 96 * 1024; + constexpr uint32_t PONG_BUF_1_OFFSET = 112 * 1024; + constexpr uint32_t PONG_BUF_2_OFFSET = 128 * 1024; + constexpr uint32_t PONG_BUF_3_OFFSET = 144 * 1024; + constexpr uint32_t PING_WIDE_IO_BUF_OFFSET = 16 * 1024; + constexpr uint32_t PONG_WIDE_IO_BUF_OFFSET = 24 * 1024; + constexpr uint32_t PING_G_BUF_OFFSET = 160 * 1024; + constexpr uint32_t PONG_G_BUF_OFFSET = 161 * 1024; + constexpr uint32_t PING_G_SUB_BUF_OFFSET = 162 * 1024; + constexpr uint32_t PONG_G_SUB_BUF_OFFSET = 163 * 1024; + constexpr uint32_t PING_G_INPUT_BUF_OFFSET = 164 * 1024; + constexpr uint32_t PONG_G_INPUT_BUF_OFFSET = 165 * 1024; + constexpr uint32_t SHARE_BUF_OFFSET = 166 * 1024; + + calcUbTensor = resource.ubBuf.template GetBufferByByte(CALC_BUF_OFFSET); + + uUbTensor_ping = resource.ubBuf.template GetBufferByByte(PING_BUF_2_OFFSET); + wsUbTensor_ping = resource.ubBuf.template GetBufferByByte(PING_BUF_0_OFFSET); + gUbTensor_ping = resource.ubBuf.template GetBufferByByte(PING_G_BUF_OFFSET); + gLastUbTensor_ping = resource.ubBuf.template GetBufferByByte(PING_G_SUB_BUF_OFFSET); + gInputUbTensor_ping = resource.ubBuf.template GetBufferByByte(PING_G_SUB_BUF_OFFSET); + vNewOutputUbTensor_ping = resource.ubBuf.template GetBufferByByte(PING_BUF_2_OFFSET); + vNewDecayUbTensor_ping = resource.ubBuf.template GetBufferByByte(PING_BUF_2_OFFSET); + wideIoUbTensor_ping = resource.ubBuf.template GetBufferByByte(PING_WIDE_IO_BUF_OFFSET); + + uUbTensor_pong = resource.ubBuf.template GetBufferByByte(PONG_BUF_2_OFFSET); + wsUbTensor_pong = resource.ubBuf.template GetBufferByByte(PONG_BUF_0_OFFSET); + gUbTensor_pong = resource.ubBuf.template GetBufferByByte(PONG_G_BUF_OFFSET); + gLastUbTensor_pong = resource.ubBuf.template GetBufferByByte(PONG_G_SUB_BUF_OFFSET); + gInputUbTensor_pong = resource.ubBuf.template GetBufferByByte(PONG_G_SUB_BUF_OFFSET); + vNewOutputUbTensor_pong = resource.ubBuf.template GetBufferByByte(PONG_BUF_2_OFFSET); + vNewDecayUbTensor_pong = resource.ubBuf.template GetBufferByByte(PONG_BUF_2_OFFSET); + wideIoUbTensor_pong = resource.ubBuf.template GetBufferByByte(PONG_WIDE_IO_BUF_OFFSET); + + gBrcbUbTensor_ = resource.ubBuf.template GetBufferByByte(SHARE_BUF_OFFSET); + } + + CATLASS_DEVICE + ~BlockEpilogue() {} + + template + CATLASS_DEVICE void CopyGmToUb(AscendC::LocalTensor dst, AscendC::GlobalTensor src, uint32_t rows, + uint32_t cols, uint32_t srcStride) + { + if (cols == srcStride) { + AscendC::DataCopy(dst, src, rows * cols); + return; + } + AscendC::DataCopyExtParams copyParams{static_cast(rows), + static_cast(cols * sizeof(Element)), + static_cast((srcStride - cols) * sizeof(Element)), 0, 0}; + AscendC::DataCopyPadExtParams padParams{false, 0, 0, 0}; + AscendC::DataCopyPad(dst, src, copyParams, padParams); + } + + template + CATLASS_DEVICE void CopyUbToGm(AscendC::GlobalTensor dst, AscendC::LocalTensor src, uint32_t rows, + uint32_t cols, uint32_t dstStride) + { + if (cols == dstStride) { + AscendC::DataCopy(dst, src, rows * cols); + return; + } + AscendC::DataCopyExtParams copyParams{static_cast(rows), + static_cast(cols * sizeof(Element)), 0, + static_cast((dstStride - cols) * sizeof(Element)), 0}; + AscendC::DataCopyPad(dst, src, copyParams); + } + + CATLASS_DEVICE + void PrepareG(AscendC::LocalTensor gUbTensor, AscendC::LocalTensor gLastUbTensor, + AscendC::LocalTensor gInputUbTensor, + AscendC::GlobalTensor gInputThisSubBlock, uint32_t mActual, uint32_t pingpongFlag) + { + AscendC::WaitFlag(EVENT_ID3 + pingpongFlag); + if constexpr (!scalarGated) { + AscendC::Duplicate(gUbTensor, 1.0f, mActual); + AscendC::PipeBarrier(); + return; + } + if (mActual == 1) { + AscendC::Duplicate(gUbTensor, 1.0f, 1); + AscendC::PipeBarrier(); + return; + } + if constexpr (std::is_same::value) { + AscendC::DataCopyParams gUbParams{1, (uint16_t)(mActual * sizeof(float)), 0, 0}; + AscendC::DataCopyPadParams gUbPadParams{false, 0, 0, 0}; + AscendC::DataCopyPad(gUbTensor, gInputThisSubBlock, gUbParams, gUbPadParams); + } else { + AscendC::DataCopyParams gUbParams{1, (uint16_t)(mActual * sizeof(GElementInput)), 0, 0}; + AscendC::DataCopyPadParams gUbPadParams{false, 0, 0, 0}; + AscendC::DataCopyPad(gInputUbTensor, gInputThisSubBlock, gUbParams, gUbPadParams); + } + AscendC::SetFlag(EVENT_ID3 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID3 + pingpongFlag); + if constexpr (!std::is_same::value) { + AscendC::Cast(gUbTensor, gInputUbTensor, AscendC::RoundMode::CAST_NONE, mActual); + AscendC::PipeBarrier(); + } + + AscendC::SetFlag(EVENT_ID3 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID3 + pingpongFlag); + float inputVal = gUbTensor.GetValue(mActual - 1); + AscendC::SetFlag(EVENT_ID3 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID3 + pingpongFlag); + + AscendC::PipeBarrier(); + AscendC::Duplicate(gLastUbTensor, inputVal, mActual); + AscendC::PipeBarrier(); + + AscendC::Sub(gUbTensor, gLastUbTensor, gUbTensor, mActual); + AscendC::PipeBarrier(); + if constexpr (useExp2) { + AscendC::Muls(gUbTensor, gUbTensor, LN2, mActual); + AscendC::PipeBarrier(); + } + AscendC::Exp(gUbTensor, gUbTensor, mActual); + AscendC::PipeBarrier(); + } + + CATLASS_DEVICE + void ApplyRowScale(AscendC::LocalTensor matrix, AscendC::LocalTensor rowScale, + uint32_t rowScaleOffset, uint32_t rows, uint32_t cols) + { + __ubuf__ float *matrixAddr = reinterpret_cast<__ubuf__ float *>(matrix.GetPhyAddr()); + __ubuf__ float *rowScaleAddr = reinterpret_cast<__ubuf__ float *>(rowScale.GetPhyAddr()); + AscendC::VF_CALL(matrixAddr, rowScaleAddr, rowScaleOffset, + static_cast(rows), static_cast(cols)); + AscendC::PipeBarrier(); + } + + CATLASS_DEVICE + void ComputeVNew(AscendC::LocalTensor workspace, AscendC::LocalTensor uInput, uint32_t count) + { + __ubuf__ float *workspaceAddr = reinterpret_cast<__ubuf__ float *>(workspace.GetPhyAddr()); + __ubuf__ UElementInput *uInputAddr = reinterpret_cast<__ubuf__ UElementInput *>(uInput.GetPhyAddr()); + AscendC::VF_CALL>(workspaceAddr, uInputAddr, count); + AscendC::PipeBarrier(); + } + + CATLASS_DEVICE + void operator()(AscendC::GlobalTensor vnewOutput, + AscendC::GlobalTensor vnewdecayOutput, + AscendC::LocalTensor l1VUpdate, AscendC::GlobalTensor gInput, + AscendC::GlobalTensor uInput, AscendC::GlobalTensor wsInput, + AscendC::GlobalTensor gkInput, AscendC::GlobalTensor kInput, + AscendC::GlobalTensor kDecayWorkspace, uint32_t chunkSize, uint32_t kHeadDim, + uint32_t vBlockDim, uint32_t vHeadDim, Arch::CrossCoreFlag cube1Done, Arch::CrossCoreFlag vec1Done, + bool isInitialState, bool isFinalState, bool storeFinalState, bool waitWsFromMte3, bool isPing, + bool cube1AlreadyWaited, bool useDirectFp32Ub, uint64_t directUbFreeFlagBegin, + uint64_t directUbReadyFlagBegin, uint64_t directUbFlagStride) + { + static constexpr uint32_t ROW_TILE = 16; + uint32_t mActual = chunkSize; + uint32_t nvActual = vBlockDim; + uint32_t nkActual = kHeadDim; + uint32_t inputStride = vHeadDim; + + uint32_t subBlockIdx = AscendC::GetSubBlockIdx(); + uint32_t subBlockNum = AscendC::GetSubBlockNum(); + uint32_t rowsPerSubBlock = CeilDiv(mActual, subBlockNum); + uint32_t rowBegin = subBlockIdx * rowsPerSubBlock; + uint32_t rowEnd = rowBegin + rowsPerSubBlock; + if (rowEnd > mActual) { + rowEnd = mActual; + } + uint32_t pingpongFlag = isPing ? 0 : pongBaseEvent; + if (rowBegin >= mActual) { + if (useDirectFp32Ub) { + uint32_t directUbSlot = isPing ? 0 : 1; + uint64_t directUbFlagOffset = directUbSlot + AscendC::GetSubBlockIdx() * directUbFlagStride; + AscendC::CrossCoreWaitFlag<0x4, PIPE_V>(directUbReadyFlagBegin + directUbFlagOffset); + AscendC::CrossCoreSetFlag<0x4, PIPE_V>(directUbFreeFlagBegin + directUbFlagOffset); + } else if (!cube1AlreadyWaited) { + Arch::CrossCoreWaitFlag(cube1Done); + } + // A zero-row AIV lane still owns the EVENT0 hand-off consumed by V2. + if (waitWsFromMte3) { + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + } + Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(vec1Done); + return; + } + AscendC::ResetMask(); + + AscendC::GlobalTensor gInputThisSubBlock = gInput; + + AscendC::LocalTensor uUbTensor = isPing ? uUbTensor_ping : uUbTensor_pong; + AscendC::LocalTensor wsUbTensor = isPing ? wsUbTensor_ping : wsUbTensor_pong; + AscendC::LocalTensor gUbTensor = isPing ? gUbTensor_ping : gUbTensor_pong; + AscendC::LocalTensor gLastUbTensor = isPing ? gLastUbTensor_ping : gLastUbTensor_pong; + AscendC::LocalTensor gInputUbTensor = isPing ? gInputUbTensor_ping : gInputUbTensor_pong; + AscendC::LocalTensor vNewOutputUbTensor = + isPing ? vNewOutputUbTensor_ping : vNewOutputUbTensor_pong; + AscendC::LocalTensor vNewDecayUbTensor = + isPing ? vNewDecayUbTensor_ping : vNewDecayUbTensor_pong; + if (nvActual > 128) { + uUbTensor = isPing ? wideIoUbTensor_ping : wideIoUbTensor_pong; + vNewOutputUbTensor = isPing ? wideIoUbTensor_ping : wideIoUbTensor_pong; + vNewDecayUbTensor = isPing ? wideIoUbTensor_ping : wideIoUbTensor_pong; + } + + if (rowBegin < rowEnd && nvActual <= 128 && nvActual == inputStride) { + uint32_t mActualThisSubBlock = rowEnd - rowBegin; + AscendC::GlobalTensor vnewOutputThisSubBlock = vnewOutput[rowBegin * inputStride]; + AscendC::GlobalTensor uInputThisSubBlock = uInput[rowBegin * inputStride]; + AscendC::GlobalTensor wsInputThisSubBlock = wsInput[rowBegin * nvActual]; + + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + AscendC::DataCopy(uUbTensor, uInputThisSubBlock, mActualThisSubBlock * nvActual); + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + if constexpr (scalarGated) { + AscendC::Cast(calcUbTensor, uUbTensor, AscendC::RoundMode::CAST_NONE, mActualThisSubBlock * nvActual); + } + + if constexpr (scalarGated) { + PrepareG(gUbTensor, gLastUbTensor, gInputUbTensor, gInputThisSubBlock, mActual, pingpongFlag); + } else { + AscendC::WaitFlag(EVENT_ID3 + pingpongFlag); + } + if (waitWsFromMte3) { + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + } else { + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + } + if (useDirectFp32Ub) { + uint32_t directUbSlot = isPing ? 0 : 1; + uint64_t directUbFlagOffset = directUbSlot + AscendC::GetSubBlockIdx() * directUbFlagStride; + AscendC::CrossCoreWaitFlag<0x4, PIPE_V>(directUbReadyFlagBegin + directUbFlagOffset); + } else { + if (!cube1AlreadyWaited) { + Arch::CrossCoreWaitFlag(cube1Done); + } + AscendC::DataCopy(wsUbTensor, wsInputThisSubBlock, mActualThisSubBlock * nvActual); + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + } + + if constexpr (scalarGated) { + AscendC::Sub(wsUbTensor, calcUbTensor, wsUbTensor, mActualThisSubBlock * nvActual); + AscendC::PipeBarrier(); + AscendC::Copy(calcUbTensor, wsUbTensor, mActualThisSubBlock * nvActual); + AscendC::PipeBarrier(); + ApplyRowScale(calcUbTensor, gUbTensor, rowBegin, mActualThisSubBlock, nvActual); + } else { + ComputeVNew(wsUbTensor, uUbTensor, mActualThisSubBlock * nvActual); + } + AscendC::SetFlag(EVENT_ID3 + pingpongFlag); + + uint32_t nvLoops = nvActual / FLOAT_NUM_PER_REPEAT; + for (uint32_t nLoop = 0; nLoop < nvLoops; nLoop++) { + uint32_t castSrcOffset = nLoop * FLOAT_NUM_PER_REPEAT; + uint32_t castDstOffset = nLoop * mActualThisSubBlock * FLOAT_NUM_PER_REPEAT; + AscendC::LocalTensor decayInput = scalarGated ? calcUbTensor : wsUbTensor; + AscendC::Cast(vNewDecayUbTensor[castDstOffset], decayInput[castSrcOffset], + AscendC::RoundMode::CAST_RINT, FLOAT_NUM_PER_REPEAT, mActualThisSubBlock, + {(uint16_t)mActualThisSubBlock, 1, 1, (uint8_t)(nvLoops * 8)}); + } + + AscendC::PipeBarrier(); + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + + uint32_t ubToL1Loops = nvActual / SIZE_16_NUM_PER_C0; + uint32_t mActualPadded = (mActual + NZ_BLOCK_SIZE - 1) / NZ_BLOCK_SIZE * NZ_BLOCK_SIZE; + AscendC::DataCopyParams intriParams; + intriParams.blockCount = ubToL1Loops; + intriParams.blockLen = mActualThisSubBlock; + intriParams.srcGap = 0; + intriParams.dstGap = mActualPadded - mActualThisSubBlock; + uint32_t l1Addr = rowBegin * SIZE_16_NUM_PER_C0; + AscendC::DataCopy(vnewdecayOutput[l1Addr], vNewDecayUbTensor, intriParams); + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + if constexpr (!kGated) { + Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(vec1Done); + } + + AscendC::Cast(vNewOutputUbTensor, wsUbTensor, AscendC::RoundMode::CAST_RINT, + mActualThisSubBlock * nvActual); + AscendC::PipeBarrier(); + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + AscendC::DataCopy(vnewOutputThisSubBlock, vNewOutputUbTensor, mActualThisSubBlock * nvActual); + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + + if constexpr (kGated) { + AscendC::GlobalTensor kInputThisSubBlock = kInput[rowBegin * nkActual]; + AscendC::GlobalTensor kDecayWorkspaceThisSubBlock = + kDecayWorkspace[rowBegin * nkActual]; + // KDA passes kg = k * exp2(g_last - gk). Keep that decay exactly once. + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + AscendC::DataCopy(vNewOutputUbTensor, kInputThisSubBlock, mActualThisSubBlock * nkActual); + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + AscendC::DataCopy(kDecayWorkspaceThisSubBlock, vNewOutputUbTensor, mActualThisSubBlock * nkActual); + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + + Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(vec1Done); + } + if (useDirectFp32Ub) { + uint32_t directUbSlot = isPing ? 0 : 1; + uint64_t directUbFlagOffset = directUbSlot + AscendC::GetSubBlockIdx() * directUbFlagStride; + AscendC::CrossCoreSetFlag<0x4, PIPE_V>(directUbFreeFlagBegin + directUbFlagOffset); + } + return; + } + + if constexpr (scalarGated) { + PrepareG(gUbTensor, gLastUbTensor, gInputUbTensor, gInputThisSubBlock, mActual, pingpongFlag); + } else { + AscendC::WaitFlag(EVENT_ID3 + pingpongFlag); + } + if (!cube1AlreadyWaited) { + Arch::CrossCoreWaitFlag(cube1Done); + } + + uint32_t mActualPadded = (mActual + NZ_BLOCK_SIZE - 1) / NZ_BLOCK_SIZE * NZ_BLOCK_SIZE; + bool waitWsThisTileFromMte3 = waitWsFromMte3; + for (uint32_t rowStart = rowBegin; rowStart < rowEnd;) { + uint32_t alignExtra = rowStart & 7; + uint32_t maxRowsThisTile = ROW_TILE - alignExtra; + uint32_t rowsThisTile = rowEnd - rowStart; + if (rowsThisTile > maxRowsThisTile) { + rowsThisTile = maxRowsThisTile; + } + uint32_t localRowStart = rowStart - rowBegin; + + AscendC::GlobalTensor vnewOutputThisTile = vnewOutput[rowStart * inputStride]; + AscendC::GlobalTensor uInputThisTile = uInput[rowStart * inputStride]; + AscendC::GlobalTensor wsInputThisTile = wsInput[rowStart * nvActual]; + AscendC::LocalTensor wsUbTensorThisTile = wsUbTensor[localRowStart * nvActual]; + + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + CopyGmToUb(uUbTensor, uInputThisTile, rowsThisTile, nvActual, inputStride); + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + if constexpr (scalarGated) { + AscendC::Cast(calcUbTensor, uUbTensor, AscendC::RoundMode::CAST_NONE, rowsThisTile * nvActual); + AscendC::PipeBarrier(); + } + + if (waitWsThisTileFromMte3) { + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + waitWsThisTileFromMte3 = false; + } else { + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + } + CopyGmToUb(wsUbTensorThisTile, wsInputThisTile, rowsThisTile, nvActual, nvActual); + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + if constexpr (scalarGated) { + AscendC::Sub(wsUbTensorThisTile, calcUbTensor, wsUbTensorThisTile, rowsThisTile * nvActual); + AscendC::PipeBarrier(); + AscendC::Copy(calcUbTensor, wsUbTensorThisTile, rowsThisTile * nvActual); + AscendC::PipeBarrier(); + ApplyRowScale(calcUbTensor, gUbTensor, rowStart, rowsThisTile, nvActual); + } else { + ComputeVNew(wsUbTensorThisTile, uUbTensor, rowsThisTile * nvActual); + } + + uint32_t nvLoops = nvActual / FLOAT_NUM_PER_REPEAT; + for (uint32_t nLoop = 0; nLoop < nvLoops; nLoop++) { + uint32_t castSrcOffset = nLoop * FLOAT_NUM_PER_REPEAT; + uint32_t castDstOffset = nLoop * rowsThisTile * FLOAT_NUM_PER_REPEAT; + AscendC::LocalTensor decayInput = scalarGated ? calcUbTensor : wsUbTensorThisTile; + AscendC::Cast(vNewDecayUbTensor[castDstOffset], decayInput[castSrcOffset], + AscendC::RoundMode::CAST_RINT, FLOAT_NUM_PER_REPEAT, rowsThisTile, + {(uint16_t)rowsThisTile, 1, 1, (uint8_t)(nvLoops * 8)}); + } + + AscendC::PipeBarrier(); + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + + uint32_t ubToL1Loops = nvActual / SIZE_16_NUM_PER_C0; + AscendC::DataCopyParams intriParams; + intriParams.blockCount = ubToL1Loops; + intriParams.blockLen = rowsThisTile; + intriParams.srcGap = 0; + intriParams.dstGap = mActualPadded - rowsThisTile; + uint32_t l1Addr = rowStart * SIZE_16_NUM_PER_C0; + AscendC::DataCopy(vnewdecayOutput[l1Addr], vNewDecayUbTensor, intriParams); + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + if constexpr (!kGated) { + if (rowStart + rowsThisTile >= rowEnd) { + Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(vec1Done); + } + } + + AscendC::Cast(vNewOutputUbTensor, wsUbTensorThisTile, AscendC::RoundMode::CAST_RINT, + rowsThisTile * nvActual); + AscendC::PipeBarrier(); + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + CopyUbToGm(vnewOutputThisTile, vNewOutputUbTensor, rowsThisTile, nvActual, inputStride); + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + rowStart += rowsThisTile; + } + + if constexpr (kGated) { + uint32_t mActualThisSubBlock = rowEnd - rowBegin; + for (uint32_t rowOffset = 0; rowOffset < mActualThisSubBlock; rowOffset += ROW_TILE) { + uint32_t rowsThisTile = Min(ROW_TILE, mActualThisSubBlock - rowOffset); + uint32_t row = rowBegin + rowOffset; + AscendC::GlobalTensor kInputThisTile = kInput[row * nkActual]; + AscendC::GlobalTensor kDecayWorkspaceThisTile = kDecayWorkspace[row * nkActual]; + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + CopyGmToUb(vNewOutputUbTensor, kInputThisTile, rowsThisTile, nkActual, nkActual); + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + CopyUbToGm(kDecayWorkspaceThisTile, vNewOutputUbTensor, rowsThisTile, nkActual, nkActual); + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + } + Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(vec1Done); + } + + if (rowBegin < rowEnd) { + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + } + AscendC::SetFlag(EVENT_ID3 + pingpongFlag); + } + +private: + uint32_t pongBaseEvent = 4; + + AscendC::LocalTensor calcUbTensor; + + AscendC::LocalTensor uUbTensor_ping; + AscendC::LocalTensor wsUbTensor_ping; + AscendC::LocalTensor gUbTensor_ping; + AscendC::LocalTensor gLastUbTensor_ping; + AscendC::LocalTensor gInputUbTensor_ping; + AscendC::LocalTensor vNewOutputUbTensor_ping; + AscendC::LocalTensor vNewDecayUbTensor_ping; + + AscendC::LocalTensor uUbTensor_pong; + AscendC::LocalTensor wsUbTensor_pong; + AscendC::LocalTensor gUbTensor_pong; + AscendC::LocalTensor gLastUbTensor_pong; + AscendC::LocalTensor gInputUbTensor_pong; + AscendC::LocalTensor vNewOutputUbTensor_pong; + AscendC::LocalTensor vNewDecayUbTensor_pong; + AscendC::LocalTensor wideIoUbTensor_ping; + AscendC::LocalTensor wideIoUbTensor_pong; + + AscendC::LocalTensor gBrcbUbTensor_; +}; +} // namespace Catlass::Epilogue::Block + +#endif diff --git a/csrc/chunk_kda_fwd/op_kernel/arch35/fwd_h/epilogue/kda_fwd_h_epilogue_policies.hpp b/csrc/chunk_kda_fwd/op_kernel/arch35/fwd_h/epilogue/kda_fwd_h_epilogue_policies.hpp new file mode 100644 index 0000000000..964e3d58ad --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/arch35/fwd_h/epilogue/kda_fwd_h_epilogue_policies.hpp @@ -0,0 +1,28 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_EPILOGUE_KDA_FWD_H_EPILOGUE_POLICIES_HPP +#define CATLASS_EPILOGUE_KDA_FWD_H_EPILOGUE_POLICIES_HPP + +#include "catlass/kda_catlass.hpp" + +namespace Catlass::Epilogue { + +struct EpilogueAtlasKdaFwdHVnew { + using ArchTag = Arch::Ascend950; +}; + +struct EpilogueAtlasKdaFwdHUpdate { + using ArchTag = Arch::Ascend950; +}; + +} // namespace Catlass::Epilogue + +#endif // CATLASS_EPILOGUE_KDA_FWD_H_EPILOGUE_POLICIES_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/arch35/fwd_h/gemm/block/block_scheduler_kda_fwd_h_arch35.hpp b/csrc/chunk_kda_fwd/op_kernel/arch35/fwd_h/gemm/block/block_scheduler_kda_fwd_h_arch35.hpp new file mode 100644 index 0000000000..c0ac3c14cc --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/arch35/fwd_h/gemm/block/block_scheduler_kda_fwd_h_arch35.hpp @@ -0,0 +1,437 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#include "catlass/kda_gemm_coord.hpp" +#include "../../../../chunk_kda_fwd_h_runtime_tiling.h" +using namespace Catlass; + +#ifndef CATLASS_GEMM_SCHEDULER_KDA_FWD_H_HPP +#define CATLASS_GEMM_SCHEDULER_KDA_FWD_H_HPP + +// constexpr uint32_t PING_PONG_STAGES = 1; +constexpr uint32_t PING_PONG_STAGES = 2; +constexpr uint32_t BYTE_SIZE_16_BIT = 2; +constexpr uint32_t BYTES_PER_C0 = 32; +constexpr uint32_t BYTE_SIZE_PER_REPEAT = 256; +constexpr uint32_t SIZE_16_NUM_PER_C0 = BYTES_PER_C0 / BYTE_SIZE_16_BIT; +constexpr uint32_t FLOAT_NUM_PER_REPEAT = BYTE_SIZE_PER_REPEAT / sizeof(float); +constexpr uint32_t NZ_BLOCK_SIZE = 16; + +template +CATLASS_DEVICE T AlignUp(T a, T b) +{ + return (b == 0) ? 0 : (a + b - 1) / b * b; +} + +template +CATLASS_DEVICE T Min(T a, T b) +{ + return (a > b) ? b : a; +} + +template +CATLASS_DEVICE T Max(T a, T b) +{ + return (a > b) ? a : b; +} + +namespace Catlass::Gemm::Block { + +struct KdaFwdHOffsets { + uint32_t hSrcOffset; + uint32_t hDstOffset; + uint32_t uvOffset; + uint32_t wkOffset; + uint32_t wOffset; + uint32_t gOffset; + uint32_t gkOffset; + uint32_t hWorkOffset; + uint32_t vWorkOffset; + uint32_t kDecayWorkOffset; + uint32_t vBlockOffset; + uint32_t vBlockDim; + uint32_t initialStateOffset; + uint32_t finalStateOffset; + bool isInitialState; + bool isFinalState; + uint32_t blockTokens; + uint32_t streamId; + // for debug + uint32_t batchIdx; + uint32_t headIdx; + uint32_t chunkIdx; +}; + +struct KdaFwdHStream { + uint32_t batchIdx; + uint32_t chunkIdx{0}; + uint32_t vHeadIdx; + uint32_t kHeadIdx; + uint32_t shapeBatchIdx; + uint32_t tokenBatchIdx; + + uint32_t chunkOffset; + uint32_t tokenOffset; + uint32_t batchChunks{0}; + uint32_t batchTokens; + uint32_t nextTaskIdx{0}; + bool active{false}; + + KdaFwdHOffsets offset; +}; + +struct KdaFwdHRunningQ { + KdaFwdHStream streams[PING_PONG_STAGES]; +}; + +struct BlockSchedulerKdaFwdH { + uint32_t batch; + uint32_t seqlen; + uint32_t kNumHead; + uint32_t vNumHead; + uint32_t kHeadDim; + uint32_t vHeadDim; + uint32_t chunkSize; + uint32_t vBlockSize{128}; + uint32_t isVariedLen; + uint32_t shapeBatch; + uint32_t tokenBatch; + uint32_t inputTokenBatch; + bool useInitialState; + bool storeFinalState; + uint32_t numSeqWorkspaceOffset; + uint32_t numChunksWorkspaceOffset; + + uint32_t taskIdx; + uint32_t taskStride; + uint32_t cubeCoreIdx; + uint32_t cubeCoreNum; + uint32_t taskNum; + uint32_t headGroups; + uint32_t totalChunks; + uint32_t totalTokens; + + KdaFwdHRunningQ runningQ; + + bool isRunning; + + AscendC::GlobalTensor gmSeqlen; + AscendC::GlobalTensor gmNumSeq; + AscendC::GlobalTensor gmNumChunks; + + Arch::CrossCoreFlag cube1Done[PING_PONG_STAGES] = {0, 1}; + Arch::CrossCoreFlag vec1Done[PING_PONG_STAGES] = {2, 3}; + Arch::CrossCoreFlag cube2Done[PING_PONG_STAGES] = {4, 5}; + Arch::CrossCoreFlag vec2Done[PING_PONG_STAGES] = {6, 7}; + + CATLASS_DEVICE + BlockSchedulerKdaFwdH() {} + + CATLASS_DEVICE + void Init(GM_ADDR cu_seqlens, GM_ADDR chunk_indices, GM_ADDR tiling, GM_ADDR user, uint32_t coreIdx, + uint32_t coreNum) + { + __gm__ ChunkKdaFwdHRuntimeTiling *__restrict kdaFwdHTilingData = + reinterpret_cast<__gm__ ChunkKdaFwdHRuntimeTiling *__restrict>(tiling); + + batch = kdaFwdHTilingData->batch; + seqlen = kdaFwdHTilingData->seqlen; + kNumHead = kdaFwdHTilingData->kNumHead; + vNumHead = kdaFwdHTilingData->vNumHead; + kHeadDim = kdaFwdHTilingData->kHeadDim; + vHeadDim = kdaFwdHTilingData->vHeadDim; + chunkSize = kdaFwdHTilingData->chunkSize; + isVariedLen = kdaFwdHTilingData->isVariedLen; + shapeBatch = kdaFwdHTilingData->shapeBatch; + tokenBatch = kdaFwdHTilingData->tokenBatch; + useInitialState = kdaFwdHTilingData->useInitialState; + storeFinalState = kdaFwdHTilingData->storeFinalState; + numSeqWorkspaceOffset = kdaFwdHTilingData->numSeqWorkspaceOffset; + numChunksWorkspaceOffset = kdaFwdHTilingData->numChunksWorkspaceOffset; + + InitRuntime(cu_seqlens, chunk_indices, user, coreIdx, coreNum); + } + + template + CATLASS_DEVICE void InitFromData(GM_ADDR cu_seqlens, GM_ADDR chunk_indices, const TilingData &tilingData, + GM_ADDR user, uint32_t coreIdx, uint32_t coreNum) + { + batch = tilingData.batch; + seqlen = tilingData.seqlen; + kNumHead = tilingData.kNumHead; + vNumHead = tilingData.vNumHead; + kHeadDim = tilingData.kHeadDim; + vHeadDim = tilingData.vHeadDim; + chunkSize = tilingData.chunkSize; + isVariedLen = tilingData.isVariedLen; + shapeBatch = tilingData.shapeBatch; + tokenBatch = tilingData.tokenBatch; + useInitialState = tilingData.useInitialState; + storeFinalState = tilingData.storeFinalState; + numSeqWorkspaceOffset = tilingData.numSeqWorkspaceOffset; + numChunksWorkspaceOffset = tilingData.numChunksWorkspaceOffset; + + InitRuntime(cu_seqlens, chunk_indices, user, coreIdx, coreNum); + } + + CATLASS_DEVICE + void InitRuntime(GM_ADDR cu_seqlens, GM_ADDR chunk_indices, GM_ADDR user, uint32_t coreIdx, uint32_t coreNum) + { + gmSeqlen.SetGlobalBuffer((__gm__ int64_t *)cu_seqlens); + gmNumSeq.SetGlobalBuffer((__gm__ int64_t *)(user + numSeqWorkspaceOffset)); + gmNumChunks.SetGlobalBuffer((__gm__ int64_t *)(user + numChunksWorkspaceOffset)); + + if (isVariedLen) { + inputTokenBatch = tokenBatch; + uint32_t actualBatch = 0; + int64_t chunkPrefix = 0; + int64_t prevSeq = 0, currSeq; + for (uint32_t b = 1; b <= inputTokenBatch; b++) { + currSeq = gmSeqlen.GetValue(b); + int64_t batchSeqLen = currSeq - prevSeq; + if (batchSeqLen > 0) { + actualBatch++; + int64_t batchChunk = (batchSeqLen + chunkSize - 1) / chunkSize; + chunkPrefix += batchChunk; + } + prevSeq = currSeq; + } + tokenBatch = actualBatch; + batch = actualBatch; + totalChunks = chunkPrefix; + totalTokens = prevSeq; + } else { + totalChunks = (seqlen + chunkSize - 1) / chunkSize; + totalTokens = seqlen; + } + + cubeCoreIdx = coreIdx; + cubeCoreNum = coreNum; + vBlockSize = vHeadDim; + taskNum = batch * vNumHead; + headGroups = vNumHead / kNumHead; + uint32_t maxTaskCntPerLoop = taskNum > cubeCoreNum ? PING_PONG_STAGES : 1; + taskStride = cubeCoreNum * maxTaskCntPerLoop; + for (uint32_t streamId = 0; streamId < PING_PONG_STAGES; ++streamId) { + auto &stream = runningQ.streams[streamId]; + stream.nextTaskIdx = + (streamId < maxTaskCntPerLoop) ? (cubeCoreIdx * maxTaskCntPerLoop + streamId) : taskNum; + stream.chunkIdx = 0; + stream.batchChunks = 0; + stream.active = false; + } + isRunning = cubeCoreIdx * maxTaskCntPerLoop < taskNum; + } + + CATLASS_DEVICE + void ResolveVarlenSequence(uint32_t compactBatchIdx, KdaFwdHStream &stream) + { + uint32_t actualBatch = 0; + int64_t chunkPrefix = 0; + int64_t prevSeq = 0; + for (uint32_t b = 1; b <= inputTokenBatch; ++b) { + int64_t currSeq = gmSeqlen.GetValue(b); + int64_t batchTokens = currSeq - prevSeq; + if (batchTokens > 0) { + int64_t batchChunks = (batchTokens + chunkSize - 1) / chunkSize; + if (actualBatch == compactBatchIdx) { + stream.chunkOffset = static_cast(chunkPrefix); + stream.batchChunks = static_cast(batchChunks); + stream.tokenOffset = static_cast(prevSeq); + stream.batchTokens = static_cast(batchTokens); + return; + } + ++actualBatch; + chunkPrefix += batchChunks; + } + prevSeq = currSeq; + } + stream.chunkOffset = 0; + stream.batchChunks = 0; + stream.tokenOffset = 0; + stream.batchTokens = 0; + } + + CATLASS_DEVICE + uint32_t GetVarlenChunkOffset(uint32_t compactBatchIdx) + { + KdaFwdHStream stream; + ResolveVarlenSequence(compactBatchIdx, stream); + return stream.chunkOffset; + } + + CATLASS_DEVICE + void InitNewStream(KdaFwdHStream &newStream) + { + newStream.batchIdx = taskIdx / vNumHead; + newStream.vHeadIdx = taskIdx % vNumHead; + newStream.kHeadIdx = newStream.vHeadIdx / headGroups; + newStream.shapeBatchIdx = isVariedLen ? 0 : newStream.batchIdx; + newStream.tokenBatchIdx = isVariedLen ? newStream.batchIdx : 0; + if (isVariedLen) { + ResolveVarlenSequence(newStream.tokenBatchIdx, newStream); + } else { + newStream.chunkOffset = 0; + newStream.batchChunks = totalChunks; + newStream.tokenOffset = 0; + newStream.batchTokens = totalTokens; + } + newStream.chunkIdx = 0; + } + + CATLASS_DEVICE + void AssignNextStream(uint32_t streamId) + { + auto &stream = runningQ.streams[streamId]; + taskIdx = stream.nextTaskIdx; + if (taskIdx >= taskNum) { + stream.active = false; + stream.batchChunks = 0; + return; + } + + stream.nextTaskIdx += taskStride; + InitNewStream(stream); + stream.active = stream.batchChunks > 0; + if (stream.active) { + UpdateTask(streamId); + } + } + + CATLASS_DEVICE + void UpdateTask(uint32_t streamId) + { + auto &stream = runningQ.streams[streamId]; + auto &offset = stream.offset; + + offset.isInitialState = stream.chunkIdx == 0; + offset.isFinalState = stream.chunkIdx == (stream.batchChunks - 1); + uint32_t vBlockOffset = 0; + uint32_t vBlockDim = vBlockSize; + offset.initialStateOffset = (stream.batchIdx * vNumHead + stream.vHeadIdx) * kHeadDim * vHeadDim + vBlockOffset; + offset.finalStateOffset = (stream.batchIdx * vNumHead + stream.vHeadIdx) * kHeadDim * vHeadDim + vBlockOffset; + offset.hSrcOffset = (stream.shapeBatchIdx * vNumHead * totalChunks + stream.vHeadIdx * totalChunks + + stream.chunkOffset + stream.chunkIdx) * + kHeadDim * vHeadDim + + vBlockOffset; + offset.hDstOffset = offset.hSrcOffset + kHeadDim * vHeadDim; + if (storeFinalState && offset.isFinalState) { + offset.hDstOffset = offset.hSrcOffset; + } + offset.uvOffset = (stream.shapeBatchIdx * vNumHead * totalTokens + stream.vHeadIdx * totalTokens + + stream.tokenOffset + stream.chunkIdx * chunkSize) * + vHeadDim + + vBlockOffset; + offset.wkOffset = (stream.shapeBatchIdx * kNumHead * totalTokens + stream.kHeadIdx * totalTokens + + stream.tokenOffset + stream.chunkIdx * chunkSize) * + kHeadDim; + offset.wOffset = (stream.shapeBatchIdx * vNumHead * totalTokens + stream.vHeadIdx * totalTokens + + stream.tokenOffset + stream.chunkIdx * chunkSize) * + kHeadDim; + offset.gOffset = stream.shapeBatchIdx * vNumHead * totalTokens + stream.vHeadIdx * totalTokens + + stream.tokenOffset + stream.chunkIdx * chunkSize; + offset.gkOffset = (stream.shapeBatchIdx * vNumHead * totalTokens + stream.vHeadIdx * totalTokens + + stream.tokenOffset + stream.chunkIdx * chunkSize) * + kHeadDim; + offset.hWorkOffset = (cubeCoreIdx * PING_PONG_STAGES + streamId) * kHeadDim * vBlockSize; + offset.vWorkOffset = (cubeCoreIdx * PING_PONG_STAGES + streamId) * chunkSize * vBlockSize; + offset.kDecayWorkOffset = (cubeCoreIdx * PING_PONG_STAGES + streamId) * chunkSize * kHeadDim; + offset.vBlockOffset = vBlockOffset; + offset.vBlockDim = vBlockDim; + offset.blockTokens = offset.isFinalState ? (stream.batchTokens - stream.chunkIdx * chunkSize) : chunkSize; + offset.streamId = streamId; + offset.batchIdx = stream.batchIdx; + offset.headIdx = stream.vHeadIdx; + offset.chunkIdx = stream.chunkIdx; + } + + CATLASS_DEVICE + void InitTasks() + { + isRunning = false; + for (uint32_t streamId = 0; streamId < PING_PONG_STAGES; ++streamId) { + auto &stream = runningQ.streams[streamId]; + if (stream.active) { + stream.chunkIdx += 1; + if (stream.chunkIdx >= stream.batchChunks) { + stream.active = false; + stream.batchChunks = 0; + } + } + if (!stream.active) { + AssignNextStream(streamId); + } else { + UpdateTask(streamId); + } + if (stream.active) { + isRunning = true; + } + } + } + + CATLASS_DEVICE + const KdaFwdHStream &GetStream(uint32_t i) const { return runningQ.streams[i]; } + + CATLASS_DEVICE + uint32_t GetStreamId(uint32_t i) const { return i; } + + CATLASS_DEVICE + const KdaFwdHOffsets &GetCurTaskOffsets(const KdaFwdHStream &stream) const { return stream.offset; } + + CATLASS_DEVICE + bool StreamIsDone(const KdaFwdHStream &stream) const { return !stream.active; } + + CATLASS_DEVICE + bool NeedProcessStage2(const KdaFwdHStream &stream) { return storeFinalState || !stream.offset.isFinalState; } +}; + +struct BlockSchedulerKdaFwdHCube : public BlockSchedulerKdaFwdH { + CATLASS_DEVICE + BlockSchedulerKdaFwdHCube() {} + + CATLASS_DEVICE + void Init(GM_ADDR cu_seqlens, GM_ADDR chunk_indices, GM_ADDR tiling, GM_ADDR user) + { + BlockSchedulerKdaFwdH::Init(cu_seqlens, chunk_indices, tiling, user, AscendC::GetBlockIdx(), + AscendC::GetBlockNum()); + } + + template + CATLASS_DEVICE void InitFromData(GM_ADDR cu_seqlens, GM_ADDR chunk_indices, const TilingData &tilingData, + GM_ADDR user) + { + BlockSchedulerKdaFwdH::InitFromData(cu_seqlens, chunk_indices, tilingData, user, AscendC::GetBlockIdx(), + AscendC::GetBlockNum()); + } +}; + +struct BlockSchedulerKdaFwdHVec : public BlockSchedulerKdaFwdH { + CATLASS_DEVICE + BlockSchedulerKdaFwdHVec() {} + + CATLASS_DEVICE + void Init(GM_ADDR cu_seqlens, GM_ADDR chunk_indices, GM_ADDR tiling, GM_ADDR user) + { + BlockSchedulerKdaFwdH::Init(cu_seqlens, chunk_indices, tiling, user, + AscendC::GetBlockIdx() / AscendC::GetSubBlockNum(), AscendC::GetBlockNum()); + } + + template + CATLASS_DEVICE void InitFromData(GM_ADDR cu_seqlens, GM_ADDR chunk_indices, const TilingData &tilingData, + GM_ADDR user) + { + BlockSchedulerKdaFwdH::InitFromData(cu_seqlens, chunk_indices, tilingData, user, + AscendC::GetBlockIdx() / AscendC::GetSubBlockNum(), AscendC::GetBlockNum()); + } +}; + +} // namespace Catlass::Gemm::Block + +#endif // CATLASS_GEMM_SCHEDULER_KDA_FWD_H_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/arch35/fwd_h/gemm/kernel/kda_fwd_h_kernel_arch35.hpp b/csrc/chunk_kda_fwd/op_kernel/arch35/fwd_h/gemm/kernel/kda_fwd_h_kernel_arch35.hpp new file mode 100644 index 0000000000..3239a21bd0 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/arch35/fwd_h/gemm/kernel/kda_fwd_h_kernel_arch35.hpp @@ -0,0 +1,1128 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#define CATLASS_ARCH 3510 + +#include "catlass/arch/kda_arch.hpp" +#include "catlass/arch/kda_cross_core_sync.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/kda_catlass.hpp" +#include "catlass/kda_debug.hpp" +#include "../block/block_scheduler_kda_fwd_h_arch35.hpp" +#include "catlass/epilogue/block/kda_block_epilogue.hpp" +#include "../../epilogue/block/block_epilogue_kda_fwdh_update_arch35.hpp" +#include "../../epilogue/block/block_epilogue_kda_fwdh_vnew_arch35.hpp" +#include "catlass/gemm/block/kda_block_mmad.hpp" +#include "../../../../kernel_utils/block/block_mmad_pingpong_tla_kernel_utils.hpp" +#include "../../../../kernel_utils/block/block_mmad_pingpong_tla_multi.hpp" +#include "../../../../kernel_utils/block/block_mmad_pingpong_tla_preloadA_l1B.hpp" +#include "catlass/gemm/block/kda_block_swizzle.hpp" +#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" +#include "catlass/gemm/kda_gemm_type.hpp" +#include "catlass/layout/kda_layout.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "tla/kda_tensor.hpp" +#include "tla/kda_tla_layout.hpp" +#include "tla/kda_tensor.hpp" + +using _0 = tla::Int<0>; +using _1 = tla::Int<1>; +using _2 = tla::Int<2>; +using _4 = tla::Int<4>; +using _8 = tla::Int<8>; +using _16 = tla::Int<16>; +using _32 = tla::Int<32>; +using _64 = tla::Int<64>; +using _128 = tla::Int<128>; +using _256 = tla::Int<256>; +using _512 = tla::Int<512>; +using _1024 = tla::Int<1024>; +using _2048 = tla::Int<2048>; +using _4096 = tla::Int<4096>; +using _8192 = tla::Int<8192>; +using _16384 = tla::Int<16384>; +using _32768 = tla::Int<32768>; +using _65536 = tla::Int<65536>; + +#include "kernel_operator.h" +using namespace Catlass; +using namespace tla; + +namespace Catlass::Gemm::Kernel { + +struct KdaFwdHTileShapes128 { + using L1TileShape = tla::Shape<_128, _128, _128>; + using L0TileShape = L1TileShape; +}; + +struct KdaFwdHTileShapes256 { + using L1TileShape = tla::Shape<_128, _256, _128>; + using L0TileShape = tla::Shape<_128, _256, _64>; +}; + +template +struct KdaFwdHGateTag { + static constexpr bool value = KGated; + static constexpr bool scalarGated = ScalarGated; + static constexpr bool useExp2 = UseExp2; +}; + +template +class KdaFwdHKernel { +public: + using ArchTag = Arch::Ascend950; + using CubeScheduler = typename Catlass::Gemm::Block::BlockSchedulerKdaFwdHCube; + using VecScheduler = typename Catlass::Gemm::Block::BlockSchedulerKdaFwdHVec; + + using DispatchPolicyTlaMulti = Gemm::MmadPingpongTlaMulti; + using DispatchPolicyTlaTail = Gemm::MmadPingpongTlaMulti; + using DispatchPolicyDirectUb = Common::MmadPingpong; + using DispatchPolicyTlaPreloadAL1B = Gemm::MmadPingpongTlaPreloadAL1B; + using L1TileShapeVTla = typename TileShapes::L1TileShape; + using L0TileShapeVTla = typename TileShapes::L0TileShape; + + using WType = Gemm::GemmType; + using HType = Gemm::GemmType; + using VworkType = Gemm::GemmType; + using KType = Gemm::GemmType; + using HworkType = Gemm::GemmType; + using VType = Gemm::GemmType; + using GType = Gemm::GemmType; + using UType = Gemm::GemmType; + using FinalStateType = Gemm::GemmType; + using VUpdateType = Gemm::GemmType; + + // cube 1 + using TileCopyWH = Catlass::Gemm::Tile::PackedTileCopyTla; + using TileCopyWHDirectUb = Common::Tile::PackedTileCopyTlaToUB; + using BlockMmadWH = Gemm::Block::BlockMmadTla; + using BlockMmadWHTail = Gemm::Block::BlockMmadTla; + using BlockMmadWHDirectUb = Common::BlockMmadTla; + + // cube 2 + using TileCopyKV = Catlass::Gemm::Tile::PackedTileCopyTla; + using TileCopyKVDirectUb = Common::Tile::PackedTileCopyTlaToUB; + using TileMmadKV = Gemm::Tile::TileMmadTla; + using BlockMmadKV = Gemm::Block::BlockMmadTla; + using BlockMmadKVTail = Gemm::Block::BlockMmadTla; + using BlockMmadKVDirectUb = Common::BlockMmadTla; + + // vec 1 + using DispatchPolicyKdaFwdHVnew = Epilogue::EpilogueAtlasKdaFwdHVnew; + using GateTag = KdaFwdHGateTag; + using EpilogueKdaFwdHVnew = Epilogue::Block::BlockEpilogue; + + // vec 2 + using DispatchPolicyKdaFwdHUpdate = Epilogue::EpilogueAtlasKdaFwdHUpdate; + using EpilogueKdaFwdHUpdate = Epilogue::Block::BlockEpilogue; + + using KdaFwdHOffsets = Catlass::Gemm::Block::KdaFwdHOffsets; + + using ElementK = INPUT_TYPE; + using ElementW = INPUT_TYPE; + using ElementU = INPUT_TYPE; + using ElementG = G_TYPE; + using ElementH = INPUT_TYPE; + using ElementV = INPUT_TYPE; + using ElementVUpdate = INPUT_TYPE; + using ElementVWork = WORKSPACE_TYPE; + using ElementHWork = WORKSPACE_TYPE; + using ElementInitialState = STATE_TYPE; + using ElementFinalState = STATE_TYPE; + + using LayoutW = Catlass::layout::RowMajor; + using LayoutH = Catlass::layout::RowMajor; + using LayoutV = Catlass::layout::RowMajor; + using LayoutK = Catlass::layout::ColumnMajor; + using LayoutVUpdate = typename VUpdateType::Layout; + + static constexpr uint64_t DIRECT_UB_FREE_FLAG_BEGIN = 1; + static constexpr uint64_t DIRECT_UB_READY_FLAG_BEGIN = 6; + static constexpr uint64_t DIRECT_UB_FLAG_STRIDE = 16; + static constexpr uint32_t DIRECT_UB_STAGES = 2; + static constexpr uint32_t DIRECT_VEC_NUM = 2; + + uint32_t batch; + uint32_t seqlen; + uint32_t kNumHead; + uint32_t vNumHead; + uint32_t kHeadDim; + uint32_t vHeadDim; + uint32_t chunkSize; + bool useInitialState; + bool storeFinalState; + uint32_t isVariedLen; + uint32_t shapeBatch; + uint32_t tokenBatch; + uint32_t vWorkspaceOffset; + uint32_t vUpdateWorkspaceOffset; + uint32_t hWorkspaceOffset; + uint32_t numSeqWorkspaceOffset; + uint32_t numChunksWorkspaceOffset; + uint32_t kDecayWorkspaceOffset; + bool useDirectFp32Ub; + + AscendC::GlobalTensor gmK; + AscendC::GlobalTensor gmW; + AscendC::GlobalTensor gmU; + AscendC::GlobalTensor gmG; + AscendC::GlobalTensor gmGk; + AscendC::GlobalTensor gmInitialState; + AscendC::GlobalTensor gmH; + AscendC::GlobalTensor gmV; + AscendC::GlobalTensor gmFinalState; + AscendC::GlobalTensor gmVWorkspace; + AscendC::GlobalTensor gmVUpdateWorkspace; + AscendC::GlobalTensor gmHWorkspace; + AscendC::GlobalTensor gmKDecayWorkspace; + + AscendC::GlobalTensor gmSeqlen; + AscendC::GlobalTensor gmNumSeq; + AscendC::GlobalTensor gmNumChunks; + + AscendC::LocalTensor ubHUpdatePing; + AscendC::LocalTensor ubHUpdatePong; + AscendC::LocalTensor ubVWorkPing; + AscendC::LocalTensor ubVWorkPong; + + AscendC::LocalTensor l1VUpdatePing; + AscendC::LocalTensor l1VUpdatePong; + + CubeScheduler cubeBlockScheduler; + VecScheduler vecBlockScheduler; + + Arch::Resource resource; + + __aicore__ inline KdaFwdHKernel() {} + + __aicore__ inline void Init(GM_ADDR k, GM_ADDR w, GM_ADDR u, GM_ADDR g, GM_ADDR gk, GM_ADDR inital_state, + GM_ADDR cu_seqlens, GM_ADDR chunk_indices, GM_ADDR h, GM_ADDR v_new, + GM_ADDR final_state, GM_ADDR tiling, GM_ADDR user) + { + __gm__ ChunkKdaFwdHRuntimeTiling *__restrict kdaFwdHTilingData = + reinterpret_cast<__gm__ ChunkKdaFwdHRuntimeTiling *__restrict>(tiling); + + batch = kdaFwdHTilingData->batch; + seqlen = kdaFwdHTilingData->seqlen; + kNumHead = kdaFwdHTilingData->kNumHead; + vNumHead = kdaFwdHTilingData->vNumHead; + kHeadDim = kdaFwdHTilingData->kHeadDim; + vHeadDim = kdaFwdHTilingData->vHeadDim; + chunkSize = kdaFwdHTilingData->chunkSize; + useInitialState = kdaFwdHTilingData->useInitialState; + storeFinalState = kdaFwdHTilingData->storeFinalState; + isVariedLen = kdaFwdHTilingData->isVariedLen; + shapeBatch = kdaFwdHTilingData->shapeBatch; + tokenBatch = kdaFwdHTilingData->tokenBatch; + vWorkspaceOffset = kdaFwdHTilingData->vWorkspaceOffset; + vUpdateWorkspaceOffset = kdaFwdHTilingData->vUpdateWorkspaceOffset; + hWorkspaceOffset = kdaFwdHTilingData->hWorkspaceOffset; + numSeqWorkspaceOffset = kdaFwdHTilingData->numSeqWorkspaceOffset; + numChunksWorkspaceOffset = kdaFwdHTilingData->numChunksWorkspaceOffset; + kDecayWorkspaceOffset = kdaFwdHTilingData->kDecayWorkspaceOffset; + uint64_t effectiveTaskCount = static_cast(isVariedLen ? tokenBatch : shapeBatch) * vNumHead; + useDirectFp32Ub = std::is_same::value && chunkSize <= 64 && kHeadDim == 128 && + vHeadDim == 128 && effectiveTaskCount >= AscendC::GetBlockNum(); + + gmK.SetGlobalBuffer((__gm__ ElementK *)k); + gmW.SetGlobalBuffer((__gm__ ElementW *)w); + gmU.SetGlobalBuffer((__gm__ ElementU *)u); + gmG.SetGlobalBuffer((__gm__ ElementG *)(scalarGated ? g : gk)); + gmGk.SetGlobalBuffer((__gm__ ElementG *)(kGated ? gk : g)); + gmInitialState.SetGlobalBuffer((__gm__ ElementInitialState *)inital_state); + gmH.SetGlobalBuffer((__gm__ ElementH *)h); + gmV.SetGlobalBuffer((__gm__ ElementV *)v_new); + gmFinalState.SetGlobalBuffer((__gm__ ElementFinalState *)final_state); + gmVWorkspace.SetGlobalBuffer((__gm__ ElementVWork *)(user + vWorkspaceOffset)); + gmVUpdateWorkspace.SetGlobalBuffer((__gm__ ElementV *)(user + vUpdateWorkspaceOffset)); + gmHWorkspace.SetGlobalBuffer((__gm__ ElementHWork *)(user + hWorkspaceOffset)); + gmKDecayWorkspace.SetGlobalBuffer((__gm__ ElementK *)(user + kDecayWorkspaceOffset)); + + gmSeqlen.SetGlobalBuffer((__gm__ int64_t *)cu_seqlens); + gmNumSeq.SetGlobalBuffer((__gm__ int64_t *)(user + numSeqWorkspaceOffset)); + gmNumChunks.SetGlobalBuffer((__gm__ int64_t *)(user + numChunksWorkspaceOffset)); + + ubHUpdatePing = resource.ubBuf.template GetBufferByByte(32 * 1024); + ubHUpdatePong = resource.ubBuf.template GetBufferByByte(96 * 1024); + ubVWorkPing = resource.ubBuf.template GetBufferByByte(32 * 1024); + ubVWorkPong = resource.ubBuf.template GetBufferByByte(96 * 1024); + + l1VUpdatePing = resource.l1Buf.template GetBufferByByte(0); + l1VUpdatePong = resource.l1Buf.template GetBufferByByte(chunkSize * vHeadDim * sizeof(ElementV)); + + if ASCEND_IS_AIC { + cubeBlockScheduler.Init(cu_seqlens, chunk_indices, tiling, user); + } + + if ASCEND_IS_AIV { + vecBlockScheduler.Init(cu_seqlens, chunk_indices, tiling, user); + } + } + + template + __aicore__ inline void InitFromData(GM_ADDR k, GM_ADDR w, GM_ADDR u, GM_ADDR g, GM_ADDR gk, GM_ADDR inital_state, + GM_ADDR cu_seqlens, GM_ADDR chunk_indices, GM_ADDR h, GM_ADDR v_new, + GM_ADDR final_state, const TilingData &tilingData, GM_ADDR user) + { + batch = tilingData.batch; + seqlen = tilingData.seqlen; + kNumHead = tilingData.kNumHead; + vNumHead = tilingData.vNumHead; + kHeadDim = tilingData.kHeadDim; + vHeadDim = tilingData.vHeadDim; + chunkSize = tilingData.chunkSize; + useInitialState = tilingData.useInitialState; + storeFinalState = tilingData.storeFinalState; + isVariedLen = tilingData.isVariedLen; + shapeBatch = tilingData.shapeBatch; + tokenBatch = tilingData.tokenBatch; + vWorkspaceOffset = tilingData.vWorkspaceOffset; + vUpdateWorkspaceOffset = tilingData.vUpdateWorkspaceOffset; + hWorkspaceOffset = tilingData.hWorkspaceOffset; + numSeqWorkspaceOffset = tilingData.numSeqWorkspaceOffset; + numChunksWorkspaceOffset = tilingData.numChunksWorkspaceOffset; + kDecayWorkspaceOffset = tilingData.kDecayWorkspaceOffset; + uint64_t effectiveTaskCount = static_cast(isVariedLen ? tokenBatch : shapeBatch) * vNumHead; + useDirectFp32Ub = std::is_same::value && chunkSize <= 64 && kHeadDim == 128 && + vHeadDim == 128 && effectiveTaskCount >= AscendC::GetBlockNum(); + + gmK.SetGlobalBuffer((__gm__ ElementK *)k); + gmW.SetGlobalBuffer((__gm__ ElementW *)w); + gmU.SetGlobalBuffer((__gm__ ElementU *)u); + gmG.SetGlobalBuffer((__gm__ ElementG *)(scalarGated ? g : gk)); + gmInitialState.SetGlobalBuffer((__gm__ ElementInitialState *)inital_state); + gmH.SetGlobalBuffer((__gm__ ElementH *)h); + gmV.SetGlobalBuffer((__gm__ ElementV *)v_new); + gmFinalState.SetGlobalBuffer((__gm__ ElementFinalState *)final_state); + gmVWorkspace.SetGlobalBuffer((__gm__ ElementVWork *)(user + vWorkspaceOffset)); + gmVUpdateWorkspace.SetGlobalBuffer((__gm__ ElementV *)(user + vUpdateWorkspaceOffset)); + gmHWorkspace.SetGlobalBuffer((__gm__ ElementHWork *)(user + hWorkspaceOffset)); + gmGk.SetGlobalBuffer((__gm__ ElementG *)(kGated ? gk : g)); + gmKDecayWorkspace.SetGlobalBuffer((__gm__ ElementK *)(user + kDecayWorkspaceOffset)); + gmSeqlen.SetGlobalBuffer((__gm__ int64_t *)cu_seqlens); + gmNumSeq.SetGlobalBuffer((__gm__ int64_t *)(user + numSeqWorkspaceOffset)); + gmNumChunks.SetGlobalBuffer((__gm__ int64_t *)(user + numChunksWorkspaceOffset)); + + ubHUpdatePing = resource.ubBuf.template GetBufferByByte(32 * 1024); + ubHUpdatePong = resource.ubBuf.template GetBufferByByte(96 * 1024); + ubVWorkPing = resource.ubBuf.template GetBufferByByte(32 * 1024); + ubVWorkPong = resource.ubBuf.template GetBufferByByte(96 * 1024); + + l1VUpdatePing = resource.l1Buf.template GetBufferByByte(0); + l1VUpdatePong = resource.l1Buf.template GetBufferByByte(chunkSize * vHeadDim * sizeof(ElementV)); + + if ASCEND_IS_AIC { + cubeBlockScheduler.InitFromData(cu_seqlens, chunk_indices, tilingData, user); + } + if ASCEND_IS_AIV { + vecBlockScheduler.InitFromData(cu_seqlens, chunk_indices, tilingData, user); + } + } + + // Tail helpers borrow the stream's V_MTE2 free token and restore it before returning. + __aicore__ inline void ComputeTailVWorkspace(const KdaFwdHOffsets &offsets, uint32_t tailEventId) + { + uint32_t subBlockIdx = AscendC::GetSubBlockIdx(); + uint32_t subBlockNum = AscendC::GetSubBlockNum(); + uint32_t rowsPerSubBlock = CeilDiv(offsets.blockTokens, subBlockNum); + uint32_t rowBegin = subBlockIdx * rowsPerSubBlock; + uint32_t rowEnd = Min(rowBegin + rowsPerSubBlock, offsets.blockTokens); + if (rowBegin >= rowEnd) { + return; + } + AscendC::ResetMask(); + AscendC::WaitFlag(tailEventId); + + if (offsets.vBlockDim == vHeadDim && vHeadDim <= 256) { + constexpr uint32_t TAIL_H_INPUT_OFFSET = 0; + constexpr uint32_t TAIL_W_INPUT_OFFSET = 166 * 1024; + constexpr uint32_t TAIL_OUTPUT_OFFSET = 168 * 1024; + uint32_t rowCount = rowEnd - rowBegin; + AscendC::LocalTensor hInputUb = + resource.ubBuf.template GetBufferByByte(TAIL_H_INPUT_OFFSET); + AscendC::LocalTensor wInputUb = + resource.ubBuf.template GetBufferByByte(TAIL_W_INPUT_OFFSET); + AscendC::LocalTensor outputUb = resource.ubBuf.template GetBufferByByte(TAIL_OUTPUT_OFFSET); + AscendC::DataCopy(hInputUb, gmH[offsets.hSrcOffset], kHeadDim * offsets.vBlockDim); + AscendC::DataCopy(wInputUb, gmW[offsets.wOffset + rowBegin * kHeadDim], rowCount * kHeadDim); + AscendC::SetFlag(tailEventId); + AscendC::WaitFlag(tailEventId); + __ubuf__ float *outputAddr = reinterpret_cast<__ubuf__ float *>(outputUb.GetPhyAddr()); + __ubuf__ ElementW *wInputAddr = reinterpret_cast<__ubuf__ ElementW *>(wInputUb.GetPhyAddr()); + __ubuf__ ElementH *hInputAddr = reinterpret_cast<__ubuf__ ElementH *>(hInputUb.GetPhyAddr()); + AscendC::VF_CALL< + Catlass::Epilogue::Block::detail::ComputeKdaTailVWorkspaceRegbaseDualIssue>( + outputAddr, wInputAddr, hInputAddr, static_cast(rowCount), + static_cast(offsets.vBlockDim), static_cast(kHeadDim)); + AscendC::PipeBarrier(); + AscendC::SetFlag(tailEventId); + AscendC::WaitFlag(tailEventId); + AscendC::DataCopy(gmVWorkspace[offsets.vWorkOffset + rowBegin * offsets.vBlockDim], outputUb, + rowCount * offsets.vBlockDim); + AscendC::SetFlag(tailEventId); + AscendC::WaitFlag(tailEventId); + AscendC::SetFlag(tailEventId); + return; + } + + constexpr uint32_t TAIL_INPUT_OFFSET = 166 * 1024; + constexpr uint32_t TAIL_FLOAT_OFFSET = 167 * 1024; + constexpr uint32_t TAIL_ACCUM_OFFSET = 168 * 1024; + constexpr uint32_t TAIL_WEIGHT_INPUT_OFFSET = 169 * 1024; + constexpr uint32_t TAIL_WEIGHT_FLOAT_OFFSET = 170 * 1024; + AscendC::LocalTensor inputUb = resource.ubBuf.template GetBufferByByte(TAIL_INPUT_OFFSET); + AscendC::LocalTensor floatUb = resource.ubBuf.template GetBufferByByte(TAIL_FLOAT_OFFSET); + AscendC::LocalTensor accumUb = resource.ubBuf.template GetBufferByByte(TAIL_ACCUM_OFFSET); + AscendC::LocalTensor weightInputUb = + resource.ubBuf.template GetBufferByByte(TAIL_WEIGHT_INPUT_OFFSET); + AscendC::LocalTensor weightFloatUb = + resource.ubBuf.template GetBufferByByte(TAIL_WEIGHT_FLOAT_OFFSET); + + for (uint32_t tokenRow = rowBegin; tokenRow < rowEnd; ++tokenRow) { + AscendC::DataCopy(weightInputUb, gmW[offsets.wOffset + tokenRow * kHeadDim], kHeadDim); + AscendC::SetFlag(tailEventId); + AscendC::WaitFlag(tailEventId); + AscendC::Cast(weightFloatUb, weightInputUb, AscendC::RoundMode::CAST_NONE, kHeadDim); + AscendC::PipeBarrier(); + AscendC::SetFlag(tailEventId); + AscendC::WaitFlag(tailEventId); + + AscendC::Duplicate(accumUb, 0.0f, offsets.vBlockDim); + AscendC::PipeBarrier(); + for (uint32_t kIdx = 0; kIdx < kHeadDim; ++kIdx) { + AscendC::DataCopy(inputUb, gmH[offsets.hSrcOffset + kIdx * vHeadDim], offsets.vBlockDim); + AscendC::SetFlag(tailEventId); + AscendC::WaitFlag(tailEventId); + AscendC::Cast(floatUb, inputUb, AscendC::RoundMode::CAST_NONE, offsets.vBlockDim); + AscendC::PipeBarrier(); + float weight = weightFloatUb.GetValue(kIdx); + AscendC::SetFlag(tailEventId); + AscendC::WaitFlag(tailEventId); + AscendC::Muls(floatUb, floatUb, weight, offsets.vBlockDim); + AscendC::PipeBarrier(); + AscendC::Add(accumUb, accumUb, floatUb, offsets.vBlockDim); + AscendC::PipeBarrier(); + AscendC::SetFlag(tailEventId); + AscendC::WaitFlag(tailEventId); + } + AscendC::SetFlag(tailEventId); + AscendC::WaitFlag(tailEventId); + AscendC::SetFlag(tailEventId); + AscendC::WaitFlag(tailEventId); + AscendC::DataCopy(gmVWorkspace[offsets.vWorkOffset + tokenRow * offsets.vBlockDim], accumUb, + offsets.vBlockDim); + AscendC::SetFlag(tailEventId); + AscendC::WaitFlag(tailEventId); + } + AscendC::SetFlag(tailEventId); + } + + __aicore__ inline void ComputeTailHWorkspace(const KdaFwdHOffsets &offsets, uint32_t tailEventId) + { + uint32_t subBlockIdx = AscendC::GetSubBlockIdx(); + uint32_t subBlockNum = AscendC::GetSubBlockNum(); + uint32_t rowsPerSubBlock = CeilDiv(kHeadDim, subBlockNum); + uint32_t rowBegin = subBlockIdx * rowsPerSubBlock; + uint32_t rowEnd = Min(rowBegin + rowsPerSubBlock, kHeadDim); + AscendC::ResetMask(); + AscendC::WaitFlag(tailEventId); + + constexpr uint32_t TAIL_INPUT_OFFSET = 166 * 1024; + constexpr uint32_t TAIL_FLOAT_OFFSET = 167 * 1024; + constexpr uint32_t TAIL_ACCUM_OFFSET = 168 * 1024; + constexpr uint32_t TAIL_WEIGHT_INPUT_OFFSET = 169 * 1024; + constexpr uint32_t TAIL_WEIGHT_FLOAT_OFFSET = 170 * 1024; + AscendC::LocalTensor inputUb = resource.ubBuf.template GetBufferByByte(TAIL_INPUT_OFFSET); + AscendC::LocalTensor floatUb = resource.ubBuf.template GetBufferByByte(TAIL_FLOAT_OFFSET); + AscendC::LocalTensor accumUb = resource.ubBuf.template GetBufferByByte(TAIL_ACCUM_OFFSET); + AscendC::LocalTensor weightInputUb = + resource.ubBuf.template GetBufferByByte(TAIL_WEIGHT_INPUT_OFFSET); + AscendC::LocalTensor weightFloatUb = + resource.ubBuf.template GetBufferByByte(TAIL_WEIGHT_FLOAT_OFFSET); + + for (uint32_t kRow = rowBegin; kRow < rowEnd; ++kRow) { + AscendC::Duplicate(accumUb, 0.0f, offsets.vBlockDim); + AscendC::PipeBarrier(); + for (uint32_t tokenRow = 0; tokenRow < offsets.blockTokens; ++tokenRow) { + AscendC::DataCopy(weightInputUb, gmKDecayWorkspace[offsets.kDecayWorkOffset + tokenRow * kHeadDim], + kHeadDim); + AscendC::SetFlag(tailEventId); + AscendC::WaitFlag(tailEventId); + AscendC::Cast(weightFloatUb, weightInputUb, AscendC::RoundMode::CAST_NONE, kHeadDim); + AscendC::PipeBarrier(); + AscendC::SetFlag(tailEventId); + AscendC::WaitFlag(tailEventId); + AscendC::DataCopy(inputUb, gmVUpdateWorkspace[offsets.vWorkOffset + tokenRow * offsets.vBlockDim], + offsets.vBlockDim); + AscendC::SetFlag(tailEventId); + AscendC::WaitFlag(tailEventId); + AscendC::Cast(floatUb, inputUb, AscendC::RoundMode::CAST_NONE, offsets.vBlockDim); + AscendC::PipeBarrier(); + float weight = weightFloatUb.GetValue(kRow); + AscendC::SetFlag(tailEventId); + AscendC::WaitFlag(tailEventId); + AscendC::Muls(floatUb, floatUb, weight, offsets.vBlockDim); + AscendC::PipeBarrier(); + AscendC::Add(accumUb, accumUb, floatUb, offsets.vBlockDim); + AscendC::PipeBarrier(); + AscendC::SetFlag(tailEventId); + AscendC::WaitFlag(tailEventId); + AscendC::SetFlag(tailEventId); + AscendC::WaitFlag(tailEventId); + } + AscendC::SetFlag(tailEventId); + AscendC::WaitFlag(tailEventId); + AscendC::DataCopy(gmHWorkspace[offsets.hWorkOffset + kRow * offsets.vBlockDim], accumUb, offsets.vBlockDim); + AscendC::SetFlag(tailEventId); + AscendC::WaitFlag(tailEventId); + } + AscendC::SetFlag(tailEventId); + } + + __aicore__ inline void ComputeKdaTailFinalState(const KdaFwdHOffsets &offsets, uint32_t tailEventId) + { + static_assert(kGated && !scalarGated && useExp2 && std::is_same::value, + "The fused tail state path is specific to FP32 KDA state"); + constexpr uint32_t TAIL_STATE_OFFSET = 0; + constexpr uint32_t TAIL_K_INPUT_OFFSET = 166 * 1024; + constexpr uint32_t TAIL_V_INPUT_OFFSET = 170 * 1024; + constexpr uint32_t TAIL_GATE_INPUT_OFFSET = 174 * 1024; + constexpr uint32_t TAIL_GATE_SCALE_OFFSET = 175 * 1024; + + uint32_t subBlockIdx = AscendC::GetSubBlockIdx(); + uint32_t subBlockNum = AscendC::GetSubBlockNum(); + uint32_t rowsPerSubBlock = CeilDiv(kHeadDim, subBlockNum); + uint32_t rowBegin = subBlockIdx * rowsPerSubBlock; + uint32_t rowEnd = Min(rowBegin + rowsPerSubBlock, kHeadDim); + if (rowBegin >= rowEnd) { + return; + } + + AscendC::ResetMask(); + AscendC::WaitFlag(tailEventId); + AscendC::LocalTensor stateUb = resource.ubBuf.template GetBufferByByte(TAIL_STATE_OFFSET); + AscendC::LocalTensor kInputUb = + resource.ubBuf.template GetBufferByByte(TAIL_K_INPUT_OFFSET); + AscendC::LocalTensor vInputUb = + resource.ubBuf.template GetBufferByByte(TAIL_V_INPUT_OFFSET); + AscendC::LocalTensor gateInputUb = + resource.ubBuf.template GetBufferByByte(TAIL_GATE_INPUT_OFFSET); + AscendC::LocalTensor gateScaleUb = + resource.ubBuf.template GetBufferByByte(TAIL_GATE_SCALE_OFFSET); + + uint32_t gateOffset = offsets.gkOffset + (offsets.blockTokens - 1) * kHeadDim + rowBegin; + uint32_t rowCount = rowEnd - rowBegin; + uint32_t stateElements = rowCount * offsets.vBlockDim; + uint32_t kElements = offsets.blockTokens * kHeadDim; + uint32_t vElements = offsets.blockTokens * offsets.vBlockDim; + + if (offsets.isInitialState) { + if (useInitialState) { + AscendC::DataCopy(stateUb, gmInitialState[offsets.initialStateOffset + rowBegin * offsets.vBlockDim], + stateElements); + } else { + AscendC::Duplicate(stateUb, 0.0f, stateElements); + } + } else { + AscendC::DataCopy(stateUb, gmFinalState[offsets.finalStateOffset + rowBegin * offsets.vBlockDim], + stateElements); + } + AscendC::DataCopy(kInputUb, gmK[offsets.wkOffset], kElements); + AscendC::DataCopy(vInputUb, gmV[offsets.uvOffset], vElements); + AscendC::DataCopy(gateInputUb, gmGk[gateOffset], rowCount); + AscendC::SetFlag(tailEventId); + AscendC::WaitFlag(tailEventId); + + __ubuf__ float *stateAddr = reinterpret_cast<__ubuf__ float *>(stateUb.GetPhyAddr()); + __ubuf__ ElementK *kInputAddr = reinterpret_cast<__ubuf__ ElementK *>(kInputUb.GetPhyAddr()); + __ubuf__ ElementV *vInputAddr = reinterpret_cast<__ubuf__ ElementV *>(vInputUb.GetPhyAddr()); + __ubuf__ ElementG *gateInputAddr = reinterpret_cast<__ubuf__ ElementG *>(gateInputUb.GetPhyAddr()); + __ubuf__ float *gateScaleAddr = reinterpret_cast<__ubuf__ float *>(gateScaleUb.GetPhyAddr()); + AscendC::VF_CALL>( + gateScaleAddr, gateInputAddr, static_cast(rowCount)); + AscendC::VF_CALL< + Catlass::Epilogue::Block::detail::ComputeKdaTailFinalStateRegbaseDualIssue>( + stateAddr, kInputAddr, vInputAddr, gateScaleAddr, static_cast(rowCount), + static_cast(offsets.vBlockDim), static_cast(offsets.blockTokens), + static_cast(kHeadDim), static_cast(rowBegin)); + AscendC::PipeBarrier(); + AscendC::SetFlag(tailEventId); + AscendC::WaitFlag(tailEventId); + AscendC::DataCopy(gmFinalState[offsets.finalStateOffset + rowBegin * offsets.vBlockDim], stateUb, + stateElements); + AscendC::SetFlag(tailEventId); + AscendC::WaitFlag(tailEventId); + AscendC::SetFlag(tailEventId); + } + + __aicore__ inline void Process() + { + // FwdH can run after another stage in a megakernel. Start its AIC/AIV + // handshake only after every core has retired the preceding stage. + AscendC::SyncAll(); + + if ASCEND_IS_AIC { + uint32_t coreIdx = AscendC::GetBlockIdx(); + uint32_t coreNum = vecBlockScheduler.cubeCoreNum; + + BlockMmadWH blockMmadWH(resource, + chunkSize * cubeBlockScheduler.vBlockSize * sizeof(ElementV) * PING_PONG_STAGES); + BlockMmadKV blockMmadKV(resource, + chunkSize * cubeBlockScheduler.vBlockSize * sizeof(ElementV) * PING_PONG_STAGES); + BlockMmadWHTail blockMmadWHTail( + resource, chunkSize * cubeBlockScheduler.vBlockSize * sizeof(ElementV) * PING_PONG_STAGES); + BlockMmadKVTail blockMmadKVTail( + resource, chunkSize * cubeBlockScheduler.vBlockSize * sizeof(ElementV) * PING_PONG_STAGES); + bool useBoundedMmad = isVariedLen || (seqlen % chunkSize != 0); + + auto wLayout = + tla::MakeLayout(shapeBatch * kNumHead * cubeBlockScheduler.totalTokens, kHeadDim); + auto hLayout = tla::MakeLayout( + shapeBatch * vNumHead * cubeBlockScheduler.totalChunks * kHeadDim, vHeadDim); + + auto kLayout = + tla::MakeLayout(kHeadDim, shapeBatch * kNumHead * cubeBlockScheduler.totalTokens); + auto hworkLayout = tla::MakeLayout(kHeadDim, cubeBlockScheduler.vBlockSize); + + AscendC::SyncAll(); + uint32_t currStage = 0; // 0: C1, 1: C2 + while (cubeBlockScheduler.isRunning) { + if (currStage == 0) { + /* C1: v_work = w @ h[i] */ + cubeBlockScheduler.InitTasks(); + if (useDirectFp32Ub) { + BlockMmadWHDirectUb blockMmadWHDirectUb( + resource, chunkSize * cubeBlockScheduler.vBlockSize * sizeof(ElementV) * PING_PONG_STAGES); + for (uint32_t i = 0; i < PING_PONG_STAGES; ++i) { + uint32_t streamId = cubeBlockScheduler.GetStreamId(i); + const auto &stream = cubeBlockScheduler.GetStream(i); + if (cubeBlockScheduler.StreamIsDone(stream)) { + continue; + } + + const KdaFwdHOffsets &cube1Offsets = cubeBlockScheduler.GetCurTaskOffsets(stream); + Arch::CrossCoreWaitFlag(cubeBlockScheduler.vec2Done[streamId]); + if (cube1Offsets.blockTokens < 16) { + Arch::CrossCoreSetFlag<0x2, PIPE_MTE2>(cubeBlockScheduler.cube1Done[streamId]); + continue; + } + int64_t cube1OffsetW = cube1Offsets.wOffset; + int64_t cube1OffsetH = cube1Offsets.hSrcOffset; + auto tensorW = tla::MakeTensor(gmW[cube1OffsetW], wLayout, Catlass::Arch::PositionGM{}); + auto tensorH = tla::MakeTensor(gmH[cube1OffsetH], hLayout, Catlass::Arch::PositionGM{}); + GemmCoord cube1Shape{cube1Offsets.blockTokens, cube1Offsets.vBlockDim, kHeadDim}; + auto tensorBlockW = + GetTile(tensorW, tla::MakeCoord(0, 0), tla::MakeShape(cube1Shape.m(), cube1Shape.k())); + auto tensorBlockH = + GetTile(tensorH, tla::MakeCoord(0, 0), tla::MakeShape(cube1Shape.k(), cube1Shape.n())); + + auto ubLayout = tla::MakeLayout(cube1Shape.m(), cube1Shape.n()); + auto tensorUbPing = tla::MakeTensor(ubVWorkPing, ubLayout, Catlass::Arch::PositionUB{}); + auto tensorUbPong = tla::MakeTensor(ubVWorkPong, ubLayout, Catlass::Arch::PositionUB{}); + using UbTensor = decltype(tensorUbPing); + UbTensor tensorUbList[BlockMmadWHDirectUb::MAX_CUBE_VEC_SYNC_NUM]; + for (uint32_t ubIdx = 0; ubIdx < BlockMmadWHDirectUb::MAX_CUBE_VEC_SYNC_NUM; ++ubIdx) { + tensorUbList[ubIdx] = (ubIdx & 1U) ? tensorUbPong : tensorUbPing; + } + uint32_t ubListId = streamId; + uint32_t rowsPerSubBlock = CeilDiv(cube1Shape.m(), DIRECT_VEC_NUM); + blockMmadWHDirectUb(tensorBlockW, tensorBlockH, tensorUbList, cube1Shape, rowsPerSubBlock, + 0, DIRECT_UB_FREE_FLAG_BEGIN, DIRECT_UB_READY_FLAG_BEGIN, ubListId, + DIRECT_VEC_NUM, DIRECT_UB_STAGES); + } + } else if (useBoundedMmad) { + for (uint32_t i = 0; i < PING_PONG_STAGES; ++i) { + uint32_t streamId = cubeBlockScheduler.GetStreamId(i); + const auto &stream = cubeBlockScheduler.GetStream(i); + if (cubeBlockScheduler.StreamIsDone(stream)) { + continue; + } + + const KdaFwdHOffsets &cube1Offsets = cubeBlockScheduler.GetCurTaskOffsets(stream); + Arch::CrossCoreWaitFlag(cubeBlockScheduler.vec2Done[streamId]); + if (cube1Offsets.blockTokens < 16) { + Arch::CrossCoreSetFlag<0x2, PIPE_MTE2>(cubeBlockScheduler.cube1Done[streamId]); + continue; + } + auto vLayout = tla::MakeLayout(cube1Offsets.blockTokens, + cube1Offsets.vBlockDim); + auto tensorW = + tla::MakeTensor(gmW[cube1Offsets.wOffset], wLayout, Catlass::Arch::PositionGM{}); + auto tensorH = + tla::MakeTensor(gmH[cube1Offsets.hSrcOffset], hLayout, Catlass::Arch::PositionGM{}); + auto tensorV = tla::MakeTensor(gmVWorkspace[cube1Offsets.vWorkOffset], vLayout, + Catlass::Arch::PositionGM{}); + GemmCoord cube1Shape{cube1Offsets.blockTokens, cube1Offsets.vBlockDim, kHeadDim}; + auto tensorBlockW = + GetTile(tensorW, tla::MakeCoord(0, 0), tla::MakeShape(cube1Shape.m(), cube1Shape.k())); + auto tensorBlockH = + GetTile(tensorH, tla::MakeCoord(0, 0), tla::MakeShape(cube1Shape.k(), cube1Shape.n())); + auto tensorBlockV = + GetTile(tensorV, tla::MakeCoord(0, 0), tla::MakeShape(cube1Shape.m(), cube1Shape.n())); + + blockMmadWHTail.preSetFlags(); + blockMmadWHTail(tensorBlockW, tensorBlockH, tensorBlockV, cube1Shape, EmptyClass{}, true); + blockMmadWHTail.finalWaitFlags(); + AscendC::PipeBarrier(); + Arch::CrossCoreSetFlag<0x2, PIPE_FIX>(cubeBlockScheduler.cube1Done[streamId]); + } + } else { + blockMmadWH.preSetFlags(); + for (uint32_t i = 0; i < PING_PONG_STAGES; ++i) { + uint32_t streamId = cubeBlockScheduler.GetStreamId(i); + const auto &stream = cubeBlockScheduler.GetStream(i); + if (cubeBlockScheduler.StreamIsDone(stream)) { + continue; + } + + const KdaFwdHOffsets &cube1Offsets = cubeBlockScheduler.GetCurTaskOffsets(stream); + Arch::CrossCoreWaitFlag(cubeBlockScheduler.vec2Done[streamId]); + if (cube1Offsets.blockTokens < 16) { + Arch::CrossCoreSetFlag<0x2, PIPE_MTE2>(cubeBlockScheduler.cube1Done[streamId]); + continue; + } + auto vLayout = tla::MakeLayout(cube1Offsets.blockTokens, + cube1Offsets.vBlockDim); + int64_t cube1OffsetW = cube1Offsets.wOffset; + int64_t cube1OffsetH = cube1Offsets.hSrcOffset; + int64_t cube1OffsetVwork = cube1Offsets.vWorkOffset; + auto tensorW = tla::MakeTensor(gmW[cube1OffsetW], wLayout, Catlass::Arch::PositionGM{}); + auto tensorH = tla::MakeTensor(gmH[cube1OffsetH], hLayout, Catlass::Arch::PositionGM{}); + auto tensorV = + tla::MakeTensor(gmVWorkspace[cube1OffsetVwork], vLayout, Catlass::Arch::PositionGM{}); + GemmCoord cube1Shape{cube1Offsets.blockTokens, cube1Offsets.vBlockDim, kHeadDim}; + auto tensorBlockW = + GetTile(tensorW, tla::MakeCoord(0, 0), tla::MakeShape(cube1Shape.m(), cube1Shape.k())); + auto tensorBlockH = + GetTile(tensorH, tla::MakeCoord(0, 0), tla::MakeShape(cube1Shape.k(), cube1Shape.n())); + auto tensorBlockV = + GetTile(tensorV, tla::MakeCoord(0, 0), tla::MakeShape(cube1Shape.m(), cube1Shape.n())); + + blockMmadWH(tensorBlockW, tensorBlockH, tensorBlockV, cube1Shape); + Arch::CrossCoreSetFlag<0x2, PIPE_FIX>(cubeBlockScheduler.cube1Done[streamId]); + } + blockMmadWH.finalWaitFlags(); + } + } else { + /* C2: h[i+1] = k.T @ v_work */ + if (useDirectFp32Ub) { + BlockMmadKVDirectUb blockMmadKVDirectUb( + resource, chunkSize * cubeBlockScheduler.vBlockSize * sizeof(ElementV) * PING_PONG_STAGES); + for (uint32_t i = 0; i < PING_PONG_STAGES; ++i) { + uint32_t streamId = cubeBlockScheduler.GetStreamId(i); + const auto &stream = cubeBlockScheduler.GetStream(i); + if (cubeBlockScheduler.StreamIsDone(stream)) { + continue; + } + const KdaFwdHOffsets &cube2Offsets = cubeBlockScheduler.GetCurTaskOffsets(stream); + Arch::CrossCoreWaitFlag(cubeBlockScheduler.vec1Done[streamId]); + + if (cubeBlockScheduler.NeedProcessStage2(stream)) { + if (cube2Offsets.blockTokens < 16) { + Arch::CrossCoreSetFlag<0x2, PIPE_MTE2>(cubeBlockScheduler.cube2Done[streamId]); + continue; + } + int64_t cube2OffsetK = kGated ? cube2Offsets.kDecayWorkOffset : cube2Offsets.wkOffset; + int64_t cube2OffsetVwork = cube2Offsets.vWorkOffset; + auto tensorK = + kGated ? tla::MakeTensor(gmKDecayWorkspace[cube2OffsetK], kLayout, + Catlass::Arch::PositionGM{}) : + tla::MakeTensor(gmK[cube2OffsetK], kLayout, Catlass::Arch::PositionGM{}); + auto vUpdateLayout = tla::MakeLayout( + cube2Offsets.blockTokens, cube2Offsets.vBlockDim); + auto tensorVwork = tla::MakeTensor(gmVUpdateWorkspace[cube2OffsetVwork], vUpdateLayout, + Catlass::Arch::PositionGM{}); + GemmCoord cube2Shape{kHeadDim, cube2Offsets.vBlockDim, cube2Offsets.blockTokens}; + auto tensorBlockK = GetTile(tensorK, tla::MakeCoord(0, 0), + tla::MakeShape(cube2Shape.m(), cube2Shape.k())); + auto tensorBlockVwork = GetTile(tensorVwork, tla::MakeCoord(0, 0), + tla::MakeShape(cube2Shape.k(), cube2Shape.n())); + + auto ubLayout = tla::MakeLayout(cube2Shape.m(), cube2Shape.n()); + auto tensorUbPing = + tla::MakeTensor(ubHUpdatePing, ubLayout, Catlass::Arch::PositionUB{}); + auto tensorUbPong = + tla::MakeTensor(ubHUpdatePong, ubLayout, Catlass::Arch::PositionUB{}); + using UbTensor = decltype(tensorUbPing); + UbTensor tensorUbList[BlockMmadKVDirectUb::MAX_CUBE_VEC_SYNC_NUM]; + for (uint32_t ubIdx = 0; ubIdx < BlockMmadKVDirectUb::MAX_CUBE_VEC_SYNC_NUM; ++ubIdx) { + tensorUbList[ubIdx] = (ubIdx & 1U) ? tensorUbPong : tensorUbPing; + } + uint32_t ubListId = streamId; + uint32_t rowsPerSubBlock = CeilDiv(cube2Shape.m(), DIRECT_VEC_NUM); + blockMmadKVDirectUb(tensorBlockK, tensorBlockVwork, tensorUbList, cube2Shape, + rowsPerSubBlock, 0, DIRECT_UB_FREE_FLAG_BEGIN, + DIRECT_UB_READY_FLAG_BEGIN, ubListId, DIRECT_VEC_NUM, + DIRECT_UB_STAGES); + } + } + } else if (useBoundedMmad) { + for (uint32_t i = 0; i < PING_PONG_STAGES; ++i) { + uint32_t streamId = cubeBlockScheduler.GetStreamId(i); + const auto &stream = cubeBlockScheduler.GetStream(i); + if (cubeBlockScheduler.StreamIsDone(stream)) { + continue; + } + const KdaFwdHOffsets &cube2Offsets = cubeBlockScheduler.GetCurTaskOffsets(stream); + Arch::CrossCoreWaitFlag(cubeBlockScheduler.vec1Done[streamId]); + + if (cubeBlockScheduler.NeedProcessStage2(stream)) { + if (cube2Offsets.blockTokens < 16) { + Arch::CrossCoreSetFlag<0x2, PIPE_MTE2>(cubeBlockScheduler.cube2Done[streamId]); + continue; + } + int64_t cube2OffsetK = kGated ? cube2Offsets.kDecayWorkOffset : cube2Offsets.wkOffset; + auto tensorK = + kGated ? tla::MakeTensor(gmKDecayWorkspace[cube2OffsetK], kLayout, + Catlass::Arch::PositionGM{}) : + tla::MakeTensor(gmK[cube2OffsetK], kLayout, Catlass::Arch::PositionGM{}); + auto vUpdateLayout = tla::MakeLayout( + cube2Offsets.blockTokens, cube2Offsets.vBlockDim); + auto tensorVwork = tla::MakeTensor(gmVUpdateWorkspace[cube2Offsets.vWorkOffset], + vUpdateLayout, Catlass::Arch::PositionGM{}); + auto tensorHwork = tla::MakeTensor(gmHWorkspace[cube2Offsets.hWorkOffset], hworkLayout, + Catlass::Arch::PositionGM{}); + GemmCoord cube2Shape{kHeadDim, cube2Offsets.vBlockDim, cube2Offsets.blockTokens}; + auto tensorBlockK = GetTile(tensorK, tla::MakeCoord(0, 0), + tla::MakeShape(cube2Shape.m(), cube2Shape.k())); + auto tensorBlockVwork = GetTile(tensorVwork, tla::MakeCoord(0, 0), + tla::MakeShape(cube2Shape.k(), cube2Shape.n())); + auto tensorBlockHwork = GetTile(tensorHwork, tla::MakeCoord(0, 0), + tla::MakeShape(cube2Shape.m(), cube2Shape.n())); + + blockMmadKVTail.preSetFlags(); + blockMmadKVTail(tensorBlockK, tensorBlockVwork, tensorBlockHwork, cube2Shape, + EmptyClass{}, true); + blockMmadKVTail.finalWaitFlags(); + AscendC::PipeBarrier(); + } + Arch::CrossCoreSetFlag<0x2, PIPE_FIX>(cubeBlockScheduler.cube2Done[streamId]); + } + } else { + blockMmadKV.preSetFlags(); + for (uint32_t i = 0; i < PING_PONG_STAGES; ++i) { + uint32_t streamId = cubeBlockScheduler.GetStreamId(i); + const auto &stream = cubeBlockScheduler.GetStream(i); + if (cubeBlockScheduler.StreamIsDone(stream)) { + continue; + } + const KdaFwdHOffsets &cube2Offsets = cubeBlockScheduler.GetCurTaskOffsets(stream); + Arch::CrossCoreWaitFlag(cubeBlockScheduler.vec1Done[streamId]); + + if (cubeBlockScheduler.NeedProcessStage2(stream)) { + if (cube2Offsets.blockTokens < 16) { + Arch::CrossCoreSetFlag<0x2, PIPE_MTE2>(cubeBlockScheduler.cube2Done[streamId]); + continue; + } + // step 3: h[i+1] = k.T @ v_work + int64_t cube2OffsetK = kGated ? cube2Offsets.kDecayWorkOffset : cube2Offsets.wkOffset; + int64_t cube2OffsetVwork = cube2Offsets.vWorkOffset; + auto tensorK = + kGated ? tla::MakeTensor(gmKDecayWorkspace[cube2OffsetK], kLayout, + Catlass::Arch::PositionGM{}) : + tla::MakeTensor(gmK[cube2OffsetK], kLayout, Catlass::Arch::PositionGM{}); + auto vUpdateLayout = tla::MakeLayout( + cube2Offsets.blockTokens, cube2Offsets.vBlockDim); + auto tensorVwork = tla::MakeTensor(gmVUpdateWorkspace[cube2OffsetVwork], vUpdateLayout, + Catlass::Arch::PositionGM{}); + auto tensorHwork = tla::MakeTensor(gmHWorkspace[cube2Offsets.hWorkOffset], hworkLayout, + Catlass::Arch::PositionGM{}); + GemmCoord cube2Shape{kHeadDim, cube2Offsets.vBlockDim, cube2Offsets.blockTokens}; + auto tensorBlockK = GetTile(tensorK, tla::MakeCoord(0, 0), + tla::MakeShape(cube2Shape.m(), cube2Shape.k())); + auto tensorBlockVwork = GetTile(tensorVwork, tla::MakeCoord(0, 0), + tla::MakeShape(cube2Shape.k(), cube2Shape.n())); + auto tensorBlockHwork = GetTile(tensorHwork, tla::MakeCoord(0, 0), + tla::MakeShape(cube2Shape.m(), cube2Shape.n())); + + blockMmadKV(tensorBlockK, tensorBlockVwork, tensorBlockHwork, cube2Shape); + } + Arch::CrossCoreSetFlag<0x2, PIPE_FIX>(cubeBlockScheduler.cube2Done[streamId]); + } + blockMmadKV.finalWaitFlags(); + } + } + currStage ^= 0x01; + } + Arch::CrossCoreWaitFlag(cubeBlockScheduler.vec2Done[0]); + Arch::CrossCoreWaitFlag(cubeBlockScheduler.vec2Done[1]); + if (useDirectFp32Ub) { + for (uint32_t slot = 0; slot < DIRECT_UB_STAGES; ++slot) { + AscendC::CrossCoreWaitFlag<0x4, PIPE_FIX>(DIRECT_UB_FREE_FLAG_BEGIN + slot); + AscendC::CrossCoreWaitFlag<0x4, PIPE_FIX>(DIRECT_UB_FREE_FLAG_BEGIN + DIRECT_UB_FLAG_STRIDE + slot); + } + } + } + + if ASCEND_IS_AIV { + uint32_t subBlockIdx = AscendC::GetSubBlockIdx(); + uint32_t subBlockNum = AscendC::GetSubBlockNum(); + uint32_t coreIdx = AscendC::GetBlockIdx() / subBlockNum; + uint32_t coreNum = AscendC::GetBlockNum(); + uint32_t taskCount = (isVariedLen ? vecBlockScheduler.tokenBatch : shapeBatch) * vNumHead; + uint32_t tasksPerCore = taskCount > coreNum ? PING_PONG_STAGES : 1; + uint32_t taskStride = coreNum * tasksPerCore; + uint32_t rowsPerSubBlock = (kHeadDim + subBlockNum - 1) / subBlockNum; + uint32_t rowBegin = subBlockIdx * rowsPerSubBlock; + uint32_t rowEnd = Min(rowBegin + rowsPerSubBlock, kHeadDim); + uint32_t hRowsPerTile = (32 * 1024) / (vHeadDim * sizeof(ElementH)); + uint32_t stateRowsPerTile = (64 * 1024) / (vHeadDim * sizeof(ElementInitialState)); + uint32_t rowsPerTile = Min(hRowsPerTile, stateRowsPerTile); + uint32_t totalChunks = isVariedLen ? vecBlockScheduler.totalChunks : ((seqlen + chunkSize - 1) / chunkSize); + uint32_t stateBlockSize = kHeadDim * vHeadDim; + uint32_t pingpongFlag = 1; + AscendC::LocalTensor stateUbTensorPing = + resource.ubBuf.template GetBufferByByte(0); + AscendC::LocalTensor stateUbTensorPong = + resource.ubBuf.template GetBufferByByte(96 * 1024); + AscendC::LocalTensor hUbTensorPing = resource.ubBuf.template GetBufferByByte(64 * 1024); + AscendC::LocalTensor hUbTensorPong = + resource.ubBuf.template GetBufferByByte(160 * 1024); + AscendC::SetFlag(EVENT_ID0); + AscendC::SetFlag(EVENT_ID1); + for (uint32_t slot = 0; slot < tasksPerCore; ++slot) { + for (uint32_t taskIdx = coreIdx * tasksPerCore + slot; taskIdx < taskCount; taskIdx += taskStride) { + uint32_t batchIdx = taskIdx / vNumHead; + uint32_t vHeadIdx = taskIdx % vNumHead; + uint32_t chunkOffset = isVariedLen ? vecBlockScheduler.GetVarlenChunkOffset(batchIdx) : 0; + uint32_t shapeBatchIdx = isVariedLen ? 0 : batchIdx; + uint32_t hBaseOffset = + (shapeBatchIdx * vNumHead * totalChunks + vHeadIdx * totalChunks + chunkOffset) * + stateBlockSize; + uint32_t initialStateBaseOffset = taskIdx * stateBlockSize; + for (uint32_t rowOffset = rowBegin; rowOffset < rowEnd; rowOffset += rowsPerTile) { + uint32_t rowsThisTile = Min(rowsPerTile, rowEnd - rowOffset); + uint32_t stateTileElems = rowsThisTile * vHeadDim; + uint32_t hOffset = hBaseOffset + rowOffset * vHeadDim; + AscendC::LocalTensor stateUbTensor = + pingpongFlag ? stateUbTensorPing : stateUbTensorPong; + AscendC::LocalTensor hUbTensor = pingpongFlag ? hUbTensorPing : hUbTensorPong; + auto eventId = pingpongFlag ? EVENT_ID1 : EVENT_ID0; + AscendC::WaitFlag(eventId); + if (useInitialState) { + uint32_t initialStateOffset = initialStateBaseOffset + rowOffset * vHeadDim; + if constexpr (!std::is_same::value) { + AscendC::DataCopy(stateUbTensor, gmInitialState[initialStateOffset], stateTileElems); + AscendC::SetFlag(eventId); + AscendC::WaitFlag(eventId); + AscendC::Cast(hUbTensor, stateUbTensor, AscendC::RoundMode::CAST_RINT, stateTileElems); + AscendC::SetFlag(eventId); + AscendC::WaitFlag(eventId); + AscendC::DataCopy(gmH[hOffset], hUbTensor, stateTileElems); + } else { + AscendC::DataCopy(stateUbTensor, gmInitialState[initialStateOffset], stateTileElems); + AscendC::SetFlag(eventId); + AscendC::WaitFlag(eventId); + AscendC::DataCopy(gmH[hOffset], stateUbTensor, stateTileElems); + } + } else { + AscendC::Duplicate(hUbTensor, static_cast(0), stateTileElems); + AscendC::SetFlag(eventId); + AscendC::WaitFlag(eventId); + AscendC::DataCopy(gmH[hOffset], hUbTensor, stateTileElems); + } + AscendC::SetFlag(eventId); + pingpongFlag = 1 - pingpongFlag; + } + } + } + AscendC::WaitFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID1); + + AscendC::SyncAll(); + + if (useDirectFp32Ub) { + for (uint32_t slot = 0; slot < DIRECT_UB_STAGES; ++slot) { + uint64_t flagOffset = slot + subBlockIdx * DIRECT_UB_FLAG_STRIDE; + AscendC::CrossCoreSetFlag<0x4, PIPE_V>(DIRECT_UB_FREE_FLAG_BEGIN + flagOffset); + } + } + Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(vecBlockScheduler.vec2Done[0]); + Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(vecBlockScheduler.vec2Done[1]); + + EpilogueKdaFwdHVnew epilogueKdaFwdHVnew(resource); + EpilogueKdaFwdHUpdate epilogueKdaFwdHUpdate(resource); + uint32_t pongBaseEvent = 4; + + if (storeFinalState && std::is_same::value) { + AscendC::SetFlag(EVENT_ID0); // preset final_state + AscendC::SetFlag(EVENT_ID0 + pongBaseEvent); + AscendC::SetFlag(EVENT_ID2); // preset h + AscendC::SetFlag(EVENT_ID2 + pongBaseEvent); + } else { + AscendC::SetFlag(EVENT_ID0); // preset h_update + AscendC::SetFlag(EVENT_ID0 + pongBaseEvent); + AscendC::SetFlag(EVENT_ID2); // preset h + AscendC::SetFlag(EVENT_ID2 + pongBaseEvent); + } + AscendC::SetFlag(EVENT_ID1); // preset u + AscendC::SetFlag(EVENT_ID1 + pongBaseEvent); + AscendC::SetFlag(EVENT_ID3); // preset g + AscendC::SetFlag(EVENT_ID3 + pongBaseEvent); + AscendC::SetFlag(EVENT_ID0); // preset h_update + AscendC::SetFlag(EVENT_ID0 + pongBaseEvent); + AscendC::SetFlag(EVENT_ID2); // preset h + AscendC::SetFlag(EVENT_ID2 + pongBaseEvent); + uint32_t currStage = 0; // 0: V1, 1: V2 + bool event0FromMte3[PING_PONG_STAGES] = {false, false}; + bool event2FromMte3[PING_PONG_STAGES] = { + !(storeFinalState && std::is_same::value), + !(storeFinalState && std::is_same::value)}; + while (vecBlockScheduler.isRunning) { + if (currStage == 0) { + /* V1: + * gmV = gmU - gmVWorkspace + * g_buf = gmG[-1] - gmG + * g_buf = exp(g_buf) + * gmVWorkspace = g_buf * gmV + */ + vecBlockScheduler.InitTasks(); + for (uint32_t i = 0; i < PING_PONG_STAGES; ++i) { + uint32_t streamId = vecBlockScheduler.GetStreamId(i); + const auto &stream = vecBlockScheduler.GetStream(i); + if (vecBlockScheduler.StreamIsDone(stream)) { + continue; + } + const KdaFwdHOffsets &vec1Offsets = vecBlockScheduler.GetCurTaskOffsets(stream); + AscendC::LocalTensor l1VUpdate = (i == 0) ? l1VUpdatePing : l1VUpdatePong; + bool tailVectorPath = vec1Offsets.blockTokens < 16; + if (tailVectorPath) { + Arch::CrossCoreWaitFlag(vecBlockScheduler.cube1Done[streamId]); + ComputeTailVWorkspace(vec1Offsets, EVENT_ID3 + (i == 0 ? 0 : pongBaseEvent)); + } + bool waitWsFromMte3 = storeFinalState && std::is_same::value && + event0FromMte3[streamId]; + bool useDirectForTask = useDirectFp32Ub && !tailVectorPath; + epilogueKdaFwdHVnew( + gmV[vec1Offsets.uvOffset], gmVUpdateWorkspace[vec1Offsets.vWorkOffset], l1VUpdate, + gmG[vec1Offsets.gOffset], gmU[vec1Offsets.uvOffset], gmVWorkspace[vec1Offsets.vWorkOffset], + gmGk[vec1Offsets.gkOffset], gmK[vec1Offsets.wkOffset], + gmKDecayWorkspace[vec1Offsets.kDecayWorkOffset], vec1Offsets.blockTokens, kHeadDim, + vec1Offsets.vBlockDim, vHeadDim, vecBlockScheduler.cube1Done[streamId], + vecBlockScheduler.vec1Done[streamId], vec1Offsets.isInitialState, vec1Offsets.isFinalState, + storeFinalState, waitWsFromMte3, (i == 0), tailVectorPath, useDirectForTask, + DIRECT_UB_FREE_FLAG_BEGIN, DIRECT_UB_READY_FLAG_BEGIN, DIRECT_UB_FLAG_STRIDE); + if (storeFinalState && std::is_same::value) { + event0FromMte3[streamId] = false; + } + } + } else { + /* V2: h[i+1] += h_work if i < num_chunks - 1 else None */ + for (uint32_t i = 0; i < PING_PONG_STAGES; ++i) { + uint32_t streamId = vecBlockScheduler.GetStreamId(i); + const auto &stream = vecBlockScheduler.GetStream(i); + if (vecBlockScheduler.StreamIsDone(stream)) { + continue; + } + const KdaFwdHOffsets &vec2Offsets = vecBlockScheduler.GetCurTaskOffsets(stream); + if (vecBlockScheduler.NeedProcessStage2(stream)) { + bool tailVectorPath = vec2Offsets.blockTokens < 16; + bool fuseTailFinalState = false; + if constexpr (kGated && !scalarGated && useExp2 && + std::is_same::value) { + fuseTailFinalState = tailVectorPath && vec2Offsets.isFinalState && storeFinalState; + } + if (tailVectorPath) { + Arch::CrossCoreWaitFlag(vecBlockScheduler.cube2Done[streamId]); + if (fuseTailFinalState) { + ComputeKdaTailFinalState(vec2Offsets, EVENT_ID3 + (i == 0 ? 0 : pongBaseEvent)); + } else { + ComputeTailHWorkspace(vec2Offsets, EVENT_ID3 + (i == 0 ? 0 : pongBaseEvent)); + } + } + if (!fuseTailFinalState) { + if (storeFinalState && std::is_same::value) { + if (vec2Offsets.isInitialState && event2FromMte3[streamId]) { + uint32_t event2Id = EVENT_ID2 + (i == 0 ? 0 : pongBaseEvent); + AscendC::WaitFlag(event2Id); + AscendC::SetFlag(event2Id); + event2FromMte3[streamId] = false; + } + // Update always writes the FP32 state through MTE3, + // including intermediate chunks. + event0FromMte3[streamId] = true; + event2FromMte3[streamId] = !vec2Offsets.isFinalState; + } + epilogueKdaFwdHUpdate( + gmH[vec2Offsets.hDstOffset], gmFinalState[vec2Offsets.finalStateOffset], + gmG[vec2Offsets.gOffset], gmH[vec2Offsets.hSrcOffset], + gmHWorkspace[vec2Offsets.hWorkOffset], gmGk[vec2Offsets.gkOffset], + gmInitialState[vec2Offsets.initialStateOffset], vec2Offsets.blockTokens, kHeadDim, + vec2Offsets.vBlockDim, vHeadDim, vecBlockScheduler.cube2Done[streamId], + vec2Offsets.isInitialState, vec2Offsets.isFinalState, storeFinalState, + useInitialState, (i == 0), tailVectorPath, useDirectFp32Ub && !tailVectorPath, + DIRECT_UB_FREE_FLAG_BEGIN, DIRECT_UB_READY_FLAG_BEGIN, DIRECT_UB_FLAG_STRIDE); + } + } else { + if (!useDirectFp32Ub) { + Arch::CrossCoreWaitFlag(vecBlockScheduler.cube2Done[streamId]); + } + } + Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(vecBlockScheduler.vec2Done[streamId]); + } + } + currStage ^= 0x01; + } + if (storeFinalState && std::is_same::value) { + if (event0FromMte3[0]) { + AscendC::WaitFlag(EVENT_ID0); + } else { + AscendC::WaitFlag(EVENT_ID0); + } + if (event0FromMte3[1]) { + AscendC::WaitFlag(EVENT_ID0 + pongBaseEvent); + } else { + AscendC::WaitFlag(EVENT_ID0 + pongBaseEvent); + } + if (event2FromMte3[0]) { + AscendC::WaitFlag(EVENT_ID2); + } else { + AscendC::WaitFlag(EVENT_ID2); + } + if (event2FromMte3[1]) { + AscendC::WaitFlag(EVENT_ID2 + pongBaseEvent); + } else { + AscendC::WaitFlag(EVENT_ID2 + pongBaseEvent); + } + } else { + AscendC::WaitFlag(EVENT_ID0); // preset h_update + AscendC::WaitFlag(EVENT_ID0 + pongBaseEvent); + AscendC::WaitFlag(EVENT_ID2); // preset h + AscendC::WaitFlag(EVENT_ID2 + pongBaseEvent); + } + AscendC::WaitFlag(EVENT_ID1); // preset u + AscendC::WaitFlag(EVENT_ID1 + pongBaseEvent); + AscendC::WaitFlag(EVENT_ID3); // preset g + AscendC::WaitFlag(EVENT_ID3 + pongBaseEvent); + AscendC::WaitFlag(EVENT_ID0); // drain h_update + AscendC::WaitFlag(EVENT_ID0 + pongBaseEvent); + AscendC::WaitFlag(EVENT_ID2); // drain h + AscendC::WaitFlag(EVENT_ID2 + pongBaseEvent); + } + } +}; + +} // namespace Catlass::Gemm::Kernel diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/arch/kda_arch.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/arch/kda_arch.hpp new file mode 100644 index 0000000000..36e70c0969 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/arch/kda_arch.hpp @@ -0,0 +1,52 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_ARCH_HPP +#define CATLASS_KDA_ARCH_HPP + +#include "catlass/kda_catlass.hpp" + +namespace Catlass::Arch { + +struct AtlasA2 { + static constexpr uint32_t BIAS_SIZE = 1024; + static constexpr uint32_t FIXBUF_SIZE = 7 * 1024; + static constexpr uint32_t UB_SIZE = 192 * 1024; + static constexpr uint32_t L1_SIZE = 512 * 1024; + static constexpr uint32_t L0A_SIZE = 64 * 1024; + static constexpr uint32_t L0B_SIZE = 64 * 1024; + static constexpr uint32_t L0C_SIZE = 128 * 1024; +}; + +// Ascend950 +struct Ascend950 { + static constexpr uint32_t BIAS_SIZE = 4 * 1024; + static constexpr uint32_t FIXBUF_SIZE = 16 * 1024; + static constexpr uint32_t UB_SIZE = 248 * 1024; + static constexpr uint32_t L1_SIZE = 512 * 1024; + static constexpr uint32_t L0A_SIZE = 64 * 1024; + static constexpr uint32_t L0B_SIZE = 64 * 1024; + static constexpr uint32_t L0C_SIZE = 256 * 1024; +}; + +template +using PositionType = std::integral_constant; + +using PositionGM = PositionType; +using PositionL1 = PositionType; +using PositionL0A = PositionType; +using PositionL0B = PositionType; +using PositionL0C = PositionType; +using PositionBias = PositionType; +using PositionUB = PositionType; + +} // namespace Catlass::Arch + +#endif // CATLASS_KDA_ARCH_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/arch/kda_cross_core_sync.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/arch/kda_cross_core_sync.hpp new file mode 100644 index 0000000000..8b074cec46 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/arch/kda_cross_core_sync.hpp @@ -0,0 +1,126 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_ARCH_KDA_CROSS_CORE_SYNC_HPP +#define CATLASS_ARCH_KDA_CROSS_CORE_SYNC_HPP + +#include "catlass/kda_catlass.hpp" + +namespace Catlass::Arch { + +// A certain cross core flag can be continuously set up to 15 times without waiting. In scenarios where there is only +// one-way synchronization between cores and the synchronization operation is executed multiple times, if the core +// performing the set operation runs faster, it may result in the flag being set more than 15 times consecutively, +// leading to a system freeze. To prevent such a freeze, we need to wait for the waiting cores to perform a reverse +// synchronization operation after executing the set operation MAX_REVERSE_DEPTH times. +constexpr uint32_t MAX_REVERSE_DEPTH = 15; + +using FlagID = uint16_t; +constexpr FlagID AIV_INTER_BLOCK_BARRIER = 8; +constexpr FlagID AIC_INTER_BLOCK_BARRIER = 9; +constexpr FlagID AIV_INTER_SUBBLOCK_BARRIER = 10; +constexpr FlagID FFTS_MAX_FLAG = 7; + +struct CrossCoreFlag { + CATLASS_DEVICE + CrossCoreFlag() + : id(0) + {} + + CATLASS_DEVICE + CrossCoreFlag(FlagID id) + : id(id) + {} + + FlagID id; +}; + +template +struct CrossCoreFlagWithReverse { + CATLASS_DEVICE + CrossCoreFlagWithReverse() + : id(0), + reverseId(0) + {} + + CATLASS_DEVICE + CrossCoreFlagWithReverse(FlagID id, FlagID reverseId) + : id(id), + reverseId(reverseId) + {} + + FlagID id; + FlagID reverseId; + uint32_t count{0}; +}; + +template +struct BarrierFlag { + static_assert(MODE != 255 && MODE == 255, "Unsupported cross core barrier flag, can not find the specialization."); +}; + +template <> +struct BarrierFlag<0x0, AscendC::AIV> { + static constexpr FlagID ID = AIV_INTER_BLOCK_BARRIER; +}; + +template <> +struct BarrierFlag<0x0, AscendC::AIC> { + static constexpr FlagID ID = AIC_INTER_BLOCK_BARRIER; +}; + +template <> +struct BarrierFlag<0x1, AscendC::AIV> { + static constexpr FlagID ID = AIV_INTER_SUBBLOCK_BARRIER; +}; + +template +CATLASS_DEVICE void CrossCoreBarrier() +{ + constexpr FlagID flagId = BarrierFlag::ID; + AscendC::CrossCoreSetFlag(flagId); + AscendC::CrossCoreWaitFlag(flagId); +} + +template +CATLASS_DEVICE void CrossCoreSetFlag(CrossCoreFlag &flag) +{ + AscendC::CrossCoreSetFlag(flag.id); +} + +CATLASS_DEVICE +void CrossCoreWaitFlag(CrossCoreFlag &flag) +{ + AscendC::CrossCoreWaitFlag(flag.id); +} + +template +CATLASS_DEVICE void CrossCoreSetFlagWithReverse(CrossCoreFlagWithReverse &flag) +{ + AscendC::CrossCoreSetFlag(flag.id); + if (++flag.count >= REVERSE_DEPTH) { + AscendC::CrossCoreWaitFlag(flag.reverseId); + flag.count = 0; + } +} + +template +CATLASS_DEVICE void CrossCoreWaitFlagWithReverse(CrossCoreFlagWithReverse &flag) +{ + AscendC::CrossCoreWaitFlag(flag.id); + if (++flag.count >= REVERSE_DEPTH) { + AscendC::CrossCoreSetFlag(flag.reverseId); + flag.count = 0; + } +} + +} // namespace Catlass::Arch + +#endif // CATLASS_ARCH_KDA_CROSS_CORE_SYNC_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/arch/kda_local_tensor_buffer.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/arch/kda_local_tensor_buffer.hpp new file mode 100644 index 0000000000..8e02f4974f --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/arch/kda_local_tensor_buffer.hpp @@ -0,0 +1,230 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_LOCAL_TENSOR_BUFFER_HPP +#define CATLASS_KDA_LOCAL_TENSOR_BUFFER_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_arch.hpp" + +namespace Catlass::Arch { + +struct LocalTensorBufferBase { +public: + template + CATLASS_DEVICE + AscendC::LocalTensor + GetBufferByByte(const uint32_t offset) const + { + return tensor[offset].template ReinterpretCast(); + } + +protected: + CATLASS_DEVICE + LocalTensorBufferBase() = default; + + AscendC::LocalTensor tensor; +}; + +template < + class ArchTag, + AscendC::TPosition Position> +struct LocalTensorBuffer { + static_assert(DEPENDENT_FALSE, "Unsupported local tensor buffer, can not find the specialization."); +}; + +/// Partial specialization for TPosition::A1 +template +struct LocalTensorBuffer : LocalTensorBufferBase { +public: + static constexpr AscendC::TPosition Position = AscendC::TPosition::A1; + + CATLASS_DEVICE + LocalTensorBuffer() + { + AscendC::TBuf tbufA1; + GetTPipePtr()->InitBuffer(tbufA1, ArchTag::L1_SIZE); + tensor = tbufA1.Get(); + } +}; + +/////////////////////////////////////////////////////////// + +/// Partial specialization for TPosition::A2 +template +struct LocalTensorBuffer : LocalTensorBufferBase { +public: + static constexpr AscendC::TPosition Position = AscendC::TPosition::A2; + + CATLASS_DEVICE + LocalTensorBuffer() + { + AscendC::TBuf tbufA2; + GetTPipePtr()->InitBuffer(tbufA2, ArchTag::L0A_SIZE); + tensor = tbufA2.Get(); + } +}; + +/////////////////////////////////////////////////////////// + +/// Partial specialization for TPosition::B1 +template +struct LocalTensorBuffer : LocalTensorBufferBase { +public: + static constexpr AscendC::TPosition Position = AscendC::TPosition::B1; + + CATLASS_DEVICE + LocalTensorBuffer() + { + AscendC::TBuf tbufB1; + GetTPipePtr()->InitBuffer(tbufB1, ArchTag::L1_SIZE); + tensor = tbufB1.Get(); + } +}; + +/////////////////////////////////////////////////////////// + +/// Partial specialization for TPosition::B2 +template +struct LocalTensorBuffer : LocalTensorBufferBase { +public: + static constexpr AscendC::TPosition Position = AscendC::TPosition::B2; + + CATLASS_DEVICE + LocalTensorBuffer() + { + AscendC::TBuf tbufB2; + GetTPipePtr()->InitBuffer(tbufB2, ArchTag::L0B_SIZE); + tensor = tbufB2.Get(); + } +}; + +/////////////////////////////////////////////////////////// + +/// Partial specialization for TPosition::C1 +template +struct LocalTensorBuffer : LocalTensorBufferBase { +public: + static constexpr AscendC::TPosition Position = AscendC::TPosition::C1; + + CATLASS_DEVICE + LocalTensorBuffer() + { + AscendC::TBuf tbufC1; + GetTPipePtr()->InitBuffer(tbufC1, ArchTag::L1_SIZE); + tensor = tbufC1.Get(); + } +}; + +/////////////////////////////////////////////////////////// + +/// Partial specialization for TPosition::C2 +template +struct LocalTensorBuffer : LocalTensorBufferBase { +public: + static constexpr AscendC::TPosition Position = AscendC::TPosition::C2; + + CATLASS_DEVICE + LocalTensorBuffer() + { + AscendC::TBuf tbufC2; + GetTPipePtr()->InitBuffer(tbufC2, ArchTag::BIAS_SIZE); + tensor = tbufC2.Get(); + } +}; + +/////////////////////////////////////////////////////////// + +/// Partial specialization for TPosition::CO1 +template +struct LocalTensorBuffer : LocalTensorBufferBase { +public: + static constexpr AscendC::TPosition Position = AscendC::TPosition::CO1; + + CATLASS_DEVICE + LocalTensorBuffer() + { + AscendC::TBuf tbufCO1; + GetTPipePtr()->InitBuffer(tbufCO1, ArchTag::L0C_SIZE); + tensor = tbufCO1.Get(); + } +}; + +/////////////////////////////////////////////////////////// + +/// Partial specialization for TPosition::C2PIPE2GM +template +struct LocalTensorBuffer : LocalTensorBufferBase { +public: + static constexpr AscendC::TPosition Position = AscendC::TPosition::C2PIPE2GM; + + CATLASS_DEVICE + LocalTensorBuffer() + { + AscendC::TBuf tbufC2PIPE2GM; + GetTPipePtr()->InitBuffer(tbufC2PIPE2GM, ArchTag::FIXBUF_SIZE); + tensor = tbufC2PIPE2GM.Get(); + } +}; + +/////////////////////////////////////////////////////////// + +/// Partial specialization for TPosition::VECIN +template +struct LocalTensorBuffer : LocalTensorBufferBase { +public: + static constexpr AscendC::TPosition Position = AscendC::TPosition::VECIN; + + CATLASS_DEVICE + LocalTensorBuffer() + { + AscendC::TBuf tbufVECIN; + GetTPipePtr()->InitBuffer(tbufVECIN, ArchTag::UB_SIZE); + tensor = tbufVECIN.Get(); + } +}; + +/////////////////////////////////////////////////////////// + +/// Partial specialization for TPosition::VECOUT +template +struct LocalTensorBuffer : LocalTensorBufferBase { +public: + static constexpr AscendC::TPosition Position = AscendC::TPosition::VECOUT; + + CATLASS_DEVICE + LocalTensorBuffer() + { + AscendC::TBuf tbufVECOUT; + GetTPipePtr()->InitBuffer(tbufVECOUT, ArchTag::UB_SIZE); + tensor = tbufVECOUT.Get(); + } +}; + +/////////////////////////////////////////////////////////// + +/// Partial specialization for TPosition::VECCALC +template +struct LocalTensorBuffer : LocalTensorBufferBase { +public: + static constexpr AscendC::TPosition Position = AscendC::TPosition::VECCALC; + + CATLASS_DEVICE + LocalTensorBuffer() + { + AscendC::TBuf tbufVECCALC; + GetTPipePtr()->InitBuffer(tbufVECCALC, ArchTag::UB_SIZE); + tensor = tbufVECCALC.Get(); + } +}; + +} // namespace Catlass::Arch + +#endif // CATLASS_KDA_LOCAL_TENSOR_BUFFER_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/arch/kda_resource.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/arch/kda_resource.hpp new file mode 100644 index 0000000000..9a521c24ce --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/arch/kda_resource.hpp @@ -0,0 +1,43 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_RESOURCE_HPP +#define CATLASS_KDA_RESOURCE_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_local_tensor_buffer.hpp" + +namespace Catlass::Arch { + +template +struct Resource { +public: + AscendC::TPipe pipe; + + LocalTensorBuffer l1Buf; + LocalTensorBuffer l0ABuf; + LocalTensorBuffer l0BBuf; + LocalTensorBuffer btBuf; + LocalTensorBuffer l0CBuf; + LocalTensorBuffer ubBuf; + LocalTensorBuffer fpBuf; + + CATLASS_DEVICE + Resource() + { + // The initialization of AscendC::Tpipe will insert some synchronization interfaces, + // which may conflict with the usage by users. Therefore, the "destroy" interface is used for releasing. + pipe.Destroy(); + } +}; + +} // namespace Catlass::Arch + +#endif // CATLASS_KDA_RESOURCE_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/detail/kda_alignment.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/detail/kda_alignment.hpp new file mode 100644 index 0000000000..f6821a9a40 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/detail/kda_alignment.hpp @@ -0,0 +1,103 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_ALIGNMENT_HPP +#define CATLASS_KDA_ALIGNMENT_HPP + +#include "catlass/detail/kda_macros.hpp" +#include "tla/numeric/integral_constant.hpp" + +template +CATLASS_HOST_DEVICE constexpr T RoundUp(const T &val) +{ + static_assert(ALIGN != 0, "ALIGN must not be 0"); + return (val + ALIGN - 1) / ALIGN * ALIGN; +} + +template +CATLASS_HOST_DEVICE constexpr auto RoundUp(T const &val, U const &align) +{ + if constexpr (tla::is_static::value && tla::is_static::value) { // Int, Int + constexpr uint32_t res = (T::value + U::value - 1) / U::value * U::value; + return tla::Int{}; + } else if constexpr (tla::is_static::value) { // Int, int + return (T::value + align - 1) / align * align; + } else if constexpr (tla::is_static::value) { // int, Int + return (val + U::value - 1) / U::value * U::value; + } else { // int, int + return (val + align - 1) / align * align; + } +} + +template +CATLASS_HOST_DEVICE constexpr T RoundDown(const T val) +{ + static_assert(ALIGN != 0, "ALIGN must not be 0"); + return val / ALIGN * ALIGN; +} + +template +CATLASS_HOST_DEVICE constexpr auto RoundDown(T const &val, U const &align) +{ + if constexpr (tla::is_static::value && tla::is_static::value) { // Int, Int + constexpr uint32_t res = T::value / U::value * U::value; + return tla::Int{}; + } else if constexpr (tla::is_static::value) { // Int, int + return T::value / align * align; + } else if constexpr (tla::is_static::value) { // int, Int + return val / U::value * U::value; + } else { // int, int + return val / align * align; + } +} + +template +CATLASS_HOST_DEVICE constexpr T CeilDiv(const T dividend) +{ + static_assert(DIVISOR != 0, "DIVISOR must not be 0"); + return (dividend + DIVISOR - 1) / DIVISOR; +} + +template +CATLASS_HOST_DEVICE constexpr auto CeilDiv(T const ÷nd, U const &divisor) +{ + if constexpr (tla::is_static::value && tla::is_static::value) { // Int, Int + constexpr uint32_t res = (T::value + U::value - 1) / U::value; + return tla::Int{}; + } else if constexpr (tla::is_static::value) { // Int, int + return (T::value + divisor - 1) / divisor; + } else if constexpr (tla::is_static::value) { // int, Int + return (dividend + U::value - 1) / U::value; + } else { // int, int + return (dividend + divisor - 1) / divisor; + } +} + +template +CATLASS_HOST_DEVICE constexpr auto Max(T const &a, U const &b) +{ + if (a > b) { + return a; + } else { + return b; + } +} + +template +CATLASS_HOST_DEVICE constexpr auto Min(T const &a, U const &b) +{ + if (a < b) { + return a; + } else { + return b; + } +} + +#endif // CATLASS_KDA_ALIGNMENT_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/detail/kda_callback.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/detail/kda_callback.hpp new file mode 100644 index 0000000000..91456611c9 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/detail/kda_callback.hpp @@ -0,0 +1,61 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_DETAIL_KDA_CALLBACK_HPP +#define CATLASS_DETAIL_KDA_CALLBACK_HPP + +#include "catlass/detail/kda_macros.hpp" + +/// @brief Callback is an alternative to std::function, providing a general carrier +/// of callable structure with no parameters and no return value. Compared with function pointers +/// of type void (*)(), Callback can carry lambda expressions with captures, and does not need to +/// pay attention to the captured content. It should be noted that Callback itself does not store +/// the callable structure it carries like std::function, so it is necessary to ensure +/// that it is used within the life cycle of the callable structure. +struct Callback { + void const *func{nullptr}; + void (*caller)(void const *){nullptr}; + + Callback() = default; + + CATLASS_DEVICE + void operator()() const + { + if (func) { + caller(func); + } + } + + CATLASS_DEVICE + operator bool() const + { + return func != nullptr; + } +}; + +template +CATLASS_DEVICE void FuncWrapper(void const *func) +{ + (*static_cast(func))(); +} + +// Use this to make a callback +template +CATLASS_DEVICE + Callback + MakeCallback(Func *func) +{ + Callback callback; + callback.func = func; + callback.caller = &FuncWrapper; + return callback; +} + +#endif // CATLASS_DETAIL_KDA_CALLBACK_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/detail/kda_dependent_false.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/detail/kda_dependent_false.hpp new file mode 100644 index 0000000000..80613e908d --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/detail/kda_dependent_false.hpp @@ -0,0 +1,20 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_DEPENDENT_FALSE_HPP +#define CATLASS_KDA_DEPENDENT_FALSE_HPP + +template +constexpr bool DEPENDENT_BOOL_VALUE = VALUE; + +template +constexpr bool DEPENDENT_FALSE = DEPENDENT_BOOL_VALUE; + +#endif // CATLASS_KDA_DEPENDENT_FALSE_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/detail/kda_macros.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/detail/kda_macros.hpp new file mode 100644 index 0000000000..7d47093730 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/detail/kda_macros.hpp @@ -0,0 +1,26 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_MACROS_HPP +#define CATLASS_KDA_MACROS_HPP + +#if defined(__CCE__) +#include +#endif + +#define CATLASS_DEVICE __forceinline__ __aicore__ +#ifdef __CCE__ +#define CATLASS_HOST_DEVICE __forceinline__[host, aicore] +#else +#define CATLASS_HOST_DEVICE +#endif +#define CATLASS_GLOBAL __global__ __aicore__ + +#endif // CATLASS_KDA_MACROS_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/detail/kda_tag_to_layout.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/detail/kda_tag_to_layout.hpp new file mode 100644 index 0000000000..119cdba6fd --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/detail/kda_tag_to_layout.hpp @@ -0,0 +1,86 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_TAG_TO_LAYOUT_HPP +#define CATLASS_KDA_TAG_TO_LAYOUT_HPP + +#include "catlass/layout/kda_layout.hpp" +#include "tla/kda_tla_layout.hpp" + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +namespace Catlass::detail { +//////////////////////////////////////////////////////////////////////////////////////////////////// +// For each Catlass::layout, provides its corresponding tla layout types +template +struct TagToLayout { + using type = LayoutTag; +}; + +template +struct TagToLayout { + using type = tla::Layout, tla::Stride>>; +}; + +template +struct TagToLayout { + using type = tla::Layout, tla::Stride, int64_t>>; +}; + +template +struct TagToLayout { + using type = tla::Layout, tla::Stride>>; +}; + +template +struct TagToLayout { + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(Element); + using type = tla::Layout< + tla::Shape, uint32_t>, tla::Shape, uint32_t>>, + tla::Stride, tla::Int>, + tla::Stride, int64_t>>>; +}; + +template +struct TagToLayout { + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(Element); + using type = tla::Layout< + tla::Shape, uint32_t>, tla::Shape, uint32_t>>, + tla::Stride, int64_t>, + tla::Stride, tla::Int>>>; +}; + +template +struct TagToLayout { + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(Element); + using type = tla::Layout< + tla::Shape, uint32_t>, tla::Shape, uint32_t>>, + tla::Stride, int64_t>, + tla::Stride, tla::Int>>>; +}; + +// Convenience aliases +template +using TagToLayout_t = typename TagToLayout::type; + +constexpr uint32_t ELE_NUM_PER_FRACTAL_L0C = 256; +using LayoutL0C = tla::Layout< + tla::Shape, uint32_t>, tla::Shape, uint32_t>>, + tla::Stride, tla::Int>, + tla::Stride, int64_t>>>; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace Catlass::detail + +#endif // CATLASS_KDA_TAG_TO_LAYOUT_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_elemwise_no_source.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_elemwise_no_source.hpp new file mode 100644 index 0000000000..06fe13078a --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_elemwise_no_source.hpp @@ -0,0 +1,185 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_EPILOGUE_BLOCK_EPILOGUE_ELEMWISE_NO_SOURCE_HPP +#define CATLASS_EPILOGUE_BLOCK_EPILOGUE_ELEMWISE_NO_SOURCE_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/epilogue/kda_dispatch_policy.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/kda_matrix_coord.hpp" +#include "catlass/layout/kda_layout.hpp" +#include "catlass/epilogue/tile/kda_tile_cast.hpp" + +namespace Catlass::Epilogue::Block { +template < + class CType_, + class DType_, + class TileElemWiseEpilogue_, + class TileCopy_> +class BlockEpilogue< + EpilogueAtlasA2ElemWiseNoSource, + CType_, + DType_, + TileElemWiseEpilogue_, + TileCopy_> { +public: + // Type aliases + using DispatchPolicy = EpilogueAtlasA2ElemWiseNoSource; + using ArchTag = typename DispatchPolicy::ArchTag; + using ElementC = typename CType_::Element; + using LayoutC = typename CType_::Layout; + + using ElementD = typename DType_::Element; + using LayoutD = typename DType_::Layout; + using TileElemWiseEpilogue = TileElemWiseEpilogue_; + using CopyGmToUbC = typename TileCopy_::CopyGmToUbC; + using CopyUbToGmD = typename TileCopy_::CopyUbToGmD; + + static constexpr uint32_t COMPUTE_LENGTH = TileElemWiseEpilogue::COMPUTE_LENGTH; + static constexpr uint32_t OPERANDS_NUM = DispatchPolicy::OPERANDS_NUM; + + using ElementCompute = ElementC; + + using LayoutComputeInUb = layout::RowMajor; + + // Check the element type of C and D + static_assert(std::is_same_v, + "Element type of C must be float"); + // Check the layout type of C and D + static_assert(std::is_same_v && std::is_same_v, + "Layout type of C, D must be RowMajor"); + + // Check if ArchTag is matched + static_assert(std::is_same_v, + "Tile epilogue's ArchTag mismatch"); + // Check if compute length is valid + static_assert(COMPUTE_LENGTH * (OPERANDS_NUM * sizeof(ElementC) + sizeof(ElementD)) <= ArchTag::UB_SIZE, + "UB out of bounds"); + + // Epilogue params definition + struct Params { + GM_ADDR ptrC; + LayoutC layoutC; + GM_ADDR ptrD; + LayoutD layoutD; + + CATLASS_HOST_DEVICE + Params() {} + + CATLASS_HOST_DEVICE + Params(GM_ADDR ptrC_, LayoutD const &layoutC_, GM_ADDR ptrD_, LayoutD const &layoutD_) + : ptrC(ptrC_), + layoutC(layoutC_), + ptrD(ptrD_), + layoutD(layoutD_) + {} + }; + + CATLASS_DEVICE + BlockEpilogue(Arch::Resource &resource, Params const ¶ms) + : params(params) + { + ubC = resource.ubBuf.template GetBufferByByte(0); + ubD = resource.ubBuf.template GetBufferByByte(COMPUTE_LENGTH * sizeof(ElementC)); + if constexpr (!std::is_same_v) { + ubCompute = resource.ubBuf.template GetBufferByByte( + COMPUTE_LENGTH * (sizeof(ElementC) + sizeof(ElementD))); + } else { + // When ElementCompute is same as ElementD, it is not necessary to do Cast from ubCompute to ubD. + // Use the same buffer. + ubCompute = ubD; + } + AscendC::SetFlag(EVENT_ID0); + AscendC::SetFlag(EVENT_ID0); + } + + CATLASS_DEVICE + ~BlockEpilogue() + { + AscendC::WaitFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + } + + CATLASS_DEVICE + void operator()( + GemmCoord const &blockShapeMNK, + GemmCoord const &blockCoordMNK, + GemmCoord const &actualBlockShapeMNK, + AscendC::GlobalTensor const &gmBlockC, + LayoutD const &layoutBlockC) + { + // Calculate the offset of the current block + MatrixCoord blockShape = blockShapeMNK.GetCoordMN(); + MatrixCoord blockCoord = blockCoordMNK.GetCoordMN(); + MatrixCoord actualBlockShape = actualBlockShapeMNK.GetCoordMN(); + MatrixCoord blockOffset = blockCoord * blockShape; + + // Calculate the offset and the shape of the current subblock + MatrixCoord subblockShape{ + CeilDiv(actualBlockShape.row(), static_cast(AscendC::GetSubBlockNum())), + actualBlockShape.column()}; + + MatrixCoord subblockCoord{AscendC::GetSubBlockIdx(), 0}; + MatrixCoord actualSubblockShape = MatrixCoord::Min( + subblockShape, actualBlockShape - subblockCoord * subblockShape); + MatrixCoord subblockOffset = subblockCoord * subblockShape; + + // Get the data and layout of C + auto gmSubblockC = gmBlockC[layoutBlockC.GetOffset(subblockOffset)]; + auto layoutSubblockC = layoutBlockC.GetTileLayout(actualSubblockShape); + + // Get the data and layout of D + AscendC::GlobalTensor gmD; + gmD.SetGlobalBuffer(reinterpret_cast<__gm__ ElementD *>(params.ptrD)); + auto gmSubblockD = gmD[params.layoutD.GetOffset(blockOffset + subblockOffset)]; + auto layoutSubblockD = params.layoutD.GetTileLayout(actualSubblockShape); + + // Get the layout on UB + auto ubTileStride = MakeCoord(static_cast(blockShape.column()), 1L); + LayoutComputeInUb layoutComputeInUb{actualSubblockShape, ubTileStride}; + LayoutComputeInUb layoutComputeOutUb{actualSubblockShape, ubTileStride}; + // Copy the data of C + AscendC::WaitFlag(EVENT_ID0); + copyGmToUbC(ubC, gmSubblockC, layoutComputeInUb, layoutSubblockC); + AscendC::SetFlag(EVENT_ID0); + + // Perform epilogue calculation + AscendC::WaitFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + tileEpilogue(ubCompute, ubC); + if constexpr (!std::is_same_v) { + AscendC::Cast(ubD, ubCompute, AscendC::RoundMode::CAST_NONE, COMPUTE_LENGTH); + } + AscendC::SetFlag(EVENT_ID0); + AscendC::SetFlag(EVENT_ID0); + + // Copy the data of D + AscendC::WaitFlag(EVENT_ID0); + copyUbToGmD(gmSubblockD, ubD, layoutSubblockD, layoutComputeOutUb); + AscendC::SetFlag(EVENT_ID0); + } + +private: + Params params; + + AscendC::LocalTensor ubC; + AscendC::LocalTensor ubCompute; + AscendC::LocalTensor ubD; + + TileElemWiseEpilogue tileEpilogue; + CopyGmToUbC copyGmToUbC; + CopyUbToGmD copyUbToGmD; +}; + +} // namespace Catlass::Epilogue::Block + +#endif // CATLASS_EPILOGUE_BLOCK_EPILOGUE_ELEMWISE_NO_SOURCE_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_fa_rescale_o.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_fa_rescale_o.hpp new file mode 100644 index 0000000000..5f7aba7c2a --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_fa_rescale_o.hpp @@ -0,0 +1,337 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_FA_RESCALE_O_HPP +#define CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_FA_RESCALE_O_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/epilogue/kda_dispatch_policy.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/kda_matrix_coord.hpp" +#include "catlass/epilogue/tile/kda_epilogue_tile_copy.hpp" + +namespace Catlass::Epilogue::Block { + +template < + class OutputType_, + class InputType_> +class BlockEpilogue< + EpilogueAtlasA2FARescaleO, + OutputType_, + InputType_> { +public: + // Type aliases + using DispatchPolicy = EpilogueAtlasA2FARescaleO; + using ArchTag = typename DispatchPolicy::ArchTag; + using ElementOutput = typename OutputType_::Element; + using ElementInput = typename InputType_::Element; + + using LayoutOutput = typename OutputType_::Layout; + using LayoutInput = typename InputType_::Layout; + + using CopyGmToUbInput = Tile::CopyGm2Ub; + using CopyUbToGmOutput = Tile::CopyUb2Gm; + + static constexpr uint32_t HALF_ELENUM_PER_BLK = 16; + static constexpr uint32_t FLOAT_ELENUM_PER_BLK = 8; + static constexpr uint32_t HALF_ELENUM_PER_VECCALC = 128; + static constexpr uint32_t FLOAT_ELENUM_PER_VECCALC = 64; + static constexpr uint32_t UB_TILE_SIZE = 16384; // 64 * 128 * 2B + static constexpr uint32_t UB_LINE_SIZE = 512; // 128 * 2 * 2B + static constexpr uint32_t HALF_ELENUM_PER_LINE = 256; // 128 * 2 + static constexpr uint32_t FLOAT_ELENUM_PER_LINE = 128; // 128 + static constexpr uint32_t MULTIPLIER = 2; + + CATLASS_DEVICE + BlockEpilogue(Arch::Resource &resource) + { + constexpr uint32_t LO_UB_TENSOR_OFFSET = 5 * UB_TILE_SIZE; + constexpr uint32_t DM_UB_TENSOR_OFFSET = 7 * UB_TILE_SIZE + 4 * UB_LINE_SIZE; + constexpr uint32_t LL_UB_TENSOR_OFFSET = 7 * UB_TILE_SIZE + 6 * UB_LINE_SIZE; + constexpr uint32_t GL_UB_TENSOR_OFFSET = 7 * UB_TILE_SIZE + 10 * UB_LINE_SIZE; + constexpr uint32_t TV_UB_TENSOR_OFFSET = 7 * UB_TILE_SIZE + 11 * UB_LINE_SIZE; + constexpr uint32_t GO_UB_TENSOR_OFFSET = 8 * UB_TILE_SIZE; + + loUbTensor = resource.ubBuf.template GetBufferByByte(LO_UB_TENSOR_OFFSET); + dmUbTensor = resource.ubBuf.template GetBufferByByte(DM_UB_TENSOR_OFFSET); + llUbTensor = resource.ubBuf.template GetBufferByByte(LL_UB_TENSOR_OFFSET); + glUbTensor = resource.ubBuf.template GetBufferByByte(GL_UB_TENSOR_OFFSET); + tvUbTensor = resource.ubBuf.template GetBufferByByte(TV_UB_TENSOR_OFFSET); + goUbTensor = resource.ubBuf.template GetBufferByByte(GO_UB_TENSOR_OFFSET); + + AscendC::SetFlag(EVENT_ID1); + AscendC::SetFlag(EVENT_ID0); + } + + CATLASS_DEVICE + ~BlockEpilogue() + { + AscendC::WaitFlag(EVENT_ID1); + AscendC::WaitFlag(EVENT_ID0); + } + + CATLASS_DEVICE + void SetMask(int32_t len) + { + const int32_t MAX_MASK_LEN = 128; + const int32_t HALF_MASK_LEN = 64; + if (len >= MAX_MASK_LEN) { + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + return; + } + int32_t highMask = len - HALF_MASK_LEN > 0 ? len - HALF_MASK_LEN : 0; + int32_t lowMask = len - HALF_MASK_LEN >= 0 ? HALF_MASK_LEN : len; + if (len < HALF_MASK_LEN) { + AscendC::SetVectorMask(0x0, ((uint64_t)1 << lowMask) - 1); + } else { + AscendC::SetVectorMask(((uint64_t)1 << highMask) - 1, 0xffffffffffffffff); + } + } + + CATLASS_DEVICE + void subCoreCompute( + AscendC::GlobalTensor gOutput, + AscendC::GlobalTensor gInput, + const LayoutOutput &layoutOutput, + const LayoutInput &layoutInput, + uint32_t nIdx, uint32_t isLast) + { + uint32_t subM = layoutInput.shape(0); + uint32_t k = layoutInput.shape(1); + uint32_t kRound = layoutInput.stride(0); + uint32_t strideQO = layoutOutput.stride(0); + uint32_t subMAligned128 = (subM + HALF_ELENUM_PER_VECCALC - 1) / HALF_ELENUM_PER_VECCALC; + uint32_t subMAligned64 = (subM + FLOAT_ELENUM_PER_VECCALC - 1) / FLOAT_ELENUM_PER_VECCALC; + uint32_t subMRound = (subM + HALF_ELENUM_PER_BLK - 1) / HALF_ELENUM_PER_BLK * HALF_ELENUM_PER_BLK; + + // Get the layout on UB + LayoutInput layoutInUb(subM, k, kRound); + + if (subM > 0) { + AscendC::WaitFlag(EVENT_ID1); + // Copy O + copyGmToUbInput(loUbTensor, gInput, layoutInUb, layoutInput); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + // 更新 L 和 O + if (nIdx != 0) { + // dm32 = castfp16to32(dm) + AscendC::Cast( + tvUbTensor, + dmUbTensor[nIdx % MULTIPLIER * HALF_ELENUM_PER_LINE], + AscendC::RoundMode::CAST_NONE, + (uint64_t)0, + subMAligned64, + AscendC::UnaryRepeatParams(1, 1, 8, 4)); + AscendC::PipeBarrier(); + // dm32_block = brcb(dm32) + AscendC::Brcb( + tvUbTensor.ReinterpretCast()[HALF_ELENUM_PER_VECCALC], + tvUbTensor.ReinterpretCast(), + subMRound / FLOAT_ELENUM_PER_BLK, + AscendC::BrcbRepeatParams(1, 8)); + AscendC::PipeBarrier(); + // dm32 = exp(dm32) + AscendC::Exp( + tvUbTensor, + tvUbTensor, + (uint64_t)0, + subMAligned64, + AscendC::UnaryRepeatParams(1, 1, 8, 8)); + AscendC::PipeBarrier(); + // gl = dm * gl + AscendC::Mul( + glUbTensor, + tvUbTensor, + glUbTensor, + (uint64_t)0, + subMAligned64, + AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); + AscendC::PipeBarrier(); + // gl = ll + gl + AscendC::Add( + glUbTensor, + glUbTensor, + llUbTensor[nIdx % MULTIPLIER * FLOAT_ELENUM_PER_LINE], + (uint64_t)0, + subMAligned64, + AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); + AscendC::PipeBarrier(); + // dm32_block = exp(dm32_block) + AscendC::Exp( + tvUbTensor[HALF_ELENUM_PER_VECCALC], + tvUbTensor[HALF_ELENUM_PER_VECCALC], + (uint64_t)0, + (subM * FLOAT_ELENUM_PER_BLK + FLOAT_ELENUM_PER_VECCALC - 1) / FLOAT_ELENUM_PER_VECCALC, + AscendC::UnaryRepeatParams(1, 1, 8, 8)); + AscendC::PipeBarrier(); + if (goFlag == 1) { + AscendC::WaitFlag(EVENT_ID0); + goFlag = 0; + } + // go = go * dm32_block + for (uint32_t vmulIdx = 0; vmulIdx < k / FLOAT_ELENUM_PER_VECCALC; vmulIdx++) { + AscendC::Mul( + goUbTensor[vmulIdx * FLOAT_ELENUM_PER_VECCALC], + goUbTensor[vmulIdx * FLOAT_ELENUM_PER_VECCALC], + tvUbTensor[HALF_ELENUM_PER_VECCALC], + (uint64_t)0, + subM, + AscendC::BinaryRepeatParams( + 1, 1, 0, kRound / FLOAT_ELENUM_PER_BLK, kRound / FLOAT_ELENUM_PER_BLK, 1)); + AscendC::PipeBarrier(); + } + if (k % FLOAT_ELENUM_PER_VECCALC > 0) { + SetMask(k % FLOAT_ELENUM_PER_VECCALC); + AscendC::Mul( + goUbTensor[k / FLOAT_ELENUM_PER_VECCALC * FLOAT_ELENUM_PER_VECCALC], + goUbTensor[k / FLOAT_ELENUM_PER_VECCALC * FLOAT_ELENUM_PER_VECCALC], + tvUbTensor[HALF_ELENUM_PER_VECCALC], + (uint64_t)0, + subM, + AscendC::BinaryRepeatParams( + 1, 1, 0, kRound / FLOAT_ELENUM_PER_BLK, kRound / FLOAT_ELENUM_PER_BLK, 1)); + AscendC::PipeBarrier(); + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + // go = lo + go + AscendC::Add( + goUbTensor, + goUbTensor, + loUbTensor, + (uint64_t)0, + (subM * kRound + FLOAT_ELENUM_PER_VECCALC - 1) / FLOAT_ELENUM_PER_VECCALC, + AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); + AscendC::PipeBarrier(); + } else { + // gl = ll + AscendC::DataCopy(glUbTensor, llUbTensor[nIdx % MULTIPLIER * FLOAT_ELENUM_PER_LINE], + AscendC::DataCopyParams(1, subMRound / FLOAT_ELENUM_PER_BLK, 0, 0)); + AscendC::PipeBarrier(); + if (goFlag == 1) { + AscendC::WaitFlag(EVENT_ID0); + goFlag = 0; + } + AscendC::DataCopy(goUbTensor, loUbTensor, + AscendC::DataCopyParams(1, subM * kRound / FLOAT_ELENUM_PER_BLK, 0, 0)); + AscendC::PipeBarrier(); + } + AscendC::SetFlag(EVENT_ID1); + if (isLast) { + // gl = castfp32to16(gl) + AscendC::Cast( + glUbTensor.ReinterpretCast(), + glUbTensor, + AscendC::RoundMode::CAST_NONE, + (uint64_t)0, + subMAligned64, + AscendC::UnaryRepeatParams(1, 1, 4, 8)); + AscendC::PipeBarrier(); + // go = castfp32to16(go) + AscendC::Cast( + goUbTensor.ReinterpretCast(), + goUbTensor, + AscendC::RoundMode::CAST_NONE, + (uint64_t)0, + (subM * kRound + FLOAT_ELENUM_PER_VECCALC - 1) / FLOAT_ELENUM_PER_VECCALC, + AscendC::UnaryRepeatParams(1, 1, 4, 8)); + AscendC::PipeBarrier(); + // gl_block = brcb(gl) + AscendC::Brcb( + tvUbTensor.ReinterpretCast(), + glUbTensor.ReinterpretCast(), + subMRound / FLOAT_ELENUM_PER_BLK, + AscendC::BrcbRepeatParams(1, 8)); + AscendC::PipeBarrier(); + // go = go / gl_block + for (uint32_t vdivIdx = 0; vdivIdx < k / HALF_ELENUM_PER_VECCALC; vdivIdx++) { + AscendC::Div( + goUbTensor.ReinterpretCast()[vdivIdx * HALF_ELENUM_PER_VECCALC], + goUbTensor.ReinterpretCast()[vdivIdx * HALF_ELENUM_PER_VECCALC], + tvUbTensor.ReinterpretCast(), + (uint64_t)0, + subM, + AscendC::BinaryRepeatParams( + 1, 1, 0, kRound / HALF_ELENUM_PER_BLK, kRound / HALF_ELENUM_PER_BLK, 1)); + } + if (k % HALF_ELENUM_PER_VECCALC > 0) { + SetMask(k % HALF_ELENUM_PER_VECCALC); + AscendC::Div( + goUbTensor.ReinterpretCast()[k / HALF_ELENUM_PER_VECCALC * HALF_ELENUM_PER_VECCALC], + goUbTensor.ReinterpretCast()[k / HALF_ELENUM_PER_VECCALC * HALF_ELENUM_PER_VECCALC], + tvUbTensor.ReinterpretCast(), + (uint64_t)0, + subM, + AscendC::BinaryRepeatParams( + 1, 1, 0, kRound / HALF_ELENUM_PER_BLK, kRound / HALF_ELENUM_PER_BLK, 1)); + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + // copy O to GM + AscendC::SetFlag(EVENT_ID1); + AscendC::WaitFlag(EVENT_ID1); + copyUbToGmOutput(gOutput, goUbTensor.ReinterpretCast(), layoutOutput, layoutInUb); + if (goFlag == 0) { + AscendC::SetFlag(EVENT_ID0); + goFlag = 1; + } + } + } + } + + CATLASS_DEVICE + void operator()( + AscendC::GlobalTensor gOutput, + AscendC::GlobalTensor gInput, + const LayoutOutput &layoutOutput, + const LayoutInput &layoutInput, + GemmCoord actualBlockShape, + uint32_t nIdx, uint32_t isLast) + { + uint32_t mActual = actualBlockShape.m(); + uint32_t nActual = actualBlockShape.n(); + + uint32_t subBlockIdx = AscendC::GetSubBlockIdx(); + uint32_t subBlockNum = AscendC::GetSubBlockNum(); + + uint32_t mActualPerSubBlock = CeilDiv(mActual, subBlockNum); + uint32_t mActualThisSubBlock = (subBlockIdx == 0) ? mActualPerSubBlock : (mActual - mActualPerSubBlock); + uint32_t mOffset = subBlockIdx * mActualPerSubBlock; + uint32_t nOffset = 0; + + int64_t offsetOutput = layoutOutput.GetOffset(MatrixCoord(mOffset, nOffset)); + auto gOutputThisSubBlock = gOutput[offsetOutput]; + auto layoutOutputThisSubBlock = layoutOutput.GetTileLayout(MatrixCoord(mActualThisSubBlock, nActual)); + + int64_t offsetInput = layoutInput.GetOffset(MatrixCoord(mOffset, nOffset)); + auto gInputThisSubBlock = gInput[offsetInput]; + auto layoutInputThisSubBlock = layoutInput.GetTileLayout(MatrixCoord(mActualThisSubBlock, nActual)); + + subCoreCompute(gOutputThisSubBlock, gInputThisSubBlock, + layoutOutputThisSubBlock, layoutInputThisSubBlock, + nIdx, isLast); + } + +private: + uint32_t goFlag = 1; + AscendC::LocalTensor loUbTensor; + AscendC::LocalTensor dmUbTensor; + AscendC::LocalTensor llUbTensor; + AscendC::LocalTensor glUbTensor; + AscendC::LocalTensor tvUbTensor; + AscendC::LocalTensor goUbTensor; + + CopyGmToUbInput copyGmToUbInput; + CopyUbToGmOutput copyUbToGmOutput; +}; + +} // namespace Catlass::Epilogue::Block + +#endif // CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_FA_RESCALE_O_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_fa_rescale_o_ascend950.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_fa_rescale_o_ascend950.hpp new file mode 100644 index 0000000000..954f6f4c56 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_fa_rescale_o_ascend950.hpp @@ -0,0 +1,216 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_RESCALE_O_A5 +#define CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_RESCALE_O_A5 + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/epilogue/kda_dispatch_policy.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/kda_matrix_coord.hpp" +#include "tla/kda_tensor.hpp" +#include "tla/kda_tla_layout.hpp" + +namespace Catlass::Epilogue::Block { + +template < + class L1TileShape_, + class OType_, + class OTmpType_> +class BlockEpilogue< + EpilogueAscend950FARescaleO, + L1TileShape_, + OType_, + OTmpType_> { +public: + using DispatchPolicy = EpilogueAscend950FARescaleO; + using ArchTag = typename DispatchPolicy::ArchTag; + + using ElementO = typename OType_::Element; + using ElementOTmp = typename OTmpType_::Element; + using LayoutTagO = typename OType_::Layout; + using LayoutTagOTmp = typename OTmpType_::Layout; + using L1TileShape = L1TileShape_; + + static constexpr uint32_t S1_BASE_SIZE = tla::get<0>(L1TileShape{}); + static constexpr uint32_t D_BASE_SIZE = tla::get<2>(L1TileShape{}); + static constexpr uint32_t VEC2_UB_SIZE = S1_BASE_SIZE / 2 * D_BASE_SIZE * sizeof(ElementOTmp); + + CATLASS_DEVICE + BlockEpilogue(Arch::Resource &resource, uint32_t &ubBufAddrStart) + { + vf2OutUb = resource.ubBuf.template GetBufferByByte(ubBufAddrStart); + ubBufAddrStart += VEC2_UB_SIZE; + AscendC::SetFlag(eventOMTE3V); + } + + CATLASS_DEVICE + ~BlockEpilogue() + { + AscendC::WaitFlag(eventOMTE3V); + } + + template + CATLASS_DEVICE void operator()( + TensorDst &attenOutGm, + const AscendC::LocalTensor &expMaxUb, + const AscendC::LocalTensor &sumUb, + TensorSrc &bmm2Res, + bool isFirstLoop, + bool isLastUpdate, + uint64_t MM2_RES_INTRA_EVENT) + { + AscendC::WaitFlag(eventOMTE3V); + uint32_t m = tla::get<0>(bmm2Res.shape()); + uint32_t n = tla::get<1>(bmm2Res.shape()); + constexpr int16_t vlSize = static_cast(AscendC::GetVecLen() / sizeof(ElementOTmp)); + int16_t nLoops = AscendC::CeilDivision(n, vlSize) - 1; + uint32_t tailN = (n - 1) % vlSize + 1; + + __ubuf__ float *vec2ResUbAddr = (__ubuf__ ElementOTmp *)vf2OutUb.GetPhyAddr(); + __ubuf__ float *bmm2UbAddr = (__ubuf__ ElementOTmp *)bmm2Res.data().GetPhyAddr(); + __ubuf__ float *expMaxUbAddr = (__ubuf__ ElementOTmp *)expMaxUb.GetPhyAddr(); + __ubuf__ float *sumUbAddr = (__ubuf__ ElementOTmp *)sumUb.GetPhyAddr(); + + if (isFirstLoop) { + DataCopy(vf2OutUb, bmm2Res.data(), m * n); + } else if (!isLastUpdate) { + FlashUpdateNew(vec2ResUbAddr, bmm2UbAddr, expMaxUbAddr, m, nLoops, tailN); + } else { + FlashUpdateLastNew( + vec2ResUbAddr, bmm2UbAddr, expMaxUbAddr, sumUbAddr, m, nLoops, tailN); + } + AscendC::CrossCoreSetFlag(MM2_RES_INTRA_EVENT); + + if (isFirstLoop && isLastUpdate) { + LastDivNew(vec2ResUbAddr, bmm2UbAddr, sumUbAddr, m, nLoops, tailN); + } + if (isLastUpdate) { + AscendC::LocalTensor attenOut; + attenOut.SetAddr(vf2OutUb.address_); + AscendC::Cast(attenOut, vf2OutUb, AscendC::RoundMode::CAST_ROUND, m * D_BASE_SIZE); + AscendC::SetFlag(eventOVMTE3); + AscendC::WaitFlag(eventOVMTE3); + auto layoutUb = tla::MakeLayout( + tla::MakeShape(m, n), + tla::MakeStride(D_BASE_SIZE, tla::Int<1>{})); + auto attenOutUb = tla::MakeTensor(attenOut, layoutUb, Arch::PositionUB{}); + using CopyUbToGmO = Tile::CopyUb2GmTla; + CopyUbToGmO copyUbToGmO; + copyUbToGmO(attenOutGm, attenOutUb); + } + AscendC::SetFlag(eventOMTE3V); + } + +private: + AscendC::LocalTensor vf2OutUb; + static constexpr int32_t SYNC_MODE = 4; + static constexpr uint16_t FLOAT_REP_SIZE = 64; + static constexpr int32_t eventOVMTE3 = 3; + static constexpr int32_t eventOMTE3V = 3; + + template + __simd_vf__ inline void FlashUpdateNew(__ubuf__ T *updateUb, __ubuf__ T *curUb, __ubuf__ T *expMaxUb, + uint16_t m, uint16_t nLoops, uint32_t tailN) + { + using namespace AscendC::MicroAPI; + RegTensor expMaxVreg; + RegTensor preSrcVreg; + RegTensor curSrcVreg; + RegTensor mulVreg; + RegTensor addVreg; + + MaskReg pregFull = CreateMask(); + MaskReg pregTailN = UpdateMask(tailN); + + for (uint16_t i = 0; i < m; ++i) { + LoadAlign(expMaxVreg, expMaxUb + i); + for (uint16_t j = 0; j < nLoops; ++j) { + LoadAlign(preSrcVreg, updateUb + i * DBaseSize + j * FLOAT_REP_SIZE); + LoadAlign(curSrcVreg, curUb + i * DBaseSize + j * FLOAT_REP_SIZE); + Mul(mulVreg, expMaxVreg, preSrcVreg, pregFull); + Add(addVreg, mulVreg, curSrcVreg, pregFull); + StoreAlign( + updateUb + i * DBaseSize + j * FLOAT_REP_SIZE, addVreg, pregFull); + } + LoadAlign(preSrcVreg, updateUb + i * DBaseSize + nLoops * FLOAT_REP_SIZE); + LoadAlign(curSrcVreg, curUb + i * DBaseSize + nLoops * FLOAT_REP_SIZE); + Mul(mulVreg, expMaxVreg, preSrcVreg, pregTailN); + Add(addVreg, mulVreg, curSrcVreg, pregTailN); + StoreAlign( + updateUb + i * DBaseSize + nLoops * FLOAT_REP_SIZE, addVreg, pregTailN); + } + } + template + __simd_vf__ inline void FlashUpdateLastNew(__ubuf__ T *updateUb, __ubuf__ T *curUb, __ubuf__ T *expMaxUb, + __ubuf__ T *expSumUb, uint16_t m, uint16_t nLoops, uint32_t tailN) + { + using namespace AscendC::MicroAPI; + RegTensor expMaxVreg; + RegTensor preSrcVreg; + RegTensor curSrcVreg; + RegTensor mulVreg; + RegTensor addVreg; + RegTensor divDstVreg; + RegTensor expSumVreg; + + MaskReg pregFull = CreateMask(); + MaskReg pregTailN = UpdateMask(tailN); + + for (uint16_t i = 0; i < m; ++i) { + LoadAlign(expMaxVreg, expMaxUb + i); + LoadAlign(expSumVreg, expSumUb + i); + for (uint16_t j = 0; j < nLoops; ++j) { + LoadAlign(preSrcVreg, updateUb + i * DBaseSize + j * FLOAT_REP_SIZE); + LoadAlign(curSrcVreg, curUb + i * DBaseSize + j * FLOAT_REP_SIZE); + Mul(mulVreg, expMaxVreg, preSrcVreg, pregFull); + Add(addVreg, mulVreg, curSrcVreg, pregFull); + Div(divDstVreg, addVreg, expSumVreg, pregFull); + StoreAlign( + updateUb + i * DBaseSize + j * FLOAT_REP_SIZE, divDstVreg, pregFull); + } + LoadAlign(preSrcVreg, updateUb + i * DBaseSize + nLoops * FLOAT_REP_SIZE); + LoadAlign(curSrcVreg, curUb + i * DBaseSize + nLoops * FLOAT_REP_SIZE); + Mul(mulVreg, expMaxVreg, preSrcVreg, pregTailN); + Add(addVreg, mulVreg, curSrcVreg, pregTailN); + Div(divDstVreg, addVreg, expSumVreg, pregTailN); + StoreAlign( + updateUb + i * DBaseSize + nLoops * FLOAT_REP_SIZE, divDstVreg, pregTailN); + } + } + template + __simd_vf__ inline void LastDivNew(__ubuf__ T *updateUb, __ubuf__ T *curUb, __ubuf__ T *expSumUb, + uint16_t m, uint16_t nLoops, uint32_t tailN) + { + using namespace AscendC::MicroAPI; + RegTensor curSrcVreg; + RegTensor divDstVreg; + RegTensor expSumVreg; + MaskReg pregFull = CreateMask(); + MaskReg pregTailN = UpdateMask(tailN); + for (uint16_t i = 0; i < m; ++i) { + LoadAlign(expSumVreg, expSumUb + i); + for (uint16_t j = 0; j < nLoops; ++j) { + LoadAlign(curSrcVreg, curUb + i * DBaseSize + j * FLOAT_REP_SIZE); + Div(divDstVreg, curSrcVreg, expSumVreg, pregFull); + StoreAlign( + updateUb + i * DBaseSize + j * FLOAT_REP_SIZE, divDstVreg, pregFull); + } + LoadAlign(curSrcVreg, curUb + i * DBaseSize + nLoops * FLOAT_REP_SIZE); + Div(divDstVreg, curSrcVreg, expSumVreg, pregTailN); + StoreAlign( + updateUb + i * DBaseSize + nLoops * FLOAT_REP_SIZE, divDstVreg, pregTailN); + } + } +}; +} // namespace Catlass::Epilogue::Block + +#endif // CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_RESCALE_O_A5 diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_fa_softmax_ascend950.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_fa_softmax_ascend950.hpp new file mode 100644 index 0000000000..60c4563a9b --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_fa_softmax_ascend950.hpp @@ -0,0 +1,405 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_FA_SOFTMAX_A5 +#define CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_FA_SOFTMAX_A5 + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_cross_core_sync.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/epilogue/kda_dispatch_policy.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/kda_matrix_coord.hpp" +#include "tla/kda_tensor.hpp" +#include "tla/kda_tla_layout.hpp" + +namespace Catlass::Epilogue::Block { + +template < + class L1TileShape_, + class PType_, + class SType_, + class MaskType_, + bool ATTENTION_MASK_FLAG_> +class BlockEpilogue< + EpilogueAscend950FASoftmax, + L1TileShape_, + PType_, + SType_, + MaskType_> { +public: + using DispatchPolicy = EpilogueAscend950FASoftmax; + using ArchTag = typename DispatchPolicy::ArchTag; + using L1TileShape = L1TileShape_; + using ElementP = typename PType_::Element; + using ElementS = typename SType_::Element; + using ElementMask = typename MaskType_::Element; + using LayoutTagMask = typename MaskType_::Layout; + using LayoutTagP = typename PType_::Layout; + + static constexpr uint32_t S1_BASE_SIZE = tla::get<0>(L1TileShape{}); + static constexpr uint32_t S2_BASE_SIZE = tla::get<1>(L1TileShape{}); + static constexpr bool ATTENTION_MASK_FLAG = ATTENTION_MASK_FLAG_; + static constexpr int32_t HALF_S1_BASE_SIZE = S1_BASE_SIZE >> 1; + static constexpr int32_t HALF_VEC_SIZE = HALF_S1_BASE_SIZE * sizeof(ElementS); + static constexpr int32_t HALF_MASK_BLOCK_SIZE = HALF_S1_BASE_SIZE * S2_BASE_SIZE * sizeof(ElementMask); + static constexpr int32_t HALF_SCM_BLOCK_SIZE = HALF_S1_BASE_SIZE * S2_BASE_SIZE * sizeof(ElementP); + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(ElementP); + + CATLASS_DEVICE + BlockEpilogue(Arch::Resource &resource, ElementS scaleValue_, uint32_t &ubBufAddrStart) + { + int32_t eventVMTE2 = 0; + int32_t eventMTE2V = 0; + int32_t eventMTE3V = 0; + int32_t eventVMTE3 = 0; + scaleValue = scaleValue_; + expSumUb = resource.ubBuf.template GetBufferByByte(ubBufAddrStart); + ubBufAddrStart += HALF_VEC_SIZE; + nowMaxUb = resource.ubBuf.template GetBufferByByte(ubBufAddrStart); + ubBufAddrStart += HALF_VEC_SIZE; + for (uint32_t i = 0; i < TASK_NUM2; i++) { + vf1OutUbList[i] = resource.ubBuf.template GetBufferByByte(ubBufAddrStart); + ubBufAddrStart += HALF_SCM_BLOCK_SIZE; + if constexpr (ATTENTION_MASK_FLAG) { + attenMaskUbList[i] = resource.ubBuf.template GetBufferByByte(ubBufAddrStart); + ubBufAddrStart += HALF_MASK_BLOCK_SIZE; + eventUbMaskVMTE2List[i] = eventVMTE2++; + eventUbMaskMTE2VList[i] = eventMTE2V++; + AscendC::SetFlag(eventUbMaskVMTE2List[i]); + } + } + + for (uint32_t i = 0; i < TASK_NUM3; i++) { + eventUbPMTE3VList[i] = eventMTE3V++; + eventUbPVMTE3List[i] = eventVMTE3++; + AscendC::SetFlag(eventUbPMTE3VList[i]); + } + } + + CATLASS_DEVICE + ~BlockEpilogue() + { + if constexpr (ATTENTION_MASK_FLAG) { + for (uint32_t i = 0; i < TASK_NUM2; i++) { + AscendC::WaitFlag(eventUbMaskVMTE2List[i]); + } + } + for (uint32_t i = 0; i < TASK_NUM3; i++) { + AscendC::WaitFlag(eventUbPMTE3VList[i]); + } + } + + template + CATLASS_DEVICE void operator()( + TensorDst &vf1OutL1, + AscendC::LocalTensor &sumUb, + AscendC::LocalTensor &lastMaxUb, + AscendC::LocalTensor &expMaxUb, + TensorSrc &bmm1Res, + TensorMask attenMaskGm, + bool isUpdate, + int8_t taskIdMod2, + int8_t taskIdMod3, + uint64_t MM1_RES_INTRA_EVENT, + uint64_t SYNC_V1_C2_FLAG) + { + uint32_t m = tla::get<0>(bmm1Res.shape()); + uint32_t n = tla::get<1>(bmm1Res.shape()); + uint32_t blockStride = AscendC::CeilDivision(m, ELE_NUM_PER_C0) * ELE_NUM_PER_C0; + auto vf1UbLayout = tla::MakeLayout( + tla::MakeShape( + tla::MakeShape(m, tla::Int<1>{}), + tla::MakeShape( + tla::Int{}, + AscendC::CeilDivision(S2_BASE_SIZE, tla::Int{}))), + tla::MakeStride( + tla::MakeStride(tla::Int{}, tla::Int{} * m), + tla::MakeStride(tla::Int<1>{}, blockStride * ELE_NUM_PER_C0))); + auto vf1OutUb = tla::MakeTensor( + vf1OutUbList[taskIdMod2], + vf1UbLayout, + Arch::PositionUB{}); + + constexpr int16_t vlSize = static_cast(AscendC::GetVecLen() / sizeof(ElementS)); + uint32_t tailN = (n - 1) % vlSize + 1; + uint32_t tailM = m; + + __ubuf__ ElementP *outputAddr = (__ubuf__ ElementP *)vf1OutUb.data().GetPhyAddr(); + __ubuf__ ElementS *inputAddr = (__ubuf__ ElementS *)bmm1Res.data().GetPhyAddr(); + __ubuf__ ElementMask *maskUbAddr; + __ubuf__ ElementMask *maskUbUnrollAddr; + __ubuf__ ElementS *lastMaxUbAddr; + __ubuf__ ElementS *nowMaxUbAddr; + __ubuf__ ElementS *expSumUbAddr; + if constexpr (ATTENTION_MASK_FLAG) { + AscendC::WaitFlag(taskIdMod2); + auto layoutUb = tla::MakeLayout( + tla::MakeShape(m, n), + tla::MakeStride(S2_BASE_SIZE, tla::Int<1>{})); + auto attenMaskUb = tla::MakeTensor( + attenMaskUbList[taskIdMod2], + layoutUb, + Arch::PositionUB{}); + using CopyGmToUbMask = Tile::CopyGm2UbTla; + CopyGmToUbMask copyGm2UbMask; + copyGm2UbMask(attenMaskUb, attenMaskGm); + AscendC::SetFlag(taskIdMod2); + AscendC::WaitFlag(taskIdMod2); + maskUbAddr = (__ubuf__ ElementMask *)attenMaskUb.data().GetPhyAddr(); + maskUbUnrollAddr = (__ubuf__ ElementMask *)(attenMaskUb.data().GetPhyAddr() + FLOAT_REP_SIZE); + } + + if (likely(isUpdate)) { + lastMaxUbAddr = (__ubuf__ ElementS *)lastMaxUb.GetPhyAddr(); + nowMaxUbAddr = (__ubuf__ ElementS *)nowMaxUb.GetPhyAddr(); + expSumUbAddr = (__ubuf__ ElementS *)expSumUb.GetPhyAddr(); + } else { + nowMaxUbAddr = (__ubuf__ ElementS *)lastMaxUb.GetPhyAddr(); + expSumUbAddr = (__ubuf__ ElementS *)sumUb.GetPhyAddr(); + } + + if (n == 128) { + ComputeMaskandScale( + inputAddr, maskUbAddr, maskUbUnrollAddr, nowMaxUbAddr, m, tailN, scaleValue); + } else if (n <= 64) { + ComputeMaskandScale( + inputAddr, maskUbAddr, maskUbUnrollAddr, nowMaxUbAddr, m, tailN, scaleValue); + } else { + ComputeMaskandScale( + inputAddr, maskUbAddr, maskUbUnrollAddr, nowMaxUbAddr, m, tailN, scaleValue); + } + + if constexpr (ATTENTION_MASK_FLAG) { + AscendC::SetFlag(taskIdMod2); + } + + if (likely(isUpdate)) { + UpdateMax(nowMaxUbAddr, lastMaxUbAddr, tailM); + } + AscendC::WaitFlag(taskIdMod3); + if (unlikely(n > 64)) { + ComputeExpSubSum( + outputAddr, inputAddr, nowMaxUbAddr, expSumUbAddr, m, blockStride); + } else { + ComputeExpSubSum( + outputAddr, inputAddr, nowMaxUbAddr, expSumUbAddr, m, blockStride); + } + + AscendC::SetFlag(taskIdMod2); + AscendC::WaitFlag(taskIdMod2); + AscendC::CrossCoreSetFlag(MM1_RES_INTRA_EVENT); + + if (likely(m != 0)) { + using CopyUbToL1P = Tile::CopyUb2L1Tla; + CopyUbToL1P copyUb2L1P; + copyUb2L1P(vf1OutL1, vf1OutUb); + } + AscendC::SetFlag(taskIdMod3); + AscendC::CrossCoreSetFlag(SYNC_V1_C2_FLAG); + + if (likely(isUpdate)) { + __ubuf__ ElementS *sumUbAddr = (__ubuf__ ElementS *)sumUb.GetPhyAddr(); + __ubuf__ ElementS *expMaxUbAddr = (__ubuf__ ElementS *)expMaxUb.GetPhyAddr(); + UpdateExpSumAndExpMax( + sumUbAddr, expMaxUbAddr, lastMaxUbAddr, expSumUbAddr, nowMaxUbAddr, tailM); + } + } + +private: + static constexpr uint32_t TASK_NUM2 = 2; + static constexpr uint32_t TASK_NUM3 = 3; + static constexpr int32_t SYNC_MODE = 4; + static constexpr int64_t BLOCK_BYTES = 32; + static constexpr uint32_t REPEAT_STRIDE = 1; + static constexpr float MIN_VALUE = -3e38f; + static constexpr uint16_t FLOAT_REP_SIZE = 64; + static constexpr uint16_t DOUBLE_FLOAT_REP_SIZE = 128; + + int32_t eventUbMaskVMTE2List[TASK_NUM2]; + int32_t eventUbMaskMTE2VList[TASK_NUM2]; + int32_t eventUbPMTE3VList[TASK_NUM3]; + int32_t eventUbPVMTE3List[TASK_NUM3]; + + ElementS scaleValue; + AscendC::LocalTensor attenMaskUbList[TASK_NUM2]; + AscendC::LocalTensor vf1OutUbList[TASK_NUM2]; + AscendC::LocalTensor expSumUb; + AscendC::LocalTensor nowMaxUb; + + enum class NRangeIndex { + N0_64 = 0, + N65_127, + N128, + N128_INF + }; + + template + __simd_vf__ inline void ComputeMaskandScale(__ubuf__ ElementS *srcUb, __ubuf__ uint8_t *maskUb, __ubuf__ uint8_t *maskUbUnroll, + __ubuf__ ElementS *newMaxUb, uint16_t m, uint32_t tailN, ElementS dScale) + { + using namespace AscendC::MicroAPI; + RegTensor minVreg; + RegTensor srcVreg0; + RegTensor srcVreg1; + RegTensor maxVreg; + RegTensor maxTmpVreg; + UnalignReg maxUreg; + MaskReg pregCompare0; + MaskReg pregCompare1; + MaskReg pregFull = CreateMask(); + MaskReg pregTailN = UpdateMask(tailN); + + if constexpr (NRange < NRangeIndex::N128 || HasAtten) { + Duplicate(minVreg, MIN_VALUE); + } + for (uint16_t i = 0; i < m; ++i) { + if constexpr (NRange > NRangeIndex::N0_64) { + LoadAlign(srcVreg0, srcUb + i * S2BaseSize); + LoadAlign(srcVreg1, srcUb + i * S2BaseSize + FLOAT_REP_SIZE); + Muls(srcVreg0, srcVreg0, dScale, pregFull); + Muls(srcVreg1, srcVreg1, dScale, pregTailN); + if constexpr (HasAtten) { + LoadAlign( + pregCompare0, (__ubuf__ uint32_t *&)maskUb, S2BaseSize); + LoadAlign( + pregCompare1, (__ubuf__ uint32_t *&)maskUbUnroll, S2BaseSize); + Select(srcVreg0, minVreg, srcVreg0, pregCompare0); + Select(srcVreg1, minVreg, srcVreg1, pregCompare1); + } + if constexpr (NRange < NRangeIndex::N128) { + Select(srcVreg1, srcVreg1, minVreg, pregTailN); + } + StoreAlign(srcUb + i * S2BaseSize, srcVreg0, pregFull); + StoreAlign(srcUb + i * S2BaseSize + FLOAT_REP_SIZE, srcVreg1, pregFull); + Max(maxTmpVreg, srcVreg0, srcVreg1, pregFull); + ReduceMax(maxVreg, maxTmpVreg, pregFull); + } else { + LoadAlign(srcVreg0, srcUb + i * S2BaseSize); + Muls(srcVreg0, srcVreg0, dScale, pregTailN); + if constexpr (HasAtten) { + LoadAlign( + pregCompare0, (__ubuf__ uint32_t *&)maskUb, S2BaseSize); + Select(srcVreg0, minVreg, srcVreg0, pregCompare0); + } + Select(srcVreg0, srcVreg0, minVreg, pregTailN); + StoreAlign(srcUb + i * S2BaseSize, srcVreg0, pregFull); + ReduceMax(maxVreg, srcVreg0, pregFull); + } + StoreUnAlign(newMaxUb, maxVreg, maxUreg, 1); + } + StoreUnAlignPost(newMaxUb, maxUreg, 0); + } + + template + __simd_vf__ inline void UpdateMax( + __ubuf__ ElementS *nowMaxUb, __ubuf__ ElementS *lastMaxUb, uint32_t tailM) + { + using namespace AscendC::MicroAPI; + RegTensor nowMaxVreg; + RegTensor lastMaxVreg; + RegTensor maxVreg; + + MaskReg pregTailM = UpdateMask(tailM); + LoadAlign(lastMaxVreg, lastMaxUb); + LoadAlign(nowMaxVreg, nowMaxUb); + Max(maxVreg, nowMaxVreg, lastMaxVreg, pregTailM); + StoreAlign(nowMaxUb, maxVreg, pregTailM); + } + + template + __simd_vf__ inline void ComputeExpSubSum(__ubuf__ ElementP *expUb, __ubuf__ ElementS *srcUb, + __ubuf__ ElementS *nowMaxUb, __ubuf__ ElementS *expSumUb, uint16_t m, uint32_t blockStride) + { + using namespace AscendC::MicroAPI; + constexpr static CastTrait castTraitZero = { + RegLayout::ZERO, + SatMode::SAT, + MaskMergeMode::ZEROING, + AscendC::RoundMode::CAST_ROUND, + }; + + constexpr static CastTrait castTraitOne = { + RegLayout::ONE, + SatMode::SAT, + MaskMergeMode::ZEROING, + AscendC::RoundMode::CAST_ROUND, + }; + RegTensor srcVreg0; + RegTensor srcVreg1; + RegTensor expVreg; + RegTensor expVreg0; + RegTensor expVreg1; + RegTensor expSumVreg; + RegTensor maxVreg; + + RegTensor expDstVreg0; + RegTensor expDstVreg1; + RegTensor expDstVreg; + + UnalignReg expSumUreg; + + MaskReg pregFull = CreateMask(); + MaskReg pregFull16 = CreateMask(); + for (uint16_t i = 0; i < m; ++i) { + LoadAlign(maxVreg, nowMaxUb + i); + if constexpr (NRange > NRangeIndex::N0_64) { + LoadAlign(srcVreg0, srcVreg1, srcUb + i * S2BaseSize); + ExpSub(expVreg0, srcVreg0, maxVreg, pregFull); + ExpSub(expVreg1, srcVreg1, maxVreg, pregFull); + Add(expVreg, expVreg0, expVreg1, pregFull); + Cast(expDstVreg0, expVreg0, pregFull); + Cast(expDstVreg1, expVreg1, pregFull); + Or((RegTensor &)expDstVreg, + (RegTensor &)expDstVreg0, (RegTensor &)expDstVreg1, + pregFull16); + StoreAlign( + expUb, expDstVreg, blockStride, REPEAT_STRIDE, pregFull16); + } else { + LoadAlign(srcVreg0, srcUb + i * S2BaseSize); + ExpSub(expVreg, srcVreg0, maxVreg, pregFull); + Cast(expDstVreg, expVreg, pregFull); + DeInterleave(expDstVreg0, expDstVreg1, expDstVreg, expDstVreg); + StoreAlign( + expUb, expDstVreg0, blockStride, REPEAT_STRIDE, pregFull16); + } + ReduceSum(expSumVreg, expVreg, pregFull); + StoreUnAlign(expSumUb, expSumVreg, expSumUreg, 1); + } + StoreUnAlignPost(expSumUb, expSumUreg, 0); + } + + template + __simd_vf__ inline void UpdateExpSumAndExpMax(__ubuf__ ElementS *sumUb, __ubuf__ ElementS *expMaxUb, + __ubuf__ ElementS *maxUb, __ubuf__ ElementS *expSumUb, __ubuf__ ElementS *nowMaxUb, uint32_t tailM) + { + using namespace AscendC::MicroAPI; + RegTensor nowMaxVreg; + RegTensor lastMaxVreg; + RegTensor expMaxVreg; + RegTensor lastExpSumVreg; + RegTensor brcExpSumVreg; + MaskReg pregTailM = UpdateMask(tailM); + LoadAlign(lastMaxVreg, maxUb); + LoadAlign(nowMaxVreg, nowMaxUb); + ExpSub(expMaxVreg, lastMaxVreg, nowMaxVreg, pregTailM); + StoreAlign(expMaxUb, expMaxVreg, pregTailM); + StoreAlign(maxUb, nowMaxVreg, pregTailM); + + LoadAlign(lastExpSumVreg, sumUb); + LoadAlign(brcExpSumVreg, expSumUb); + MulDstAdd(expMaxVreg, lastExpSumVreg, brcExpSumVreg, pregTailM); + StoreAlign(sumUb, expMaxVreg, pregTailM); + } +}; + +} // namespace Catlass::Epilogue::Block + +#endif // CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_FA_SOFTMAX_A5 diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_fixpipe.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_fixpipe.hpp new file mode 100644 index 0000000000..eaf873b4bd --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_fixpipe.hpp @@ -0,0 +1,83 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_EPILOGUE_BLOCK_EPILOGUE_FIXPIPE_HPP +#define CATLASS_EPILOGUE_BLOCK_EPILOGUE_FIXPIPE_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/epilogue/kda_dispatch_policy.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/kda_matrix_coord.hpp" +#include "catlass/epilogue/tile/kda_epilogue_tile_copy.hpp" +#include "catlass/gemm/kda_gemm_helper.hpp" +#include "catlass/epilogue/tile/kda_copy_ub_to_gm_tla.hpp" + +namespace Catlass::Epilogue::Block { + +template < + class L0TileShape_, + class ElementOut_, + class ElementIn_, + bool SPLIT_M_> +class BlockEpilogue< + EpilogueAscend950Fixpipe, + L0TileShape_, + ElementOut_, + ElementIn_> { +public: + using DispatchPolicy = EpilogueAscend950Fixpipe; + using ArchTag = typename DispatchPolicy::ArchTag; + using ElementOut = ElementOut_; + using ElementIn = ElementIn_; + + static constexpr bool SPLIT_M = DispatchPolicy::SPLIT_M; + static constexpr int64_t ML0_ = tla::get<0>(L0TileShape_{}); + static constexpr int64_t NL0_ = tla::get<1>(L0TileShape_{}); + + CATLASS_DEVICE + BlockEpilogue() = default; + + template + CATLASS_DEVICE void operator()(TensorOut &tensorOut, TensorIn &tensorIn) + { + if constexpr (!SPLIT_M && AscendC::GetSubBlockIdx() > 0) { + return; + } + + int64_t blockShapeM = tla::get<0>(tensorOut.shape()); + int64_t blockShapeN = tla::get<1>(tensorOut.shape()); + int64_t halfBlockShapeM = CeilDiv(blockShapeM, AscendC::GetTaskRation()); + if constexpr (SPLIT_M) { + blockShapeM = (static_cast(blockShapeM) & 1UL) > 0UL ? + (halfBlockShapeM - AscendC::GetSubBlockIdx()) : + halfBlockShapeM; + } + + // real copy data size + int64_t copySize = blockShapeM * blockShapeN; + if (copySize <= 0) { + return; + } + + auto tensorTileOut = GetTile( + tensorOut, + tla::MakeCoord(halfBlockShapeM * (AscendC::GetSubBlockIdx() & 0x1), 0), + tla::MakeShape(blockShapeM, blockShapeN)); + + using CopyUbToGm = Epilogue::Tile::CopyUb2GmTla; + CopyUbToGm copyUbToGm; + copyUbToGm(tensorTileOut, tensorIn); + } +}; + +} // namespace Catlass::Epilogue::Block + +#endif // CATLASS_EPILOGUE_BLOCK_EPILOGUE_FIXPIPE_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_mla_fd_rescale_o.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_mla_fd_rescale_o.hpp new file mode 100644 index 0000000000..d57ce3cdb7 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_mla_fd_rescale_o.hpp @@ -0,0 +1,307 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_MLA_FD_RESCALE_O_HPP +#define CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_MLA_FD_RESCALE_O_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/epilogue/kda_dispatch_policy.hpp" +#include "catlass/epilogue/tile/kda_epilogue_tile_copy.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/kda_matrix_coord.hpp" + +namespace Catlass::Epilogue::Block { + +template < + class OutputType_, + class InputType_, + uint32_t ComputeEleNum_> +class BlockEpilogue< + EpilogueAtlasA2MLAFDRescaleO, + OutputType_, + InputType_> { +public: + // Type aliases + using DispatchPolicy = EpilogueAtlasA2MLAFDRescaleO; + using ArchTag = typename DispatchPolicy::ArchTag; + using ElementOutput = typename OutputType_::Element; + using ElementInput = typename InputType_::Element; + using LayoutOutput = typename OutputType_::Layout; + using LayoutInput = typename InputType_::Layout; + + static constexpr uint32_t KV_SPLIT_MAX = DispatchPolicy::KV_SPLIT_MAX; + static constexpr uint32_t HEADS_PROCESS_MAX = DispatchPolicy::HEADS_PROCESS_MAX; + static constexpr uint32_t COMPUTE_ELE_NUM = DispatchPolicy::COMPUTE_ELE_NUM; + static constexpr uint32_t FLOAT_ELENUM_PER_VECCALC = 64; + static constexpr uint32_t FLOAT_BLOCK_SIZE = 8; + static constexpr uint32_t STAGES = 2; + + CATLASS_DEVICE + BlockEpilogue(Arch::Resource &resource, uint32_t kvSplitCoreNum_) + { + kvSplitCoreNum = kvSplitCoreNum_; + + uint32_t ubOffset = 0; + oIn[0] = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += COMPUTE_ELE_NUM * sizeof(float); + oIn[1] = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += COMPUTE_ELE_NUM * sizeof(float); + oTemp[0] = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += COMPUTE_ELE_NUM * sizeof(float); + oTemp[1] = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += COMPUTE_ELE_NUM * sizeof(float); + oSum = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += COMPUTE_ELE_NUM * sizeof(float); + out = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += COMPUTE_ELE_NUM * sizeof(ElementOutput); + lIn = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += KV_SPLIT_MAX * HEADS_PROCESS_MAX * sizeof(float); + lExp = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += KV_SPLIT_MAX * HEADS_PROCESS_MAX * sizeof(float); + lMax = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += HEADS_PROCESS_MAX * sizeof(float); + lSum = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += HEADS_PROCESS_MAX * sizeof(float); + lBrcb[0] = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += HEADS_PROCESS_MAX * FLOAT_BLOCK_SIZE * sizeof(float); + lBrcb[1] = resource.ubBuf.template GetBufferByByte(ubOffset); + + AscendC::SetFlag(EVENT_ID0); + AscendC::SetFlag(EVENT_ID0); + AscendC::SetFlag(EVENT_ID1); + AscendC::SetFlag(EVENT_ID2); + } + CATLASS_DEVICE + ~BlockEpilogue() + { + AscendC::WaitFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID1); + AscendC::WaitFlag(EVENT_ID2); + } + + CATLASS_DEVICE + void SetMask(int32_t len) + { + constexpr int32_t MAX_MASK_LEN = 128; + constexpr int32_t HALF_MASK_LEN = 64; + if (len >= MAX_MASK_LEN) { + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + return; + } + int32_t highMask = len - HALF_MASK_LEN > 0 ? len - HALF_MASK_LEN : 0; + int32_t lowMask = len - HALF_MASK_LEN >= 0 ? HALF_MASK_LEN : len; + if (len < HALF_MASK_LEN) { + AscendC::SetVectorMask(0x0, ((uint64_t)1 << lowMask) - 1); + } else { + AscendC::SetVectorMask(((uint64_t)1 << highMask) - 1, 0xffffffffffffffff); + } + } + + CATLASS_DEVICE + void operator()( + AscendC::GlobalTensor gOutput, + AscendC::GlobalTensor gOCoreTmp, + AscendC::GlobalTensor gl, + uint32_t actualHeads, uint32_t headsProcess, uint32_t headSize) + { + uint32_t kvSplitRound = (kvSplitCoreNum + FLOAT_BLOCK_SIZE - 1) / FLOAT_BLOCK_SIZE * FLOAT_BLOCK_SIZE; + + AscendC::WaitFlag(EVENT_ID2); + AscendC::DataCopyPad( + lIn, gl, + AscendC::DataCopyExtParams( + actualHeads, kvSplitCoreNum * sizeof(ElementInput), 0, + (KV_SPLIT_MAX - kvSplitCoreNum) / FLOAT_BLOCK_SIZE, 0), + AscendC::DataCopyPadExtParams(false, 0, 0, 0)); + + AscendC::SetFlag(EVENT_ID2); + AscendC::WaitFlag(EVENT_ID2); + + SetMask(kvSplitCoreNum); + AscendC::WholeReduceMax( + lMax, lIn, (int32_t)0, actualHeads, 1, 1, 8, + AscendC::ReduceOrder::ORDER_ONLY_VALUE); + AscendC::PipeBarrier(); + + for (uint32_t i = 0; i < kvSplitRound / FLOAT_BLOCK_SIZE; i++) { + AscendC::Brcb( + lExp[i * FLOAT_BLOCK_SIZE], + lMax, + (headsProcess + FLOAT_BLOCK_SIZE - 1) / FLOAT_BLOCK_SIZE, + AscendC::BrcbRepeatParams(KV_SPLIT_MAX / FLOAT_BLOCK_SIZE, + 8 * KV_SPLIT_MAX / FLOAT_BLOCK_SIZE)); + } + AscendC::PipeBarrier(); + + SetMask(kvSplitCoreNum); + AscendC::Sub( + lExp, + lIn, + lExp, + (uint64_t)0, + actualHeads, + AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); + AscendC::PipeBarrier(); + + AscendC::Exp( + lExp, + lExp, + (uint64_t)0, + actualHeads, + AscendC::UnaryRepeatParams(1, 1, 8, 8)); + AscendC::PipeBarrier(); + + AscendC::RepeatReduceSum(lSum, lExp, actualHeads, 0, 0, 1, 1, 8); + AscendC::PipeBarrier(); + + AscendC::Ln(lSum, lSum, (headsProcess + FLOAT_BLOCK_SIZE - 1) / FLOAT_BLOCK_SIZE * FLOAT_BLOCK_SIZE); + AscendC::PipeBarrier(); + + AscendC::Add(lSum, lSum, lMax, (headsProcess + FLOAT_BLOCK_SIZE - 1) / FLOAT_BLOCK_SIZE * FLOAT_BLOCK_SIZE); + AscendC::PipeBarrier(); + + for (uint32_t i = 0; i < kvSplitRound / FLOAT_BLOCK_SIZE; i++) { + AscendC::Brcb( + lExp[i * FLOAT_BLOCK_SIZE], + lSum, + (headsProcess + FLOAT_BLOCK_SIZE - 1) / FLOAT_BLOCK_SIZE, + AscendC::BrcbRepeatParams(KV_SPLIT_MAX / FLOAT_BLOCK_SIZE, + 8 * KV_SPLIT_MAX / FLOAT_BLOCK_SIZE)); + } + AscendC::PipeBarrier(); + + SetMask(kvSplitCoreNum); + AscendC::Sub( + lExp, + lIn, + lExp, + (uint64_t)0, + actualHeads, + AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); + AscendC::PipeBarrier(); + AscendC::SetFlag(EVENT_ID2); + + AscendC::Exp( + lExp, + lExp, + (uint64_t)0, + actualHeads, + AscendC::UnaryRepeatParams(1, 1, 8, 8)); + AscendC::PipeBarrier(); + + // preload + AscendC::WaitFlag(EVENT_ID0); + AscendC::DataCopyPad( + oIn[0], gOCoreTmp, + AscendC::DataCopyExtParams( + actualHeads, headSize * sizeof(ElementInput), + (kvSplitCoreNum * headSize - headSize) * sizeof(ElementInput), 0, 0), + AscendC::DataCopyPadExtParams(false, 0, 0, 0)); + AscendC::SetFlag(EVENT_ID0); + + AscendC::SetFlag(EVENT_ID2); + AscendC::WaitFlag(EVENT_ID2); + + SetMask(FLOAT_ELENUM_PER_VECCALC); + uint32_t bufferId = 0; + for (uint32_t i = 0; i < kvSplitCoreNum; i++) { + // load next o + if (i < kvSplitCoreNum - 1) { + uint32_t nextBufferId = 1 - bufferId; + AscendC::WaitFlag(oInEventList[nextBufferId]); + AscendC::DataCopyPad( + oIn[nextBufferId], gOCoreTmp[(i + 1) * headSize], + AscendC::DataCopyExtParams( + actualHeads, headSize * sizeof(ElementInput), + (kvSplitCoreNum * headSize - headSize) * sizeof(ElementInput), 0, 0), + AscendC::DataCopyPadExtParams(false, 0, 0, 0)); + AscendC::SetFlag(oInEventList[nextBufferId]); + } + + AscendC::PipeBarrier(); + for (uint32_t j = 0; j < actualHeads; j++) { + float a = lExp[j * KV_SPLIT_MAX + i].GetValue(0); + AscendC::SetFlag(oTempEventList[bufferId]); + AscendC::WaitFlag(oTempEventList[bufferId]); + AscendC::Duplicate( + lBrcb[bufferId][j * FLOAT_BLOCK_SIZE], a, uint64_t(0), 1, 0, 0); + } + AscendC::PipeBarrier(); + + // caculate current o + AscendC::WaitFlag(oInEventList[bufferId]); + uint32_t loops = (headSize + FLOAT_ELENUM_PER_VECCALC - 1) / FLOAT_ELENUM_PER_VECCALC; + if (i > 0) { + for (uint32_t j = 0; j < loops; j++) { + AscendC::Mul( + oTemp[bufferId][j * FLOAT_ELENUM_PER_VECCALC], lBrcb[bufferId], + oIn[bufferId][j * FLOAT_ELENUM_PER_VECCALC], (uint64_t)0, + actualHeads, + AscendC::BinaryRepeatParams( + 1, 0, 1, headSize / FLOAT_BLOCK_SIZE, 1, + headSize / FLOAT_BLOCK_SIZE)); + } + } else { + for (uint32_t j = 0; j < loops; j++) { + AscendC::Mul( + oSum[j * FLOAT_ELENUM_PER_VECCALC], lBrcb[bufferId], + oIn[bufferId][j * FLOAT_ELENUM_PER_VECCALC], (uint64_t)0, + actualHeads, + AscendC::BinaryRepeatParams( + 1, 0, 1, headSize / FLOAT_BLOCK_SIZE, 1, + headSize / FLOAT_BLOCK_SIZE)); + } + } + AscendC::PipeBarrier(); + AscendC::SetFlag(oInEventList[bufferId]); + + if (i > 0) { + AscendC::Add(oSum, oSum, oTemp[bufferId], actualHeads * headSize); + } + AscendC::PipeBarrier(); + bufferId = 1 - bufferId; + } + + AscendC::WaitFlag(EVENT_ID0); + if (std::is_same::value) { + AscendC::Cast(out, oSum, AscendC::RoundMode::CAST_RINT, actualHeads * headSize); + } else { + AscendC::Cast(out, oSum, AscendC::RoundMode::CAST_NONE, actualHeads * headSize); + } + + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + AscendC::DataCopyPad( + gOutput, out, + AscendC::DataCopyExtParams( + actualHeads, headSize * sizeof(ElementOutput), 0, 0, 0)); + AscendC::SetFlag(EVENT_ID0); + } + +private: + uint32_t kvSplitCoreNum = 1; + AscendC::LocalTensor out; + AscendC::LocalTensor oIn[STAGES]; + AscendC::LocalTensor oTemp[STAGES]; + AscendC::LocalTensor lBrcb[STAGES]; + AscendC::LocalTensor oSum; + AscendC::LocalTensor lIn; + AscendC::LocalTensor lExp; + AscendC::LocalTensor lTrans; + AscendC::LocalTensor lMax; + AscendC::LocalTensor lSum; + + int32_t oTempEventList[STAGES] = {0, 1}; + int32_t oInEventList[STAGES] = {0, 1}; +}; +} // namespace Catlass::Epilogue::Block +#endif // CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_MLA_FD_RESCALE_O_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_mla_rescale_o.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_mla_rescale_o.hpp new file mode 100644 index 0000000000..b1723cc657 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_mla_rescale_o.hpp @@ -0,0 +1,479 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_MLA_RESCALE_O_HPP +#define CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_MLA_RESCALE_O_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/epilogue/kda_dispatch_policy.hpp" +#include "catlass/epilogue/tile/kda_epilogue_tile_copy.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/kda_matrix_coord.hpp" + +namespace Catlass::Epilogue::Block { + +template < + class OutputType_, + class UpdateType_, + class InputType_> +class BlockEpilogue< + EpilogueAtlasA2MLARescaleO, + OutputType_, + UpdateType_, + InputType_> { +public: + // Type aliases + using DispatchPolicy = EpilogueAtlasA2MLARescaleO; + using ArchTag = typename DispatchPolicy::ArchTag; + + using ElementOutput = typename OutputType_::Element; + using ElementUpdate = typename UpdateType_::Element; + using ElementInput = typename InputType_::Element; + + using LayoutOutput = typename OutputType_::Layout; + using LayoutUpdate = typename UpdateType_::Layout; + using LayoutInput = typename InputType_::Layout; + + static constexpr uint32_t HALF_ELENUM_PER_BLK = 16; + static constexpr uint32_t HALF_ELENUM_PER_VECCALC = 128; + static constexpr uint32_t FLOAT_ELENUM_PER_VECCALC = 64; + static constexpr uint32_t HALF_ELENUM_PER_LINE = 256; + static constexpr uint32_t FLOAT_ELENUM_PER_LINE = 128; + static constexpr uint32_t MULTIPLIER = 2; + static constexpr uint32_t FLOAT_BLOCK_SIZE = 8; + static constexpr uint32_t FLOAT_VECTOR_SIZE = 64; + static constexpr uint32_t UB_UINT8_LINE_SIZE = 512; + static constexpr uint32_t UB_UINT8_BLOCK_SIZE_MLA = 16384; + static constexpr uint32_t ROW_WISE_CYCLE_TILE = 8; + static constexpr uint32_t HALF_DM_UB_SIZE = 128; + static constexpr uint32_t HALF_LL_UB_SIZE = 256; + static constexpr uint32_t VECTOR_SIZE = 128; + + CATLASS_DEVICE + BlockEpilogue(Arch::Resource &resource, uint32_t kvSplitCoreNum_) + { + // Allocate UB space + constexpr uint32_t LO_UB_TENSOR_OFFSET = 4 * UB_UINT8_BLOCK_SIZE_MLA; + constexpr uint32_t DM_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA + 6 * UB_UINT8_LINE_SIZE; + constexpr uint32_t LL_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA + 9 * UB_UINT8_LINE_SIZE; + constexpr uint32_t GL_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA + 15 * UB_UINT8_LINE_SIZE; + constexpr uint32_t GO_UB_TENSOR_OFFSET = 8 * UB_UINT8_BLOCK_SIZE_MLA; + constexpr uint32_t TV_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE_MLA; + constexpr uint32_t HM_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA + 1 * UB_UINT8_LINE_SIZE; + constexpr uint32_t GM_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA + 13 * UB_UINT8_LINE_SIZE; + + kvSplitCoreNum = kvSplitCoreNum_; + loUbTensor = resource.ubBuf.template GetBufferByByte(LO_UB_TENSOR_OFFSET); + dmUbTensor = resource.ubBuf.template GetBufferByByte(DM_UB_TENSOR_OFFSET); + llUbTensor = resource.ubBuf.template GetBufferByByte(LL_UB_TENSOR_OFFSET); + glUbTensor = resource.ubBuf.template GetBufferByByte(GL_UB_TENSOR_OFFSET); + tvUbTensor = resource.ubBuf.template GetBufferByByte(TV_UB_TENSOR_OFFSET); + goUbTensor32 = resource.ubBuf.template GetBufferByByte(GO_UB_TENSOR_OFFSET); + goUbTensor16 = resource.ubBuf.template GetBufferByByte(GO_UB_TENSOR_OFFSET); + hmUbTensor = resource.ubBuf.template GetBufferByByte(HM_UB_TENSOR_OFFSET); + gmUbTensor = resource.ubBuf.template GetBufferByByte(GM_UB_TENSOR_OFFSET); + } + + CATLASS_DEVICE + ~BlockEpilogue() + { + } + + CATLASS_DEVICE + void SetMask(int32_t len) + { + uint64_t mask = 0; + uint64_t one = 1; + uint64_t temp = len % FLOAT_VECTOR_SIZE; + for (int64_t i = 0; i < temp; i++) { + mask |= one << i; + } + if (len == VECTOR_SIZE) { + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + } else if (len >= FLOAT_VECTOR_SIZE) { + AscendC::SetVectorMask(mask, (uint64_t)-1); + } else { + AscendC::SetVectorMask(0x0, mask); + } + } + + CATLASS_DEVICE + void SubCoreCompute( + AscendC::GlobalTensor gInput, + AscendC::GlobalTensor gUpdate, + AscendC::GlobalTensor gOutput, + AscendC::GlobalTensor gOCoreTmp, + AscendC::GlobalTensor gl, + const LayoutInput &layoutInput, + const LayoutOutput &layoutOutput, + const LayoutUpdate &layoutUpdate, + uint32_t nIdx, + uint32_t isLastNTile, + uint32_t needRowLoop, + uint32_t rowLoopIdx, + uint32_t proTokenIdx, + uint32_t proTokenNum, + uint32_t epiTokenNum, + uint32_t integralHeadNum, + uint32_t rescaleOPingPongFlag, + uint32_t &glFlag) + { + uint32_t curRowNum = layoutInput.shape(0); + uint32_t embed = layoutInput.shape(1); + uint32_t embedRound = layoutInput.stride(0); + uint32_t strideQO = layoutOutput.stride(0); + uint32_t tokenNumPerHead = layoutOutput.shape(0); + uint32_t curRowNumAligned64 = (curRowNum + FLOAT_ELENUM_PER_VECCALC - 1) / FLOAT_ELENUM_PER_VECCALC; + uint32_t curRowNumRound = (curRowNum + HALF_ELENUM_PER_BLK - 1) / HALF_ELENUM_PER_BLK * HALF_ELENUM_PER_BLK; + uint64_t dmUbOffsetCurCycle = (uint64_t)(rescaleOPingPongFlag * HALF_DM_UB_SIZE + + rowLoopIdx * ROW_WISE_CYCLE_TILE); + uint64_t llUbOffsetCurCycle = (uint64_t)(rescaleOPingPongFlag * HALF_LL_UB_SIZE + + rowLoopIdx * ROW_WISE_CYCLE_TILE); + uint32_t oUbOffset = oPingPangFlag * ROW_WISE_CYCLE_TILE * embedRound; + AscendC::WaitFlag(oPingPangFlag); + if ((nIdx - 1) != 0) { + AscendC::DataCopy( + loUbTensor[oUbOffset], + gInput, + AscendC::DataCopyParams( + 1, curRowNum * embedRound / FLOAT_BLOCK_SIZE, 0, 0)); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + } + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + AscendC::WaitFlag(oPingPangFlag + 4); + if ((nIdx - 1) != 0) { + // *** dm = exp(dm) + if (rowLoopIdx == 0) { + AscendC::Exp( + dmUbTensor[dmUbOffsetCurCycle], + dmUbTensor[dmUbOffsetCurCycle], + (uint64_t)0, + curRowNumAligned64, + AscendC::UnaryRepeatParams(1, 1, 8, 8)); + AscendC::PipeBarrier(); + AscendC::Mul( + glUbTensor, + dmUbTensor[dmUbOffsetCurCycle], + glUbTensor, + (uint64_t)0, + curRowNumAligned64, + AscendC::BinaryRepeatParams( + 1, 1, 1, 8, 8, 8)); + AscendC::PipeBarrier(); + AscendC::Add( + glUbTensor, + glUbTensor, + llUbTensor[llUbOffsetCurCycle], + (uint64_t)0, + curRowNumAligned64, + AscendC::BinaryRepeatParams( + 1, 1, 1, 8, 8, 8)); + AscendC::PipeBarrier(); + } + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + AscendC::Brcb( + tvUbTensor.ReinterpretCast(), + dmUbTensor[dmUbOffsetCurCycle].ReinterpretCast(), + curRowNumRound / FLOAT_BLOCK_SIZE, + AscendC::BrcbRepeatParams(1, 8)); + AscendC::PipeBarrier(); + if (needRowLoop) { + AscendC::DataCopy( + goUbTensor32[oUbOffset], + gUpdate, + AscendC::DataCopyParams( + 1, curRowNum * embedRound / FLOAT_BLOCK_SIZE, 0, 0)); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + } + // *** go = go * dm_block + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + for (uint32_t mulIdx = 0; mulIdx < embed / FLOAT_VECTOR_SIZE; ++mulIdx) { + AscendC::Mul( + goUbTensor32[oUbOffset + mulIdx * FLOAT_VECTOR_SIZE], + goUbTensor32[oUbOffset + mulIdx * FLOAT_VECTOR_SIZE], + tvUbTensor, + (uint64_t)0, + curRowNum, + AscendC::BinaryRepeatParams( + 1, 1, 0, embedRound / FLOAT_BLOCK_SIZE, embedRound / FLOAT_BLOCK_SIZE, 1)); + } + if (embed % FLOAT_VECTOR_SIZE > 0) { + SetMask(embed % FLOAT_VECTOR_SIZE); + AscendC::Mul( + goUbTensor32[oUbOffset + embed / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + goUbTensor32[oUbOffset + embed / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + tvUbTensor, + (uint64_t)0, + curRowNum, + AscendC::BinaryRepeatParams( + 1, 1, 0, embedRound / FLOAT_BLOCK_SIZE, embedRound / FLOAT_BLOCK_SIZE, 1)); + + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + AscendC::PipeBarrier(); + // *** go = lo + go + AscendC::Add( + goUbTensor32[oUbOffset], + goUbTensor32[oUbOffset], + loUbTensor[oUbOffset], + (uint64_t)0, + (curRowNum * embedRound + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, + AscendC::BinaryRepeatParams( + 1, 1, 1, 8, 8, 8)); + + AscendC::PipeBarrier(); + } else { + // *** gl = ll + if (rowLoopIdx == 0) { + AscendC::DataCopy( + glUbTensor, + llUbTensor[llUbOffsetCurCycle], + AscendC::DataCopyParams( + 1, 64 / FLOAT_BLOCK_SIZE, 0, 0)); + AscendC::PipeBarrier(); + } + AscendC::DataCopy( + goUbTensor32[oUbOffset], + gInput, + AscendC::DataCopyParams( + 1, curRowNum * embedRound / FLOAT_BLOCK_SIZE, 0, 0)); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + } + AscendC::SetFlag(oPingPangFlag); + + if (isLastNTile) { + AscendC::Brcb( + tvUbTensor.ReinterpretCast(), + glUbTensor.ReinterpretCast()[rowLoopIdx * ROW_WISE_CYCLE_TILE], + curRowNumRound / FLOAT_BLOCK_SIZE, + AscendC::BrcbRepeatParams(1, 8)); + AscendC::PipeBarrier(); + // *** go = go / gl_block + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + for (uint32_t divIdx = 0; divIdx < embed / FLOAT_VECTOR_SIZE; ++divIdx) { + AscendC::Div( + goUbTensor32[oUbOffset + divIdx * FLOAT_VECTOR_SIZE], + goUbTensor32[oUbOffset + divIdx * FLOAT_VECTOR_SIZE], + tvUbTensor, + (uint64_t)0, + curRowNum, + AscendC::BinaryRepeatParams( + 1, 1, 0, embedRound / FLOAT_BLOCK_SIZE, embedRound / FLOAT_BLOCK_SIZE, 1)); + } + + if (embed % FLOAT_VECTOR_SIZE > 0) { + SetMask(embed % FLOAT_VECTOR_SIZE); + AscendC::Div( + goUbTensor32[oUbOffset + embed / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + goUbTensor32[oUbOffset + embed / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + tvUbTensor, + (uint64_t)0, + curRowNum, + AscendC::BinaryRepeatParams( + 1, 1, 0, embedRound / FLOAT_BLOCK_SIZE, embedRound / FLOAT_BLOCK_SIZE, 1)); + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); // fix hidden_size=96 + } + AscendC::PipeBarrier(); + + if (kvSplitCoreNum != 1) { + // log(l) + AscendC::Ln( + tvUbTensor, + tvUbTensor, + (uint64_t)0, + curRowNum, + AscendC::UnaryRepeatParams(1, 1, 8, 8)); + AscendC::PipeBarrier(); + AscendC::Brcb( + hmUbTensor.ReinterpretCast(), + gmUbTensor.ReinterpretCast()[rowLoopIdx * ROW_WISE_CYCLE_TILE], + curRowNumRound / FLOAT_BLOCK_SIZE, + AscendC::BrcbRepeatParams(1, 8)); + AscendC::PipeBarrier(); + // logf(lse_sum) + lse_max + AscendC::Add( + tvUbTensor, + tvUbTensor, + hmUbTensor, + (uint64_t)0, + curRowNum, + AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); + AscendC::PipeBarrier(); + + AscendC::SetFlag(EVENT_ID2); + AscendC::WaitFlag(EVENT_ID2); + AscendC::DataCopyPad(gl, tvUbTensor, + AscendC::DataCopyExtParams(curRowNum, 4, 0, (kvSplitCoreNum - 1) * 4, 0)); + + if (glFlag == 0) { + AscendC::SetFlag(EVENT_ID2); + glFlag = 1; + } + uint32_t srcGap = ((embed % 16 <= 8) && (embed % 16 > 0)) ? 1 : 0; + AscendC::DataCopyPad(gOCoreTmp, goUbTensor32[oUbOffset], + AscendC::DataCopyExtParams(curRowNum, embed * 4, srcGap, (kvSplitCoreNum - 1) * embed * 4, 0)); + AscendC::SetFlag(EVENT_ID3); + AscendC::WaitFlag(EVENT_ID3); + AscendC::SetFlag(EVENT_ID1); + AscendC::WaitFlag(EVENT_ID1); + } else { + // *** go = castfp32to16(go) + if (std::is_same::value) { + AscendC::Cast( + goUbTensor16[oUbOffset * 2], goUbTensor32[oUbOffset], + AscendC::RoundMode::CAST_RINT, (uint64_t)0, + (curRowNum * embedRound + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, + AscendC::UnaryRepeatParams(1, 1, 4, 8)); + } else { + AscendC::Cast( + goUbTensor16[oUbOffset * 2], goUbTensor32[oUbOffset], + AscendC::RoundMode::CAST_NONE, (uint64_t)0, + (curRowNum * embedRound + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, + AscendC::UnaryRepeatParams(1, 1, 4, 8)); + } + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + // ********************* move O to GM ************************ + if (tokenNumPerHead == 1) { + AscendC::DataCopyPad(gOutput, goUbTensor16[oUbOffset * 2], + AscendC::DataCopyExtParams(curRowNum, embed * 2, 0, 0, 0)); + } else { + uint32_t innerOGmOffset = 0; + uint32_t inner_go_ubuf_offset = oUbOffset * 2; + if (proTokenNum != 0) { + AscendC::DataCopyPad(gOutput[innerOGmOffset + proTokenIdx * strideQO], + goUbTensor16[inner_go_ubuf_offset], + AscendC::DataCopyExtParams(proTokenNum, + embed * 2, 0, (strideQO - embed) * 2, 0)); + innerOGmOffset += embed; + inner_go_ubuf_offset += proTokenNum * embed; + } + for (uint32_t qN_idx = 0; qN_idx < integralHeadNum; qN_idx++) { + AscendC::DataCopyPad(gOutput[innerOGmOffset], + goUbTensor16[inner_go_ubuf_offset], + AscendC::DataCopyExtParams(tokenNumPerHead, + embed * 2, 0, (strideQO - embed) * 2, 0)); + innerOGmOffset += embed; + inner_go_ubuf_offset += tokenNumPerHead * embed; + } + if (epiTokenNum != 0) { + AscendC::DataCopyPad(gOutput[innerOGmOffset], + goUbTensor16[inner_go_ubuf_offset], + AscendC::DataCopyExtParams(epiTokenNum, + embed * 2, 0, (strideQO - embed) * 2, 0)); + } + } + } + } else if (needRowLoop) { + AscendC::SetFlag(EVENT_ID5); + AscendC::WaitFlag(EVENT_ID5); + AscendC::DataCopy( + gUpdate, + goUbTensor32[oUbOffset], + AscendC::DataCopyParams( + 1, curRowNum * embedRound / FLOAT_BLOCK_SIZE, 0, 0)); + } + AscendC::SetFlag(oPingPangFlag + 4); + if (needRowLoop) { + oPingPangFlag = 1 - oPingPangFlag; + } + } + + CATLASS_DEVICE + void operator()( + AscendC::GlobalTensor gInput, + AscendC::GlobalTensor gUpdate, + AscendC::GlobalTensor gOutput, + AscendC::GlobalTensor gOCoreTmp, + AscendC::GlobalTensor gl, + const LayoutInput &layoutInput, + const LayoutOutput &layoutOutput, + const LayoutUpdate &layoutUpdate, + GemmCoord actualBlockShape, + uint32_t nIdx, + uint32_t isLastNTile, + uint32_t curHeadNum, + uint32_t rescaleOPingPongFlag, + uint32_t &glFlag) + { + uint32_t tokenNumPerHead = layoutOutput.shape(0); + uint32_t embed = layoutInput.shape(1); + uint32_t rowActual = actualBlockShape.m(); // curHeadNum * tokenNumPerHead + uint32_t columnActual = actualBlockShape.n(); // embed + + uint32_t subBlockIdx = AscendC::GetSubBlockIdx(); + uint32_t subBlockNum = AscendC::GetSubBlockNum(); + + uint32_t curHeadSplitSubBlock = curHeadNum / subBlockNum; + uint32_t curHeadThisSubBlock = (subBlockIdx == 0) ? curHeadSplitSubBlock : (curHeadNum - curHeadSplitSubBlock); + + uint32_t rowActualThisSubBlock = curHeadThisSubBlock * tokenNumPerHead; + uint32_t rowOffsetSubBlock = subBlockIdx * curHeadSplitSubBlock * tokenNumPerHead; + uint32_t outOffsetSubBlock = subBlockIdx * curHeadSplitSubBlock * embed; + + if (rowActualThisSubBlock > 0) { + uint32_t rowLoop = (rowActualThisSubBlock + ROW_WISE_CYCLE_TILE - 1) / ROW_WISE_CYCLE_TILE; + uint32_t needRowLoop = (rowLoop > 1) ? 1 : 0; + // The rows of each cycle consist of multiple heads with several tokens. + // There are several integral heads, one prologue head, one epilogue head. + uint32_t proTokenIdx = 0; // the token idx of the start token of the prologue part + uint32_t proTokenNum = 0; // the token num of the prologue part + uint32_t epiTokenNum = 0; // the token num of the epilogue part + uint32_t integralHeadNum = 0; // the number of integral heads within a cycle + for (uint32_t rowLoopIdx = 0; rowLoopIdx < rowLoop; rowLoopIdx++) { + uint32_t rowOffsetLoop = rowLoopIdx * ROW_WISE_CYCLE_TILE; + uint32_t rowOffsetCurCycle = rowOffsetSubBlock + rowOffsetLoop; + uint32_t rowActualCurCycle = (rowLoopIdx == (rowLoop - 1)) ? + rowActualThisSubBlock - rowLoopIdx * ROW_WISE_CYCLE_TILE : + ROW_WISE_CYCLE_TILE; + int64_t offsetInput = layoutInput.GetOffset(MatrixCoord(rowOffsetCurCycle, 0)); + auto gInputThisCurCycle = gInput[offsetInput]; + auto layoutInputCurCycle = layoutInput.GetTileLayout(MatrixCoord(rowActualCurCycle, columnActual)); + int64_t offsetOutput = rowLoopIdx * ROW_WISE_CYCLE_TILE / tokenNumPerHead * embed + outOffsetSubBlock; + auto gOutputCurCycle = gOutput[offsetOutput]; + auto layoutOutputCurCycle = layoutOutput; + int64_t offsetUpdate = layoutUpdate.GetOffset(MatrixCoord(rowOffsetCurCycle, 0)); + auto gUpdateCurCycle = gUpdate[offsetUpdate]; + auto layoutUpdateCurCycle = layoutUpdate.GetTileLayout(MatrixCoord(rowActualCurCycle, columnActual)); + proTokenIdx = epiTokenNum; + proTokenNum = (tokenNumPerHead - epiTokenNum) % tokenNumPerHead; + integralHeadNum = (rowActualCurCycle - proTokenNum) / tokenNumPerHead; + epiTokenNum = rowActualCurCycle - proTokenNum - integralHeadNum * tokenNumPerHead; + SubCoreCompute(gInputThisCurCycle, gUpdateCurCycle, gOutputCurCycle, + gOCoreTmp[rowOffsetLoop * embed * kvSplitCoreNum], + gl[rowOffsetLoop * kvSplitCoreNum], + layoutInputCurCycle, layoutOutputCurCycle, layoutUpdateCurCycle, + nIdx, isLastNTile, needRowLoop, rowLoopIdx, + proTokenIdx, proTokenNum, epiTokenNum, integralHeadNum, rescaleOPingPongFlag, glFlag); + } + } + } + +private: + uint32_t kvSplitCoreNum = 1; + uint32_t oPingPangFlag = 0; + AscendC::LocalTensor goUbTensor16; + AscendC::LocalTensor loUbTensor; + AscendC::LocalTensor dmUbTensor; + AscendC::LocalTensor llUbTensor; + AscendC::LocalTensor glUbTensor; + AscendC::LocalTensor tvUbTensor; + AscendC::LocalTensor goUbTensor32; + AscendC::LocalTensor hmUbTensor; + AscendC::LocalTensor gmUbTensor; +}; + +} // namespace Catlass::Epilogue::Block + +#endif // CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_MLA_RESCALE_O_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_mla_tp1_rescale_o.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_mla_tp1_rescale_o.hpp new file mode 100644 index 0000000000..5ac4da6989 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_mla_tp1_rescale_o.hpp @@ -0,0 +1,341 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_MLA_TP1_RESCALE_O_HPP +#define CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_MLA_TP1_RESCALE_O_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/epilogue/kda_dispatch_policy.hpp" +#include "catlass/epilogue/tile/kda_epilogue_tile_copy.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/kda_matrix_coord.hpp" + +namespace Catlass::Epilogue::Block { + +template +class BlockEpilogue { +public: + // Type aliases + using DispatchPolicy = EpilogueAtlasA2MLATP1RescaleO; + using ArchTag = typename DispatchPolicy::ArchTag; + + using ElementOutput = typename OutputType_::Element; + using ElementUpdate = typename UpdateType_::Element; + using ElementInput = typename InputType_::Element; + + using LayoutOutput = typename OutputType_::Layout; + using LayoutUpdate = typename UpdateType_::Layout; + using LayoutInput = typename InputType_::Layout; + + static constexpr uint32_t HALF_ELENUM_PER_BLK = 16; + static constexpr uint32_t HALF_ELENUM_PER_VECCALC = 128; + static constexpr uint32_t FLOAT_ELENUM_PER_VECCALC = 64; + static constexpr uint32_t HALF_ELENUM_PER_LINE = 256; + static constexpr uint32_t FLOAT_ELENUM_PER_LINE = 128; + static constexpr uint32_t MULTIPLIER = 2; + static constexpr uint32_t FLOAT_BLOCK_SIZE = 8; + static constexpr uint32_t FLOAT_VECTOR_SIZE = 64; + static constexpr uint32_t UB_UINT8_LINE_SIZE = 512; + static constexpr uint32_t UB_UINT8_BLOCK_SIZE_MLA = 16384; + static constexpr uint32_t ROW_WISE_CYCLE_TILE = 8; + static constexpr uint32_t HALF_DM_UB_SIZE = 64; + static constexpr uint32_t HALF_LL_UB_SIZE = 256; + static constexpr uint32_t VECTOR_SIZE = 128; + static constexpr uint32_t NUM4 = 4; + + CATLASS_DEVICE + BlockEpilogue(Arch::Resource &resource, uint32_t kvSplitCoreNum_ = 1) + { + // Allocate UB space + constexpr uint32_t LO_UB_TENSOR_OFFSET = 4 * UB_UINT8_BLOCK_SIZE_MLA; + constexpr uint32_t DM_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA + 6 * UB_UINT8_LINE_SIZE; + constexpr uint32_t GL_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA + 16 * UB_UINT8_LINE_SIZE; + constexpr uint32_t GO_UB_TENSOR_OFFSET = 8 * UB_UINT8_BLOCK_SIZE_MLA; + constexpr uint32_t TV_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE_MLA; + constexpr uint32_t HM_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA + 1 * UB_UINT8_LINE_SIZE; + constexpr uint32_t GM_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA + 14 * UB_UINT8_LINE_SIZE; + + kvSplitCoreNum = kvSplitCoreNum_; + loUbTensor = resource.ubBuf.template GetBufferByByte(LO_UB_TENSOR_OFFSET); + dmUbTensor = resource.ubBuf.template GetBufferByByte(DM_UB_TENSOR_OFFSET); + glUbTensor = resource.ubBuf.template GetBufferByByte(GL_UB_TENSOR_OFFSET); + tvUbTensor = resource.ubBuf.template GetBufferByByte(TV_UB_TENSOR_OFFSET); + goUbTensor16 = resource.ubBuf.template GetBufferByByte(GO_UB_TENSOR_OFFSET); + goUbTensor32 = resource.ubBuf.template GetBufferByByte(GO_UB_TENSOR_OFFSET); + hmUbTensor = resource.ubBuf.template GetBufferByByte(HM_UB_TENSOR_OFFSET); + gmUbTensor = resource.ubBuf.template GetBufferByByte(GM_UB_TENSOR_OFFSET); + } + + CATLASS_DEVICE + ~BlockEpilogue() {} + + CATLASS_DEVICE + void SetMask(int32_t len) + { + uint64_t mask = 0; + uint64_t one = 1; + uint64_t temp = len % FLOAT_VECTOR_SIZE; + for (int64_t i = 0; i < temp; i++) { + mask |= one << i; + } + + if (len == VECTOR_SIZE) { + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + } else if (len >= FLOAT_VECTOR_SIZE) { + AscendC::SetVectorMask(mask, (uint64_t)-1); + } else { + AscendC::SetVectorMask(0x0, mask); + } + } + + CATLASS_DEVICE + void SetkvSplitCoreNum(uint32_t kvSplitCoreNum_) + { + kvSplitCoreNum = kvSplitCoreNum_; + } + + CATLASS_DEVICE + void SubCoreCompute(AscendC::GlobalTensor gInput, AscendC::GlobalTensor gUpdate, + AscendC::GlobalTensor gOutput, AscendC::GlobalTensor gOCoreTmp, + AscendC::GlobalTensor gl, const LayoutInput &layoutInput, + const LayoutOutput &layoutOutput, const LayoutUpdate &layoutUpdate, uint32_t nIdx, + uint32_t isLastNTile, uint32_t needRowLoop, uint32_t rowLoopIdx, uint32_t rescaleOPingPongFlag, + uint32_t &glFlag) + { + uint32_t curRowNum = layoutInput.shape(0); + uint32_t embed = layoutInput.shape(1); + uint32_t embedRound = layoutInput.stride(0); + uint32_t curRowNumRound = RoundUp(curRowNum); + uint64_t dmUbOffsetCurCycle = + (uint64_t)(rescaleOPingPongFlag * HALF_DM_UB_SIZE + rowLoopIdx * ROW_WISE_CYCLE_TILE); + uint32_t oUbOffset = oPingPangFlag * ROW_WISE_CYCLE_TILE * embedRound; + AscendC::WaitFlag(oPingPangFlag); + if (nIdx != NUM4) { + AscendC::DataCopy(loUbTensor[oUbOffset], gInput, + AscendC::DataCopyParams(1, curRowNum * embedRound / FLOAT_BLOCK_SIZE, 0, 0)); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + } + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + AscendC::WaitFlag(oPingPangFlag + 4); + if (nIdx != NUM4) { + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + AscendC::Brcb(tvUbTensor.ReinterpretCast(), + dmUbTensor[dmUbOffsetCurCycle].ReinterpretCast(), curRowNumRound / FLOAT_BLOCK_SIZE, + AscendC::BrcbRepeatParams(1, 8)); + AscendC::PipeBarrier(); + if (needRowLoop) { + AscendC::DataCopy(goUbTensor32[oUbOffset], gUpdate, + AscendC::DataCopyParams(1, curRowNum * embedRound / FLOAT_BLOCK_SIZE, 0, 0)); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + } + // *** go = go * dm_block + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + for (uint32_t vmul_idx = 0; vmul_idx < embed / FLOAT_VECTOR_SIZE; ++vmul_idx) { + AscendC::Mul(goUbTensor32[oUbOffset + vmul_idx * FLOAT_VECTOR_SIZE], + goUbTensor32[oUbOffset + vmul_idx * FLOAT_VECTOR_SIZE], + tvUbTensor, (uint64_t)0, + curRowNum, + AscendC::BinaryRepeatParams(1, 1, 0, embedRound / FLOAT_BLOCK_SIZE, + embedRound / FLOAT_BLOCK_SIZE, 1)); + } + if (embed % FLOAT_VECTOR_SIZE > 0) { + SetMask(embed % FLOAT_VECTOR_SIZE); + AscendC::Mul(goUbTensor32[oUbOffset + embed / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + goUbTensor32[oUbOffset + embed / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + tvUbTensor, + (uint64_t)0, curRowNum, + AscendC::BinaryRepeatParams(1, 1, 0, embedRound / FLOAT_BLOCK_SIZE, + embedRound / FLOAT_BLOCK_SIZE, 1)); + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + AscendC::PipeBarrier(); + // *** go = lo + go + AscendC::Add(goUbTensor32[oUbOffset], goUbTensor32[oUbOffset], + loUbTensor[oUbOffset], (uint64_t)0, + (curRowNum * embedRound + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, + AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); + AscendC::PipeBarrier(); + } else { + // *** gl = ll + AscendC::DataCopy(goUbTensor32[oUbOffset], gInput, + AscendC::DataCopyParams(1, curRowNum * embedRound / FLOAT_BLOCK_SIZE, 0, 0)); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + } + AscendC::SetFlag(oPingPangFlag); + + if (isLastNTile) { + // *** gl_block = expand_to_block(gl), 存放于 tv + AscendC::Brcb(tvUbTensor.ReinterpretCast(), + glUbTensor.ReinterpretCast()[rowLoopIdx * ROW_WISE_CYCLE_TILE], + curRowNumRound / FLOAT_BLOCK_SIZE, + AscendC::BrcbRepeatParams(1, 8)); + AscendC::PipeBarrier(); + // *** go = go / gl_block + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + for (uint32_t vdiv_idx = 0; vdiv_idx < embed / FLOAT_VECTOR_SIZE; ++vdiv_idx) { + AscendC::Div(goUbTensor32[oUbOffset + vdiv_idx * FLOAT_VECTOR_SIZE], + goUbTensor32[oUbOffset + vdiv_idx * FLOAT_VECTOR_SIZE], + tvUbTensor, (uint64_t)0, + curRowNum, + AscendC::BinaryRepeatParams(1, 1, 0, embedRound / FLOAT_BLOCK_SIZE, + embedRound / FLOAT_BLOCK_SIZE, 1)); + } + if (embed % FLOAT_VECTOR_SIZE > 0) { + SetMask(embed % FLOAT_VECTOR_SIZE); + AscendC::Div(goUbTensor32[oUbOffset + embed / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + goUbTensor32[oUbOffset + embed / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + tvUbTensor, + (uint64_t)0, curRowNum, + AscendC::BinaryRepeatParams(1, 1, 0, embedRound / FLOAT_BLOCK_SIZE, + embedRound / FLOAT_BLOCK_SIZE, 1)); + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + AscendC::PipeBarrier(); + + if (kvSplitCoreNum != 1) { + // log(l) + AscendC::Ln( + tvUbTensor, + tvUbTensor, + (uint64_t)0, + curRowNum, + AscendC::UnaryRepeatParams(1, 1, 8, 8)); + AscendC::PipeBarrier(); + AscendC::Brcb( + hmUbTensor.ReinterpretCast(), + gmUbTensor.ReinterpretCast()[rowLoopIdx * ROW_WISE_CYCLE_TILE], + curRowNumRound / FLOAT_BLOCK_SIZE, + AscendC::BrcbRepeatParams(1, 8)); + AscendC::PipeBarrier(); + // logf(lse_sum) + lse_max + AscendC::Add( + tvUbTensor, + tvUbTensor, + hmUbTensor, + (uint64_t)0, + curRowNum, + AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); + AscendC::PipeBarrier(); + + AscendC::SetFlag(EVENT_ID2); + AscendC::WaitFlag(EVENT_ID2); + AscendC::DataCopyPad(gl, tvUbTensor, + AscendC::DataCopyExtParams(curRowNum, 4, 0, (kvSplitCoreNum - 1) * 4, 0)); + + if (glFlag == 0) { + AscendC::SetFlag(EVENT_ID2); + glFlag = 1; + } + uint32_t srcGap = ((embed % 16 <= 8) && (embed % 16 > 0)) ? 1 : 0; + AscendC::DataCopyPad(gOCoreTmp, goUbTensor32[oUbOffset], + AscendC::DataCopyExtParams(curRowNum, embed * 4, srcGap, (kvSplitCoreNum - 1) * embed * 4, 0)); + AscendC::SetFlag(EVENT_ID3); + AscendC::WaitFlag(EVENT_ID3); + AscendC::SetFlag(EVENT_ID1); + AscendC::WaitFlag(EVENT_ID1); + } else { + // *** go = castfp32to16(go) + if (std::is_same::value) { + AscendC::Cast( + goUbTensor16[oUbOffset * 2], goUbTensor32[oUbOffset], + AscendC::RoundMode::CAST_RINT, (uint64_t)0, + (curRowNum * embedRound + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, + AscendC::UnaryRepeatParams(1, 1, 4, 8)); + } else { + AscendC::Cast( + goUbTensor16[oUbOffset * 2], goUbTensor32[oUbOffset], + AscendC::RoundMode::CAST_NONE, (uint64_t)0, + (curRowNum * embedRound + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, + AscendC::UnaryRepeatParams(1, 1, 4, 8)); + } + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + + // ********************* move O to GM ************************ + AscendC::DataCopyPad(gOutput, goUbTensor16[oUbOffset * 2], + AscendC::DataCopyExtParams(curRowNum, embed * 2, 0, 0, 0)); + } + } else if (needRowLoop) { + AscendC::SetFlag(EVENT_ID5); + AscendC::WaitFlag(EVENT_ID5); + AscendC::DataCopy(gUpdate, goUbTensor32[oUbOffset], + AscendC::DataCopyParams(1, curRowNum * embedRound / FLOAT_BLOCK_SIZE, 0, 0)); + } + AscendC::SetFlag(oPingPangFlag + 4); + oPingPangFlag = 1 - oPingPangFlag; + } + + CATLASS_DEVICE + void operator()(AscendC::GlobalTensor gInput, AscendC::GlobalTensor gUpdate, + AscendC::GlobalTensor gOutput, AscendC::GlobalTensor gOCoreTmp, + AscendC::GlobalTensor gl, const LayoutInput &layoutInput, + const LayoutUpdate &layoutUpdate, const LayoutOutput &layoutOutput, GemmCoord actualBlockShape, + uint32_t nIdx, uint32_t isLastNTile, uint32_t rescaleOPingPongFlag, uint32_t &glFlag) + { + uint32_t embed = layoutInput.shape(1); + uint32_t rowActual = actualBlockShape.m(); + uint32_t columnActual = actualBlockShape.n(); + + uint32_t subBlockIdx = AscendC::GetSubBlockIdx(); + uint32_t subBlockNum = AscendC::GetSubBlockNum(); + + uint32_t curRowSplitSubBlock = rowActual / subBlockNum; + uint32_t rowActualThisSubBlock = (subBlockIdx == 0) ? curRowSplitSubBlock : (rowActual - curRowSplitSubBlock); + uint32_t rowOffsetSubBlock = subBlockIdx * curRowSplitSubBlock; + + if (rowActualThisSubBlock > 0) { + uint32_t rowLoop = (rowActualThisSubBlock + ROW_WISE_CYCLE_TILE - 1) / ROW_WISE_CYCLE_TILE; + uint32_t needRowLoop = (rowLoop > 1) ? 1 : 0; + for (uint32_t rowLoopIdx = 0; rowLoopIdx < rowLoop; rowLoopIdx++) { + uint32_t rowOffsetLoop = rowLoopIdx * ROW_WISE_CYCLE_TILE; + uint32_t rowOffsetCurCycle = rowOffsetSubBlock + rowOffsetLoop; + uint32_t rowActualCurCycle = + (rowLoopIdx == (rowLoop - 1)) ? rowActualThisSubBlock - rowOffsetLoop : ROW_WISE_CYCLE_TILE; + int64_t offsetInput = layoutInput.GetOffset(MatrixCoord(rowOffsetCurCycle, 0)); + auto gInputThisCurCycle = gInput[offsetInput]; + auto layoutInputCurCycle = layoutInput.GetTileLayout(MatrixCoord(rowActualCurCycle, columnActual)); + + int64_t offsetUpdate = layoutUpdate.GetOffset(MatrixCoord(rowOffsetCurCycle, 0)); + auto gUpdateCurCycle = gUpdate[offsetUpdate]; + auto layoutUpdateCurCycle = layoutUpdate.GetTileLayout(MatrixCoord(rowActualCurCycle, columnActual)); + + int64_t offsetOutput = layoutOutput.GetOffset(MatrixCoord(rowOffsetCurCycle, 0)); + auto gOutputCurCycle = gOutput[offsetOutput]; + auto layoutOutputCurCycle = layoutOutput.GetTileLayout(MatrixCoord(rowActualCurCycle, columnActual)); + + SubCoreCompute(gInputThisCurCycle, gUpdateCurCycle, gOutputCurCycle, + gOCoreTmp[rowOffsetLoop * embed * kvSplitCoreNum], + gl[rowOffsetLoop * kvSplitCoreNum], layoutInputCurCycle, + layoutOutputCurCycle, layoutUpdateCurCycle, nIdx, isLastNTile, needRowLoop, rowLoopIdx, + rescaleOPingPongFlag, glFlag); + } + } + } + +private: + uint32_t kvSplitCoreNum = 1; + uint32_t oPingPangFlag = 0; + AscendC::LocalTensor loUbTensor; + AscendC::LocalTensor dmUbTensor; + AscendC::LocalTensor glUbTensor; + AscendC::LocalTensor tvUbTensor; + AscendC::LocalTensor goUbTensor16; + AscendC::LocalTensor goUbTensor32; + AscendC::LocalTensor hmUbTensor; + AscendC::LocalTensor gmUbTensor; +}; + +} // namespace Catlass::Epilogue::Block + +#endif // CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_MLA_TP1_RESCALE_O_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_online_softmax_no_mask.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_online_softmax_no_mask.hpp new file mode 100644 index 0000000000..a783da27e3 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_online_softmax_no_mask.hpp @@ -0,0 +1,877 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_ONLINE_SOFTMAX_NO_MASK_HPP +#define CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_ONLINE_SOFTMAX_NO_MASK_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_cross_core_sync.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/epilogue/kda_dispatch_policy.hpp" +#include "catlass/epilogue/tile/kda_epilogue_tile_copy.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/kda_matrix_coord.hpp" + +namespace Catlass::Epilogue::Block { + +template < + class OutputType_, + class InputType_, + class MaskType_> +class BlockEpilogue< + EpilogueAtlasA2OnlineSoftmax, + OutputType_, + InputType_, + MaskType_> { +public: + using DispatchPolicy = EpilogueAtlasA2OnlineSoftmax; + using ArchTag = typename DispatchPolicy::ArchTag; + using ElementOutput = typename OutputType_::Element; + using ElementInput = typename InputType_::Element; + using ElementMask = typename MaskType_::Element; + + using LayoutOutput = typename OutputType_::Layout; + using LayoutInput = typename InputType_::Layout; + using LayoutMask = typename MaskType_::Layout; + + static constexpr uint32_t FLOAT_BLOCK_SIZE = 8; + static constexpr uint32_t FLOAT_VECTOR_SIZE = 64; + static constexpr uint32_t HALF_VECTOR_SIZE = 128; + static constexpr uint32_t BLOCK_SIZE = 16; + static constexpr uint32_t UB_UINT8_VECTOR_SIZE = 1024; + static constexpr uint32_t UB_UINT8_BLOCK_SIZE = 16384; + static constexpr uint32_t VECTOR_SIZE = 128; + static constexpr uint32_t MAX_UB_S_ELEM_NUM = 8192; + + static constexpr uint32_t REDUCE_UB_SIZE = 1024; + static constexpr uint32_t ROW_OPS_SPEC_MASK_32 = 32; + static constexpr uint32_t ROW_OPS_SPEC_MASK_4 = 4; + static constexpr uint32_t MAX_ROW_NUM_SUB_CORE = 128; + static constexpr int64_t UB_FLOAT_LINE_SIZE = 64; + enum class MaskCategory { + NO_MASK = 0, + CAUSAL_MASK = 1 + }; + CATLASS_DEVICE + BlockEpilogue(Arch::Resource &resource, float scaleValue_) + { + // Allocate UB space + constexpr uint32_t LS_UB_TENSOR_OFFSET = 0; + constexpr uint32_t LP_UB_TENSOR_OFFSET = 4 * UB_UINT8_BLOCK_SIZE; + constexpr uint32_t MASK32_UB_TENSOR_OFFSET = 4 * UB_UINT8_BLOCK_SIZE; + + constexpr uint32_t TV_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE; + constexpr uint32_t LM_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 8 * UB_UINT8_VECTOR_SIZE; + + constexpr uint32_t HM_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 9 * UB_UINT8_VECTOR_SIZE; + constexpr uint32_t GM_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 10 * UB_UINT8_VECTOR_SIZE; + constexpr uint32_t LL_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 11 * UB_UINT8_VECTOR_SIZE; + constexpr uint32_t GL_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 12 * UB_UINT8_VECTOR_SIZE; + constexpr uint32_t DM_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 13 * UB_UINT8_VECTOR_SIZE; + + constexpr uint32_t MASK_UB_TENSOR_OFFSET = 11 * UB_UINT8_BLOCK_SIZE; + + scaleValue = scaleValue_; + lsUbTensor = resource.ubBuf.template GetBufferByByte(LS_UB_TENSOR_OFFSET); + lpUbTensor = resource.ubBuf.template GetBufferByByte(LP_UB_TENSOR_OFFSET); + maskUbTensor = resource.ubBuf.template GetBufferByByte(MASK_UB_TENSOR_OFFSET); + maskUbTensor32 = resource.ubBuf.template GetBufferByByte(MASK32_UB_TENSOR_OFFSET); + lmUbTensor = resource.ubBuf.template GetBufferByByte(LM_UB_TENSOR_OFFSET); + hmUbTensor = resource.ubBuf.template GetBufferByByte(HM_UB_TENSOR_OFFSET); + gmUbTensor = resource.ubBuf.template GetBufferByByte(GM_UB_TENSOR_OFFSET); + dmUbTensor = resource.ubBuf.template GetBufferByByte(DM_UB_TENSOR_OFFSET); + llUbTensor = resource.ubBuf.template GetBufferByByte(LL_UB_TENSOR_OFFSET); + tvUbTensor = resource.ubBuf.template GetBufferByByte(TV_UB_TENSOR_OFFSET); + glUbTensor = resource.ubBuf.template GetBufferByByte(GL_UB_TENSOR_OFFSET); + } + + CATLASS_DEVICE + ~BlockEpilogue() {} + + template + CATLASS_DEVICE T Min(T a, T b) + { + return (a > b) ? b : a; + } + + CATLASS_DEVICE + void SetVecMask(int32_t len) + { + uint64_t mask = 0; + uint64_t one = 1; + uint64_t temp = len % FLOAT_VECTOR_SIZE; + for (int64_t i = 0; i < temp; i++) { + mask |= one << i; + } + + if (len == VECTOR_SIZE || len == 0) { + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + } else if (len >= FLOAT_VECTOR_SIZE) { + AscendC::SetVectorMask(mask, (uint64_t)-1); + } else { + AscendC::SetVectorMask(0x0, mask); + } + } + + CATLASS_DEVICE + void SetBlockReduceMask(int32_t len) + { + if (len > 8 || len < 1) { + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + return; + } + uint64_t subMask = ((uint64_t)1 << len) - 1; + uint64_t maskValue = (subMask << 48) + (subMask << 32) + (subMask << 16) + subMask + (subMask << 56) + + (subMask << 40) + (subMask << 24) + (subMask << 8); + AscendC::SetVectorMask(maskValue, maskValue); + } + + CATLASS_DEVICE + void RowsumSPECTILE512(const AscendC::LocalTensor &srcUb, const AscendC::LocalTensor &rowsumUb, + const AscendC::LocalTensor &tvUbTensor, uint32_t numRowsRound, uint32_t numElems, + uint32_t numElemsAligned) + { + AscendC::BlockReduceSum(tvUbTensor, srcUb, numRowsRound * numElemsAligned / FLOAT_VECTOR_SIZE, 0, + 1, 1, 8); + AscendC::PipeBarrier(); + + AscendC::BlockReduceSum(tvUbTensor[REDUCE_UB_SIZE], tvUbTensor, + numRowsRound * numElemsAligned / FLOAT_BLOCK_SIZE / FLOAT_VECTOR_SIZE, 0, + 1, 1, 8); + AscendC::PipeBarrier(); + AscendC::BlockReduceSum(rowsumUb, tvUbTensor[REDUCE_UB_SIZE], + numRowsRound * numElemsAligned / FLOAT_VECTOR_SIZE / FLOAT_VECTOR_SIZE, 0, + 1, 1, 8); + AscendC::PipeBarrier(); + } + + CATLASS_DEVICE + void RowsumSPECTILE256(const AscendC::LocalTensor &srcUb, const AscendC::LocalTensor &rowsumUb, + const AscendC::LocalTensor &tvUbTensor, uint32_t numRowsRound, uint32_t numElems, + uint32_t numElemsAligned) + { + AscendC::BlockReduceSum(tvUbTensor, srcUb, numRowsRound * numElemsAligned / FLOAT_VECTOR_SIZE, 0, + 1, 1, 8); + AscendC::PipeBarrier(); + SetVecMask(ROW_OPS_SPEC_MASK_32); + AscendC::BlockReduceSum(tvUbTensor[REDUCE_UB_SIZE], tvUbTensor, numRowsRound, 0, 1, 1, 4); + AscendC::PipeBarrier(); + SetBlockReduceMask(ROW_OPS_SPEC_MASK_4); + AscendC::BlockReduceSum( + rowsumUb, tvUbTensor[REDUCE_UB_SIZE], + (numRowsRound * FLOAT_BLOCK_SIZE + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, 0, 1, 1, 8); + AscendC::PipeBarrier(); + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + + CATLASS_DEVICE + void RowsumTAILTILE(const AscendC::LocalTensor &srcUb, const AscendC::LocalTensor &rowsumUb, + const AscendC::LocalTensor &tvUbTensor, uint32_t numRowsRound, uint32_t numElems, + uint32_t numElemsAligned) + { + if (numElems >= FLOAT_VECTOR_SIZE) { + AscendC::BlockReduceSum(tvUbTensor, srcUb, numRowsRound, 0, 1, 1, + numElemsAligned / FLOAT_BLOCK_SIZE); + AscendC::PipeBarrier(); + AscendC::BlockReduceSum( + rowsumUb, tvUbTensor, (numRowsRound * FLOAT_BLOCK_SIZE + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, 0, + 1, 1, 8); + AscendC::PipeBarrier(); + for (uint64_t rowSumIdx = 1; rowSumIdx < (uint64_t)numElems / FLOAT_VECTOR_SIZE; ++rowSumIdx) { + AscendC::BlockReduceSum(tvUbTensor, srcUb[rowSumIdx * FLOAT_VECTOR_SIZE], numRowsRound, 0, + 1, 1, numElemsAligned / FLOAT_BLOCK_SIZE); + AscendC::PipeBarrier(); + AscendC::BlockReduceSum( + tvUbTensor[REDUCE_UB_SIZE], tvUbTensor, + (numRowsRound * FLOAT_BLOCK_SIZE + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, 0, 1, 1, 8); + AscendC::PipeBarrier(); + SetVecMask(numRowsRound); + AscendC::Add(rowsumUb, rowsumUb, tvUbTensor[REDUCE_UB_SIZE], (uint64_t)0, 1, + AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); + AscendC::PipeBarrier(); + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + } + if (numElems % FLOAT_VECTOR_SIZE > 0) { + SetVecMask(numElems % FLOAT_VECTOR_SIZE); + AscendC::BlockReduceSum(tvUbTensor, srcUb[numElems / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + numRowsRound, 0, 1, 1, numElemsAligned / FLOAT_BLOCK_SIZE); + AscendC::PipeBarrier(); + SetBlockReduceMask((numElems % FLOAT_VECTOR_SIZE + FLOAT_BLOCK_SIZE - 1) / FLOAT_BLOCK_SIZE); + if (numElems < FLOAT_VECTOR_SIZE) { + AscendC::BlockReduceSum( + rowsumUb, tvUbTensor, (numRowsRound * FLOAT_BLOCK_SIZE + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, + 0, 1, 1, 8); + AscendC::PipeBarrier(); + } else { + AscendC::BlockReduceSum( + tvUbTensor[REDUCE_UB_SIZE], tvUbTensor, + (numRowsRound * FLOAT_BLOCK_SIZE + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, 0, 1, 1, 8); + AscendC::PipeBarrier(); + SetVecMask(numRowsRound); + AscendC::Add(rowsumUb, rowsumUb, tvUbTensor[REDUCE_UB_SIZE], (uint64_t)0, 1, + AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); + AscendC::PipeBarrier(); + } + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + } + + CATLASS_DEVICE + void RowmaxSPECTILE512(const AscendC::LocalTensor &srcUb, const AscendC::LocalTensor &rowmaxUb, + const AscendC::LocalTensor &tvUbTensor, uint32_t numRowsRound, uint32_t numElems, + uint32_t numElemsAligned) + { + AscendC::BlockReduceMax(tvUbTensor, srcUb, numRowsRound * numElemsAligned / FLOAT_VECTOR_SIZE, 0, + 1, 1, 8); + AscendC::PipeBarrier(); + AscendC::BlockReduceMax(tvUbTensor[REDUCE_UB_SIZE], tvUbTensor, + numRowsRound * numElemsAligned / FLOAT_BLOCK_SIZE / FLOAT_VECTOR_SIZE, 0, + 1, 1, 8); + AscendC::PipeBarrier(); + AscendC::BlockReduceMax(rowmaxUb, tvUbTensor[REDUCE_UB_SIZE], + numRowsRound * numElemsAligned / FLOAT_VECTOR_SIZE / FLOAT_VECTOR_SIZE, 0, + 1, 1, 8); + AscendC::PipeBarrier(); + } + + CATLASS_DEVICE + void RowmaxSPECTILE256(const AscendC::LocalTensor &srcUb, const AscendC::LocalTensor &rowmaxUb, + const AscendC::LocalTensor &tvUbTensor, uint32_t numRowsRound, uint32_t numElems, + uint32_t numElemsAligned) + { + AscendC::BlockReduceMax(tvUbTensor, srcUb, numRowsRound * numElemsAligned / FLOAT_VECTOR_SIZE, 0, + 1, 1, 8); + AscendC::PipeBarrier(); + SetVecMask(ROW_OPS_SPEC_MASK_32); + AscendC::BlockReduceMax(tvUbTensor[REDUCE_UB_SIZE], tvUbTensor, numRowsRound, 0, 1, 1, 4); + AscendC::PipeBarrier(); + SetBlockReduceMask(ROW_OPS_SPEC_MASK_4); + AscendC::BlockReduceMax( + rowmaxUb, tvUbTensor[REDUCE_UB_SIZE], + (numRowsRound * FLOAT_BLOCK_SIZE + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, 0, 1, 1, 8); + AscendC::PipeBarrier(); + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + + CATLASS_DEVICE + void RowmaxTAILTILE(const AscendC::LocalTensor &srcUb, const AscendC::LocalTensor &rowmaxUb, + const AscendC::LocalTensor &tvUbTensor, uint32_t numRowsRound, uint32_t numElems, + uint32_t numElemsAligned) + { + if (numElems >= FLOAT_VECTOR_SIZE) { + AscendC::BlockReduceMax(tvUbTensor, srcUb, numRowsRound, 0, 1, 1, + numElemsAligned / FLOAT_BLOCK_SIZE); + AscendC::PipeBarrier(); + AscendC::BlockReduceMax( + rowmaxUb, tvUbTensor, (numRowsRound * FLOAT_BLOCK_SIZE + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, 0, + 1, 1, 8); + AscendC::PipeBarrier(); + for (uint64_t rowmax_idx = 1; rowmax_idx < (uint64_t)numElems / FLOAT_VECTOR_SIZE; ++rowmax_idx) { + AscendC::BlockReduceMax(tvUbTensor, srcUb[rowmax_idx * FLOAT_VECTOR_SIZE], numRowsRound, + 0, 1, 1, numElemsAligned / FLOAT_BLOCK_SIZE); + AscendC::PipeBarrier(); + AscendC::BlockReduceMax( + tvUbTensor[REDUCE_UB_SIZE], tvUbTensor, + (numRowsRound * FLOAT_BLOCK_SIZE + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, 0, 1, 1, 8); + AscendC::PipeBarrier(); + SetVecMask(numRowsRound); + AscendC::Max(rowmaxUb, rowmaxUb, tvUbTensor[REDUCE_UB_SIZE], (uint64_t)0, 1, + AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); + AscendC::PipeBarrier(); + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + } + if (numElems % FLOAT_VECTOR_SIZE > 0) { + SetVecMask(numElems % FLOAT_VECTOR_SIZE); + AscendC::BlockReduceMax(tvUbTensor, srcUb[numElems / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + numRowsRound, 0, 1, 1, numElemsAligned / FLOAT_BLOCK_SIZE); + AscendC::PipeBarrier(); + SetBlockReduceMask((numElems % FLOAT_VECTOR_SIZE + FLOAT_BLOCK_SIZE - 1) / FLOAT_BLOCK_SIZE); + if (numElems < FLOAT_VECTOR_SIZE) { + AscendC::BlockReduceMax( + rowmaxUb, tvUbTensor, (numRowsRound * FLOAT_BLOCK_SIZE + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, + 0, 1, 1, 8); + AscendC::PipeBarrier(); + } else { + AscendC::BlockReduceMax( + tvUbTensor[REDUCE_UB_SIZE], tvUbTensor, + (numRowsRound * FLOAT_BLOCK_SIZE + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, 0, 1, 1, 8); + AscendC::PipeBarrier(); + SetVecMask(numRowsRound); + AscendC::Max(rowmaxUb, rowmaxUb, tvUbTensor[REDUCE_UB_SIZE], (uint64_t)0, 1, + AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); + AscendC::PipeBarrier(); + } + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + } + + CATLASS_DEVICE + void CopySGmToUb( + AscendC::GlobalTensor gInput, + uint32_t sUbOffset, + uint32_t rowNumCurLoop, + uint32_t columnNumRound, + uint32_t columnNumPad) + { + // input S + AscendC::DataCopy( + lsUbTensor[sUbOffset], gInput, + AscendC::DataCopyParams( + rowNumCurLoop, columnNumRound / FLOAT_BLOCK_SIZE, + (columnNumPad - columnNumRound) / FLOAT_BLOCK_SIZE, 0)); + } + + CATLASS_DEVICE + void CopyMaskGmToUb( + AscendC::GlobalTensor gMask, + uint32_t columnNum, + uint32_t columnNumRound, + uint32_t maskStride, + uint32_t qSBlockSize, + uint32_t proTokenIdx, + uint32_t proTokenNum, + uint32_t integralHeadNum, + uint32_t epiTokenNum) + { + uint32_t innerUbRowOffset = 0; + if (proTokenNum != 0) { + AscendC::DataCopyPad( + maskUbTensor[innerUbRowOffset], gMask[proTokenIdx * maskStride], + AscendC::DataCopyExtParams( + proTokenNum, columnNum * 2, (maskStride - columnNum) * 2, 0, 0), + AscendC::DataCopyPadExtParams(false, 0, 0, 0)); + innerUbRowOffset += proTokenNum * columnNumRound; + } + for (uint32_t headIdx = 0; headIdx < integralHeadNum; headIdx++) { + AscendC::DataCopyPad( + maskUbTensor[innerUbRowOffset], gMask, + AscendC::DataCopyExtParams( + qSBlockSize, columnNum * 2, (maskStride - columnNum) * 2, 0, 0), + AscendC::DataCopyPadExtParams(false, 0, 0, 0)); + innerUbRowOffset += qSBlockSize * columnNumRound; + } + if (epiTokenNum != 0) { + AscendC::DataCopyPad( + maskUbTensor[innerUbRowOffset], gMask, + AscendC::DataCopyExtParams( + epiTokenNum, columnNum * 2, (maskStride - columnNum) * 2, 0, 0), + AscendC::DataCopyPadExtParams(false, 0, 0, 0)); + } + } + + CATLASS_DEVICE + void ScaleS( + uint32_t sUbOffset, + uint32_t rowNumCurLoop, + uint32_t columnNumRound) + { + // *** ls = scaleValue * ls + AscendC::Muls( + lsUbTensor[sUbOffset], lsUbTensor[sUbOffset], scaleValue, (uint64_t)0, + (rowNumCurLoop * columnNumRound + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, + AscendC::UnaryRepeatParams(1, 1, 8, 8)); + + AscendC::PipeBarrier(); + } + + CATLASS_DEVICE + void UpCastMask( + uint32_t rowNumCurLoop, + uint32_t columnNumRound) + { + AscendC::Cast( + maskUbTensor32, maskUbTensor, AscendC::RoundMode::CAST_NONE, (uint64_t)0, + (rowNumCurLoop * columnNumRound + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, + AscendC::UnaryRepeatParams(1, 1, 8, 4)); + AscendC::PipeBarrier(); + } + + CATLASS_DEVICE + void ApplyMask( + uint32_t sUbOffset, + uint32_t rowNumCurLoop, + uint32_t columnNumRound) + { + AscendC::Muls( + maskUbTensor32, maskUbTensor32, (float)-3e38, (uint64_t)0, + (rowNumCurLoop * columnNumRound + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, + AscendC::UnaryRepeatParams(1, 1, 8, 8)); + AscendC::PipeBarrier(); + AscendC::Add( + lsUbTensor[sUbOffset], lsUbTensor[sUbOffset], maskUbTensor32, + (uint64_t)0, (rowNumCurLoop * columnNumRound + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, + AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); + AscendC::PipeBarrier(); + } + + CATLASS_DEVICE + void CalcLocalRowMax( + uint32_t sUbOffset, + uint32_t rowNumCurLoopRound, + uint32_t columnNum, + uint32_t columnNumRound, + uint32_t rowOffset) + { + if (columnNum == 512) { + RowmaxSPECTILE512( + lsUbTensor[sUbOffset], lmUbTensor[rowOffset], + tvUbTensor, rowNumCurLoopRound, columnNum, columnNumRound); + } else if (columnNum == 256) { + RowmaxSPECTILE256( + lsUbTensor[sUbOffset], lmUbTensor[rowOffset], + tvUbTensor, rowNumCurLoopRound, columnNum, columnNumRound); + } else { + RowmaxTAILTILE( + lsUbTensor[sUbOffset], lmUbTensor[rowOffset], + tvUbTensor, rowNumCurLoopRound, columnNum, columnNumRound); + } + } + + CATLASS_DEVICE + void UpdateGlobalRowMax( + uint32_t rowNumCurLoop, + uint32_t rowNumCurLoopRound, + uint32_t columnNum, + uint32_t columnNumRound, + uint32_t dmUbOffsetCurCycle, + uint32_t rowOffset, + uint32_t isFirstStackTile) + { + if (isFirstStackTile) { + AscendC::DataCopy( + hmUbTensor[rowOffset], lmUbTensor[rowOffset], + AscendC::DataCopyParams(1, rowNumCurLoopRound / FLOAT_BLOCK_SIZE, 0, 0)); + AscendC::PipeBarrier(); + } else { + SetVecMask(rowNumCurLoop); + // *** hm = vmax(lm, gm) + AscendC::Max( + hmUbTensor[rowOffset], lmUbTensor[rowOffset], gmUbTensor[rowOffset], (uint64_t)0, + 1, AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); + + AscendC::PipeBarrier(); + // *** dm = gm - hm + AscendC::Sub( + dmUbTensor[dmUbOffsetCurCycle], + gmUbTensor[rowOffset], hmUbTensor[rowOffset], (uint64_t)0, 1, + AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); + + AscendC::PipeBarrier(); + // *** dm = exp(dm) + AscendC::Exp( + dmUbTensor[dmUbOffsetCurCycle], + dmUbTensor[dmUbOffsetCurCycle], + (uint64_t)0, 1, AscendC::UnaryRepeatParams(1, 1, 8, 8)); + } + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + AscendC::PipeBarrier(); + // *** gm = hm + AscendC::DataCopy(gmUbTensor[rowOffset], hmUbTensor[rowOffset], + AscendC::DataCopyParams(1, rowNumCurLoopRound / FLOAT_BLOCK_SIZE, 0, 0)); + AscendC::PipeBarrier(); + } + + CATLASS_DEVICE + void CalcExp( + uint32_t sUbOffset, + uint32_t rowNumCurLoop, + uint32_t rowNumCurLoopRound, + uint32_t columnNum, + uint32_t columnNumRound, + uint32_t rowOffset) + { + // *** hm_block = expand_to_block(hm), 存放于 tv + AscendC::Brcb( + tvUbTensor.template ReinterpretCast(), + hmUbTensor[rowOffset].template ReinterpretCast(), rowNumCurLoopRound / FLOAT_BLOCK_SIZE, + AscendC::BrcbRepeatParams(1, 8)); + AscendC::PipeBarrier(); + // *** ls = ls - hm_block + for (uint32_t subIdx = 0; subIdx < columnNum / FLOAT_VECTOR_SIZE; ++subIdx) { + AscendC::Sub( + lsUbTensor[sUbOffset][subIdx * FLOAT_VECTOR_SIZE], lsUbTensor[sUbOffset][subIdx * FLOAT_VECTOR_SIZE], + tvUbTensor, (uint64_t)0, rowNumCurLoop, + AscendC::BinaryRepeatParams(1, 1, 0, columnNumRound / FLOAT_BLOCK_SIZE, columnNumRound / FLOAT_BLOCK_SIZE, 1)); + } + if (columnNum % FLOAT_VECTOR_SIZE > 0) { + SetVecMask(columnNum % FLOAT_VECTOR_SIZE); + AscendC::Sub( + lsUbTensor[sUbOffset][columnNum / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + lsUbTensor[sUbOffset][columnNum / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], tvUbTensor, (uint64_t)0, rowNumCurLoop, + AscendC::BinaryRepeatParams(1, 1, 0, columnNumRound / FLOAT_BLOCK_SIZE, columnNumRound / FLOAT_BLOCK_SIZE, 1)); + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + AscendC::PipeBarrier(); + // *** ls = exp(ls) + AscendC::Exp(lsUbTensor[sUbOffset], lsUbTensor[sUbOffset], (uint64_t)0, + (rowNumCurLoop * columnNumRound + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, + AscendC::UnaryRepeatParams(1, 1, 8, 8)); + AscendC::PipeBarrier(); + } + + CATLASS_DEVICE + void CalcLocalRowSum( + uint32_t sUbOffset, + uint32_t rowNumCurLoopRound, + uint32_t columnNum, + uint32_t columnNumRound, + uint32_t rowOffset) + { + // *** ll = rowsum(ls32) + if (columnNum == 512) { + RowsumSPECTILE512( + lsUbTensor[sUbOffset], llUbTensor[rowOffset], tvUbTensor, + rowNumCurLoopRound, columnNum, columnNumRound); + } else if (columnNum == 256) { + RowsumSPECTILE256( + lsUbTensor[sUbOffset], llUbTensor[rowOffset], tvUbTensor, + rowNumCurLoopRound, columnNum, columnNumRound); + } else { + RowsumTAILTILE( + lsUbTensor[sUbOffset], llUbTensor[rowOffset], tvUbTensor, + rowNumCurLoopRound, columnNum, columnNumRound); + } + } + + CATLASS_DEVICE + void UpdateGlobalRowSum( + uint32_t sUbOffset, + uint32_t rowNumCurLoop, + uint32_t rowNumCurLoopRound, + uint32_t dmUbOffsetCurCycle, + uint32_t rowOffset, + uint32_t isFirstStackTile) + { + if (isFirstStackTile) { + // *** gl = ll + AscendC::DataCopy( + glUbTensor[rowOffset], llUbTensor[rowOffset], + AscendC::DataCopyParams(1, rowNumCurLoopRound / FLOAT_BLOCK_SIZE, 0, 0)); + AscendC::PipeBarrier(); + } else { + SetVecMask(rowNumCurLoop); + // *** gl = dm * gl + AscendC::Mul( + glUbTensor[rowOffset], dmUbTensor[dmUbOffsetCurCycle], + glUbTensor[rowOffset], (uint64_t)0, 1, AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); + AscendC::PipeBarrier(); + // *** gl = ll + gl + AscendC::Add(glUbTensor[rowOffset], glUbTensor[rowOffset], llUbTensor[rowOffset], (uint64_t)0, + 1, AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); + AscendC::PipeBarrier(); + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + } + + CATLASS_DEVICE + void DownCastP( + uint32_t sUbOffset, + uint32_t rowNumCurLoop, + uint32_t columnNumRound) + { + // *** lp = castfp32to16(ls) + if (std::is_same::value) { + AscendC::Cast( + lpUbTensor[sUbOffset], lsUbTensor[sUbOffset], AscendC::RoundMode::CAST_RINT, (uint64_t)0, + (rowNumCurLoop * columnNumRound + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, AscendC::UnaryRepeatParams(1, 1, 4, 8)); + } else { + AscendC::Cast( + lpUbTensor[sUbOffset], lsUbTensor[sUbOffset], AscendC::RoundMode::CAST_NONE, (uint64_t)0, + (rowNumCurLoop * columnNumRound + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, AscendC::UnaryRepeatParams(1, 1, 4, 8)); + } + } + + CATLASS_DEVICE + void CopyPUbToGm( + AscendC::GlobalTensor gOutput, + uint32_t sUbOffset, + uint32_t rowNumCurLoop, + uint32_t columnNumRound, + uint32_t columnNumPad) + { + if (columnNumRound == columnNumPad) { + AscendC::DataCopy( + gOutput, lpUbTensor[sUbOffset], + AscendC::DataCopyParams( + 1, rowNumCurLoop * columnNumRound / BLOCK_SIZE, + 0, 0)); + } else { + AscendC::DataCopy( + gOutput, lpUbTensor[sUbOffset], + AscendC::DataCopyParams( + rowNumCurLoop, columnNumRound / BLOCK_SIZE, + 0, (columnNumPad - columnNumRound) / BLOCK_SIZE)); + } + } + + template + CATLASS_DEVICE void SubCoreCompute(AscendC::GlobalTensor gOutput, + const LayoutOutput &layoutOutput, + uint32_t rowOffset, uint32_t isFirstStackTile, + uint32_t columnNumRound, uint32_t pingpongFlag, + uint32_t curStackTileMod) + { + uint32_t rowNumCurLoop = layoutOutput.shape(0); + uint32_t rowNumCurLoopRound = RoundUp(rowNumCurLoop, FLOAT_BLOCK_SIZE); + uint32_t columnNum = layoutOutput.shape(1); + // Align colNum to 16, for both float&half compute© + uint32_t columnNumPad = layoutOutput.stride(0); + uint32_t sUbOffset = pingpongFlag * MAX_UB_S_ELEM_NUM; + uint32_t dmUbOffsetCurCycle = curStackTileMod * MAX_ROW_NUM_SUB_CORE + rowOffset; + + CalcLocalRowMax(sUbOffset, rowNumCurLoopRound, columnNum, columnNumRound, rowOffset); + UpdateGlobalRowMax(rowNumCurLoop, rowNumCurLoopRound, columnNum, columnNumRound, dmUbOffsetCurCycle, rowOffset, isFirstStackTile); + + CalcExp(sUbOffset, rowNumCurLoop, rowNumCurLoopRound, columnNum, columnNumRound, rowOffset); + if constexpr (maskCat == MaskCategory::NO_MASK) { + AscendC::WaitFlag(pingpongFlag); + } + + DownCastP(sUbOffset, rowNumCurLoop, columnNumRound); + AscendC::SetFlag(pingpongFlag); + + CalcLocalRowSum(sUbOffset, rowNumCurLoopRound, columnNum, columnNumRound, rowOffset); + AscendC::SetFlag(pingpongFlag); + + AscendC::WaitFlag(pingpongFlag); + CopyPUbToGm(gOutput, sUbOffset, rowNumCurLoop, columnNumRound, columnNumPad); + if constexpr (maskCat == MaskCategory::NO_MASK) { + AscendC::SetFlag(pingpongFlag); + } else if constexpr (maskCat == MaskCategory::CAUSAL_MASK) { + AscendC::SetFlag(EVENT_ID0); + AscendC::SetFlag(EVENT_ID1); + } + UpdateGlobalRowSum(sUbOffset, rowNumCurLoop, rowNumCurLoopRound, dmUbOffsetCurCycle, rowOffset, isFirstStackTile); + } + + CATLASS_DEVICE + void operator()(AscendC::GlobalTensor gOutput, AscendC::GlobalTensor gInput, + const LayoutOutput &layoutOutput, const LayoutInput &layoutInput, GemmCoord actualBlockShape, + uint32_t isFirstStackTile, uint32_t qSBlockSize, uint32_t qNBlockSize, uint32_t curStackTileMod) + { + uint32_t rowNum = actualBlockShape.m(); + uint32_t columnNum = actualBlockShape.n(); + uint32_t columnNumRound = RoundUp(columnNum, BLOCK_SIZE); + uint32_t columnNumPad = layoutInput.stride(0); + + uint32_t subBlockIdx = AscendC::GetSubBlockIdx(); + uint32_t subBlockNum = AscendC::GetSubBlockNum(); + + uint32_t qNSplitSubBlock = qNBlockSize / subBlockNum; + uint32_t qNThisSubBlock = (qNBlockSize == 1) ? + 0 : + (subBlockIdx == 1) ? (qNBlockSize - qNSplitSubBlock) : + qNSplitSubBlock; + uint32_t rowSplitSubBlock = (qNBlockSize == 1) ? + (qSBlockSize / 2) : + (qSBlockSize * qNSplitSubBlock); + uint32_t rowActualThisSubBlock = (subBlockIdx == 1) ? + (rowNum - rowSplitSubBlock) : + rowSplitSubBlock; + uint32_t rowOffsetThisSubBlock = subBlockIdx * rowSplitSubBlock; + uint32_t maxRowNumPerLoop = MAX_UB_S_ELEM_NUM / columnNumRound; + uint32_t rowNumTile = RoundDown(maxRowNumPerLoop, FLOAT_BLOCK_SIZE); + uint32_t rowLoopNum = CeilDiv(rowActualThisSubBlock, rowNumTile); + uint32_t preLoad = 1; + + for (uint32_t rowLoopIdx = 0; rowLoopIdx < rowLoopNum + preLoad; rowLoopIdx++) { + if (rowLoopIdx < rowLoopNum) { + uint32_t pingpongFlag = rowLoopIdx % 2; + uint32_t rowOffsetCurLoop = rowLoopIdx * rowNumTile; + uint32_t rowOffsetIoGm = rowOffsetCurLoop + rowOffsetThisSubBlock; + uint32_t rowNumCurLoop = (rowLoopIdx == rowLoopNum - 1) ? + (rowActualThisSubBlock - rowOffsetCurLoop) : + rowNumTile; + // loop 0 mask load before cross core sync + + int64_t offsetInput = layoutInput.GetOffset(MatrixCoord(rowOffsetIoGm, 0)); + auto gInputCurLoop = gInput[offsetInput]; + + AscendC::WaitFlag(pingpongFlag); + CopySGmToUb(gInputCurLoop, (pingpongFlag * MAX_UB_S_ELEM_NUM), rowNumCurLoop, columnNumRound, columnNumPad); + AscendC::SetFlag(pingpongFlag); + } + if (rowLoopIdx >= preLoad) { + uint32_t delayedRowLoopIdx = rowLoopIdx - preLoad; + uint32_t pingpongFlag = delayedRowLoopIdx % 2; + uint32_t rowOffsetCurLoop = delayedRowLoopIdx * rowNumTile; + uint32_t rowOffsetIoGm = rowOffsetCurLoop + rowOffsetThisSubBlock; + uint32_t rowNumCurLoop = (delayedRowLoopIdx == rowLoopNum - 1) ? + (rowActualThisSubBlock - rowOffsetCurLoop) : + rowNumTile; + + int64_t offsetOutput = layoutOutput.GetOffset(MatrixCoord(rowOffsetIoGm, 0)); + auto gOutputCurLoop = gOutput[offsetOutput]; + auto layoutOutputCurLoop = layoutOutput.GetTileLayout(MatrixCoord(rowNumCurLoop, columnNum)); + AscendC::WaitFlag(pingpongFlag); + ScaleS((pingpongFlag * MAX_UB_S_ELEM_NUM), rowNumCurLoop, columnNumRound); + SubCoreCompute( + gOutputCurLoop, + layoutOutputCurLoop, + rowOffsetCurLoop, + isFirstStackTile, + columnNumRound, + pingpongFlag, + curStackTileMod); + } + } + } + + CATLASS_DEVICE + void operator()(AscendC::GlobalTensor gOutput, AscendC::GlobalTensor gInput, + AscendC::GlobalTensor gMask, + const LayoutOutput &layoutOutput, const LayoutInput &layoutInput, const LayoutInput &layoutMask, + GemmCoord actualBlockShape, + uint32_t isFirstStackTile, uint32_t qSBlockSize, uint32_t qNBlockSize, uint32_t curStackTileMod, + Arch::CrossCoreFlag qkReady) + { + uint32_t rowNum = actualBlockShape.m(); + uint32_t columnNum = actualBlockShape.n(); + uint32_t columnNumRound = RoundUp(columnNum, BLOCK_SIZE); + uint32_t columnNumPad = layoutInput.stride(0); + uint32_t maskStride = layoutMask.stride(0); + uint32_t subBlockIdx = AscendC::GetSubBlockIdx(); + uint32_t subBlockNum = AscendC::GetSubBlockNum(); + + uint32_t qNSplitSubBlock = qNBlockSize / subBlockNum; + uint32_t qNThisSubBlock = (qNBlockSize == 1) ? + 0 : + (subBlockIdx == 1) ? (qNBlockSize - qNSplitSubBlock) : + qNSplitSubBlock; + uint32_t rowSplitSubBlock = (qNBlockSize == 1) ? + (qSBlockSize / 2) : + (qSBlockSize * qNSplitSubBlock); + uint32_t rowActualThisSubBlock = (subBlockIdx == 1) ? + (rowNum - rowSplitSubBlock) : + rowSplitSubBlock; + uint32_t rowOffsetThisSubBlock = subBlockIdx * rowSplitSubBlock; + + uint32_t tokenNumPerHeadThisSubBlock = Min(qSBlockSize, rowActualThisSubBlock); + + uint32_t maskOffsetThisSubBlock = (qNBlockSize == 1) ? + rowOffsetThisSubBlock : + 0; + int64_t offsetMask = layoutMask.GetOffset(MatrixCoord(maskOffsetThisSubBlock, 0)); + auto gMaskThisSubBlock = gMask[offsetMask]; + auto layoutMaskThisSubBlock = layoutMask; + + uint32_t maxRowNumPerLoop = MAX_UB_S_ELEM_NUM / columnNumRound; + uint32_t rowNumTile = RoundDown(maxRowNumPerLoop, FLOAT_BLOCK_SIZE); + uint32_t rowLoopNum = CeilDiv(rowActualThisSubBlock, rowNumTile); + uint32_t preLoad = 1; + + if (rowActualThisSubBlock == 0) { + Arch::CrossCoreWaitFlag(qkReady); + return; + } + + for (uint32_t rowLoopIdx = 0; rowLoopIdx < rowLoopNum + preLoad; rowLoopIdx++) { + if (rowLoopIdx < rowLoopNum) { + uint32_t pingpongFlag = rowLoopIdx % 2; + uint32_t rowOffsetCurLoop = rowLoopIdx * rowNumTile; + uint32_t rowOffsetIoGm = rowOffsetCurLoop + rowOffsetThisSubBlock; + uint32_t rowNumCurLoop = (rowLoopIdx == rowLoopNum - 1) ? + (rowActualThisSubBlock - rowOffsetCurLoop) : + rowNumTile; + // loop 0 mask load before cross core sync + if (rowLoopIdx == 0) { + // the token idx of the start token of the prologue part + uint32_t proTokenIdx = rowOffsetCurLoop % tokenNumPerHeadThisSubBlock; + // the token num of the prologue part + uint32_t proTokenNum = Min(rowNumCurLoop, (tokenNumPerHeadThisSubBlock - proTokenIdx)) % + tokenNumPerHeadThisSubBlock; + // the token num of the epilogue part + uint32_t integralHeadNum = (rowNumCurLoop - proTokenNum) / tokenNumPerHeadThisSubBlock; + // the number of integral heads within a cycle + uint32_t epiTokenNum = rowNumCurLoop - proTokenNum - integralHeadNum * tokenNumPerHeadThisSubBlock; + AscendC::WaitFlag(EVENT_ID2); + CopyMaskGmToUb(gMaskThisSubBlock, columnNum, columnNumRound, maskStride, + tokenNumPerHeadThisSubBlock, proTokenIdx, proTokenNum, integralHeadNum, epiTokenNum); + AscendC::SetFlag(EVENT_ID2); + Arch::CrossCoreWaitFlag(qkReady); + } + int64_t offsetInput = layoutInput.GetOffset(MatrixCoord(rowOffsetIoGm, 0)); + auto gInputCurLoop = gInput[offsetInput]; + AscendC::WaitFlag(pingpongFlag); + CopySGmToUb(gInputCurLoop, (pingpongFlag * MAX_UB_S_ELEM_NUM), rowNumCurLoop, columnNumRound, columnNumPad); + AscendC::SetFlag(pingpongFlag); + } + if (rowLoopIdx >= preLoad) { + uint32_t delayedRowLoopIdx = rowLoopIdx - preLoad; + uint32_t pingpongFlag = delayedRowLoopIdx % 2; + uint32_t rowOffsetCurLoop = delayedRowLoopIdx * rowNumTile; + uint32_t rowNumCurLoop = (delayedRowLoopIdx == rowLoopNum - 1) ? + (rowActualThisSubBlock - rowOffsetCurLoop) : + rowNumTile; + + AscendC::WaitFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID1); + AscendC::WaitFlag(EVENT_ID2); + UpCastMask(rowNumCurLoop, columnNumRound); + AscendC::WaitFlag(pingpongFlag); + ScaleS((pingpongFlag * MAX_UB_S_ELEM_NUM), rowNumCurLoop, columnNumRound); + ApplyMask((pingpongFlag * MAX_UB_S_ELEM_NUM), rowNumCurLoop, columnNumRound); + AscendC::SetFlag(EVENT_ID2); + // next loop mask load + if (rowLoopIdx < rowLoopNum) { + uint32_t rowOffsetCurLoop = rowLoopIdx * rowNumTile; + uint32_t rowNumCurLoop = (rowLoopIdx == rowLoopNum - 1) ? + (rowActualThisSubBlock - rowOffsetCurLoop) : + rowNumTile; + // the token idx of the start token of the prologue part + uint32_t proTokenIdx = rowOffsetCurLoop % tokenNumPerHeadThisSubBlock; + // the token num of the prologue part + uint32_t proTokenNum = Min(rowNumCurLoop, (tokenNumPerHeadThisSubBlock - proTokenIdx)) % + tokenNumPerHeadThisSubBlock; + // the number of integral heads within a cycle + uint32_t integralHeadNum = (rowNumCurLoop - proTokenNum) / tokenNumPerHeadThisSubBlock; + // the token num of the epilogue part + uint32_t epiTokenNum = rowNumCurLoop - proTokenNum - integralHeadNum * tokenNumPerHeadThisSubBlock; + AscendC::WaitFlag(EVENT_ID2); + CopyMaskGmToUb(gMaskThisSubBlock, columnNum, columnNumRound, maskStride, + tokenNumPerHeadThisSubBlock, proTokenIdx, proTokenNum, integralHeadNum, epiTokenNum); + AscendC::SetFlag(EVENT_ID2); + } + // online softmax vectorized compute + uint32_t rowOffsetIoGm = rowOffsetCurLoop + rowOffsetThisSubBlock; + int64_t offsetOutput = layoutOutput.GetOffset(MatrixCoord(rowOffsetIoGm, 0)); + auto gOutputCurLoop = gOutput[offsetOutput]; + auto layoutOutputCurLoop = layoutOutput.GetTileLayout(MatrixCoord(rowNumCurLoop, columnNum)); + SubCoreCompute( + gOutputCurLoop, + layoutOutputCurLoop, + rowOffsetCurLoop, + isFirstStackTile, + columnNumRound, + pingpongFlag, + curStackTileMod); + } + } + } + +private: + float scaleValue; + AscendC::LocalTensor lsUbTensor; + AscendC::LocalTensor lpUbTensor; + AscendC::LocalTensor maskUbTensor; + AscendC::LocalTensor maskUbTensor32; + AscendC::LocalTensor lmUbTensor; + AscendC::LocalTensor hmUbTensor; + AscendC::LocalTensor gmUbTensor; + AscendC::LocalTensor dmUbTensor; + AscendC::LocalTensor llUbTensor; + AscendC::LocalTensor tvUbTensor; + AscendC::LocalTensor glUbTensor; +}; +} // namespace Catlass::Epilogue::Block + +#endif // CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_ONLINE_SOFTMAX_NO_MASK_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_per_group_per_block.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_per_group_per_block.hpp new file mode 100644 index 0000000000..e84dcbd790 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_per_group_per_block.hpp @@ -0,0 +1,863 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_EPILOGUE_BLOCK_EPILOGUE_PER_GROUP_PER_BLOCK_HPP +#define CATLASS_EPILOGUE_BLOCK_EPILOGUE_PER_GROUP_PER_BLOCK_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/epilogue/kda_dispatch_policy.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/kda_matrix_coord.hpp" +#include "tla/kda_tensor.hpp" +#include "tla/kda_tla_layout.hpp" + +namespace Catlass::Epilogue::Block { + +#define QBMM_BLOCK_EPILOGUE_PERTILE_CLASS_LOCAL_PARAMS \ + template < \ + class L0TileShape_, class DataTypeOut_, class DataTypeIn_, class DataTypeBias_, \ + class DataTypeX1Scale_, class DataTypeX2Scale_, class LayoutX1Scale_, class LayoutX2Scale_> +#define QBMM_BLOCK_EPILOGUE_PERTILE_FUNC_LOCAL_PARAMS \ + BlockEpiloguePertile, L0TileShape_, DataTypeOut_, DataTypeIn_, DataTypeBias_, DataTypeX1Scale_, DataTypeX2Scale_, LayoutX1Scale_, \ + LayoutX2Scale_ + +template < + class L0TileShape_, + class DataTypeOut_, + class DataTypeIn_, + class DataTypeBias_, + class DataTypeX1Scale_, + class DataTypeX2Scale_, + class LayoutX1Scale_, + class LayoutX2Scale_> +class BlockEpilogue< + BlockEpiloguePertile, + L0TileShape_, + DataTypeOut_, + DataTypeIn_, + DataTypeBias_, + DataTypeX1Scale_, + DataTypeX2Scale_, + LayoutX1Scale_, + LayoutX2Scale_> { +public: + using DispatchPolicy = BlockEpiloguePertile; + using ArchTag = typename DispatchPolicy::ArchTag; + CATLASS_DEVICE BlockEpilogue() + { + if ASCEND_IS_AIV { + AscendC::SetFlag(0); + AscendC::SetFlag(1); + if constexpr (!AscendC::IsSameType::value) { + AscendC::SetFlag(0); + AscendC::SetFlag(1); + } + } + } + + CATLASS_DEVICE ~BlockEpilogue() + { + AscendC::WaitFlag(0); + AscendC::WaitFlag(1); + if ASCEND_IS_AIV { + if constexpr (!AscendC::IsSameType::value) { + AscendC::WaitFlag(0); + AscendC::WaitFlag(1); + } + if (isBias_) { + AscendC::WaitFlag(biasPingPongID_); + AscendC::WaitFlag(biasPingPongID_ ^ 1); + } + } + } + + struct Arguments { + GM_ADDR outGmAddr{nullptr}; + GM_ADDR x2ScaleGmAddr{nullptr}; + GM_ADDR x1ScaleGmAddr{nullptr}; + GM_ADDR biasGmAddr{nullptr}; + uint32_t baseM; + uint32_t baseN; + uint32_t baseK; + uint32_t groupSizeM = 1U; + uint32_t groupSizeN = 128U; + uint32_t groupSizeK = 128U; + uint32_t isBias = 0; + Arguments() = default; + }; + + // params + using Params = Arguments; + using YType = DataTypeOut_; + using CType = DataTypeIn_; + using BiasType = float; + using X2ScaleType = DataTypeX2Scale_; + using X1ScaleType = DataTypeX1Scale_; + using LayoutX1Scale = LayoutX1Scale_; + using LayoutX2Scale = LayoutX2Scale_; + using CalcType = float; + using BlockCoord = tla::Coord; + + constexpr static uint32_t Y_IDX = 0; + constexpr static uint32_t X2SCALE_IDX = 1; + constexpr static uint32_t X1SCALE_IDX = 2; + constexpr static uint32_t BIAS_IDX = 3; + constexpr static uint32_t QBMM_BUFFER_NUM = 2; + constexpr static uint16_t QBMM_AIV_SYNC_AIC_FLAG = 6; + constexpr static uint16_t QBMM_AIC_SYNC_AIV_FLAG = 8; + constexpr static uint8_t QBMM_AIC_SYNC_AIV_MODE = 4; + constexpr static uint64_t QBMM_MAX_STEP_SCALEA_K = 16; + constexpr static uint32_t QBMM_UB_ALIGN_SIZE = 32; + constexpr static uint32_t UB_TWO_BANK_ELEMS_B32 = 128U; + constexpr static int64_t PER_BLOCK_SIZE = 128LL; + constexpr static uint32_t UB_SUB_BANK_NUM = 2U; + constexpr static uint32_t UB_SUB_BANK_ELEMS_B32 = 64U; + constexpr static uint32_t UB_ALIGN_SIZE = 32U; + constexpr static uint32_t QBMM_BMM_BLOCK_NUM = 16; + + struct PertileUBParam { + bool CopyOutWithSplitN = false; + uint16_t ndNum; + uint64_t singleM; + uint64_t singleN; + uint64_t validM; + uint32_t validN[UB_SUB_BANK_NUM]; + uint64_t offsetScaleM; + uint64_t offsetScaleN[UB_SUB_BANK_NUM]; + uint64_t offsetY[UB_SUB_BANK_NUM]; + uint64_t offsetBias[UB_SUB_BANK_NUM]; + }; + +public: + CATLASS_DEVICE void Init(const Params *params); + template + CATLASS_DEVICE void operator()(const TensorUb &tensorUbPing); + CATLASS_DEVICE void UpdateGlobalAddr(const BlockCoord &baseOffset); + CATLASS_DEVICE void UpdateParamsForNextProblem(const GemmCoord &problemShape); + CATLASS_DEVICE auto GetL0c2UbPingTensor(); + CATLASS_DEVICE auto GetL0c2UbPongTensor(); + +private: + CATLASS_DEVICE void ProcessAivSingleKPertile(int64_t x1ScaleOffset, + __gm__ X2ScaleType *x2ScaleAddr[UB_SUB_BANK_NUM]); + + CATLASS_DEVICE void ProcessAivSingleKPerblock(int64_t x1ScaleOffset, + __gm__ X2ScaleType *x2ScaleAddr[UB_SUB_BANK_NUM]); + CATLASS_DEVICE void CopyInBias(); + template + CATLASS_DEVICE __ubuf__ T *CopyInX1Scale(uint64_t srcOffset, uint64_t m, uint64_t k); + template + CATLASS_DEVICE T CopyInX1ScalePertile(__gm__ T *src, uint64_t offset); + template + CATLASS_DEVICE void CopyInX2Scale(T x2Scale[UB_SUB_BANK_NUM], __gm__ T *src[UB_SUB_BANK_NUM], uint64_t offset); + CATLASS_DEVICE int64_t CalcX1OffsetPerGroup(); + CATLASS_DEVICE void CalcX2OffsetPerGroup(int64_t x2ScaleOffset[UB_SUB_BANK_NUM]); + template + CATLASS_DEVICE __ubuf__ T *GetX1ScaleUbAddrPerGroup(int64_t x1ScaleOffset, uint64_t kOffset, uint64_t kElem); + template + __simd_vf__ void AivPerTensor(__ubuf__ CalcType *dst, __ubuf__ CType *l0cOut, __ubuf__ X1ScaleType *x1Scale, + uint16_t mSize, uint32_t nSize0, uint32_t nSize1, uint16_t kSize, + X2ScaleType x2Scale0, X2ScaleType x2Scale1, uint64_t x1ScaleKIdxInCache); + template + __simd_vf__ void AivPerTensor(__ubuf__ CType *dst, __ubuf__ CType *l0cOut, X1ScaleType x1Scale, + uint16_t mSize, uint32_t nSize0, uint32_t nSize1, X2ScaleType x2Scale0, + X2ScaleType x2Scale1); + template + __simd_vf__ void AddBias( + __ubuf__ CalcType *mmAdd, __ubuf__ BiasType *bias, uint16_t mSize, uint32_t nSize0, uint32_t nSize1); + CATLASS_DEVICE void AivPostProcess(const AscendC::LocalTensor &mmAddUb); + CATLASS_DEVICE void CopyOut(const AscendC::LocalTensor &ubRes, uint16_t eventId, uint16_t blkCount, + uint32_t blkLen, uint32_t srcStride, uint32_t dstStride, uint64_t yOffset); + CATLASS_DEVICE void CastAndCopyOut(const AscendC::LocalTensor &mmAddUb); + CATLASS_DEVICE void UpdatePertileUBValidMN(); + CATLASS_DEVICE void UpdatePertileUBParam(); + CATLASS_DEVICE void WaitForCube(uint16_t crossPingPongID) + { + AscendC::CrossCoreWaitFlag(QBMM_AIV_SYNC_AIC_FLAG + crossPingPongID); + } + CATLASS_DEVICE void NotifyCube(uint16_t crossPingPongID) + { + AscendC::CrossCoreSetFlag(QBMM_AIC_SYNC_AIV_FLAG + crossPingPongID); + } + + AscendC::GlobalTensor biasGlobal_; + AscendC::GlobalTensor cGlobal_; + AscendC::GlobalTensor x1ScaleGlobal_; + __gm__ X1ScaleType *x1ScaleGlobalPerblock_; + __gm__ X2ScaleType *x2ScaleGlobal_; + AscendC::LocalTensor mmResPing_; + AscendC::LocalTensor mmResPong_; + AscendC::LocalTensor ubResPing_; + AscendC::LocalTensor ubResPong_; + AscendC::LocalTensor mmAddUb_; + AscendC::LocalTensor x1ScaleUbPing_; + AscendC::LocalTensor x1ScaleUbPong_; + AscendC::LocalTensor biasUbPing_; + AscendC::LocalTensor biasUbPong_; + +private: + const Params *params_; + PertileUBParam ubParams_; + GemmCoord problemShape_{}; + GemmCoord actualSingleShape_{}; + BlockCoord baseOffset_{0, 0, 0, 0}; + BlockCoord blockCoord_{0, 0, 0, 0}; + + uint64_t scaleM_ = 0; + uint64_t scaleN_ = 0; + uint64_t scaleK_ = 0; + uint32_t subBlockIdx_; + uint16_t crossPingPongID_ = 0; + uint16_t x1ScalePingPongID_ = 0; + uint16_t biasPingPongID_ = 2; + bool isPergroup_ = false; + bool isBias_ = false; + bool needAivSet = false; +}; + +QBMM_BLOCK_EPILOGUE_PERTILE_CLASS_LOCAL_PARAMS +CATLASS_DEVICE void BlockEpilogue::Init(const Params *params) +{ + if ASCEND_IS_AIC { + return; + } + + params_ = params; + subBlockIdx_ = AscendC::GetSubBlockIdx(); + constexpr uint32_t elems = UB_TWO_BANK_ELEMS_B32 * PER_BLOCK_SIZE; + constexpr uint32_t addUbOffset = elems * UB_SUB_BANK_NUM * sizeof(CType); // l0c res 128KB + mmAddUb_ = AscendC::LocalTensor(AscendC::TPosition::VECCALC, addUbOffset, elems); + constexpr uint32_t afterAddOffset = addUbOffset + elems * sizeof(CalcType); // ub add res 64KB + if constexpr (!AscendC::IsSameType::value) { + ubResPing_ = AscendC::LocalTensor(AscendC::TPosition::VECCALC, afterAddOffset, elems); + ubResPong_ = ubResPing_[elems / QBMM_BUFFER_NUM]; + } + + isPergroup_ = params_->groupSizeM == 1; + if (isPergroup_) { + constexpr uint32_t x1ScaleUbOffset = + (AscendC::IsSameType::value) ? afterAddOffset : afterAddOffset + elems * sizeof(YType); + x1ScaleUbPing_ = AscendC::LocalTensor(AscendC::TPosition::VECCALC, x1ScaleUbOffset, + PER_BLOCK_SIZE * QBMM_MAX_STEP_SCALEA_K); + x1ScaleUbPong_ = AscendC::LocalTensor( + AscendC::TPosition::VECCALC, x1ScaleUbOffset + PER_BLOCK_SIZE * QBMM_MAX_STEP_SCALEA_K * sizeof(X1ScaleType), + PER_BLOCK_SIZE * QBMM_MAX_STEP_SCALEA_K); + if (params->isBias == 1) { + AscendC::SetFlag( + biasPingPongID_); + AscendC::SetFlag(biasPingPongID_ ^ 1); + isBias_ = true; + constexpr uint32_t afterx1ScaleUbOffset = + x1ScaleUbOffset + 2 * PER_BLOCK_SIZE * QBMM_MAX_STEP_SCALEA_K * sizeof(X1ScaleType); + biasUbPing_ = + AscendC::LocalTensor(AscendC::TPosition::VECCALC, afterx1ScaleUbOffset, PER_BLOCK_SIZE); + biasUbPong_ = AscendC::LocalTensor( + AscendC::TPosition::VECCALC, afterx1ScaleUbOffset + PER_BLOCK_SIZE * sizeof(BiasType), PER_BLOCK_SIZE); + } + } +} + +QBMM_BLOCK_EPILOGUE_PERTILE_CLASS_LOCAL_PARAMS +CATLASS_DEVICE void BlockEpilogue::UpdateParamsForNextProblem( + const GemmCoord &problemShape) +{ + problemShape_ = problemShape; + + scaleM_ = CeilDiv(problemShape_.m(), params_->groupSizeM); + scaleN_ = CeilDiv(problemShape_.n(), params_->groupSizeN); + scaleK_ = CeilDiv(problemShape_.k(), params_->groupSizeK); +} + +QBMM_BLOCK_EPILOGUE_PERTILE_CLASS_LOCAL_PARAMS +CATLASS_DEVICE void BlockEpilogue::UpdateGlobalAddr( + const BlockCoord &baseOffset) +{ + if ASCEND_IS_AIV { + x1ScaleGlobal_.SetGlobalBuffer((__gm__ X1ScaleType *)params_->x1ScaleGmAddr + tla::get(baseOffset)); + x1ScaleGlobalPerblock_ = (__gm__ X1ScaleType *)params_->x1ScaleGmAddr + tla::get(baseOffset); + x2ScaleGlobal_ = (__gm__ X2ScaleType *)params_->x2ScaleGmAddr + tla::get(baseOffset); + cGlobal_.SetGlobalBuffer((__gm__ YType *)params_->outGmAddr + tla::get(baseOffset)); + biasGlobal_.SetGlobalBuffer((__gm__ BiasType *)params_->biasGmAddr + tla::get(baseOffset)); + } +} + +QBMM_BLOCK_EPILOGUE_PERTILE_CLASS_LOCAL_PARAMS +CATLASS_DEVICE int64_t BlockEpilogue::CalcX1OffsetPerGroup() +{ + int64_t x1ScaleOffset = tla::get(blockCoord_) * scaleK_; + if (subBlockIdx_ == 1) { + x1ScaleOffset += (ubParams_.offsetScaleM * scaleK_); + } + return x1ScaleOffset; +} + +QBMM_BLOCK_EPILOGUE_PERTILE_CLASS_LOCAL_PARAMS +CATLASS_DEVICE void BlockEpilogue::CalcX2OffsetPerGroup( + int64_t x2ScaleOffset[UB_SUB_BANK_NUM]) +{ + x2ScaleOffset[0] = tla::get(blockCoord_) / params_->groupSizeN + ubParams_.offsetScaleN[0]; + x2ScaleOffset[1] = tla::get(blockCoord_) / params_->groupSizeN + ubParams_.offsetScaleN[1]; +} + +QBMM_BLOCK_EPILOGUE_PERTILE_CLASS_LOCAL_PARAMS +template +CATLASS_DEVICE __ubuf__ T *BlockEpilogue::CopyInX1Scale( + uint64_t srcOffset, uint64_t m, uint64_t k) +{ + AscendC::DataCopyExtParams x1ScaleGm2UbParams; + AscendC::DataCopyPadExtParams padParams; + x1ScaleGm2UbParams.blockCount = m; + x1ScaleGm2UbParams.blockLen = k * sizeof(T); + x1ScaleGm2UbParams.srcStride = (scaleK_ - k) * sizeof(T); + + auto x1ScaleUb = x1ScalePingPongID_ == 0 ? &x1ScaleUbPing_ : &x1ScaleUbPong_; + AscendC::WaitFlag(x1ScalePingPongID_); + AscendC::DataCopyPad(*x1ScaleUb, x1ScaleGlobal_[srcOffset], x1ScaleGm2UbParams, padParams); + AscendC::SetFlag(x1ScalePingPongID_); + AscendC::WaitFlag(x1ScalePingPongID_); + return reinterpret_cast<__ubuf__ T *>(x1ScaleUb->GetPhyAddr()); +} +QBMM_BLOCK_EPILOGUE_PERTILE_CLASS_LOCAL_PARAMS +CATLASS_DEVICE void BlockEpilogue::CopyInBias() +{ + AscendC::DataCopyExtParams bias4N0Gm2UbParams; + AscendC::DataCopyExtParams bias4N1Gm2UbParams; + AscendC::DataCopyPadExtParams padParams; + uint32_t validN0Align = ubParams_.validN[0] * sizeof(BiasType); + bias4N0Gm2UbParams.blockCount = 1; + bias4N0Gm2UbParams.blockLen = validN0Align; + bias4N0Gm2UbParams.srcStride = 0; + bias4N0Gm2UbParams.dstStride = 0; + AscendC::WaitFlag(biasPingPongID_); + auto biasUb = biasPingPongID_ == 2 ? biasUbPing_ : biasUbPong_; + AscendC::DataCopyPad(biasUb, biasGlobal_[ubParams_.offsetBias[0]], bias4N0Gm2UbParams, padParams); + if (ubParams_.validN[1] != 0) { + uint32_t validN1Align = ubParams_.validN[1] * sizeof(BiasType); + bias4N1Gm2UbParams.blockCount = 1; + bias4N1Gm2UbParams.blockLen = validN1Align; + bias4N1Gm2UbParams.srcStride = 0; + bias4N1Gm2UbParams.dstStride = 0; + AscendC::DataCopyPad( + biasUb[RoundUp(ubParams_.validN[0], UB_ALIGN_SIZE / sizeof(BiasType))], biasGlobal_[ubParams_.offsetBias[1]], + bias4N1Gm2UbParams, padParams); + } + AscendC::SetFlag(biasPingPongID_); +} +QBMM_BLOCK_EPILOGUE_PERTILE_CLASS_LOCAL_PARAMS +template +CATLASS_DEVICE T BlockEpilogue::CopyInX1ScalePertile( + __gm__ T *src, uint64_t offset) +{ + return src[offset]; +} + +QBMM_BLOCK_EPILOGUE_PERTILE_CLASS_LOCAL_PARAMS +template +CATLASS_DEVICE void BlockEpilogue::CopyInX2Scale( + T x2Scale[UB_SUB_BANK_NUM], __gm__ T *src[UB_SUB_BANK_NUM], uint64_t offset) +{ + x2Scale[0] = src[0][offset * scaleN_]; + x2Scale[1] = src[1][offset * scaleN_]; +} + +QBMM_BLOCK_EPILOGUE_PERTILE_CLASS_LOCAL_PARAMS +CATLASS_DEVICE void BlockEpilogue::UpdatePertileUBValidMN() +{ + int64_t actualN = actualSingleShape_.n(); + if (ubParams_.CopyOutWithSplitN) { + ubParams_.validM = ubParams_.singleM; + uint64_t subBlockIdxOffset = AscendC::GetSubBlockIdx() * ubParams_.singleN; + uint64_t ndNumN = 2 * ubParams_.singleN + subBlockIdxOffset; // 2: the nSize of 2 ND, base is 2 ND + ubParams_.validN[0] = actualN < subBlockIdxOffset ? 0 : Min(ubParams_.singleN, actualN - subBlockIdxOffset); + ubParams_.validN[1] = actualN < ndNumN ? 0 : Min(ubParams_.singleN, actualN - ndNumN); + } else { + if (AscendC::GetSubBlockIdx() == 0) { + ubParams_.validM = ubParams_.singleM; + } else { + ubParams_.validM = actualSingleShape_.m() - ubParams_.singleM; + } + ubParams_.validN[0] = Min(ubParams_.singleN, static_cast(actualN)); + ubParams_.validN[1] = actualN < ubParams_.singleN ? 0 : Min(ubParams_.singleN, actualN - ubParams_.singleN); + } +} + +QBMM_BLOCK_EPILOGUE_PERTILE_CLASS_LOCAL_PARAMS +CATLASS_DEVICE void BlockEpilogue::UpdatePertileUBParam() +{ + ubParams_.CopyOutWithSplitN = + actualSingleShape_.n() > params_->groupSizeN || actualSingleShape_.m() == 1; + uint32_t fixpipeN = 0; + if (ubParams_.CopyOutWithSplitN) { + // (m * n/2) is written to 2 UB, n must be multiples of 32 + // | AIV0 singleN | AIV1 singleN | AIV0 singleN | AIV1 singleN |, max(singleN) = 64 + ubParams_.ndNum = actualSingleShape_.n() > UB_TWO_BANK_ELEMS_B32 ? 2 : 1; // 2: 2 ND + int64_t alignedNBase = + actualSingleShape_.n() > PER_BLOCK_SIZE ? PER_BLOCK_SIZE : AscendC::ONE_BLK_SIZE * ubParams_.ndNum; + fixpipeN = RoundUp(actualSingleShape_.n(), static_cast(alignedNBase)) / ubParams_.ndNum; + ubParams_.singleN = fixpipeN / static_cast(AscendC::GetTaskRation()); + ubParams_.singleM = actualSingleShape_.m(); + } else { + // (m/2 * n) is written to 2 UB, m must be multiples of 2 + // | AIV0 singleN | AIV0 singleN | + // | AIV1 singleN | AIV1 singleN | + ubParams_.ndNum = actualSingleShape_.n() > UB_SUB_BANK_ELEMS_B32 ? 2 : 1; // 2: 2 ND + fixpipeN = RoundUp(actualSingleShape_.n(), static_cast(AscendC::BLOCK_CUBE) * ubParams_.ndNum) / + ubParams_.ndNum; + ubParams_.singleN = fixpipeN; + ubParams_.singleM = CeilDiv(actualSingleShape_.m(), AscendC::GetTaskRation()); + } + + UpdatePertileUBValidMN(); + int64_t offsetM = 0; + int64_t offsetN0 = 0; + int64_t offsetN1 = 0; + + if (ubParams_.CopyOutWithSplitN) { + offsetN0 = ubParams_.validN[0] == 0 ? 0 : AscendC::GetSubBlockIdx() * ubParams_.singleN; + offsetN1 = ubParams_.validN[1] == 0 ? offsetN0 : offsetN0 + UB_SUB_BANK_NUM * ubParams_.singleN; + } else { + if (AscendC::GetSubBlockIdx() == 1) { + offsetM += ubParams_.singleM; + } + offsetN1 = ubParams_.validN[1] == 0 ? 0 : ubParams_.singleN; + } + + ubParams_.offsetScaleM = offsetM / params_->groupSizeM; + ubParams_.offsetScaleN[0] = offsetN0 / params_->groupSizeN; + ubParams_.offsetScaleN[1] = offsetN1 / params_->groupSizeN; + ubParams_.offsetY[0] = tla::get(blockCoord_) + offsetM * problemShape_.n() + offsetN0; + ubParams_.offsetY[1] = tla::get(blockCoord_) + offsetM * problemShape_.n() + offsetN1; + ubParams_.offsetBias[0] = tla::get(blockCoord_) + offsetN0; + ubParams_.offsetBias[1] = tla::get(blockCoord_) + offsetN1; +} + +QBMM_BLOCK_EPILOGUE_PERTILE_CLASS_LOCAL_PARAMS +template +CATLASS_DEVICE void BlockEpilogue::operator()( + const TensorUb &tensorUb) +{ + actualSingleShape_ = GemmCoord{tla::get<0>(tensorUb.shape()), + tla::get<1>(tensorUb.shape()), + problemShape_.k()}; + blockCoord_ = BlockCoord{tla::get<0>(tensorUb.coord()) * problemShape_.n() + tla::get<1>(tensorUb.coord()), + tla::get<1>(tensorUb.coord()), + tla::get<0>(tensorUb.coord()), 0}; + + constexpr uint32_t elems = UB_TWO_BANK_ELEMS_B32 * PER_BLOCK_SIZE; + mmResPing_ = tensorUb.data(); + mmResPong_ = tensorUb.data()[elems]; + UpdatePertileUBParam(); + int64_t x1ScaleOffset = CalcX1OffsetPerGroup(); // one block same x1Scale + int64_t x2ScaleOffset[UB_SUB_BANK_NUM] = {0}; // maybe 2 different x2Scale + CalcX2OffsetPerGroup(x2ScaleOffset); + __gm__ X2ScaleType *x2ScaleAddr[UB_SUB_BANK_NUM] = {x2ScaleGlobal_ + x2ScaleOffset[0], + x2ScaleGlobal_ + x2ScaleOffset[1]}; + + if (isPergroup_) { + ProcessAivSingleKPertile(x1ScaleOffset, x2ScaleAddr); + } else { + ProcessAivSingleKPerblock(x1ScaleOffset, x2ScaleAddr); + } +} + +QBMM_BLOCK_EPILOGUE_PERTILE_CLASS_LOCAL_PARAMS +template +CATLASS_DEVICE __ubuf__ T * +BlockEpilogue::GetX1ScaleUbAddrPerGroup( + int64_t x1ScaleOffset, uint64_t kOffset, uint64_t kElem) +{ + uint64_t scaleX1GmOffset; + scaleX1GmOffset = x1ScaleOffset + kOffset; + return CopyInX1Scale(scaleX1GmOffset, ubParams_.validM, kElem); +} + +QBMM_BLOCK_EPILOGUE_PERTILE_CLASS_LOCAL_PARAMS +CATLASS_DEVICE void BlockEpilogue::ProcessAivSingleKPertile( + int64_t x1ScaleOffset, __gm__ X2ScaleType *x2ScaleAddr[UB_SUB_BANK_NUM]) +{ + auto mmAddUbAddr = reinterpret_cast<__ubuf__ CType *>(mmAddUb_.GetPhyAddr()); + const uint16_t x1ScaleKElem = Min(QBMM_MAX_STEP_SCALEA_K, scaleK_); + uint64_t kElem; + __ubuf__ X1ScaleType *x1ScaleUbAddr; + X2ScaleType x2Scale[UB_SUB_BANK_NUM]; + + for (uint64_t kb = 0, kOffset = 0; kb < problemShape_.k(); kb += params_->baseK, kOffset++) { + CopyInX2Scale(x2Scale, x2ScaleAddr, kOffset); + uint64_t x1ScaleKRem = kOffset % x1ScaleKElem; + if (x1ScaleKRem == 0) { + kElem = Min(static_cast(x1ScaleKElem), scaleK_ - kOffset); + x1ScaleUbAddr = GetX1ScaleUbAddrPerGroup(x1ScaleOffset, kOffset, kElem); + } + + WaitForCube(crossPingPongID_); + auto mmUbInputAddr = crossPingPongID_ == 0 ? reinterpret_cast(mmResPing_.GetPhyAddr()) : + reinterpret_cast(mmResPong_.GetPhyAddr()); + if (kb == 0) { + if (ubParams_.ndNum == 1) { + AivPerTensor((__ubuf__ CalcType *)mmAddUbAddr, (__ubuf__ CType *)mmUbInputAddr, x1ScaleUbAddr, + ubParams_.validM, ubParams_.validN[0], ubParams_.validN[1], kElem, x2Scale[0], + x2Scale[1], x1ScaleKRem); + } else { + AivPerTensor((__ubuf__ CalcType *)mmAddUbAddr, (__ubuf__ CType *)mmUbInputAddr, x1ScaleUbAddr, + ubParams_.validM, ubParams_.validN[0], ubParams_.validN[1], kElem, x2Scale[0], + x2Scale[1], x1ScaleKRem); + } + } else { + if (ubParams_.ndNum == 1) { + AivPerTensor((__ubuf__ CalcType *)mmAddUbAddr, (__ubuf__ CType *)mmUbInputAddr, x1ScaleUbAddr, + ubParams_.validM, ubParams_.validN[0], ubParams_.validN[1], kElem, x2Scale[0], + x2Scale[1], x1ScaleKRem); + } else { + AivPerTensor((__ubuf__ CalcType *)mmAddUbAddr, (__ubuf__ CType *)mmUbInputAddr, x1ScaleUbAddr, + ubParams_.validM, ubParams_.validN[0], ubParams_.validN[1], kElem, x2Scale[0], + x2Scale[1], x1ScaleKRem); + } + } + if (x1ScaleKRem == x1ScaleKElem - 1 || kOffset == scaleK_ - 1) { + AscendC::SetFlag(x1ScalePingPongID_); + x1ScalePingPongID_ = (x1ScalePingPongID_ + 1) & 1; + } + NotifyCube(crossPingPongID_); + needAivSet = needAivSet || crossPingPongID_ == 1; + crossPingPongID_ = (crossPingPongID_ + 1) & 1; + } + if (isBias_) { + CopyInBias(); + } + AivPostProcess(mmAddUb_); +} + +QBMM_BLOCK_EPILOGUE_PERTILE_CLASS_LOCAL_PARAMS +CATLASS_DEVICE void BlockEpilogue::ProcessAivSingleKPerblock( + int64_t x1ScaleOffset, __gm__ X2ScaleType *x2ScaleAddr[UB_SUB_BANK_NUM]) +{ + auto mmAddUbAddr = reinterpret_cast<__ubuf__ CType *>(mmAddUb_.GetPhyAddr()); + auto x1ScaleAddr = x1ScaleGlobalPerblock_ + x1ScaleOffset; + X2ScaleType x2Scale[UB_SUB_BANK_NUM]; + for (uint64_t kb = 0, kOffset = 0; kb < problemShape_.k(); kb += params_->baseK, kOffset++) { + CopyInX2Scale(x2Scale, x2ScaleAddr, kOffset); + X1ScaleType x1Scale = CopyInX1ScalePertile(x1ScaleAddr, kOffset); + WaitForCube(crossPingPongID_); + auto mmUbInputAddr = crossPingPongID_ == 0 ? reinterpret_cast(mmResPing_.GetPhyAddr()) : + reinterpret_cast(mmResPong_.GetPhyAddr()); + if (kb == 0) { + if (ubParams_.ndNum == 1) { + AivPerTensor((__ubuf__ CType *)mmAddUbAddr, (__ubuf__ CType *)mmUbInputAddr, x1Scale, + ubParams_.validM, ubParams_.validN[0], ubParams_.validN[1], x2Scale[0], + x2Scale[1]); + } else { + AivPerTensor((__ubuf__ CType *)mmAddUbAddr, (__ubuf__ CType *)mmUbInputAddr, x1Scale, + ubParams_.validM, ubParams_.validN[0], ubParams_.validN[1], x2Scale[0], + x2Scale[1]); + } + } else { + if (ubParams_.ndNum == 1) { + AivPerTensor((__ubuf__ CType *)mmAddUbAddr, (__ubuf__ CType *)mmUbInputAddr, x1Scale, + ubParams_.validM, ubParams_.validN[0], ubParams_.validN[1], x2Scale[0], + x2Scale[1]); + } else { + AivPerTensor((__ubuf__ CType *)mmAddUbAddr, (__ubuf__ CType *)mmUbInputAddr, x1Scale, + ubParams_.validM, ubParams_.validN[0], ubParams_.validN[1], x2Scale[0], + x2Scale[1]); + } + } + NotifyCube(crossPingPongID_); + crossPingPongID_ = (crossPingPongID_ + 1) & 1; + } + AivPostProcess(mmAddUb_); +} + +QBMM_BLOCK_EPILOGUE_PERTILE_CLASS_LOCAL_PARAMS +template +__simd_vf__ void BlockEpilogue::AivPerTensor( + __ubuf__ CalcType *dst, __ubuf__ CType *l0cOut, __ubuf__ X1ScaleType *x1Scale, uint16_t mSize, uint32_t nSize0, + uint32_t nSize1, uint16_t kSize, X2ScaleType x2Scale0, X2ScaleType x2Scale1, uint64_t x1ScaleKIdxInCache) +{ + uint16_t alignM = RoundUp(mSize, QBMM_UB_ALIGN_SIZE / sizeof(X1ScaleType)); + uint16_t alignK = RoundUp(kSize, QBMM_UB_ALIGN_SIZE / sizeof(X1ScaleType)); + AscendC::MicroAPI::RegTensor x1ScaleReg, muledScaleReg; + AscendC::MicroAPI::RegTensor l0cOutReg; + AscendC::MicroAPI::RegTensor addReg; + AscendC::MicroAPI::RegTensor ResReg, mulScaleOutReg; + constexpr static AscendC::MicroAPI::CastTrait ctInt322Fp32 = { + AscendC::MicroAPI::RegLayout::UNKNOWN, AscendC::MicroAPI::SatMode::UNKNOWN, + AscendC::MicroAPI::MaskMergeMode::ZEROING, AscendC::RoundMode::CAST_RINT}; + { + for (uint16_t mIdx = 0; mIdx < mSize; mIdx++) { + AscendC::MicroAPI::LoadAlign( + x1ScaleReg, x1Scale + mIdx * alignK + x1ScaleKIdxInCache); + + // 1 ND + uint32_t elementNum = nSize0; + AscendC::MicroAPI::MaskReg maskN = AscendC::MicroAPI::UpdateMask(elementNum); + // copy input from ub to register, addr of ub should align to 32B + uint32_t offset = mIdx * UB_TWO_BANK_ELEMS_B32; + uint32_t l0cOutOffset = mIdx * (nSize0 + nSize1); + AscendC::MicroAPI::LoadAlign(l0cOutReg, l0cOut + l0cOutOffset); + // l0c_out * scale + AscendC::MicroAPI::Muls(muledScaleReg, x1ScaleReg, x2Scale0, maskN); + if constexpr (AscendC::IsSameType::value) { + AscendC::MicroAPI::RegTensor l0cOutRegAfterCast{}; + AscendC::MicroAPI::Cast(l0cOutRegAfterCast, l0cOutReg, maskN); + AscendC::MicroAPI::Mul(mulScaleOutReg, l0cOutRegAfterCast, muledScaleReg, maskN); + } else { + AscendC::MicroAPI::Mul(mulScaleOutReg, l0cOutReg, muledScaleReg, maskN); + } + uint32_t dstUbOffset = offset; + if constexpr (isFirstKLoop) { + AscendC::MicroAPI::StoreAlign(dst + dstUbOffset, + mulScaleOutReg, maskN); + } else { + AscendC::MicroAPI::LoadAlign(addReg, dst + dstUbOffset); + AscendC::MicroAPI::Add(ResReg, mulScaleOutReg, addReg, maskN); + // copy out from register to ub + AscendC::MicroAPI::StoreAlign(dst + dstUbOffset, + ResReg, maskN); + } + // 2 ND + if constexpr (ndNum == 1) { + continue; + } + elementNum = nSize1; + maskN = AscendC::MicroAPI::UpdateMask(elementNum); + // copy input from ub to register, addr of ub should align to 32B + l0cOutOffset = l0cOutOffset + nSize0; // + distance of 2 NDs + AscendC::MicroAPI::LoadAlign(l0cOutReg, l0cOut + l0cOutOffset); + // l0c_out * scale + AscendC::MicroAPI::Muls(muledScaleReg, x1ScaleReg, x2Scale1, maskN); + if constexpr (AscendC::IsSameType::value) { + AscendC::MicroAPI::RegTensor l0cOutRegAfterCast{}; + AscendC::MicroAPI::Cast(l0cOutRegAfterCast, l0cOutReg, maskN); + AscendC::MicroAPI::Mul(mulScaleOutReg, l0cOutRegAfterCast, muledScaleReg, maskN); + } else { + AscendC::MicroAPI::Mul(mulScaleOutReg, l0cOutReg, muledScaleReg, maskN); + } + dstUbOffset = offset + nSize0; + if constexpr (isFirstKLoop) { + AscendC::MicroAPI::StoreAlign(dst + dstUbOffset, + mulScaleOutReg, maskN); + } else { + AscendC::MicroAPI::LoadAlign(addReg, dst + dstUbOffset); + AscendC::MicroAPI::Add(ResReg, mulScaleOutReg, addReg, maskN); + // copy out from register to ub + AscendC::MicroAPI::StoreAlign(dst + dstUbOffset, + ResReg, maskN); + } + } + } +} + +QBMM_BLOCK_EPILOGUE_PERTILE_CLASS_LOCAL_PARAMS +template +__simd_vf__ void BlockEpilogue::AivPerTensor( + __ubuf__ CType *dst, __ubuf__ CType *l0cOut, X1ScaleType x1Scale, uint16_t mSize, uint32_t nSize0, uint32_t nSize1, + X2ScaleType x2Scale0, X2ScaleType x2Scale1) +{ + { + for (uint16_t mIdx = 0; mIdx < mSize; mIdx++) { + // 1 ND + X2ScaleType scaleMul = x1Scale * x2Scale0; + uint32_t elementNum = nSize0; + AscendC::MicroAPI::MaskReg maskN = AscendC::MicroAPI::UpdateMask(elementNum); + AscendC::MicroAPI::RegTensor l0cOutReg; + AscendC::MicroAPI::RegTensor addReg; + AscendC::MicroAPI::RegTensor ResReg, mulScaleOutReg; + // copy input from ub to register, addr of ub should align to 32B + uint32_t offset = mIdx * UB_TWO_BANK_ELEMS_B32; + uint32_t l0cOutOffset = offset; + AscendC::MicroAPI::LoadAlign(l0cOutReg, l0cOut + offset); + // l0c_out * scale + AscendC::MicroAPI::Muls(mulScaleOutReg, l0cOutReg, scaleMul, maskN); + uint32_t dstUbOffset = offset; + if constexpr (isFirstKLoop) { + AscendC::MicroAPI::StoreAlign(dst + dstUbOffset, + mulScaleOutReg, maskN); + } else { + AscendC::MicroAPI::LoadAlign(addReg, dst + dstUbOffset); + AscendC::MicroAPI::Add(ResReg, mulScaleOutReg, addReg, maskN); + // copy out from register to ub + AscendC::MicroAPI::StoreAlign(dst + dstUbOffset, + ResReg, maskN); + } + // 2 ND + if constexpr (ndNum == 1) { + continue; + } + scaleMul = x1Scale * x2Scale1; + elementNum = nSize1; + maskN = AscendC::MicroAPI::UpdateMask(elementNum); + // copy input from ub to register, addr of ub should align to 32B + l0cOutOffset = offset + nSize0; // + distance of 2 NDs + AscendC::MicroAPI::LoadAlign(l0cOutReg, l0cOut + l0cOutOffset); + // l0c_out * scale + AscendC::MicroAPI::Muls(mulScaleOutReg, l0cOutReg, scaleMul, maskN); + dstUbOffset = offset + nSize0; + if constexpr (isFirstKLoop) { + AscendC::MicroAPI::StoreAlign(dst + dstUbOffset, + mulScaleOutReg, maskN); + } else { + AscendC::MicroAPI::LoadAlign(addReg, dst + dstUbOffset); + AscendC::MicroAPI::Add(ResReg, mulScaleOutReg, addReg, maskN); + // copy out from register to ub + AscendC::MicroAPI::StoreAlign(dst + dstUbOffset, + ResReg, maskN); + } + } + } +} + +QBMM_BLOCK_EPILOGUE_PERTILE_CLASS_LOCAL_PARAMS +template +__simd_vf__ void BlockEpilogue::AddBias( + __ubuf__ CalcType *mmAdd, __ubuf__ BiasType *bias, uint16_t mSize, uint32_t nSize0, uint32_t nSize1) +{ + { + uint32_t mmAddOffset{}; + for (uint16_t mIdx = 0; mIdx < mSize; mIdx++) { + AscendC::MicroAPI::RegTensor biasReg{}; + AscendC::MicroAPI::RegTensor mmAddReg{}; + AscendC::MicroAPI::DataCopy(biasReg, bias); + mmAddOffset = mIdx * PER_BLOCK_SIZE; + uint32_t elementNum = nSize0; + AscendC::MicroAPI::MaskReg maskN = AscendC::MicroAPI::UpdateMask(elementNum); + AscendC::MicroAPI::DataCopy(mmAddReg, mmAdd + mmAddOffset); + AscendC::MicroAPI::Add(mmAddReg, mmAddReg, biasReg, maskN); + AscendC::MicroAPI::DataCopy( + mmAdd + mmAddOffset, mmAddReg, maskN); + if constexpr (ndNum == 1) { + continue; + } + elementNum = nSize1; + maskN = AscendC::MicroAPI::UpdateMask(elementNum); + mmAddOffset += nSize0; + AscendC::MicroAPI::DataCopy(biasReg, bias + nSize0); + AscendC::MicroAPI::DataCopy(mmAddReg, mmAdd + mmAddOffset); + AscendC::MicroAPI::Add(mmAddReg, mmAddReg, biasReg, maskN); + AscendC::MicroAPI::DataCopy( + mmAdd + mmAddOffset, mmAddReg, maskN); + } + } +} + +QBMM_BLOCK_EPILOGUE_PERTILE_CLASS_LOCAL_PARAMS +CATLASS_DEVICE void BlockEpilogue::AivPostProcess( + const AscendC::LocalTensor &mmAddUb) +{ + if (ubParams_.validM == 0) { + return; + } + if constexpr (AscendC::IsSameType::value) { + if (isBias_) { + auto biasUb = biasPingPongID_ == 2 ? biasUbPing_ : biasUbPong_; + AscendC::WaitFlag(biasPingPongID_); + auto mmAddUbAddr = reinterpret_cast<__ubuf__ CalcType *>(mmAddUb.GetPhyAddr()); + auto biasUbAddr = reinterpret_cast<__ubuf__ BiasType *>(biasUb.GetPhyAddr()); + if (ubParams_.ndNum == 1) { + AddBias<1U>( + (__ubuf__ CalcType *)mmAddUbAddr, (__ubuf__ BiasType *)biasUbAddr, ubParams_.validM, + ubParams_.validN[0], ubParams_.validN[1]); + } else if (ubParams_.ndNum == 2) { + AddBias<2U>( + (__ubuf__ CalcType *)mmAddUbAddr, (__ubuf__ BiasType *)biasUbAddr, ubParams_.validM, + ubParams_.validN[0], ubParams_.validN[1]); + } + AscendC::SetFlag(biasPingPongID_); + biasPingPongID_ = biasPingPongID_ ^ 1; + } + } + if constexpr (AscendC::IsSameType::value) { + // mov optimize in splitM, 0~63 + 64 ~127 -> 0~127 + AscendC::SetFlag(0); + AscendC::WaitFlag(0); + if (ubParams_.ndNum == 2 && !ubParams_.CopyOutWithSplitN) { // 2: 2ND, opt branch with splitM + uint32_t sumN = ubParams_.validN[0] + ubParams_.validN[1]; + CopyOut(mmAddUb, 0, ubParams_.validM, sumN, UB_TWO_BANK_ELEMS_B32 - sumN, problemShape_.n() - sumN, + ubParams_.offsetY[0]); + } else { + for (uint64_t ndIdx = 0; ndIdx < ubParams_.ndNum; ndIdx++) { + if (ubParams_.validN[ndIdx] > 0) { + CopyOut(mmAddUb[ndIdx * ubParams_.validN[0]], 0, ubParams_.validM, ubParams_.validN[ndIdx], + UB_TWO_BANK_ELEMS_B32 - ubParams_.validN[ndIdx], + problemShape_.n() - ubParams_.validN[ndIdx], ubParams_.offsetY[ndIdx]); + } + } + } + AscendC::SetFlag(0); + AscendC::WaitFlag(0); + } else { + AscendC::PipeBarrier(); + CastAndCopyOut(mmAddUb); + } +} + +QBMM_BLOCK_EPILOGUE_PERTILE_CLASS_LOCAL_PARAMS +CATLASS_DEVICE void BlockEpilogue::CopyOut( + const AscendC::LocalTensor &ubRes, uint16_t eventId, uint16_t blkCount, uint32_t blkLen, uint32_t srcStride, + uint32_t dstStride, uint64_t yOffset) +{ + AscendC::DataCopyExtParams copyParams{ + blkCount, static_cast(blkLen * sizeof(YType)), + static_cast(srcStride * sizeof(YType) / AscendC::ONE_BLK_SIZE), + static_cast(dstStride * sizeof(YType)), 0}; + AscendC::DataCopyPad(cGlobal_[yOffset], ubRes, copyParams); +} + +QBMM_BLOCK_EPILOGUE_PERTILE_CLASS_LOCAL_PARAMS +CATLASS_DEVICE void BlockEpilogue::CastAndCopyOut( + const AscendC::LocalTensor &mmAddUb) +{ + if (ubParams_.ndNum == 2 && !ubParams_.CopyOutWithSplitN) { // 2: 2ND, opt branch with splitM + uint32_t sumN = ubParams_.validN[0] + ubParams_.validN[1]; + uint32_t mSizePing = CeilDiv(ubParams_.validM, static_cast(QBMM_BUFFER_NUM)); + uint32_t mSize[QBMM_BUFFER_NUM] = {mSizePing, static_cast(ubParams_.validM - mSizePing)}; + for (uint32_t mDbIdx = 0; mDbIdx < QBMM_BUFFER_NUM; ++mDbIdx) { + if (mSize[mDbIdx] > 0 && sumN > 0) { + auto ubRes = mDbIdx == 0 ? &ubResPing_ : &ubResPong_; + AscendC::WaitFlag(mDbIdx); + AscendC::Cast( + *ubRes, mmAddUb[mDbIdx * mSizePing * UB_TWO_BANK_ELEMS_B32], AscendC::RoundMode::CAST_RINT, + mSize[mDbIdx] * UB_TWO_BANK_ELEMS_B32); + AscendC::SetFlag(mDbIdx); + AscendC::WaitFlag(mDbIdx); + CopyOut( + *ubRes, mDbIdx, mSize[mDbIdx], sumN, UB_TWO_BANK_ELEMS_B32 - sumN, problemShape_.n() - sumN, + ubParams_.offsetY[0] + mDbIdx * mSizePing * problemShape_.n()); + AscendC::SetFlag(mDbIdx); + } + } + } else { + for (uint64_t ndIdx = 0; ndIdx < ubParams_.ndNum; ndIdx++) { + auto ubRes = ndIdx == 0 ? &ubResPing_ : &ubResPong_; + if (ubParams_.validN[ndIdx] > 0) { + AscendC::UnaryRepeatParams repeatParam; + repeatParam.srcBlkStride = 1; + repeatParam.dstBlkStride = 1; + // write continuously + repeatParam.dstRepStride = CeilDiv(ubParams_.singleN, AscendC::ONE_BLK_SIZE / sizeof(YType)); + // srcStride is 16(512B / 32B), subBank0 256B one repeat + repeatParam.srcRepStride = QBMM_BMM_BLOCK_NUM; + AscendC::WaitFlag(ndIdx); + AscendC::Cast(*ubRes, mmAddUb[ubParams_.validN[0] * ndIdx], AscendC::RoundMode::CAST_RINT, + ubParams_.validN[ndIdx], ubParams_.validM, repeatParam); + AscendC::SetFlag(ndIdx); + AscendC::WaitFlag(ndIdx); + CopyOut(*ubRes, ndIdx, ubParams_.validM, ubParams_.validN[ndIdx], + ubParams_.singleN - ubParams_.validN[ndIdx], + problemShape_.n() - ubParams_.validN[ndIdx], ubParams_.offsetY[ndIdx]); + AscendC::SetFlag(ndIdx); + } + } + } +} + +} // namespace Catlass::Epilogue::Block + +#endif // CATLASS_EPILOGUE_BLOCK_EPILOGUE_PER_GROUP_PER_BLOCK_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_per_token_dequant_tla.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_per_token_dequant_tla.hpp new file mode 100644 index 0000000000..cdf6e9847a --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_per_token_dequant_tla.hpp @@ -0,0 +1,784 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_EPILOGUE_BLOCK_EPILOGUE_PER_TOKEN_DEQUANT_TLA_HPP +#define CATLASS_EPILOGUE_BLOCK_EPILOGUE_PER_TOKEN_DEQUANT_TLA_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/epilogue/kda_dispatch_policy.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/gemm/kda_gemm_type.hpp" +#include "catlass/kda_matrix_coord.hpp" +#include "catlass/layout/kda_layout.hpp" +#include "catlass/detail/kda_callback.hpp" +#include "tla/kda_tla_layout.hpp" +#include "tla/kda_tensor.hpp" + +namespace Catlass::Epilogue::Block { + +template < + uint32_t UB_STAGES_, + class ElementC_, + class ElementScale_, + class ElementPerTokenScale_, + class ElementD_, + class TileRowBroadcastMul_, + class TileBroadcastOneBlk_, + class TileOneBlkColumnBroadcastMul_, + class TileCopy_, + class EpilogueTileSwizzle_> +class BlockEpilogue< + EpilogueAtlasA2PerTokenDequantTla, + ElementC_, + ElementScale_, + ElementPerTokenScale_, + ElementD_, + TileRowBroadcastMul_, + TileBroadcastOneBlk_, + TileOneBlkColumnBroadcastMul_, + TileCopy_, + EpilogueTileSwizzle_> { +public: + using DispatchPolicy = EpilogueAtlasA2PerTokenDequantTla; + using ArchTag = typename DispatchPolicy::ArchTag; + static constexpr uint32_t UB_STAGES = UB_STAGES_; + using TileCopy = TileCopy_; + + // Data infos + using ElementC = ElementC_; + using LayoutC = typename TileCopy::LayoutC; + using ElementScale = ElementScale_; + using LayoutScale = typename TileCopy::LayoutX; + using ElementPerTokenScale = ElementPerTokenScale_; + using LayoutPerTokenScale = typename TileCopy::LayoutY; + using ElementD = ElementD_; + using LayoutD = typename TileCopy::LayoutD; + + // Check data infos + static_assert( + std::is_same_v && (std::is_same_v || std::is_same_v) && + std::is_same_v && std::is_same_v, + "The element type template parameters of BlockEpilogue are wrong"); + static_assert( + tla::detail::isRowMajor::value && tla::detail::isRowMajor::value, + "The layout template parameters of BlockEpilogue are wrong"); + + // Tile compute ops + using TileRowBroadcastMul = TileRowBroadcastMul_; + using TileBroadcastOneBlk = TileBroadcastOneBlk_; + using TileOneBlkColumnBroadcastMul = TileOneBlkColumnBroadcastMul_; + + using EpilogueTileSwizzle = EpilogueTileSwizzle_; + + using TileShape = typename TileRowBroadcastMul::TileShape; + + static constexpr uint32_t ELE_NUM_PER_BLK = BYTE_PER_BLK / sizeof(float); + + static_assert( + TileShape::ROW == TileBroadcastOneBlk::COMPUTE_LENGTH && + std::is_same_v, + "TileShape must be consistent for all tile compute ops"); + + static_assert( + (UB_STAGES * (TileShape::COUNT * sizeof(ElementC) + TileShape::COLUMN * sizeof(ElementScale) + TileShape::ROW * sizeof(ElementPerTokenScale) + TileShape::COUNT * sizeof(ElementD)) + (TileShape::COUNT + TileShape::COLUMN + TileShape::COUNT + TileShape::ROW) * sizeof(float) + TileShape::ROW * BYTE_PER_BLK) <= ArchTag::UB_SIZE, + "TileShape is too large to fit in UB"); + + struct Params { + GM_ADDR ptrScale{nullptr}; + LayoutScale layoutScale{}; + GM_ADDR ptrPerTokenScale{nullptr}; + LayoutPerTokenScale layoutPerTokenScale{}; + GM_ADDR ptrD{nullptr}; + LayoutD layoutD{}; + Params() = default; + }; + + CATLASS_DEVICE + BlockEpilogue(Arch::Resource const &resource, Params const ¶ms = Params{}, uint32_t sharedUbSize = 0) + : params(params) + { + uint32_t ubOffset = sharedUbSize; + int32_t eventVMTE2 = 0; + int32_t eventMTE2V = 0; + int32_t eventMTE3V = 0; + int32_t eventVMTE3 = 0; + for (uint32_t i = 0; i < UB_STAGES; ++i) { + ubCList[i] = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += TileShape::COUNT * sizeof(ElementC); + ubScaleList[i] = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += TileShape::COLUMN * sizeof(ElementScale); + ubPerTokenScaleList[i] = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += TileShape::ROW * sizeof(ElementPerTokenScale); + ubDList[i] = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += TileShape::COUNT * sizeof(ElementD); + + eventUbCVMTE2List[i] = eventVMTE2++; + eventUbCMTE2VList[i] = eventMTE2V++; + eventUbScaleVMTE2List[i] = eventVMTE2++; + eventUbScaleMTE2VList[i] = eventMTE2V++; + eventUbPerTokenScaleVMTE2List[i] = eventVMTE2++; + eventUbPerTokenScaleMTE2VList[i] = eventMTE2V++; + eventUbDMTE3VList[i] = eventMTE3V++; + eventUbDVMTE3List[i] = eventVMTE3++; + + AscendC::SetFlag(eventUbCVMTE2List[i]); + AscendC::SetFlag(eventUbScaleVMTE2List[i]); + AscendC::SetFlag(eventUbPerTokenScaleVMTE2List[i]); + AscendC::SetFlag(eventUbDMTE3VList[i]); + } + ubCFp32 = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += TileShape::COUNT * sizeof(float); + ubScaleFp32 = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += TileShape::COLUMN * sizeof(float); + ubMul = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += TileShape::COUNT * sizeof(float); + ubPerTokenScaleFp32 = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += TileShape::ROW * sizeof(float); + ubPerTokenScaleFp32Brcb = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += TileShape::ROW * BYTE_PER_BLK; + ubPerTokenMul = ubMul; + } + + CATLASS_DEVICE + ~BlockEpilogue() + { + for (uint32_t i = 0; i < UB_STAGES; ++i) { + AscendC::WaitFlag(eventUbCVMTE2List[i]); + AscendC::WaitFlag(eventUbScaleVMTE2List[i]); + AscendC::WaitFlag(eventUbPerTokenScaleVMTE2List[i]); + AscendC::WaitFlag(eventUbDMTE3VList[i]); + } + } + + CATLASS_DEVICE + void UpdateParams(Params const ¶ms_) + { + params = params_; + } + + template + CATLASS_DEVICE void operator()( + TensorC &tensorBlockC, TensorScale &tensorBlockScale, TensorPerTokenScaleScale &tensorBlockPerTokenScaleScale, + TensorD &tensorBlockD, GemmCoord const &actualBlockShapeMNK, Callback &&callback = Callback{}) + { + if (actualBlockShapeMNK.k() == 0) { + return; + } + callback(); + + MatrixCoord actualBlockShape = actualBlockShapeMNK.GetCoordMN(); + + using CopyGmToUbC = typename TileCopy::template CopyGmToUbC; + using CopyGmToUbScale = typename TileCopy::template CopyGmToUbX; + using CopyGmToUbPerTokenScale = typename TileCopy::template CopyGmToUbY; + using CopyUbToGmD = typename TileCopy::template CopyUbToGmD; + CopyGmToUbC copyGmToUbC; + CopyGmToUbScale copyGmToUbScale; + CopyGmToUbPerTokenScale copyGmToUbPerTokenScale; + CopyUbToGmD copyUbToGmD; + + auto ubTileStride = static_cast(TileShape::COLUMN); + auto ubTileStrideRow = static_cast(TileShape::ROW); + auto tileShape = TileShape::ToCoord(); + EpilogueTileSwizzle epilogueTileSwizzle(actualBlockShape, tileShape); + uint32_t tileLoops = epilogueTileSwizzle.GetLoops(); + uint32_t subblockIdx = AscendC::GetSubBlockIdx(); + uint32_t subblockNum = AscendC::GetSubBlockNum(); + for (uint32_t loopIdx = subblockIdx; loopIdx < tileLoops; loopIdx += subblockNum) { + auto tileCoord = epilogueTileSwizzle.GetTileCoord(loopIdx); + auto actualTileShape = epilogueTileSwizzle.GetActualTileShape(tileCoord); + MatrixCoord tileOffsetInBlock = tileCoord * tileShape; + auto tileOffsetInBlockRow = tileOffsetInBlock.row(); + auto tileOffsetInBlockColumn = tileOffsetInBlock.column(); + + // build tensor C block in GM + auto tensorSubBlockC = GetTile( + tensorBlockC, tla::MakeCoord(tileOffsetInBlockRow, tileOffsetInBlockColumn), + tla::MakeShape(actualTileShape.row(), actualTileShape.column())); + // build tensor C block in UB + auto &ubC = ubCList[ubListId]; + auto layoutUbC = tla::MakeLayout( + tla::MakeShape(actualTileShape.row(), actualTileShape.column()), tla::MakeStride(ubTileStride, tla::Int<1>{})); + auto tensorUbC = tla::MakeTensor(ubC, layoutUbC, Arch::PositionUB{}); + // copy tensor C from GM to UB + AscendC::WaitFlag(eventUbCVMTE2List[ubListId]); + copyGmToUbC(tensorUbC, tensorSubBlockC); + AscendC::SetFlag(eventUbCMTE2VList[ubListId]); + + /////////////////////////////////////////////////// + // build tensor Scale block in GM + auto tensorSubBlockScale = GetTile( + tensorBlockScale, tla::MakeCoord(0, tileOffsetInBlockColumn), + tla::MakeShape(tla::Int<1>{}, actualTileShape.column())); + // build tensor Scale block in UB + auto &ubScale = ubScaleList[ubListId]; + auto layoutUbScale = tla::MakeLayout( + tla::MakeShape(tla::Int<1>{}, actualTileShape.column()), tla::MakeStride(ubTileStride, tla::Int<1>{})); + auto tensorUbScale = tla::MakeTensor(ubScale, layoutUbScale, Arch::PositionUB{}); + // copy tensor Scale from GM to UB + AscendC::WaitFlag(eventUbScaleVMTE2List[ubListId]); + copyGmToUbScale(tensorUbScale, tensorSubBlockScale); + AscendC::SetFlag(eventUbScaleMTE2VList[ubListId]); + + /////////////////////////////////////////////////// + // build tensor PerTokenScale block in GM + auto tensorSubBlockPerTokenScale = GetTile( + tensorBlockPerTokenScaleScale, tla::MakeCoord(0, tileOffsetInBlockRow), + tla::MakeShape(1, actualTileShape.row())); + // build tensor PerTokenScale block in UB + auto &ubPerTokenScale = ubPerTokenScaleList[ubListId]; + auto layoutUbPerTokenScale = tla::MakeLayout( + tla::MakeShape(tla::Int<1>{}, actualTileShape.row()), tla::MakeStride(ubTileStrideRow, tla::Int<1>{})); + auto tensorUbPerTokenScale = tla::MakeTensor(ubPerTokenScale, layoutUbPerTokenScale, Arch::PositionUB{}); + // copy tensor PerTokenScale from GM to UB + AscendC::WaitFlag(eventUbPerTokenScaleVMTE2List[ubListId]); + copyGmToUbPerTokenScale(tensorUbPerTokenScale, tensorSubBlockPerTokenScale); + AscendC::SetFlag(eventUbPerTokenScaleMTE2VList[ubListId]); + + /////////////////////////////////////////////////// + auto tensorUbScaleFp32 = tla::MakeTensor(ubScaleFp32, layoutUbScale, Arch::PositionUB{}); + auto tensorUbPerTokenScaleFp32 = tla::MakeTensor(ubPerTokenScaleFp32, layoutUbPerTokenScale, Arch::PositionUB{}); + auto tensorUbCFp32 = tla::MakeTensor(ubCFp32, layoutUbC, Arch::PositionUB{}); + auto tensorUbMul = tla::MakeTensor(ubMul, layoutUbC, Arch::PositionUB{}); + + auto layoutUbPerTokenScaleBrcb = tla::MakeLayout( + tla::MakeShape(actualTileShape.row(), ELE_NUM_PER_BLK), tla::MakeStride(ELE_NUM_PER_BLK, tla::Int<1>{})); + auto tensorUbPerTokenScaleFp32Brcb = tla::MakeTensor(ubPerTokenScaleFp32Brcb, layoutUbPerTokenScaleBrcb, Arch::PositionUB{}); + auto tensorUbPerTokenMul = tla::MakeTensor(ubPerTokenMul, layoutUbC, Arch::PositionUB{}); + + AscendC::WaitFlag(eventUbCMTE2VList[ubListId]); + AscendC::Cast(ubCFp32, ubC, AscendC::RoundMode::CAST_RINT, TileShape::COUNT); + AscendC::SetFlag(eventUbCVMTE2List[ubListId]); + + AscendC::WaitFlag(eventUbScaleMTE2VList[ubListId]); + AscendC::Cast(ubScaleFp32, ubScale, AscendC::RoundMode::CAST_NONE, TileShape::COLUMN); + AscendC::SetFlag(eventUbScaleVMTE2List[ubListId]); + + AscendC::WaitFlag(eventUbPerTokenScaleMTE2VList[ubListId]); + AscendC::Cast(ubPerTokenScaleFp32, ubPerTokenScale, AscendC::RoundMode::CAST_NONE, TileShape::ROW); + AscendC::SetFlag(eventUbPerTokenScaleVMTE2List[ubListId]); + + AscendC::PipeBarrier(); + tileRowBroadcastMul(tensorUbMul, tensorUbCFp32, tensorUbScaleFp32); + tileBroadcastOneBlk(tensorUbPerTokenScaleFp32Brcb, tensorUbPerTokenScaleFp32); + AscendC::PipeBarrier(); + tileOneBlkColumnBroadcastMul(tensorUbPerTokenMul, tensorUbMul, tensorUbPerTokenScaleFp32Brcb); + AscendC::PipeBarrier(); + + auto &ubD = ubDList[ubListId]; + + AscendC::WaitFlag(eventUbDMTE3VList[ubListId]); + AscendC::Cast(ubD, ubPerTokenMul, AscendC::RoundMode::CAST_RINT, TileShape::COUNT); + AscendC::SetFlag(eventUbDVMTE3List[ubListId]); + + // build tensor D block in GM + auto tensorSubBlockD = GetTile( + tensorBlockD, tla::MakeCoord(tileOffsetInBlockRow, tileOffsetInBlockColumn), + tla::MakeShape(actualTileShape.row(), actualTileShape.column())); + // build tensor D block in UB + auto tensorUbD = tla::MakeTensor(ubD, layoutUbC, Arch::PositionUB{}); + + AscendC::WaitFlag(eventUbDVMTE3List[ubListId]); + copyUbToGmD(tensorSubBlockD, tensorUbD); + AscendC::SetFlag(eventUbDMTE3VList[ubListId]); + + ubListId = (ubListId + 1 < UB_STAGES) ? (ubListId + 1) : 0; + } + } + +private: + Params params; + + AscendC::LocalTensor ubCList[UB_STAGES]; + AscendC::LocalTensor ubScaleList[UB_STAGES]; + AscendC::LocalTensor ubPerTokenScaleList[UB_STAGES]; + AscendC::LocalTensor ubDList[UB_STAGES]; + + int32_t eventUbCVMTE2List[UB_STAGES]; + int32_t eventUbCMTE2VList[UB_STAGES]; + int32_t eventUbScaleVMTE2List[UB_STAGES]; + int32_t eventUbScaleMTE2VList[UB_STAGES]; + int32_t eventUbPerTokenScaleVMTE2List[UB_STAGES]; + int32_t eventUbPerTokenScaleMTE2VList[UB_STAGES]; + int32_t eventUbDMTE3VList[UB_STAGES]; + int32_t eventUbDVMTE3List[UB_STAGES]; + + uint32_t ubListId{0}; + + AscendC::LocalTensor ubCFp32; + AscendC::LocalTensor ubScaleFp32; + AscendC::LocalTensor ubMul; + AscendC::LocalTensor ubPerTokenScaleFp32; + AscendC::LocalTensor ubPerTokenScaleFp32Brcb; + AscendC::LocalTensor ubPerTokenMul; + + TileRowBroadcastMul tileRowBroadcastMul; + TileBroadcastOneBlk tileBroadcastOneBlk; + TileOneBlkColumnBroadcastMul tileOneBlkColumnBroadcastMul; +}; + +template < + uint32_t UB_STAGES_, + class ElementC_, + class ElementD_, + class TileRowBroadcastMul_, + class TileBroadcastOneBlk_, + class TileOneBlkColumnBroadcastMul_, + class TileCopy_, + class EpilogueTileSwizzle_> +class BlockEpilogue< + EpilogueAtlasA2PerTokenDequantTla, + ElementC_, + float, + float, + ElementD_, + TileRowBroadcastMul_, + TileBroadcastOneBlk_, + TileOneBlkColumnBroadcastMul_, + TileCopy_, + EpilogueTileSwizzle_> { +public: + using DispatchPolicy = EpilogueAtlasA2PerTokenDequantTla; + using ArchTag = typename DispatchPolicy::ArchTag; + static constexpr uint32_t UB_STAGES = UB_STAGES_; + using TileCopy = TileCopy_; + + // Data infos + using ElementC = ElementC_; + using LayoutC = typename TileCopy::LayoutC; + using ElementScale = float; + using LayoutScale = typename TileCopy::LayoutX; + using ElementPerTokenScale = float; + using LayoutPerTokenScale = typename TileCopy::LayoutY; + using ElementD = ElementD_; + using LayoutD = typename TileCopy::LayoutD; + + // Check data infos + static_assert( + std::is_same_v && (std::is_same_v || std::is_same_v || + std::is_same_v), + "The element type template parameters of BlockEpilogue are wrong"); + static_assert( + tla::detail::isRowMajor::value && tla::detail::isRowMajor::value, + "The layout template parameters of BlockEpilogue are wrong"); + + // Tile compute ops + using TileRowBroadcastMul = TileRowBroadcastMul_; + using TileBroadcastOneBlk = TileBroadcastOneBlk_; + using TileOneBlkColumnBroadcastMul = TileOneBlkColumnBroadcastMul_; + + using EpilogueTileSwizzle = EpilogueTileSwizzle_; + + using TileShape = typename TileRowBroadcastMul::TileShape; + + static constexpr uint32_t ELE_NUM_PER_BLK = BYTE_PER_BLK / sizeof(float); + + static_assert( + TileShape::ROW == TileBroadcastOneBlk::COMPUTE_LENGTH && + std::is_same_v, + "TileShape must be consistent for all tile compute ops"); + + static_assert( + (UB_STAGES * (TileShape::COUNT * sizeof(ElementC) + TileShape::COLUMN * sizeof(ElementScale) + TileShape::ROW * sizeof(ElementPerTokenScale) + TileShape::COUNT * sizeof(ElementD)) + (TileShape::COUNT + TileShape::COUNT) * sizeof(float) + TileShape::ROW * BYTE_PER_BLK) <= ArchTag::UB_SIZE, + "TileShape is too large to fit in UB"); + + struct Params { + GM_ADDR ptrScale{nullptr}; + LayoutScale layoutScale{}; + GM_ADDR ptrPerTokenScale{nullptr}; + LayoutPerTokenScale layoutPerTokenScale{}; + GM_ADDR ptrD{nullptr}; + LayoutD layoutD{}; + Params() = default; + }; + + CATLASS_DEVICE + BlockEpilogue(Arch::Resource const &resource, Params const ¶ms = Params{}, uint32_t sharedUbSize = 0) + : params(params) + { + uint32_t ubOffset = sharedUbSize; + int32_t eventVMTE2 = 0; + int32_t eventMTE2V = 0; + int32_t eventMTE3V = 0; + int32_t eventVMTE3 = 0; + for (uint32_t i = 0; i < UB_STAGES; ++i) { + ubCList[i] = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += TileShape::COUNT * sizeof(ElementC); + ubScaleList[i] = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += TileShape::COLUMN * sizeof(ElementScale); + ubPerTokenScaleList[i] = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += TileShape::ROW * sizeof(ElementPerTokenScale); + ubDList[i] = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += TileShape::COUNT * sizeof(ElementD); + + eventUbCVMTE2List[i] = eventVMTE2++; + eventUbCMTE2VList[i] = eventMTE2V++; + eventUbScaleVMTE2List[i] = eventVMTE2++; + eventUbScaleMTE2VList[i] = eventMTE2V++; + eventUbPerTokenScaleVMTE2List[i] = eventVMTE2++; + eventUbPerTokenScaleMTE2VList[i] = eventMTE2V++; + eventUbDMTE3VList[i] = eventMTE3V++; + eventUbDVMTE3List[i] = eventVMTE3++; + + AscendC::SetFlag(eventUbCVMTE2List[i]); + AscendC::SetFlag(eventUbScaleVMTE2List[i]); + AscendC::SetFlag(eventUbPerTokenScaleVMTE2List[i]); + AscendC::SetFlag(eventUbDMTE3VList[i]); + } + ubCFp32 = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += TileShape::COUNT * sizeof(float); + ubMul = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += TileShape::COUNT * sizeof(float); + ubPerTokenScaleBrcb = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += TileShape::ROW * BYTE_PER_BLK; + ubPerTokenMul = ubCFp32; + } + + CATLASS_DEVICE + ~BlockEpilogue() + { + for (uint32_t i = 0; i < UB_STAGES; ++i) { + AscendC::WaitFlag(eventUbCVMTE2List[i]); + AscendC::WaitFlag(eventUbScaleVMTE2List[i]); + AscendC::WaitFlag(eventUbPerTokenScaleVMTE2List[i]); + AscendC::WaitFlag(eventUbDMTE3VList[i]); + } + } + + CATLASS_DEVICE + void UpdateParams(Params const ¶ms_) + { + params = params_; + } + + template + CATLASS_DEVICE void operator()( + TensorC &tensorBlockC, TensorScale &tensorBlockScale, TensorPerTokenScaleScale &tensorBlockPerTokenScaleScale, + TensorD &tensorBlockD, GemmCoord const &actualBlockShapeMNK, Callback &&callback = Callback{}) + { + if (actualBlockShapeMNK.k() == 0) { + return; + } + callback(); + + MatrixCoord actualBlockShape = actualBlockShapeMNK.GetCoordMN(); + + using CopyGmToUbC = typename TileCopy::template CopyGmToUbC; + using CopyGmToUbScale = typename TileCopy::template CopyGmToUbX; + using CopyGmToUbPerTokenScale = typename TileCopy::template CopyGmToUbY; + using CopyUbToGmD = typename TileCopy::template CopyUbToGmD; + CopyGmToUbC copyGmToUbC; + CopyGmToUbScale copyGmToUbScale; + CopyGmToUbPerTokenScale copyGmToUbPerTokenScale; + CopyUbToGmD copyUbToGmD; + + auto ubTileStride = static_cast(TileShape::COLUMN); + auto ubTileStrideRow = static_cast(TileShape::ROW); + auto tileShape = TileShape::ToCoord(); + EpilogueTileSwizzle epilogueTileSwizzle(actualBlockShape, tileShape); + uint32_t tileLoops = epilogueTileSwizzle.GetLoops(); + uint32_t subblockIdx = AscendC::GetSubBlockIdx(); + uint32_t subblockNum = AscendC::GetSubBlockNum(); + for (uint32_t loopIdx = subblockIdx; loopIdx < tileLoops; loopIdx += subblockNum) { + auto tileCoord = epilogueTileSwizzle.GetTileCoord(loopIdx); + auto actualTileShape = epilogueTileSwizzle.GetActualTileShape(tileCoord); + MatrixCoord tileOffsetInBlock = tileCoord * tileShape; + auto tileOffsetInBlockRow = tileOffsetInBlock.row(); + auto tileOffsetInBlockColumn = tileOffsetInBlock.column(); + + // build tensor C block in GM + auto tensorSubBlockC = GetTile( + tensorBlockC, tla::MakeCoord(tileOffsetInBlockRow, tileOffsetInBlockColumn), + tla::MakeShape(actualTileShape.row(), actualTileShape.column())); + // build tensor C block in UB + auto &ubC = ubCList[ubListId]; + auto layoutUbC = tla::MakeLayout( + tla::MakeShape(actualTileShape.row(), actualTileShape.column()), tla::MakeStride(ubTileStride, tla::Int<1>{})); + auto tensorUbC = tla::MakeTensor(ubC, layoutUbC, Arch::PositionUB{}); + // copy tensor C from GM to UB + AscendC::WaitFlag(eventUbCVMTE2List[ubListId]); + copyGmToUbC(tensorUbC, tensorSubBlockC); + AscendC::SetFlag(eventUbCMTE2VList[ubListId]); + + /////////////////////////////////////////////////// + // build tensor Scale block in GM + auto tensorSubBlockScale = GetTile( + tensorBlockScale, tla::MakeCoord(0, tileOffsetInBlockColumn), + tla::MakeShape(tla::Int<1>{}, actualTileShape.column())); + // build tensor Scale block in UB + auto &ubScale = ubScaleList[ubListId]; + auto layoutUbScale = tla::MakeLayout( + tla::MakeShape(tla::Int<1>{}, actualTileShape.column()), tla::MakeStride(ubTileStride, tla::Int<1>{})); + auto tensorUbScale = tla::MakeTensor(ubScale, layoutUbScale, Arch::PositionUB{}); + // copy tensor Scale from GM to UB + AscendC::WaitFlag(eventUbScaleVMTE2List[ubListId]); + copyGmToUbScale(tensorUbScale, tensorSubBlockScale); + AscendC::SetFlag(eventUbScaleMTE2VList[ubListId]); + + /////////////////////////////////////////////////// + // build tensor PerTokenScale block in GM + auto tensorSubBlockPerTokenScale = GetTile( + tensorBlockPerTokenScaleScale, tla::MakeCoord(0, tileOffsetInBlockRow), + tla::MakeShape(tla::Int<1>{}, actualTileShape.row())); + // build tensor PerTokenScale block in UB + auto &ubPerTokenScale = ubPerTokenScaleList[ubListId]; + auto layoutUbPerTokenScale = tla::MakeLayout( + tla::MakeShape(tla::Int<1>{}, actualTileShape.row()), tla::MakeStride(ubTileStrideRow, tla::Int<1>{})); + auto tensorUbPerTokenScale = tla::MakeTensor(ubPerTokenScale, layoutUbPerTokenScale, Arch::PositionUB{}); + // copy tensor PerTokenScale from GM to UB + AscendC::WaitFlag(eventUbPerTokenScaleVMTE2List[ubListId]); + copyGmToUbPerTokenScale(tensorUbPerTokenScale, tensorSubBlockPerTokenScale); + AscendC::SetFlag(eventUbPerTokenScaleMTE2VList[ubListId]); + + /////////////////////////////////////////////////// + auto tensorUbCFp32 = tla::MakeTensor(ubCFp32, layoutUbC, Arch::PositionUB{}); + auto tensorUbMul = tla::MakeTensor(ubMul, layoutUbC, Arch::PositionUB{}); + + auto layoutUbPerTokenScaleBrcb = tla::MakeLayout( + tla::MakeShape(actualTileShape.row(), ELE_NUM_PER_BLK), tla::MakeStride(ELE_NUM_PER_BLK, tla::Int<1>{})); + auto tensorUbPerTokenScaleBrcb = tla::MakeTensor(ubPerTokenScaleBrcb, layoutUbPerTokenScaleBrcb, Arch::PositionUB{}); + auto tensorUbPerTokenMul = tla::MakeTensor(ubPerTokenMul, layoutUbC, Arch::PositionUB{}); + + AscendC::WaitFlag(eventUbCMTE2VList[ubListId]); + AscendC::Cast(ubCFp32, ubC, AscendC::RoundMode::CAST_RINT, TileShape::COUNT); + AscendC::SetFlag(eventUbCVMTE2List[ubListId]); + + AscendC::PipeBarrier(); + AscendC::WaitFlag(eventUbScaleMTE2VList[ubListId]); + tileRowBroadcastMul(tensorUbMul, tensorUbCFp32, tensorUbScale); + AscendC::SetFlag(eventUbScaleVMTE2List[ubListId]); + + AscendC::PipeBarrier(); + AscendC::WaitFlag(eventUbPerTokenScaleMTE2VList[ubListId]); + tileBroadcastOneBlk(tensorUbPerTokenScaleBrcb, tensorUbPerTokenScale); + AscendC::SetFlag(eventUbPerTokenScaleVMTE2List[ubListId]); + + // build tensor D block in UB + auto &ubD = ubDList[ubListId]; + auto tensorUbD = tla::MakeTensor(ubD, layoutUbC, Arch::PositionUB{}); + + if constexpr (std::is_same_v) { + AscendC::PipeBarrier(); + AscendC::WaitFlag(eventUbDMTE3VList[ubListId]); + tileOneBlkColumnBroadcastMul(tensorUbD, tensorUbMul, tensorUbPerTokenScaleBrcb); + AscendC::SetFlag(eventUbDVMTE3List[ubListId]); + } else if constexpr (std::is_same_v) { + AscendC::PipeBarrier(); + tileOneBlkColumnBroadcastMul(tensorUbPerTokenMul, tensorUbMul, tensorUbPerTokenScaleBrcb); + AscendC::PipeBarrier(); + + AscendC::WaitFlag(eventUbDMTE3VList[ubListId]); + AscendC::Cast(ubD, ubPerTokenMul, AscendC::RoundMode::CAST_NONE, TileShape::COUNT); + AscendC::SetFlag(eventUbDVMTE3List[ubListId]); + } else { // std::is_same_v + AscendC::PipeBarrier(); + tileOneBlkColumnBroadcastMul(tensorUbPerTokenMul, tensorUbMul, tensorUbPerTokenScaleBrcb); + AscendC::PipeBarrier(); + + AscendC::WaitFlag(eventUbDMTE3VList[ubListId]); + AscendC::Cast(ubD, ubPerTokenMul, AscendC::RoundMode::CAST_RINT, TileShape::COUNT); + AscendC::SetFlag(eventUbDVMTE3List[ubListId]); + } + + // build tensor D block in GM + auto tensorSubBlockD = GetTile( + tensorBlockD, tla::MakeCoord(tileOffsetInBlockRow, tileOffsetInBlockColumn), + tla::MakeShape(actualTileShape.row(), actualTileShape.column())); + + AscendC::WaitFlag(eventUbDVMTE3List[ubListId]); + copyUbToGmD(tensorSubBlockD, tensorUbD); + AscendC::SetFlag(eventUbDMTE3VList[ubListId]); + + ubListId = (ubListId + 1 < UB_STAGES) ? (ubListId + 1) : 0; + } + } + +private: + Params params; + + AscendC::LocalTensor ubCList[UB_STAGES]; + AscendC::LocalTensor ubScaleList[UB_STAGES]; + AscendC::LocalTensor ubPerTokenScaleList[UB_STAGES]; + AscendC::LocalTensor ubDList[UB_STAGES]; + + int32_t eventUbCVMTE2List[UB_STAGES]; + int32_t eventUbCMTE2VList[UB_STAGES]; + int32_t eventUbScaleVMTE2List[UB_STAGES]; + int32_t eventUbScaleMTE2VList[UB_STAGES]; + int32_t eventUbPerTokenScaleVMTE2List[UB_STAGES]; + int32_t eventUbPerTokenScaleMTE2VList[UB_STAGES]; + int32_t eventUbDMTE3VList[UB_STAGES]; + int32_t eventUbDVMTE3List[UB_STAGES]; + + uint32_t ubListId{0}; + + AscendC::LocalTensor ubCFp32; + AscendC::LocalTensor ubMul; + AscendC::LocalTensor ubPerTokenScaleBrcb; + AscendC::LocalTensor ubPerTokenMul; + + TileRowBroadcastMul tileRowBroadcastMul; + TileBroadcastOneBlk tileBroadcastOneBlk; + TileOneBlkColumnBroadcastMul tileOneBlkColumnBroadcastMul; +}; + +template < + uint32_t UB_STAGES_, + class TileShape_, + class ElementSrc_, + class ElementScale_, + class ElementPerToken_, + class ElementDst_, + class TilePerTokenDequant_, + class TileCopy_> +class BlockEpilogue< + EpilogueAscend950PerTokenDequantTla, + TileShape_, + ElementSrc_, + ElementScale_, + ElementPerToken_, + ElementDst_, + TilePerTokenDequant_, + TileCopy_> { +public: + using DispatchPolicy = EpilogueAscend950PerTokenDequantTla; + using ArchTag = typename DispatchPolicy::ArchTag; + using TileCopy = TileCopy_; + + using ElementSrc = ElementSrc_; + using ElementScale = ElementScale_; + using LayoutScale = typename TileCopy::LayoutTagX; + using ElementPerToken = ElementPerToken_; + using LayoutPerToken = typename TileCopy::LayoutTagY; + using ElementDst = ElementDst_; + using LayoutDst = typename TileCopy::LayoutTagD; + + using TilePerTokenDequant = TilePerTokenDequant_; + using TileShape = TileShape_; + + static constexpr uint32_t UB_STAGES = UB_STAGES_; + static constexpr int16_t N_BASE_SIZE = static_cast(TileShape::COLUMN); + static constexpr uint16_t CV_RATIO = 2; + static constexpr uint32_t BLOCK_SIZE = TileShape::COUNT / CV_RATIO; + + static_assert( + UB_STAGES * (BLOCK_SIZE * sizeof(ElementSrc) + TileShape::COLUMN / CV_RATIO * sizeof(ElementScale) + TileShape::ROW / CV_RATIO * sizeof(ElementPerToken)) <= ArchTag::UB_SIZE, + "TileShape is too large to fit in UB"); + + CATLASS_DEVICE + BlockEpilogue(Arch::Resource const &resource, uint32_t &ubOffset) + { + int32_t eventVMTE2 = 0; + int32_t eventMTE2V = 0; + int32_t eventMTE3V = 0; + int32_t eventVMTE3 = 0; + for (uint32_t i = 0; i < UB_STAGES; ++i) { + ubScaleList[i] = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += TileShape::COLUMN * sizeof(ElementScale); + ubPerTokenList[i] = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += TileShape::ROW * sizeof(ElementPerToken); + + eventUbScaleVMTE2List[i] = eventVMTE2++; + eventUbScaleMTE2VList[i] = eventMTE2V++; + eventUbPerTokenVMTE2List[i] = eventVMTE2++; + eventUbPerTokenMTE2VList[i] = eventMTE2V++; + eventUbDMTE3VList[i] = eventMTE3V++; + eventUbDVMTE3List[i] = eventVMTE3++; + + AscendC::SetFlag(eventUbScaleVMTE2List[i]); + AscendC::SetFlag(eventUbPerTokenVMTE2List[i]); + AscendC::SetFlag(eventUbDMTE3VList[i]); + } + } + + CATLASS_DEVICE + ~BlockEpilogue() + { + for (uint32_t i = 0; i < UB_STAGES; ++i) { + AscendC::WaitFlag(eventUbScaleVMTE2List[i]); + AscendC::WaitFlag(eventUbPerTokenVMTE2List[i]); + AscendC::WaitFlag(eventUbDMTE3VList[i]); + } + } + + template + CATLASS_DEVICE void operator()( + TensorDst gmDequantOut, + TensorSrc ubGmmRes, + TensorScale gmScale, + TensorPerToken gmPerToken, + uint32_t ubListId) + { + using CopyGmToUbScale = typename TileCopy::template CopyGmToUbX; + using CopyGmToUbPerToken = typename TileCopy::template CopyGmToUbY; + using CopyUbToGmDequant = typename TileCopy::template CopyUbToGmD; + + CopyGmToUbScale copyGmToUbScale; + CopyGmToUbPerToken copyGmToUbPerToken; + CopyUbToGmDequant copyUbToGmDequant; + TilePerTokenDequant tilePerTokenDequant; + + uint32_t m = tla::get<0>(ubGmmRes.shape()); + uint32_t n = tla::get<1>(ubGmmRes.shape()); + + auto scaleLayout = tla::MakeLayout(n); + auto ubScale = tla::MakeTensor( + ubScaleList[ubListId], + scaleLayout, + Arch::PositionUB{}); + + auto perTokenLayout = tla::MakeLayout(m); + auto ubPerToken = tla::MakeTensor( + ubPerTokenList[ubListId], + perTokenLayout, + Arch::PositionUB{}); + + auto ubDequantOut = tla::MakeTensor( + ubGmmRes.data().template ReinterpretCast(), + ubGmmRes.layout(), + Arch::PositionUB{}); + + AscendC::WaitFlag(eventUbScaleVMTE2List[ubListId]); + copyGmToUbScale(ubScale, gmScale); + AscendC::SetFlag(eventUbScaleMTE2VList[ubListId]); + + AscendC::WaitFlag(eventUbPerTokenVMTE2List[ubListId]); + copyGmToUbPerToken(ubPerToken, gmPerToken); + AscendC::SetFlag(eventUbPerTokenMTE2VList[ubListId]); + + AscendC::WaitFlag(eventUbScaleMTE2VList[ubListId]); + AscendC::WaitFlag(eventUbPerTokenMTE2VList[ubListId]); + AscendC::WaitFlag(eventUbDMTE3VList[ubListId]); + tilePerTokenDequant(ubDequantOut, ubGmmRes, ubScale, ubPerToken); + AscendC::SetFlag(eventUbDVMTE3List[ubListId]); + AscendC::SetFlag(eventUbScaleVMTE2List[ubListId]); + AscendC::SetFlag(eventUbPerTokenVMTE2List[ubListId]); + + AscendC::WaitFlag(eventUbDVMTE3List[ubListId]); + copyUbToGmDequant(gmDequantOut, ubDequantOut); + AscendC::SetFlag(eventUbDMTE3VList[ubListId]); + } + +private: + AscendC::LocalTensor ubScaleList[UB_STAGES]; + AscendC::LocalTensor ubPerTokenList[UB_STAGES]; + + int32_t eventUbScaleVMTE2List[UB_STAGES]; + int32_t eventUbScaleMTE2VList[UB_STAGES]; + int32_t eventUbPerTokenVMTE2List[UB_STAGES]; + int32_t eventUbPerTokenMTE2VList[UB_STAGES]; + int32_t eventUbDMTE3VList[UB_STAGES]; + int32_t eventUbDVMTE3List[UB_STAGES]; +}; + +} // namespace Catlass::Epilogue::Block + +#endif // CATLASS_EPILOGUE_BLOCK_EPILOGUE_PER_TOKEN_DEQUANT_TLA_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_rescale_o_no_split_row.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_rescale_o_no_split_row.hpp new file mode 100644 index 0000000000..657ace9526 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_rescale_o_no_split_row.hpp @@ -0,0 +1,313 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_RESCALE_O_NO_SPLIT_ROW_HPP +#define CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_RESCALE_O_NO_SPLIT_ROW_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/epilogue/kda_dispatch_policy.hpp" +#include "catlass/epilogue/tile/kda_epilogue_tile_copy.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/kda_matrix_coord.hpp" + +namespace Catlass::Epilogue::Block { + +template < + class OutputType_, + class InputType_, + class UpdateType_> +class BlockEpilogue< + EpilogueAtlasA2RescaleO, + OutputType_, + InputType_, + UpdateType_> { +public: + // Type aliases + using DispatchPolicy = EpilogueAtlasA2RescaleO; + using ArchTag = typename DispatchPolicy::ArchTag; + + using ElementOutput = typename OutputType_::Element; + using ElementInput = typename InputType_::Element; + using ElementUpdate = typename UpdateType_::Element; + + using LayoutOutput = typename OutputType_::Layout; + using LayoutInput = typename InputType_::Layout; + using LayoutUpdate = typename UpdateType_::Layout; + + static constexpr uint32_t HALF_ELENUM_PER_BLK = 16; + static constexpr uint32_t BLOCK_SIZE = 16; + static constexpr uint32_t HALF_ELENUM_PER_VECCALC = 128; + static constexpr uint32_t FLOAT_ELENUM_PER_VECCALC = 64; + static constexpr uint32_t HALF_ELENUM_PER_LINE = 256; + static constexpr uint32_t FLOAT_ELENUM_PER_LINE = 128; + static constexpr uint32_t MULTIPLIER = 2; + static constexpr uint32_t FLOAT_BLOCK_SIZE = 8; + static constexpr uint32_t FLOAT_VECTOR_SIZE = 64; + static constexpr uint32_t UB_UINT8_VECTOR_SIZE = 1024; + static constexpr uint32_t UB_UINT8_BLOCK_SIZE = 16384; + static constexpr uint32_t HALF_DM_UB_SIZE = 64; + static constexpr uint32_t HALF_LL_UB_SIZE = 256; + static constexpr uint32_t VECTOR_SIZE = 128; + static constexpr uint32_t NUM4 = 4; + static constexpr uint32_t MAX_UB_O_ELEM_NUM = 4096; + static constexpr uint32_t MAX_ROW_NUM_SUB_CORE = 128; + + CATLASS_DEVICE + BlockEpilogue(Arch::Resource &resource) + { + // Allocate UB space + constexpr uint32_t LO_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE; + constexpr uint32_t GO_UB_TENSOR_OFFSET = 8 * UB_UINT8_BLOCK_SIZE; + constexpr uint32_t TV_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE; + + constexpr uint32_t HM_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 9 * UB_UINT8_VECTOR_SIZE; + constexpr uint32_t GL_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 12 * UB_UINT8_VECTOR_SIZE; + constexpr uint32_t DM_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 13 * UB_UINT8_VECTOR_SIZE; + + loUbTensor = resource.ubBuf.template GetBufferByByte(LO_UB_TENSOR_OFFSET); + dmUbTensor = resource.ubBuf.template GetBufferByByte(DM_UB_TENSOR_OFFSET); + glUbTensor = resource.ubBuf.template GetBufferByByte(GL_UB_TENSOR_OFFSET); + tvUbTensor = resource.ubBuf.template GetBufferByByte(TV_UB_TENSOR_OFFSET); + goUbTensor16 = resource.ubBuf.template GetBufferByByte(GO_UB_TENSOR_OFFSET); + goUbTensor32 = resource.ubBuf.template GetBufferByByte(GO_UB_TENSOR_OFFSET); + hmUbTensor = resource.ubBuf.template GetBufferByByte(HM_UB_TENSOR_OFFSET); + } + + CATLASS_DEVICE + ~BlockEpilogue() + { + } + + CATLASS_DEVICE + void SetMask(int32_t len) + { + uint64_t mask = 0; + uint64_t one = 1; + uint64_t temp = len % FLOAT_VECTOR_SIZE; + for (int64_t i = 0; i < temp; i++) { + mask |= one << i; + } + + if (len == VECTOR_SIZE) { + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + } else if (len >= FLOAT_VECTOR_SIZE) { + AscendC::SetVectorMask(mask, (uint64_t)-1); + } else { + AscendC::SetVectorMask(0x0, mask); + } + } + + CATLASS_DEVICE + void CopyOToGm( + AscendC::GlobalTensor gOutput, + uint32_t curRowNum, uint32_t qSBlockSize, uint32_t embed, + uint32_t embedRound, uint32_t qNThisSubBlock, uint32_t oHiddenSize) + { + if (qNThisSubBlock == 0) { + AscendC::DataCopyPad( + gOutput, + goUbTensor16, + AscendC::DataCopyExtParams(curRowNum, embed * 2, 0, (oHiddenSize - embed) * 2, 0)); + } else { + for (uint32_t qNIdx = 0; qNIdx < qNThisSubBlock; qNIdx++) { + AscendC::DataCopyPad( + gOutput[qNIdx * embed], + goUbTensor16[qNIdx * embedRound * qSBlockSize], + AscendC::DataCopyExtParams(qSBlockSize, embed * 2, 0, (oHiddenSize - embed) * 2, 0)); + } + } + } + + CATLASS_DEVICE + void SubCoreCompute( + AscendC::GlobalTensor gOutput, + AscendC::GlobalTensor gInput, + const LayoutOutput &layoutOutput, + const LayoutInput &layoutInput, + uint32_t qNThisSubBlock, + uint32_t isFirstStackTile, uint32_t isLastStackTile, uint32_t curStackTileMod) + { + uint32_t curRowNum = layoutInput.shape(0); + uint32_t embed = layoutInput.shape(1); + uint32_t embedRound = layoutInput.stride(0); + uint32_t curRowNumRound = RoundUp(curRowNum, FLOAT_BLOCK_SIZE); + uint32_t qSBlockSize = layoutOutput.shape(0); + uint32_t oHiddenSize = layoutOutput.shape(1); + uint32_t dmUbOffsetCurStackTile = + curStackTileMod * MAX_ROW_NUM_SUB_CORE; + + AscendC::WaitFlag(EVENT_ID3); + if (!isFirstStackTile) { + AscendC::DataCopy(loUbTensor, gInput, + AscendC::DataCopyParams(1, curRowNum * embedRound / FLOAT_BLOCK_SIZE, 0, 0)); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + AscendC::Brcb(tvUbTensor.ReinterpretCast(), + dmUbTensor[dmUbOffsetCurStackTile].ReinterpretCast(), curRowNumRound / FLOAT_BLOCK_SIZE, + AscendC::BrcbRepeatParams(1, 8)); + AscendC::PipeBarrier(); + // *** go = go * dm_block + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + for (uint32_t vmul_idx = 0; vmul_idx < embed / FLOAT_VECTOR_SIZE; ++vmul_idx) { + AscendC::Mul(goUbTensor32[vmul_idx * FLOAT_VECTOR_SIZE], + goUbTensor32[vmul_idx * FLOAT_VECTOR_SIZE], + tvUbTensor, (uint64_t)0, + curRowNum, + AscendC::BinaryRepeatParams(1, 1, 0, embedRound / FLOAT_BLOCK_SIZE, + embedRound / FLOAT_BLOCK_SIZE, 1)); + } + if (embed % FLOAT_VECTOR_SIZE > 0) { + SetMask(embed % FLOAT_VECTOR_SIZE); + AscendC::Mul(goUbTensor32[embed / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + goUbTensor32[embed / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + tvUbTensor, + (uint64_t)0, curRowNum, + AscendC::BinaryRepeatParams(1, 1, 0, embedRound / FLOAT_BLOCK_SIZE, + embedRound / FLOAT_BLOCK_SIZE, 1)); + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + AscendC::PipeBarrier(); + // *** go = lo + go + AscendC::Add(goUbTensor32, goUbTensor32, + loUbTensor, (uint64_t)0, + (curRowNum * embedRound + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, + AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); + AscendC::PipeBarrier(); + } else { + // *** go = lo + AscendC::DataCopy(goUbTensor32, gInput, + AscendC::DataCopyParams(1, curRowNum * embedRound / FLOAT_BLOCK_SIZE, 0, 0)); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + } + AscendC::SetFlag(EVENT_ID3); + + if (isLastStackTile) { + // *** gl_block = expand_to_block(gl), 存放于 tv + AscendC::Brcb(tvUbTensor.ReinterpretCast(), + glUbTensor.ReinterpretCast(), + curRowNumRound / FLOAT_BLOCK_SIZE, + AscendC::BrcbRepeatParams(1, 8)); + AscendC::PipeBarrier(); + // *** go = go / gl_block + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + for (uint32_t vdiv_idx = 0; vdiv_idx < embed / FLOAT_VECTOR_SIZE; ++vdiv_idx) { + AscendC::Div(goUbTensor32[vdiv_idx * FLOAT_VECTOR_SIZE], + goUbTensor32[vdiv_idx * FLOAT_VECTOR_SIZE], + tvUbTensor, (uint64_t)0, + curRowNum, + AscendC::BinaryRepeatParams(1, 1, 0, embedRound / FLOAT_BLOCK_SIZE, + embedRound / FLOAT_BLOCK_SIZE, 1)); + } + if (embed % FLOAT_VECTOR_SIZE > 0) { + SetMask(embed % FLOAT_VECTOR_SIZE); + AscendC::Div(goUbTensor32[embed / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + goUbTensor32[embed / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + tvUbTensor, + (uint64_t)0, curRowNum, + AscendC::BinaryRepeatParams(1, 1, 0, embedRound / FLOAT_BLOCK_SIZE, + embedRound / FLOAT_BLOCK_SIZE, 1)); + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + AscendC::PipeBarrier(); + + // *** go = castfp32to16(go) + if (std::is_same::value) { + AscendC::Cast( + goUbTensor16, goUbTensor32, + AscendC::RoundMode::CAST_RINT, (uint64_t)0, + (curRowNum * embedRound + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, + AscendC::UnaryRepeatParams(1, 1, 4, 8)); + } else { + AscendC::Cast( + goUbTensor16, goUbTensor32, + AscendC::RoundMode::CAST_NONE, (uint64_t)0, + (curRowNum * embedRound + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, + AscendC::UnaryRepeatParams(1, 1, 4, 8)); + } + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + + // ***move O to GM + CopyOToGm( + gOutput, curRowNum, qSBlockSize, embed, + embedRound, qNThisSubBlock, oHiddenSize); + } + } + + CATLASS_DEVICE + void operator()( + AscendC::GlobalTensor gOutput, + AscendC::GlobalTensor gInput, + const LayoutOutput &layoutOutput, + const LayoutInput &layoutInput, + GemmCoord actualBlockShape, + uint32_t qSBlockSize, uint32_t qNBlockSize, + uint32_t isFirstStackTile, uint32_t isLastStackTile, uint32_t curStackTileMod) + { + uint32_t rowNum = actualBlockShape.m(); + uint32_t embed = actualBlockShape.n(); + + uint32_t subBlockIdx = AscendC::GetSubBlockIdx(); + uint32_t subBlockNum = AscendC::GetSubBlockNum(); + + uint32_t qNSplitSubBlock = qNBlockSize / subBlockNum; + uint32_t qNThisSubBlock = (qNBlockSize == 1) ? + 0 : + (subBlockIdx == 1) ? + (qNBlockSize - qNSplitSubBlock) : + qNSplitSubBlock; + uint32_t inRowSplitSubBlock = (qNBlockSize == 1) ? + (qSBlockSize / subBlockNum) : + (qSBlockSize * qNSplitSubBlock); + uint32_t inRowActualThisSubBlock = (subBlockIdx == 1) ? + (rowNum - inRowSplitSubBlock) : + inRowSplitSubBlock; + uint32_t inRowOffsetThisSubBlock = subBlockIdx * inRowSplitSubBlock; + uint32_t outRowOffsetThisSubBlock = (qNBlockSize == 1) ? + inRowOffsetThisSubBlock : + 0; + uint32_t outColOffsetThisSubBlock = (qNBlockSize == 1) ? + 0 : + subBlockIdx * qNSplitSubBlock * embed; + + if (inRowActualThisSubBlock > 0) { + int64_t offsetOutput = + layoutOutput.GetOffset(MatrixCoord(outRowOffsetThisSubBlock, outColOffsetThisSubBlock)); + auto gOutputThisSubBlock = gOutput[offsetOutput]; + auto layoutOutputThisSubBlock = layoutOutput; + + int64_t offsetInput = layoutInput.GetOffset(MatrixCoord(inRowOffsetThisSubBlock, 0)); + auto gInputThisSubBlock = gInput[offsetInput]; + auto layoutInputThisSubBlock = layoutInput.GetTileLayout(MatrixCoord(inRowActualThisSubBlock, embed)); + SubCoreCompute( + gOutputThisSubBlock, + gInputThisSubBlock, + layoutOutputThisSubBlock, + layoutInputThisSubBlock, + qNThisSubBlock, isFirstStackTile, isLastStackTile, curStackTileMod); + } + } + +private: + AscendC::LocalTensor loUbTensor; + AscendC::LocalTensor dmUbTensor; + AscendC::LocalTensor hmUbTensor; + AscendC::LocalTensor glUbTensor; + AscendC::LocalTensor tvUbTensor; + AscendC::LocalTensor goUbTensor16; + AscendC::LocalTensor goUbTensor32; +}; +} // namespace Catlass::Epilogue::Block + +#endif // CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_RESCALE_O_NO_SPLIT_ROW_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_w4a4_per_token_per_channel_dequant.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_w4a4_per_token_per_channel_dequant.hpp new file mode 100644 index 0000000000..2d2e768c55 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_w4a4_per_token_per_channel_dequant.hpp @@ -0,0 +1,252 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ +#ifndef CATLASS_EPILOGUE_BLOCK_EPILOGUE_W4A4_HPP +#define CATLASS_EPILOGUE_BLOCK_EPILOGUE_W4A4_HPP + +#include "catlass/arch/kda_resource.hpp" +#include "catlass/kda_catlass.hpp" +#include "catlass/detail/kda_callback.hpp" +#include "catlass/epilogue/kda_dispatch_policy.hpp" +#include "catlass/gemm/kda_gemm_type.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/layout/kda_layout.hpp" +#include "catlass/kda_matrix_coord.hpp" + +namespace Catlass::Epilogue::Block { + +template < + class CType_, + class PerTokenScaleType_, + class DType_, + class TileBroadcastOneBlk_, + class TileOneBlkColumnBroadcastMul_, + class TileCopy_, + class EpilogueTileSwizzle_> +class BlockEpilogue< + EpilogueAtlasA2W4A4PerTokenPerChannelDequant, + CType_, + PerTokenScaleType_, + DType_, + TileBroadcastOneBlk_, + TileOneBlkColumnBroadcastMul_, + TileCopy_, + EpilogueTileSwizzle_> { +public: + using DispatchPolicy = EpilogueAtlasA2W4A4PerTokenPerChannelDequant; + using ArchTag = typename DispatchPolicy::ArchTag; + + // Data infos + using ElementC = typename CType_::Element; + using LayoutC = typename CType_::Layout; + using ElementPerTokenScale = typename PerTokenScaleType_::Element; + using LayoutPerTokenScale = typename PerTokenScaleType_::Layout; + using ElementD = typename DType_::Element; + using LayoutD = typename DType_::Layout; + + // Tile compute ops + using TileBroadcastOneBlk = TileBroadcastOneBlk_; + using TileOneBlkColumnBroadcastMul = TileOneBlkColumnBroadcastMul_; + + // Tile copy + using CopyGmToUbC = typename TileCopy_::CopyGmToUbC; + using CopyGmToUbPerTokenScale = typename TileCopy_::CopyGmToUbPerTokenScale; + using CopyUbToGmD = typename TileCopy_::CopyUbToGmD; + + using EpilogueTileSwizzle = EpilogueTileSwizzle_; + using TileShape = typename TileOneBlkColumnBroadcastMul::TileShape; + + // Check + static_assert( + std::is_same_v && std::is_same_v && std::is_same_v, + "Accumulator type (`ElementC`) should be half, output type (`ElementD`) should be bfloat16_t while " + "perTokenScale type should be float"); + static_assert( + std::is_same_v && std::is_same_v && std::is_same_v, + "The layout template parameters of BlockEpilogue are wrong"); + + static_assert( + TileShape::ROW == TileBroadcastOneBlk::COMPUTE_LENGTH, + "TileShape must be consistent for all tile compute ops"); + + static_assert(TileShape::ROW % 16 == 0, "Epilogue TileShape::ROW must be divisible by 16"); + + static_assert( + /*ubC*/ TileShape::COUNT * sizeof(ElementC) + /*ubPerTokenScale*/ TileShape::ROW * sizeof(ElementPerTokenScale) + /*ubD*/ TileShape::COUNT * sizeof(ElementD) + /*ubCFp32 & ubPerTokenScaleBrcb*/ 2 * TileShape::COUNT * sizeof(float) <= ArchTag::UB_SIZE, + "TileShape is too large to fit in UB"); + + struct Params { + __gm__ ElementPerTokenScale *ptrPerTokenScale{nullptr}; + LayoutPerTokenScale layoutPerTokenScale{}; + __gm__ ElementD *ptrD{nullptr}; + LayoutD layoutD{}; + + CATLASS_DEVICE + Params() {}; + + CATLASS_DEVICE + Params( + __gm__ ElementPerTokenScale *ptrPerTokenScale_, + LayoutPerTokenScale const &layoutPerTokenScale_, + __gm__ ElementD *ptrD_, + LayoutD const &layoutD_) + : ptrPerTokenScale(ptrPerTokenScale_), + layoutPerTokenScale(layoutPerTokenScale_), + ptrD(ptrD_), + layoutD(layoutD_) + { + } + }; + + CATLASS_DEVICE + BlockEpilogue(Arch::Resource const &resource, Params const ¶ms = Params{}) + : params(params) + { + size_t ubCOffset = 0; + size_t ubPerTokenScaleOffset = ubCOffset + TileShape::COUNT * sizeof(ElementC); + size_t ubDOffset = ubPerTokenScaleOffset + TileShape::ROW * sizeof(ElementPerTokenScale); + size_t ubOffset1 = ubDOffset + TileShape::COUNT * sizeof(ElementD); // for `ubCFp32` & `ubDFp32` + size_t ubOffset2 = ubOffset1 + TileShape::COUNT * sizeof(float); // for `ubPerTokenScaleBrcb` + // for `ubPerTokenScaleBrcb` + + // Load from GM + ubC = resource.ubBuf.template GetBufferByByte(ubCOffset); + ubPerTokenScale = resource.ubBuf.template GetBufferByByte(ubPerTokenScaleOffset); + ubD = resource.ubBuf.template GetBufferByByte(ubDOffset); + + // Caculation (with perToken scale) + ubCFp32 = resource.ubBuf.template GetBufferByByte(ubOffset1); + ubPerTokenScaleBrcb = resource.ubBuf.template GetBufferByByte(ubOffset2); + ubDFp32 = resource.ubBuf.template GetBufferByByte(ubOffset1); + + AscendC::SetFlag(eventID); + } + + CATLASS_DEVICE + ~BlockEpilogue() + { + AscendC::WaitFlag(eventID); + } + + CATLASS_DEVICE + void UpdateParams(Params const ¶ms_) + { + params = params_; + } + + CATLASS_DEVICE + void operator()( + GemmCoord const &blockShapeMNK, + GemmCoord const &blockCoordMNK, + GemmCoord const &actualBlockShapeMNK, + AscendC::GlobalTensor const &gmBlockC, + LayoutC const &layoutBlockC, + Callback &&callback = Callback{}) + { + if (actualBlockShapeMNK.k() == 0) { + return; + } + callback(); + + // Calculate the offset of the current block + MatrixCoord blockShape = blockShapeMNK.GetCoordMN(); + MatrixCoord blockCoord = blockCoordMNK.GetCoordMN(); + MatrixCoord actualBlockShape = actualBlockShapeMNK.GetCoordMN(); + MatrixCoord blockOffset = blockCoord * blockShape; + + AscendC::GlobalTensor gmPerTokenScale; + gmPerTokenScale.SetGlobalBuffer(params.ptrPerTokenScale); + AscendC::GlobalTensor gmD; + gmD.SetGlobalBuffer(params.ptrD); + + auto ubTileStride = MakeCoord(static_cast(TileShape::COLUMN), 1L); + auto tileShape = TileShape::ToCoord(); + EpilogueTileSwizzle epilogueTileSwizzle(actualBlockShape, tileShape); + uint32_t tileLoops = epilogueTileSwizzle.GetLoops(); + uint32_t subblockIdx = AscendC::GetSubBlockIdx(); + uint32_t subblockNum = AscendC::GetSubBlockNum(); + for (uint32_t loopIdx = subblockIdx; loopIdx < tileLoops; loopIdx += subblockNum) { + auto tileCoord = epilogueTileSwizzle.GetTileCoord(loopIdx); + auto actualTileShape = epilogueTileSwizzle.GetActualTileShape(tileCoord); + auto tileOffsetInBlock = tileCoord * tileShape; + auto tileOffset = blockOffset + tileOffsetInBlock; + + // for GMc + auto gmTileC = gmBlockC[layoutBlockC.GetOffset(tileOffsetInBlock)]; + auto layoutGmTileC = layoutBlockC.GetTileLayout(actualTileShape); + + LayoutC layoutUbC{actualTileShape, ubTileStride}; + + auto calcNum = actualTileShape.row() * TileShape::COLUMN; + + AscendC::WaitFlag(eventID); + copyGmToUbC(ubC, gmTileC, layoutUbC, layoutGmTileC); + + auto perTokenScaleTileOffset = tileOffset.template GetCoordByAxis<0>(); + auto perTokenScaleTileShape = actualTileShape.template GetCoordByAxis<0>(); + + auto gmTilePerTokenScale = gmPerTokenScale[params.layoutPerTokenScale.GetOffset(perTokenScaleTileOffset)]; + auto layoutGmTilePerTokenScale = params.layoutPerTokenScale.GetTileLayout(perTokenScaleTileShape); + + auto layoutUbPerTokenScale = LayoutPerTokenScale::template MakeLayoutInUb( + perTokenScaleTileShape); + + copyGmToUbPerTokenScale( + ubPerTokenScale, gmTilePerTokenScale, layoutUbPerTokenScale, layoutGmTilePerTokenScale); + AscendC::SetFlag(eventID); + + AscendC::WaitFlag(eventID); + AscendC::Cast(/*float*/ ubCFp32, /*half*/ ubC, AscendC::RoundMode::CAST_NONE, calcNum); + AscendC::PipeBarrier(); + + tileBroadcastOneBlk(ubPerTokenScaleBrcb, ubPerTokenScale); + AscendC::PipeBarrier(); + + tileOneBlkColumnBroadcastMul(ubDFp32, ubCFp32, ubPerTokenScaleBrcb); + AscendC::PipeBarrier(); + + AscendC::Cast(ubD, ubDFp32, AscendC::RoundMode::CAST_RINT, calcNum); + AscendC::SetFlag(eventID); + + auto resTileShape = MatrixCoord{actualTileShape.row(), actualTileShape.column()}; + LayoutD layoutUbD{resTileShape, ubTileStride}; + + auto gmTileD = gmD[params.layoutD.GetOffset(tileOffset)]; + auto layoutGmTileD = params.layoutD.GetTileLayout(resTileShape); + + AscendC::WaitFlag(eventID); + copyUbToGmD(gmTileD, ubD, layoutGmTileD, layoutUbD); + AscendC::SetFlag(eventID); + } + } + +private: + Params params; + + AscendC::LocalTensor ubC; + AscendC::LocalTensor ubPerTokenScale; + AscendC::LocalTensor ubD; + AscendC::LocalTensor ubCFp32; + AscendC::LocalTensor ubPerTokenScaleBrcb; + AscendC::LocalTensor ubDFp32; + // [gmTileC] -> ubC --(Cast:float)--> ubCFp32 --(Muls)--> ubDFp32 --(Cast:bf16)--> ubD -> [gmTileD] + + int32_t eventID{0}; + + TileBroadcastOneBlk tileBroadcastOneBlk; + TileOneBlkColumnBroadcastMul tileOneBlkColumnBroadcastMul; + + CopyGmToUbC copyGmToUbC; + CopyGmToUbPerTokenScale copyGmToUbPerTokenScale; + CopyUbToGmD copyUbToGmD; +}; + +} // namespace Catlass::Epilogue::Block + +#endif // CATLASS_EPILOGUE_BLOCK_EPILOGUE_W4A4_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/kda_block_epilogue.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/kda_block_epilogue.hpp new file mode 100644 index 0000000000..47185bcf5c --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/kda_block_epilogue.hpp @@ -0,0 +1,51 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_BLOCK_EPILOGUE_HPP +#define CATLASS_KDA_BLOCK_EPILOGUE_HPP + +#include "catlass/kda_catlass.hpp" + +namespace Catlass::Epilogue::Block { + +template < + class DispatchPolicy, + class... Args> +class BlockEpilogue { + static_assert(DEPENDENT_FALSE, "Could not find an epilogue specialization"); +}; + +} // namespace Catlass::Epilogue::Block + +#include "catlass/epilogue/block/block_epilogue_elemwise_no_source.hpp" +#include "catlass/epilogue/block/kda_block_epilogue_elemwise_one_source.hpp" +#include "catlass/epilogue/block/kda_block_epilogue_fa_softmax.hpp" +#include "catlass/epilogue/block/block_epilogue_fa_rescale_o.hpp" +#include "catlass/epilogue/block/kda_block_epilogue_mla_softmax.hpp" +#include "catlass/epilogue/block/block_epilogue_mla_rescale_o.hpp" +#include "catlass/epilogue/block/block_epilogue_mla_fd_rescale_o.hpp" +#include "catlass/epilogue/block/kda_block_epilogue_per_token_dequant.hpp" +#include "catlass/epilogue/block/block_epilogue_per_token_dequant_tla.hpp" +#include "catlass/epilogue/block/kda_block_epilogue_gemm.hpp" +#include "catlass/epilogue/block/kda_block_epilogue_gemv.hpp" +#include "catlass/epilogue/block/kda_block_epilogue_mla_tp1_softmax.hpp" +#include "catlass/epilogue/block/block_epilogue_mla_tp1_rescale_o.hpp" +#include "catlass/epilogue/block/block_epilogue_online_softmax_no_mask.hpp" +#include "catlass/epilogue/block/block_epilogue_rescale_o_no_split_row.hpp" +#include "catlass/epilogue/block/block_epilogue_w4a4_per_token_per_channel_dequant.hpp" + +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 3510) +#include "catlass/epilogue/block/block_epilogue_fa_softmax_ascend950.hpp" +#include "catlass/epilogue/block/block_epilogue_fa_rescale_o_ascend950.hpp" +#include "catlass/epilogue/block/block_epilogue_fixpipe.hpp" +#include "catlass/epilogue/block/block_epilogue_per_group_per_block.hpp" +#endif + +#endif // CATLASS_KDA_BLOCK_EPILOGUE_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/kda_block_epilogue_elemwise_one_source.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/kda_block_epilogue_elemwise_one_source.hpp new file mode 100644 index 0000000000..1f2b6e4283 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/kda_block_epilogue_elemwise_one_source.hpp @@ -0,0 +1,185 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_BLOCK_EPILOGUE_ELEMWISE_ONE_SOURCE_HPP +#define CATLASS_KDA_BLOCK_EPILOGUE_ELEMWISE_ONE_SOURCE_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/epilogue/kda_dispatch_policy.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/kda_matrix_coord.hpp" +#include "catlass/layout/kda_layout.hpp" + +namespace Catlass::Epilogue::Block { + +template < + class CType_, + class XType_, + class DType_, + class TileElemWiseEpilogue_, + class TileCopy_> +class BlockEpilogue< + EpilogueAtlasA2ElemWiseOneSource, + CType_, + XType_, + DType_, + TileElemWiseEpilogue_, + TileCopy_> { +public: + // Type aliases + using DispatchPolicy = EpilogueAtlasA2ElemWiseOneSource; + using ArchTag = typename DispatchPolicy::ArchTag; + using ElementC = typename CType_::Element; + using LayoutC = typename CType_::Layout; + using ElementX = typename XType_::Element; + using LayoutX = typename XType_::Layout; + using ElementD = typename DType_::Element; + using LayoutD = typename DType_::Layout; + using TileElemWiseEpilogue = TileElemWiseEpilogue_; + using CopyGmToUbC = typename TileCopy_::CopyGmToUbC; + using CopyGmToUbX = typename TileCopy_::CopyGmToUbX; + using CopyUbToGmD = typename TileCopy_::CopyUbToGmD; + + static constexpr uint32_t COMPUTE_LENGTH = TileElemWiseEpilogue::COMPUTE_LENGTH; + static constexpr uint32_t OPERANDS_NUM = DispatchPolicy::OPERANDS_NUM; + + // Check the element type of C, X and D + static_assert(std::is_same_v && std::is_same_v, + "Element type of C, X and D must be the same"); + using ElementCompute = ElementD; + + // Check the layout type of C, X and D + static_assert(std::is_same_v && std::is_same_v && + std::is_same_v, + "Layout type of C, X and D must be RowMajor"); + using LayoutComputeInUb = layout::RowMajor; + + // Check if ArchTag is matched + static_assert(std::is_same_v, "Tile epilogue's ArchTag mismatch"); + // Check if compute length is valid + static_assert(COMPUTE_LENGTH * OPERANDS_NUM * sizeof(ElementCompute) <= ArchTag::UB_SIZE, "UB out of bounds"); + + // Epilogue params definition + struct Params { + GM_ADDR ptrX; + LayoutX layoutX; + GM_ADDR ptrD; + LayoutD layoutD; + + CATLASS_HOST_DEVICE + Params() {} + + CATLASS_HOST_DEVICE + Params(GM_ADDR ptrX_, LayoutX const &layoutX_, GM_ADDR ptrD_, LayoutD const &layoutD_) + : ptrX(ptrX_), + layoutX(layoutX_), + ptrD(ptrD_), + layoutD(layoutD_) + {} + }; + + CATLASS_DEVICE + BlockEpilogue(Arch::Resource &resource, Params const ¶ms) + : params(params) + { + ubC = resource.ubBuf.template GetBufferByByte(0); + ubX = resource.ubBuf.template GetBufferByByte(COMPUTE_LENGTH * sizeof(ElementC)); + ubD = resource.ubBuf.template GetBufferByByte( + COMPUTE_LENGTH * sizeof(ElementC) + COMPUTE_LENGTH * sizeof(ElementX)); + + AscendC::SetFlag(EVENT_ID0); + AscendC::SetFlag(EVENT_ID0); + } + + CATLASS_DEVICE + ~BlockEpilogue() + { + AscendC::WaitFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + } + + CATLASS_DEVICE + void operator()( + GemmCoord const &blockShapeMNK, + GemmCoord const &blockCoordMNK, + GemmCoord const &actualBlockShapeMNK, + AscendC::GlobalTensor const &gmBlockC, + LayoutX const &layoutBlockC) + { + // Calculate the offset of the current block + MatrixCoord blockShape = blockShapeMNK.GetCoordMN(); + MatrixCoord blockCoord = blockCoordMNK.GetCoordMN(); + MatrixCoord actualBlockShape = actualBlockShapeMNK.GetCoordMN(); + MatrixCoord blockOffset = blockCoord * blockShape; + + // Calculate the offset and the shape of the current subblock + MatrixCoord subblockShape{ + CeilDiv(actualBlockShape.row(), static_cast(AscendC::GetSubBlockNum())), + actualBlockShape.column()}; + MatrixCoord subblockCoord{AscendC::GetSubBlockIdx(), 0}; + MatrixCoord actualSubblockShape = MatrixCoord::Min( + subblockShape, actualBlockShape - subblockCoord * subblockShape); + MatrixCoord subblockOffset = subblockCoord * subblockShape; + + // Get the data and layout of C + auto gmSubblockC = gmBlockC[layoutBlockC.GetOffset(subblockOffset)]; + auto layoutSubblockC = layoutBlockC.GetTileLayout(actualSubblockShape); + + // Get the data and layout of X + AscendC::GlobalTensor gmX; + gmX.SetGlobalBuffer(reinterpret_cast<__gm__ ElementX *>(params.ptrX)); + auto gmSubblockX = gmX[params.layoutX.GetOffset(blockOffset + subblockOffset)]; + auto layoutSubblockX = params.layoutX.GetTileLayout(actualSubblockShape); + + // Get the data and layout of D + AscendC::GlobalTensor gmD; + gmD.SetGlobalBuffer(reinterpret_cast<__gm__ ElementD *>(params.ptrD)); + auto gmSubblockD = gmD[params.layoutD.GetOffset(blockOffset + subblockOffset)]; + auto layoutSubblockD = params.layoutD.GetTileLayout(actualSubblockShape); + + // Get the layout on UB + auto layoutComputeInUb = LayoutComputeInUb::template MakeLayoutInUb(actualSubblockShape); + + // Copy the data of C and X + AscendC::WaitFlag(EVENT_ID0); + copyGmToUbC(ubC, gmSubblockC, layoutComputeInUb, layoutSubblockC); + copyGmToUbX(ubX, gmSubblockX, layoutComputeInUb, layoutSubblockX); + AscendC::SetFlag(EVENT_ID0); + + // Perform epilogue calculation + AscendC::WaitFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + tileEpilogue(ubD, ubC, ubX); + AscendC::SetFlag(EVENT_ID0); + AscendC::SetFlag(EVENT_ID0); + + // Copy the data of D + AscendC::WaitFlag(EVENT_ID0); + copyUbToGmD(gmSubblockD, ubD, layoutSubblockD, layoutComputeInUb); + AscendC::SetFlag(EVENT_ID0); + } + +private: + Params params; + + AscendC::LocalTensor ubC; + AscendC::LocalTensor ubX; + AscendC::LocalTensor ubD; + + TileElemWiseEpilogue tileEpilogue; + CopyGmToUbC copyGmToUbC; + CopyGmToUbX copyGmToUbX; + CopyUbToGmD copyUbToGmD; +}; + +} // namespace Catlass::Epilogue::Block + +#endif // CATLASS_KDA_BLOCK_EPILOGUE_ELEMWISE_ONE_SOURCE_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/kda_block_epilogue_fa_softmax.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/kda_block_epilogue_fa_softmax.hpp new file mode 100644 index 0000000000..8b1ecfa270 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/kda_block_epilogue_fa_softmax.hpp @@ -0,0 +1,403 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_BLOCK_EPILOGUE_FA_SOFTMAX_HPP +#define CATLASS_KDA_BLOCK_EPILOGUE_FA_SOFTMAX_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_cross_core_sync.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/epilogue/kda_dispatch_policy.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/kda_matrix_coord.hpp" +#include "catlass/epilogue/tile/kda_epilogue_tile_copy.hpp" + +namespace Catlass::Epilogue::Block { + +template < + class OutputType_, + class InputType_, + class MaskType_> +class BlockEpilogue< + EpilogueAtlasA2FASoftmax, + OutputType_, + InputType_, + MaskType_> { +public: + // Type aliases + using DispatchPolicy = EpilogueAtlasA2FASoftmax; + using ArchTag = typename DispatchPolicy::ArchTag; + using ElementOutput = typename OutputType_::Element; + using ElementInput = typename InputType_::Element; + using ElementMask = typename MaskType_::Element; + + using LayoutOutput = typename OutputType_::Layout; + using LayoutInput = typename InputType_::Layout; + using LayoutMask = typename MaskType_::Layout; + + using CopyGmToUbInput = Tile::CopyGm2Ub; + using CopyGmToUbMask = Tile::CopyGm2Ub; + using CopyUbToGmOutput = Tile::CopyUb2Gm; + + static constexpr uint32_t HALF_ELENUM_PER_BLK = 16; + static constexpr uint32_t FLOAT_ELENUM_PER_BLK = 8; + static constexpr uint32_t HALF_ELENUM_PER_VECCALC = 128; + static constexpr uint32_t FLOAT_ELENUM_PER_VECCALC = 64; + static constexpr uint32_t UB_TILE_SIZE = 16384; // 64 * 128 * 2B + static constexpr uint32_t UB_LINE_SIZE = 512; // 128 * 2 * 2B + static constexpr uint32_t HALF_ELENUM_PER_LINE = 256; // 128 * 2 + static constexpr uint32_t FLOAT_ELENUM_PER_LINE = 128; // 128 + static constexpr uint32_t MULTIPLIER = 2; + + CATLASS_DEVICE + BlockEpilogue(Arch::Resource &resource, half tor_) + { + constexpr uint32_t LS32_UB_TENSOR_OFFSET = 2 * UB_TILE_SIZE; + constexpr uint32_t MASK_UB_TENSOR_OFFSET = 4 * UB_TILE_SIZE; + constexpr uint32_t LM_UB_TENSOR_OFFSET = 7 * UB_TILE_SIZE; + constexpr uint32_t HM_UB_TENSOR_OFFSET = 7 * UB_TILE_SIZE + 1 * UB_LINE_SIZE; + constexpr uint32_t GM_UB_TENSOR_OFFSET = 7 * UB_TILE_SIZE + 2 * UB_LINE_SIZE; + constexpr uint32_t DM_UB_TENSOR_OFFSET = 7 * UB_TILE_SIZE + 4 * UB_LINE_SIZE; + constexpr uint32_t LL_UB_TENSOR_OFFSET = 7 * UB_TILE_SIZE + 6 * UB_LINE_SIZE; + constexpr uint32_t TV_UB_TENSOR_OFFSET = 7 * UB_TILE_SIZE + 11 * UB_LINE_SIZE; + + tor = tor_; + lsUbTensor = resource.ubBuf.template GetBufferByByte(0); + lpUbTensor = resource.ubBuf.template GetBufferByByte(0); + ls32UbTensor = resource.ubBuf.template GetBufferByByte(LS32_UB_TENSOR_OFFSET); + maskUbTensor = resource.ubBuf.template GetBufferByByte(MASK_UB_TENSOR_OFFSET); + lmUbTensor = resource.ubBuf.template GetBufferByByte(LM_UB_TENSOR_OFFSET); + hmUbTensor = resource.ubBuf.template GetBufferByByte(HM_UB_TENSOR_OFFSET); + gmUbTensor = resource.ubBuf.template GetBufferByByte(GM_UB_TENSOR_OFFSET); + dmUbTensor = resource.ubBuf.template GetBufferByByte(DM_UB_TENSOR_OFFSET); + llUbTensor = resource.ubBuf.template GetBufferByByte(LL_UB_TENSOR_OFFSET); + tvUbTensor = resource.ubBuf.template GetBufferByByte(TV_UB_TENSOR_OFFSET); + + AscendC::SetFlag(EVENT_ID0); + AscendC::SetFlag(EVENT_ID0); + AscendC::SetFlag(EVENT_ID1); + } + + CATLASS_DEVICE + ~BlockEpilogue() + { + AscendC::WaitFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID1); + } + + CATLASS_DEVICE + void SetMask(int32_t len) + { + const int32_t MAX_MASK_LEN = 128; + const int32_t HALF_MASK_LEN = 64; + if (len >= MAX_MASK_LEN) { + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + return; + } + int32_t highMask = len - HALF_MASK_LEN > 0 ? len - HALF_MASK_LEN : 0; + int32_t lowMask = len - HALF_MASK_LEN >= 0 ? HALF_MASK_LEN : len; + if (len < HALF_MASK_LEN) { + AscendC::SetVectorMask(0x0, ((uint64_t)1 << lowMask) - 1); + } else { + AscendC::SetVectorMask(((uint64_t)1 << highMask) - 1, 0xffffffffffffffff); + } + } + + CATLASS_DEVICE + void SetVcgMask(int32_t len) + { + const int32_t MAX_LEN = 16; + if (len > MAX_LEN) { + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + return; + } + uint64_t subMask = ((uint64_t)1 << len) - 1; + uint64_t maskValue = (subMask << 48) + (subMask << 32) + (subMask << 16) + subMask; + AscendC::SetVectorMask(maskValue, maskValue); + } + + CATLASS_DEVICE + void subCoreCompute( + AscendC::GlobalTensor gOutput, + AscendC::GlobalTensor gInput, + AscendC::GlobalTensor gMask, + const LayoutOutput &layoutOutput, + const LayoutInput &layoutInput, + const LayoutMask &layoutMask, + uint32_t nIdx, Arch::CrossCoreFlag qkReady) + { + uint32_t subM = layoutInput.shape(0); + uint32_t qkN = layoutInput.shape(1); + uint32_t qkNRound = layoutInput.stride(0); + uint32_t maxSeqlen = layoutMask.stride(0); + uint32_t offset = pingpongFlag * UB_TILE_SIZE / sizeof(ElementInput); + uint32_t subMAligned128 = (subM + HALF_ELENUM_PER_VECCALC - 1) / HALF_ELENUM_PER_VECCALC; + uint32_t subMRound = (subM + HALF_ELENUM_PER_BLK - 1) / HALF_ELENUM_PER_BLK * HALF_ELENUM_PER_BLK; + + // Get the layout on UB + auto layoutInUb = LayoutInput::template MakeLayoutInUb(MatrixCoord{subM, qkN}); + + if (subM > 0) { + // Copy mask + AscendC::WaitFlag(EVENT_ID0); + copyGmToUbMask(maskUbTensor, gMask, layoutInUb, layoutMask); + } + Arch::CrossCoreWaitFlag(qkReady); + if (subM > 0) { + AscendC::WaitFlag(pingpongFlag); + // Copy QK + copyGmToUbInput(lsUbTensor[offset], gInput, layoutInUb, layoutInput); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + + // ls = tor * ls + AscendC::Muls( + lsUbTensor[offset], + lsUbTensor[offset], + tor, + (uint64_t)0, + (subM * qkNRound + HALF_ELENUM_PER_VECCALC - 1) / HALF_ELENUM_PER_VECCALC, + AscendC::UnaryRepeatParams(1, 1, 8, 8)); + AscendC::PipeBarrier(); + // ls = ls + mask + AscendC::Add( + lsUbTensor[offset], + lsUbTensor[offset], + maskUbTensor, + (uint64_t)0, + (subM * qkNRound + HALF_ELENUM_PER_VECCALC - 1) / HALF_ELENUM_PER_VECCALC, + AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); + AscendC::PipeBarrier(); + AscendC::SetFlag(EVENT_ID0); + // lm = rowmax(ls) + if (qkN <= HALF_ELENUM_PER_VECCALC) { + SetMask(qkN); + AscendC::BlockReduceMax( + tvUbTensor.ReinterpretCast(), + lsUbTensor[offset], + subM, + 0, + 2, + 1, + qkNRound / HALF_ELENUM_PER_BLK); + AscendC::PipeBarrier(); + SetVcgMask(qkNRound / HALF_ELENUM_PER_BLK); + AscendC::BlockReduceMax( + lmUbTensor, + tvUbTensor.ReinterpretCast(), + (subM * HALF_ELENUM_PER_BLK + HALF_ELENUM_PER_VECCALC - 1) / HALF_ELENUM_PER_VECCALC, + 0, + 1, + 1, + 8); + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + AscendC::PipeBarrier(); + if (nIdx == 0) { + // hm = lm + AscendC::DataCopy( + hmUbTensor, lmUbTensor, AscendC::DataCopyParams(1, subMRound / HALF_ELENUM_PER_BLK, 0, 0)); + AscendC::PipeBarrier(); + } else { + // hm = vmax(lm, gm) + AscendC::Max( + hmUbTensor, + lmUbTensor, + gmUbTensor, + (uint64_t)0, + subMAligned128, + AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); + AscendC::PipeBarrier(); + // dm = gm - hm + AscendC::Sub( + dmUbTensor[nIdx % MULTIPLIER * HALF_ELENUM_PER_LINE], + gmUbTensor, + hmUbTensor, + (uint64_t)0, + subMAligned128, + AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); + AscendC::PipeBarrier(); + } + // gm = hm + AscendC::DataCopy( + gmUbTensor, hmUbTensor, AscendC::DataCopyParams(1, subMRound / HALF_ELENUM_PER_BLK, 0, 0)); + AscendC::PipeBarrier(); + // hm_block = brcb(hm), 存放于tv + AscendC::Brcb( + tvUbTensor.ReinterpretCast(), + hmUbTensor.ReinterpretCast(), + subMRound / FLOAT_ELENUM_PER_BLK, + AscendC::BrcbRepeatParams(1, 8)); + AscendC::PipeBarrier(); + // ls = ls - hm_block + for (uint32_t vsubIdx = 0; vsubIdx < qkN / HALF_ELENUM_PER_VECCALC; vsubIdx++) { + AscendC::Sub( + lsUbTensor[offset + vsubIdx * HALF_ELENUM_PER_VECCALC], + lsUbTensor[offset + vsubIdx * HALF_ELENUM_PER_VECCALC], + tvUbTensor.ReinterpretCast(), + (uint64_t)0, + subM, + AscendC::BinaryRepeatParams( + 1, 1, 0, qkNRound / HALF_ELENUM_PER_BLK, qkNRound / HALF_ELENUM_PER_BLK, 1)); + } + if (qkN % HALF_ELENUM_PER_VECCALC > 0) { + SetMask(qkN % HALF_ELENUM_PER_VECCALC); + AscendC::Sub( + lsUbTensor[offset + qkN / HALF_ELENUM_PER_VECCALC * HALF_ELENUM_PER_VECCALC], + lsUbTensor[offset + qkN / HALF_ELENUM_PER_VECCALC * HALF_ELENUM_PER_VECCALC], + tvUbTensor.ReinterpretCast(), + (uint64_t)0, + subM, + AscendC::BinaryRepeatParams( + 1, 1, 0, qkNRound / HALF_ELENUM_PER_BLK, qkNRound / HALF_ELENUM_PER_BLK, 1)); + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + AscendC::PipeBarrier(); + // ls32 = castfp16to32(ls) + AscendC::Cast( + ls32UbTensor, + lsUbTensor[offset], + AscendC::RoundMode::CAST_NONE, + (uint64_t)0, + (subM * qkNRound + FLOAT_ELENUM_PER_VECCALC - 1) / FLOAT_ELENUM_PER_VECCALC, + AscendC::UnaryRepeatParams(1, 1, 8, 4)); + AscendC::PipeBarrier(); + // ls32 = exp(ls32) + AscendC::Exp( + ls32UbTensor, + ls32UbTensor, + (uint64_t)0, + (subM * qkNRound + FLOAT_ELENUM_PER_VECCALC - 1) / FLOAT_ELENUM_PER_VECCALC, + AscendC::UnaryRepeatParams(1, 1, 8, 8)); + AscendC::PipeBarrier(); + // lp = castfp32to16(ls) + AscendC::Cast( + lpUbTensor[offset], + ls32UbTensor, + AscendC::RoundMode::CAST_NONE, + (uint64_t)0, + (subM * qkNRound + FLOAT_ELENUM_PER_VECCALC - 1) / FLOAT_ELENUM_PER_VECCALC, + AscendC::UnaryRepeatParams(1, 1, 4, 8)); + AscendC::PipeBarrier(); + AscendC::SetFlag(EVENT_ID0); + // ll = rowsum(ls32) + if (qkN <= FLOAT_ELENUM_PER_VECCALC) { + SetMask(qkN); + AscendC::RepeatReduceSum( + llUbTensor[nIdx % MULTIPLIER * FLOAT_ELENUM_PER_LINE], + ls32UbTensor, + subM, 0, 0, 1, 1, qkNRound / FLOAT_ELENUM_PER_BLK); + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + } else { + for (uint32_t vaddIdx = 1; vaddIdx < qkN / FLOAT_ELENUM_PER_VECCALC; vaddIdx++) { + AscendC::Add( + ls32UbTensor, + ls32UbTensor, + ls32UbTensor[vaddIdx * FLOAT_ELENUM_PER_VECCALC], + (uint64_t)0, + subM, + AscendC::BinaryRepeatParams( + 1, 1, 1, + qkNRound / FLOAT_ELENUM_PER_BLK, + qkNRound / FLOAT_ELENUM_PER_BLK, + qkNRound / FLOAT_ELENUM_PER_BLK)); + AscendC::PipeBarrier(); + } + if (qkN % FLOAT_ELENUM_PER_VECCALC > 0) { + SetMask(qkN % FLOAT_ELENUM_PER_VECCALC); + AscendC::Add( + ls32UbTensor, + ls32UbTensor, + ls32UbTensor[qkN / FLOAT_ELENUM_PER_VECCALC * FLOAT_ELENUM_PER_VECCALC], + (uint64_t)0, + subM, + AscendC::BinaryRepeatParams( + 1, 1, 1, + qkNRound / FLOAT_ELENUM_PER_BLK, + qkNRound / FLOAT_ELENUM_PER_BLK, + qkNRound / FLOAT_ELENUM_PER_BLK)); + AscendC::PipeBarrier(); + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + AscendC::RepeatReduceSum( + llUbTensor[nIdx % MULTIPLIER * FLOAT_ELENUM_PER_LINE], + ls32UbTensor, + subM, 0, 0, 1, 1, qkNRound / FLOAT_ELENUM_PER_BLK); + } + AscendC::PipeBarrier(); + AscendC::WaitFlag(EVENT_ID0); + copyUbToGmOutput(gOutput, lpUbTensor[offset], layoutOutput, layoutInUb); + AscendC::SetFlag(pingpongFlag); + pingpongFlag = 1 - pingpongFlag; + } + } + + CATLASS_DEVICE + void operator()( + AscendC::GlobalTensor gOutput, + AscendC::GlobalTensor gInput, + AscendC::GlobalTensor gMask, + const LayoutOutput &layoutOutput, + const LayoutInput &layoutInput, + const LayoutMask &layoutMask, + GemmCoord actualBlockShape, + uint32_t nIdx, Arch::CrossCoreFlag qkReady) + { + uint32_t mActual = actualBlockShape.m(); + uint32_t nActual = actualBlockShape.n(); + + uint32_t subBlockIdx = AscendC::GetSubBlockIdx(); + uint32_t subBlockNum = AscendC::GetSubBlockNum(); + + uint32_t mActualPerSubBlock = CeilDiv(mActual, subBlockNum); + uint32_t mActualThisSubBlock = (subBlockIdx == 0) ? mActualPerSubBlock : (mActual - mActualPerSubBlock); + uint32_t mOffset = subBlockIdx * mActualPerSubBlock; + uint32_t nOffset = 0; + + int64_t offsetOutput = layoutOutput.GetOffset(MatrixCoord(mOffset, nOffset)); + auto gOutputThisSubBlock = gOutput[offsetOutput]; + auto layoutOutputThisSubBlock = layoutOutput.GetTileLayout(MatrixCoord(mActualThisSubBlock, nActual)); + + int64_t offsetInput = layoutInput.GetOffset(MatrixCoord(mOffset, nOffset)); + auto gInputThisSubBlock = gInput[offsetInput]; + auto layoutInputThisSubBlock = layoutInput.GetTileLayout(MatrixCoord(mActualThisSubBlock, nActual)); + + int64_t offsetMask = layoutMask.GetOffset(MatrixCoord(mOffset, nOffset)); + auto gMaskThisSubBlock = gMask[offsetMask]; + auto layoutMaskThisSubBlock = layoutMask.GetTileLayout(MatrixCoord(mActualThisSubBlock, nActual)); + + subCoreCompute(gOutputThisSubBlock, gInputThisSubBlock, gMaskThisSubBlock, + layoutOutputThisSubBlock, layoutInputThisSubBlock, layoutMaskThisSubBlock, + nIdx, qkReady); + } + +private: + half tor; + uint32_t pingpongFlag = 0; + AscendC::LocalTensor lsUbTensor; + AscendC::LocalTensor lpUbTensor; + AscendC::LocalTensor ls32UbTensor; + AscendC::LocalTensor maskUbTensor; + AscendC::LocalTensor lmUbTensor; + AscendC::LocalTensor hmUbTensor; + AscendC::LocalTensor gmUbTensor; + AscendC::LocalTensor dmUbTensor; + AscendC::LocalTensor llUbTensor; + AscendC::LocalTensor tvUbTensor; + + CopyGmToUbInput copyGmToUbInput; + CopyGmToUbMask copyGmToUbMask; + CopyUbToGmOutput copyUbToGmOutput; +}; + +} // namespace Catlass::Epilogue::Block + +#endif // CATLASS_KDA_BLOCK_EPILOGUE_FA_SOFTMAX_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/kda_block_epilogue_gemm.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/kda_block_epilogue_gemm.hpp new file mode 100644 index 0000000000..7e7d043ec1 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/kda_block_epilogue_gemm.hpp @@ -0,0 +1,218 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_BLOCK_EPILOGUE_GEMM_HPP +#define CATLASS_KDA_BLOCK_EPILOGUE_GEMM_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/epilogue/kda_dispatch_policy.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/kda_matrix_coord.hpp" +#include "catlass/epilogue/tile/kda_epilogue_tile_copy.hpp" +#include "catlass/gemm/kda_gemm_helper.hpp" + +namespace Catlass::Epilogue::Block { + +template < + class CType_, + class XType_, + class DType_, + class TileElemWiseEpilogueAdd_, + class TileElemWiseEpilogueMuls_, + class TileElemWiseCastD_, + class TileCopy_> +class BlockEpilogue< + EpilogueAtlasA2Gemm, + CType_, + XType_, + DType_, + TileElemWiseEpilogueAdd_, + TileElemWiseEpilogueMuls_, + TileElemWiseCastD_, + TileCopy_> { +public: + // Type aliases + using DispatchPolicy = EpilogueAtlasA2Gemm; + using ArchTag = typename DispatchPolicy::ArchTag; + using ElementC = typename CType_::Element; + using LayoutC = typename CType_::Layout; + using ElementX = typename XType_::Element; + using LayoutX = typename XType_::Layout; + using ElementD = typename DType_::Element; + using LayoutD = typename DType_::Layout; + using TileElemWiseEpilogueAdd = TileElemWiseEpilogueAdd_; + using TileElemWiseEpilogueMuls = TileElemWiseEpilogueMuls_; + using TileElemWiseCastD = TileElemWiseCastD_; + using CopyGmToUbC = typename TileCopy_::CopyGmToUbC; + using CopyGmToUbX = typename TileCopy_::CopyGmToUbX; + using CopyUbToGmD = typename TileCopy_::CopyUbToGmD; + + const uint32_t SubNum = AscendC::GetSubBlockNum(); + const uint32_t UBSize = ArchTag::UB_SIZE; + static constexpr bool isNeedCast = !std::is_same::value; + static constexpr uint32_t COMPUTE_LENGTH = TileElemWiseEpilogueAdd::COMPUTE_LENGTH; + + using ElementCompute = typename Catlass::Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; + using ElementScalar = ElementCompute; + + // Check if ArchTag is matched + static_assert(std::is_same_v, "Tile epilogue's ArchTag mismatch"); + static_assert(std::is_same_v, "Tile epilogue's ArchTag mismatch"); + static_assert(std::is_same_v, "LayoutC only support RowMajor yet!"); + + struct Params { + ElementScalar alpha; + ElementScalar beta; + GM_ADDR ptrX; + LayoutX layoutX; + GM_ADDR ptrD; + LayoutD layoutD; + + CATLASS_HOST_DEVICE + Params() {} + + CATLASS_HOST_DEVICE + Params(ElementScalar alpha_, ElementScalar beta_, GM_ADDR ptrX_, LayoutX layoutX_, GM_ADDR ptrD_, LayoutD layoutD_) + : alpha(alpha_), + beta(beta_), + ptrX(ptrX_), + layoutX(layoutX_), + ptrD(ptrD_), + layoutD(layoutD_) + {} + }; + + CATLASS_DEVICE + BlockEpilogue(Arch::Resource &resource, GemmCoord blockShape_, Params const ¶ms_, uint32_t ubByteStart = 0) + : blockShapeMNK(blockShape_), + params(params_) + { + uint32_t maxMPerBlock = blockShapeMNK.m(); + uint32_t maxNPerBlock = blockShapeMNK.n(); + uint32_t tileSize = maxMPerBlock * maxNPerBlock / SubNum; + uint32_t ubCSize = tileSize * sizeof(ElementC); + uint32_t ubXSize = tileSize * sizeof(ElementX); + uint32_t ubDSize = tileSize * sizeof(ElementD); + uint32_t ubXCastSize = tileSize * sizeof(ElementCompute); + uint32_t ubDCastSize = tileSize * sizeof(ElementCompute); + ubCTensor = resource.ubBuf.template GetBufferByByte(ubByteStart); + ubByteStart += ubCSize; + ubXTensor = resource.ubBuf.template GetBufferByByte(ubByteStart); + ubByteStart += ubXSize; + ubDTensor = resource.ubBuf.template GetBufferByByte(ubByteStart); + ubByteStart += ubDSize; + if constexpr (isNeedCast) { + ubXTensorCast = resource.ubBuf.template GetBufferByByte(ubByteStart); + ubByteStart += ubXCastSize; + ubDTensorCast = resource.ubBuf.template GetBufferByByte(ubByteStart); + ; + ubByteStart += ubDCastSize; + } + AscendC::SetFlag(EVENT_ID0); + AscendC::SetFlag(EVENT_ID1); + AscendC::SetFlag(EVENT_ID0); + } + + CATLASS_DEVICE + ~BlockEpilogue() + { + AscendC::WaitFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID1); + AscendC::WaitFlag(EVENT_ID0); + } + + CATLASS_DEVICE + void operator()( + GemmCoord const &actualShapeMNK, GemmCoord const &blockCoordMNK, + AscendC::GlobalTensor const &gmBlockC, LayoutC const &layoutC, + uint64_t const &offset) + { + AscendC::GlobalTensor gmBlockX; + gmBlockX.SetGlobalBuffer(reinterpret_cast<__gm__ ElementX *>(params.ptrX)); + AscendC::GlobalTensor gmBlockD; + gmBlockD.SetGlobalBuffer(reinterpret_cast<__gm__ ElementD *>(params.ptrD)); + MatrixCoord blockShapeMN = blockShapeMNK.GetCoordMN(); + MatrixCoord blockCoordMN = blockCoordMNK.GetCoordMN(); + MatrixCoord actualShapeMN = actualShapeMNK.GetCoordMN(); + MatrixCoord blockOffset = blockCoordMN * blockShapeMN; + MatrixCoord subblockShape{ + CeilDiv(actualShapeMN.row(), SubNum), + actualShapeMN.column()}; + MatrixCoord subblockCoord{AscendC::GetSubBlockIdx(), 0}; + MatrixCoord actualSubblockShape = MatrixCoord::Min( + subblockShape, actualShapeMN - subblockCoord * subblockShape); + MatrixCoord subblockOffset = subblockCoord * subblockShape; + LayoutC layoutInUb{blockShapeMN.row() / SubNum, blockShapeMN.column()}; + AscendC::WaitFlag(EVENT_ID1); + auto layoutTileX = params.layoutX.GetTileLayout(actualSubblockShape); + auto layoutXInUb = layoutInUb.GetTileLayout(actualSubblockShape); + auto gmTileX = gmBlockX[offset + params.layoutX.GetOffset(blockOffset + subblockOffset)]; + copyGmToUbX(ubXTensor, gmTileX, layoutXInUb, layoutTileX); + AscendC::SetFlag(EVENT_ID1); + AscendC::WaitFlag(EVENT_ID1); + if constexpr (isNeedCast) { + AscendC::Cast(ubXTensorCast, ubXTensor, AscendC::RoundMode::CAST_NONE, COMPUTE_LENGTH); + AscendC::PipeBarrier(); + tileElemWiseEpilogueMuls(ubXTensorCast, ubXTensorCast, (ElementCompute)params.beta); + } else { + tileElemWiseEpilogueMuls(ubXTensor, ubXTensor, (ElementX)params.beta); + } + AscendC::WaitFlag(EVENT_ID0); + auto layoutTileC = layoutC.GetTileLayout(actualSubblockShape); + auto layoutCInUb = layoutInUb.GetTileLayout(actualSubblockShape); + auto gmTileC = gmBlockC[offset + layoutC.GetOffset(blockOffset + subblockOffset)]; + copyGmToUbC(ubCTensor, gmTileC, layoutCInUb, layoutTileC); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + tileElemWiseEpilogueMuls(ubCTensor, ubCTensor, (ElementC)params.alpha); + AscendC::WaitFlag(EVENT_ID0); + AscendC::PipeBarrier(); + if constexpr (isNeedCast) { + tileElemWiseEpilogueAdd(ubDTensorCast, ubCTensor, ubXTensorCast); + } else { + tileElemWiseEpilogueAdd(ubDTensor, ubCTensor, ubXTensor); + } + AscendC::SetFlag(EVENT_ID0); + AscendC::SetFlag(EVENT_ID1); + AscendC::PipeBarrier(); + if constexpr (isNeedCast) { + tileElemWiseCastD(ubDTensor, ubDTensorCast); + } + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + auto layoutDInGm = params.layoutD.GetTileLayout(actualSubblockShape); + auto layoutTileD = layoutInUb.GetTileLayout(actualSubblockShape); + auto gmTileD = gmBlockD[offset + params.layoutD.GetOffset(blockOffset + subblockOffset)]; + copyUbToGmD(gmTileD, ubDTensor, layoutDInGm, layoutTileD); + AscendC::SetFlag(EVENT_ID0); + } + +private: + GemmCoord blockShapeMNK; + Params params; + + AscendC::LocalTensor ubCTensor; + AscendC::LocalTensor ubXTensor; + AscendC::LocalTensor ubDTensor; + AscendC::LocalTensor ubXTensorCast; + AscendC::LocalTensor ubDTensorCast; + + CopyGmToUbC copyGmToUbC; + CopyGmToUbX copyGmToUbX; + CopyUbToGmD copyUbToGmD; + + TileElemWiseEpilogueAdd tileElemWiseEpilogueAdd; + TileElemWiseEpilogueMuls tileElemWiseEpilogueMuls; + TileElemWiseCastD tileElemWiseCastD; +}; +} // namespace Catlass::Epilogue::Block + +#endif // CATLASS_KDA_BLOCK_EPILOGUE_GEMM_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/kda_block_epilogue_gemv.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/kda_block_epilogue_gemv.hpp new file mode 100644 index 0000000000..4f58adc934 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/kda_block_epilogue_gemv.hpp @@ -0,0 +1,222 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_BLOCK_EPILOGUE_GEMV_HPP +#define CATLASS_KDA_BLOCK_EPILOGUE_GEMV_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/epilogue/kda_dispatch_policy.hpp" +#include "catlass/gemm/kda_gemm_helper.hpp" +#include "catlass/gemv/kda_gemv_helper.hpp" +#include "catlass/kda_gemv_coord.hpp" +#include "catlass/layout/kda_layout.hpp" +#include "catlass/kda_matrix_coord.hpp" + +namespace Catlass::Epilogue::Block { + +template < + class CType_, + class YType_, + class ZType_, + class TileElemWiseEpilogueAdd_, + class TileElemWiseEpilogueMuls_, + class TileCopy_> +class BlockEpilogue< + EpilogueAtlasA2Gemv, + CType_, + YType_, + ZType_, + TileElemWiseEpilogueAdd_, + TileElemWiseEpilogueMuls_, + TileCopy_> { +public: + using DispatchPolicy = EpilogueAtlasA2Gemv; + using ArchTag = typename DispatchPolicy::ArchTag; + + using ElementC = typename CType_::Element; + using LayoutC = typename CType_::Layout; + using ElementY = typename YType_::Element; + using LayoutY = typename YType_::Layout; + using ElementZ = typename ZType_::Element; + using LayoutZ = typename ZType_::Layout; + + using TileElemWiseEpilogueAdd = TileElemWiseEpilogueAdd_; + using TileElemWiseEpilogueMuls = TileElemWiseEpilogueMuls_; + + using CopyGmToUbY = typename TileCopy_::CopyGmToUbC; + using CopyGmToubC = typename TileCopy_::CopyGmToUbX; + using CopyUbToGmZ = typename TileCopy_::CopyUbToGmD; + + static constexpr uint32_t COMPUTE_LENGTH = TileElemWiseEpilogueMuls::COMPUTE_LENGTH; + + static constexpr bool isNeedCast = !std::is_same::value; + + using ElementCompute = typename Catlass::Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; + using ElementScalar = ElementCompute; + using TensorCoord = layout::VectorLayout::TensorCoord; + + // check the layout of Y, C and Z + static_assert(std::is_same_v && std::is_same_v && + std::is_same_v, + "Layout type of Y, C and Z must be VectorLayout"); + + using LayoutComputeInUb = layout::VectorLayout; + + // Check if ArchTag is matched + static_assert(std::is_same_v, "Tile epilogue's ArchTag mismatch"); + + struct Params { + ElementScalar alpha; + ElementScalar beta; + GM_ADDR ptrY; + LayoutY layoutY; + GM_ADDR ptrZ; + LayoutZ layoutZ; + + // Methods + CATLASS_HOST_DEVICE + Params() {} + + CATLASS_HOST_DEVICE + Params(ElementScalar alpha_, ElementScalar beta_, GM_ADDR ptrY_, LayoutC layoutY_, GM_ADDR ptrZ_, LayoutZ layoutZ_) + : alpha(alpha_), + beta(beta_), + ptrY(ptrY_), + layoutY(layoutY_), + ptrZ(ptrZ_), + layoutZ(layoutZ_) + {} + }; + + CATLASS_DEVICE + BlockEpilogue(Arch::Resource &resource, Params const ¶ms) + : params(params) + { + ubC = resource.ubBuf.template GetBufferByByte(0); + ubY = resource.ubBuf.template GetBufferByByte(COMPUTE_LENGTH * sizeof(ElementC)); + ubYCast = resource.ubBuf.template GetBufferByByte(COMPUTE_LENGTH * sizeof(ElementC)); + ubZ = resource.ubBuf.template GetBufferByByte( + COMPUTE_LENGTH * sizeof(ElementY) + COMPUTE_LENGTH * sizeof(ElementC)); + ubZCast = resource.ubBuf.template GetBufferByByte( + COMPUTE_LENGTH * sizeof(ElementY) + COMPUTE_LENGTH * sizeof(ElementC)); + AscendC::SetFlag(EVENT_ID0); + } + + CATLASS_DEVICE + ~BlockEpilogue() + { + AscendC::WaitFlag(EVENT_ID0); + } + + CATLASS_DEVICE + void operator()( + TensorCoord const &blockOffsetMN, + TensorCoord const &actualBlockShapeMN, + AscendC::GlobalTensor const &gmBlockC, + LayoutC const &layoutBlockC) + { + TensorCoord actualBlockShape = actualBlockShapeMN; + TensorCoord blockOffset = blockOffsetMN; + + TensorCoord subblockShape{ + CeilDiv(actualBlockShape[0], static_cast(AscendC::GetSubBlockNum()))}; + TensorCoord subblockCoord{static_cast(AscendC::GetSubBlockIdx())}; + + TensorCoord actualSubblockShape = TensorCoord::Min(subblockShape, actualBlockShape - subblockCoord * subblockShape); + TensorCoord subblockOffset = subblockCoord * subblockShape; + + // Get the data and layout of C + auto gmSubblockC = gmBlockC[layoutBlockC.GetOffset(subblockOffset)]; + auto layoutSubblockC = layoutBlockC.GetTileLayout(actualSubblockShape); + + // Get the data and layout of y + AscendC::GlobalTensor gmY; + gmY.SetGlobalBuffer(reinterpret_cast<__gm__ ElementY *>(params.ptrY)); + auto gmSubblockY = gmY[params.layoutY.GetOffset(blockOffset + subblockOffset)]; + auto layoutSubblockY = params.layoutY.GetTileLayout(actualSubblockShape); + + // Get the data and layout of Z + AscendC::GlobalTensor gmZ; + gmZ.SetGlobalBuffer(reinterpret_cast<__gm__ ElementZ *>(params.ptrZ)); + auto gmSubblockZ = gmZ[params.layoutZ.GetOffset(blockOffset + subblockOffset)]; + auto layoutSubblockZ = params.layoutZ.GetTileLayout(actualSubblockShape); + + // get the layout on UB + auto layoutComputeInUb = LayoutComputeInUb::template MakeLayoutInUb(actualSubblockShape); + + // load C(A*x) from gm to ub + AscendC::WaitFlag(EVENT_ID0); + copyGmToubC(ubC, gmSubblockC, layoutComputeInUb, layoutSubblockC); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + + // compute C * alpha + tileEpilogueMul(ubC, ubC, params.alpha); + + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + + // load Y from gm to ub + copyGmToUbY(ubY, gmSubblockY, layoutComputeInUb, layoutSubblockY); + + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + + // compute Y * beta + if constexpr (isNeedCast) { + AscendC::Cast(ubYCast, ubY, AscendC::RoundMode::CAST_NONE, COMPUTE_LENGTH); + AscendC::PipeBarrier(); + tileEpilogueMul(ubYCast, ubYCast, params.beta); + AscendC::PipeBarrier(); + } else { + tileEpilogueMul(ubY, ubY, params.beta); + AscendC::PipeBarrier(); + } + + if constexpr (isNeedCast) { + tileEpilogueAdd(ubZCast, ubC, ubYCast); + } else { + tileEpilogueAdd(ubZ, ubC, ubY); + } + + if constexpr (isNeedCast) { + AscendC::PipeBarrier(); + AscendC::Cast(ubZ, ubZCast, AscendC::RoundMode::CAST_RINT, COMPUTE_LENGTH); + AscendC::PipeBarrier(); + } + + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + + copyUbToGmZ(gmSubblockZ, ubZ, layoutSubblockZ, layoutComputeInUb); + AscendC::SetFlag(EVENT_ID0); + }; + +private: + Params params; + + AscendC::LocalTensor ubY; + AscendC::LocalTensor ubYCast; + AscendC::LocalTensor ubC; + AscendC::LocalTensor ubZ; + AscendC::LocalTensor ubZCast; + + TileElemWiseEpilogueAdd tileEpilogueAdd; + TileElemWiseEpilogueMuls tileEpilogueMul; + + CopyGmToUbY copyGmToUbY; + CopyGmToubC copyGmToubC; + CopyUbToGmZ copyUbToGmZ; +}; + +} // namespace Catlass::Epilogue::Block + +#endif // CATLASS_KDA_BLOCK_EPILOGUE_GEMV_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/kda_block_epilogue_mla_softmax.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/kda_block_epilogue_mla_softmax.hpp new file mode 100644 index 0000000000..46d939296b --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/kda_block_epilogue_mla_softmax.hpp @@ -0,0 +1,441 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_BLOCK_EPILOGUE_MLA_SOFTMAX_HPP +#define CATLASS_KDA_BLOCK_EPILOGUE_MLA_SOFTMAX_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_cross_core_sync.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/epilogue/kda_dispatch_policy.hpp" +#include "catlass/epilogue/tile/kda_epilogue_tile_copy.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/kda_matrix_coord.hpp" + +namespace Catlass::Epilogue::Block { + +template < + class OutputType_, + class InputType_, + class MaskType_> +class BlockEpilogue< + EpilogueAtlasA2MLASoftmax, + OutputType_, + InputType_, + MaskType_> { +public: + // Type aliases + using DispatchPolicy = EpilogueAtlasA2MLASoftmax; + using ArchTag = typename DispatchPolicy::ArchTag; + using ElementOutput = typename OutputType_::Element; + using ElementInput = typename InputType_::Element; + using ElementMask = typename MaskType_::Element; + + using LayoutOutput = typename OutputType_::Layout; + using LayoutInput = typename InputType_::Layout; + using LayoutMask = typename MaskType_::Layout; + + using CopyGmToUbInput = Tile::CopyGm2Ub; + using CopyGmToUbMask = Tile::CopyGm2Ub; + using CopyUbToGmOutput = Tile::CopyUb2Gm; + + static constexpr uint32_t HALF_ELENUM_PER_BLK = 16; + static constexpr uint32_t FLOAT_BLOCK_SIZE = 8; + static constexpr uint32_t HALF_ELENUM_PER_VECCALC = 128; + static constexpr uint32_t FLOAT_VECTOR_SIZE = 64; + static constexpr uint32_t UB_TILE_SIZE = 16384; // 64 * 128 * 2B + static constexpr uint32_t UB_LINE_SIZE = 512; // 128 * 2 * 2B + static constexpr uint32_t HALF_ELENUM_PER_LINE = 256; // 128 * 2 + static constexpr uint32_t FLOAT_ELENUM_PER_LINE = 128; // 128 + static constexpr uint32_t MULTIPLIER = 2; + static constexpr uint32_t HALF_VECTOR_SIZE = 128; + static constexpr uint32_t BLOCK_SIZE = 16; + static constexpr uint32_t T_BLOCK_SIZE = 32 / 2; + static constexpr uint32_t UB_UINT8_LINE_SIZE = 512; + static constexpr uint32_t UB_UINT8_BLOCK_SIZE_MLA = 16384; + static constexpr uint32_t HALF_DM_UB_SIZE = 128; + static constexpr uint32_t VECTOR_SIZE = 128; + static constexpr uint32_t HALF_LL_UB_SIZE = 256; + + CATLASS_DEVICE + BlockEpilogue(Arch::Resource &resource, half tor_, uint32_t kvSplitCoreNum_) + { + // Allocate UB space + constexpr uint32_t LS_UB_TENSOR_OFFSET = 0; + constexpr uint32_t LP_UB_TENSOR_OFFSET = 2 * UB_UINT8_BLOCK_SIZE_MLA; + constexpr uint32_t LM_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA; + constexpr uint32_t HM_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA + 1 * UB_UINT8_LINE_SIZE; + constexpr uint32_t DM_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA + 6 * UB_UINT8_LINE_SIZE; + constexpr uint32_t LL_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA + 9 * UB_UINT8_LINE_SIZE; + constexpr uint32_t GM_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA + 13 * UB_UINT8_LINE_SIZE; + constexpr uint32_t TV_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE_MLA; + + tor = tor_; + kvSplitCoreNum = kvSplitCoreNum_; + tvUbTensor16 = resource.ubBuf.template GetBufferByByte(LP_UB_TENSOR_OFFSET); + lpUbTensor32 = resource.ubBuf.template GetBufferByByte(LP_UB_TENSOR_OFFSET); + lsUbTensor = resource.ubBuf.template GetBufferByByte(LS_UB_TENSOR_OFFSET); + lmUbTensor = resource.ubBuf.template GetBufferByByte(LM_UB_TENSOR_OFFSET); + hmUbTensor = resource.ubBuf.template GetBufferByByte(HM_UB_TENSOR_OFFSET); + gmUbTensor = resource.ubBuf.template GetBufferByByte(GM_UB_TENSOR_OFFSET); + dmUbTensor = resource.ubBuf.template GetBufferByByte(DM_UB_TENSOR_OFFSET); + llUbTensor = resource.ubBuf.template GetBufferByByte(LL_UB_TENSOR_OFFSET); + tvUbTensor = resource.ubBuf.template GetBufferByByte(TV_UB_TENSOR_OFFSET); + } + + CATLASS_DEVICE + ~BlockEpilogue() + { + } + + CATLASS_DEVICE + void SetMask(int32_t len) + { + uint64_t mask = 0; + uint64_t one = 1; + uint64_t temp = len % FLOAT_VECTOR_SIZE; + for (int64_t i = 0; i < temp; i++) { + mask |= one << i; + } + if (len == VECTOR_SIZE) { + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + } else if (len >= FLOAT_VECTOR_SIZE) { + AscendC::SetVectorMask(mask, (uint64_t)-1); + } else { + AscendC::SetVectorMask(0x0, mask); + } + } + + CATLASS_DEVICE + void ReduceSumRepeatM( + const AscendC::LocalTensor &dst, + const AscendC::LocalTensor &src, + uint32_t curRowNum, + uint32_t kSeqTile, + uint32_t kSeqTileRound) + { + if (kSeqTile <= FLOAT_VECTOR_SIZE) { + SetMask(kSeqTile); + AscendC::RepeatReduceSum(dst, src, curRowNum, 0, 0, 1, 1, kSeqTileRound / FLOAT_BLOCK_SIZE); + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + } else { + for (uint32_t rowsum_idx = 1; rowsum_idx < kSeqTile / FLOAT_VECTOR_SIZE; ++rowsum_idx) { + AscendC::Add( + src, + src, + src[rowsum_idx * FLOAT_VECTOR_SIZE], + (uint64_t)0, + curRowNum, + AscendC::BinaryRepeatParams( + 1, 1, 1, kSeqTileRound / FLOAT_BLOCK_SIZE, + kSeqTileRound / FLOAT_BLOCK_SIZE, kSeqTileRound / FLOAT_BLOCK_SIZE)); + AscendC::PipeBarrier(); + } + if (kSeqTile % FLOAT_VECTOR_SIZE > 0) { + SetMask(kSeqTile % FLOAT_VECTOR_SIZE); + AscendC::Add( + src, + src, + src[kSeqTile / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + (uint64_t)0, + curRowNum, + AscendC::BinaryRepeatParams( + 1, 1, 1, kSeqTileRound / FLOAT_BLOCK_SIZE, + kSeqTileRound / FLOAT_BLOCK_SIZE, kSeqTileRound / FLOAT_BLOCK_SIZE)); + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + AscendC::PipeBarrier(); + AscendC::RepeatReduceSum(dst, src, curRowNum, 0, 0, 1, 1, kSeqTileRound / FLOAT_BLOCK_SIZE); + } + } + + CATLASS_DEVICE + void TensorSubValueRepeatM( + const AscendC::LocalTensor &dst, + const AscendC::LocalTensor &src, + const AscendC::LocalTensor &MaxTensor, + const AscendC::LocalTensor &tempMaxTensor, + uint32_t curRowNum, + uint32_t subMRound, + uint32_t kSeqTile, + uint32_t kSeqTileRound) + { + AscendC::Brcb(tempMaxTensor.ReinterpretCast(), MaxTensor.ReinterpretCast(), + subMRound / FLOAT_BLOCK_SIZE, AscendC::BrcbRepeatParams(1, 8)); + AscendC::PipeBarrier(); + for (uint32_t subIdx = 0; subIdx < kSeqTile / FLOAT_VECTOR_SIZE; ++subIdx) { + AscendC::Sub( + dst[subIdx * FLOAT_VECTOR_SIZE], + src[subIdx * FLOAT_VECTOR_SIZE], + tempMaxTensor, + (uint64_t)0, + curRowNum, + AscendC::BinaryRepeatParams( + 1, 1, 0, kSeqTileRound / FLOAT_BLOCK_SIZE, kSeqTileRound / FLOAT_BLOCK_SIZE, 1)); + } + if (kSeqTile % FLOAT_VECTOR_SIZE > 0) { + SetMask(kSeqTile % FLOAT_VECTOR_SIZE); + AscendC::Sub( + dst[kSeqTile / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + src[kSeqTile / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + tempMaxTensor, + (uint64_t)0, + curRowNum, + AscendC::BinaryRepeatParams( + 1, 1, 0, kSeqTileRound / FLOAT_BLOCK_SIZE, kSeqTileRound / FLOAT_BLOCK_SIZE, 1)); + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + AscendC::PipeBarrier(); + } + + CATLASS_DEVICE + void ReduceMaxRepeatM( + const AscendC::LocalTensor &dst, + const AscendC::LocalTensor &src, + const AscendC::LocalTensor &tempTensor, + uint32_t curRowNum, + uint32_t kSeqTile, + uint32_t kSeqTileRound) + { + if (kSeqTile <= FLOAT_VECTOR_SIZE) { + SetMask(kSeqTile); + AscendC::WholeReduceMax( + dst, src, (int32_t)0, curRowNum, 1, 1, + kSeqTileRound / FLOAT_BLOCK_SIZE, AscendC::ReduceOrder::ORDER_ONLY_VALUE); + } else { + AscendC::DataCopy(tempTensor, src, AscendC::DataCopyParams(curRowNum, HALF_VECTOR_SIZE / BLOCK_SIZE, (kSeqTileRound - FLOAT_VECTOR_SIZE) / FLOAT_BLOCK_SIZE, 0)); + AscendC::PipeBarrier(); + for (uint32_t rowmaxIdx = 1; rowmaxIdx < kSeqTile / FLOAT_VECTOR_SIZE; ++rowmaxIdx) { + AscendC::Max( + tempTensor, + tempTensor, + src[rowmaxIdx * FLOAT_VECTOR_SIZE], + (uint64_t)0, + curRowNum, + AscendC::BinaryRepeatParams( + 1, 1, 1, 8, 8, kSeqTileRound / FLOAT_BLOCK_SIZE)); + AscendC::PipeBarrier(); + } + if (kSeqTile % FLOAT_VECTOR_SIZE > 0) { + SetMask(kSeqTile % FLOAT_VECTOR_SIZE); + AscendC::Max( + tempTensor, + tempTensor, + src[kSeqTile / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + (uint64_t)0, + curRowNum, + AscendC::BinaryRepeatParams( + 1, 1, 1, 8, 8, kSeqTileRound / FLOAT_BLOCK_SIZE)); + } + AscendC::PipeBarrier(); + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + AscendC::WholeReduceMax( + dst, tempTensor, (int32_t)0, curRowNum, 1, 1, 8, AscendC::ReduceOrder::ORDER_ONLY_VALUE); + } + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + AscendC::PipeBarrier(); + } + + CATLASS_DEVICE + void SubCoreCompute( + AscendC::GlobalTensor gOutput, + AscendC::GlobalTensor gInput, + const LayoutOutput &layoutOutput, + const LayoutInput &layoutInput, + uint32_t nIdx, + uint32_t softmaxPingPongFlag, + uint32_t &glFlag) + { + uint32_t curRowNum = layoutInput.shape(0); + uint32_t kSeqTile = layoutInput.shape(1); + uint32_t kSeqTileRound = layoutInput.stride(0); + uint32_t subMRound = (curRowNum + 16 - 1) / 16 * 16; + uint32_t sub_m_d64 = (curRowNum + 63) / 64; // up aligned to 128 + uint64_t dmUbOffsetCurCycle = (uint64_t)(softmaxPingPongFlag * HALF_DM_UB_SIZE); + uint64_t llUbOffsetCurCycle = (uint64_t)(softmaxPingPongFlag * HALF_LL_UB_SIZE); + AscendC::WaitFlag(EVENT_ID2); + AscendC::DataCopy(lsUbTensor, gInput, + AscendC::DataCopyParams(1, curRowNum * kSeqTileRound / FLOAT_BLOCK_SIZE, 0, 0)); + + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + + // muls scale_value + for (uint32_t mulsIdx = 0; mulsIdx < kSeqTile / FLOAT_VECTOR_SIZE; ++mulsIdx) { + AscendC::Muls( + lsUbTensor[mulsIdx * FLOAT_VECTOR_SIZE], + lsUbTensor[mulsIdx * FLOAT_VECTOR_SIZE], + tor, + (uint64_t)0, + curRowNum, + AscendC::UnaryRepeatParams(1, 1, kSeqTileRound / FLOAT_BLOCK_SIZE, kSeqTileRound / FLOAT_BLOCK_SIZE)); + } + if (kSeqTile % FLOAT_VECTOR_SIZE > 0) { + SetMask(kSeqTile % FLOAT_VECTOR_SIZE); + AscendC::Muls( + lsUbTensor[kSeqTile / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + lsUbTensor[kSeqTile / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + tor, + (uint64_t)0, + curRowNum, + AscendC::UnaryRepeatParams(1, 1, kSeqTileRound / FLOAT_BLOCK_SIZE, kSeqTileRound / FLOAT_BLOCK_SIZE)); + + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + AscendC::PipeBarrier(); + + // *** lm = rowmax(ls) + ReduceMaxRepeatM(lmUbTensor, lsUbTensor, lpUbTensor32, curRowNum, kSeqTile, kSeqTileRound); + + if (nIdx != 0) { + AscendC::Max( + hmUbTensor, + lmUbTensor, + gmUbTensor, + (uint64_t)0, + sub_m_d64, + AscendC::BinaryRepeatParams( + 1, 1, 1, 8, 8, 8)); + AscendC::PipeBarrier(); + AscendC::Sub( + dmUbTensor[dmUbOffsetCurCycle], + gmUbTensor, + hmUbTensor, + (uint64_t)0, + sub_m_d64, + AscendC::BinaryRepeatParams( + 1, 1, 1, 8, 8, 8)); + AscendC::PipeBarrier(); + } else { + AscendC::DataCopy(hmUbTensor, lmUbTensor, AscendC::DataCopyParams(1, subMRound / FLOAT_BLOCK_SIZE, 0, 0)); + AscendC::PipeBarrier(); + } + // *** gm = hm + AscendC::DataCopy(gmUbTensor, hmUbTensor, AscendC::DataCopyParams(1, subMRound / FLOAT_BLOCK_SIZE, 0, 0)); + AscendC::PipeBarrier(); + + if (kvSplitCoreNum != 1) { + if (nIdx == 0) { + if (glFlag == 1) { + AscendC::WaitFlag(EVENT_ID2); + glFlag = 0; + } + } + } + + // *** ls = ls - hm_block + TensorSubValueRepeatM(lsUbTensor, lsUbTensor, + hmUbTensor, tvUbTensor, + curRowNum, subMRound, kSeqTile, kSeqTileRound); + + AscendC::Exp( + lsUbTensor, + lsUbTensor, + (uint64_t)0, + (curRowNum * kSeqTileRound + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, + AscendC::UnaryRepeatParams(1, 1, 8, 8)); + + AscendC::PipeBarrier(); + // *** lp = castfp32to16(ls) + if (std::is_same::value) { + AscendC::Cast( + tvUbTensor16, lsUbTensor, AscendC::RoundMode::CAST_RINT, (uint64_t)0, + (curRowNum * kSeqTileRound + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, + AscendC::UnaryRepeatParams(1, 1, 4, 8)); + } else { + AscendC::Cast( + tvUbTensor16, lsUbTensor, AscendC::RoundMode::CAST_NONE, (uint64_t)0, + (curRowNum * kSeqTileRound + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, + AscendC::UnaryRepeatParams(1, 1, 4, 8)); + } + + AscendC::PipeBarrier(); + + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + + uint16_t blockCount = 1; + uint16_t blockLen = curRowNum * kSeqTileRound / T_BLOCK_SIZE; + uint16_t srcStride = 0; + uint16_t dstStride = 0; + + AscendC::DataCopy(gOutput, + tvUbTensor16, + AscendC::DataCopyParams( + blockCount, // blockCount + blockLen, // blockLen + srcStride, // srcGap + dstStride)); + + // *** ll = rowsum(ls32) + ReduceSumRepeatM(llUbTensor[llUbOffsetCurCycle], lsUbTensor, curRowNum, kSeqTile, kSeqTileRound); + AscendC::SetFlag(EVENT_ID2); + AscendC::PipeBarrier(); + } + + CATLASS_DEVICE + void operator()( + AscendC::GlobalTensor gOutput, + AscendC::GlobalTensor gInput, + const LayoutOutput &layoutOutput, + const LayoutInput &layoutInput, + GemmCoord actualBlockShape, + uint32_t nIdx, + uint32_t curHeadNum, + uint32_t softmaxPingPongFlag, + uint32_t &glFlag) + { + uint32_t rowActual = actualBlockShape.m(); + uint32_t nActual = actualBlockShape.n(); + uint32_t tokenNumPerHead = rowActual / curHeadNum; + + uint32_t subBlockIdx = AscendC::GetSubBlockIdx(); + uint32_t subBlockNum = AscendC::GetSubBlockNum(); + + uint32_t curHeadSplitSubBlock = curHeadNum / subBlockNum; + uint32_t curHeadThisSubBlock = (subBlockIdx == 0) ? curHeadSplitSubBlock : (curHeadNum - curHeadSplitSubBlock); + + uint32_t rowActualThisSubBlock = curHeadThisSubBlock * tokenNumPerHead; + uint32_t rowOffsetSubBlock = subBlockIdx * curHeadSplitSubBlock * tokenNumPerHead; + + if (rowActualThisSubBlock > 0) { + int64_t offsetInput = layoutInput.GetOffset(MatrixCoord(rowOffsetSubBlock, 0)); + auto gInputThisSubBlock = gInput[offsetInput]; + auto layoutInputThisSubBlock = layoutInput.GetTileLayout(MatrixCoord(rowActualThisSubBlock, nActual)); + int64_t offsetOutput = layoutOutput.GetOffset(MatrixCoord(rowOffsetSubBlock, 0)); + auto gOutputThisSubBlock = gOutput[offsetOutput]; + auto layoutOutputThisSubBlock = layoutOutput.GetTileLayout(MatrixCoord(rowActualThisSubBlock, nActual)); + SubCoreCompute(gOutputThisSubBlock, gInputThisSubBlock, layoutOutputThisSubBlock, layoutInputThisSubBlock, + nIdx, softmaxPingPongFlag, glFlag); + } + } + +private: + float tor; + uint32_t pingpongFlag = 0; + uint32_t kvSplitCoreNum = 1; + AscendC::LocalTensor tvUbTensor16; + AscendC::LocalTensor lpUbTensor32; + AscendC::LocalTensor lsUbTensor; + AscendC::LocalTensor lmUbTensor; + AscendC::LocalTensor hmUbTensor; + AscendC::LocalTensor gmUbTensor; + AscendC::LocalTensor dmUbTensor; + AscendC::LocalTensor llUbTensor; + AscendC::LocalTensor tvUbTensor; + + CopyGmToUbInput copyGmToUbInput; + CopyGmToUbMask copyGmToUbMask; + CopyUbToGmOutput copyUbToGmOutput; +}; + +} // namespace Catlass::Epilogue::Block + +#endif // CATLASS_KDA_BLOCK_EPILOGUE_MLA_SOFTMAX_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/kda_block_epilogue_mla_tp1_softmax.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/kda_block_epilogue_mla_tp1_softmax.hpp new file mode 100644 index 0000000000..1ade25fd59 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/kda_block_epilogue_mla_tp1_softmax.hpp @@ -0,0 +1,507 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_BLOCK_EPILOGUE_MLA_TP1_SOFTMAX_HPP +#define CATLASS_KDA_BLOCK_EPILOGUE_MLA_TP1_SOFTMAX_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_cross_core_sync.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/epilogue/kda_dispatch_policy.hpp" +#include "catlass/epilogue/tile/kda_epilogue_tile_copy.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/kda_matrix_coord.hpp" + +namespace Catlass::Epilogue::Block { + +template < + class OutputType_, + class InputType_, + class MaskType_> +class BlockEpilogue< + EpilogueAtlasA2MLATP1Softmax, + OutputType_, + InputType_, + MaskType_> { +public: + // Type aliases + using DispatchPolicy = EpilogueAtlasA2MLATP1Softmax; + using ArchTag = typename DispatchPolicy::ArchTag; + using ElementOutput = typename OutputType_::Element; + using ElementInput = typename InputType_::Element; + using ElementMask = typename MaskType_::Element; + + using LayoutOutput = typename OutputType_::Layout; + using LayoutInput = typename InputType_::Layout; + using LayoutMask = typename MaskType_::Layout; + + using CopyGmToUbInput = Tile::CopyGm2Ub; + using CopyGmToUbMask = Tile::CopyGm2Ub; + using CopyUbToGmOutput = Tile::CopyUb2Gm; + + static constexpr uint32_t FLOAT_BLOCK_SIZE = 8; + static constexpr uint32_t FLOAT_VECTOR_SIZE = 64; + static constexpr uint32_t HALF_VECTOR_SIZE = 128; + static constexpr uint32_t BLOCK_SIZE = 16; + static constexpr uint32_t UB_UINT8_LINE_SIZE = 512; + static constexpr uint32_t UB_UINT8_BLOCK_SIZE_MLA = 16384; + static constexpr uint32_t VECTOR_SIZE = 128; + + static constexpr uint32_t REDUCE_UB_SIZE = 1024; + static constexpr uint32_t ROW_OPS_SPEC_MASK_32 = 32; + static constexpr uint32_t ROW_OPS_SPEC_MASK_4 = 4; + static constexpr uint32_t S_BLOCK_STACK = 4; + static constexpr int64_t UB_FLOAT_LINE_SIZE = 64; + static constexpr uint32_t M_SLICE = 16; + static constexpr uint32_t QK_READY_ID = 1; + + CATLASS_DEVICE + BlockEpilogue(Arch::Resource &resource, half tor_, uint32_t kvSplitCoreNum_ = 1) + { + // Allocate UB space + constexpr uint32_t LS_UB_TENSOR_OFFSET = 0; + constexpr uint32_t LP_UB_TENSOR_OFFSET = 0; + constexpr uint32_t LM_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA; + constexpr uint32_t HM_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA + 1 * UB_UINT8_LINE_SIZE; + constexpr uint32_t DM_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA + 6 * UB_UINT8_LINE_SIZE; + constexpr uint32_t LL_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA + 10 * UB_UINT8_LINE_SIZE; + constexpr uint32_t GM_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA + 14 * UB_UINT8_LINE_SIZE; + constexpr uint32_t GL_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA + 16 * UB_UINT8_LINE_SIZE; + constexpr uint32_t TV_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE_MLA; + + tor = tor_; + kvSplitCoreNum = kvSplitCoreNum_; + lsUbTensor = resource.ubBuf.template GetBufferByByte(LS_UB_TENSOR_OFFSET); + lpUbTensor = resource.ubBuf.template GetBufferByByte(LP_UB_TENSOR_OFFSET); + lmUbTensor = resource.ubBuf.template GetBufferByByte(LM_UB_TENSOR_OFFSET); + hmUbTensor = resource.ubBuf.template GetBufferByByte(HM_UB_TENSOR_OFFSET); + gmUbTensor = resource.ubBuf.template GetBufferByByte(GM_UB_TENSOR_OFFSET); + dmUbTensor = resource.ubBuf.template GetBufferByByte(DM_UB_TENSOR_OFFSET); + llUbTensor = resource.ubBuf.template GetBufferByByte(LL_UB_TENSOR_OFFSET); + tvUbTensor = resource.ubBuf.template GetBufferByByte(TV_UB_TENSOR_OFFSET); + glUbTensor = resource.ubBuf.template GetBufferByByte(GL_UB_TENSOR_OFFSET); + } + + CATLASS_DEVICE + ~BlockEpilogue() {} + + CATLASS_DEVICE + void SetVecMask(int32_t len) + { + uint64_t mask = 0; + uint64_t one = 1; + uint64_t temp = len % FLOAT_VECTOR_SIZE; + for (int64_t i = 0; i < temp; i++) { + mask |= one << i; + } + + if (len == VECTOR_SIZE || len == 0) { + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + } else if (len >= FLOAT_VECTOR_SIZE) { + AscendC::SetVectorMask(mask, (uint64_t)-1); + } else { + AscendC::SetVectorMask(0x0, mask); + } + } + + CATLASS_DEVICE + void SetkvSplitCoreNum(uint32_t kvSplitCoreNum_) + { + kvSplitCoreNum = kvSplitCoreNum_; + } + + CATLASS_DEVICE + void SetBlockReduceMask(int32_t len) + { + if (len > 8 || len < 1) { + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + return; + } + uint64_t subMask = ((uint64_t)1 << len) - 1; + uint64_t maskValue = (subMask << 48) + (subMask << 32) + (subMask << 16) + subMask + (subMask << 56) + + (subMask << 40) + (subMask << 24) + (subMask << 8); + AscendC::SetVectorMask(maskValue, maskValue); + } + + CATLASS_DEVICE + void RowsumSPECTILE512(const AscendC::LocalTensor &srcUb, const AscendC::LocalTensor &rowsumUb, + const AscendC::LocalTensor &tvUbTensor, uint32_t numRowsRound, uint32_t numElems, + uint32_t numElemsAligned) + { + AscendC::BlockReduceSum(tvUbTensor, srcUb, numRowsRound * numElemsAligned / FLOAT_VECTOR_SIZE, 0, + 1, 1, 8); + AscendC::PipeBarrier(); + + AscendC::BlockReduceSum(tvUbTensor[REDUCE_UB_SIZE], tvUbTensor, + numRowsRound * numElemsAligned / FLOAT_BLOCK_SIZE / FLOAT_VECTOR_SIZE, 0, + 1, 1, 8); + AscendC::PipeBarrier(); + AscendC::BlockReduceSum(rowsumUb, tvUbTensor[REDUCE_UB_SIZE], + numRowsRound * numElemsAligned / FLOAT_VECTOR_SIZE / FLOAT_VECTOR_SIZE, 0, + 1, 1, 8); + AscendC::PipeBarrier(); + } + + CATLASS_DEVICE + void RowsumSPECTILE256(const AscendC::LocalTensor &srcUb, const AscendC::LocalTensor &rowsumUb, + const AscendC::LocalTensor &tvUbTensor, uint32_t numRowsRound, uint32_t numElems, + uint32_t numElemsAligned) + { + AscendC::BlockReduceSum(tvUbTensor, srcUb, numRowsRound * numElemsAligned / FLOAT_VECTOR_SIZE, 0, + 1, 1, 8); + AscendC::PipeBarrier(); + SetVecMask(ROW_OPS_SPEC_MASK_32); + AscendC::BlockReduceSum(tvUbTensor[REDUCE_UB_SIZE], tvUbTensor, numRowsRound, 0, 1, 1, 4); + AscendC::PipeBarrier(); + SetBlockReduceMask(ROW_OPS_SPEC_MASK_4); + AscendC::BlockReduceSum( + rowsumUb, tvUbTensor[REDUCE_UB_SIZE], + (numRowsRound * FLOAT_BLOCK_SIZE + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, 0, 1, 1, 8); + AscendC::PipeBarrier(); + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + + CATLASS_DEVICE + void RowsumTAILTILE(const AscendC::LocalTensor &srcUb, const AscendC::LocalTensor &rowsumUb, + const AscendC::LocalTensor &tvUbTensor, uint32_t numRowsRound, uint32_t numElems, + uint32_t numElemsAligned) + { + if (numElems >= FLOAT_VECTOR_SIZE) { + AscendC::BlockReduceSum(tvUbTensor, srcUb, numRowsRound, 0, 1, 1, + numElemsAligned / FLOAT_BLOCK_SIZE); + AscendC::PipeBarrier(); + AscendC::BlockReduceSum( + rowsumUb, tvUbTensor, (numRowsRound * FLOAT_BLOCK_SIZE + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, 0, + 1, 1, 8); + AscendC::PipeBarrier(); + for (uint64_t rowSumIdx = 1; rowSumIdx < (uint64_t)numElems / FLOAT_VECTOR_SIZE; ++rowSumIdx) { + AscendC::BlockReduceSum(tvUbTensor, srcUb[rowSumIdx * FLOAT_VECTOR_SIZE], numRowsRound, 0, + 1, 1, numElemsAligned / FLOAT_BLOCK_SIZE); + AscendC::PipeBarrier(); + AscendC::BlockReduceSum( + tvUbTensor[REDUCE_UB_SIZE], tvUbTensor, + (numRowsRound * FLOAT_BLOCK_SIZE + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, 0, 1, 1, 8); + AscendC::PipeBarrier(); + SetVecMask(numRowsRound); + AscendC::Add(rowsumUb, rowsumUb, tvUbTensor[REDUCE_UB_SIZE], (uint64_t)0, 1, + AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); + AscendC::PipeBarrier(); + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + } + if (numElems % FLOAT_VECTOR_SIZE > 0) { + SetVecMask(numElems % FLOAT_VECTOR_SIZE); + AscendC::BlockReduceSum(tvUbTensor, srcUb[numElems / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + numRowsRound, 0, 1, 1, numElemsAligned / FLOAT_BLOCK_SIZE); + AscendC::PipeBarrier(); + SetBlockReduceMask((numElems % FLOAT_VECTOR_SIZE + FLOAT_BLOCK_SIZE - 1) / FLOAT_BLOCK_SIZE); + if (numElems < FLOAT_VECTOR_SIZE) { + AscendC::BlockReduceSum( + rowsumUb, tvUbTensor, (numRowsRound * FLOAT_BLOCK_SIZE + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, + 0, 1, 1, 8); + AscendC::PipeBarrier(); + } else { + AscendC::BlockReduceSum( + tvUbTensor[REDUCE_UB_SIZE], tvUbTensor, + (numRowsRound * FLOAT_BLOCK_SIZE + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, 0, 1, 1, 8); + AscendC::PipeBarrier(); + SetVecMask(numRowsRound); + AscendC::Add(rowsumUb, rowsumUb, tvUbTensor[REDUCE_UB_SIZE], (uint64_t)0, 1, + AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); + AscendC::PipeBarrier(); + } + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + } + + CATLASS_DEVICE + void RowmaxSPECTILE512(const AscendC::LocalTensor &srcUb, const AscendC::LocalTensor &rowmaxUb, + const AscendC::LocalTensor &tvUbTensor, uint32_t numRowsRound, uint32_t numElems, + uint32_t numElemsAligned) + { + AscendC::BlockReduceMax(tvUbTensor, srcUb, numRowsRound * numElemsAligned / FLOAT_VECTOR_SIZE, 0, + 1, 1, 8); + AscendC::PipeBarrier(); + AscendC::BlockReduceMax(tvUbTensor[REDUCE_UB_SIZE], tvUbTensor, + numRowsRound * numElemsAligned / FLOAT_BLOCK_SIZE / FLOAT_VECTOR_SIZE, 0, + 1, 1, 8); + AscendC::PipeBarrier(); + AscendC::BlockReduceMax(rowmaxUb, tvUbTensor[REDUCE_UB_SIZE], + numRowsRound * numElemsAligned / FLOAT_VECTOR_SIZE / FLOAT_VECTOR_SIZE, 0, + 1, 1, 8); + AscendC::PipeBarrier(); + } + + CATLASS_DEVICE + void RowmaxSPECTILE256(const AscendC::LocalTensor &srcUb, const AscendC::LocalTensor &rowmaxUb, + const AscendC::LocalTensor &tvUbTensor, uint32_t numRowsRound, uint32_t numElems, + uint32_t numElemsAligned) + { + AscendC::BlockReduceMax(tvUbTensor, srcUb, numRowsRound * numElemsAligned / FLOAT_VECTOR_SIZE, 0, + 1, 1, 8); + AscendC::PipeBarrier(); + SetVecMask(ROW_OPS_SPEC_MASK_32); + AscendC::BlockReduceMax(tvUbTensor[REDUCE_UB_SIZE], tvUbTensor, numRowsRound, 0, 1, 1, 4); + AscendC::PipeBarrier(); + SetBlockReduceMask(ROW_OPS_SPEC_MASK_4); + AscendC::BlockReduceMax( + rowmaxUb, tvUbTensor[REDUCE_UB_SIZE], + (numRowsRound * FLOAT_BLOCK_SIZE + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, 0, 1, 1, 8); + AscendC::PipeBarrier(); + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + + CATLASS_DEVICE + void RowmaxTAILTILE(const AscendC::LocalTensor &srcUb, const AscendC::LocalTensor &rowmaxUb, + const AscendC::LocalTensor &tvUbTensor, uint32_t numRowsRound, uint32_t numElems, + uint32_t numElemsAligned) + { + if (numElems >= FLOAT_VECTOR_SIZE) { + AscendC::BlockReduceMax(tvUbTensor, srcUb, numRowsRound, 0, 1, 1, + numElemsAligned / FLOAT_BLOCK_SIZE); + AscendC::PipeBarrier(); + AscendC::BlockReduceMax( + rowmaxUb, tvUbTensor, (numRowsRound * FLOAT_BLOCK_SIZE + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, 0, + 1, 1, 8); + AscendC::PipeBarrier(); + for (uint64_t rowmax_idx = 1; rowmax_idx < (uint64_t)numElems / FLOAT_VECTOR_SIZE; ++rowmax_idx) { + AscendC::BlockReduceMax(tvUbTensor, srcUb[rowmax_idx * FLOAT_VECTOR_SIZE], numRowsRound, + 0, 1, 1, numElemsAligned / FLOAT_BLOCK_SIZE); + AscendC::PipeBarrier(); + AscendC::BlockReduceMax( + tvUbTensor[REDUCE_UB_SIZE], tvUbTensor, + (numRowsRound * FLOAT_BLOCK_SIZE + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, 0, 1, 1, 8); + AscendC::PipeBarrier(); + SetVecMask(numRowsRound); + AscendC::Max(rowmaxUb, rowmaxUb, tvUbTensor[REDUCE_UB_SIZE], (uint64_t)0, 1, + AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); + AscendC::PipeBarrier(); + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + } + if (numElems % FLOAT_VECTOR_SIZE > 0) { + SetVecMask(numElems % FLOAT_VECTOR_SIZE); + AscendC::BlockReduceMax(tvUbTensor, srcUb[numElems / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + numRowsRound, 0, 1, 1, numElemsAligned / FLOAT_BLOCK_SIZE); + AscendC::PipeBarrier(); + SetBlockReduceMask((numElems % FLOAT_VECTOR_SIZE + FLOAT_BLOCK_SIZE - 1) / FLOAT_BLOCK_SIZE); + if (numElems < FLOAT_VECTOR_SIZE) { + AscendC::BlockReduceMax( + rowmaxUb, tvUbTensor, (numRowsRound * FLOAT_BLOCK_SIZE + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, + 0, 1, 1, 8); + AscendC::PipeBarrier(); + } else { + AscendC::BlockReduceMax( + tvUbTensor[REDUCE_UB_SIZE], tvUbTensor, + (numRowsRound * FLOAT_BLOCK_SIZE + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, 0, 1, 1, 8); + AscendC::PipeBarrier(); + SetVecMask(numRowsRound); + AscendC::Max(rowmaxUb, rowmaxUb, tvUbTensor[REDUCE_UB_SIZE], (uint64_t)0, 1, + AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); + AscendC::PipeBarrier(); + } + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + } + + CATLASS_DEVICE + void SubCoreCompute(AscendC::GlobalTensor gOutput, AscendC::GlobalTensor gInput, + uint32_t m, uint32_t nReal, uint32_t nStride, uint32_t pingpongFlag, uint32_t rowOffset, + uint32_t sUbOffset, uint32_t nIdx, uint32_t &glFlag) + { + uint32_t round_m = (m + FLOAT_BLOCK_SIZE - 1) / FLOAT_BLOCK_SIZE * FLOAT_BLOCK_SIZE; + AscendC::WaitFlag(pingpongFlag); + // input QK + AscendC::DataCopy(lsUbTensor[sUbOffset], gInput, AscendC::DataCopyParams(m, nStride / FLOAT_BLOCK_SIZE, 0, 0)); + + AscendC::SetFlag(pingpongFlag); + AscendC::WaitFlag(pingpongFlag); + + // *** ls = tor * ls + AscendC::Muls(lsUbTensor[sUbOffset], lsUbTensor[sUbOffset], tor, (uint64_t)0, + (m * nStride + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, + AscendC::UnaryRepeatParams(1, 1, 8, 8)); + + AscendC::PipeBarrier(); + + if (kvSplitCoreNum != 1) { + if (nIdx == 0) { + if (glFlag == 1) { + AscendC::WaitFlag(EVENT_ID2); + glFlag = 0; + } + } + } + + if (nReal == 512) { + RowmaxSPECTILE512(lsUbTensor[sUbOffset], lmUbTensor[rowOffset], tvUbTensor, round_m, nReal, nStride); + } else if (nReal == 256) { + RowmaxSPECTILE256(lsUbTensor[sUbOffset], lmUbTensor[rowOffset], tvUbTensor, round_m, nReal, nStride); + } else { + RowmaxTAILTILE(lsUbTensor[sUbOffset], lmUbTensor[rowOffset], tvUbTensor, round_m, nReal, nStride); + } + + if (nIdx == 0) { + AscendC::DataCopy(hmUbTensor[rowOffset], lmUbTensor[rowOffset], + AscendC::DataCopyParams(1, round_m / FLOAT_BLOCK_SIZE, 0, 0)); + AscendC::PipeBarrier(); + } else { + SetVecMask(m); + // *** hm = vmax(lm, gm) + AscendC::Max(hmUbTensor[rowOffset], lmUbTensor[rowOffset], gmUbTensor[rowOffset], (uint64_t)0, + 1, AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); + + AscendC::PipeBarrier(); + // *** dm = gm - hm + AscendC::Sub(dmUbTensor[((nIdx / S_BLOCK_STACK) % 2) * UB_FLOAT_LINE_SIZE + rowOffset], + gmUbTensor[rowOffset], hmUbTensor[rowOffset], (uint64_t)0, 1, + AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); + + AscendC::PipeBarrier(); + // *** dm = exp(dm) + AscendC::Exp(dmUbTensor[((nIdx / S_BLOCK_STACK) % 2) * UB_FLOAT_LINE_SIZE + rowOffset], + dmUbTensor[((nIdx / S_BLOCK_STACK) % 2) * UB_FLOAT_LINE_SIZE + rowOffset], + (uint64_t)0, 1, AscendC::UnaryRepeatParams(1, 1, 8, 8)); + } + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + AscendC::PipeBarrier(); + // *** gm = hm + AscendC::DataCopy(gmUbTensor[rowOffset], hmUbTensor[rowOffset], + AscendC::DataCopyParams(1, round_m / FLOAT_BLOCK_SIZE, 0, 0)); + AscendC::PipeBarrier(); + // *** hm_block = expand_to_block(hm), 存放于 tv + AscendC::Brcb(tvUbTensor.template ReinterpretCast(), + hmUbTensor[rowOffset].template ReinterpretCast(), round_m / FLOAT_BLOCK_SIZE, + AscendC::BrcbRepeatParams(1, 8)); + AscendC::PipeBarrier(); + // *** ls = ls - hm_block + for (uint32_t subIdx = 0; subIdx < nReal / FLOAT_VECTOR_SIZE; ++subIdx) { + AscendC::Sub( + lsUbTensor[sUbOffset][subIdx * FLOAT_VECTOR_SIZE], lsUbTensor[sUbOffset][subIdx * FLOAT_VECTOR_SIZE], + tvUbTensor, (uint64_t)0, m, + AscendC::BinaryRepeatParams(1, 1, 0, nStride / FLOAT_BLOCK_SIZE, nStride / FLOAT_BLOCK_SIZE, 1)); + } + if (nReal % FLOAT_VECTOR_SIZE > 0) { + SetVecMask(nReal % FLOAT_VECTOR_SIZE); + AscendC::Sub( + lsUbTensor[sUbOffset][nReal / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + lsUbTensor[sUbOffset][nReal / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], tvUbTensor, (uint64_t)0, m, + AscendC::BinaryRepeatParams(1, 1, 0, nStride / FLOAT_BLOCK_SIZE, nStride / FLOAT_BLOCK_SIZE, 1)); + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + AscendC::PipeBarrier(); + + // *** ls = exp(ls) + AscendC::Exp(lsUbTensor[sUbOffset], lsUbTensor[sUbOffset], (uint64_t)0, + (m * nStride + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, + AscendC::UnaryRepeatParams(1, 1, 8, 8)); + AscendC::PipeBarrier(); + // *** ll = rowsum(ls32) + if (nReal == 512) { + RowsumSPECTILE512(lsUbTensor[sUbOffset], llUbTensor[rowOffset], tvUbTensor, round_m, nReal, nStride); + } else if (nReal == 256) { + RowsumSPECTILE256(lsUbTensor[sUbOffset], llUbTensor[rowOffset], tvUbTensor, round_m, nReal, nStride); + } else { + RowsumTAILTILE(lsUbTensor[sUbOffset], llUbTensor[rowOffset], tvUbTensor, round_m, nReal, nStride); + } + + // *** lp = castfp32to16(ls) + if (std::is_same::value) { + AscendC::Cast( + lpUbTensor[sUbOffset * 2], lsUbTensor[sUbOffset], AscendC::RoundMode::CAST_RINT, (uint64_t)0, + (m * nStride + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, AscendC::UnaryRepeatParams(1, 1, 4, 8)); + } else { + AscendC::Cast( + lpUbTensor[sUbOffset * 2], lsUbTensor[sUbOffset], AscendC::RoundMode::CAST_NONE, (uint64_t)0, + (m * nStride + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, AscendC::UnaryRepeatParams(1, 1, 4, 8)); + } + + AscendC::SetFlag(pingpongFlag); + AscendC::WaitFlag(pingpongFlag); + AscendC::DataCopy(gOutput, lpUbTensor[sUbOffset * 2], AscendC::DataCopyParams(m, nStride * 2 / 32, 0, 0)); + AscendC::SetFlag(pingpongFlag); + if (nIdx == 0) { + // *** gl = ll + AscendC::DataCopy(glUbTensor[rowOffset], llUbTensor[rowOffset], + AscendC::DataCopyParams(1, round_m / FLOAT_BLOCK_SIZE, 0, 0)); + AscendC::PipeBarrier(); + } else { + SetVecMask(m); + // *** gl = dm * gl + AscendC::Mul( + glUbTensor[rowOffset], dmUbTensor[((nIdx / S_BLOCK_STACK) % 2) * UB_FLOAT_LINE_SIZE + rowOffset], + glUbTensor[rowOffset], (uint64_t)0, 1, AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); + AscendC::PipeBarrier(); + // *** gl = ll + gl + AscendC::Add(glUbTensor[rowOffset], glUbTensor[rowOffset], llUbTensor[rowOffset], (uint64_t)0, + 1, AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); + AscendC::PipeBarrier(); + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + } + + CATLASS_DEVICE + void operator()(AscendC::GlobalTensor gOutput, AscendC::GlobalTensor gInput, + const LayoutOutput &layoutOutput, const LayoutInput &layoutInput, GemmCoord actualBlockShape, + uint32_t nIdx, uint32_t &glFlag) + { + uint32_t cur_head_num = actualBlockShape.m(); + uint32_t qkN = actualBlockShape.n(); + uint32_t qkRoundN = layoutInput.stride(0); + uint32_t pingpongFlag = 0; + + uint32_t subBlockIdx = AscendC::GetSubBlockIdx(); + uint32_t subBlockNum = AscendC::GetSubBlockNum(); + + uint32_t subM = (subBlockIdx == 1) ? (cur_head_num - cur_head_num / 2) : cur_head_num / 2; + + uint32_t mEnd = (subM + M_SLICE - 1) / M_SLICE; + + for (uint32_t mInd = 0; mInd < mEnd; mInd++) { + uint32_t rowOffset = mInd * M_SLICE; + uint32_t currM = mInd == mEnd - 1 ? subM - rowOffset : M_SLICE; + uint32_t sUbOffset = pingpongFlag * 8192; + int64_t offsetOutput = rowOffset * qkRoundN; + auto gOutputThisSubBlock = gOutput[offsetOutput]; + int64_t offsetInput = rowOffset * qkRoundN; + auto gInputThisSubBlock = gInput[offsetInput]; + if (mInd == 0) { + Arch::CrossCoreWaitFlag(qkReady); + } + if (currM == 0) { + continue; + } + SubCoreCompute(gOutputThisSubBlock, gInputThisSubBlock, currM, qkN, qkRoundN, pingpongFlag, rowOffset, + sUbOffset, nIdx, glFlag); + pingpongFlag = 1 - pingpongFlag; + } + } + +private: + float tor; + uint32_t pingpongFlag = 0; + uint32_t kvSplitCoreNum = 1; + + AscendC::LocalTensor lsUbTensor; + AscendC::LocalTensor lpUbTensor; + AscendC::LocalTensor lmUbTensor; + AscendC::LocalTensor hmUbTensor; + AscendC::LocalTensor gmUbTensor; + AscendC::LocalTensor dmUbTensor; + AscendC::LocalTensor llUbTensor; + AscendC::LocalTensor tvUbTensor; + AscendC::LocalTensor glUbTensor; + + Arch::CrossCoreFlag qkReady{QK_READY_ID}; +}; + +} // namespace Catlass::Epilogue::Block + +#endif // CATLASS_KDA_BLOCK_EPILOGUE_MLA_TP1_SOFTMAX_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/kda_block_epilogue_per_token_dequant.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/kda_block_epilogue_per_token_dequant.hpp new file mode 100644 index 0000000000..53c7cd6636 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/kda_block_epilogue_per_token_dequant.hpp @@ -0,0 +1,683 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_BLOCK_EPILOGUE_PER_TOKEN_DEQUANT_HPP +#define CATLASS_KDA_BLOCK_EPILOGUE_PER_TOKEN_DEQUANT_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/epilogue/kda_dispatch_policy.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/gemm/kda_gemm_type.hpp" +#include "catlass/kda_matrix_coord.hpp" +#include "catlass/layout/kda_layout.hpp" +#include "catlass/detail/kda_callback.hpp" + +namespace Catlass::Epilogue::Block { + +template < + uint32_t UB_STAGES_, + class CType_, + class ScaleType_, + class PerTokenScaleType_, + class DType_, + class TileRowBroadcastMul_, + class TileBroadcastOneBlk_, + class TileOneBlkColumnBroadcastMul_, + class TileCopy_, + class EpilogueTileSwizzle_> +class BlockEpilogue< + EpilogueAtlasA2PerTokenDequant, + CType_, + ScaleType_, + PerTokenScaleType_, + DType_, + TileRowBroadcastMul_, + TileBroadcastOneBlk_, + TileOneBlkColumnBroadcastMul_, + TileCopy_, + EpilogueTileSwizzle_> { +public: + using DispatchPolicy = EpilogueAtlasA2PerTokenDequant; + using ArchTag = typename DispatchPolicy::ArchTag; + static constexpr uint32_t UB_STAGES = UB_STAGES_; + + // Data infos + using ElementC = typename CType_::Element; + using LayoutC = typename CType_::Layout; + using ElementScale = typename ScaleType_::Element; + using LayoutScale = typename ScaleType_::Layout; + using ElementPerTokenScale = typename PerTokenScaleType_::Element; + using LayoutPerTokenScale = typename PerTokenScaleType_::Layout; + using ElementD = typename DType_::Element; + using LayoutD = typename DType_::Layout; + + // Check data infos + static_assert( + std::is_same_v && (std::is_same_v || std::is_same_v) && + std::is_same_v && std::is_same_v, + "The element type template parameters of BlockEpilogue are wrong"); + static_assert( + std::is_same_v && std::is_same_v && + std::is_same_v && std::is_same_v, + "The layout template parameters of BlockEpilogue are wrong"); + + // Tile compute ops + using TileRowBroadcastMul = TileRowBroadcastMul_; + using TileBroadcastOneBlk = TileBroadcastOneBlk_; + using TileOneBlkColumnBroadcastMul = TileOneBlkColumnBroadcastMul_; + + // Tile copy + using CopyGmToUbC = typename TileCopy_::CopyGmToUbC; + using CopyGmToUbScale = typename TileCopy_::CopyGmToUbX; + using CopyGmToUbPerTokenScale = typename TileCopy_::CopyGmToUbY; + using CopyUbToGmD = typename TileCopy_::CopyUbToGmD; + + using EpilogueTileSwizzle = EpilogueTileSwizzle_; + + using TileShape = typename TileRowBroadcastMul::TileShape; + + static_assert( + TileShape::ROW == TileBroadcastOneBlk::COMPUTE_LENGTH && + std::is_same_v, + "TileShape must be consistent for all tile compute ops"); + + static_assert( + (UB_STAGES * (TileShape::COUNT * sizeof(ElementC) + TileShape::COLUMN * sizeof(ElementScale) + TileShape::ROW * sizeof(ElementPerTokenScale) + TileShape::COUNT * sizeof(ElementD)) + (TileShape::COUNT + TileShape::COLUMN + TileShape::COUNT + TileShape::ROW) * sizeof(float) + TileShape::ROW * BYTE_PER_BLK) <= ArchTag::UB_SIZE, + "TileShape is too large to fit in UB"); + + struct Params { + __gm__ ElementScale *ptrScale{nullptr}; + LayoutScale layoutScale{}; + __gm__ ElementPerTokenScale *ptrPerTokenScale{nullptr}; + LayoutPerTokenScale layoutPerTokenScale{}; + __gm__ ElementD *ptrD{nullptr}; + LayoutD layoutD{}; + + CATLASS_DEVICE + Params() {}; + + CATLASS_DEVICE + Params( + __gm__ ElementScale *ptrScale_, LayoutScale const &layoutScale_, + __gm__ ElementPerTokenScale *ptrPerTokenScale_, LayoutPerTokenScale const &layoutPerTokenScale_, + __gm__ ElementD *ptrD_, LayoutD const &layoutD_) + : ptrScale(ptrScale_), + layoutScale(layoutScale_), + ptrPerTokenScale(ptrPerTokenScale_), + layoutPerTokenScale(layoutPerTokenScale_), + ptrD(ptrD_), + layoutD(layoutD_) + {} + }; + + CATLASS_DEVICE + BlockEpilogue(Arch::Resource const &resource, Params const ¶ms = Params{}) + : params(params) + { + size_t ubOffset = 0; + int32_t eventVMTE2 = 0; + int32_t eventMTE2V = 0; + int32_t eventMTE3V = 0; + int32_t eventVMTE3 = 0; + for (uint32_t i = 0; i < UB_STAGES; ++i) { + ubCList[i] = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += TileShape::COUNT * sizeof(ElementC); + ubScaleList[i] = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += TileShape::COLUMN * sizeof(ElementScale); + ubPerTokenScaleList[i] = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += TileShape::ROW * sizeof(ElementPerTokenScale); + ubDList[i] = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += TileShape::COUNT * sizeof(ElementD); + + eventUbCVMTE2List[i] = eventVMTE2++; + eventUbCMTE2VList[i] = eventMTE2V++; + eventUbScaleVMTE2List[i] = eventVMTE2++; + eventUbScaleMTE2VList[i] = eventMTE2V++; + eventUbPerTokenScaleVMTE2List[i] = eventVMTE2++; + eventUbPerTokenScaleMTE2VList[i] = eventMTE2V++; + eventUbDMTE3VList[i] = eventMTE3V++; + eventUbDVMTE3List[i] = eventVMTE3++; + + AscendC::SetFlag(eventUbCVMTE2List[i]); + AscendC::SetFlag(eventUbScaleVMTE2List[i]); + AscendC::SetFlag(eventUbPerTokenScaleVMTE2List[i]); + AscendC::SetFlag(eventUbDMTE3VList[i]); + } + ubCFp32 = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += TileShape::COUNT * sizeof(float); + ubScaleFp32 = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += TileShape::COLUMN * sizeof(float); + ubMul = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += TileShape::COUNT * sizeof(float); + ubPerTokenScaleFp32 = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += TileShape::ROW * sizeof(float); + ubPerTokenScaleFp32Brcb = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += TileShape::ROW * BYTE_PER_BLK; + ubPerTokenMul = ubMul; + } + + CATLASS_DEVICE + ~BlockEpilogue() + { + for (uint32_t i = 0; i < UB_STAGES; ++i) { + AscendC::WaitFlag(eventUbCVMTE2List[i]); + AscendC::WaitFlag(eventUbScaleVMTE2List[i]); + AscendC::WaitFlag(eventUbPerTokenScaleVMTE2List[i]); + AscendC::WaitFlag(eventUbDMTE3VList[i]); + } + } + + CATLASS_DEVICE + void UpdateParams(Params const ¶ms_) + { + params = params_; + } + + CATLASS_DEVICE + void operator()( + GemmCoord const &blockShapeMNK, + GemmCoord const &blockCoordMNK, + GemmCoord const &actualBlockShapeMNK, + AscendC::GlobalTensor const &gmBlockC, + LayoutC const &layoutBlockC, Callback &&callback = Callback{}) + { + MatrixCoord blockShape = blockShapeMNK.GetCoordMN(); + MatrixCoord blockCoord = blockCoordMNK.GetCoordMN(); + MatrixCoord actualBlockShape = actualBlockShapeMNK.GetCoordMN(); + MatrixCoord blockOffset = blockCoord * blockShape; + + auto ubTileStride = MakeCoord(static_cast(TileShape::COLUMN), 1L); + auto tileShape = TileShape::ToCoord(); + EpilogueTileSwizzle epilogueTileSwizzle(actualBlockShape, tileShape); + uint32_t tileLoops = epilogueTileSwizzle.GetLoops(); + uint32_t subblockIdx = AscendC::GetSubBlockIdx(); + uint32_t subblockNum = AscendC::GetSubBlockNum(); + + if (actualBlockShapeMNK.k() == 0) { + AscendC::GlobalTensor gmD; + gmD.SetGlobalBuffer(params.ptrD); + + ElementD zeroD{}; + for (uint32_t loopIdx = subblockIdx; loopIdx < tileLoops; loopIdx += subblockNum) { + auto tileCoord = epilogueTileSwizzle.GetTileCoord(loopIdx); + auto actualTileShape = epilogueTileSwizzle.GetActualTileShape(tileCoord); + auto tileOffset = blockOffset + tileCoord * tileShape; + + auto &ubD = ubDList[ubListId]; + LayoutD layoutUbD{actualTileShape, ubTileStride}; + + AscendC::WaitFlag(eventUbDMTE3VList[ubListId]); + AscendC::Duplicate(ubD, zeroD, TileShape::COUNT); + AscendC::SetFlag(eventUbDVMTE3List[ubListId]); + + auto gmTileD = gmD[params.layoutD.GetOffset(tileOffset)]; + auto layoutGmTileD = params.layoutD.GetTileLayout(actualTileShape); + + AscendC::WaitFlag(eventUbDVMTE3List[ubListId]); + copyUbToGmD(gmTileD, ubD, layoutGmTileD, layoutUbD); + AscendC::SetFlag(eventUbDMTE3VList[ubListId]); + + ubListId = (ubListId + 1 < UB_STAGES) ? (ubListId + 1) : 0; + } + return; + } + callback(); + + AscendC::GlobalTensor gmScale; + gmScale.SetGlobalBuffer(params.ptrScale); + AscendC::GlobalTensor gmPerTokenScale; + gmPerTokenScale.SetGlobalBuffer(params.ptrPerTokenScale); + AscendC::GlobalTensor gmD; + gmD.SetGlobalBuffer(params.ptrD); + + for (uint32_t loopIdx = subblockIdx; loopIdx < tileLoops; loopIdx += subblockNum) { + auto tileCoord = epilogueTileSwizzle.GetTileCoord(loopIdx); + auto actualTileShape = epilogueTileSwizzle.GetActualTileShape(tileCoord); + auto tileOffsetInBlock = tileCoord * tileShape; + auto tileOffset = blockOffset + tileOffsetInBlock; + + auto gmTileC = gmBlockC[layoutBlockC.GetOffset(tileOffsetInBlock)]; + auto layoutGmTileC = layoutBlockC.GetTileLayout(actualTileShape); + + auto &ubC = ubCList[ubListId]; + LayoutC layoutUbC{actualTileShape, ubTileStride}; + + AscendC::WaitFlag(eventUbCVMTE2List[ubListId]); + copyGmToUbC(ubC, gmTileC, layoutUbC, layoutGmTileC); + AscendC::SetFlag(eventUbCMTE2VList[ubListId]); + + auto scaleTileOffset = tileOffset.template GetCoordByAxis<1>(); + auto scaleTileShape = actualTileShape.template GetCoordByAxis<1>(); + + auto gmTileScale = gmScale[params.layoutScale.GetOffset(scaleTileOffset)]; + auto layoutGmTileScale = params.layoutScale.GetTileLayout(scaleTileShape); + + auto &ubScale = ubScaleList[ubListId]; + auto layoutUbScale = LayoutScale::template MakeLayoutInUb(scaleTileShape); + + AscendC::WaitFlag(eventUbScaleVMTE2List[ubListId]); + copyGmToUbScale(ubScale, gmTileScale, layoutUbScale, layoutGmTileScale); + AscendC::SetFlag(eventUbScaleMTE2VList[ubListId]); + + auto perTokenScaleTileOffset = tileOffset.template GetCoordByAxis<0>(); + auto perTokenScaleTileShape = actualTileShape.template GetCoordByAxis<0>(); + + auto gmTilePerTokenScale = gmPerTokenScale[params.layoutPerTokenScale.GetOffset(perTokenScaleTileOffset)]; + auto layoutGmTilePerTokenScale = params.layoutPerTokenScale.GetTileLayout(perTokenScaleTileShape); + + auto &ubPerTokenScale = ubPerTokenScaleList[ubListId]; + auto layoutUbPerTokenScale = LayoutScale::template MakeLayoutInUb( + perTokenScaleTileShape); + + AscendC::WaitFlag(eventUbPerTokenScaleVMTE2List[ubListId]); + copyGmToUbPerTokenScale(ubPerTokenScale, gmTilePerTokenScale, layoutUbPerTokenScale, + layoutGmTilePerTokenScale); + AscendC::SetFlag(eventUbPerTokenScaleMTE2VList[ubListId]); + + AscendC::WaitFlag(eventUbCMTE2VList[ubListId]); + AscendC::Cast(ubCFp32, ubC, AscendC::RoundMode::CAST_RINT, TileShape::COUNT); + AscendC::SetFlag(eventUbCVMTE2List[ubListId]); + + AscendC::WaitFlag(eventUbScaleMTE2VList[ubListId]); + AscendC::Cast(ubScaleFp32, ubScale, AscendC::RoundMode::CAST_NONE, TileShape::COLUMN); + AscendC::SetFlag(eventUbScaleVMTE2List[ubListId]); + + AscendC::WaitFlag(eventUbPerTokenScaleMTE2VList[ubListId]); + AscendC::Cast(ubPerTokenScaleFp32, ubPerTokenScale, AscendC::RoundMode::CAST_NONE, TileShape::ROW); + AscendC::SetFlag(eventUbPerTokenScaleVMTE2List[ubListId]); + + AscendC::PipeBarrier(); + tileRowBroadcastMul(ubMul, ubCFp32, ubScaleFp32); + tileBroadcastOneBlk(ubPerTokenScaleFp32Brcb, ubPerTokenScaleFp32); + AscendC::PipeBarrier(); + tileOneBlkColumnBroadcastMul(ubPerTokenMul, ubMul, ubPerTokenScaleFp32Brcb); + AscendC::PipeBarrier(); + + auto &ubD = ubDList[ubListId]; + LayoutD layoutUbD{actualTileShape, ubTileStride}; + + AscendC::WaitFlag(eventUbDMTE3VList[ubListId]); + AscendC::Cast(ubD, ubPerTokenMul, AscendC::RoundMode::CAST_RINT, TileShape::COUNT); + AscendC::SetFlag(eventUbDVMTE3List[ubListId]); + + auto gmTileD = gmD[params.layoutD.GetOffset(tileOffset)]; + auto layoutGmTileD = params.layoutD.GetTileLayout(actualTileShape); + + AscendC::WaitFlag(eventUbDVMTE3List[ubListId]); + copyUbToGmD(gmTileD, ubD, layoutGmTileD, layoutUbD); + AscendC::SetFlag(eventUbDMTE3VList[ubListId]); + + ubListId = (ubListId + 1 < UB_STAGES) ? (ubListId + 1) : 0; + } + } + +private: + Params params; + + AscendC::LocalTensor ubCList[UB_STAGES]; + AscendC::LocalTensor ubScaleList[UB_STAGES]; + AscendC::LocalTensor ubPerTokenScaleList[UB_STAGES]; + AscendC::LocalTensor ubDList[UB_STAGES]; + + int32_t eventUbCVMTE2List[UB_STAGES]; + int32_t eventUbCMTE2VList[UB_STAGES]; + int32_t eventUbScaleVMTE2List[UB_STAGES]; + int32_t eventUbScaleMTE2VList[UB_STAGES]; + int32_t eventUbPerTokenScaleVMTE2List[UB_STAGES]; + int32_t eventUbPerTokenScaleMTE2VList[UB_STAGES]; + int32_t eventUbDMTE3VList[UB_STAGES]; + int32_t eventUbDVMTE3List[UB_STAGES]; + + uint32_t ubListId{0}; + + AscendC::LocalTensor ubCFp32; + AscendC::LocalTensor ubScaleFp32; + AscendC::LocalTensor ubMul; + AscendC::LocalTensor ubPerTokenScaleFp32; + AscendC::LocalTensor ubPerTokenScaleFp32Brcb; + AscendC::LocalTensor ubPerTokenMul; + + TileRowBroadcastMul tileRowBroadcastMul; + TileBroadcastOneBlk tileBroadcastOneBlk; + TileOneBlkColumnBroadcastMul tileOneBlkColumnBroadcastMul; + + CopyGmToUbC copyGmToUbC; + CopyGmToUbScale copyGmToUbScale; + CopyGmToUbPerTokenScale copyGmToUbPerTokenScale; + CopyUbToGmD copyUbToGmD; +}; + +template < + uint32_t UB_STAGES_, + class CType_, + class LayoutScale_, + class LayoutPerTokenScale_, + class DType_, + class TileRowBroadcastMul_, + class TileBroadcastOneBlk_, + class TileOneBlkColumnBroadcastMul_, + class TileCopy_, + class EpilogueTileSwizzle_> +class BlockEpilogue< + EpilogueAtlasA2PerTokenDequant, + CType_, + Gemm::GemmType, + Gemm::GemmType, + DType_, + TileRowBroadcastMul_, + TileBroadcastOneBlk_, + TileOneBlkColumnBroadcastMul_, + TileCopy_, + EpilogueTileSwizzle_> { +public: + using DispatchPolicy = EpilogueAtlasA2PerTokenDequant; + using ArchTag = typename DispatchPolicy::ArchTag; + static constexpr uint32_t UB_STAGES = UB_STAGES_; + + // Data infos + using ElementC = typename CType_::Element; + using LayoutC = typename CType_::Layout; + using ElementScale = float; + using LayoutScale = LayoutScale_; + using ElementPerTokenScale = float; + using LayoutPerTokenScale = LayoutPerTokenScale_; + using ElementD = typename DType_::Element; + using LayoutD = typename DType_::Layout; + + // Check data infos + static_assert( + std::is_same_v && (std::is_same_v || std::is_same_v), + "The element type template parameters of BlockEpilogue are wrong"); + static_assert( + std::is_same_v && std::is_same_v && + std::is_same_v && std::is_same_v, + "The layout template parameters of BlockEpilogue are wrong"); + + // Tile compute ops + using TileRowBroadcastMul = TileRowBroadcastMul_; + using TileBroadcastOneBlk = TileBroadcastOneBlk_; + using TileOneBlkColumnBroadcastMul = TileOneBlkColumnBroadcastMul_; + + // Tile copy + using CopyGmToUbC = typename TileCopy_::CopyGmToUbC; + using CopyGmToUbScale = typename TileCopy_::CopyGmToUbX; + using CopyGmToUbPerTokenScale = typename TileCopy_::CopyGmToUbY; + using CopyUbToGmD = typename TileCopy_::CopyUbToGmD; + + using EpilogueTileSwizzle = EpilogueTileSwizzle_; + + using TileShape = typename TileRowBroadcastMul::TileShape; + + static_assert( + TileShape::ROW == TileBroadcastOneBlk::COMPUTE_LENGTH && + std::is_same_v, + "TileShape must be consistent for all tile compute ops"); + + static_assert( + (UB_STAGES * (TileShape::COUNT * sizeof(ElementC) + TileShape::COLUMN * sizeof(ElementScale) + TileShape::ROW * sizeof(ElementPerTokenScale) + TileShape::COUNT * sizeof(ElementD)) + (TileShape::COUNT + TileShape::COUNT) * sizeof(float) + TileShape::ROW * BYTE_PER_BLK) <= ArchTag::UB_SIZE, + "TileShape is too large to fit in UB"); + + struct Params { + __gm__ ElementScale *ptrScale{nullptr}; + LayoutScale layoutScale{}; + __gm__ ElementPerTokenScale *ptrPerTokenScale{nullptr}; + LayoutPerTokenScale layoutPerTokenScale{}; + __gm__ ElementD *ptrD{nullptr}; + LayoutD layoutD{}; + + CATLASS_DEVICE + Params() {}; + + CATLASS_DEVICE + Params( + __gm__ ElementScale *ptrScale_, LayoutScale const &layoutScale_, + __gm__ ElementPerTokenScale *ptrPerTokenScale_, LayoutPerTokenScale const &layoutPerTokenScale_, + __gm__ ElementD *ptrD_, LayoutD const &layoutD_) + : ptrScale(ptrScale_), + layoutScale(layoutScale_), + ptrPerTokenScale(ptrPerTokenScale_), + layoutPerTokenScale(layoutPerTokenScale_), + ptrD(ptrD_), + layoutD(layoutD_) + {} + }; + + CATLASS_DEVICE + BlockEpilogue(Arch::Resource const &resource, Params const ¶ms = Params{}) + : params(params) + { + size_t ubOffset = 0; + int32_t eventVMTE2 = 0; + int32_t eventMTE2V = 0; + int32_t eventMTE3V = 0; + int32_t eventVMTE3 = 0; + for (uint32_t i = 0; i < UB_STAGES; ++i) { + ubCList[i] = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += TileShape::COUNT * sizeof(ElementC); + ubScaleList[i] = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += TileShape::COLUMN * sizeof(ElementScale); + ubPerTokenScaleList[i] = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += TileShape::ROW * sizeof(ElementPerTokenScale); + ubDList[i] = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += TileShape::COUNT * sizeof(ElementD); + + eventUbCVMTE2List[i] = eventVMTE2++; + eventUbCMTE2VList[i] = eventMTE2V++; + eventUbScaleVMTE2List[i] = eventVMTE2++; + eventUbScaleMTE2VList[i] = eventMTE2V++; + eventUbPerTokenScaleVMTE2List[i] = eventVMTE2++; + eventUbPerTokenScaleMTE2VList[i] = eventMTE2V++; + eventUbDMTE3VList[i] = eventMTE3V++; + eventUbDVMTE3List[i] = eventVMTE3++; + + AscendC::SetFlag(eventUbCVMTE2List[i]); + AscendC::SetFlag(eventUbScaleVMTE2List[i]); + AscendC::SetFlag(eventUbPerTokenScaleVMTE2List[i]); + AscendC::SetFlag(eventUbDMTE3VList[i]); + } + ubCFp32 = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += TileShape::COUNT * sizeof(float); + ubMul = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += TileShape::COUNT * sizeof(float); + ubPerTokenScaleBrcb = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += TileShape::ROW * BYTE_PER_BLK; + ubPerTokenMul = ubCFp32; + } + + CATLASS_DEVICE + ~BlockEpilogue() + { + for (uint32_t i = 0; i < UB_STAGES; ++i) { + AscendC::WaitFlag(eventUbCVMTE2List[i]); + AscendC::WaitFlag(eventUbScaleVMTE2List[i]); + AscendC::WaitFlag(eventUbPerTokenScaleVMTE2List[i]); + AscendC::WaitFlag(eventUbDMTE3VList[i]); + } + } + + CATLASS_DEVICE + void UpdateParams(Params const ¶ms_) + { + params = params_; + } + + CATLASS_DEVICE + void operator()( + GemmCoord const &blockShapeMNK, + GemmCoord const &blockCoordMNK, + GemmCoord const &actualBlockShapeMNK, + AscendC::GlobalTensor const &gmBlockC, + LayoutC const &layoutBlockC, Callback &&callback = Callback{}) + { + MatrixCoord blockShape = blockShapeMNK.GetCoordMN(); + MatrixCoord blockCoord = blockCoordMNK.GetCoordMN(); + MatrixCoord actualBlockShape = actualBlockShapeMNK.GetCoordMN(); + MatrixCoord blockOffset = blockCoord * blockShape; + + auto ubTileStride = MakeCoord(static_cast(TileShape::COLUMN), 1L); + auto tileShape = TileShape::ToCoord(); + EpilogueTileSwizzle epilogueTileSwizzle(actualBlockShape, tileShape); + uint32_t tileLoops = epilogueTileSwizzle.GetLoops(); + uint32_t subblockIdx = AscendC::GetSubBlockIdx(); + uint32_t subblockNum = AscendC::GetSubBlockNum(); + + if (actualBlockShapeMNK.k() == 0) { + AscendC::GlobalTensor gmD; + gmD.SetGlobalBuffer(params.ptrD); + + ElementD zeroD{}; + for (uint32_t loopIdx = subblockIdx; loopIdx < tileLoops; loopIdx += subblockNum) { + auto tileCoord = epilogueTileSwizzle.GetTileCoord(loopIdx); + auto actualTileShape = epilogueTileSwizzle.GetActualTileShape(tileCoord); + auto tileOffset = blockOffset + tileCoord * tileShape; + + auto &ubD = ubDList[ubListId]; + LayoutD layoutUbD{actualTileShape, ubTileStride}; + + AscendC::WaitFlag(eventUbDMTE3VList[ubListId]); + AscendC::Duplicate(ubD, zeroD, TileShape::COUNT); + AscendC::SetFlag(eventUbDVMTE3List[ubListId]); + + auto gmTileD = gmD[params.layoutD.GetOffset(tileOffset)]; + auto layoutGmTileD = params.layoutD.GetTileLayout(actualTileShape); + + AscendC::WaitFlag(eventUbDVMTE3List[ubListId]); + copyUbToGmD(gmTileD, ubD, layoutGmTileD, layoutUbD); + AscendC::SetFlag(eventUbDMTE3VList[ubListId]); + + ubListId = (ubListId + 1 < UB_STAGES) ? (ubListId + 1) : 0; + } + return; + } + callback(); + + AscendC::GlobalTensor gmScale; + gmScale.SetGlobalBuffer(params.ptrScale); + AscendC::GlobalTensor gmPerTokenScale; + gmPerTokenScale.SetGlobalBuffer(params.ptrPerTokenScale); + AscendC::GlobalTensor gmD; + gmD.SetGlobalBuffer(params.ptrD); + + for (uint32_t loopIdx = subblockIdx; loopIdx < tileLoops; loopIdx += subblockNum) { + auto tileCoord = epilogueTileSwizzle.GetTileCoord(loopIdx); + auto actualTileShape = epilogueTileSwizzle.GetActualTileShape(tileCoord); + auto tileOffsetInBlock = tileCoord * tileShape; + auto tileOffset = blockOffset + tileOffsetInBlock; + + auto gmTileC = gmBlockC[layoutBlockC.GetOffset(tileOffsetInBlock)]; + auto layoutGmTileC = layoutBlockC.GetTileLayout(actualTileShape); + + auto &ubC = ubCList[ubListId]; + LayoutC layoutUbC{actualTileShape, ubTileStride}; + + AscendC::WaitFlag(eventUbCVMTE2List[ubListId]); + copyGmToUbC(ubC, gmTileC, layoutUbC, layoutGmTileC); + AscendC::SetFlag(eventUbCMTE2VList[ubListId]); + + auto scaleTileOffset = tileOffset.template GetCoordByAxis<1>(); + auto scaleTileShape = actualTileShape.template GetCoordByAxis<1>(); + + auto gmTileScale = gmScale[params.layoutScale.GetOffset(scaleTileOffset)]; + auto layoutGmTileScale = params.layoutScale.GetTileLayout(scaleTileShape); + + auto &ubScale = ubScaleList[ubListId]; + auto layoutUbScale = LayoutScale::template MakeLayoutInUb(scaleTileShape); + + AscendC::WaitFlag(eventUbScaleVMTE2List[ubListId]); + copyGmToUbScale(ubScale, gmTileScale, layoutUbScale, layoutGmTileScale); + AscendC::SetFlag(eventUbScaleMTE2VList[ubListId]); + + auto perTokenScaleTileOffset = tileOffset.template GetCoordByAxis<0>(); + auto perTokenScaleTileShape = actualTileShape.template GetCoordByAxis<0>(); + + auto gmTilePerTokenScale = gmPerTokenScale[params.layoutPerTokenScale.GetOffset(perTokenScaleTileOffset)]; + auto layoutGmTilePerTokenScale = params.layoutPerTokenScale.GetTileLayout(perTokenScaleTileShape); + + auto &ubPerTokenScale = ubPerTokenScaleList[ubListId]; + auto layoutUbPerTokenScale = LayoutScale::template MakeLayoutInUb( + perTokenScaleTileShape); + + AscendC::WaitFlag(eventUbPerTokenScaleVMTE2List[ubListId]); + copyGmToUbPerTokenScale(ubPerTokenScale, gmTilePerTokenScale, layoutUbPerTokenScale, + layoutGmTilePerTokenScale); + AscendC::SetFlag(eventUbPerTokenScaleMTE2VList[ubListId]); + + AscendC::WaitFlag(eventUbCMTE2VList[ubListId]); + AscendC::Cast(ubCFp32, ubC, AscendC::RoundMode::CAST_RINT, TileShape::COUNT); + AscendC::SetFlag(eventUbCVMTE2List[ubListId]); + + AscendC::WaitFlag(eventUbScaleMTE2VList[ubListId]); + tileRowBroadcastMul(ubMul, ubCFp32, ubScale); + AscendC::SetFlag(eventUbScaleVMTE2List[ubListId]); + + AscendC::WaitFlag(eventUbPerTokenScaleMTE2VList[ubListId]); + tileBroadcastOneBlk(ubPerTokenScaleBrcb, ubPerTokenScale); + AscendC::SetFlag(eventUbPerTokenScaleVMTE2List[ubListId]); + + AscendC::PipeBarrier(); + tileOneBlkColumnBroadcastMul(ubPerTokenMul, ubMul, ubPerTokenScaleBrcb); + AscendC::PipeBarrier(); + + auto &ubD = ubDList[ubListId]; + LayoutD layoutUbD{actualTileShape, ubTileStride}; + + AscendC::WaitFlag(eventUbDMTE3VList[ubListId]); + AscendC::Cast(ubD, ubPerTokenMul, AscendC::RoundMode::CAST_RINT, TileShape::COUNT); + AscendC::SetFlag(eventUbDVMTE3List[ubListId]); + + auto gmTileD = gmD[params.layoutD.GetOffset(tileOffset)]; + auto layoutGmTileD = params.layoutD.GetTileLayout(actualTileShape); + + AscendC::WaitFlag(eventUbDVMTE3List[ubListId]); + copyUbToGmD(gmTileD, ubD, layoutGmTileD, layoutUbD); + AscendC::SetFlag(eventUbDMTE3VList[ubListId]); + + ubListId = (ubListId + 1 < UB_STAGES) ? (ubListId + 1) : 0; + } + } + +private: + Params params; + + AscendC::LocalTensor ubCList[UB_STAGES]; + AscendC::LocalTensor ubScaleList[UB_STAGES]; + AscendC::LocalTensor ubPerTokenScaleList[UB_STAGES]; + AscendC::LocalTensor ubDList[UB_STAGES]; + + int32_t eventUbCVMTE2List[UB_STAGES]; + int32_t eventUbCMTE2VList[UB_STAGES]; + int32_t eventUbScaleVMTE2List[UB_STAGES]; + int32_t eventUbScaleMTE2VList[UB_STAGES]; + int32_t eventUbPerTokenScaleVMTE2List[UB_STAGES]; + int32_t eventUbPerTokenScaleMTE2VList[UB_STAGES]; + int32_t eventUbDMTE3VList[UB_STAGES]; + int32_t eventUbDVMTE3List[UB_STAGES]; + + uint32_t ubListId{0}; + + AscendC::LocalTensor ubCFp32; + AscendC::LocalTensor ubMul; + AscendC::LocalTensor ubPerTokenScaleBrcb; + AscendC::LocalTensor ubPerTokenMul; + + TileRowBroadcastMul tileRowBroadcastMul; + TileBroadcastOneBlk tileBroadcastOneBlk; + TileOneBlkColumnBroadcastMul tileOneBlkColumnBroadcastMul; + + CopyGmToUbC copyGmToUbC; + CopyGmToUbScale copyGmToUbScale; + CopyGmToUbPerTokenScale copyGmToUbPerTokenScale; + CopyUbToGmD copyUbToGmD; +}; + +} // namespace Catlass::Epilogue::Block + +#endif // CATLASS_KDA_BLOCK_EPILOGUE_PER_TOKEN_DEQUANT_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/kda_dispatch_policy.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/kda_dispatch_policy.hpp new file mode 100644 index 0000000000..f417f91865 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/kda_dispatch_policy.hpp @@ -0,0 +1,143 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_DISPATCH_POLICY_HPP +#define CATLASS_KDA_DISPATCH_POLICY_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_arch.hpp" + +namespace Catlass::Epilogue { + +// For AtlasA2, an element wise epilogue of the form D = C + X, where X is an additional source +struct EpilogueAtlasA2ElemWiseOneSource { + using ArchTag = Arch::AtlasA2; + // Number of operands. Including C, X, and D 3 operands + static constexpr uint32_t OPERANDS_NUM = 3; +}; + +struct EpilogueAtlasA2ElemWiseNoSource { + using ArchTag = Arch::AtlasA2; + // Number of operands. Including C, D 2 operands + static constexpr uint32_t OPERANDS_NUM = 2; +}; + +// For AtlasA2, FA Softmax +struct EpilogueAtlasA2FASoftmax { + using ArchTag = Arch::AtlasA2; +}; + +// For AtlasA2, FA RescaleO +struct EpilogueAtlasA2FARescaleO { + using ArchTag = Arch::AtlasA2; +}; + +// For AtlasA2, MLA Softmax +struct EpilogueAtlasA2MLASoftmax { + using ArchTag = Arch::AtlasA2; +}; + +// For AtlasA2, FA Infer online Softmax no mask +struct EpilogueAtlasA2OnlineSoftmax { + using ArchTag = Arch::AtlasA2; +}; + +// For AtlasA2, FA Infer RescaleO no split row +struct EpilogueAtlasA2RescaleO { + using ArchTag = Arch::AtlasA2; +}; + +// For Ascend950, FA Infer online Softmax +template +struct EpilogueAscend950FASoftmax { + using ArchTag = Arch::Ascend950; + static constexpr bool ATTENTION_MASK_FLAG = ATTENTION_MASK_FLAG_; +}; + +// For Ascend950, FA Infer RescaleO +struct EpilogueAscend950FARescaleO { + using ArchTag = Arch::Ascend950; +}; + +// For AtlasA2, MLA RescaleO +struct EpilogueAtlasA2MLARescaleO { + using ArchTag = Arch::AtlasA2; +}; + +// For AtlasA2, MLA FD RescaleO +template +struct EpilogueAtlasA2MLAFDRescaleO { + using ArchTag = Arch::AtlasA2; + static constexpr uint32_t KV_SPLIT_MAX = 64; + static constexpr uint32_t HEADS_PROCESS_MAX = 16; + static constexpr uint32_t COMPUTE_ELE_NUM = COMPUTE_ELE_NUM_; +}; + +// For AtlasA2, MLA TP1 Softmax +struct EpilogueAtlasA2MLATP1Softmax { + using ArchTag = Arch::AtlasA2; +}; + +// For AtlasA2, MLA TP1 RescaleO +struct EpilogueAtlasA2MLATP1RescaleO { + using ArchTag = Arch::AtlasA2; +}; + +// For AtlasA2, per token dequant +template +struct EpilogueAtlasA2PerTokenDequant { + using ArchTag = Arch::AtlasA2; + static constexpr uint32_t UB_STAGES = UB_STAGES_; +}; + +// For AtlasA2, per token dequant tla version +template +struct EpilogueAtlasA2PerTokenDequantTla { + using ArchTag = Arch::AtlasA2; + static constexpr uint32_t UB_STAGES = UB_STAGES_; +}; + +// For Ascend950, per token dequant +template +struct EpilogueAscend950PerTokenDequantTla { + using ArchTag = Arch::Ascend950; + static constexpr uint32_t UB_STAGES = UB_STAGES_; +}; + +// For Ascend950, perGroup + perBlock dequant +struct BlockEpiloguePertile { + using ArchTag = Arch::Ascend950; +}; + +// For AtlasA2, W4A4 epilogue process +struct EpilogueAtlasA2W4A4PerTokenPerChannelDequant { + using ArchTag = Arch::AtlasA2; +}; +//////////////////////////// +/// new add +// For AtlasA2, GEMM +struct EpilogueAtlasA2Gemm { + using ArchTag = Arch::AtlasA2; +}; + +// For AtlasA2, GEMV +struct EpilogueAtlasA2Gemv { + using ArchTag = Arch::AtlasA2; +}; +/////////////////////////// +// For Ascend950, fixpipe-opti +template +struct EpilogueAscend950Fixpipe { + using ArchTag = Arch::Ascend950; + static constexpr bool SPLIT_M = SPLIT_M_; +}; +} // namespace Catlass::Epilogue + +#endif // CATLASS_KDA_DISPATCH_POLICY_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/tile/copy_ub_to_l1_tla.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/tile/copy_ub_to_l1_tla.hpp new file mode 100644 index 0000000000..90f2a76673 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/tile/copy_ub_to_l1_tla.hpp @@ -0,0 +1,69 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_EPILOGUE_TILE_COPY_UB_TO_L1_TLA_HPP +#define CATLASS_EPILOGUE_TILE_COPY_UB_TO_L1_TLA_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_arch.hpp" +#include "tla/kda_tensor.hpp" +#include "tla/kda_tla_layout.hpp" + +namespace Catlass::Epilogue::Tile { + +template < + class ArchTag, + class TensorSrc, + class TensorDst, + class Enable = void> +struct CopyUb2L1Tla { + static_assert(DEPENDENT_FALSE, "Unsupported CopyUb2L1Tla, can not find the specialization."); +}; + +/// Partial specialization for Ascend950, zN in and zN out. +template +struct CopyUb2L1Tla, LayoutSrc, CoordSrc, AscendC::TPosition::VECCALC>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::A1>, + std::enable_if_t::value && + tla::detail::iszN::value>> { + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(ElementSrc); + + // Methods + + CATLASS_DEVICE + CopyUb2L1Tla() = default; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert(tla::detail::iszNUnAlign::value && + tla::detail::iszN::value && + TensorSrc::position == AscendC::TPosition::VECCALC && + TensorDst::position == AscendC::TPosition::A1, + "The input parameters do not match. TensorSrc must be UB and zN, " + "while TensorDst must be L1 and zN"); + + int64_t srcShape = tla::get<0, 0>(srcTensor.shape()); + AscendC::DataCopyParams dataCopyParams( + tla::get<1, 1>(srcTensor.shape()), + srcShape, + (tla::get<1, 1>(srcTensor.stride()) / ELE_NUM_PER_C0 - srcShape), + (tla::get<1, 1>(dstTensor.stride()) / ELE_NUM_PER_C0 - srcShape)); + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + AscendC::DataCopy(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], dataCopyParams); + } +}; + +} // namespace Catlass::Epilogue::Tile + +#endif // CATLASS_EPILOGUE_TILE_COPY_UB_TO_L1_TLA_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/tile/kda_copy_gm_to_ub.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/tile/kda_copy_gm_to_ub.hpp new file mode 100644 index 0000000000..683f7cefbb --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/tile/kda_copy_gm_to_ub.hpp @@ -0,0 +1,181 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_COPY_GM_TO_UB_HPP +#define CATLASS_KDA_COPY_GM_TO_UB_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_arch.hpp" +#include "catlass/layout/kda_layout.hpp" +#include "catlass/gemm/kda_gemm_type.hpp" + +namespace Catlass::Epilogue::Tile { + +template < + class ArchTag, + class GmType> +struct CopyGm2Ub { + static_assert(DEPENDENT_FALSE, "Unsupported copy gm to ub, can not find the specialization."); +}; + +template +struct CopyGm2Ub> { + using LayoutSrc = layout::RowMajor; + using LayoutDst = layout::RowMajor; + + static constexpr uint32_t ELE_NUM_PER_BLK = BYTE_PER_BLK / sizeof(Element); + + CATLASS_DEVICE + CopyGm2Ub() = default; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + layout::RowMajor const &layoutDst, + layout::RowMajor const &layoutSrc) + { + AscendC::DataCopyExtParams dataCopyParams( + layoutSrc.shape(0), + layoutSrc.shape(1) * sizeof(Element), + (layoutSrc.stride(0) - layoutSrc.shape(1)) * sizeof(Element), + (layoutDst.stride(0) - layoutDst.shape(1)) / ELE_NUM_PER_BLK, + 0); + AscendC::DataCopyPadExtParams padParams(false, 0, 0, 0); + AscendC::DataCopyPad(dstTensor, srcTensor, dataCopyParams, padParams); + }; +}; + +template +struct CopyGm2Ub> { + using LayoutSrc = layout::VectorLayout; + using LayoutDst = layout::VectorLayout; + + static constexpr uint32_t ELE_NUM_PER_BLK = BYTE_PER_BLK / sizeof(Element); + + CATLASS_DEVICE + CopyGm2Ub() = default; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + layout::VectorLayout const &layoutDst, + layout::VectorLayout const &layoutSrc) + { + AscendC::DataCopyExtParams dataCopyParams( + 1, + layoutSrc.shape(0) * sizeof(Element), + 0, + 0, + 0); + AscendC::DataCopyPadExtParams padParams(false, 0, 0, 0); + AscendC::DataCopyPad(dstTensor, srcTensor, dataCopyParams, padParams); + }; +}; + +/// @brief This copy instruction used to copy per token scale from GM to UB. +/// Copy the scale of shape (m,1) on GM to the first column of shape (m,n) on UB, +/// and pad the first block of each row (i.e. pad to shape (m,8) when element type is float). +/// @tparam ArchTag: Architecture tag. +/// @tparam GmType: Type of data on GM. +template < + class ArchTag, + class GmType> +struct CopyPerTokenScale2Ub { + static_assert(std::is_same_v, + "Unsupported layout for CopyPerTokenScale2Ub."); + + using Element = typename GmType::Element; + using LayoutSrc = typename GmType::Layout; + using LayoutDst = layout::RowMajor; + + static constexpr uint32_t ELE_NUM_PER_BLK = BYTE_PER_BLK / sizeof(Element); + + CATLASS_DEVICE + CopyPerTokenScale2Ub() = default; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, + LayoutSrc const &layoutSrc) + { + AscendC::DataCopyExtParams dataCopyParams; + AscendC::DataCopyPadExtParams padParams; + + dataCopyParams.blockCount = layoutSrc.shape(0); + dataCopyParams.blockLen = layoutSrc.shape(1) * sizeof(Element); // per token scale has only one column + dataCopyParams.srcStride = 0; + dataCopyParams.dstStride = (layoutDst.stride(0) - layoutDst.shape(1)) / ELE_NUM_PER_BLK; + // Pad the data to the complete block + padParams.isPad = true; + padParams.leftPadding = 0; + padParams.rightPadding = 0; + + AscendC::DataCopyPad(dstTensor, srcTensor, dataCopyParams, padParams); + } +}; + +template < + class ArchTag, + class GmType> +struct CopyGm2UbAligned { + static_assert(DEPENDENT_FALSE, "Unsupported copy gm to ub aligned, can not find the specialization."); +}; + +template +struct CopyGm2UbAligned> { + using LayoutSrc = layout::RowMajor; + using LayoutDst = layout::RowMajor; + + static constexpr uint32_t ELE_NUM_PER_BLK = BYTE_PER_BLK / sizeof(Element); + static constexpr uint32_t BLOCK_LEN_LIMIT = 65536; + static constexpr uint32_t MAX_REPEAT = 4095; + static constexpr uint32_t STRIDE_LIMIT = 65536; + + CATLASS_DEVICE + CopyGm2UbAligned() = default; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + layout::RowMajor const &layoutDst, + layout::RowMajor const &layoutSrc) + { + uint32_t rows = layoutSrc.shape(0); + uint32_t cols = layoutSrc.shape(1); + uint32_t srcStride = (layoutSrc.stride(0) - layoutSrc.shape(1)) / ELE_NUM_PER_BLK; + uint32_t dstStride = (layoutDst.stride(0) - layoutDst.shape(1)) / ELE_NUM_PER_BLK; + + if ((layoutSrc.shape(1) == layoutSrc.stride(0)) && (layoutDst.shape(1) == layoutDst.stride(0))) { + DataCopy(dstTensor, srcTensor, rows * cols); + } else if (srcStride < STRIDE_LIMIT && dstStride < STRIDE_LIMIT && (cols / ELE_NUM_PER_BLK) < BLOCK_LEN_LIMIT) { + uint32_t rLoops = CeilDiv(rows, MAX_REPEAT); + for (uint32_t i = 0; i < rLoops; ++i) { + uint32_t rActual = (i < rLoops - 1) ? MAX_REPEAT : rows - i * MAX_REPEAT; + AscendC::DataCopyParams dataCopyParams( + rActual, cols / ELE_NUM_PER_BLK, srcStride, dstStride); + DataCopy(dstTensor[i * MAX_REPEAT * layoutDst.stride(0)], + srcTensor[i * MAX_REPEAT * layoutSrc.stride(0)], dataCopyParams); + } + } else { + for (uint32_t i = 0; i < rows; ++i) { + DataCopy(dstTensor[i * layoutDst.stride(0)], srcTensor[i * layoutSrc.stride(0)], cols); + } + } + }; +}; + +} // namespace Catlass::Epilogue::Tile + +#endif diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/tile/kda_copy_gm_to_ub_tla.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/tile/kda_copy_gm_to_ub_tla.hpp new file mode 100644 index 0000000000..380a7464f9 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/tile/kda_copy_gm_to_ub_tla.hpp @@ -0,0 +1,143 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_COPY_GM_TO_UB_TLA_HPP +#define CATLASS_KDA_COPY_GM_TO_UB_TLA_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_arch.hpp" +#include "tla/kda_tensor.hpp" +#include "tla/kda_tla_layout.hpp" + +namespace Catlass::Epilogue::Tile { + +template < + class ArchTag, + class TensorSrc, + class TensorDst, + class Enable = void> +struct CopyGm2UbTla { + static_assert(DEPENDENT_FALSE, "Unsupported CopyGm2UbTla, can not find the specialization."); +}; + +/// Partial specialization for AtlasA2, RowMajor in and RowMajor out. +template +struct CopyGm2UbTla, LayoutSrc, CoordSrc, AscendC::TPosition::GM>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::VECCALC>, + std::enable_if_t::value && + tla::detail::isRowMajor::value>> { + static constexpr uint32_t ELE_NUM_PER_BLK = BYTE_PER_BLK / sizeof(ElementSrc); + + // Methods + + CATLASS_DEVICE + CopyGm2UbTla() = default; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert(tla::detail::isRowMajor::value && + tla::detail::isRowMajor::value && + TensorSrc::position == AscendC::TPosition::GM && + TensorDst::position == AscendC::TPosition::VECCALC, + "The input parameters do not match. TensorSrc must be GM and RowMajor, " + "while TensorDst must be UB and RowMajor"); + + AscendC::DataCopyExtParams dataCopyParams( + tla::get<0>(srcTensor.shape()), + tla::get<1>(srcTensor.shape()) * sizeof(ElementSrc), + (tla::get<0>(srcTensor.stride()) - tla::get<1>(srcTensor.shape())) * sizeof(ElementSrc), + (tla::get<0>(dstTensor.stride()) - tla::get<1>(dstTensor.shape())) / ELE_NUM_PER_BLK, + 0); + AscendC::DataCopyPadExtParams padParams(false, 0, 0, 0); + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + AscendC::DataCopyPad(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], dataCopyParams, padParams); + }; +}; + +/// Partial specialization for Ascend950, Vector in and Vector out. +template +struct CopyGm2UbTla, LayoutSrc, CoordSrc, AscendC::TPosition::GM>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::VECCALC>, + std::enable_if_t::value && + tla::detail::isVector::value>> { + static constexpr uint32_t ELE_NUM_PER_BLK = BYTE_PER_BLK / sizeof(ElementSrc); + + // Methods + + CATLASS_DEVICE + CopyGm2UbTla() = default; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert(tla::detail::isVector::value && + tla::detail::isVector::value && + TensorSrc::position == AscendC::TPosition::GM && + TensorDst::position == AscendC::TPosition::VECCALC, + "The input parameters do not match. TensorSrc must be GM and Vector, " + "while TensorDst must be UB and Vector"); + + AscendC::DataCopyExtParams dataCopyParams( + 1, + tla::get<0>(srcTensor.shape()) * sizeof(ElementSrc), + 0, + 0, + 0); + AscendC::DataCopyPadExtParams padParams(false, 0, 0, 0); + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + AscendC::DataCopyPad(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], dataCopyParams, padParams); + }; +}; + +/// Partial specialization for Ascend950, RowMajor in and RowMajor out. +template +struct CopyGm2UbTla, LayoutSrc, CoordSrc, AscendC::TPosition::GM>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::VECCALC>, + std::enable_if_t::value && + tla::detail::isRowMajor::value>> { + static constexpr uint32_t ELE_NUM_PER_BLK = BYTE_PER_BLK / sizeof(ElementSrc); + + // Methods + + CATLASS_DEVICE + CopyGm2UbTla() = default; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert(tla::detail::isRowMajor::value && + tla::detail::isRowMajor::value && + TensorSrc::position == AscendC::TPosition::GM && + TensorDst::position == AscendC::TPosition::VECCALC, + "The input parameters do not match. TensorSrc must be GM and RowMajor, " + "while TensorDst must be UB and RowMajor"); + + AscendC::DataCopyExtParams dataCopyParams( + tla::get<0>(srcTensor.shape()), + tla::get<1>(srcTensor.shape()) * sizeof(ElementSrc), + (tla::get<0>(srcTensor.stride()) - tla::get<1>(srcTensor.shape())) * sizeof(ElementSrc), + (tla::get<0>(dstTensor.stride()) - tla::get<1>(dstTensor.shape())) / ELE_NUM_PER_BLK, + 0); + AscendC::DataCopyPadExtParams padParams(false, 0, 0, 0); + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + AscendC::DataCopyPad(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], dataCopyParams, padParams); + }; +}; + +} // namespace Catlass::Epilogue::Tile + +#endif // CATLASS_KDA_COPY_GM_TO_UB_TLA_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/tile/kda_copy_ub_to_gm.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/tile/kda_copy_ub_to_gm.hpp new file mode 100644 index 0000000000..8fc513fcf9 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/tile/kda_copy_ub_to_gm.hpp @@ -0,0 +1,136 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_COPY_UB_TO_GM_HPP +#define CATLASS_KDA_COPY_UB_TO_GM_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_arch.hpp" +#include "catlass/layout/kda_layout.hpp" +#include "catlass/gemm/kda_gemm_type.hpp" + +namespace Catlass::Epilogue::Tile { + +template < + class ArchTag, + class GmType> +struct CopyUb2Gm { + static_assert(DEPENDENT_FALSE, "Unsupported copy ub to gm, can not find the specialization."); +}; + +template +struct CopyUb2Gm> { + using LayoutDst = layout::RowMajor; + using LayoutSrc = layout::RowMajor; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + + CATLASS_DEVICE + CopyUb2Gm() = default; + + CATLASS_DEVICE + void operator()( + AscendC::GlobalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + layout::RowMajor const &layoutDst, + layout::RowMajor const &layoutSrc) + { + AscendC::DataCopyExtParams dataCopyParams( + layoutDst.shape(0), + layoutDst.shape(1) * sizeof(Element), + (layoutSrc.stride(0) - layoutSrc.shape(1)) / ELE_NUM_PER_C0, + (layoutDst.stride(0) - layoutDst.shape(1)) * sizeof(Element), + 0); + AscendC::DataCopyPad(dstTensor, srcTensor, dataCopyParams); + } +}; + +// new add vectorlayout version +template +struct CopyUb2Gm> { + using LayoutSrc = layout::VectorLayout; + using LayoutDst = layout::VectorLayout; + + static constexpr uint32_t ELE_NUM_PER_BLK = BYTE_PER_BLK / sizeof(Element); + + CATLASS_DEVICE + CopyUb2Gm() = default; + + CATLASS_DEVICE + void operator()( + AscendC::GlobalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + layout::VectorLayout const &layoutDst, + layout::VectorLayout const &layoutSrc) + { + AscendC::DataCopyExtParams dataCopyParams( + 1, + layoutDst.shape(0) * sizeof(Element), + 0, + 0, + 0); + AscendC::DataCopyPad(dstTensor, srcTensor, dataCopyParams); + }; +}; + +template < + class ArchTag, + class GmType> +struct CopyUb2GmAligned { + static_assert(DEPENDENT_FALSE, "Unsupported copy ub to gm aligned, can not find the specialization."); +}; + +template +struct CopyUb2GmAligned> { + using LayoutSrc = layout::RowMajor; + using LayoutDst = layout::RowMajor; + + static constexpr uint32_t ELE_NUM_PER_BLK = BYTE_PER_BLK / sizeof(Element); + static constexpr uint32_t BLOCK_LEN_LIMIT = 65536; + static constexpr uint32_t MAX_REPEAT = 4095; + static constexpr uint32_t STRIDE_LIMIT = 65536; + + CATLASS_DEVICE + CopyUb2GmAligned() = default; + + CATLASS_DEVICE + void operator()( + AscendC::GlobalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + layout::RowMajor const &layoutDst, + layout::RowMajor const &layoutSrc) + { + uint32_t rows = layoutDst.shape(0); + uint32_t cols = layoutDst.shape(1); + uint32_t srcStride = (layoutSrc.stride(0) - layoutSrc.shape(1)) / ELE_NUM_PER_BLK; + uint32_t dstStride = (layoutDst.stride(0) - layoutDst.shape(1)) / ELE_NUM_PER_BLK; + + if ((layoutSrc.shape(1) == layoutSrc.stride(0)) && (layoutDst.shape(1) == layoutDst.stride(0))) { + DataCopy(dstTensor, srcTensor, rows * cols); + } else if (srcStride < STRIDE_LIMIT && dstStride < STRIDE_LIMIT && (cols / ELE_NUM_PER_BLK) < BLOCK_LEN_LIMIT) { + uint32_t rLoops = CeilDiv(rows, MAX_REPEAT); + for (uint32_t i = 0; i < rLoops; ++i) { + uint32_t rActual = (i < rLoops - 1) ? MAX_REPEAT : rows - i * MAX_REPEAT; + AscendC::DataCopyParams dataCopyParams( + rActual, cols / ELE_NUM_PER_BLK, srcStride, dstStride); + DataCopy(dstTensor[i * MAX_REPEAT * layoutDst.stride(0)], + srcTensor[i * MAX_REPEAT * layoutSrc.stride(0)], dataCopyParams); + } + } else { + for (uint32_t i = 0; i < rows; ++i) { + DataCopy(dstTensor[i * layoutDst.stride(0)], srcTensor[i * layoutSrc.stride(0)], cols); + } + } + }; +}; + +} // namespace Catlass::Epilogue::Tile + +#endif diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/tile/kda_copy_ub_to_gm_tla.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/tile/kda_copy_ub_to_gm_tla.hpp new file mode 100644 index 0000000000..504b79a3d9 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/tile/kda_copy_ub_to_gm_tla.hpp @@ -0,0 +1,104 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_COPY_UB_TO_GM_TLA_HPP +#define CATLASS_KDA_COPY_UB_TO_GM_TLA_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_arch.hpp" +#include "tla/kda_tensor.hpp" +#include "tla/kda_tla_layout.hpp" + +namespace Catlass::Epilogue::Tile { + +template < + class ArchTag, + class TensorSrc, + class TensorDst, + class Enable = void> +struct CopyUb2GmTla { + static_assert(DEPENDENT_FALSE, "Unsupported CopyUb2GmTla, can not find the specialization."); +}; + +/// Partial specialization for AtlasA2, RowMajor in and RowMajor out. +template +struct CopyUb2GmTla, LayoutSrc, CoordSrc, AscendC::TPosition::VECCALC>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::GM>, + std::enable_if_t::value && + tla::detail::isRowMajor::value>> { + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(ElementSrc); + + // Methods + + CATLASS_DEVICE + CopyUb2GmTla() = default; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert(tla::detail::isRowMajor::value && + tla::detail::isRowMajor::value && + TensorSrc::position == AscendC::TPosition::VECCALC && + TensorDst::position == AscendC::TPosition::GM, + "The input parameters do not match. TensorSrc must be UB and RowMajor, " + "while TensorDst must be GM and RowMajor"); + + AscendC::DataCopyExtParams dataCopyParams( + tla::get<0>(dstTensor.shape()), + tla::get<1>(dstTensor.shape()) * sizeof(ElementSrc), + (tla::get<0>(srcTensor.stride()) - tla::get<1>(srcTensor.shape())) / ELE_NUM_PER_C0, + (tla::get<0>(dstTensor.stride()) - tla::get<1>(dstTensor.shape())) * sizeof(ElementSrc), + 0); + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + AscendC::DataCopyPad(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], dataCopyParams); + }; +}; + +/// Partial specialization for Ascend950, RowMajor in and RowMajor out. +template +struct CopyUb2GmTla, LayoutSrc, CoordSrc, AscendC::TPosition::VECCALC>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::GM>, + std::enable_if_t::value && + tla::detail::isRowMajor::value>> { + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(ElementSrc); + + // Methods + + CATLASS_DEVICE + CopyUb2GmTla() = default; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert(tla::detail::isRowMajor::value && + tla::detail::isRowMajor::value && + TensorSrc::position == AscendC::TPosition::VECCALC && + TensorDst::position == AscendC::TPosition::GM, + "The input parameters do not match. TensorSrc must be UB and RowMajor, " + "while TensorDst must be GM and RowMajor"); + + AscendC::DataCopyExtParams dataCopyParams( + tla::get<0>(dstTensor.shape()), + tla::get<1>(dstTensor.shape()) * sizeof(ElementSrc), + (tla::get<0>(srcTensor.stride()) - tla::get<1>(srcTensor.shape())) / ELE_NUM_PER_C0, + (tla::get<0>(dstTensor.stride()) - tla::get<1>(dstTensor.shape())) * sizeof(ElementSrc), + 0); + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + AscendC::DataCopyPad(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], dataCopyParams); + }; +}; + +} // namespace Catlass::Epilogue::Tile + +#endif // CATLASS_KDA_COPY_UB_TO_GM_TLA_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/tile/kda_epilogue_tile_copy.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/tile/kda_epilogue_tile_copy.hpp new file mode 100644 index 0000000000..74cd66dc2f --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/tile/kda_epilogue_tile_copy.hpp @@ -0,0 +1,192 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_EPILOGUE_TILE_COPY_HPP +#define CATLASS_KDA_EPILOGUE_TILE_COPY_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_arch.hpp" +#include "catlass/detail/kda_tag_to_layout.hpp" +#include "catlass/epilogue/tile/kda_copy_gm_to_ub.hpp" +#include "catlass/epilogue/tile/kda_copy_ub_to_gm.hpp" +#include "catlass/epilogue/tile/kda_copy_gm_to_ub_tla.hpp" +#include "catlass/epilogue/tile/kda_copy_ub_to_gm_tla.hpp" +#include "tla/kda_tensor.hpp" + +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 3510) +#include "catlass/epilogue/tile/copy_ub_to_l1_tla.hpp" +#endif + +namespace Catlass::Epilogue::Tile { + +template < + /// Tag indicating architecture + class ArchTag, + class... Args> +struct TileCopy { + static_assert(DEPENDENT_FALSE, "Unsupported tile copy, can not find the specialization."); +}; + +template < + class ArchTag, + /// GemmType for C matrix operand + class CType, + /// GemmType for D matrix operand + class DType> +struct TileCopy { + using ElementC = typename CType::Element; + using ElementD = typename DType::Element; + + using CopyGmToUbC = CopyGm2Ub; + using CopyUbToGmD = CopyUb2Gm; +}; + +template < + class ArchTag, + /// GemmType for C matrix operand + class CType, + /// GemmType for X matrix operand + class XType, + /// GemmType for D matrix operand + class DType> +struct TileCopy { + using ElementC = typename CType::Element; + using ElementX = typename XType::Element; + using ElementD = typename DType::Element; + + using CopyGmToUbC = CopyGm2Ub; + using CopyGmToUbX = CopyGm2Ub; + using CopyUbToGmD = CopyUb2Gm; +}; + +template < + class ArchTag, + class CType, + class XType, + class YType, + class DType> +struct TileCopy { + using ElementC = typename CType::Element; + using ElementX = typename XType::Element; + using ElementY = typename YType::Element; + using ElementD = typename DType::Element; + + using CopyGmToUbC = CopyGm2Ub; + using CopyGmToUbX = CopyGm2Ub; + using CopyGmToUbY = CopyGm2Ub; + using CopyUbToGmD = CopyUb2Gm; +}; + +template < + class ArchTag, + class CType, + class XType, + class YType, + class DType> +struct TileCopyBf16 { + using ElementC = typename CType::Element; + using ElementX = bfloat16_t; + using ElementY = bfloat16_t; + using ElementD = bfloat16_t; + + using CopyGmToUbC = CopyGm2Ub; + using CopyGmToUbX = CopyGm2Ub>; + using CopyGmToUbY = CopyGm2Ub>; + using CopyUbToGmD = CopyUb2Gm>; +}; + +template < + class ArchTag, + class CType, + class ScaleType, + class PerTokenScaleType, + class DType> +struct TileCopyPerTokenDequant { + using ElementC = typename CType::Element; + using ElementScale = typename ScaleType::Element; + using ElementPerTokenScale = typename PerTokenScaleType::Element; + using ElementD = typename DType::Element; + + using CopyGmToUbC = CopyGm2Ub; + using CopyGmToUbScale = CopyGm2Ub; + using CopyGmToUbPerTokenScale = CopyPerTokenScale2Ub; + using CopyUbToGmD = CopyUb2Gm; +}; + +template < + class ArchTag, + class CType, + class PerTokenScaleType, + class DType> +struct TileCopyW4A4Gemm { + using ElementC = typename CType::Element; + using ElementPerTokenScale = typename PerTokenScaleType::Element; + using ElementD = typename DType::Element; + + using CopyGmToUbC = CopyGm2Ub; + using CopyGmToUbPerTokenScale = CopyGm2Ub; + using CopyUbToGmD = CopyUb2Gm; +}; + +template < + class ArchTag, + /// GemmType for C matrix operand + class ElementC_, + class LayoutTagC_, + /// GemmType for X matrix operand + class ElementX_, + class LayoutTagX_, + /// GemmType for Y matrix operand + class ElementY_, + class LayoutTagY_, + /// GemmType for D matrix operand + class ElementD_, + class LayoutTagD_> +struct TileCopyDequantTla { + using ElementC = ElementC_; + using LayoutTagC = LayoutTagC_; + using LayoutC = detail::TagToLayout_t; + using TensorUbC = + tla::Tensor, LayoutC, tla::Coord, AscendC::TPosition::VECCALC>; + + using ElementX = ElementX_; + using LayoutTagX = LayoutTagX_; + using LayoutX = detail::TagToLayout_t; + using TensorUbX = + tla::Tensor, LayoutX, tla::Coord, AscendC::TPosition::VECCALC>; + + using ElementY = ElementY_; + using LayoutTagY = LayoutTagY_; + using LayoutY = detail::TagToLayout_t; + using TensorUbY = + tla::Tensor, LayoutY, tla::Coord, AscendC::TPosition::VECCALC>; + + using ElementD = ElementD_; + using LayoutTagD = LayoutTagD_; + using LayoutD = detail::TagToLayout_t; + using TensorUbD = + tla::Tensor, LayoutD, tla::Coord, AscendC::TPosition::VECCALC>; + + template + using CopyGmToUbC = CopyGm2UbTla; + + template + using CopyGmToUbX = CopyGm2UbTla; + + template + using CopyGmToUbY = CopyGm2UbTla; + + template + using CopyUbToGmD = CopyUb2GmTla; +}; + +} // namespace Catlass::Epilogue::Tile + +#endif // CATLASS_KDA_EPILOGUE_TILE_COPY_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/tile/kda_tile_cast.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/tile/kda_tile_cast.hpp new file mode 100644 index 0000000000..38e2e43c0e --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/tile/kda_tile_cast.hpp @@ -0,0 +1,46 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_TILE_CAST_HPP +#define CATLASS_KDA_TILE_CAST_HPP + +#include "catlass/kda_catlass.hpp" + +namespace Catlass::Epilogue::Tile { + +template < + /// Tag indicating architecture + class ArchTag_, + /// Compute data type + class DstType_, + class SrcType_, + /// Length of the compute buffer + class TileShape_> +struct TileCast { + using ArchTag = ArchTag_; + using ElementDst = typename DstType_::Element; + using ElementSrc = typename SrcType_::Element; + using TileShape = TileShape_; + + CATLASS_DEVICE + TileCast() {} + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &ubOut, + AscendC::LocalTensor const &ubIn) + { + AscendC::Cast(ubOut, ubIn, AscendC::RoundMode::CAST_RINT, TileShape::COUNT); + } +}; + +} // namespace Catlass::Epilogue::Tile + +#endif diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_dynamic_aiv.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_dynamic_aiv.hpp new file mode 100644 index 0000000000..0817c0772c --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_dynamic_aiv.hpp @@ -0,0 +1,505 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_GEMM_BLOCK_MMAD_DYNAMIC_AIV_HPP +#define CATLASS_GEMM_BLOCK_MMAD_DYNAMIC_AIV_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/kda_coord.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" +#include "catlass/gemm/kda_gemm_helper.hpp" +#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" +#include "catlass/gemm/tile/kda_tile_muls.hpp" + +namespace Catlass::Gemm::Block { +template < + uint32_t SCALAR_BUFFER_ELE_NUM_, + uint32_t STAGES_, + class AType_, + class BType_, + class CType_, + class BiasType_, + class TileCopy_, + class TileVmuls_> +struct BlockMmadAiv< + MmadAtlasA2DynamicAiv, + AType_, + BType_, + CType_, + BiasType_, + TileCopy_, + TileVmuls_> { +public: + // Type Aliases + using DispatchPolicy = MmadAtlasA2DynamicAiv; + using ArchTag = typename DispatchPolicy::ArchTag; + using ElementA = typename AType_::Element; + using LayoutA = typename AType_::Layout; + using ElementB = typename BType_::Element; + using LayoutB = typename BType_::Layout; + using ElementC = typename CType_::Element; + using LayoutC = typename CType_::Layout; + using CopyGmToUbA = typename TileCopy_::CopyGmToUbA; + using CopyGmToUbB = typename TileCopy_::CopyGmToUbB; + using CopyUbToGmC = typename TileCopy_::CopyUbToGmC; + using TileVmuls = TileVmuls_; + using TensorCoord = layout::VectorLayout::TensorCoord; + + static constexpr uint32_t UBA_ALIGN = BYTE_PER_BLK / sizeof(ElementA); + static constexpr uint32_t UBB_ALIGN = BYTE_PER_BLK / sizeof(ElementB); + static constexpr uint32_t STAGES = DispatchPolicy::STAGES; + static constexpr uint32_t SCALAR_BUFFER_ELE_NUM = DispatchPolicy::SCALAR_BUFFER_ELE_NUM; + + // Check LayoutC + static_assert(std::is_same_v, "LayoutC only support RowMajor yet!"); + + /// Construct + CATLASS_DEVICE + BlockMmadAiv(Arch::Resource &resource, MatrixCoord taskShape, uint32_t ubBufAddrStart = 0) + { + uint32_t ubASize = SCALAR_BUFFER_ELE_NUM * sizeof(ElementA); + uint32_t ubBSize = taskShape.column() * sizeof(ElementB); + uint32_t ubCSize = ArchTag::UB_SIZE / STAGES - ubASize - ubBSize; + + ubTileNum = ubCSize / ubBSize; + ubTileSize = ubBSize / sizeof(ElementB); + + uint32_t ubAOffset = ubBufAddrStart; + uint32_t ubBOffset = ubBufAddrStart + ubASize * STAGES; + uint32_t ubCOffset = ubBOffset + ubBSize * STAGES; + // Init buffers + for (uint32_t i = 0; i < STAGES; i++) { + ubATensorList[i] = resource.ubBuf.template GetBufferByByte(ubAOffset + ubASize * i); + ubBTensorList[i] = resource.ubBuf.template GetBufferByByte(ubBOffset + ubBSize * i); + ubCTensorList[i] = resource.ubBuf.template GetBufferByByte(ubCOffset + ubCSize * i); + + inEventIds[i] = i; + outEventIds[i] = i; + AscendC::SetFlag(inEventIds[i]); + AscendC::SetFlag(outEventIds[i]); + } + } + + /// Destructor + CATLASS_DEVICE + ~BlockMmadAiv() + { + for (uint32_t i = 0; i < STAGES; i++) { + AscendC::WaitFlag(inEventIds[i]); + AscendC::WaitFlag(outEventIds[i]); + } + } + + /// Perform a block-scoped matrix multiply-accumulate + CATLASS_DEVICE + void operator()( + AscendC::GlobalTensor const &gmBlockA, LayoutA const &layoutA, + AscendC::GlobalTensor const &gmBlockB, LayoutB const &layoutB, + AscendC::GlobalTensor const &gmBlockC, LayoutC const &layoutC, + GemmCoord const &actualShape) + { + AscendC::WaitFlag((event_t)(inEventIds[UbInStageId])); + + uint32_t mRound = RoundUp(actualShape.m()); + ElementA scalarA[SCALAR_BUFFER_ELE_NUM]; + LayoutA layoutAInUb{mRound}; + LayoutB layoutBInUb{ubTileSize}; + LayoutC layoutCInUb{ubTileNum, ubTileSize}; + + // load B + auto layoutTileB = layoutB.GetTileLayout(TensorCoord(actualShape.n())); + copyGmToUbB(ubBTensorList[UbInStageId], gmBlockB[0], layoutBInUb, layoutTileB); + AscendC::SetFlag((event_t)(inEventIds[UbInStageId])); + + // load A + auto layoutTileA = layoutA.GetTileLayout(TensorCoord(actualShape.m())); + copyGmToUbA(ubATensorList[UbInStageId], gmBlockA[0], layoutAInUb, layoutTileA); + AscendC::SetFlag((event_t)(inEventIds[UbInStageId])); + AscendC::WaitFlag((event_t)(inEventIds[UbInStageId])); + for (uint32_t i = 0; i < actualShape.m(); i++) { + scalarA[i] = ubATensorList[UbInStageId].GetValue(i); + } + AscendC::SetFlag((event_t)(inEventIds[UbInStageId])); + AscendC::WaitFlag((event_t)(inEventIds[UbInStageId])); + AscendC::WaitFlag((event_t)(inEventIds[UbInStageId])); + + auto mTiles = CeilDiv(actualShape.m(), ubTileNum); + uint32_t lastTileM = actualShape.m() - (mTiles - 1) * ubTileNum; + // tile vmuls && write c to gm + for (uint32_t i = 0; i < mTiles; i++) { + auto actualTileM = (i == mTiles - 1) ? lastTileM : ubTileNum; + AscendC::WaitFlag((event_t)(outEventIds[UbOutStageId])); + for (uint32_t j = 0; j < actualTileM; j++) { + tileVmuls(ubCTensorList[UbOutStageId][j * ubTileSize], ubBTensorList[UbInStageId], + scalarA[i * ubTileNum + j], actualShape.n()); + } + AscendC::PipeBarrier(); + AscendC::SetFlag((event_t)(outEventIds[UbOutStageId])); + AscendC::WaitFlag((event_t)(outEventIds[UbOutStageId])); + MatrixCoord gmTileCOffset{i * ubTileNum, 0}; + auto gmTileC = gmBlockC[layoutC.GetOffset(gmTileCOffset)]; + LayoutC layoutCBlock = layoutC.GetTileLayout(MakeCoord(actualTileM, actualShape.n())); + LayoutC lastLayoutCInUb = layoutCInUb.GetTileLayout(MakeCoord(actualTileM, actualShape.n())); + copyUbToGmC(gmTileC, ubCTensorList[UbOutStageId], layoutCBlock, lastLayoutCInUb); + AscendC::SetFlag((event_t)(outEventIds[UbOutStageId])); + UbOutStageId = (UbOutStageId + 1 < STAGES) ? (UbOutStageId + 1) : 0; + } + AscendC::SetFlag((event_t)(inEventIds[UbInStageId])); + UbInStageId = (UbInStageId + 1 < STAGES) ? (UbInStageId + 1) : 0; + } + +protected: + /// Data members + AscendC::LocalTensor ubATensorList[STAGES]; + AscendC::LocalTensor ubBTensorList[STAGES]; + AscendC::LocalTensor ubCTensorList[STAGES]; + + CopyGmToUbA copyGmToUbA; + CopyGmToUbB copyGmToUbB; + CopyUbToGmC copyUbToGmC; + TileVmuls tileVmuls; + + int32_t inEventIds[STAGES]; + int32_t outEventIds[STAGES]; + + uint32_t UbOutStageId{0}; + uint32_t UbInStageId{0}; + + uint32_t ubTileNum{0}; + uint32_t ubTileSize{0}; +}; + +template < + uint32_t SCALAR_BUFFER_ELE_NUM_, + bool IS_TILE_M_, + class AType_, + class BType_, + class CType_, + class BiasType_, + class TileCopy_, + class TileVmuls_> +struct BlockMmadAiv< + MmadAtlasA2DynamicAivSimple, + AType_, + BType_, + CType_, + BiasType_, + TileCopy_, + TileVmuls_> { +public: + // Type Aliases + using DispatchPolicy = MmadAtlasA2DynamicAivSimple; + using ArchTag = typename DispatchPolicy::ArchTag; + using ElementA = typename AType_::Element; + using LayoutA = typename AType_::Layout; + using ElementB = typename BType_::Element; + using LayoutB = typename BType_::Layout; + using ElementC = typename CType_::Element; + using LayoutC = typename CType_::Layout; + + using CopyGmToUbA = typename TileCopy_::CopyGmToUbA; + using CopyGmToUbB = typename TileCopy_::CopyGmToUbB; + using CopyUbToGmC = typename TileCopy_::CopyUbToGmC; + using TileVmuls = TileVmuls_; + using TensorCoord = layout::VectorLayout::TensorCoord; + + static constexpr uint32_t UBA_ALIGN = BYTE_PER_BLK / sizeof(ElementA); + static constexpr uint32_t UBB_ALIGN = BYTE_PER_BLK / sizeof(ElementB); + static constexpr uint32_t SCALAR_BUFFER_ELE_NUM = DispatchPolicy::SCALAR_BUFFER_ELE_NUM; + static constexpr bool IS_TILE_M = DispatchPolicy::IS_TILE_M; + + // Check LayoutC + static_assert(std::is_same_v, "LayoutC only support RowMajor yet!"); + + /// Construct + CATLASS_DEVICE + BlockMmadAiv(Arch::Resource &resource, MatrixCoord taskShape, uint32_t ubBufAddrStart = 0) + { + uint32_t ubASize = SCALAR_BUFFER_ELE_NUM * sizeof(ElementA); + uint32_t ubBSize = taskShape.column() * sizeof(ElementB); + uint32_t ubCSize = ArchTag::UB_SIZE - ubASize - ubBSize; + + ubTileNum = ubCSize / ubBSize; + ubTileSize = ubBSize / sizeof(ElementB); + + uint32_t ubAOffset = ubBufAddrStart; + uint32_t ubBOffset = ubBufAddrStart + ubASize; + uint32_t ubCOffset = ubBOffset + ubBSize; + // Init buffers + ubATensor = resource.ubBuf.template GetBufferByByte(ubAOffset); + ubBTensor = resource.ubBuf.template GetBufferByByte(ubBOffset); + ubCTensor = resource.ubBuf.template GetBufferByByte(ubCOffset); + + AscendC::SetFlag(EVENT_ID0); + AscendC::SetFlag(EVENT_ID1); + AscendC::SetFlag(EVENT_ID0); + } + + /// Destructor + CATLASS_DEVICE + ~BlockMmadAiv() + { + AscendC::WaitFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID1); + AscendC::WaitFlag(EVENT_ID0); + } + + /// Perform a block-scoped matrix multiply-accumulate + CATLASS_DEVICE + void operator()( + AscendC::GlobalTensor const &gmBlockA, LayoutA const &layoutA, + AscendC::GlobalTensor const &gmBlockB, LayoutB const &layoutB, + AscendC::GlobalTensor const &gmBlockC, LayoutC const &layoutC, + GemmCoord const &actualShape) + { + uint32_t mRound = RoundUp(actualShape.m()); + + AscendC::WaitFlag((event_t)(EVENT_ID0)); + AscendC::WaitFlag((event_t)(EVENT_ID1)); + AscendC::WaitFlag((event_t)(EVENT_ID0)); + + ElementA scalarA[SCALAR_BUFFER_ELE_NUM]; + LayoutA layoutAInUb{mRound}; + LayoutB layoutBInUb{ubTileSize}; + LayoutC layoutCInUb{ubTileNum, ubTileSize}; + + // load B + auto layoutTileB = layoutB.GetTileLayout(TensorCoord(actualShape.n())); + copyGmToUbB(ubBTensor, gmBlockB[0], layoutBInUb, layoutTileB); + AscendC::SetFlag((event_t)(EVENT_ID1)); + // load A + auto layoutTileA = layoutA.GetTileLayout(TensorCoord(actualShape.m())); + copyGmToUbA(ubATensor, gmBlockA[0], layoutAInUb, layoutTileA); + AscendC::SetFlag((event_t)(EVENT_ID0)); + AscendC::SetFlag((event_t)(EVENT_ID0)); + AscendC::WaitFlag((event_t)(EVENT_ID0)); + if constexpr (IS_TILE_M) { + for (uint32_t i = 0; i < actualShape.m(); i++) { + scalarA[i] = ubATensor.GetValue(i); + } + } else { + scalarA[0] = ubATensor.GetValue(0); + } + AscendC::WaitFlag((event_t)(EVENT_ID0)); + AscendC::WaitFlag((event_t)(EVENT_ID1)); + AscendC::SetFlag((event_t)(EVENT_ID0)); + AscendC::WaitFlag((event_t)(EVENT_ID0)); + + // compute && write c to gm + LayoutC layoutCBlock = layoutC.GetTileLayout(MakeCoord(actualShape.m(), actualShape.n())); + LayoutC lastLayoutCInUb = layoutCInUb.GetTileLayout(MakeCoord(actualShape.m(), actualShape.n())); + if constexpr (IS_TILE_M) { + for (uint32_t j = 0; j < actualShape.m(); j++) { + tileVmuls(ubCTensor[j * ubTileSize], ubBTensor, scalarA[j], actualShape.n()); + } + } else { + tileVmuls(ubCTensor[0], ubBTensor, scalarA[0], actualShape.n()); + } + + AscendC::SetFlag((event_t)(EVENT_ID0)); + AscendC::SetFlag((event_t)(EVENT_ID1)); + AscendC::SetFlag((event_t)(EVENT_ID0)); + AscendC::WaitFlag((event_t)(EVENT_ID0)); + copyUbToGmC(gmBlockC[0], ubCTensor, layoutCBlock, lastLayoutCInUb); + AscendC::SetFlag((event_t)(EVENT_ID0)); + } + +protected: + /// Data members + AscendC::LocalTensor ubATensor; + AscendC::LocalTensor ubBTensor; + AscendC::LocalTensor ubCTensor; + + CopyGmToUbA copyGmToUbA; + CopyGmToUbB copyGmToUbB; + CopyUbToGmC copyUbToGmC; + TileVmuls tileVmuls; + + uint32_t ubTileNum{0}; + uint32_t ubTileSize{0}; +}; + +template < + uint32_t SCALAR_BUFFER_ELE_NUM_, + class AType_, + class BType_, + class CType_, + class BiasType_, + class TileCopy_, + class TileVmuls_> +struct BlockMmadAiv< + MmadAtlasA2DynamicAivTrans, + AType_, + BType_, + CType_, + BiasType_, + TileCopy_, + TileVmuls_> { +public: + // Type Aliases + using DispatchPolicy = MmadAtlasA2DynamicAivTrans; + using ArchTag = typename DispatchPolicy::ArchTag; + using ElementA = typename AType_::Element; + using LayoutA = typename AType_::Layout; + using ElementB = typename BType_::Element; + using LayoutB = typename BType_::Layout; + using ElementC = typename CType_::Element; + using LayoutC = typename CType_::Layout; + using TileVmuls = TileVmuls_; + using TensorCoord = layout::VectorLayout::TensorCoord; + + using CopyGmToUbB = typename TileCopy_::CopyGmToUbB; + using CopyGmToUbA = typename TileCopy_::CopyGmToUbA; + using CopyUbToGmC = typename TileCopy_::CopyUbToGmC; + + static constexpr uint32_t UBA_ALIGN = BYTE_PER_BLK / sizeof(ElementA); + static constexpr uint32_t UBB_ALIGN = BYTE_PER_BLK / sizeof(ElementB); + static constexpr uint32_t SCALAR_BUFFER_ELE_NUM = DispatchPolicy::SCALAR_BUFFER_ELE_NUM; + static constexpr uint16_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(ElementC); + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(ElementC); + static constexpr uint32_t VNCHW_SIZE = 16; + + // Check LayoutC + static_assert(std::is_same_v, "LayoutC only support RowMajor yet!"); + + /// Construct + CATLASS_DEVICE + BlockMmadAiv(Arch::Resource &resource, MatrixCoord taskShape, uint32_t ubBufAddrStart = 0) + { + ubTileSize = RoundUp(taskShape.row()); + ubTileNum = RoundUp(taskShape.column()); + uint32_t ubBSize = SCALAR_BUFFER_ELE_NUM * sizeof(ElementB); + uint32_t ubASize = ubTileSize * sizeof(ElementA); + uint32_t ubCSize = ubTileNum * ubTileSize * sizeof(ElementA); + uint32_t ubTransposeCSize = ubTileNum * ubTileSize * sizeof(ElementA); + + uint32_t ubBOffset = ubBufAddrStart; + uint32_t ubAOffset = ubBufAddrStart + ubBSize; + uint32_t ubTransposeCOffset = ubAOffset + ubASize; + uint32_t ubCOffset = ubTransposeCOffset + ubTransposeCSize; + // Init buffers + ubBTensor = resource.ubBuf.template GetBufferByByte(ubBOffset); + ubATensor = resource.ubBuf.template GetBufferByByte(ubAOffset); + ubTransposeCTensor = resource.ubBuf.template GetBufferByByte(ubTransposeCOffset); + ubCTensor = resource.ubBuf.template GetBufferByByte(ubCOffset); + + AscendC::SetFlag(EVENT_ID0); + AscendC::SetFlag(EVENT_ID1); + AscendC::SetFlag(EVENT_ID2); + } + + /// Destructor + CATLASS_DEVICE + ~BlockMmadAiv() + { + AscendC::WaitFlag(EVENT_ID2); + AscendC::WaitFlag(EVENT_ID1); + AscendC::WaitFlag(EVENT_ID0); + } + + /// Perform a block-scoped matrix multiply-accumulate + CATLASS_DEVICE + void operator()( + AscendC::GlobalTensor const &gmBlockA, LayoutA const &layoutA, + AscendC::GlobalTensor const &gmBlockB, LayoutB const &layoutB, + AscendC::GlobalTensor const &gmBlockC, LayoutC const &layoutC, + GemmCoord const &actualShape) + { + AscendC::WaitFlag((event_t)(EVENT_ID0)); + AscendC::WaitFlag((event_t)(EVENT_ID1)); + + ElementA scalarB[SCALAR_BUFFER_ELE_NUM]; + LayoutA layoutBInUb{SCALAR_BUFFER_ELE_NUM}; + LayoutB layoutAInUb{ubTileSize}; + + // load A + auto layoutTileA = layoutA.GetTileLayout(TensorCoord(actualShape.m())); + copyGmToUbA(ubATensor, gmBlockA[0], layoutAInUb, layoutTileA); + AscendC::SetFlag((event_t)(EVENT_ID1)); + // load B + auto layoutTileB = layoutB.GetTileLayout(TensorCoord(actualShape.n())); + copyGmToUbB(ubBTensor, gmBlockB[0], layoutBInUb, layoutTileB); + AscendC::SetFlag((event_t)(EVENT_ID0)); + AscendC::SetFlag((event_t)(EVENT_ID0)); + AscendC::WaitFlag((event_t)(EVENT_ID0)); + for (uint32_t i = 0; i < actualShape.n(); i++) { + scalarB[i] = ubBTensor.GetValue(i); + } + AscendC::WaitFlag((event_t)(EVENT_ID0)); + AscendC::WaitFlag((event_t)(EVENT_ID1)); + AscendC::SetFlag((event_t)(EVENT_ID0)); + AscendC::WaitFlag((event_t)(EVENT_ID0)); + + // compute && write c to ub + for (uint32_t i = 0; i < actualShape.n(); i++) { + tileVmuls(ubCTensor[i * ubTileSize], ubATensor, scalarB[i], actualShape.m()); + } + AscendC::PipeBarrier(); + AscendC::SetFlag((event_t)(EVENT_ID0)); + AscendC::SetFlag((event_t)(EVENT_ID1)); + + LayoutC layoutCInUb{ubTileNum, ubTileSize}; + LayoutC layoutTranposeCInUb{ubTileSize, ubTileNum}; + auto actualLayoutC = layoutC.GetTileLayout(MakeCoord(actualShape.m(), actualShape.n())); + auto actualTransLayoutC = layoutTranposeCInUb.GetTileLayout(MakeCoord(actualShape.m(), actualShape.n())); + + // transpose 16*16 block + auto mTiles = CeilDiv(actualShape.n()); + auto usedBufferNum = CeilDiv(actualShape.m()); + AscendC::WaitFlag((event_t)(EVENT_ID2)); + uint64_t dstLocalList[VNCHW_SIZE]; + uint64_t srcLocalList[VNCHW_SIZE]; + for (uint32_t i = 0; i < mTiles; ++i) { + for (uint32_t j = 0; j < VNCHW_SIZE; ++j) { + srcLocalList[j] = ubCTensor.GetPhyAddr() + i * C0_NUM_PER_FRACTAL * ubTileSize * sizeof(ElementC) + j * ubTileSize * sizeof(ElementC); + if constexpr (std::is_same_v) { + dstLocalList[j] = ubTransposeCTensor.GetPhyAddr() + i * C0_NUM_PER_FRACTAL * sizeof(ElementC) + j * ubTileNum * sizeof(ElementC); + } + if constexpr (std::is_same_v) { + dstLocalList[j] = ubTransposeCTensor.GetPhyAddr() + i * C0_NUM_PER_FRACTAL * sizeof(ElementC) + j / 2 * ubTileNum * sizeof(ElementC) + j % 2 * BYTE_PER_C0; + } + } + if (usedBufferNum > 1) { + AscendC::TransDataTo5HDParams params(false, false, usedBufferNum, ubTileNum, 1); + AscendC::TransDataTo5HD(dstLocalList, srcLocalList, params); + } else { + AscendC::TransDataTo5HDParams paramsRepeatOne{false, false, 1, 0, 0}; + AscendC::TransDataTo5HD(dstLocalList, srcLocalList, paramsRepeatOne); + } + } + + AscendC::SetFlag((event_t)(EVENT_ID0)); + AscendC::WaitFlag((event_t)(EVENT_ID0)); + // write to gmC + copyUbToGmC(gmBlockC[0], ubTransposeCTensor, actualLayoutC, actualTransLayoutC); + AscendC::SetFlag((event_t)(EVENT_ID2)); + } + +protected: + /// Data members + AscendC::LocalTensor ubATensor; + AscendC::LocalTensor ubBTensor; + AscendC::LocalTensor ubCTensor; + AscendC::LocalTensor ubTransposeCTensor; + + CopyGmToUbA copyGmToUbA; + CopyGmToUbB copyGmToUbB; + CopyUbToGmC copyUbToGmC; + TileVmuls tileVmuls; + + uint32_t stageId{0}; + uint32_t ubTileNum{0}; + uint32_t ubTileSize{0}; +}; + +} // namespace Catlass::Gemm::Block +#endif // CATLASS_GEMM_BLOCK_MMAD_DYNAMIC_AIV_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_dynamic_common.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_dynamic_common.hpp new file mode 100644 index 0000000000..6d61d88158 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_dynamic_common.hpp @@ -0,0 +1,353 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_GEMM_BLOCK_BLOCK_MMAD_DYNAMIC_COMMON_HPP +#define CATLASS_GEMM_BLOCK_BLOCK_MMAD_DYNAMIC_COMMON_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/kda_coord.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" +#include "catlass/gemm/kda_gemm_helper.hpp" +#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" +#include "catlass/gemm/tile/kda_tile_mmad.hpp" + +namespace Catlass::Gemm::Block { + +template +struct BlockMmad, L1TileShape_, L0TileShape_, AType_, + BType_, CType_, BiasType_, TileCopy_, TileMmad_> { +public: + // Type Aliases + using DispatchPolicy = MmadAtlasA2DynamicCommon; + using ArchTag = typename DispatchPolicy::ArchTag; + using ElementA = typename AType_::Element; + using LayoutA = typename AType_::Layout; + using ElementB = typename BType_::Element; + using LayoutB = typename BType_::Layout; + using ElementC = typename CType_::Element; + using LayoutC = typename CType_::Layout; + using TileMmad = TileMmad_; + using CopyGmToL1A = typename TileCopy_::CopyGmToL1A; + using CopyGmToL1B = typename TileCopy_::CopyGmToL1B; + using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; + using CopyL0CToGm = typename TileCopy_::CopyL0CToGm; + using ElementAccumulator = + typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; + using LayoutAInL1 = typename CopyL1ToL0A::LayoutSrc; + using LayoutBInL1 = typename CopyL1ToL0B::LayoutSrc; + using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst; + using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst; + using LayoutCInL0 = layout::zN; + + using L1AAlignHelper = Gemm::helper::L1AlignHelper; + using L1BAlignHelper = Gemm::helper::L1AlignHelper; + + static constexpr bool ENABLE_UNIT_FLAG = DispatchPolicy::ENABLE_UNIT_FLAG; + static constexpr bool ENABLE_SHUFFLE_K = DispatchPolicy::ENABLE_SHUFFLE_K; + static constexpr uint32_t STAGES = DispatchPolicy::STAGES; + + static constexpr uint32_t L0A_PINGPONG_BUF_SIZE = ArchTag::L0A_SIZE / STAGES; + static constexpr uint32_t L0B_PINGPONG_BUF_SIZE = ArchTag::L0B_SIZE / STAGES; + + // Check LayoutC + static_assert(std::is_same_v, "LayoutC only support RowMajor yet!"); + + /// Construct + CATLASS_DEVICE + BlockMmad(GemmCoord const &l1TileShape_, Arch::Resource &resource, uint32_t l1BufAddrStart = 0) + : l1TileShape(l1TileShape_) + { + uint32_t l1ASize = l1TileShape.m() * l1TileShape.k() * sizeof(ElementA); + uint32_t l1BSize = l1TileShape.k() * l1TileShape.n() * sizeof(ElementB); + + kPartLenMax = min(L0A_PINGPONG_BUF_SIZE / sizeof(ElementA) / l1TileShape.m() / L1AAlignHelper::ELE_NUM_PER_C0 * + L1AAlignHelper::ELE_NUM_PER_C0, + L0B_PINGPONG_BUF_SIZE / sizeof(ElementB) / l1TileShape.n() / L1BAlignHelper::ELE_NUM_PER_C0 * + L1BAlignHelper::ELE_NUM_PER_C0); + + if constexpr (std::is_same_v && std::is_same_v) { + kPartLenMax = RoundDown(kPartLenMax); + } + + uint32_t l1AOffset = l1BufAddrStart; + uint32_t l1BOffset = l1BufAddrStart + l1ASize * STAGES; + // Init buffers + for (uint32_t i = 0; i < STAGES; i++) { + l1ATensorList[i] = resource.l1Buf.template GetBufferByByte(l1AOffset + l1ASize * i); + l1BTensorList[i] = resource.l1Buf.template GetBufferByByte(l1BOffset + l1BSize * i); + l0ATensorList[i] = resource.l0ABuf.template GetBufferByByte(L0A_PINGPONG_BUF_SIZE * i); + l0BTensorList[i] = resource.l0BBuf.template GetBufferByByte(L0B_PINGPONG_BUF_SIZE * i); + + l1AEventList[i] = i; + l1BEventList[i] = i + STAGES; + l0AEventList[i] = i; + l0BEventList[i] = i + STAGES; + AscendC::SetFlag(l1AEventList[i]); + AscendC::SetFlag(l1BEventList[i]); + AscendC::SetFlag(l0AEventList[i]); + AscendC::SetFlag(l0BEventList[i]); + } + l0CTensor = resource.l0CBuf.template GetBufferByByte(0); + AscendC::SetFlag(EVENT_ID0); + } + + /// Destructor + CATLASS_DEVICE + ~BlockMmad() + { + for (uint32_t i = 0; i < STAGES; i++) { + AscendC::WaitFlag(l1AEventList[i]); + AscendC::WaitFlag(l1BEventList[i]); + AscendC::WaitFlag(l0AEventList[i]); + AscendC::WaitFlag(l0BEventList[i]); + } + AscendC::WaitFlag(EVENT_ID0); + } + + /// Perform a block-scoped matrix multiply-accumulate + CATLASS_DEVICE + void operator()(AscendC::GlobalTensor const &gmBlockA, LayoutA const &layoutA, + AscendC::GlobalTensor const &gmBlockB, LayoutB const &layoutB, + AscendC::GlobalTensor const &gmBlockC, LayoutC const &layoutC, + AscendC::GlobalTensor const &gmNextBlockA, AscendC::GlobalTensor const &gmNextBlockB, + GemmCoord const &actualShape, GemmCoord const &actualShapeNext, bool isFirstBlock, bool hasNextBlock) + { + uint32_t mRound = RoundUp(actualShape.m()); + uint32_t nRound = RoundUp(actualShape.n()); + + auto layoutAInL1 = LayoutAInL1::template MakeLayout(l1TileShape.m(), l1TileShape.k()); + auto layoutBInL1 = LayoutBInL1::template MakeLayout(l1TileShape.k(), l1TileShape.n()); + auto layoutInL0C = LayoutCInL0::MakeLayoutInL0C(MakeCoord(mRound, nRound)); + + uint32_t kTileCount = CeilDiv(actualShape.k(), l1TileShape.k()); + + if constexpr (!ENABLE_UNIT_FLAG) { + AscendC::WaitFlag(EVENT_ID0); + } + uint32_t startTileIdx = 0; + if constexpr (ENABLE_SHUFFLE_K) { + startTileIdx = AscendC::GetBlockIdx(); + } + uint32_t firstTileIdx = startTileIdx % kTileCount; + uint32_t lastTileIdx = (startTileIdx + kTileCount - 1) % kTileCount; + uint32_t kActual = + (firstTileIdx < kTileCount - 1) ? l1TileShape.k() : (actualShape.k() - firstTileIdx * l1TileShape.k()); + + // k loop + for (uint32_t kLoopIdx = 0; kLoopIdx < kTileCount; kLoopIdx++) { + uint32_t shuffleKIdx = (startTileIdx + kLoopIdx) % kTileCount; + // Load first matrix A tile in total kernel loop from GM to L1 + if (shuffleKIdx == firstTileIdx && isFirstBlock) { + MatrixCoord gmTileACoord{0, shuffleKIdx * l1TileShape.k()}; + MatrixCoord gmTileBCoord{shuffleKIdx * l1TileShape.k(), 0}; + auto gmTileA = gmBlockA[layoutA.GetOffset(gmTileACoord)]; + auto gmTileB = gmBlockB[layoutB.GetOffset(gmTileBCoord)]; + // Load first matrix A tile from GM to L1 + AscendC::WaitFlag(l1AEventList[l1ListId]); + auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShape.m(), kActual)); + copyGmToL1A(l1ATensorList[l1ListId], gmTileA, layoutAInL1, layoutTileA); + AscendC::SetFlag(l1AEventList[l1ListId]); + + // Load first matrix B tile from GM to L1 + AscendC::WaitFlag(l1BEventList[l1ListId]); + auto layoutTileB = layoutB.GetTileLayout(MakeCoord(kActual, actualShape.n())); + copyGmToL1B(l1BTensorList[l1ListId], gmTileB, layoutBInL1, layoutTileB); + AscendC::SetFlag(l1BEventList[l1ListId]); + } + + uint32_t l1ListIdNext = (l1ListId + 1 < STAGES) ? (l1ListId + 1) : 0; + uint32_t kActualNext{0}; + + // preload next tile from GM to L1 + if (shuffleKIdx != lastTileIdx) { + uint32_t shuffleKIdxNext = (startTileIdx + kLoopIdx + 1) % kTileCount; + // Get L1 tensor for next stage + auto l1ATensor = l1ATensorList[l1ListIdNext]; + auto l1BTensor = l1BTensorList[l1ListIdNext]; + // Get GM tensor for next stage + kActualNext = (shuffleKIdxNext < kTileCount - 1) ? l1TileShape.k() : (actualShape.k() - shuffleKIdxNext * l1TileShape.k()); + MatrixCoord gmTileACoord{0, shuffleKIdxNext * l1TileShape.k()}; + MatrixCoord gmTileBCoord{shuffleKIdxNext * l1TileShape.k(), 0}; + auto gmTileA = gmBlockA[layoutA.GetOffset(gmTileACoord)]; + auto gmTileB = gmBlockB[layoutB.GetOffset(gmTileBCoord)]; + + // load next matrix A tile from GM to L1 + AscendC::WaitFlag(l1AEventList[l1ListIdNext]); + auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShape.m(), kActualNext)); + copyGmToL1A(l1ATensor, gmTileA, layoutAInL1, layoutTileA); + AscendC::SetFlag(l1AEventList[l1ListIdNext]); + + // load next matrix B tile from GM to L1 + AscendC::WaitFlag(l1BEventList[l1ListIdNext]); + auto layoutTileB = layoutB.GetTileLayout(MakeCoord(kActualNext, actualShape.n())); + copyGmToL1B(l1BTensor, gmTileB, layoutBInL1, layoutTileB); + AscendC::SetFlag(l1BEventList[l1ListIdNext]); + } + if (shuffleKIdx == lastTileIdx && hasNextBlock) { + uint32_t kTileCount = CeilDiv(actualShapeNext.k(), l1TileShape.k()); + uint32_t firstTileIdx = startTileIdx % kTileCount; + // Get L1 tensor for next stage + auto l1ATensor = l1ATensorList[l1ListIdNext]; + auto l1BTensor = l1BTensorList[l1ListIdNext]; + // Get GM tensor for next stage + kActualNext = (firstTileIdx < kTileCount - 1) ? l1TileShape.k() : (actualShapeNext.k() - firstTileIdx * l1TileShape.k()); + MatrixCoord gmTileACoord{0, firstTileIdx * l1TileShape.k()}; + MatrixCoord gmTileBCoord{firstTileIdx * l1TileShape.k(), 0}; + auto gmTileA = gmNextBlockA[layoutA.GetOffset(gmTileACoord)]; + auto gmTileB = gmNextBlockB[layoutB.GetOffset(gmTileBCoord)]; + // load next matrix A tile from GM to L1 + AscendC::WaitFlag(l1AEventList[l1ListIdNext]); + auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShapeNext.m(), kActualNext)); + copyGmToL1A(l1ATensor, gmTileA, layoutAInL1, layoutTileA); + AscendC::SetFlag(l1AEventList[l1ListIdNext]); + + // load next matrix B tile from GM to L1 + AscendC::WaitFlag(l1BEventList[l1ListIdNext]); + auto layoutTileB = layoutB.GetTileLayout(MakeCoord(kActualNext, actualShapeNext.n())); + copyGmToL1B(l1BTensor, gmTileB, layoutBInL1, layoutTileB); + AscendC::SetFlag(l1BEventList[l1ListIdNext]); + } + + // Get L1 tensor for current stage + auto l1ATensor = l1ATensorList[l1ListId]; + auto l1BTensor = l1BTensorList[l1ListId]; + + // Get the loop nums on L0 + uint32_t kPartLoop = CeilDiv(kActual, kPartLenMax); + + uint32_t l0ABufId = 0; + uint32_t l0BBufId = 0; + + for (int kPartIdx = 0; kPartIdx < kPartLoop; kPartIdx++) { + uint32_t kPartActual = + (kPartIdx < kPartLoop - 1) ? kPartLenMax : (kActual - kPartIdx * kPartLenMax); + + // Locate the current tile on L0A + auto l0ATile = l0ATensorList[l0ABufId]; + LayoutAInL0 layoutAInL0 = LayoutAInL0::template MakeLayout(mRound, kPartActual); + // Locate the current tile of matrix A on L1 + MatrixCoord l1ACoord{0, kPartIdx * kPartLenMax}; + auto l1ATile = l1ATensor[layoutAInL1.GetOffset(l1ACoord)]; + + AscendC::WaitFlag(l0AEventList[l0ABufId]); + if (kPartIdx == 0) { + AscendC::WaitFlag(l1AEventList[l1ListId]); + } + + copyL1ToL0A(l0ATile, l1ATile, layoutAInL0, layoutAInL1); + + if (kPartIdx == kPartLoop - 1) { + AscendC::SetFlag(l1AEventList[l1ListId]); + } + + // Locate the current tile on L0B + auto l0BTile = l0BTensorList[l0BBufId]; + LayoutBInL0 layoutBInL0 = LayoutBInL0::template MakeLayout(kPartActual, nRound); + // Locate the current tile of matrix B on L1 + MatrixCoord l1BCoord{kPartIdx * kPartLenMax, 0}; + auto l1BTile = l1BTensor[layoutBInL1.GetOffset(l1BCoord)]; + + // Wait for mmad finished + AscendC::WaitFlag(l0BEventList[l0BBufId]); + // If the current tile is the first one on the k&n axis, wait for loading matrix B from GM to L1 + if (kPartIdx == 0) { + AscendC::WaitFlag(l1BEventList[l1ListId]); + } + + // Load current tile from L1 to L0B + copyL1ToL0B(l0BTile, l1BTile, layoutBInL0, layoutBInL1); + + // If the current tile is the last one on the k&n axis, notify to load matrix B from GM to L1 + if (kPartIdx == kPartLoop - 1) { + AscendC::SetFlag(l1BEventList[l1ListId]); + } + // Notify to do mmad + AscendC::SetFlag(EVENT_ID0); + + // Locate the current tile on L0C + MatrixCoord l0CCoord{static_cast(0U), static_cast(0U)}; + auto l0CTile = l0CTensor[layoutInL0C.GetOffset(l0CCoord)]; + + // Compute the matrix multiplication on L0A and L0B and write the result to the accumulator + // Wait for loading L0B + AscendC::WaitFlag(EVENT_ID0); + + // If the current tile is the first tile on the k axis, the accumulator needs to be reset to 0 + bool initC = ((kLoopIdx == 0) && (kPartIdx == 0)); + // If the unit flag is enabled, the unit flag is set according to the calculation progress + uint8_t unitFlag = 0b00; + if constexpr (ENABLE_UNIT_FLAG) { + if ((kLoopIdx == kTileCount - 1) && (kPartIdx == kPartLoop - 1)) { + unitFlag = 0b11; + } else { + unitFlag = 0b10; + } + } + // Perform calculation operations + tileMmad(l0CTile, l0ATile, l0BTile, mRound, nRound, kPartActual, initC, unitFlag); + + // Notify to move the next L0B tile + AscendC::SetFlag(l0BEventList[l0BBufId]); + + l0BBufId = (l0BBufId + 1 < STAGES) ? (l0BBufId + 1) : 0; + + AscendC::SetFlag(l0AEventList[l0ABufId]); + l0ABufId = (l0ABufId + 1 < STAGES) ? (l0ABufId + 1) : 0; + } + + l1ListId = l1ListIdNext; + kActual = kActualNext; + } + + // copy block out + LayoutC layoutBlock = layoutC.GetTileLayout(actualShape.GetCoordMN()); + + if constexpr (!ENABLE_UNIT_FLAG) { + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + copyL0CToGm(gmBlockC, l0CTensor, layoutBlock, layoutInL0C); + AscendC::SetFlag(EVENT_ID0); + } else { + copyL0CToGm(gmBlockC, l0CTensor, layoutBlock, layoutInL0C, 0b11); + } + } + +protected: + /// Data members + AscendC::LocalTensor l1ATensorList[STAGES]; + AscendC::LocalTensor l1BTensorList[STAGES]; + AscendC::LocalTensor l0ATensorList[STAGES]; + AscendC::LocalTensor l0BTensorList[STAGES]; + AscendC::LocalTensor l0CTensor; + + int32_t l1AEventList[STAGES]; + int32_t l1BEventList[STAGES]; + int32_t l0AEventList[STAGES]; + int32_t l0BEventList[STAGES]; + + uint32_t l1ListId{0}; + GemmCoord l1TileShape; + uint32_t kPartLenMax; + + TileMmad tileMmad; + CopyGmToL1A copyGmToL1A; + CopyGmToL1B copyGmToL1B; + CopyL1ToL0A copyL1ToL0A; + CopyL1ToL0B copyL1ToL0B; + CopyL0CToGm copyL0CToGm; +}; + +} // namespace Catlass::Gemm::Block + +#endif // CATLASS_GEMM_BLOCK_BLOCK_MMAD_DYNAMIC_COMMON_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_dynamic_preload_async_with_callback.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_dynamic_preload_async_with_callback.hpp new file mode 100644 index 0000000000..c0d7093aa5 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_dynamic_preload_async_with_callback.hpp @@ -0,0 +1,441 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_GEMM_BLOCK_BLOCK_MMAD_DYNAMIC_PRELOAD_ASYNC_WITH_CALLBACK_HPP +#define CATLASS_GEMM_BLOCK_BLOCK_MMAD_DYNAMIC_PRELOAD_ASYNC_WITH_CALLBACK_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/kda_coord.hpp" +#include "catlass/detail/kda_callback.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" +#include "catlass/gemm/kda_gemm_helper.hpp" +#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" +#include "catlass/gemm/tile/kda_tile_mmad.hpp" + +namespace Catlass::Gemm::Block { + +template < + uint32_t PRELOAD_STAGES_, + uint32_t L1_STAGES_, + uint32_t L0A_STAGES_, + uint32_t L0B_STAGES_, + uint32_t L0C_STAGES_, + bool ENABLE_UNIT_FLAG_, + bool ENABLE_SHUFFLE_K_, + class L1TileShape_, + class L0TileShape_, + class AType_, + class BType_, + class CType_, + class BiasType_, + class TileCopy_, + class TileMmad_> +struct BlockMmad< + MmadAtlasA2DynamicPreloadAsyncWithCallback< + PRELOAD_STAGES_, + L1_STAGES_, + L0A_STAGES_, + L0B_STAGES_, + L0C_STAGES_, + ENABLE_UNIT_FLAG_, + ENABLE_SHUFFLE_K_>, + L1TileShape_, + L0TileShape_, + AType_, + BType_, + CType_, + BiasType_, + TileCopy_, + TileMmad_> { +public: + // Type Aliases + using DispatchPolicy = MmadAtlasA2DynamicPreloadAsyncWithCallback< + PRELOAD_STAGES_, + L1_STAGES_, + L0A_STAGES_, + L0B_STAGES_, + L0C_STAGES_, + ENABLE_UNIT_FLAG_, + ENABLE_SHUFFLE_K_>; + using ArchTag = typename DispatchPolicy::ArchTag; + using ElementA = typename AType_::Element; + using LayoutA = typename AType_::Layout; + using ElementB = typename BType_::Element; + using LayoutB = typename BType_::Layout; + using ElementC = typename CType_::Element; + using LayoutC = typename CType_::Layout; + using TileMmad = TileMmad_; + using CopyGmToL1A = typename TileCopy_::CopyGmToL1A; + using CopyGmToL1B = typename TileCopy_::CopyGmToL1B; + using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; + using CopyL0CToGm = typename TileCopy_::CopyL0CToGm; + using ElementAccumulator = + typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; + using LayoutAInL1 = typename CopyL1ToL0A::LayoutSrc; + using LayoutBInL1 = typename CopyL1ToL0B::LayoutSrc; + using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst; + using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst; + using LayoutCInL0 = layout::zN; + + using L1AAlignHelper = Gemm::helper::L1AlignHelper; + using L1BAlignHelper = Gemm::helper::L1AlignHelper; + + static constexpr uint32_t PRELOAD_STAGES = DispatchPolicy::PRELOAD_STAGES; + static constexpr uint32_t L1_STAGES = DispatchPolicy::L1_STAGES; + static constexpr uint32_t L0A_STAGES = DispatchPolicy::L0A_STAGES; + static constexpr uint32_t L0B_STAGES = DispatchPolicy::L0B_STAGES; + static constexpr uint32_t L0C_STAGES = DispatchPolicy::L0C_STAGES; + + static constexpr bool ENABLE_UNIT_FLAG = DispatchPolicy::ENABLE_UNIT_FLAG; + static constexpr bool ENABLE_SHUFFLE_K = DispatchPolicy::ENABLE_SHUFFLE_K; + + // Check LayoutC + static_assert(std::is_same_v, "LayoutC only support RowMajor yet!"); + + CATLASS_DEVICE + BlockMmad(GemmCoord const &l1TileShape_, GemmCoord const &l0TileShape_, Arch::Resource &resource, uint32_t l1BufAddrStart = 0) + : l1TileShape(l1TileShape_), + l0TileShape(l0TileShape_) + { + if ASCEND_IS_AIV { + return; + } + + uint32_t l1ASize = l1TileShape.m() * l1TileShape.k() * sizeof(ElementA); + uint32_t l1BSize = l1TileShape.k() * l1TileShape.n() * sizeof(ElementB); + + uint32_t l0ASize = l0TileShape.m() * l0TileShape.k() * sizeof(ElementA); + uint32_t l0BSize = l0TileShape.k() * l0TileShape.n() * sizeof(ElementB); + uint32_t l0CSize = l0TileShape.m() * l0TileShape.n() * sizeof(ElementAccumulator); + + uint32_t l1AOffset = l1BufAddrStart; + uint32_t l1BOffset = l1BufAddrStart + l1ASize * L1_STAGES; + for (uint32_t i = 0; i < L1_STAGES; ++i) { + l1ATensorList[i] = resource.l1Buf.template GetBufferByByte(l1AOffset + l1ASize * i); + l1BTensorList[i] = resource.l1Buf.template GetBufferByByte(l1BOffset + l1BSize * i); + l1AEventList[i] = i; + l1BEventList[i] = i + L1_STAGES; + AscendC::SetFlag(l1AEventList[i]); + AscendC::SetFlag(l1BEventList[i]); + } + + for (uint32_t i = 0; i < L0A_STAGES; ++i) { + l0ATensorList[i] = resource.l0ABuf.template GetBufferByByte(l0ASize * i); + l0AEventList[i] = i; + AscendC::SetFlag(l0AEventList[i]); + } + + for (uint32_t i = 0; i < L0B_STAGES; ++i) { + l0BTensorList[i] = resource.l0BBuf.template GetBufferByByte(l0BSize * i); + l0BEventList[i] = i + L0A_STAGES; + AscendC::SetFlag(l0BEventList[i]); + } + + for (uint32_t i = 0; i < L0C_STAGES; ++i) { + l0CTensorList[i] = resource.l0CBuf.template GetBufferByByte(l0CSize * i); + l0CEventList[i] = i; + AscendC::SetFlag(l0CEventList[i]); + } + } + + CATLASS_DEVICE + ~BlockMmad() + { + if ASCEND_IS_AIV { + return; + } + + SynchronizeBlock(); + for (uint32_t i = 0; i < L1_STAGES; ++i) { + AscendC::WaitFlag(l1AEventList[i]); + AscendC::WaitFlag(l1BEventList[i]); + } + for (uint32_t i = 0; i < L0A_STAGES; ++i) { + AscendC::WaitFlag(l0AEventList[i]); + } + for (uint32_t i = 0; i < L0B_STAGES; ++i) { + AscendC::WaitFlag(l0BEventList[i]); + } + for (uint32_t i = 0; i < L0C_STAGES; ++i) { + AscendC::WaitFlag(l0CEventList[i]); + } + } + + CATLASS_DEVICE + void operator()( + AscendC::GlobalTensor const &gmBlockA, LayoutA const &layoutA, + AscendC::GlobalTensor const &gmBlockB, LayoutB const &layoutB, + AscendC::GlobalTensor const &gmBlockC, LayoutC const &layoutC, + GemmCoord const &actualShape, + Callback const &callbackBeforeFixpipe, Callback const &callbackAfterFixpipe) + { + uint32_t kTileCount = CeilDiv(actualShape.k(), l1TileShape.k()); + + uint32_t mRound = RoundUp(actualShape.m()); + uint32_t nRound = RoundUp(actualShape.n()); + + uint32_t startTileIdx = 0; + if constexpr (ENABLE_SHUFFLE_K) { + startTileIdx = AscendC::GetBlockIdx() % kTileCount; + } + + for (uint32_t kLoopIdx = 0; kLoopIdx < kTileCount; ++kLoopIdx) { + uint32_t kTileIdx = (startTileIdx + kLoopIdx < kTileCount) ? + (startTileIdx + kLoopIdx) : + (startTileIdx + kLoopIdx - kTileCount); + + uint32_t kActual = (kTileIdx < kTileCount - 1) ? + l1TileShape.k() : + (actualShape.k() - kTileIdx * l1TileShape.k()); + + // Emission load instruction from GM to L1 + MatrixCoord gmTileAOffset{0, kTileIdx * l1TileShape.k()}; + MatrixCoord gmTileBOffset{kTileIdx * l1TileShape.k(), 0}; + auto gmTileA = gmBlockA[layoutA.GetOffset(gmTileAOffset)]; + auto gmTileB = gmBlockB[layoutB.GetOffset(gmTileBOffset)]; + // Load first matrix A tile from GM to L1 + AscendC::WaitFlag(l1AEventList[l1ListId]); + auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShape.m(), kActual)); + auto l1a_layout = LayoutAInL1::template MakeLayout(l1TileShape.m(), l1TileShape.k()); + copyGmToL1A(l1ATensorList[l1ListId], gmTileA, l1a_layout, layoutTileA); + AscendC::SetFlag(l1AEventList[l1ListId]); + // Load first matrix B tile from GM to L1 + AscendC::WaitFlag(l1BEventList[l1ListId]); + auto layoutTileB = layoutB.GetTileLayout(MakeCoord(kActual, actualShape.n())); + auto l1b_layout = LayoutBInL1::template MakeLayout(l1TileShape.k(), l1TileShape.n()); + copyGmToL1B(l1BTensorList[l1ListId], gmTileB, l1b_layout, layoutTileB); + AscendC::SetFlag(l1BEventList[l1ListId]); + + // If the number of preload instructions reaches the upper limit, perform an mmad calculation on L1 tile + if (preloadCount == PRELOAD_STAGES) { + L1TileMmad(l1TileMmadParamsList[l1TileMmadParamsId]); + } + + // Store the current load status + uint32_t preloadL1TileMmadParamsId = (l1TileMmadParamsId + preloadCount < PRELOAD_STAGES) ? + (l1TileMmadParamsId + preloadCount) : + (l1TileMmadParamsId + preloadCount - PRELOAD_STAGES); + auto &l1TileMmadParams = l1TileMmadParamsList[preloadL1TileMmadParamsId]; + l1TileMmadParams.l1ListId = l1ListId; + l1TileMmadParams.mRound = mRound; + l1TileMmadParams.nRound = nRound; + l1TileMmadParams.kActual = kActual; + l1TileMmadParams.isKLoopFirst = (kLoopIdx == 0); + l1TileMmadParams.isKLoopLast = (kLoopIdx == kTileCount - 1); + if (kLoopIdx == kTileCount - 1) { + l1TileMmadParams.gmBlockC = gmBlockC; + l1TileMmadParams.layoutCInGm = layoutC.GetTileLayout(actualShape.GetCoordMN()); + l1TileMmadParams.callbackBeforeFixpipe = callbackBeforeFixpipe; + l1TileMmadParams.callbackAfterFixpipe = callbackAfterFixpipe; + } + + if (preloadCount < PRELOAD_STAGES) { + ++preloadCount; + } else { + l1TileMmadParamsId = (l1TileMmadParamsId + 1 < PRELOAD_STAGES) ? (l1TileMmadParamsId + 1) : 0; + } + l1ListId = (l1ListId + 1 < L1_STAGES) ? (l1ListId + 1) : 0; + } + } + + CATLASS_DEVICE + void SynchronizeBlock() + { + while (preloadCount > 0) { + L1TileMmad(l1TileMmadParamsList[l1TileMmadParamsId]); + l1TileMmadParamsId = (l1TileMmadParamsId + 1 < PRELOAD_STAGES) ? (l1TileMmadParamsId + 1) : 0; + --preloadCount; + } + } + +private: + struct L1TileMmadParams { + uint32_t l1ListId; + uint32_t mRound; + uint32_t nRound; + uint32_t kActual; + bool isKLoopFirst; + bool isKLoopLast; + AscendC::GlobalTensor gmBlockC; + LayoutC layoutCInGm; + Callback callbackBeforeFixpipe; + Callback callbackAfterFixpipe; + + CATLASS_DEVICE + L1TileMmadParams() = default; + }; + + CATLASS_DEVICE + void InitL1(Arch::Resource &resource, uint32_t l1BufAddrStart) + { + uint32_t l1ASize = l1TileShape.m() * l1TileShape.k() * sizeof(ElementA); + uint32_t l1BSize = l1TileShape.k() * l1TileShape.n() * sizeof(ElementB); + uint32_t l1AOffset = l1BufAddrStart; + uint32_t l1BOffset = l1BufAddrStart + l1ASize * L1_STAGES; + for (uint32_t i = 0; i < L1_STAGES; ++i) { + l1ATensorList[i] = resource.l1Buf.template GetBufferByByte(l1AOffset + l1ASize * i); + l1BTensorList[i] = resource.l1Buf.template GetBufferByByte(l1BOffset + l1BSize * i); + l1AEventList[i] = i; + l1BEventList[i] = i + L1_STAGES; + AscendC::SetFlag(l1AEventList[i]); + AscendC::SetFlag(l1BEventList[i]); + } + } + + CATLASS_DEVICE + void L1TileMmad(L1TileMmadParams const ¶ms) + { + uint32_t mPartLoop = CeilDiv(params.mRound, l0TileShape.m()); + uint32_t nPartLoop = CeilDiv(params.nRound, l0TileShape.n()); + uint32_t kPartLoop = CeilDiv(params.kActual, l0TileShape.k()); + + auto &l1ATensor = l1ATensorList[params.l1ListId]; + auto &l1BTensor = l1BTensorList[params.l1ListId]; + + auto &l0CTensor = l0CTensorList[l0CListId]; + LayoutCInL0 layoutCInL0 = LayoutCInL0::MakeLayoutInL0C(MakeCoord(params.mRound, params.nRound)); + + if constexpr (!ENABLE_UNIT_FLAG) { + if (params.isKLoopFirst) { + AscendC::WaitFlag(l0CEventList[l0CListId]); + } + } + + for (uint32_t mPartIdx = 0; mPartIdx < mPartLoop; ++mPartIdx) { + uint32_t mPartActual = (mPartIdx < mPartLoop - 1) ? + l0TileShape.m() : + (params.mRound - mPartIdx * l0TileShape.m()); + + for (uint32_t kPartIdx = 0; kPartIdx < kPartLoop; ++kPartIdx) { + uint32_t kPartActual = (kPartIdx < kPartLoop - 1) ? + l0TileShape.k() : + (params.kActual - kPartIdx * l0TileShape.k()); + + auto &l0ATile = l0ATensorList[l0AListId]; + auto layoutAInL0 = LayoutAInL0::template MakeLayout(mPartActual, kPartActual); + auto l1AOffset = MakeCoord(mPartIdx, kPartIdx) * MakeCoord(l0TileShape.m(), l0TileShape.k()); + auto l1a_layout = LayoutAInL1::template MakeLayout(l1TileShape.m(), l1TileShape.k()); + auto l1ATile = l1ATensor[l1a_layout.GetOffset(l1AOffset)]; + + AscendC::WaitFlag(l0AEventList[l0AListId]); + if ((mPartIdx == 0) && (kPartIdx == 0)) { + AscendC::WaitFlag(l1AEventList[params.l1ListId]); + } + copyL1ToL0A(l0ATile, l1ATile, layoutAInL0, l1a_layout); + if ((mPartIdx == mPartLoop - 1) && (kPartIdx == kPartLoop - 1)) { + AscendC::SetFlag(l1AEventList[params.l1ListId]); + } + + for (uint32_t nPartIdx = 0; nPartIdx < nPartLoop; ++nPartIdx) { + uint32_t nPartActual = (nPartIdx < nPartLoop - 1) ? + l0TileShape.n() : + (params.nRound - nPartIdx * l0TileShape.n()); + + auto &l0BTile = l0BTensorList[l0BListId]; + auto layoutBInL0 = LayoutBInL0::template MakeLayout(kPartActual, nPartActual); + auto l1BOffset = MakeCoord(kPartIdx, nPartIdx) * MakeCoord(l0TileShape.k(), l0TileShape.n()); + auto l1b_layout = LayoutBInL1::template MakeLayout(l1TileShape.k(), l1TileShape.n()); + auto l1BTile = l1BTensor[l1b_layout.GetOffset(l1BOffset)]; + + AscendC::WaitFlag(l0BEventList[l0BListId]); + if ((kPartIdx == 0) && (nPartIdx == 0)) { + AscendC::WaitFlag(l1BEventList[params.l1ListId]); + } + copyL1ToL0B(l0BTile, l1BTile, layoutBInL0, l1b_layout); + if ((kPartIdx == kPartLoop - 1) && (nPartIdx == nPartLoop - 1)) { + AscendC::SetFlag(l1BEventList[params.l1ListId]); + } + + AscendC::SetFlag(EVENT_ID0); + + auto l0COffset = MakeCoord(mPartIdx, nPartIdx) * MakeCoord(l0TileShape.m(), l0TileShape.n()); + auto l0CTile = l0CTensor[layoutCInL0.GetOffset(l0COffset)]; + + AscendC::WaitFlag(EVENT_ID0); + // If the current tile is the first tile on the k axis, the accumulator needs to be reset to 0 + bool initC = (params.isKLoopFirst && (kPartIdx == 0)); + // If the unit flag is enabled, the unit flag is set according to the calculation progress + uint8_t unitFlag = 0b00; + if constexpr (ENABLE_UNIT_FLAG) { + if (params.isKLoopLast && + (mPartIdx == mPartLoop - 1) && (kPartIdx == kPartLoop - 1) && (nPartIdx == nPartLoop - 1)) { + unitFlag = 0b11; + } else { + unitFlag = 0b10; + } + } + tileMmad(l0CTile, l0ATile, l0BTile, mPartActual, nPartActual, kPartActual, initC, unitFlag); + + AscendC::SetFlag(l0BEventList[l0BListId]); + l0BListId = (l0BListId + 1 < L0B_STAGES) ? (l0BListId + 1) : 0; + } + AscendC::SetFlag(l0AEventList[l0AListId]); + l0AListId = (l0AListId + 1 < L0A_STAGES) ? (l0AListId + 1) : 0; + } + } + + if (params.isKLoopLast) { + auto layoutCInGm = params.layoutCInGm; + + params.callbackBeforeFixpipe(); + + if constexpr (!ENABLE_UNIT_FLAG) { + AscendC::SetFlag(l0CEventList[l0CListId]); + AscendC::WaitFlag(l0CEventList[l0CListId]); + copyL0CToGm(params.gmBlockC, l0CTensor, layoutCInGm, layoutCInL0); + AscendC::SetFlag(l0CEventList[l0CListId]); + } else { + copyL0CToGm(params.gmBlockC, l0CTensor, layoutCInGm, layoutCInL0, 0b11); + } + l0CListId = (l0CListId + 1 < L0C_STAGES) ? (l0CListId + 1) : 0; + + params.callbackAfterFixpipe(); + } + } + +protected: + AscendC::LocalTensor l1ATensorList[L1_STAGES]; + AscendC::LocalTensor l1BTensorList[L1_STAGES]; + int32_t l1AEventList[L1_STAGES]; + int32_t l1BEventList[L1_STAGES]; + uint32_t l1ListId{0}; + + AscendC::LocalTensor l0ATensorList[L0A_STAGES]; + int32_t l0AEventList[L0A_STAGES]; + uint32_t l0AListId{0}; + + AscendC::LocalTensor l0BTensorList[L0B_STAGES]; + int32_t l0BEventList[L0B_STAGES]; + uint32_t l0BListId{0}; + + AscendC::LocalTensor l0CTensorList[L0C_STAGES_]; + int32_t l0CEventList[L0C_STAGES_]; + uint32_t l0CListId{0}; + + L1TileMmadParams l1TileMmadParamsList[PRELOAD_STAGES]; + uint32_t l1TileMmadParamsId{0}; + uint32_t preloadCount{0}; + GemmCoord l1TileShape; + GemmCoord l0TileShape; + + TileMmad tileMmad; + CopyGmToL1A copyGmToL1A; + CopyGmToL1B copyGmToL1B; + CopyL1ToL0A copyL1ToL0A; + CopyL1ToL0B copyL1ToL0B; + CopyL0CToGm copyL0CToGm; +}; + +} // namespace Catlass::Gemm::Block + +#endif // CATLASS_GEMM_BLOCK_BLOCK_MMAD_DYNAMIC_PRELOAD_ASYNC_WITH_CALLBACK_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_dynamic_single_core_splitk.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_dynamic_single_core_splitk.hpp new file mode 100644 index 0000000000..bacb903288 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_dynamic_single_core_splitk.hpp @@ -0,0 +1,391 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_GEMM_BLOCK_MMAD_DYNAMIC_SINGLE_CORE_SPLITK_HPP +#define CATLASS_GEMM_BLOCK_MMAD_DYNAMIC_SINGLE_CORE_SPLITK_HPP + +#include "catlass/arch/kda_resource.hpp" +#include "catlass/kda_catlass.hpp" +#include "catlass/kda_coord.hpp" +#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" +#include "catlass/gemm/kda_gemm_helper.hpp" +#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" +#include "catlass/gemm/tile/kda_tile_mmad.hpp" +#include "catlass/kda_gemm_coord.hpp" + +namespace Catlass::Gemm::Block { + +template < + uint32_t L1A_STAGES_, + uint32_t L1B_STAGES_, + uint32_t L0C_STAGES_, + bool ENABLE_UNIT_FLAG_, + class L1TileShape_, + class L0TileShape_, + class AType_, + class BType_, + class CType_, + class BiasType_, + class TileCopy_, + class TileMmad_> +struct BlockMmad< + MmadAtlasA2DynamicSingleCoreSplitk, + L1TileShape_, + L0TileShape_, + AType_, + BType_, + CType_, + BiasType_, + TileCopy_, + TileMmad_> { +public: + // Type Aliases + using DispatchPolicy = MmadAtlasA2DynamicSingleCoreSplitk; + using ArchTag = typename DispatchPolicy::ArchTag; + using ElementA = typename AType_::Element; + using LayoutA = typename AType_::Layout; + using ElementB = typename BType_::Element; + using LayoutB = typename BType_::Layout; + using ElementC = typename CType_::Element; + using LayoutC = typename CType_::Layout; + using TileMmad = TileMmad_; + using CopyGmToL1A = typename TileCopy_::CopyGmToL1A; + using CopyGmToL1B = typename TileCopy_::CopyGmToL1B; + using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; + using CopyL0CToGm = typename TileCopy_::CopyL0CToGm; + using ElementAccumulator = typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; + using LayoutAInL1 = typename CopyL1ToL0A::LayoutSrc; + using LayoutBInL1 = typename CopyL1ToL0B::LayoutSrc; + using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst; + using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst; + using LayoutCInL0 = layout::zN; + + using L1AAlignHelper = Gemm::helper::L1AlignHelper; + using L1BAlignHelper = Gemm::helper::L1AlignHelper; + + static constexpr uint32_t L1A_STAGES = DispatchPolicy::L1A_STAGES; + static constexpr uint32_t L1B_STAGES = DispatchPolicy::L1B_STAGES; + static constexpr uint32_t L0AB_STAGES = DispatchPolicy::L0AB_STAGES; + static constexpr uint32_t L0C_STAGES = DispatchPolicy::L0C_STAGES; + static constexpr bool ENABLE_UNIT_FLAG = DispatchPolicy::ENABLE_UNIT_FLAG; + + static constexpr uint32_t L0A_PINGPONG_BUF_SIZE = ArchTag::L0A_SIZE / L0AB_STAGES; + static constexpr uint32_t L0B_PINGPONG_BUF_SIZE = ArchTag::L0B_SIZE / L0AB_STAGES; + static constexpr uint32_t L0C_PINGPONG_BUF_SIZE = ArchTag::L0C_SIZE / L0C_STAGES; + + // Check LayoutC + static_assert(std::is_same_v, "LayoutC only support RowMajor yet!"); + + /// Construct + CATLASS_DEVICE + BlockMmad( + GemmCoord const &l1TileShape_, + GemmCoord const &l0TileShape_, + Arch::Resource &resource, + uint32_t l1BufAddrStart = 0) + : l1TileShape(l1TileShape_), + l0TileShape(l0TileShape_) + { + if ASCEND_IS_AIV { + return; + } + + uint32_t l1ASize = l1TileShape.m() * l1TileShape.k() * sizeof(ElementA); + uint32_t l1BSize = l1TileShape.k() * l1TileShape.n() * sizeof(ElementB); + + uint32_t l1AOffset = l1BufAddrStart; + uint32_t l1BOffset = l1BufAddrStart + l1ASize * L1A_STAGES; + // Init buffers + for (uint32_t i = 0; i < L1A_STAGES; i++) { + l1ATensorList[i] = resource.l1Buf.template GetBufferByByte(l1AOffset + l1ASize * i); + l1AEventList[i] = i; + AscendC::SetFlag(l1AEventList[i]); + } + for (uint32_t i = 0; i < L1B_STAGES; i++) { + l1BTensorList[i] = resource.l1Buf.template GetBufferByByte(l1BOffset + l1BSize * i); + l1BEventList[i] = i + L1A_STAGES; + AscendC::SetFlag(l1BEventList[i]); + } + + for (uint32_t i = 0; i < L0AB_STAGES; i++) { + l0ATensorList[i] = resource.l0ABuf.template GetBufferByByte(L0A_PINGPONG_BUF_SIZE * i); + l0BTensorList[i] = resource.l0BBuf.template GetBufferByByte(L0B_PINGPONG_BUF_SIZE * i); + l0AEventList[i] = i; + l0BEventList[i] = i + L0AB_STAGES; + AscendC::SetFlag(l0AEventList[i]); + AscendC::SetFlag(l0BEventList[i]); + } + for (uint32_t i = 0; i < L0C_STAGES; i++) { + l0CTensorList[i] = resource.l0CBuf.template GetBufferByByte(L0C_PINGPONG_BUF_SIZE * i); + l0CEventList[i] = i; + AscendC::SetFlag(l0CEventList[i]); + } + } + + /// Destructor + CATLASS_DEVICE + ~BlockMmad() + { + if ASCEND_IS_AIV { + return; + } + + for (uint32_t i = 0; i < L1A_STAGES; i++) { + AscendC::WaitFlag(l1AEventList[i]); + } + for (uint32_t i = 0; i < L1B_STAGES; i++) { + AscendC::WaitFlag(l1BEventList[i]); + } + + for (uint32_t i = 0; i < L0AB_STAGES; i++) { + AscendC::WaitFlag(l0AEventList[i]); + AscendC::WaitFlag(l0BEventList[i]); + } + for (uint32_t i = 0; i < L0C_STAGES; i++) { + AscendC::WaitFlag(l0CEventList[i]); + } + } + + /// Perform a block-scoped matrix multiply-accumulate + CATLASS_DEVICE + void operator()( + AscendC::GlobalTensor const &gmBlockA, + LayoutA const &layoutA, + AscendC::GlobalTensor const &gmBlockB, + LayoutB const &layoutB, + AscendC::GlobalTensor const &gmBlockC, + LayoutC const &layoutC, + AscendC::GlobalTensor const &gmNextBlockA, + AscendC::GlobalTensor const &gmNextBlockB, + GemmCoord const &actualShape, + GemmCoord const &actualShapeNext, + bool needLoadNextA, + bool needLoadNextB, + bool atomicAdd) + { + auto layoutAInL1 = LayoutAInL1::template MakeLayout(l1TileShape.m(), l1TileShape.k()); + auto layoutBInL1 = LayoutBInL1::template MakeLayout(l1TileShape.k(), l1TileShape.n()); + + bool isFirstBlock = !preLoadedA && !preLoadedB; + bool isLastBlock = !needLoadNextA && !needLoadNextB; + uint32_t l1AListIdNext = (l1AListId + 1 < L1A_STAGES) ? (l1AListId + 1) : 0; + uint32_t l1BListIdNext = (l1BListId + 1 < L1B_STAGES) ? (l1BListId + 1) : 0; + + if (isFirstBlock) { + // Load matrix A tile from GM to L1 + AscendC::WaitFlag(l1AEventList[l1AListId]); + auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShape.m(), actualShape.k())); + copyGmToL1A(l1ATensorList[l1AListId], gmBlockA, layoutAInL1, layoutTileA); + AscendC::SetFlag(l1AEventList[l1AListId]); + + // Load matrix B tile from GM to L1 + AscendC::WaitFlag(l1BEventList[l1BListId]); + auto layoutTileB = layoutB.GetTileLayout(MakeCoord(actualShape.k(), actualShape.n())); + copyGmToL1B(l1BTensorList[l1BListId], gmBlockB, layoutBInL1, layoutTileB); + AscendC::SetFlag(l1BEventList[l1BListId]); + } + + if (needLoadNextA && (L1A_STAGES > 1)) { + AscendC::WaitFlag(l1AEventList[l1AListIdNext]); + auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShapeNext.m(), actualShapeNext.k())); + copyGmToL1A(l1ATensorList[l1AListIdNext], gmNextBlockA, layoutAInL1, layoutTileA); + AscendC::SetFlag(l1AEventList[l1AListIdNext]); + } + if (needLoadNextB && (L1B_STAGES > 1)) { + AscendC::WaitFlag(l1BEventList[l1BListIdNext]); + auto layoutTileB = layoutB.GetTileLayout(MakeCoord(actualShapeNext.k(), actualShapeNext.n())); + copyGmToL1B(l1BTensorList[l1BListIdNext], gmNextBlockB, layoutBInL1, layoutTileB); + AscendC::SetFlag(l1BEventList[l1BListIdNext]); + } + + // Get L1 tensor for current stage + auto l1ATensor = l1ATensorList[l1AListId]; + auto l1BTensor = l1BTensorList[l1BListId]; + + // Atomic set once in mmad level + if (atomicAdd) { + AscendC::SetAtomicAdd(); + } else { + AscendC::SetAtomicNone(); + } + + uint32_t l1mLoops = CeilDiv(actualShape.m(), l0TileShape.m()); + uint32_t l1nLoops = CeilDiv(actualShape.n(), l0TileShape.n()); + for (uint32_t l1mLoopIdx = 0; l1mLoopIdx < l1mLoops; l1mLoopIdx++) { + for (uint32_t l1nLoopIdx = 0; l1nLoopIdx < l1nLoops; l1nLoopIdx++) { + // layoutAInL0 + uint32_t mActual = (l1mLoopIdx < l1mLoops - 1) ? l0TileShape.m() : (actualShape.m() - l1mLoopIdx * l0TileShape.m()); + uint32_t nActual = (l1nLoopIdx < l1nLoops - 1) ? l0TileShape.n() : (actualShape.n() - l1nLoopIdx * l0TileShape.n()); + uint32_t mRound = RoundUp(mActual); + uint32_t nRound = RoundUp(nActual); + auto layoutInL0C = LayoutCInL0::MakeLayoutInL0C(MakeCoord(mRound, nRound)); + + uint32_t l0ABufId = 0; + uint32_t l0BBufId = 0; + + // wait/set conditions in block/l1mLoops/l1nLoops level + bool mnLoopStart = (l1mLoopIdx == 0 && l1nLoopIdx == 0); + bool mnLoopEnd = (l1mLoopIdx == l1mLoops - 1) && (l1nLoopIdx == l1nLoops - 1); + bool waitFlagA = (isFirstBlock || preLoadedA) && mnLoopStart; + bool setFlagA = (needLoadNextA || isLastBlock) && mnLoopEnd; + bool waitFlagB = (isFirstBlock || preLoadedB) && mnLoopStart; + bool setFlagB = (needLoadNextB || isLastBlock) && mnLoopEnd; + + if constexpr (!ENABLE_UNIT_FLAG) { + AscendC::WaitFlag(l0CEventList[l0CListId]); + } + + // Get the loop nums on L0 + uint32_t l1kLoops = CeilDiv(actualShape.k(), l0TileShape.k()); + for (uint32_t l1kLoopIdx = 0; l1kLoopIdx < l1kLoops; l1kLoopIdx++) { + uint32_t kActual = (l1kLoopIdx < l1kLoops - 1) ? l0TileShape.k() : (actualShape.k() - l1kLoopIdx * l0TileShape.k()); + + // Locate the current tile on L0A + auto l0ATile = l0ATensorList[l0ABufId]; + LayoutAInL0 layoutAInL0 = LayoutAInL0::template MakeLayout(mRound, kActual); + // Locate the current tile of matrix A on L1 + MatrixCoord l1ACoord{l1mLoopIdx * l0TileShape.m(), l1kLoopIdx * l0TileShape.k()}; + auto l1ATile = l1ATensor[layoutAInL1.GetOffset(l1ACoord)]; + + AscendC::WaitFlag(l0AEventList[l0ABufId]); + if (waitFlagA && (l1kLoopIdx == 0)) { + AscendC::WaitFlag(l1AEventList[l1AListId]); + } + + copyL1ToL0A(l0ATile, l1ATile, layoutAInL0, layoutAInL1); + + if (setFlagA && (l1kLoopIdx == l1kLoops - 1)) { + AscendC::SetFlag(l1AEventList[l1AListId]); + } + + // Locate the current tile on L0B + auto l0BTile = l0BTensorList[l0BBufId]; + LayoutBInL0 layoutBInL0 = LayoutBInL0::template MakeLayout(kActual, nRound); + // Locate the current tile of matrix B on L1 + MatrixCoord l1BCoord{l1kLoopIdx * l0TileShape.k(), l1nLoopIdx * l0TileShape.n()}; + auto l1BTile = l1BTensor[layoutBInL1.GetOffset(l1BCoord)]; + + // Wait for mmad finished + AscendC::WaitFlag(l0BEventList[l0BBufId]); + // If the current tile is the first one on the k&n axis, wait for loading matrix B from GM to L1 + if (waitFlagB && (l1kLoopIdx == 0)) { + AscendC::WaitFlag(l1BEventList[l1BListId]); + } + + // Load current tile from L1 to L0B + copyL1ToL0B(l0BTile, l1BTile, layoutBInL0, layoutBInL1); + + // If the current tile is the last one on the k&n axis, notify to load matrix B from GM to L1 + if (setFlagB && (l1kLoopIdx == l1kLoops - 1)) { + AscendC::SetFlag(l1BEventList[l1BListId]); + } + // Notify to do mmad + AscendC::SetFlag(l0CEventList[l0CListId]); + + // Locate the current tile on L0C + auto l0CTile = l0CTensorList[l0CListId]; + + // Compute the matrix multiplication on L0A and L0B and write the result to the accumulator + // Wait for loading L0B + AscendC::WaitFlag(l0CEventList[l0CListId]); + + // If the current tile is the first tile on the k axis, the accumulator needs to be reset to 0 + bool initC = (l1kLoopIdx == 0); + // If the unit flag is enabled, the unit flag is set according to the calculation progress + uint8_t unitFlag = 0b00; + if constexpr (ENABLE_UNIT_FLAG) { + if (l1kLoopIdx == l1kLoops - 1) { + unitFlag = 0b11; + } else { + unitFlag = 0b10; + } + } + // Perform calculation operations + tileMmad(l0CTile, l0ATile, l0BTile, mRound, nRound, kActual, initC, unitFlag); + + // Notify to move the next L0B tile + AscendC::SetFlag(l0BEventList[l0BBufId]); + l0BBufId = (l0BBufId + 1 < L0AB_STAGES) ? (l0BBufId + 1) : 0; + + AscendC::SetFlag(l0AEventList[l0ABufId]); + l0ABufId = (l0ABufId + 1 < L0AB_STAGES) ? (l0ABufId + 1) : 0; + } + // copy block out + LayoutC layoutBlock = layoutC.GetTileLayout(MakeCoord(mActual, nActual)); + MatrixCoord gmCoordC{l1mLoopIdx * l0TileShape.m(), l1nLoopIdx * l0TileShape.n()}; + auto gmTileC = gmBlockC[layoutC.GetOffset(gmCoordC)]; + if constexpr (!ENABLE_UNIT_FLAG) { + AscendC::SetFlag(l0CEventList[l0CListId]); + AscendC::WaitFlag(l0CEventList[l0CListId]); + copyL0CToGm(gmTileC, l0CTensorList[l0CListId], layoutBlock, layoutInL0C); + AscendC::SetFlag(l0CEventList[l0CListId]); + } else { + copyL0CToGm(gmTileC, l0CTensorList[l0CListId], layoutBlock, layoutInL0C, 0b11); + } + l0CListId = (l0CListId + 1 < L0C_STAGES) ? (l0CListId + 1) : 0; + } + } + + // if L1A_STAGES/L1B_STAGES == 1, load next block after compute + if (needLoadNextA && (L1A_STAGES == 1)) { + // Load first matrix A tile from GM to L1 + AscendC::WaitFlag(l1AEventList[l1AListIdNext]); + auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShapeNext.m(), actualShapeNext.k())); + copyGmToL1A(l1ATensorList[l1AListIdNext], gmNextBlockA, layoutAInL1, layoutTileA); + AscendC::SetFlag(l1AEventList[l1AListIdNext]); + } + if (needLoadNextB && (L1B_STAGES == 1)) { + // Load first matrix B tile from GM to L1 + AscendC::WaitFlag(l1BEventList[l1BListIdNext]); + auto layoutTileB = layoutB.GetTileLayout(MakeCoord(actualShapeNext.k(), actualShapeNext.n())); + copyGmToL1B(l1BTensorList[l1BListIdNext], gmNextBlockB, layoutBInL1, layoutTileB); + AscendC::SetFlag(l1BEventList[l1BListIdNext]); + } + + preLoadedA = needLoadNextA; + preLoadedB = needLoadNextB; + l1AListId = needLoadNextA ? l1AListIdNext : l1AListId; + l1BListId = needLoadNextB ? l1BListIdNext : l1BListId; + } + +protected: + /// Data members + AscendC::LocalTensor l1ATensorList[L1A_STAGES]; + AscendC::LocalTensor l1BTensorList[L1B_STAGES]; + AscendC::LocalTensor l0ATensorList[L0AB_STAGES]; + AscendC::LocalTensor l0BTensorList[L0AB_STAGES]; + AscendC::LocalTensor l0CTensorList[L0C_STAGES]; + + int32_t l1AEventList[L1A_STAGES]; + int32_t l1BEventList[L1B_STAGES]; + int32_t l0AEventList[L0AB_STAGES]; + int32_t l0BEventList[L0AB_STAGES]; + int32_t l0CEventList[L0C_STAGES]; + + bool preLoadedA{false}; + bool preLoadedB{false}; + uint32_t l1AListId{0}; + uint32_t l1BListId{0}; + uint32_t l0CListId{0}; + GemmCoord l1TileShape; + GemmCoord l0TileShape; + + TileMmad tileMmad; + CopyGmToL1A copyGmToL1A; + CopyGmToL1B copyGmToL1B; + CopyL1ToL0A copyL1ToL0A; + CopyL1ToL0B copyL1ToL0B; + CopyL0CToGm copyL0CToGm; +}; + +} // namespace Catlass::Gemm::Block + +#endif // CATLASS_GEMM_BLOCK_MMAD_DYNAMIC_SINGLE_CORE_SPLITK_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_dynamic_small.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_dynamic_small.hpp new file mode 100644 index 0000000000..37449a8dae --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_dynamic_small.hpp @@ -0,0 +1,284 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_GEMM_BLOCK_BLOCK_MMAD_DYNAMIC_SMALL_HPP +#define CATLASS_GEMM_BLOCK_BLOCK_MMAD_DYNAMIC_SMALL_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/kda_coord.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" +#include "catlass/gemm/kda_gemm_helper.hpp" +#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" +#include "catlass/gemm/tile/kda_tile_mmad.hpp" + +namespace Catlass::Gemm::Block { + +template +struct BlockMmad, L1TileShape_, L0TileShape_, + AType_, BType_, CType_, BiasType_, TileCopy_, TileMmad_> { +public: + // Type Aliases + using DispatchPolicy = MmadAtlasA2DynamicSmall; + using ArchTag = typename DispatchPolicy::ArchTag; + + using ElementA = typename AType_::Element; + using LayoutA = typename AType_::Layout; + using ElementB = typename BType_::Element; + using LayoutB = typename BType_::Layout; + using ElementC = typename CType_::Element; + using LayoutC = typename CType_::Layout; + using TileMmad = TileMmad_; + using CopyGmToL1A = typename TileCopy_::CopyGmToL1A; + using CopyGmToL1B = typename TileCopy_::CopyGmToL1B; + using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; + using CopyL0CToGm = typename TileCopy_::CopyL0CToGm; + using ElementAccumulator = + typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; + using LayoutAInL1 = typename CopyL1ToL0A::LayoutSrc; + using LayoutBInL1 = typename CopyL1ToL0B::LayoutSrc; + using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst; + using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst; + using LayoutCInL0 = layout::zN; + + using L1AAlignHelper = Gemm::helper::L1AlignHelper; + using L1BAlignHelper = Gemm::helper::L1AlignHelper; + + static constexpr bool ENABLE_UNIT_FLAG = DispatchPolicy::ENABLE_UNIT_FLAG; + static constexpr uint32_t STAGES = DispatchPolicy::STAGES; + + static constexpr uint32_t L0A_PINGPONG_BUF_SIZE = ArchTag::L0A_SIZE / STAGES; + static constexpr uint32_t L0B_PINGPONG_BUF_SIZE = ArchTag::L0B_SIZE / STAGES; + + // Check LayoutC + static_assert(std::is_same_v, "LayoutC only support RowMajor yet!"); + + /// Construct + CATLASS_DEVICE + BlockMmad(GemmCoord const &l1TileShape_, Arch::Resource &resource, uint32_t l1BufAddrStart = 0) + : l1TileShape(l1TileShape_) + { + uint32_t l1ASize = l1TileShape.m() * l1TileShape.k() * sizeof(ElementA); + uint32_t l1BSize = l1TileShape.k() * l1TileShape.n() * sizeof(ElementB); + + kPartLenMax = min(L0A_PINGPONG_BUF_SIZE / sizeof(ElementA) / l1TileShape.m() / L1AAlignHelper::ELE_NUM_PER_C0 * + L1AAlignHelper::ELE_NUM_PER_C0, + L0B_PINGPONG_BUF_SIZE / sizeof(ElementB) / l1TileShape.n() / L1BAlignHelper::ELE_NUM_PER_C0 * + L1BAlignHelper::ELE_NUM_PER_C0); + + if constexpr (std::is_same_v && std::is_same_v) { + kPartLenMax = RoundDown(kPartLenMax); + } + + uint32_t l1AOffset = l1BufAddrStart; + uint32_t l1BOffset = l1BufAddrStart + l1ASize * STAGES; + // Init buffers + for (uint32_t i = 0; i < STAGES; i++) { + l1ATensorList[i] = resource.l1Buf.template GetBufferByByte(l1AOffset + l1ASize * i); + l1BTensorList[i] = resource.l1Buf.template GetBufferByByte(l1BOffset + l1BSize * i); + l0ATensorList[i] = resource.l0ABuf.template GetBufferByByte(L0A_PINGPONG_BUF_SIZE * i); + l0BTensorList[i] = resource.l0BBuf.template GetBufferByByte(L0B_PINGPONG_BUF_SIZE * i); + + l1AEventList[i] = i; + l1BEventList[i] = i + STAGES; + l0AEventList[i] = i; + l0BEventList[i] = i + STAGES; + AscendC::SetFlag(l1AEventList[i]); + AscendC::SetFlag(l1BEventList[i]); + AscendC::SetFlag(l0AEventList[i]); + AscendC::SetFlag(l0BEventList[i]); + } + l0CTensor = resource.l0CBuf.template GetBufferByByte(0); + if constexpr (!ENABLE_UNIT_FLAG) { + AscendC::SetFlag(EVENT_ID0); + } + } + + /// Destructor + CATLASS_DEVICE + ~BlockMmad() + { + for (uint32_t i = 0; i < STAGES; i++) { + AscendC::WaitFlag(l1AEventList[i]); + AscendC::WaitFlag(l1BEventList[i]); + AscendC::WaitFlag(l0AEventList[i]); + AscendC::WaitFlag(l0BEventList[i]); + } + if constexpr (!ENABLE_UNIT_FLAG) { + AscendC::WaitFlag(EVENT_ID0); + } + } + + /// Perform a block-scoped matrix multiply-accumulate + CATLASS_DEVICE + void operator()(AscendC::GlobalTensor const &gmBlockA, LayoutA const &layoutA, + AscendC::GlobalTensor const &gmBlockB, LayoutB const &layoutB, + AscendC::GlobalTensor const &gmBlockC, LayoutC const &layoutC, + GemmCoord const &actualShape) + { + uint32_t mRound = RoundUp(actualShape.m()); + uint32_t nRound = RoundUp(actualShape.n()); + + auto layoutAInL1 = LayoutAInL1::template MakeLayout(l1TileShape.m(), l1TileShape.k()); + auto layoutBInL1 = LayoutBInL1::template MakeLayout(l1TileShape.k(), l1TileShape.n()); + auto layoutInL0C = LayoutCInL0::MakeLayoutInL0C(MakeCoord(mRound, nRound)); + + if constexpr (!ENABLE_UNIT_FLAG) { + AscendC::WaitFlag(EVENT_ID0); + } + + uint32_t kActual = actualShape.k(); + + // k loop + { + // For small-shape matmul, the k-loop executes only once. + + // load first matrix A tile from GM to L1 + AscendC::WaitFlag(l1AEventList[l1ListId]); + auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShape.m(), kActual)); + copyGmToL1A(l1ATensorList[l1ListId], gmBlockA, layoutAInL1, layoutTileA); + AscendC::SetFlag(l1AEventList[l1ListId]); + + // load first matrix B tile from GM to L1 + AscendC::WaitFlag(l1BEventList[l1ListId]); + auto layoutTileB = layoutB.GetTileLayout(MakeCoord(kActual, actualShape.n())); + copyGmToL1B(l1BTensorList[l1ListId], gmBlockB, layoutBInL1, layoutTileB); + AscendC::SetFlag(l1BEventList[l1ListId]); + + // Get L1 tensor for current stage + auto l1ATensor = l1ATensorList[l1ListId]; + auto l1BTensor = l1BTensorList[l1ListId]; + + // Get the loop nums on L0 + uint32_t kPartLoop = CeilDiv(kActual, kPartLenMax); + + for (int kPartIdx = 0; kPartIdx < kPartLoop; kPartIdx++) { + uint32_t kPartActual = (kPartIdx < kPartLoop - 1) ? kPartLenMax : (kActual - kPartIdx * kPartLenMax); + + // Locate the current tile on L0A + auto l0ATile = l0ATensorList[l0AListId]; + LayoutAInL0 layoutAInL0 = LayoutAInL0::template MakeLayout(mRound, kPartActual); + // Locate the current tile of matrix A on L1 + MatrixCoord l1AOffset{0, kPartIdx * kPartLenMax}; + auto l1ATile = l1ATensor[layoutAInL1.GetOffset(l1AOffset)]; + + AscendC::WaitFlag(l0AEventList[l0AListId]); + if (kPartIdx == 0) { + AscendC::WaitFlag(l1AEventList[l1ListId]); + } + + // Load current tile from L1 to L0A + copyL1ToL0A(l0ATile, l1ATile, layoutAInL0, layoutAInL1); + + if (kPartIdx == kPartLoop - 1) { + AscendC::SetFlag(l1AEventList[l1ListId]); + } + + // Locate the current tile on L0B + auto l0BTile = l0BTensorList[l0BListId]; + LayoutBInL0 layoutBInL0 = LayoutBInL0::template MakeLayout(kPartActual, nRound); + // Locate the current tile of matrix B on L1 + MatrixCoord l1BOffset{kPartIdx * kPartLenMax, 0}; + auto l1BTile = l1BTensor[layoutBInL1.GetOffset(l1BOffset)]; + + // Wait for mmad finished + AscendC::WaitFlag(l0BEventList[l0BListId]); + // If the current tile is the first one on the k&n axis, wait for loading matrix B from GM to L1 + if (kPartIdx == 0) { + AscendC::WaitFlag(l1BEventList[l1ListId]); + } + + // Load current tile from L1 to L0B + copyL1ToL0B(l0BTile, l1BTile, layoutBInL0, layoutBInL1); + + // If the current tile is the last one on the k&n axis, notify to load matrix B from GM to L1 + if (kPartIdx == kPartLoop - 1) { + AscendC::SetFlag(l1BEventList[l1ListId]); + } + // Notify to do mmad + AscendC::SetFlag(EVENT_ID0); + + // Locate the current tile on L0C + auto l0CTile = l0CTensor; + + // Compute the matrix multiplication on L0A and L0B and write the result to the accumulator + // Wait for loading L0B + AscendC::WaitFlag(EVENT_ID0); + + // If the current tile is the first tile on the k axis, the accumulator needs to be reset to 0 + bool initC = (kPartIdx == 0); + // If the unit flag is enabled, the unit flag is set according to the calculation progress + uint8_t unitFlag = 0b00; + if constexpr (ENABLE_UNIT_FLAG) { + if (kPartIdx == kPartLoop - 1) { + unitFlag = 0b11; + } else { + unitFlag = 0b10; + } + } + // Perform calculation operations + tileMmad(l0CTile, l0ATile, l0BTile, mRound, nRound, kPartActual, initC, unitFlag); + + // Notify to move the next L0B tile + AscendC::SetFlag(l0BEventList[l0BListId]); + l0BListId = (l0BListId + 1 < STAGES) ? (l0BListId + 1) : 0; + + AscendC::SetFlag(l0AEventList[l0AListId]); + l0AListId = (l0AListId + 1 < STAGES) ? (l0AListId + 1) : 0; + } + } + + // copy block out + LayoutC layoutBlock = layoutC.GetTileLayout(actualShape.GetCoordMN()); + + if constexpr (!ENABLE_UNIT_FLAG) { + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + copyL0CToGm(gmBlockC, l0CTensor, layoutBlock, layoutInL0C); + AscendC::SetFlag(EVENT_ID0); + } else { + copyL0CToGm(gmBlockC, l0CTensor, layoutBlock, layoutInL0C, 0b11); + } + } + +protected: + /// Data members + AscendC::LocalTensor l1ATensorList[STAGES]; + AscendC::LocalTensor l1BTensorList[STAGES]; + AscendC::LocalTensor l0ATensorList[STAGES]; + AscendC::LocalTensor l0BTensorList[STAGES]; + AscendC::LocalTensor l0CTensor; + + int32_t l1AEventList[STAGES]; + int32_t l1BEventList[STAGES]; + int32_t l0AEventList[STAGES]; + int32_t l0BEventList[STAGES]; + + uint32_t l1ListId{0}; + uint32_t l0AListId{0}; + uint32_t l0BListId{0}; + + GemmCoord l1TileShape; + uint32_t kPartLenMax; + + TileMmad tileMmad; + CopyGmToL1A copyGmToL1A; + CopyGmToL1B copyGmToL1B; + CopyL1ToL0A copyL1ToL0A; + CopyL1ToL0B copyL1ToL0B; + CopyL0CToGm copyL0CToGm; +}; + +} // namespace Catlass::Gemm::Block + +#endif // CATLASS_GEMM_BLOCK_BLOCK_MMAD_DYNAMIC_SMALL_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_dynamic_streamk.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_dynamic_streamk.hpp new file mode 100644 index 0000000000..e5e6933ef4 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_dynamic_streamk.hpp @@ -0,0 +1,370 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_GEMM_BLOCK_BLOCK_MMAD_DYNAMIC_STREAMK_HPP +#define CATLASS_GEMM_BLOCK_BLOCK_MMAD_DYNAMIC_STREAMK_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/kda_coord.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" +#include "catlass/gemm/kda_gemm_helper.hpp" +#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" +#include "catlass/gemm/tile/kda_tile_mmad.hpp" + +namespace Catlass::Gemm::Block { + +template +struct BlockMmad, L1TileShape_, L0TileShape_, AType_, + BType_, CType_, BiasType_, TileCopy_, TileMmad_> { +public: + // Type Aliases + using DispatchPolicy = MmadAtlasA2DynamicStreamk; + using ArchTag = typename DispatchPolicy::ArchTag; + using ElementA = typename AType_::Element; + using LayoutA = typename AType_::Layout; + using ElementB = typename BType_::Element; + using LayoutB = typename BType_::Layout; + using ElementC = typename CType_::Element; + using LayoutC = typename CType_::Layout; + using TileMmad = TileMmad_; + using CopyGmToL1A = typename TileCopy_::CopyGmToL1A; + using CopyGmToL1B = typename TileCopy_::CopyGmToL1B; + using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; + using CopyL0CToGmNormalBlock = typename TileCopy_::CopyL0CToGm; + using CopyL0CToGmStreamkBlock = Gemm::Tile::CopyL0CToGm>; + using ElementAccumulator = + typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; + using LayoutAInL1 = typename CopyL1ToL0A::LayoutSrc; + using LayoutBInL1 = typename CopyL1ToL0B::LayoutSrc; + using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst; + using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst; + using LayoutCInL0 = layout::zN; + + using L1AAlignHelper = Gemm::helper::L1AlignHelper; + using L1BAlignHelper = Gemm::helper::L1AlignHelper; + + static constexpr bool ENABLE_UNIT_FLAG = DispatchPolicy::ENABLE_UNIT_FLAG; + static constexpr bool ENABLE_SHUFFLE_K = DispatchPolicy::ENABLE_SHUFFLE_K; + static constexpr uint32_t STAGES = DispatchPolicy::STAGES; + + static constexpr uint32_t L0A_PINGPONG_BUF_SIZE = ArchTag::L0A_SIZE / STAGES; + static constexpr uint32_t L0B_PINGPONG_BUF_SIZE = ArchTag::L0B_SIZE / STAGES; + + // Check LayoutC + static_assert(std::is_same_v, "LayoutC only support RowMajor yet!"); + + /// Construct + CATLASS_DEVICE + BlockMmad(GemmCoord const &l1TileShape_, Arch::Resource &resource, uint32_t l1BufAddrStart = 0) + : l1TileShape(l1TileShape_) + { + uint32_t l1ASize = l1TileShape.m() * l1TileShape.k() * sizeof(ElementA); + uint32_t l1BSize = l1TileShape.k() * l1TileShape.n() * sizeof(ElementB); + + kPartLenMax = min(L0A_PINGPONG_BUF_SIZE / sizeof(ElementA) / l1TileShape.m() / L1AAlignHelper::ELE_NUM_PER_C0 * + L1AAlignHelper::ELE_NUM_PER_C0, + L0B_PINGPONG_BUF_SIZE / sizeof(ElementB) / l1TileShape.n() / L1BAlignHelper::ELE_NUM_PER_C0 * + L1BAlignHelper::ELE_NUM_PER_C0); + + if constexpr (std::is_same_v && std::is_same_v) { + kPartLenMax = RoundDown(kPartLenMax); + } + + uint32_t l1AOffset = l1BufAddrStart; + uint32_t l1BOffset = l1BufAddrStart + l1ASize * STAGES; + // Init buffers + for (uint32_t i = 0; i < STAGES; i++) { + l1ATensorList[i] = resource.l1Buf.template GetBufferByByte(l1AOffset + l1ASize * i); + l1BTensorList[i] = resource.l1Buf.template GetBufferByByte(l1BOffset + l1BSize * i); + l0ATensorList[i] = resource.l0ABuf.template GetBufferByByte(L0A_PINGPONG_BUF_SIZE * i); + l0BTensorList[i] = resource.l0BBuf.template GetBufferByByte(L0B_PINGPONG_BUF_SIZE * i); + + l1AEventList[i] = i; + l1BEventList[i] = i + STAGES; + l0AEventList[i] = i; + l0BEventList[i] = i + STAGES; + AscendC::SetFlag(l1AEventList[i]); + AscendC::SetFlag(l1BEventList[i]); + AscendC::SetFlag(l0AEventList[i]); + AscendC::SetFlag(l0BEventList[i]); + } + l0CTensor = resource.l0CBuf.template GetBufferByByte(0); + if constexpr (!ENABLE_UNIT_FLAG) { + AscendC::SetFlag(EVENT_ID0); + } + } + + /// Destructor + CATLASS_DEVICE + ~BlockMmad() + { + for (uint32_t i = 0; i < STAGES; i++) { + AscendC::WaitFlag(l1AEventList[i]); + AscendC::WaitFlag(l1BEventList[i]); + AscendC::WaitFlag(l0AEventList[i]); + AscendC::WaitFlag(l0BEventList[i]); + } + if constexpr (!ENABLE_UNIT_FLAG) { + AscendC::WaitFlag(EVENT_ID0); + } + } + + /// Perform a block-scoped matrix multiply-accumulate + CATLASS_DEVICE + void operator()( + AscendC::GlobalTensor const &gmBlockA, LayoutA const &layoutA, + AscendC::GlobalTensor const &gmBlockB, LayoutB const &layoutB, + AscendC::GlobalTensor const &gmBlockC, LayoutC const &layoutC, + AscendC::GlobalTensor const &gmW, LayoutC const &layoutW, + AscendC::GlobalTensor const &gmNextBlockA, AscendC::GlobalTensor const &gmNextBlockB, + GemmCoord const &actualShape, GemmCoord const &actualShapeNext, + bool isFirstBlock, bool hasNextBlock, bool isSkBlock) + { + uint32_t mRound = RoundUp(actualShape.m()); + uint32_t nRound = RoundUp(actualShape.n()); + + auto layoutAInL1 = LayoutAInL1::template MakeLayout(l1TileShape.m(), l1TileShape.k()); + auto layoutBInL1 = LayoutBInL1::template MakeLayout(l1TileShape.k(), l1TileShape.n()); + auto layoutInL0C = LayoutCInL0::MakeLayoutInL0C(MakeCoord(mRound, nRound)); + + uint32_t kTileCount = CeilDiv(actualShape.k(), l1TileShape.k()); + + if constexpr (!ENABLE_UNIT_FLAG) { + AscendC::WaitFlag(EVENT_ID0); + } + uint32_t startTileIdx = 0; + if constexpr (ENABLE_SHUFFLE_K) { + startTileIdx = AscendC::GetBlockIdx(); + } + uint32_t firstTileIdx = startTileIdx % kTileCount; + uint32_t lastTileIdx = (startTileIdx + kTileCount - 1) % kTileCount; + uint32_t kActual = + (firstTileIdx < kTileCount - 1) ? l1TileShape.k() : (actualShape.k() - firstTileIdx * l1TileShape.k()); + + // k loop + for (uint32_t kLoopIdx = 0; kLoopIdx < kTileCount; kLoopIdx++) { + uint32_t shuffleKIdx = (startTileIdx + kLoopIdx) % kTileCount; + // Load first matrix A tile in total kernel loop from GM to L1 + if (shuffleKIdx == firstTileIdx && isFirstBlock) { + MatrixCoord gmTileACoord{0, shuffleKIdx * l1TileShape.k()}; + MatrixCoord gmTileBCoord{shuffleKIdx * l1TileShape.k(), 0}; + auto gmTileA = gmBlockA[layoutA.GetOffset(gmTileACoord)]; + auto gmTileB = gmBlockB[layoutB.GetOffset(gmTileBCoord)]; + // Load first matrix A tile from GM to L1 + AscendC::WaitFlag(l1AEventList[l1ListId]); + auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShape.m(), kActual)); + copyGmToL1A(l1ATensorList[l1ListId], gmTileA, layoutAInL1, layoutTileA); + AscendC::SetFlag(l1AEventList[l1ListId]); + + // Load first matrix B tile from GM to L1 + AscendC::WaitFlag(l1BEventList[l1ListId]); + auto layoutTileB = layoutB.GetTileLayout(MakeCoord(kActual, actualShape.n())); + copyGmToL1B(l1BTensorList[l1ListId], gmTileB, layoutBInL1, layoutTileB); + AscendC::SetFlag(l1BEventList[l1ListId]); + } + + uint32_t l1ListIdNext = (l1ListId + 1 < STAGES) ? (l1ListId + 1) : 0; + uint32_t kActualNext{0}; + + // preload next tile from GM to L1 + if (shuffleKIdx != lastTileIdx) { + uint32_t shuffleKIdxNext = (startTileIdx + kLoopIdx + 1) % kTileCount; + // Get L1 tensor for next stage + auto l1ATensor = l1ATensorList[l1ListIdNext]; + auto l1BTensor = l1BTensorList[l1ListIdNext]; + // Get GM tensor for next stage + kActualNext = (shuffleKIdxNext < kTileCount - 1) ? l1TileShape.k() : (actualShape.k() - shuffleKIdxNext * l1TileShape.k()); + MatrixCoord gmTileACoord{0, shuffleKIdxNext * l1TileShape.k()}; + MatrixCoord gmTileBCoord{shuffleKIdxNext * l1TileShape.k(), 0}; + auto gmTileA = gmBlockA[layoutA.GetOffset(gmTileACoord)]; + auto gmTileB = gmBlockB[layoutB.GetOffset(gmTileBCoord)]; + + // load next matrix A tile from GM to L1 + AscendC::WaitFlag(l1AEventList[l1ListIdNext]); + auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShape.m(), kActualNext)); + copyGmToL1A(l1ATensor, gmTileA, layoutAInL1, layoutTileA); + AscendC::SetFlag(l1AEventList[l1ListIdNext]); + + // load next matrix B tile from GM to L1 + AscendC::WaitFlag(l1BEventList[l1ListIdNext]); + auto layoutTileB = layoutB.GetTileLayout(MakeCoord(kActualNext, actualShape.n())); + copyGmToL1B(l1BTensor, gmTileB, layoutBInL1, layoutTileB); + AscendC::SetFlag(l1BEventList[l1ListIdNext]); + } + if (shuffleKIdx == lastTileIdx && hasNextBlock) { + uint32_t kTileCount = CeilDiv(actualShapeNext.k(), l1TileShape.k()); + uint32_t firstTileIdx = startTileIdx % kTileCount; + // Get L1 tensor for next stage + auto l1ATensor = l1ATensorList[l1ListIdNext]; + auto l1BTensor = l1BTensorList[l1ListIdNext]; + // Get GM tensor for next stage + kActualNext = (firstTileIdx < kTileCount - 1) ? l1TileShape.k() : (actualShapeNext.k() - firstTileIdx * l1TileShape.k()); + MatrixCoord gmTileACoord{0, firstTileIdx * l1TileShape.k()}; + MatrixCoord gmTileBCoord{firstTileIdx * l1TileShape.k(), 0}; + auto gmTileA = gmNextBlockA[layoutA.GetOffset(gmTileACoord)]; + auto gmTileB = gmNextBlockB[layoutB.GetOffset(gmTileBCoord)]; + // load next matrix A tile from GM to L1 + AscendC::WaitFlag(l1AEventList[l1ListIdNext]); + auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShapeNext.m(), kActualNext)); + copyGmToL1A(l1ATensor, gmTileA, layoutAInL1, layoutTileA); + AscendC::SetFlag(l1AEventList[l1ListIdNext]); + + // load next matrix B tile from GM to L1 + AscendC::WaitFlag(l1BEventList[l1ListIdNext]); + auto layoutTileB = layoutB.GetTileLayout(MakeCoord(kActualNext, actualShapeNext.n())); + copyGmToL1B(l1BTensor, gmTileB, layoutBInL1, layoutTileB); + AscendC::SetFlag(l1BEventList[l1ListIdNext]); + } + + // Get L1 tensor for current stage + auto l1ATensor = l1ATensorList[l1ListId]; + auto l1BTensor = l1BTensorList[l1ListId]; + + // Get the loop nums on L0 + uint32_t kPartLoop = CeilDiv(kActual, kPartLenMax); + + uint32_t l0ABufId = 0; + uint32_t l0BBufId = 0; + + for (int kPartIdx = 0; kPartIdx < kPartLoop; kPartIdx++) { + uint32_t kPartActual = + (kPartIdx < kPartLoop - 1) ? kPartLenMax : (kActual - kPartIdx * kPartLenMax); + + // Locate the current tile on L0A + auto l0ATile = l0ATensorList[l0ABufId]; + LayoutAInL0 layoutAInL0 = LayoutAInL0::template MakeLayout(mRound, kPartActual); + // Locate the current tile of matrix A on L1 + MatrixCoord l1ACoord{0, kPartIdx * kPartLenMax}; + auto l1ATile = l1ATensor[layoutAInL1.GetOffset(l1ACoord)]; + + AscendC::WaitFlag(l0AEventList[l0ABufId]); + if (kPartIdx == 0) { + AscendC::WaitFlag(l1AEventList[l1ListId]); + } + + copyL1ToL0A(l0ATile, l1ATile, layoutAInL0, layoutAInL1); + + if (kPartIdx == kPartLoop - 1) { + AscendC::SetFlag(l1AEventList[l1ListId]); + } + + // Locate the current tile on L0B + auto l0BTile = l0BTensorList[l0BBufId]; + LayoutBInL0 layoutBInL0 = LayoutBInL0::template MakeLayout(kPartActual, nRound); + // Locate the current tile of matrix B on L1 + MatrixCoord l1BCoord{kPartIdx * kPartLenMax, 0}; + auto l1BTile = l1BTensor[layoutBInL1.GetOffset(l1BCoord)]; + + // Wait for mmad finished + AscendC::WaitFlag(l0BEventList[l0BBufId]); + // If the current tile is the first one on the k&n axis, wait for loading matrix B from GM to L1 + if (kPartIdx == 0) { + AscendC::WaitFlag(l1BEventList[l1ListId]); + } + + // Load current tile from L1 to L0B + copyL1ToL0B(l0BTile, l1BTile, layoutBInL0, layoutBInL1); + + // If the current tile is the last one on the k&n axis, notify to load matrix B from GM to L1 + if (kPartIdx == kPartLoop - 1) { + AscendC::SetFlag(l1BEventList[l1ListId]); + } + // Notify to do mmad + AscendC::SetFlag(EVENT_ID0); + + // Locate the current tile on L0C + MatrixCoord l0CCoord{static_cast(0U), static_cast(0U)}; + auto l0CTile = l0CTensor[layoutInL0C.GetOffset(l0CCoord)]; + + // Compute the matrix multiplication on L0A and L0B and write the result to the accumulator + // Wait for loading L0B + AscendC::WaitFlag(EVENT_ID0); + + // If the current tile is the first tile on the k axis, the accumulator needs to be reset to 0 + bool initC = ((kLoopIdx == 0) && (kPartIdx == 0)); + // If the unit flag is enabled, the unit flag is set according to the calculation progress + uint8_t unitFlag = 0b00; + if constexpr (ENABLE_UNIT_FLAG) { + if ((kLoopIdx == kTileCount - 1) && (kPartIdx == kPartLoop - 1)) { + unitFlag = 0b11; + } else { + unitFlag = 0b10; + } + } + // Perform calculation operations + tileMmad(l0CTile, l0ATile, l0BTile, mRound, nRound, kPartActual, initC, unitFlag); + + // Notify to move the next L0B tile + AscendC::SetFlag(l0BEventList[l0BBufId]); + + l0BBufId = (l0BBufId + 1 < STAGES) ? (l0BBufId + 1) : 0; + + AscendC::SetFlag(l0AEventList[l0ABufId]); + l0ABufId = (l0ABufId + 1 < STAGES) ? (l0ABufId + 1) : 0; + } + + l1ListId = l1ListIdNext; + kActual = kActualNext; + } + + // copy block out + LayoutC layoutBlock = layoutC.GetTileLayout(actualShape.GetCoordMN()); + + if constexpr (!ENABLE_UNIT_FLAG) { + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + if (isSkBlock) { + copyL0CToGmStreamkBlock(gmW, l0CTensor, layoutW, layoutInL0C); + } else { + copyL0CToGmNormalBlock(gmBlockC, l0CTensor, layoutBlock, layoutInL0C); + } + AscendC::SetFlag(EVENT_ID0); + } else { + if (isSkBlock) { + copyL0CToGmStreamkBlock(gmW, l0CTensor, layoutW, layoutInL0C, 0b11); + } else { + copyL0CToGmNormalBlock(gmBlockC, l0CTensor, layoutBlock, layoutInL0C, 0b11); + } + } + } + +protected: + /// Data members + AscendC::LocalTensor l1ATensorList[STAGES]; + AscendC::LocalTensor l1BTensorList[STAGES]; + AscendC::LocalTensor l0ATensorList[STAGES]; + AscendC::LocalTensor l0BTensorList[STAGES]; + AscendC::LocalTensor l0CTensor; + + int32_t l1AEventList[STAGES]; + int32_t l1BEventList[STAGES]; + int32_t l0AEventList[STAGES]; + int32_t l0BEventList[STAGES]; + + uint32_t l1ListId{0}; + GemmCoord l1TileShape; + uint32_t kPartLenMax; + + TileMmad tileMmad; + CopyGmToL1A copyGmToL1A; + CopyGmToL1B copyGmToL1B; + CopyL1ToL0A copyL1ToL0A; + CopyL1ToL0B copyL1ToL0B; + CopyL0CToGmNormalBlock copyL0CToGmNormalBlock; + CopyL0CToGmStreamkBlock copyL0CToGmStreamkBlock; +}; + +} // namespace Catlass::Gemm::Block + +#endif // CATLASS_GEMM_BLOCK_BLOCK_MMAD_DYNAMIC_STREAMK_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_pv_head_tail.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_pv_head_tail.hpp new file mode 100644 index 0000000000..63fcde0508 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_pv_head_tail.hpp @@ -0,0 +1,315 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_GEMM_BLOCK_MMAD_PV_TAIL_HPP +#define CATLASS_GEMM_BLOCK_MMAD_PV_TAIL_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/kda_coord.hpp" +#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" +#include "catlass/gemm/kda_gemm_helper.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" +#include "catlass/gemm/tile/kda_tile_mmad.hpp" + +//////////////////////////////////////////////////////////////////// + +namespace Catlass::Gemm::Block { +//////////////////////////////////////////////////////////////////// + +template < + bool PAGED_CACHE_FLAG_, + bool ENABLE_UNIT_FLAG_, + class L1TileShape_, + class L0TileShape_, + class AType_, + class BType_, + class CType_, + class BiasType_, + class TileCopy_, + class TileMmad_> +struct BlockMmad< + MmadAtlasA2FAITailPV, + L1TileShape_, + L0TileShape_, + AType_, + BType_, + CType_, + BiasType_, + TileCopy_, + TileMmad_> { +public: + // Type Aliases + using DispatchPolicy = MmadAtlasA2FAITailPV; + using ArchTag = typename DispatchPolicy::ArchTag; + using L1TileShape = L1TileShape_; + using L0TileShape = L0TileShape_; + using ElementA = typename AType_::Element; + using LayoutA = typename AType_::Layout; + using ElementB = typename BType_::Element; + using LayoutB = typename BType_::Layout; + using ElementC = typename CType_::Element; + using LayoutC = typename CType_::Layout; + using TileMmad = TileMmad_; + using CopyGmToL1A = typename TileCopy_::CopyGmToL1A; + using CopyGmToL1B = typename TileCopy_::CopyGmToL1B; + using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; + using CopyL0CToGm = typename TileCopy_::CopyL0CToGm; + using ElementAccumulator = + typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; + using LayoutAInL1 = typename CopyL1ToL0A::LayoutSrc; + using LayoutBInL1 = typename CopyL1ToL0B::LayoutSrc; + using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst; + using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst; + using LayoutCInL0 = layout::zN; + + using L1AAlignHelper = Gemm::helper::L1AlignHelper; + using L1BAlignHelper = Gemm::helper::L1AlignHelper; + + static constexpr uint32_t STAGES = DispatchPolicy::STAGES; + static constexpr uint32_t L1A_SIZE = L1TileShape::M * L1TileShape::K * sizeof(ElementA); + static constexpr uint32_t L1B_SIZE = L1TileShape::N * L1TileShape::K * sizeof(ElementB); + static constexpr uint32_t L0A_SIZE = ArchTag::L0A_SIZE; + static constexpr uint32_t L0B_SIZE = ArchTag::L0B_SIZE; + static constexpr uint32_t L0C_SIZE = ArchTag::L0C_SIZE; + static constexpr uint32_t L0A_PINGPONG_BUF_SIZE = L0A_SIZE / STAGES; + static constexpr uint32_t L0B_PINGPONG_BUF_SIZE = L0B_SIZE / STAGES; + static constexpr uint32_t L0C_PINGPONG_BUF_SIZE = L0C_SIZE / STAGES; + static constexpr uint32_t BLOCK_SIZE = 16; + static constexpr uint32_t EMBED_SPLIT_SIZE = 128; + static constexpr uint32_t UNIT_BLOCK_STACK_NUM = 4; + static constexpr uint32_t KV_BASE_BLOCK = 512; + static constexpr uint32_t KV_SPLIT_SIZE = 128; + static constexpr uint32_t LOAB_BLOCK = 1; + + // Check LayoutC + static_assert(std::is_same_v, "LayoutC only support RowMajor yet!"); + + /// Construct + CATLASS_DEVICE + BlockMmad(Arch::Resource &resource, uint32_t l1BufAddrStart = 0) + { + // Allocate L1 memory space + l0ATensor = resource.l0ABuf.template GetBufferByByte(0); + l0BTensor = resource.l0BBuf.template GetBufferByByte(0); + for (uint32_t i = 0; i < STAGES; i++) { + l1ATensor[i] = resource.l1Buf.template GetBufferByByte(l1BufAddrStart + L1A_SIZE * LOAB_BLOCK * i); + l0CTensor[i] = resource.l0CBuf.template GetBufferByByte(L0C_PINGPONG_BUF_SIZE * i); + l1BTensor[i] = + resource.l1Buf.template GetBufferByByte(l1BufAddrStart + L1A_SIZE * 2 * 2 + L1B_SIZE * i); + } + } + + /// Destructor + CATLASS_DEVICE + ~BlockMmad() {} + + CATLASS_DEVICE + void getBlockShape(GemmCoord &actualShape, uint32_t &nowNIdx, uint32_t &kIdx, + uint32_t &nLoop, uint32_t &kLoop, uint32_t &kvSeqlen, uint32_t &embed, bool firstBlock, uint32_t maskTailS = 0) + { + uint32_t nSplitSize = KV_SPLIT_SIZE * LOAB_BLOCK; + uint32_t embedSplitSize = EMBED_SPLIT_SIZE; + if (nowNIdx + LOAB_BLOCK > nLoop - 1) { + nSplitSize = kvSeqlen - nowNIdx * KV_SPLIT_SIZE; + } + if (firstBlock && maskTailS != 0) { + nSplitSize = nSplitSize - maskTailS; + } + if (kIdx == kLoop - 1) { + embedSplitSize = embed - kIdx * EMBED_SPLIT_SIZE; + } + actualShape[1] = embedSplitSize; + actualShape[2] = nSplitSize; + } + + CATLASS_DEVICE + void getKVOffset(AscendC::GlobalTensor &gBlockTable, uint32_t &kOffset, uint32_t &nowNIdx, uint32_t &kIdx, + uint32_t &nLoop, uint32_t &kLoop, uint32_t &strideKV, uint32_t &blockSize, uint32_t maskTailS = 0) + { + if (nowNIdx >= nLoop || kIdx >= kLoop) { + kOffset = 0; + } + if constexpr (PAGED_CACHE_FLAG_) { + uint32_t blockTableId = gBlockTable.GetValue(nowNIdx); + kOffset = blockTableId * blockSize * strideKV + maskTailS * strideKV + kIdx * EMBED_SPLIT_SIZE; + } else { + kOffset = nowNIdx * KV_SPLIT_SIZE * strideKV + kIdx * EMBED_SPLIT_SIZE; + } + } + + CATLASS_DEVICE + void operator()(AscendC::GlobalTensor gA, + AscendC::GlobalTensor gB, + AscendC::GlobalTensor gC, + AscendC::GlobalTensor gBlockTable, + LayoutA layoutA, LayoutB layoutB, GemmCoord actualOriShape, + uint32_t &nIdx, uint32_t &nLoop, uint32_t &blockSize, uint32_t kvSeqlen, uint32_t strideKV, Arch::CrossCoreFlag softmaxFlag, uint32_t maskTailS, bool preloadFlag) + { + uint32_t embed = actualOriShape[1]; + uint32_t kLoop = CeilDiv(embed); + uint32_t rowNum = layoutA.shape(0); + uint32_t blockN = layoutA.shape(1); + GemmCoord actualShape{rowNum, 0, 0}; + GemmCoord actualNextShape{rowNum, 0, 0}; + uint32_t nkBlockLoop = (nLoop + LOAB_BLOCK - 1) / LOAB_BLOCK * kLoop; // gap + uint32_t nkBlockNextIdx = (nIdx + LOAB_BLOCK - 1) / LOAB_BLOCK * kLoop + 1; // gap + uint32_t gBOffset = 0; + uint32_t gBNextOffset = 0; + uint32_t nowMaskTailS = 0; + uint32_t gPOffset = 0; + for (uint32_t kIdx = 0; kIdx < kLoop; kIdx++) { + nowMaskTailS = maskTailS; + gPOffset = 0; + for (uint32_t blockStackIdx = 0; (blockStackIdx < UNIT_BLOCK_STACK_NUM) && ((nIdx + blockStackIdx) < nLoop); + blockStackIdx += LOAB_BLOCK) { + uint32_t nowNIdx = nIdx + blockStackIdx; + uint32_t kLoopNextIdx = (nkBlockNextIdx % (kLoop * UNIT_BLOCK_STACK_NUM)) / (UNIT_BLOCK_STACK_NUM / LOAB_BLOCK); + uint32_t nLoopNextIdx = (nkBlockNextIdx % (kLoop * UNIT_BLOCK_STACK_NUM)) % (UNIT_BLOCK_STACK_NUM / LOAB_BLOCK) + nkBlockNextIdx / (kLoop * UNIT_BLOCK_STACK_NUM) * UNIT_BLOCK_STACK_NUM; + uint32_t startSeqOffset = nowNIdx == nIdx ? maskTailS : 0; + uint32_t startSeqNxtOffset = nLoopNextIdx == nIdx ? maskTailS : 0; + getBlockShape(actualShape, nowNIdx, kIdx, nLoop, kLoop, kvSeqlen, embed, nowNIdx == nIdx, nowMaskTailS); + getBlockShape(actualNextShape, nLoopNextIdx, kLoopNextIdx, nLoop, kLoop, kvSeqlen, embed, nLoopNextIdx == nIdx, nowMaskTailS); + getKVOffset(gBlockTable, gBOffset, nowNIdx, kIdx, nLoop, kLoop, strideKV, blockSize, startSeqOffset); + getKVOffset(gBlockTable, gBNextOffset, nLoopNextIdx, kLoopNextIdx, nLoop, kLoop, strideKV, blockSize, startSeqNxtOffset); + bool firstItr = blockStackIdx == 0; + bool endItr = (blockStackIdx + LOAB_BLOCK > UNIT_BLOCK_STACK_NUM - 1) || (nowNIdx + LOAB_BLOCK > nLoop - 1); + bool initMmad = blockStackIdx == 0; + bool pvCVItr = firstItr && kIdx == 0; + LayoutC layoutOTmpTemp(rowNum, embed, embed); + computePV(gA[gPOffset], gB[gBOffset], gC, gB[gBNextOffset], layoutA, layoutB, layoutOTmpTemp, + actualShape, actualNextShape, nowNIdx, nkBlockNextIdx, nkBlockLoop, firstItr, endItr, initMmad, pvCVItr, softmaxFlag, preloadFlag); + gPOffset += actualShape.k(); + ++nkBlockNextIdx; + nowMaskTailS = 0; + preloadFlag = false; + } + } + } + + CATLASS_DEVICE + void computePV( + AscendC::GlobalTensor const &gA, + AscendC::GlobalTensor const &gB, + AscendC::GlobalTensor const &gC, + AscendC::GlobalTensor const &gmNextBlockB, + LayoutA layoutA, LayoutB layoutB, LayoutC layoutC, + GemmCoord actualShape, GemmCoord actualNextShape, uint32_t nowIdx, uint32_t &nkblockIdx, + uint32_t &nkblockLoop, bool firstItr, bool endItr, bool initMmad, bool pvCVItr, Arch::CrossCoreFlag softmaxFlag, bool preloadFlag = false) + { + uint32_t MActual = actualShape.m(); + uint32_t kActual = actualShape.k(); + uint32_t nActual = actualShape.n(); + uint32_t mRound = RoundUp(MActual); + uint32_t kRound = RoundUp(kActual); + uint32_t nRound = RoundUp(nActual); + LayoutAInL1 layoutAInL1 = LayoutAInL1::template MakeLayout(mRound, kActual); + LayoutAInL0 layoutAInL0 = LayoutAInL0::template MakeLayout(mRound, kActual); + LayoutBInL1 layoutBInL1 = LayoutBInL1::template MakeLayout(kActual, nActual); + LayoutBInL0 layoutBInL0 = LayoutBInL0::template MakeLayout(kActual, nActual); + uint32_t l1KvPingPongFlag = nkblockIdx % 2; + uint32_t l0ABPingPongFlag = nkblockIdx % 2; + if (nkblockIdx == 1 || preloadFlag) { + auto layoutBTile = layoutB.GetTileLayout(MakeCoord(kActual, nActual)); + AscendC::WaitFlag(l1KvPingPongFlag + 2); + copyGmToL1B(l1BTensor[l1KvPingPongFlag], gB, layoutBInL1, layoutBTile); + AscendC::SetFlag(l1KvPingPongFlag + 2); + } + + AscendC::WaitFlag(l1KvPingPongFlag + 2); + AscendC::WaitFlag(2); + AscendC::WaitFlag(3); + copyL1ToL0B(l0BTensor, l1BTensor[l1KvPingPongFlag], layoutBInL0, layoutBInL1); + AscendC::SetFlag(l1KvPingPongFlag + 2); + + if (pvCVItr) { + Arch::CrossCoreWaitFlag(softmaxFlag); + } + AscendC::WaitFlag(l1KvPingPongFlag + 4); + auto layoutATile = layoutA.GetTileLayout(MakeCoord(MActual, kActual)); + copyGmToL1A(l1ATensor[l1KvPingPongFlag], gA, layoutAInL1, layoutATile); + AscendC::SetFlag(EVENT_ID5); + + if (nkblockIdx != nkblockLoop) { + uint32_t nNextActual = actualNextShape.n(); + uint32_t kNextActual = actualNextShape.k(); + LayoutBInL1 layoutBNextInL1 = LayoutBInL1::template MakeLayout(kNextActual, nNextActual); + auto layoutNextBTile = layoutB.GetTileLayout(MakeCoord(kNextActual, nNextActual)); + AscendC::WaitFlag(1 - l1KvPingPongFlag + 2); + copyGmToL1B(l1BTensor[1 - l1KvPingPongFlag], gmNextBlockB, layoutBNextInL1, layoutNextBTile); + AscendC::SetFlag(1 - l1KvPingPongFlag + 2); + } + + AscendC::WaitFlag(EVENT_ID5); + AscendC::WaitFlag(0); + AscendC::WaitFlag(1); + copyL1ToL0A(l0ATensor, l1ATensor[l1KvPingPongFlag], layoutAInL0, layoutAInL1); + AscendC::SetFlag(l1KvPingPongFlag + 4); + + AscendC::SetFlag(l0ABPingPongFlag); + AscendC::WaitFlag(l0ABPingPongFlag); + uint8_t unitFlag = 0b00; + if constexpr (!ENABLE_UNIT_FLAG_) { + if (firstItr) { + AscendC::WaitFlag(0); + AscendC::WaitFlag(1); + } + } else { + if (endItr) { + unitFlag = 0b11; + } else { + unitFlag = 0b10; + } + } + tileMmad( + l0CTensor[0], l0ATensor, l0BTensor, + mRound, nActual, kActual, initMmad, unitFlag); + // AscendC::PipeBarrier(); + AscendC::SetFlag(0); + AscendC::SetFlag(1); + AscendC::SetFlag(2); + AscendC::SetFlag(3); + auto layoutInL0C = LayoutCInL0::MakeLayoutInL0C(MakeCoord(MActual, nActual)); + if (endItr) { + if constexpr (!ENABLE_UNIT_FLAG_) { + AscendC::SetFlag(0); + AscendC::WaitFlag(0); + copyL0CToGm(gC, l0CTensor[0], layoutC, layoutInL0C); + AscendC::SetFlag(0); + AscendC::SetFlag(1); + } else { + copyL0CToGm(gC, l0CTensor[0], layoutC, layoutInL0C, 0b11); + } + } + } + +protected: + /// Data members + AscendC::LocalTensor l1ATensor[STAGES]; + AscendC::LocalTensor l1BTensor[STAGES]; + AscendC::LocalTensor l0ATensor; + AscendC::LocalTensor l0BTensor; + AscendC::LocalTensor l0CTensor[STAGES]; + + TileMmad tileMmad; + CopyGmToL1A copyGmToL1A; + CopyGmToL1B copyGmToL1B; + CopyL1ToL0A copyL1ToL0A; + CopyL1ToL0B copyL1ToL0B; + CopyL0CToGm copyL0CToGm; +}; + +//////////////////////////////////////////////////////////////////// + +} // namespace Catlass::Gemm::Block + +#endif // CATLASS_GEMM_BLOCK_MMAD_PV_TAIL_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_pv_head_tail_tla.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_pv_head_tail_tla.hpp new file mode 100644 index 0000000000..fc17f5fda4 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_pv_head_tail_tla.hpp @@ -0,0 +1,350 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_GEMM_BLOCK_MMAD_PV_TAIL_TLA_HPP +#define CATLASS_GEMM_BLOCK_MMAD_PV_TAIL_TLA_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/kda_coord.hpp" +#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" +#include "catlass/gemm/kda_gemm_helper.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" +#include "catlass/gemm/tile/kda_tile_mmad.hpp" +#include "tla/kda_tla_layout.hpp" +#include "tla/kda_tensor.hpp" + +//////////////////////////////////////////////////////////////////// + +namespace Catlass::Gemm::Block { +//////////////////////////////////////////////////////////////////// + +template < + class ArchTag_, + bool PAGED_CACHE_FLAG_, + bool ENABLE_UNIT_FLAG_, + class L1TileShape_, + class L0TileShape_, + class ElementA_, + class ElementB_, + class ElementC_, + class ElementBias_, + class TileCopy_, + class TileMmad_> +struct BlockMmadTla< + MmadFAITailPV, + L1TileShape_, + L0TileShape_, + ElementA_, + ElementB_, + ElementC_, + ElementBias_, + TileCopy_, + TileMmad_> { +public: + // Type Aliases + using DispatchPolicy = MmadFAITailPV; + using ArchTag = typename DispatchPolicy::ArchTag; + using L1TileShape = L1TileShape_; + using L0TileShape = L0TileShape_; + using ElementA = ElementA_; + using LayoutA = typename TileCopy_::LayoutA; + using ElementB = ElementB_; + using LayoutB = typename TileCopy_::LayoutB; + using ElementC = ElementC_; + using LayoutC = typename TileCopy_::LayoutC; + + using TileMmad = TileMmad_; + + using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; + + using ElementAccumulator = typename TileCopy_::ElementAccumulator; + + using LayoutTagL1A = typename TileCopy_::LayoutTagL1A; + using LayoutTagL1B = typename TileCopy_::LayoutTagL1B; + using LayoutTagL0A = typename TileCopy_::LayoutTagL0A; + using LayoutTagL0B = typename TileCopy_::LayoutTagL0B; + + using L1AAlignHelper = typename TileCopy_::L1AAlignHelper; + using L1BAlignHelper = typename TileCopy_::L1BAlignHelper; + + static_assert(tla::is_tuple::value && tla::is_static::value, + "L1TileShape must be tla::tuple and static!"); + static_assert(tla::is_tuple::value && tla::is_static::value, + "L0TileShape must be tla::tuple and static!"); + + static constexpr bool ENABLE_UNIT_FLAG = DispatchPolicy::ENABLE_UNIT_FLAG; + static constexpr uint32_t STAGES = DispatchPolicy::STAGES; + static constexpr uint32_t L1_TILE_M = tla::get<0>(L1TileShape{}); + static constexpr uint32_t L1_TILE_N = tla::get<1>(L1TileShape{}); + static constexpr uint32_t L1_TILE_K = tla::get<2>(L1TileShape{}); + static constexpr uint32_t L0_TILE_M = tla::get<0>(L0TileShape{}); + static constexpr uint32_t L0_TILE_N = tla::get<1>(L0TileShape{}); + static constexpr uint32_t L0_TILE_K = tla::get<2>(L0TileShape{}); + + // L1 tile size + static constexpr uint32_t L1A_TILE_SIZE = L1_TILE_M * L1_TILE_K * sizeof(ElementA); + static constexpr uint32_t L1B_TILE_SIZE = L1_TILE_N * L1_TILE_K * sizeof(ElementB); + // L0 tile size + static constexpr uint32_t L0A_TILE_SIZE = L0_TILE_M * L0_TILE_K * sizeof(ElementA); + static constexpr uint32_t L0B_TILE_SIZE = L0_TILE_K * L0_TILE_N * sizeof(ElementB); + static constexpr uint32_t L0C_TILE_SIZE = L0_TILE_M * L0_TILE_N * sizeof(ElementAccumulator); + static constexpr uint32_t L0A_PINGPONG_BUF_SIZE = ArchTag::L0A_SIZE / STAGES; + static constexpr uint32_t L0B_PINGPONG_BUF_SIZE = ArchTag::L0B_SIZE / STAGES; + static constexpr uint32_t L0C_PINGPONG_BUF_SIZE = ArchTag::L0C_SIZE / STAGES; + + // Check LayoutC + static_assert(tla::detail::isRowMajor::value || + ((std::is_same_v || std::is_same_v || + std::is_same_v) && + tla::detail::iszN::value), + "LayoutC only supports zN in half or bfloat16 or float, RowMajor in all dtype yet!"); + + // Check L1TileShape + static_assert((L1A_TILE_SIZE + L1B_TILE_SIZE) * STAGES <= ArchTag::L1_SIZE, + "L1TileShape exceeding the L1 space!"); + + // Check L0TileShape + static_assert(L0A_TILE_SIZE <= ArchTag::L0A_SIZE, "L0TileShape exceeding the L0A space!"); + static_assert(L0B_TILE_SIZE <= ArchTag::L0B_SIZE, "L0TileShape exceeding the L0B space!"); + static_assert(L0C_TILE_SIZE * STAGES <= ArchTag::L0C_SIZE, "L0TileShape exceeding the L0C space!"); + + static constexpr uint32_t BLOCK_SIZE = 16; + static constexpr uint32_t EMBED_SPLIT_SIZE = 128; + static constexpr uint32_t UNIT_BLOCK_STACK_NUM = 4; + static constexpr uint32_t KV_BASE_BLOCK = 512; + static constexpr uint32_t KV_SPLIT_SIZE = 128; + static constexpr uint32_t LOAB_BLOCK = 1; + + /// Construct + CATLASS_DEVICE + BlockMmadTla(Arch::Resource &resource, uint32_t l1BufAddrStart = 0) + { + // Allocate L1 memory space + l0ATensor = resource.l0ABuf.template GetBufferByByte(0); + l0BTensor = resource.l0BBuf.template GetBufferByByte(0); + for (uint32_t i = 0; i < STAGES; i++) { + l1ATensor[i] = resource.l1Buf.template GetBufferByByte(l1BufAddrStart + L1A_TILE_SIZE * LOAB_BLOCK * i); + l0CTensor[i] = resource.l0CBuf.template GetBufferByByte(L0C_PINGPONG_BUF_SIZE * i); + l1BTensor[i] = + resource.l1Buf.template GetBufferByByte(l1BufAddrStart + L1A_TILE_SIZE * 2 * 2 + L1B_TILE_SIZE * i); + } + } + + /// Destructor + CATLASS_DEVICE + ~BlockMmadTla() {} + + CATLASS_DEVICE + void getBlockShape(GemmCoord &actualShape, uint32_t &nowNIdx, uint32_t &kIdx, + uint32_t &nLoop, uint32_t &kLoop, uint32_t &kvSeqlen, uint32_t &embed, bool firstBlock, uint32_t maskTailS = 0) + { + uint32_t nSplitSize = KV_SPLIT_SIZE * LOAB_BLOCK; + uint32_t embedSplitSize = EMBED_SPLIT_SIZE; + if (nowNIdx + LOAB_BLOCK > nLoop - 1) { + nSplitSize = kvSeqlen - nowNIdx * KV_SPLIT_SIZE; + } + if (firstBlock && maskTailS != 0) { + nSplitSize = nSplitSize - maskTailS; + } + if (kIdx == kLoop - 1) { + embedSplitSize = embed - kIdx * EMBED_SPLIT_SIZE; + } + actualShape[1] = embedSplitSize; + actualShape[2] = nSplitSize; + } + + CATLASS_DEVICE + void getKVOffset(AscendC::GlobalTensor &gBlockTable, uint32_t &kOffset, uint32_t &nowNIdx, uint32_t &kIdx, + uint32_t &nLoop, uint32_t &kLoop, uint32_t &strideKV, uint32_t &blockSize, uint32_t maskTailS = 0) + { + if (nowNIdx >= nLoop || kIdx >= kLoop) { + kOffset = 0; + } + if constexpr (PAGED_CACHE_FLAG_) { + uint32_t blockTableId = gBlockTable.GetValue(nowNIdx); + kOffset = blockTableId * blockSize * strideKV + maskTailS * strideKV + kIdx * EMBED_SPLIT_SIZE; + } else { + kOffset = nowNIdx * KV_SPLIT_SIZE * strideKV + kIdx * EMBED_SPLIT_SIZE; + } + } + + template + CATLASS_DEVICE void operator()(TensorA &tensorA, TensorB &tensorB, TensorC &tensorC, + AscendC::GlobalTensor gBlockTable, GemmCoord actualOriShape, + uint32_t &nIdx, uint32_t &nLoop, uint32_t &blockSize, uint32_t kvSeqlen, uint32_t strideKV, Arch::CrossCoreFlag softmaxFlag, uint32_t maskTailS, bool preloadFlag) + { + uint32_t embed = actualOriShape[1]; + uint32_t kLoop = CeilDiv(embed); + uint32_t rowNum = tla::get<0>(tensorA.layout().shape()); + uint32_t blockN = tla::get<1>(tensorA.layout().shape()); + GemmCoord actualShape{rowNum, 0, 0}; + GemmCoord actualNextShape{rowNum, 0, 0}; + uint32_t nkBlockLoop = (nLoop + LOAB_BLOCK - 1) / LOAB_BLOCK * kLoop; // gap + uint32_t nkBlockNextIdx = (nIdx + LOAB_BLOCK - 1) / LOAB_BLOCK * kLoop + 1; // gap + uint32_t gBOffset = 0; + uint32_t gBNextOffset = 0; + uint32_t nowMaskTailS = 0; + uint32_t gPOffset = 0; + for (uint32_t kIdx = 0; kIdx < kLoop; kIdx++) { + nowMaskTailS = maskTailS; + gPOffset = 0; + for (uint32_t blockStackIdx = 0; (blockStackIdx < UNIT_BLOCK_STACK_NUM) && ((nIdx + blockStackIdx) < nLoop); + blockStackIdx += LOAB_BLOCK) { + uint32_t nowNIdx = nIdx + blockStackIdx; + uint32_t kLoopNextIdx = (nkBlockNextIdx % (kLoop * UNIT_BLOCK_STACK_NUM)) / (UNIT_BLOCK_STACK_NUM / LOAB_BLOCK); + uint32_t nLoopNextIdx = (nkBlockNextIdx % (kLoop * UNIT_BLOCK_STACK_NUM)) % (UNIT_BLOCK_STACK_NUM / LOAB_BLOCK) + nkBlockNextIdx / (kLoop * UNIT_BLOCK_STACK_NUM) * UNIT_BLOCK_STACK_NUM; + uint32_t startSeqOffset = nowNIdx == nIdx ? maskTailS : 0; + uint32_t startSeqNxtOffset = nLoopNextIdx == nIdx ? maskTailS : 0; + getBlockShape(actualShape, nowNIdx, kIdx, nLoop, kLoop, kvSeqlen, embed, nowNIdx == nIdx, nowMaskTailS); + getBlockShape(actualNextShape, nLoopNextIdx, kLoopNextIdx, nLoop, kLoop, kvSeqlen, embed, nLoopNextIdx == nIdx, nowMaskTailS); + getKVOffset(gBlockTable, gBOffset, nowNIdx, kIdx, nLoop, kLoop, strideKV, blockSize, startSeqOffset); + getKVOffset(gBlockTable, gBNextOffset, nLoopNextIdx, kLoopNextIdx, nLoop, kLoop, strideKV, blockSize, startSeqNxtOffset); + bool firstItr = blockStackIdx == 0; + bool endItr = (blockStackIdx + LOAB_BLOCK > UNIT_BLOCK_STACK_NUM - 1) || (nowNIdx + LOAB_BLOCK > nLoop - 1); + bool initMmad = blockStackIdx == 0; + bool pvCVItr = firstItr && kIdx == 0; + auto tensorAOffset = GetTile(tensorA, tla::MakeCoord(gPOffset / rowNum, gPOffset % rowNum), tensorA.layout().shape()); + computePV(tensorAOffset, tensorB, tensorC, gBOffset, gBNextOffset, + actualShape, actualNextShape, nowNIdx, nkBlockNextIdx, nkBlockLoop, strideKV, firstItr, endItr, initMmad, pvCVItr, softmaxFlag, preloadFlag); + gPOffset += actualShape.k(); + ++nkBlockNextIdx; + nowMaskTailS = 0; + preloadFlag = false; + } + } + } + + template + CATLASS_DEVICE void computePV( + TensorA &tensorA, TensorB &tensorB, TensorC &tensorC, uint32_t gBOffset, uint32_t gBNextOffset, + GemmCoord actualShape, GemmCoord actualNextShape, uint32_t nowIdx, uint32_t &nkblockIdx, + uint32_t &nkblockLoop, uint32_t strideKV, bool firstItr, bool endItr, bool initMmad, bool pvCVItr, Arch::CrossCoreFlag softmaxFlag, bool preloadFlag = false) + { + using CopyGmToL1A = typename TileCopy_::template CopyGmToL1A; + using CopyGmToL1B = typename TileCopy_::template CopyGmToL1B; + using CopyL0CToGm = typename TileCopy_::template CopyL0CToGm; + CopyGmToL1A copyGmToL1A; + CopyGmToL1B copyGmToL1B; + CopyL0CToGm copyL0CToGm; + + uint32_t MActual = actualShape.m(); + uint32_t kActual = actualShape.k(); + uint32_t nActual = actualShape.n(); + uint32_t mRound = RoundUp(MActual); + uint32_t kRound = RoundUp(kActual); + uint32_t nRound = RoundUp(nActual); + auto layoutAInL1 = tla::MakeLayout(mRound, kActual); + auto layoutAInL0 = tla::MakeLayout(mRound, kActual); + auto layoutBInL1 = tla::MakeLayout(kActual, nActual); + auto layoutBInL0 = tla::MakeLayout(kActual, nActual); + uint32_t l1KvPingPongFlag = nkblockIdx % 2; + uint32_t l0ABPingPongFlag = nkblockIdx % 2; + auto tensorTileA = GetTile(tensorA, tla::MakeCoord(0, 0), tla::MakeShape(MActual, kActual)); + auto tensorL1A = tla::MakeTensor(l1ATensor[l1KvPingPongFlag], layoutAInL1, Arch::PositionL1{}); + auto tensorL0A = tla::MakeTensor(l0ATensor, layoutAInL0, Arch::PositionL0A{}); + auto tensorTileB = GetTile(tensorB, tla::MakeCoord(gBOffset / strideKV, gBOffset % strideKV), tla::MakeShape(kActual, nActual)); + auto tensorL1B = tla::MakeTensor(l1BTensor[l1KvPingPongFlag], layoutBInL1, Arch::PositionL1{}); + auto tensorL0B = tla::MakeTensor(l0BTensor, layoutBInL0, Arch::PositionL0B{}); + if (nkblockIdx == 1 || preloadFlag) { + AscendC::WaitFlag(l1KvPingPongFlag + 2); + copyGmToL1B(tensorL1B, tensorTileB); + AscendC::SetFlag(l1KvPingPongFlag + 2); + } + + AscendC::WaitFlag(l1KvPingPongFlag + 2); + AscendC::WaitFlag(2); + AscendC::WaitFlag(3); + copyL1ToL0B(tensorL0B, tensorL1B); + AscendC::SetFlag(l1KvPingPongFlag + 2); + + if (pvCVItr) { + Arch::CrossCoreWaitFlag(softmaxFlag); + } + AscendC::WaitFlag(l1KvPingPongFlag + 4); + copyGmToL1A(tensorL1A, tensorTileA); + AscendC::SetFlag(EVENT_ID5); + + if (nkblockIdx != nkblockLoop) { + uint32_t nNextActual = actualNextShape.n(); + uint32_t kNextActual = actualNextShape.k(); + auto tensorTileBNext = GetTile(tensorB, tla::MakeCoord(gBNextOffset / strideKV, gBNextOffset % strideKV), tla::MakeShape(kNextActual, nNextActual)); + auto layoutBNextInL1 = tla::MakeLayout(kNextActual, nNextActual); + auto tensorL1BNext = tla::MakeTensor(l1BTensor[1 - l1KvPingPongFlag], layoutBNextInL1, Arch::PositionL1{}); + AscendC::WaitFlag(1 - l1KvPingPongFlag + 2); + copyGmToL1B(tensorL1BNext, tensorTileBNext); + AscendC::SetFlag(1 - l1KvPingPongFlag + 2); + } + + AscendC::WaitFlag(EVENT_ID5); + AscendC::WaitFlag(0); + AscendC::WaitFlag(1); + copyL1ToL0A(tensorL0A, tensorL1A); + AscendC::SetFlag(l1KvPingPongFlag + 4); + + AscendC::SetFlag(l0ABPingPongFlag); + AscendC::WaitFlag(l0ABPingPongFlag); + uint8_t unitFlag = 0b00; + if constexpr (!ENABLE_UNIT_FLAG_) { + if (firstItr) { + AscendC::WaitFlag(0); + AscendC::WaitFlag(1); + } + } else { + if (endItr) { + unitFlag = 0b11; + } else { + unitFlag = 0b10; + } + } + auto layoutCTileInL0 = tla::MakeLayoutL0C(mRound, nActual); + auto tensorL0CTile = tla::MakeTensor(l0CTensor[0], layoutCTileInL0, Arch::PositionL0C{}); + tileMmad( + tensorL0CTile, tensorL0A, tensorL0B, + mRound, nActual, kActual, initMmad, unitFlag); + AscendC::SetFlag(0); + AscendC::SetFlag(1); + AscendC::SetFlag(2); + AscendC::SetFlag(3); + auto layoutCInL0 = tla::MakeLayoutL0C(MActual, nActual); + auto tensorL0C = tla::MakeTensor(l0CTensor[0], layoutCInL0, Arch::PositionL0C{}); + if (endItr) { + if constexpr (!ENABLE_UNIT_FLAG_) { + AscendC::SetFlag(0); + AscendC::WaitFlag(0); + copyL0CToGm(tensorC, tensorL0C); + AscendC::SetFlag(0); + AscendC::SetFlag(1); + } else { + copyL0CToGm(tensorC, tensorL0C, 0b11); + } + } + } + +protected: + /// Data members + AscendC::LocalTensor l1ATensor[STAGES]; + AscendC::LocalTensor l1BTensor[STAGES]; + AscendC::LocalTensor l0ATensor; + AscendC::LocalTensor l0BTensor; + AscendC::LocalTensor l0CTensor[STAGES]; + + TileMmad tileMmad; + CopyL1ToL0A copyL1ToL0A; + CopyL1ToL0B copyL1ToL0B; +}; + +//////////////////////////////////////////////////////////////////// + +} // namespace Catlass::Gemm::Block + +#endif // CATLASS_GEMM_BLOCK_MMAD_PV_TAIL_TLA_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_pv_normal.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_pv_normal.hpp new file mode 100644 index 0000000000..d8d42f9489 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_pv_normal.hpp @@ -0,0 +1,305 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_GEMM_BLOCK_MMAD_PV_HPP +#define CATLASS_GEMM_BLOCK_MMAD_PV_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/kda_coord.hpp" +#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" +#include "catlass/gemm/kda_gemm_helper.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" +#include "catlass/gemm/tile/kda_tile_mmad.hpp" + +//////////////////////////////////////////////////////////////////// + +namespace Catlass::Gemm::Block { +//////////////////////////////////////////////////////////////////// + +template < + bool PAGED_CACHE_FLAG_, + bool ENABLE_UNIT_FLAG_, + class L1TileShape_, + class L0TileShape_, + class AType_, + class BType_, + class CType_, + class BiasType_, + class TileCopy_, + class TileMmad_> +struct BlockMmad< + MmadAtlasA2FAIPV, + L1TileShape_, + L0TileShape_, + AType_, + BType_, + CType_, + BiasType_, + TileCopy_, + TileMmad_> { +public: + // Type Aliases + using DispatchPolicy = MmadAtlasA2FAIPV; + using ArchTag = typename DispatchPolicy::ArchTag; + using L1TileShape = L1TileShape_; + using L0TileShape = L0TileShape_; + using ElementA = typename AType_::Element; + using LayoutA = typename AType_::Layout; + using ElementB = typename BType_::Element; + using LayoutB = typename BType_::Layout; + using ElementC = typename CType_::Element; + using LayoutC = typename CType_::Layout; + using TileMmad = TileMmad_; + using CopyGmToL1A = typename TileCopy_::CopyGmToL1A; + using CopyGmToL1B = typename TileCopy_::CopyGmToL1B; + using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; + using CopyL0CToGm = typename TileCopy_::CopyL0CToGm; + using ElementAccumulator = + typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; + using LayoutAInL1 = typename CopyL1ToL0A::LayoutSrc; + using LayoutBInL1 = typename CopyL1ToL0B::LayoutSrc; + using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst; + using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst; + using LayoutCInL0 = layout::zN; + + using L1AAlignHelper = Gemm::helper::L1AlignHelper; + using L1BAlignHelper = Gemm::helper::L1AlignHelper; + + static constexpr uint32_t STAGES = DispatchPolicy::STAGES; + static constexpr uint32_t L1A_SIZE = L1TileShape::M * L1TileShape::K * sizeof(ElementA); + static constexpr uint32_t L1B_SIZE = L1TileShape::N * L1TileShape::K * sizeof(ElementB); + static constexpr uint32_t L0A_SIZE = ArchTag::L0A_SIZE; + static constexpr uint32_t L0B_SIZE = ArchTag::L0B_SIZE; + static constexpr uint32_t L0C_SIZE = ArchTag::L0C_SIZE; + static constexpr uint32_t L0A_PINGPONG_BUF_SIZE = L0A_SIZE / STAGES; + static constexpr uint32_t L0B_PINGPONG_BUF_SIZE = L0B_SIZE / STAGES; + static constexpr uint32_t L0C_PINGPONG_BUF_SIZE = L0C_SIZE / STAGES; + static constexpr uint32_t BLOCK_SIZE = 16; + static constexpr uint32_t EMBED_SPLIT_SIZE = 128; + static constexpr uint32_t UNIT_BLOCK_STACK_NUM = 4; + static constexpr uint32_t KV_BASE_BLOCK = 512; + static constexpr uint32_t KV_SPLIT_SIZE = 128; + static constexpr uint32_t LOAB_BLOCK = 1; + + // Check LayoutC + static_assert(std::is_same_v, "LayoutC only support RowMajor yet!"); + + /// Construct + CATLASS_DEVICE + BlockMmad(Arch::Resource &resource, uint32_t l1BufAddrStart = 0) + { + // Allocate L1 memory space + l1ATensor = resource.l1Buf.template GetBufferByByte(l1BufAddrStart); + for (uint32_t i = 0; i < STAGES; i++) { + l0ATensor[i] = resource.l0ABuf.template GetBufferByByte(L0A_PINGPONG_BUF_SIZE * i); + l0BTensor[i] = resource.l0BBuf.template GetBufferByByte(L0B_PINGPONG_BUF_SIZE * i); + l0CTensor[i] = resource.l0CBuf.template GetBufferByByte(L0C_PINGPONG_BUF_SIZE * i); + l1BTensor[i] = + resource.l1Buf.template GetBufferByByte(l1BufAddrStart + L1A_SIZE * 2 * 2 + L1B_SIZE * i); + } + } + + /// Destructor + CATLASS_DEVICE + ~BlockMmad() {} + + CATLASS_DEVICE + void getBlockShape(GemmCoord &actualShape, uint32_t &nowNIdx, uint32_t &kIdx, + uint32_t &nLoop, uint32_t &kLoop, uint32_t &kvSeqlen, uint32_t &embed, bool firstBlock, uint32_t maskTailS = 0) + { + uint32_t nSplitSize = KV_SPLIT_SIZE * LOAB_BLOCK; + uint32_t embedSplitSize = EMBED_SPLIT_SIZE; + if (nowNIdx + LOAB_BLOCK > nLoop - 1) { + nSplitSize = kvSeqlen - nowNIdx * KV_SPLIT_SIZE; + } + if (firstBlock && maskTailS != 0) { + nSplitSize = nSplitSize - maskTailS; + } + if (kIdx == kLoop - 1) { + embedSplitSize = embed - kIdx * EMBED_SPLIT_SIZE; + } + actualShape[1] = embedSplitSize; + actualShape[2] = nSplitSize; + } + + CATLASS_DEVICE + void getKVOffset(AscendC::GlobalTensor &gBlockTable, uint32_t &kOffset, uint32_t &nowNIdx, uint32_t &kIdx, + uint32_t &nLoop, uint32_t &kLoop, uint32_t &strideKV, uint32_t &blockSize, uint32_t maskTailS = 0) + { + if (nowNIdx >= nLoop || kIdx >= kLoop) { + kOffset = 0; + } + if constexpr (PAGED_CACHE_FLAG_) { + uint32_t blockTableId = gBlockTable.GetValue(nowNIdx); + kOffset = blockTableId * blockSize * strideKV + maskTailS * strideKV + kIdx * EMBED_SPLIT_SIZE; + } else { + kOffset = nowNIdx * KV_SPLIT_SIZE * strideKV + kIdx * EMBED_SPLIT_SIZE; + } + } + + CATLASS_DEVICE + void operator()(AscendC::GlobalTensor gA, + AscendC::GlobalTensor gB, + AscendC::GlobalTensor gC, + AscendC::GlobalTensor gBlockTable, + LayoutA layoutA, LayoutB layoutB, GemmCoord actualOriShape, + uint32_t &nIdx, uint32_t &nLoop, uint32_t &blockSize, uint32_t kvSeqlen, uint32_t strideKV, Arch::CrossCoreFlag softmaxFlag) + { + // Arch::CrossCoreWaitFlag(softmaxFlag); + uint32_t embed = actualOriShape[1]; + uint32_t kLoop = CeilDiv(embed); + uint32_t rowNum = layoutA.shape(0); + uint32_t blockN = layoutA.shape(1); + GemmCoord actualShape{rowNum, 0, 0}; + GemmCoord actualNextShape{rowNum, 0, 0}; + uint32_t nkBlockLoop = (nLoop + LOAB_BLOCK - 1) / LOAB_BLOCK * kLoop; // gap + uint32_t nkBlockNextIdx = (nIdx + LOAB_BLOCK - 1) / LOAB_BLOCK * kLoop + 1; // gap + uint32_t gBOffset = 0; + uint32_t gBNextOffset = 0; + for (uint32_t kIdx = 0; kIdx < kLoop; kIdx++) { + for (uint32_t blockStackIdx = 0; (blockStackIdx < UNIT_BLOCK_STACK_NUM) && ((nIdx + blockStackIdx) < nLoop); + blockStackIdx += LOAB_BLOCK) { + uint32_t nowNIdx = nIdx + blockStackIdx; + uint32_t kLoopNextIdx = (nkBlockNextIdx % (kLoop * UNIT_BLOCK_STACK_NUM)) / (UNIT_BLOCK_STACK_NUM / LOAB_BLOCK); + uint32_t nLoopNextIdx = (nkBlockNextIdx % (kLoop * UNIT_BLOCK_STACK_NUM)) % (UNIT_BLOCK_STACK_NUM / LOAB_BLOCK) + nkBlockNextIdx / (kLoop * UNIT_BLOCK_STACK_NUM) * UNIT_BLOCK_STACK_NUM; + getBlockShape(actualShape, nowNIdx, kIdx, nLoop, kLoop, kvSeqlen, embed, nowNIdx == nIdx); + getBlockShape(actualNextShape, nLoopNextIdx, kLoopNextIdx, nLoop, kLoop, kvSeqlen, embed, nowNIdx == nIdx); + getKVOffset(gBlockTable, gBOffset, nowNIdx, kIdx, nLoop, kLoop, strideKV, blockSize); + getKVOffset(gBlockTable, gBNextOffset, nLoopNextIdx, kLoopNextIdx, nLoop, kLoop, strideKV, blockSize); + bool firstItr = blockStackIdx == 0; + bool endItr = (blockStackIdx + LOAB_BLOCK > UNIT_BLOCK_STACK_NUM - 1) || (nowNIdx + LOAB_BLOCK > nLoop - 1); + bool initMmad = blockStackIdx == 0; + bool pvCVItr = firstItr && kIdx == 0; + LayoutC layoutOTmpTemp(rowNum, embed, embed); + computePV(gA, gB[gBOffset], gC, gB[gBNextOffset], layoutA, layoutB, layoutOTmpTemp, + actualShape, actualNextShape, blockStackIdx, nkBlockNextIdx, nkBlockLoop, firstItr, endItr, initMmad, pvCVItr, softmaxFlag); + ++nkBlockNextIdx; + } + } + } + + CATLASS_DEVICE + void computePV( + AscendC::GlobalTensor const &gA, + AscendC::GlobalTensor const &gB, + AscendC::GlobalTensor const &gC, + AscendC::GlobalTensor const &gmNextBlockB, + LayoutA layoutA, LayoutB layoutB, LayoutC layoutC, + GemmCoord actualShape, GemmCoord actualNextShape, uint32_t nowIdx, uint32_t &nkblockIdx, + uint32_t &nkblockLoop, bool firstItr, bool endItr, bool initMmad, bool pvCVItr, Arch::CrossCoreFlag softmaxFlag, bool preloadFlag = false) + { + uint32_t MActual = actualShape.m(); + uint32_t kActual = actualShape.k(); + uint32_t nActual = actualShape.n(); + uint32_t mRound = RoundUp(MActual); + uint32_t kRound = RoundUp(kActual); + uint32_t nRound = RoundUp(nActual); + LayoutAInL1 layoutAInL1 = LayoutAInL1::template MakeLayout(mRound, (uint32_t)512); + LayoutAInL0 layoutAInL0 = LayoutAInL0::template MakeLayout(mRound, kActual); + LayoutBInL1 layoutBInL1 = LayoutBInL1::template MakeLayout(kActual, nActual); + LayoutBInL0 layoutBInL0 = LayoutBInL0::template MakeLayout(kActual, nActual); + uint32_t l1KvPingPongFlag = nkblockIdx % 2; + uint32_t l0ABPingPongFlag = nkblockIdx % 2; + if (nkblockIdx == 1 || preloadFlag) { + auto layoutBTile = layoutB.GetTileLayout(MakeCoord(kActual, nActual)); + AscendC::WaitFlag(l1KvPingPongFlag + 2); + copyGmToL1B(l1BTensor[l1KvPingPongFlag], gB, layoutBInL1, layoutBTile); + AscendC::SetFlag(l1KvPingPongFlag + 2); + } + + AscendC::WaitFlag(l1KvPingPongFlag + 2); + AscendC::WaitFlag(l0ABPingPongFlag + 2); + copyL1ToL0B(l0BTensor[l0ABPingPongFlag], l1BTensor[l1KvPingPongFlag], layoutBInL0, layoutBInL1); + AscendC::SetFlag(l1KvPingPongFlag + 2); + + if (pvCVItr) { + Arch::CrossCoreWaitFlag(softmaxFlag); + AscendC::WaitFlag(4); + AscendC::WaitFlag(5); + auto layoutATile = layoutA.GetTileLayout(MakeCoord(MActual, (uint32_t)512)); + copyGmToL1A(l1ATensor, gA, layoutAInL1, layoutATile); + AscendC::SetFlag(4); + AscendC::WaitFlag(4); + } + + if (nkblockIdx != nkblockLoop) { + uint32_t nNextActual = actualNextShape.n(); + uint32_t kNextActual = actualNextShape.k(); + LayoutBInL1 layoutBNextInL1 = LayoutBInL1::template MakeLayout(kNextActual, nNextActual); + auto layoutNextBTile = layoutB.GetTileLayout(MakeCoord(kNextActual, nNextActual)); + AscendC::WaitFlag(1 - l1KvPingPongFlag + 2); + copyGmToL1B(l1BTensor[1 - l1KvPingPongFlag], gmNextBlockB, layoutBNextInL1, layoutNextBTile); + AscendC::SetFlag(1 - l1KvPingPongFlag + 2); + } + + AscendC::WaitFlag(l0ABPingPongFlag); + copyL1ToL0A(l0ATensor[l0ABPingPongFlag], l1ATensor[mRound * 128 * nowIdx], layoutAInL0, layoutAInL1); + + AscendC::SetFlag(l0ABPingPongFlag); + AscendC::WaitFlag(l0ABPingPongFlag); + uint8_t unitFlag = 0b00; + if constexpr (!ENABLE_UNIT_FLAG_) { + if (firstItr) { + AscendC::WaitFlag(0); + AscendC::WaitFlag(1); + } + } else { + if (endItr) { + unitFlag = 0b11; + } else { + unitFlag = 0b10; + } + } + tileMmad( + l0CTensor[0], l0ATensor[l0ABPingPongFlag], l0BTensor[l0ABPingPongFlag], + mRound, nActual, kActual, initMmad, unitFlag); + // AscendC::PipeBarrier(); + AscendC::SetFlag(l0ABPingPongFlag); + AscendC::SetFlag(l0ABPingPongFlag + 2); + auto layoutInL0C = LayoutCInL0::MakeLayoutInL0C(MakeCoord(MActual, nActual)); + if (endItr) { + AscendC::SetFlag(4); + AscendC::SetFlag(5); + if constexpr (!ENABLE_UNIT_FLAG_) { + AscendC::SetFlag(0); + AscendC::WaitFlag(0); + copyL0CToGm(gC, l0CTensor[0], layoutC, layoutInL0C); + AscendC::SetFlag(0); + AscendC::SetFlag(1); + } else { + copyL0CToGm(gC, l0CTensor[0], layoutC, layoutInL0C, 0b11); + } + } + } + +protected: + /// Data members + AscendC::LocalTensor l1ATensor; + AscendC::LocalTensor l1BTensor[STAGES]; + AscendC::LocalTensor l0ATensor[STAGES]; + AscendC::LocalTensor l0BTensor[STAGES]; + AscendC::LocalTensor l0CTensor[STAGES]; + + TileMmad tileMmad; + CopyGmToL1A copyGmToL1A; + CopyGmToL1B copyGmToL1B; + CopyL1ToL0A copyL1ToL0A; + CopyL1ToL0B copyL1ToL0B; + CopyL0CToGm copyL0CToGm; +}; + +//////////////////////////////////////////////////////////////////// + +} // namespace Catlass::Gemm::Block + +#endif // CATLASS_GEMM_BLOCK_MMAD_PV_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_pv_normal_tla.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_pv_normal_tla.hpp new file mode 100644 index 0000000000..7f98bea467 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_pv_normal_tla.hpp @@ -0,0 +1,340 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_GEMM_BLOCK_MMAD_PV_TLA_HPP +#define CATLASS_GEMM_BLOCK_MMAD_PV_TLA_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/kda_coord.hpp" +#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" +#include "catlass/gemm/kda_gemm_helper.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" +#include "catlass/gemm/tile/kda_tile_mmad.hpp" +#include "tla/kda_tla_layout.hpp" +#include "tla/kda_tensor.hpp" + +//////////////////////////////////////////////////////////////////// + +namespace Catlass::Gemm::Block { +//////////////////////////////////////////////////////////////////// + +template < + class ArchTag_, + bool PAGED_CACHE_FLAG_, + bool ENABLE_UNIT_FLAG_, + class L1TileShape_, + class L0TileShape_, + class ElementA_, + class ElementB_, + class ElementC_, + class ElementBias_, + class TileCopy_, + class TileMmad_> +struct BlockMmadTla< + MmadFAIPV, + L1TileShape_, + L0TileShape_, + ElementA_, + ElementB_, + ElementC_, + ElementBias_, + TileCopy_, + TileMmad_> { +public: + // Type Aliases + using DispatchPolicy = MmadFAIPV; + using ArchTag = typename DispatchPolicy::ArchTag; + using L1TileShape = L1TileShape_; + using L0TileShape = L0TileShape_; + using ElementA = ElementA_; + using LayoutA = typename TileCopy_::LayoutA; + using ElementB = ElementB_; + using LayoutB = typename TileCopy_::LayoutB; + using ElementC = ElementC_; + using LayoutC = typename TileCopy_::LayoutC; + + using TileMmad = TileMmad_; + + using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; + + using ElementAccumulator = typename TileCopy_::ElementAccumulator; + + using LayoutTagL1A = typename TileCopy_::LayoutTagL1A; + using LayoutTagL1B = typename TileCopy_::LayoutTagL1B; + using LayoutTagL0A = typename TileCopy_::LayoutTagL0A; + using LayoutTagL0B = typename TileCopy_::LayoutTagL0B; + + using L1AAlignHelper = typename TileCopy_::L1AAlignHelper; + using L1BAlignHelper = typename TileCopy_::L1BAlignHelper; + + static_assert(tla::is_tuple::value && tla::is_static::value, + "L1TileShape must be tla::tuple and static!"); + static_assert(tla::is_tuple::value && tla::is_static::value, + "L0TileShape must be tla::tuple and static!"); + + static constexpr bool ENABLE_UNIT_FLAG = DispatchPolicy::ENABLE_UNIT_FLAG; + static constexpr uint32_t STAGES = DispatchPolicy::STAGES; + static constexpr uint32_t L1_TILE_M = tla::get<0>(L1TileShape{}); + static constexpr uint32_t L1_TILE_N = tla::get<1>(L1TileShape{}); + static constexpr uint32_t L1_TILE_K = tla::get<2>(L1TileShape{}); + static constexpr uint32_t L0_TILE_M = tla::get<0>(L0TileShape{}); + static constexpr uint32_t L0_TILE_N = tla::get<1>(L0TileShape{}); + static constexpr uint32_t L0_TILE_K = tla::get<2>(L0TileShape{}); + + // L1 tile size + static constexpr uint32_t L1A_TILE_SIZE = L1_TILE_M * L1_TILE_K * sizeof(ElementA); + static constexpr uint32_t L1B_TILE_SIZE = L1_TILE_N * L1_TILE_K * sizeof(ElementB); + // L0 tile size + static constexpr uint32_t L0A_TILE_SIZE = L0_TILE_M * L0_TILE_K * sizeof(ElementA); + static constexpr uint32_t L0B_TILE_SIZE = L0_TILE_K * L0_TILE_N * sizeof(ElementB); + static constexpr uint32_t L0C_TILE_SIZE = L0_TILE_M * L0_TILE_N * sizeof(ElementAccumulator); + static constexpr uint32_t L0A_PINGPONG_BUF_SIZE = ArchTag::L0A_SIZE / STAGES; + static constexpr uint32_t L0B_PINGPONG_BUF_SIZE = ArchTag::L0B_SIZE / STAGES; + static constexpr uint32_t L0C_PINGPONG_BUF_SIZE = ArchTag::L0C_SIZE / STAGES; + + // Check LayoutC + static_assert(tla::detail::isRowMajor::value || + ((std::is_same_v || std::is_same_v || + std::is_same_v) && + tla::detail::iszN::value), + "LayoutC only supports zN in half or bfloat16 or float, RowMajor in all dtype yet!"); + + // Check L1TileShape + static_assert(L1A_TILE_SIZE + L1B_TILE_SIZE * STAGES <= ArchTag::L1_SIZE, + "L1TileShape exceeding the L1 space!"); + + // Check L0TileShape + static_assert(L0A_TILE_SIZE * STAGES <= ArchTag::L0A_SIZE, "L0TileShape exceeding the L0A space!"); + static_assert(L0B_TILE_SIZE * STAGES <= ArchTag::L0B_SIZE, "L0TileShape exceeding the L0B space!"); + static_assert(L0C_TILE_SIZE * STAGES <= ArchTag::L0C_SIZE, "L0TileShape exceeding the L0C space!"); + + static constexpr uint32_t BLOCK_SIZE = 16; + static constexpr uint32_t EMBED_SPLIT_SIZE = 128; + static constexpr uint32_t UNIT_BLOCK_STACK_NUM = 4; + static constexpr uint32_t KV_BASE_BLOCK = 512; + static constexpr uint32_t KV_SPLIT_SIZE = 128; + static constexpr uint32_t LOAB_BLOCK = 1; + + /// Construct + CATLASS_DEVICE + BlockMmadTla(Arch::Resource &resource, uint32_t l1BufAddrStart = 0) + { + // Allocate L1 memory space + l1ATensor = resource.l1Buf.template GetBufferByByte(l1BufAddrStart); + for (uint32_t i = 0; i < STAGES; i++) { + l0ATensor[i] = resource.l0ABuf.template GetBufferByByte(L0A_PINGPONG_BUF_SIZE * i); + l0BTensor[i] = resource.l0BBuf.template GetBufferByByte(L0B_PINGPONG_BUF_SIZE * i); + l0CTensor[i] = resource.l0CBuf.template GetBufferByByte(L0C_PINGPONG_BUF_SIZE * i); + l1BTensor[i] = + resource.l1Buf.template GetBufferByByte(l1BufAddrStart + L1A_TILE_SIZE * 2 * 2 + L1B_TILE_SIZE * i); + } + } + + /// Destructor + CATLASS_DEVICE + ~BlockMmadTla() {} + + CATLASS_DEVICE + void getBlockShape(GemmCoord &actualShape, uint32_t &nowNIdx, uint32_t &kIdx, + uint32_t &nLoop, uint32_t &kLoop, uint32_t &kvSeqlen, uint32_t &embed, bool firstBlock, uint32_t maskTailS = 0) + { + uint32_t nSplitSize = KV_SPLIT_SIZE * LOAB_BLOCK; + uint32_t embedSplitSize = EMBED_SPLIT_SIZE; + if (nowNIdx + LOAB_BLOCK > nLoop - 1) { + nSplitSize = kvSeqlen - nowNIdx * KV_SPLIT_SIZE; + } + if (firstBlock && maskTailS != 0) { + nSplitSize = nSplitSize - maskTailS; + } + if (kIdx == kLoop - 1) { + embedSplitSize = embed - kIdx * EMBED_SPLIT_SIZE; + } + actualShape[1] = embedSplitSize; + actualShape[2] = nSplitSize; + } + + CATLASS_DEVICE + void getKVOffset(AscendC::GlobalTensor &gBlockTable, uint32_t &kOffset, uint32_t &nowNIdx, uint32_t &kIdx, + uint32_t &nLoop, uint32_t &kLoop, uint32_t &strideKV, uint32_t &blockSize, uint32_t maskTailS = 0) + { + if (nowNIdx >= nLoop || kIdx >= kLoop) { + kOffset = 0; + } + if constexpr (PAGED_CACHE_FLAG_) { + uint32_t blockTableId = gBlockTable.GetValue(nowNIdx); + kOffset = blockTableId * blockSize * strideKV + maskTailS * strideKV + kIdx * EMBED_SPLIT_SIZE; + } else { + kOffset = nowNIdx * KV_SPLIT_SIZE * strideKV + kIdx * EMBED_SPLIT_SIZE; + } + } + + template + CATLASS_DEVICE void operator()(TensorA &tensorA, TensorB &tensorB, TensorC &tensorC, + AscendC::GlobalTensor gBlockTable, GemmCoord actualOriShape, + uint32_t &nIdx, uint32_t &nLoop, uint32_t &blockSize, uint32_t kvSeqlen, uint32_t strideKV, Arch::CrossCoreFlag softmaxFlag) + { + // Arch::CrossCoreWaitFlag(softmaxFlag); + uint32_t embed = actualOriShape[1]; + uint32_t kLoop = CeilDiv(embed); + uint32_t rowNum = tla::get<0>(tensorA.layout().shape()); + uint32_t blockN = tla::get<1>(tensorA.layout().shape()); + GemmCoord actualShape{rowNum, 0, 0}; + GemmCoord actualNextShape{rowNum, 0, 0}; + uint32_t nkBlockLoop = (nLoop + LOAB_BLOCK - 1) / LOAB_BLOCK * kLoop; // gap + uint32_t nkBlockNextIdx = (nIdx + LOAB_BLOCK - 1) / LOAB_BLOCK * kLoop + 1; // gap + uint32_t gBOffset = 0; + uint32_t gBNextOffset = 0; + for (uint32_t kIdx = 0; kIdx < kLoop; kIdx++) { + for (uint32_t blockStackIdx = 0; (blockStackIdx < UNIT_BLOCK_STACK_NUM) && ((nIdx + blockStackIdx) < nLoop); + blockStackIdx += LOAB_BLOCK) { + uint32_t nowNIdx = nIdx + blockStackIdx; + uint32_t kLoopNextIdx = (nkBlockNextIdx % (kLoop * UNIT_BLOCK_STACK_NUM)) / (UNIT_BLOCK_STACK_NUM / LOAB_BLOCK); + uint32_t nLoopNextIdx = (nkBlockNextIdx % (kLoop * UNIT_BLOCK_STACK_NUM)) % (UNIT_BLOCK_STACK_NUM / LOAB_BLOCK) + nkBlockNextIdx / (kLoop * UNIT_BLOCK_STACK_NUM) * UNIT_BLOCK_STACK_NUM; + getBlockShape(actualShape, nowNIdx, kIdx, nLoop, kLoop, kvSeqlen, embed, nowNIdx == nIdx); + getBlockShape(actualNextShape, nLoopNextIdx, kLoopNextIdx, nLoop, kLoop, kvSeqlen, embed, nowNIdx == nIdx); + getKVOffset(gBlockTable, gBOffset, nowNIdx, kIdx, nLoop, kLoop, strideKV, blockSize); + getKVOffset(gBlockTable, gBNextOffset, nLoopNextIdx, kLoopNextIdx, nLoop, kLoop, strideKV, blockSize); + bool firstItr = blockStackIdx == 0; + bool endItr = (blockStackIdx + LOAB_BLOCK > UNIT_BLOCK_STACK_NUM - 1) || (nowNIdx + LOAB_BLOCK > nLoop - 1); + bool initMmad = blockStackIdx == 0; + bool pvCVItr = firstItr && kIdx == 0; + computePV(tensorA, tensorB, tensorC, gBOffset, gBNextOffset, + actualShape, actualNextShape, blockStackIdx, nkBlockNextIdx, nkBlockLoop, strideKV, firstItr, endItr, initMmad, pvCVItr, softmaxFlag); + ++nkBlockNextIdx; + } + } + } + + template + CATLASS_DEVICE void computePV( + TensorA &tensorA, TensorB &tensorB, TensorC &tensorC, uint32_t gBOffset, uint32_t gBNextOffset, + GemmCoord actualShape, GemmCoord actualNextShape, uint32_t nowIdx, uint32_t &nkblockIdx, + uint32_t &nkblockLoop, uint32_t strideKV, bool firstItr, bool endItr, bool initMmad, bool pvCVItr, Arch::CrossCoreFlag softmaxFlag, bool preloadFlag = false) + { + using CopyGmToL1A = typename TileCopy_::template CopyGmToL1A; + using CopyGmToL1B = typename TileCopy_::template CopyGmToL1B; + using CopyL0CToGm = typename TileCopy_::template CopyL0CToGm; + CopyGmToL1A copyGmToL1A; + CopyGmToL1B copyGmToL1B; + CopyL0CToGm copyL0CToGm; + + uint32_t MActual = actualShape.m(); + uint32_t kActual = actualShape.k(); + uint32_t nActual = actualShape.n(); + uint32_t mRound = RoundUp(MActual); + uint32_t kRound = RoundUp(kActual); + uint32_t nRound = RoundUp(nActual); + auto layoutAInL1 = tla::MakeLayout(mRound, (uint32_t)512); + auto layoutAInL0 = tla::MakeLayout(mRound, kActual); + auto layoutBInL1 = tla::MakeLayout(kActual, nActual); + auto layoutBInL0 = tla::MakeLayout(kActual, nActual); + uint32_t l1KvPingPongFlag = nkblockIdx % 2; + uint32_t l0ABPingPongFlag = nkblockIdx % 2; + auto tensorTileA = GetTile(tensorA, tla::MakeCoord(0, 0), tla::MakeShape(MActual, (uint32_t)512)); + auto tensorL1A = tla::MakeTensor(l1ATensor, layoutAInL1, Arch::PositionL1{}); + auto tensorL0A = tla::MakeTensor(l0ATensor[l0ABPingPongFlag], layoutAInL0, Arch::PositionL0A{}); + auto tensorTileB = GetTile(tensorB, tla::MakeCoord(gBOffset / strideKV, gBOffset % strideKV), tla::MakeShape(kActual, nActual)); + auto tensorL1B = tla::MakeTensor(l1BTensor[l1KvPingPongFlag], layoutBInL1, Arch::PositionL1{}); + auto tensorL0B = tla::MakeTensor(l0BTensor[l0ABPingPongFlag], layoutBInL0, Arch::PositionL0B{}); + if (nkblockIdx == 1 || preloadFlag) { + AscendC::WaitFlag(l1KvPingPongFlag + 2); + copyGmToL1B(tensorL1B, tensorTileB); + AscendC::SetFlag(l1KvPingPongFlag + 2); + } + + AscendC::WaitFlag(l1KvPingPongFlag + 2); + AscendC::WaitFlag(l0ABPingPongFlag + 2); + copyL1ToL0B(tensorL0B, tensorL1B); + AscendC::SetFlag(l1KvPingPongFlag + 2); + + if (pvCVItr) { + Arch::CrossCoreWaitFlag(softmaxFlag); + AscendC::WaitFlag(4); + AscendC::WaitFlag(5); + copyGmToL1A(tensorL1A, tensorTileA); + AscendC::SetFlag(4); + AscendC::WaitFlag(4); + } + + if (nkblockIdx != nkblockLoop) { + uint32_t nNextActual = actualNextShape.n(); + uint32_t kNextActual = actualNextShape.k(); + auto tensorTileBNext = GetTile(tensorB, tla::MakeCoord(gBNextOffset / strideKV, gBNextOffset % strideKV), tla::MakeShape(kNextActual, nNextActual)); + auto layoutBNextInL1 = tla::MakeLayout(kNextActual, nNextActual); + auto tensorL1BNext = tla::MakeTensor(l1BTensor[1 - l1KvPingPongFlag], layoutBNextInL1, Arch::PositionL1{}); + AscendC::WaitFlag(1 - l1KvPingPongFlag + 2); + copyGmToL1B(tensorL1BNext, tensorTileBNext); + AscendC::SetFlag(1 - l1KvPingPongFlag + 2); + } + + AscendC::WaitFlag(l0ABPingPongFlag); + auto tensorL1AOffset = tla::MakeTensor(l1ATensor[mRound * 128 * nowIdx], layoutAInL1, Arch::PositionL1{}); + copyL1ToL0A(tensorL0A, tensorL1AOffset); + + AscendC::SetFlag(l0ABPingPongFlag); + AscendC::WaitFlag(l0ABPingPongFlag); + uint8_t unitFlag = 0b00; + if constexpr (!ENABLE_UNIT_FLAG_) { + if (firstItr) { + AscendC::WaitFlag(0); + AscendC::WaitFlag(1); + } + } else { + if (endItr) { + unitFlag = 0b11; + } else { + unitFlag = 0b10; + } + } + auto layoutCTileInL0 = tla::MakeLayoutL0C(mRound, nActual); + auto tensorL0CTile = tla::MakeTensor(l0CTensor[0], layoutCTileInL0, Arch::PositionL0C{}); + tileMmad( + tensorL0CTile, tensorL0A, tensorL0B, + mRound, nActual, kActual, initMmad, unitFlag); + AscendC::SetFlag(l0ABPingPongFlag); + AscendC::SetFlag(l0ABPingPongFlag + 2); + auto layoutCInL0 = tla::MakeLayoutL0C(MActual, nActual); + auto tensorL0C = tla::MakeTensor(l0CTensor[0], layoutCInL0, Arch::PositionL0C{}); + if (endItr) { + AscendC::SetFlag(4); + AscendC::SetFlag(5); + if constexpr (!ENABLE_UNIT_FLAG_) { + AscendC::SetFlag(0); + AscendC::WaitFlag(0); + copyL0CToGm(tensorC, tensorL0C); + AscendC::SetFlag(0); + AscendC::SetFlag(1); + } else { + copyL0CToGm(tensorC, tensorL0C, 0b11); + } + } + } + +protected: + /// Data members + AscendC::LocalTensor l1ATensor; + AscendC::LocalTensor l1BTensor[STAGES]; + AscendC::LocalTensor l0ATensor[STAGES]; + AscendC::LocalTensor l0BTensor[STAGES]; + AscendC::LocalTensor l0CTensor[STAGES]; + + TileMmad tileMmad; + CopyL1ToL0A copyL1ToL0A; + CopyL1ToL0B copyL1ToL0B; +}; + +//////////////////////////////////////////////////////////////////// + +} // namespace Catlass::Gemm::Block + +#endif // CATLASS_GEMM_BLOCK_MMAD_PV_TLA_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_pv_tla.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_pv_tla.hpp new file mode 100644 index 0000000000..63066b8283 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_pv_tla.hpp @@ -0,0 +1,306 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ +#ifndef CATLASS_GEMM_BLOCK_BLOCK_MMAD_FAI_PV_TLA_HPP +#define CATLASS_GEMM_BLOCK_BLOCK_MMAD_FAI_PV_TLA_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/kda_coord.hpp" +#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" +#include "catlass/gemm/kda_gemm_helper.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" +#include "catlass/gemm/tile/kda_tile_mmad.hpp" +#include "tla/kda_tensor.hpp" +#include "tla/kda_tla_layout.hpp" + +//////////////////////////////////////////////////////////////////// + +namespace Catlass::Gemm::Block { +//////////////////////////////////////////////////////////////////// + +template < + bool PAGED_CACHE_FLAG_, + bool ENABLE_UNIT_FLAG_, + class L1TileShape_, + class L0TileShape_, + class ElementA_, + class ElementB_, + class ElementC_, + class ElementBias_, + class TileCopy_, + class TileMmad_> +struct BlockMmadTla< + MmadFAIPV, + L1TileShape_, + L0TileShape_, + ElementA_, + ElementB_, + ElementC_, + ElementBias_, + TileCopy_, + TileMmad_> { +public: + using DispatchPolicy = MmadFAIPV; + using ArchTag = typename DispatchPolicy::ArchTag; + using L1TileShape = L1TileShape_; + using L0TileShape = L0TileShape_; + using TileCopy = TileCopy_; + using TileMmad = TileMmad_; + + using ElementA = ElementA_; + using ElementB = ElementB_; + using ElementC = ElementC_; + using ElementAccumulator = typename TileCopy::ElementAccumulator; + + using CopyL1ToL0A = typename TileCopy::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopy::CopyL1ToL0B; + + using LayoutTagA = typename TileCopy::LayoutTagA; + using LayoutTagB = typename TileCopy::LayoutTagB; + using LayoutTagC = typename TileCopy::LayoutTagC; + + using LayoutTagL1A = typename TileCopy::LayoutTagL1A; + using LayoutTagL1B = typename TileCopy::LayoutTagL1B; + using LayoutTagL0A = typename TileCopy::LayoutTagL0A; + using LayoutTagL0B = typename TileCopy::LayoutTagL0B; + + static constexpr uint32_t STAGES = DispatchPolicy::STAGES; + // L1 tile shape + static constexpr uint32_t L1_TILE_M = tla::get<0>(L1TileShape{}); // s1 + static constexpr uint32_t L1_TILE_K = tla::get<1>(L1TileShape{}); // s2 + static constexpr uint32_t L1_TILE_N = tla::get<2>(L1TileShape{}); // d + // L0 tile shape + static constexpr uint32_t L0_TILE_M = tla::get<0>(L0TileShape{}); // s1 + static constexpr uint32_t L0_TILE_K = tla::get<1>(L0TileShape{}); // s2 + static constexpr uint32_t L0_TILE_N = tla::get<2>(L0TileShape{}); // d + // L1 tile size + static constexpr uint32_t L1A_TILE_SIZE = L1_TILE_M * L1_TILE_K * sizeof(ElementA); + static constexpr uint32_t L1B_TILE_SIZE = L1_TILE_N * L1_TILE_K * sizeof(ElementB); + // L0 tile size + static constexpr uint32_t L0A_TILE_SIZE = L0_TILE_M * L0_TILE_K * sizeof(ElementA); + static constexpr uint32_t L0B_TILE_SIZE = L0_TILE_K * L0_TILE_N * sizeof(ElementB); + static constexpr uint32_t L0C_TILE_SIZE = L0_TILE_M * L0_TILE_N * sizeof(ElementAccumulator); + + static constexpr uint32_t BLOCK_L1_SIZE = L1B_TILE_SIZE * STAGES; + static constexpr uint32_t BLOCK_L0C_SIZE = L0C_TILE_SIZE * STAGES; + + // Check L1/L0TileShape + static_assert(L1_TILE_M == L0_TILE_M && L1_TILE_N == L0_TILE_N && L1_TILE_K == L0_TILE_K, + "The situation where the basic blocks of L1 and L0 differ on the m, n, k axes is not supported yet"); + static_assert(L0A_TILE_SIZE * STAGES <= ArchTag::L0A_SIZE, "L0TileShape exceeding the L0A space!"); + static_assert(L0B_TILE_SIZE * STAGES <= ArchTag::L0B_SIZE, "L0TileShape exceeding the L0B space!"); + +public: + /// Construct + CATLASS_DEVICE + BlockMmadTla(Arch::Resource &resource, uint32_t &l1BufAddrStart, uint32_t &l0CBufAddrStart) + { + uint32_t l1BOffset = l1BufAddrStart; + uint32_t l0COffset = l0CBufAddrStart; + + for (uint32_t i = 0; i < STAGES; i++) { + // Assign L1/L0A/L0B space for each stages + l1BTensorList_[i] = resource.l1Buf.template GetBufferByByte(l1BOffset + L1B_TILE_SIZE * i); + l0ATensorList_[i] = resource.l0ABuf.template GetBufferByByte(L0A_TILE_SIZE * i); + l0BTensorList_[i] = resource.l0BBuf.template GetBufferByByte(L0B_TILE_SIZE * i); + l0CTensorList_[i] = resource.l0CBuf.template GetBufferByByte(l0COffset + L0C_TILE_SIZE * i); + + // Assign event ID for each stages + l1BEventList_[i] = BLOCK_EVENT_ID + i + STAGES; + l0AEventList_[i] = i; + l0BEventList_[i] = i + STAGES; + l0CEventList_[i] = BLOCK_EVENT_ID + i; + + // The event id that needs to be set before the loop + AscendC::SetFlag(l1BEventList_[i]); + AscendC::SetFlag(l0CEventList_[i]); + } + + l1BufAddrStart += BLOCK_L1_SIZE; + l0CBufAddrStart += BLOCK_L0C_SIZE; + } + + /// Destructor + CATLASS_DEVICE + ~BlockMmadTla() + { + for (uint32_t i = 0; i < STAGES; i++) { + AscendC::WaitFlag(l1BEventList_[i]); + AscendC::WaitFlag(l0CEventList_[i]); + } + } + +public: + template + CATLASS_DEVICE void operator()( + TensorA &tensorA, + TensorB &tensorB, + TensorC &tensorC, + AscendC::GlobalTensor &blockTable, + Shape &actualShape, + int64_t taskId, int32_t blockSize) + { + int32_t blockM = tla::get<0>(actualShape); + int32_t blockN = tla::get<1>(actualShape); + int32_t blockK = tla::get<2>(actualShape); + + // load matrix B tile from GM to L1 + auto layoutBInL1 = tla::MakeLayout(blockK, blockN); + auto tensorL1B = tla::MakeTensor(l1BTensorList_[l1BListId_], layoutBInL1, Arch::PositionL1{}); + AscendC::WaitFlag(l1BEventList_[l1BListId_]); + CopyInL1B(tensorL1B, tensorB, blockTable, blockK, blockN, blockSize); + AscendC::SetFlag(l1BEventList_[l1BListId_]); + AscendC::WaitFlag(l1BEventList_[l1BListId_]); + + AscendC::WaitFlag(l0CEventList_[l0CListId_]); + auto layoutInL0C = tla::MakeLayoutL0C(blockM, blockN); + auto tensorL0C = tla::MakeTensor(l0CTensorList_[l0CListId_], layoutInL0C, Arch::PositionL0C{}); + + uint32_t nLoops = (blockN + L0_TILE_N - 1) / L0_TILE_N; + uint32_t tailSize = blockN % L0_TILE_N; + uint32_t tailN = tailSize ? tailSize : L0_TILE_N; + + AscendC::WaitFlag(l0AEventList_[l0AListId_]); + for (uint32_t nIdx = 0; nIdx < nLoops; nIdx++) { + uint32_t tileN = (nIdx == (nLoops - 1)) ? tailN : L0_TILE_N; + + // load matrix A tile from L1 to L0A + auto layoutAInL0 = tla::MakeLayout(blockM, blockK); + auto tensorL0A = tla::MakeTensor(l0ATensorList_[l0AListId_], layoutAInL0, Arch::PositionL0A{}); + if (nIdx == 0) { + auto tensorL1TileA = GetTile( + tensorA, + tla::MakeCoord(0, 0), + tla::MakeShape(blockM, blockK)); + copyL1ToL0A(tensorL0A, tensorL1TileA); + } + + // load matrix B tile from L1 to L0B + auto layoutBInL0 = tla::MakeLayout(blockK, tileN); + auto tensorL0B = tla::MakeTensor(l0BTensorList_[l0BListId_], layoutBInL0, Arch::PositionL0B{}); + auto tensorL1TileB = GetTile( + tensorL1B, + tla::MakeCoord(0, nIdx * L0_TILE_N), + tla::MakeShape(blockK, tileN)); + AscendC::WaitFlag(l0BEventList_[l0BListId_]); + copyL1ToL0B(tensorL0B, tensorL1TileB); + + AscendC::SetFlag(l0CEventList_[l0CListId_]); + AscendC::WaitFlag(l0CEventList_[l0CListId_]); + + // Mmad + bool initC = true; + uint32_t tileM = (blockM == 1 ? M_ALIGN : blockM); + uint32_t tileK = blockK; + auto tensorTileL0C = GetTile( + tensorL0C, + tla::MakeCoord(0, nIdx * L0_TILE_N), + tla::MakeShape(tileM, tileN)); + tileMmad(tensorTileL0C, tensorL0A, tensorL0B, tileM, tileN, tileK, initC); + + AscendC::SetFlag(l0BEventList_[l0BListId_]); + l0BListId_ = (l0BListId_ + 1 < STAGES) ? (l0BListId_ + 1) : 0; + } + AscendC::SetFlag(l0AEventList_[l0AListId_]); + l0AListId_ = (l0AListId_ + 1 < STAGES) ? (l0AListId_ + 1) : 0; + + AscendC::SetFlag(l1BEventList_[l1BListId_]); + l1BListId_ = (l1BListId_ + 1 < STAGES) ? (l1BListId_ + 1) : 0; + + AscendC::SetFlag(l0CEventList_[l0CListId_]); + AscendC::WaitFlag(l0CEventList_[l0CListId_]); + + AscendC::CrossCoreWaitFlag(MM2_RES_INTRA_EVENT[taskId]); + AscendC::CrossCoreWaitFlag(16 + MM2_RES_INTRA_EVENT[taskId]); + + // copy block out + using CopyL0CToDst = typename TileCopy::template CopyL0CToDst; + CopyL0CToDst copyL0CToDst; + copyL0CToDst(tensorC, tensorL0C); + + AscendC::SetFlag(l0CEventList_[l0CListId_]); + l0CListId_ = (l0CListId_ + 1 < STAGES) ? (l0CListId_ + 1) : 0; + } + +private: + template + CATLASS_DEVICE void CopyInL1B( + TensorL1B &tensorL1B, + TensorB &tensorB, + AscendC::GlobalTensor &blockTable, + int32_t tileK, int32_t tileN, int32_t blockSize) + { + using CopyGmToL1B = typename TileCopy::template CopyGmToL1B; + CopyGmToL1B copyGmToL1B; + + if constexpr (PAGED_CACHE_FLAG_) { + int32_t copyRowCnt = 0; + int32_t blockLoops = (tileK + blockSize - 1) / blockSize; + + for (int32_t blockIdx = 0; blockIdx < blockLoops; ++blockIdx) { + int32_t curCopyRow = blockIdx == blockLoops - 1 ? tileK - copyRowCnt : blockSize; + int32_t idInBlockTable = blockTable.GetValue(blockIdx); + + auto tensorL1TileB = GetTile( + tensorL1B, + tla::MakeCoord(copyRowCnt, 0), + tla::MakeShape(curCopyRow, tileN)); + + auto tensorTileB = GetTile( + tensorB, + tla::MakeCoord(idInBlockTable * blockSize, 0), + tla::MakeShape(curCopyRow, tileN)); + + copyGmToL1B(tensorL1TileB, tensorTileB); + copyRowCnt += curCopyRow; + } + } else { + auto tensorTileB = GetTile( + tensorB, + tla::MakeCoord(0, 0), + tla::MakeShape(tileK, tileN)); + + copyGmToL1B(tensorL1B, tensorTileB); + } + } + +private: + static constexpr uint16_t M_ALIGN = 16; + static constexpr uint32_t SYNC_MODE = 4; + static constexpr uint32_t MM2_RES_INTRA_EVENT[2] = {7, 8}; + static constexpr uint32_t BLOCK_EVENT_ID = 4; + + // Multi-stage tensors list + AscendC::LocalTensor l1ATensorList_[STAGES]; + AscendC::LocalTensor l1BTensorList_[STAGES]; + AscendC::LocalTensor l0ATensorList_[STAGES]; + AscendC::LocalTensor l0BTensorList_[STAGES]; + AscendC::LocalTensor l0CTensorList_[STAGES]; + + // Multi-stage event id list + int32_t l1AEventList_[STAGES]; + int32_t l1BEventList_[STAGES]; + int32_t l0AEventList_[STAGES]; + int32_t l0BEventList_[STAGES]; + int32_t l0CEventList_[STAGES]; + + uint32_t l1AListId_{0}; + uint32_t l1BListId_{0}; + uint32_t l0AListId_{0}; + uint32_t l0BListId_{0}; + uint32_t l0CListId_{0}; + + CopyL1ToL0A copyL1ToL0A; + CopyL1ToL0B copyL1ToL0B; + TileMmad tileMmad; +}; +} // namespace Catlass::Gemm::Block +#endif diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_qk_head_tail.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_qk_head_tail.hpp new file mode 100644 index 0000000000..08971b249f --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_qk_head_tail.hpp @@ -0,0 +1,306 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_GEMM_BLOCK_MMAD_QK_TAIL_HPP +#define CATLASS_GEMM_BLOCK_MMAD_QK_TAIL_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/kda_coord.hpp" +#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" +#include "catlass/gemm/kda_gemm_helper.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" +#include "catlass/gemm/tile/kda_tile_mmad.hpp" + +//////////////////////////////////////////////////////////////////// + +namespace Catlass::Gemm::Block { +//////////////////////////////////////////////////////////////////// + +template < + bool PAGED_CACHE_FLAG_, + bool ENABLE_UNIT_FLAG_, + class L1TileShape_, + class L0TileShape_, + class AType_, + class BType_, + class CType_, + class BiasType_, + class TileCopy_, + class TileMmad_> +struct BlockMmad< + MmadAtlasA2FAITailQK, + L1TileShape_, + L0TileShape_, + AType_, + BType_, + CType_, + BiasType_, + TileCopy_, + TileMmad_> { +public: + // Type Aliases + using DispatchPolicy = MmadAtlasA2FAITailQK; + using ArchTag = typename DispatchPolicy::ArchTag; + using L1TileShape = L1TileShape_; + using L0TileShape = L0TileShape_; + using ElementA = typename AType_::Element; + using LayoutA = typename AType_::Layout; + using ElementB = typename BType_::Element; + using LayoutB = typename BType_::Layout; + using ElementC = typename CType_::Element; + using LayoutC = typename CType_::Layout; + using TileMmad = TileMmad_; + using CopyGmToL1A = typename TileCopy_::CopyGmToL1A; + using CopyGmToL1B = typename TileCopy_::CopyGmToL1B; + using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; + using CopyL0CToGm = typename TileCopy_::CopyL0CToGm; + using ElementAccumulator = + typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; + using LayoutAInL1 = typename CopyL1ToL0A::LayoutSrc; + using LayoutBInL1 = typename CopyL1ToL0B::LayoutSrc; + using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst; + using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst; + using LayoutCInL0 = layout::zN; + + using L1AAlignHelper = Gemm::helper::L1AlignHelper; + using L1BAlignHelper = Gemm::helper::L1AlignHelper; + + static constexpr uint32_t STAGES = DispatchPolicy::STAGES; + static constexpr uint32_t L1A_SIZE = L1TileShape::M * L1TileShape::K * sizeof(ElementA); + static constexpr uint32_t L1B_SIZE = L1TileShape::N * L1TileShape::K * sizeof(ElementB); + static constexpr uint32_t L0A_SIZE = ArchTag::L0A_SIZE; + static constexpr uint32_t L0B_SIZE = ArchTag::L0B_SIZE; + static constexpr uint32_t L0C_SIZE = ArchTag::L0C_SIZE; + static constexpr uint32_t L0A_PINGPONG_BUF_SIZE = L0A_SIZE / STAGES; + static constexpr uint32_t L0B_PINGPONG_BUF_SIZE = L0B_SIZE / STAGES; + static constexpr uint32_t L0C_PINGPONG_BUF_SIZE = L0C_SIZE / STAGES; + static constexpr uint32_t BLOCK_SIZE = 16; + static constexpr uint32_t EMBED_SPLIT_SIZE = 128; + static constexpr uint32_t UNIT_BLOCK_STACK_NUM = 4; + static constexpr uint32_t KV_BASE_BLOCK = 512; + static constexpr uint32_t KV_SPLIT_SIZE = 128; + + static_assert(std::is_same_v, "LayoutC only support RowMajor yet!"); + + CATLASS_DEVICE + BlockMmad(Arch::Resource &resource, uint32_t l1BufAddrStart = 0) + { + // Allocate L1 memory space + l1ATensor = resource.l1Buf.template GetBufferByByte(l1BufAddrStart); + for (uint32_t i = 0; i < STAGES; i++) { + l1BTensor[i] = + resource.l1Buf.template GetBufferByByte(l1BufAddrStart + L1A_SIZE + L1B_SIZE * i); + l0ATensor[i] = resource.l0ABuf.template GetBufferByByte(L0A_PINGPONG_BUF_SIZE * i); + l0BTensor[i] = resource.l0BBuf.template GetBufferByByte(L0B_PINGPONG_BUF_SIZE * i); + } + l0CTensor = resource.l0CBuf.template GetBufferByByte(0); + } + + /// Destructor + CATLASS_DEVICE + ~BlockMmad() {} + + CATLASS_DEVICE + void loadQGM(AscendC::GlobalTensor gA, LayoutA layoutA, uint32_t rowNum, uint32_t &singleGroupHeads, uint32_t &qHeads) + { + uint32_t embed = layoutA.shape(1); + uint32_t rowNumRound = RoundUp(rowNum); + uint32_t tokenNumPerGroup = rowNum / singleGroupHeads; + auto layoutSingleANd = layoutA.GetTileLayout(MakeCoord(singleGroupHeads, embed)); + LayoutAInL1 layoutAInL1 = LayoutAInL1::template MakeLayout(rowNum, embed); + copyGmToL1A( + l1ATensor, gA, + layoutAInL1, layoutSingleANd, + tokenNumPerGroup, qHeads * embed, tokenNumPerGroup, BLOCK_SIZE, rowNumRound); + AscendC::SetFlag(EVENT_ID3); + AscendC::WaitFlag(EVENT_ID3); + } + + CATLASS_DEVICE + void getBlockShape(GemmCoord &actualShape, uint32_t &nowNIdx, uint32_t &kIdx, + uint32_t &nLoop, uint32_t &kLoop, uint32_t &kvSeqlen, uint32_t &embed, bool firstBlock, uint32_t maskTailS = 0) + { + uint32_t nSplitSize = KV_SPLIT_SIZE; + uint32_t embedSplitSize = EMBED_SPLIT_SIZE; + if (nowNIdx == nLoop - 1) { + nSplitSize = kvSeqlen - nowNIdx * KV_SPLIT_SIZE; + } + if (firstBlock && maskTailS != 0) { + nSplitSize = nSplitSize - maskTailS; + } + if (kIdx == kLoop - 1) { + embedSplitSize = embed - kIdx * EMBED_SPLIT_SIZE; + } + actualShape[1] = nSplitSize; + actualShape[2] = embedSplitSize; + } + + CATLASS_DEVICE + void getKVOffset(AscendC::GlobalTensor &gBlockTable, uint32_t &kOffset, uint32_t &nowNIdx, uint32_t &kIdx, + uint32_t &nLoop, uint32_t &kLoop, uint32_t &strideKV, uint32_t &blockSize, uint32_t maskTailS = 0) + { + if (nowNIdx >= nLoop || kIdx >= kLoop) { + kOffset = 0; + } + if constexpr (PAGED_CACHE_FLAG_) { + uint32_t blockTableId = gBlockTable.GetValue(nowNIdx); + kOffset = blockTableId * blockSize * strideKV + maskTailS * strideKV + kIdx * EMBED_SPLIT_SIZE; + } else { + kOffset = nowNIdx * KV_SPLIT_SIZE * strideKV + kIdx * EMBED_SPLIT_SIZE; + } + } + + CATLASS_DEVICE + void operator()(AscendC::GlobalTensor gA, + AscendC::GlobalTensor gB, + AscendC::GlobalTensor gC, + AscendC::GlobalTensor gBlockTable, + LayoutA layoutA, LayoutB layoutB, GemmCoord actualOriShape, + uint32_t &nIdx, uint32_t &nLoop, uint32_t &blockSize, uint32_t kvSeqlen, uint32_t strideKV, uint32_t maskTailS, bool preloadFlag) + { + uint32_t rowNum = actualOriShape[0]; + uint32_t embed = actualOriShape[2]; + uint32_t kLoop = CeilDiv(embed); + uint32_t nkBlockLoop = nLoop * kLoop; + GemmCoord actualShape{rowNum, 0, 0}; + GemmCoord actualNextShape{rowNum, 0, 0}; + uint32_t nkBlockNextIdx = nIdx * kLoop + 1; + uint32_t gBOffset = 0; + uint32_t gBNextOffset = 0; + uint32_t stackTile = 0; + uint32_t preStackTile = 0; + for (uint32_t blockStackIdx = 0; (blockStackIdx < UNIT_BLOCK_STACK_NUM) && ((nIdx + blockStackIdx) < nLoop); + ++blockStackIdx) { + for (uint32_t kIdx = 0; kIdx < kLoop; kIdx++) { + uint32_t nowNIdx = nIdx + blockStackIdx; + uint32_t nLoopNextIdx = nkBlockNextIdx / kLoop; + uint32_t kLoopNextIdx = nkBlockNextIdx % kLoop; + uint32_t gCOffset = stackTile; + uint32_t startSeqOffset = nowNIdx == nIdx ? maskTailS : 0; + uint32_t startSeqNxtOffset = nLoopNextIdx == nIdx ? maskTailS : 0; + getBlockShape(actualShape, nowNIdx, kIdx, nLoop, kLoop, kvSeqlen, embed, nowNIdx == nIdx, maskTailS); + getBlockShape(actualNextShape, nLoopNextIdx, kLoopNextIdx, nLoop, kLoop, kvSeqlen, embed, nLoopNextIdx == nIdx, maskTailS); + getKVOffset(gBlockTable, gBOffset, nowNIdx, kIdx, nLoop, kLoop, strideKV, blockSize, startSeqOffset); + getKVOffset(gBlockTable, gBNextOffset, nLoopNextIdx, kLoopNextIdx, nLoop, kLoop, strideKV, blockSize, startSeqNxtOffset); + bool firstItr = kIdx == 0; + bool endItr = kIdx == kLoop - 1; + bool firstQtr = blockStackIdx == 0; + bool endQItr = ((nowNIdx == nLoop - 1) || (blockStackIdx == UNIT_BLOCK_STACK_NUM - 1)) && (kIdx == kLoop - 1); + bool initMmad = kIdx == 0; + preStackTile = stackTile; + stackTile += actualShape[1]; + LayoutC layOutSTemp(rowNum, stackTile, 512); + computeQK(gA, gB[gBOffset], gC[gCOffset], gB[gBNextOffset], layoutA, layoutB, layOutSTemp, + actualShape, actualNextShape, blockStackIdx, nkBlockNextIdx, nkBlockLoop, firstItr, endItr, initMmad, firstQtr, endQItr, preloadFlag); + ++nkBlockNextIdx; + preloadFlag = false; + } + maskTailS = 0; + } + } + + CATLASS_DEVICE void computeQK( + AscendC::GlobalTensor const &gA, + AscendC::GlobalTensor const &gB, + AscendC::GlobalTensor const &gC, + AscendC::GlobalTensor const &gmNextBlockB, + LayoutA layoutA, LayoutB layoutB, LayoutC layoutC, + GemmCoord actualShape, GemmCoord actualNextShape, uint32_t nowIdx, uint32_t &nkblockIdx, + uint32_t &nkblockLoop, bool firstItr, bool endItr, bool initMmad, bool firstQItr, bool endQItr, bool preloadFlag) + { + uint32_t mActual = actualShape.m(); + uint32_t kActual = actualShape.k(); + uint32_t nActual = actualShape.n(); + uint32_t mRound = RoundUp(mActual); + uint32_t kRound = RoundUp(kActual); + uint32_t nRound = RoundUp(nActual); + LayoutAInL1 layoutAInL1 = LayoutAInL1::template MakeLayout(mRound, kActual); // embed + LayoutAInL0 layoutAInL0 = LayoutAInL0::template MakeLayout(mRound, kActual); + LayoutBInL1 layoutBInL1 = LayoutBInL1::template MakeLayout(kActual, nActual); + LayoutBInL0 layoutBInL0 = LayoutBInL0::template MakeLayout(kActual, nRound); + uint32_t locPingPongFlag = nkblockIdx % 2; + uint32_t l1KvPingPongFlag = nkblockIdx % 2; + uint32_t l0ABPingPongFlag = nkblockIdx % 2; + if (preloadFlag) { + auto layoutBTile = layoutB.GetTileLayout(MakeCoord(kActual, nActual)); + AscendC::WaitFlag(l1KvPingPongFlag); + copyGmToL1B(l1BTensor[l1KvPingPongFlag], gB, layoutBInL1, layoutBTile); + AscendC::SetFlag(l1KvPingPongFlag); + } + if (nkblockIdx != nkblockLoop) { + uint32_t nNextActual = actualNextShape.n(); + uint32_t kNextActual = actualNextShape.k(); + LayoutBInL1 layoutBNextInL1 = LayoutBInL1::template MakeLayout(kNextActual, nNextActual); + auto layoutNextBTile = layoutB.GetTileLayout(MakeCoord(kNextActual, nNextActual)); + AscendC::WaitFlag(1 - l1KvPingPongFlag); + copyGmToL1B(l1BTensor[1 - l1KvPingPongFlag], gmNextBlockB, layoutBNextInL1, layoutNextBTile); + AscendC::SetFlag(1 - l1KvPingPongFlag); + } + AscendC::WaitFlag(l0ABPingPongFlag); + copyL1ToL0A(l0ATensor[l0ABPingPongFlag], l1ATensor, layoutAInL0, layoutAInL1); + + AscendC::WaitFlag(l1KvPingPongFlag); + AscendC::WaitFlag(l0ABPingPongFlag + 2); + copyL1ToL0B(l0BTensor[l0ABPingPongFlag], l1BTensor[l1KvPingPongFlag], layoutBInL0, layoutBInL1); + AscendC::SetFlag(l1KvPingPongFlag); + + AscendC::SetFlag(l0ABPingPongFlag); + AscendC::WaitFlag(l0ABPingPongFlag); + uint8_t unitFlag = 0b00; + if constexpr (!ENABLE_UNIT_FLAG_) { + if (firstItr) { + AscendC::WaitFlag(locPingPongFlag); + } + } else { + unitFlag = 0b11; + } + tileMmad( + l0CTensor[locPingPongFlag * mRound * 128], l0ATensor[l0ABPingPongFlag], l0BTensor[l0ABPingPongFlag], + mRound, nActual, kActual, initMmad, unitFlag); + // AscendC::PipeBarrier(); + AscendC::SetFlag(l0ABPingPongFlag); + AscendC::SetFlag(l0ABPingPongFlag + 2); + auto layoutInL0C = LayoutCInL0::MakeLayoutInL0C(MakeCoord(mActual, nActual)); + if (endItr) { + if constexpr (!ENABLE_UNIT_FLAG_) { + AscendC::SetFlag(locPingPongFlag); + AscendC::WaitFlag(locPingPongFlag); + copyL0CToGm(gC, l0CTensor[locPingPongFlag * mRound * 128], layoutC, layoutInL0C); + AscendC::SetFlag(locPingPongFlag); + } else { + copyL0CToGm(gC, l0CTensor[locPingPongFlag * mRound * 128], layoutC, layoutInL0C, unitFlag); + } + } + } + +protected: + /// Data members + AscendC::LocalTensor l1ATensor; + AscendC::LocalTensor l1BTensor[STAGES]; + AscendC::LocalTensor l0ATensor[STAGES]; + AscendC::LocalTensor l0BTensor[STAGES]; + AscendC::LocalTensor l0CTensor; + + TileMmad tileMmad; + CopyGmToL1A copyGmToL1A; + CopyGmToL1B copyGmToL1B; + CopyL1ToL0A copyL1ToL0A; + CopyL1ToL0B copyL1ToL0B; + CopyL0CToGm copyL0CToGm; +}; + +//////////////////////////////////////////////////////////////////// + +} // namespace Catlass::Gemm::Block + +#endif // CATLASS_GEMM_BLOCK_MMAD_QK_TAIL_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_qk_head_tail_tla.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_qk_head_tail_tla.hpp new file mode 100644 index 0000000000..75b24d6608 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_qk_head_tail_tla.hpp @@ -0,0 +1,326 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_GEMM_BLOCK_MMAD_QK_TAIL_TLA_HPP +#define CATLASS_GEMM_BLOCK_MMAD_QK_TAIL_TLA_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/kda_coord.hpp" +#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" +#include "catlass/gemm/kda_gemm_helper.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" +#include "catlass/gemm/tile/kda_tile_mmad.hpp" +#include "tla/kda_tla_layout.hpp" +#include "tla/kda_tensor.hpp" + +//////////////////////////////////////////////////////////////////// + +namespace Catlass::Gemm::Block { +//////////////////////////////////////////////////////////////////// + +template < + class ArchTag_, + bool PAGED_CACHE_FLAG_, + bool ENABLE_UNIT_FLAG_, + class L1TileShape_, + class L0TileShape_, + class ElementA_, + class ElementB_, + class ElementC_, + class ElementBias_, + class TileCopy_, + class TileMmad_> +struct BlockMmadTla< + MmadFAITailQK, + L1TileShape_, + L0TileShape_, + ElementA_, + ElementB_, + ElementC_, + ElementBias_, + TileCopy_, + TileMmad_> { +public: + // Type Aliases + using DispatchPolicy = MmadFAITailQK; + using ArchTag = typename DispatchPolicy::ArchTag; + using L1TileShape = L1TileShape_; + using L0TileShape = L0TileShape_; + using ElementA = ElementA_; + using LayoutA = typename TileCopy_::LayoutA; + using ElementB = ElementB_; + using LayoutB = typename TileCopy_::LayoutB; + using ElementC = ElementC_; + using LayoutC = typename TileCopy_::LayoutC; + + using TileMmad = TileMmad_; + + using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; + + using ElementAccumulator = typename TileCopy_::ElementAccumulator; + + using LayoutTagL1A = typename TileCopy_::LayoutTagL1A; + using LayoutTagL1B = typename TileCopy_::LayoutTagL1B; + using LayoutTagL0A = typename TileCopy_::LayoutTagL0A; + using LayoutTagL0B = typename TileCopy_::LayoutTagL0B; + + using L1AAlignHelper = typename TileCopy_::L1AAlignHelper; + using L1BAlignHelper = typename TileCopy_::L1BAlignHelper; + + static_assert(tla::is_tuple::value && tla::is_static::value, + "L1TileShape must be tla::tuple and static!"); + static_assert(tla::is_tuple::value && tla::is_static::value, + "L0TileShape must be tla::tuple and static!"); + + static constexpr bool ENABLE_UNIT_FLAG = DispatchPolicy::ENABLE_UNIT_FLAG; + static constexpr uint32_t STAGES = DispatchPolicy::STAGES; + static constexpr uint32_t L1_TILE_M = tla::get<0>(L1TileShape{}); + static constexpr uint32_t L1_TILE_N = tla::get<1>(L1TileShape{}); + static constexpr uint32_t L1_TILE_K = tla::get<2>(L1TileShape{}); + static constexpr uint32_t L0_TILE_M = tla::get<0>(L0TileShape{}); + static constexpr uint32_t L0_TILE_N = tla::get<1>(L0TileShape{}); + static constexpr uint32_t L0_TILE_K = tla::get<2>(L0TileShape{}); + + // L1 tile size + static constexpr uint32_t L1A_TILE_SIZE = L1_TILE_M * L1_TILE_K * sizeof(ElementA); + static constexpr uint32_t L1B_TILE_SIZE = L1_TILE_N * L1_TILE_K * sizeof(ElementB); + // L0 tile size + static constexpr uint32_t L0A_TILE_SIZE = L0_TILE_M * L0_TILE_K * sizeof(ElementA); + static constexpr uint32_t L0B_TILE_SIZE = L0_TILE_K * L0_TILE_N * sizeof(ElementB); + static constexpr uint32_t L0C_TILE_SIZE = L0_TILE_M * L0_TILE_N * sizeof(ElementAccumulator); + static constexpr uint32_t L0A_PINGPONG_BUF_SIZE = ArchTag::L0A_SIZE / STAGES; + static constexpr uint32_t L0B_PINGPONG_BUF_SIZE = ArchTag::L0B_SIZE / STAGES; + static constexpr uint32_t L0C_PINGPONG_BUF_SIZE = ArchTag::L0C_SIZE / STAGES; + + // Check LayoutC + static_assert(tla::detail::isRowMajor::value || + ((std::is_same_v || std::is_same_v || + std::is_same_v) && + tla::detail::iszN::value), + "LayoutC only supports zN in half or bfloat16 or float, RowMajor in all dtype yet!"); + + // Check L1TileShape + static_assert(L1A_TILE_SIZE + L1B_TILE_SIZE * STAGES <= ArchTag::L1_SIZE, + "L1TileShape exceeding the L1 space!"); + + // Check L0TileShape + static_assert(L0A_TILE_SIZE * STAGES <= ArchTag::L0A_SIZE, "L0TileShape exceeding the L0A space!"); + static_assert(L0B_TILE_SIZE * STAGES <= ArchTag::L0B_SIZE, "L0TileShape exceeding the L0B space!"); + static_assert(L0C_TILE_SIZE <= ArchTag::L0C_SIZE, "L0TileShape exceeding the L0C space!"); + + static constexpr uint32_t BLOCK_SIZE = 16; + static constexpr uint32_t EMBED_SPLIT_SIZE = 128; + static constexpr uint32_t UNIT_BLOCK_STACK_NUM = 4; + static constexpr uint32_t KV_BASE_BLOCK = 512; + static constexpr uint32_t KV_SPLIT_SIZE = 128; + + CATLASS_DEVICE + BlockMmadTla(Arch::Resource &resource, uint32_t l1BufAddrStart = 0) + { + // Allocate L1 memory space + l1ATensor = resource.l1Buf.template GetBufferByByte(l1BufAddrStart); + for (uint32_t i = 0; i < STAGES; i++) { + l1BTensor[i] = + resource.l1Buf.template GetBufferByByte(l1BufAddrStart + L1A_TILE_SIZE + L1B_TILE_SIZE * i); + l0ATensor[i] = resource.l0ABuf.template GetBufferByByte(L0A_PINGPONG_BUF_SIZE * i); + l0BTensor[i] = resource.l0BBuf.template GetBufferByByte(L0B_PINGPONG_BUF_SIZE * i); + } + l0CTensor = resource.l0CBuf.template GetBufferByByte(0); + } + + /// Destructor + CATLASS_DEVICE + ~BlockMmadTla() {} + + CATLASS_DEVICE + void getBlockShape(GemmCoord &actualShape, uint32_t &nowNIdx, uint32_t &kIdx, + uint32_t &nLoop, uint32_t &kLoop, uint32_t &kvSeqlen, uint32_t &embed, bool firstBlock, uint32_t maskTailS = 0) + { + uint32_t nSplitSize = KV_SPLIT_SIZE; + uint32_t embedSplitSize = EMBED_SPLIT_SIZE; + if (nowNIdx == nLoop - 1) { + nSplitSize = kvSeqlen - nowNIdx * KV_SPLIT_SIZE; + } + if (firstBlock && maskTailS != 0) { + nSplitSize = nSplitSize - maskTailS; + } + if (kIdx == kLoop - 1) { + embedSplitSize = embed - kIdx * EMBED_SPLIT_SIZE; + } + actualShape[1] = nSplitSize; + actualShape[2] = embedSplitSize; + } + + CATLASS_DEVICE + void getKVOffset(AscendC::GlobalTensor &gBlockTable, uint32_t &kOffset, uint32_t &nowNIdx, uint32_t &kIdx, + uint32_t &nLoop, uint32_t &kLoop, uint32_t &strideKV, uint32_t &blockSize, uint32_t maskTailS = 0) + { + if (nowNIdx >= nLoop || kIdx >= kLoop) { + kOffset = 0; + } + if constexpr (PAGED_CACHE_FLAG_) { + uint32_t blockTableId = gBlockTable.GetValue(nowNIdx); + kOffset = blockTableId * blockSize * strideKV + maskTailS * strideKV + kIdx * EMBED_SPLIT_SIZE; + } else { + kOffset = nowNIdx * KV_SPLIT_SIZE * strideKV + kIdx * EMBED_SPLIT_SIZE; + } + } + + template + CATLASS_DEVICE void operator()(TensorA &tensorA, TensorB &tensorB, AscendC::GlobalTensor gC, + AscendC::GlobalTensor gBlockTable, GemmCoord actualOriShape, + uint32_t &nIdx, uint32_t &nLoop, uint32_t &blockSize, uint32_t kvSeqlen, uint32_t strideKV, uint32_t maskTailS, bool preloadFlag) + { + uint32_t rowNum = actualOriShape[0]; + uint32_t embed = actualOriShape[2]; + uint32_t kLoop = CeilDiv(embed); + uint32_t nkBlockLoop = nLoop * kLoop; + GemmCoord actualShape{rowNum, 0, 0}; + GemmCoord actualNextShape{rowNum, 0, 0}; + uint32_t nkBlockNextIdx = nIdx * kLoop + 1; + uint32_t gBOffset = 0; + uint32_t gBNextOffset = 0; + uint32_t stackTile = 0; + uint32_t preStackTile = 0; + for (uint32_t blockStackIdx = 0; (blockStackIdx < UNIT_BLOCK_STACK_NUM) && ((nIdx + blockStackIdx) < nLoop); + ++blockStackIdx) { + for (uint32_t kIdx = 0; kIdx < kLoop; kIdx++) { + uint32_t nowNIdx = nIdx + blockStackIdx; + uint32_t nLoopNextIdx = nkBlockNextIdx / kLoop; + uint32_t kLoopNextIdx = nkBlockNextIdx % kLoop; + uint32_t gCOffset = stackTile; + uint32_t startSeqOffset = nowNIdx == nIdx ? maskTailS : 0; + uint32_t startSeqNxtOffset = nLoopNextIdx == nIdx ? maskTailS : 0; + getBlockShape(actualShape, nowNIdx, kIdx, nLoop, kLoop, kvSeqlen, embed, nowNIdx == nIdx, maskTailS); + getBlockShape(actualNextShape, nLoopNextIdx, kLoopNextIdx, nLoop, kLoop, kvSeqlen, embed, nLoopNextIdx == nIdx, maskTailS); + getKVOffset(gBlockTable, gBOffset, nowNIdx, kIdx, nLoop, kLoop, strideKV, blockSize, startSeqOffset); + getKVOffset(gBlockTable, gBNextOffset, nLoopNextIdx, kLoopNextIdx, nLoop, kLoop, strideKV, blockSize, startSeqNxtOffset); + bool firstItr = kIdx == 0; + bool endItr = kIdx == kLoop - 1; + bool firstQtr = blockStackIdx == 0; + bool endQItr = ((nowNIdx == nLoop - 1) || (blockStackIdx == UNIT_BLOCK_STACK_NUM - 1)) && (kIdx == kLoop - 1); + bool initMmad = kIdx == 0; + preStackTile = stackTile; + stackTile += actualShape[1]; + auto layoutS = tla::MakeLayout(tla::MakeShape(rowNum, stackTile), tla::MakeStride(Int<512>{}, Int<1>{})); + auto tensorCOffset = tla::MakeTensor(gC[gCOffset], layoutS, Arch::PositionGM{}); + computeQK(tensorA, tensorB, tensorCOffset, gBOffset, gBNextOffset, + actualShape, actualNextShape, blockStackIdx, nkBlockNextIdx, nkBlockLoop, strideKV, firstItr, endItr, initMmad, firstQtr, endQItr, preloadFlag); + ++nkBlockNextIdx; + preloadFlag = false; + } + maskTailS = 0; + } + } + + template + CATLASS_DEVICE void computeQK( + TensorA &tensorA, TensorB &tensorB, TensorC &tensorC, uint32_t gBOffset, uint32_t gBNextOffset, + GemmCoord actualShape, GemmCoord actualNextShape, uint32_t nowIdx, uint32_t &nkblockIdx, + uint32_t &nkblockLoop, uint32_t strideKV, bool firstItr, bool endItr, bool initMmad, bool firstQItr, bool endQItr, bool preloadFlag) + { + using CopyGmToL1B = typename TileCopy_::template CopyGmToL1B; + using CopyL0CToGm = typename TileCopy_::template CopyL0CToGm; + CopyGmToL1B copyGmToL1B; + CopyL0CToGm copyL0CToGm; + + uint32_t mActual = actualShape.m(); + uint32_t kActual = actualShape.k(); + uint32_t nActual = actualShape.n(); + uint32_t mRound = RoundUp(mActual); + uint32_t kRound = RoundUp(kActual); + uint32_t nRound = RoundUp(nActual); + auto layoutAInL1 = tla::MakeLayout(mRound, kActual); + auto layoutAInL0 = tla::MakeLayout(mRound, kActual); + auto layoutBInL1 = tla::MakeLayout(kActual, nActual); + auto layoutBInL0 = tla::MakeLayout(kActual, nRound); + uint32_t locPingPongFlag = nkblockIdx % 2; + uint32_t l1KvPingPongFlag = nkblockIdx % 2; + uint32_t l0ABPingPongFlag = nkblockIdx % 2; + auto tensorL1A = tla::MakeTensor(l1ATensor, layoutAInL1, Arch::PositionL1{}); + auto tensorL0A = tla::MakeTensor(l0ATensor[l0ABPingPongFlag], layoutAInL0, Arch::PositionL0A{}); + auto tensorTileB = GetTile(tensorB, tla::MakeCoord(gBOffset % strideKV, gBOffset / strideKV), tla::MakeShape(kActual, nActual)); + auto tensorL1B = tla::MakeTensor(l1BTensor[l1KvPingPongFlag], layoutBInL1, Arch::PositionL1{}); + auto tensorL0B = tla::MakeTensor(l0BTensor[l0ABPingPongFlag], layoutBInL0, Arch::PositionL0B{}); + if (preloadFlag) { + AscendC::WaitFlag(l1KvPingPongFlag); + copyGmToL1B(tensorL1B, tensorTileB); + AscendC::SetFlag(l1KvPingPongFlag); + } + if (nkblockIdx != nkblockLoop) { + uint32_t nNextActual = actualNextShape.n(); + uint32_t kNextActual = actualNextShape.k(); + auto tensorTileBNext = GetTile(tensorB, tla::MakeCoord(gBNextOffset % strideKV, gBNextOffset / strideKV), tla::MakeShape(kNextActual, nNextActual)); + auto layoutBNextInL1 = tla::MakeLayout(kNextActual, nNextActual); + auto tensorL1BNext = tla::MakeTensor(l1BTensor[1 - l1KvPingPongFlag], layoutBNextInL1, Arch::PositionL1{}); + AscendC::WaitFlag(1 - l1KvPingPongFlag); + copyGmToL1B(tensorL1BNext, tensorTileBNext); + AscendC::SetFlag(1 - l1KvPingPongFlag); + } + AscendC::WaitFlag(l0ABPingPongFlag); + copyL1ToL0A(tensorL0A, tensorL1A); + + AscendC::WaitFlag(l1KvPingPongFlag); + AscendC::WaitFlag(l0ABPingPongFlag + 2); + copyL1ToL0B(tensorL0B, tensorL1B); + AscendC::SetFlag(l1KvPingPongFlag); + + AscendC::SetFlag(l0ABPingPongFlag); + AscendC::WaitFlag(l0ABPingPongFlag); + uint8_t unitFlag = 0b00; + if constexpr (!ENABLE_UNIT_FLAG_) { + if (firstItr) { + AscendC::WaitFlag(locPingPongFlag); + } + } else { + unitFlag = 0b11; + } + auto layoutCTileInL0 = tla::MakeLayoutL0C(mRound, nActual); + auto tensorL0CTile = tla::MakeTensor(l0CTensor[locPingPongFlag * mRound * 128], layoutCTileInL0, Arch::PositionL0C{}); + tileMmad( + tensorL0CTile, tensorL0A, tensorL0B, + mRound, nActual, kActual, initMmad, unitFlag); + AscendC::SetFlag(l0ABPingPongFlag); + AscendC::SetFlag(l0ABPingPongFlag + 2); + auto layoutCInL0 = tla::MakeLayoutL0C(mActual, nActual); + auto tensorL0C = tla::MakeTensor(l0CTensor[locPingPongFlag * mRound * 128], layoutCInL0, Arch::PositionL0C{}); + if (endItr) { + if constexpr (!ENABLE_UNIT_FLAG_) { + AscendC::SetFlag(locPingPongFlag); + AscendC::WaitFlag(locPingPongFlag); + copyL0CToGm(tensorC, tensorL0C); + AscendC::SetFlag(locPingPongFlag); + } else { + copyL0CToGm(tensorC, tensorL0C, unitFlag); + } + } + } + +protected: + /// Data members + AscendC::LocalTensor l1ATensor; + AscendC::LocalTensor l1BTensor[STAGES]; + AscendC::LocalTensor l0ATensor[STAGES]; + AscendC::LocalTensor l0BTensor[STAGES]; + AscendC::LocalTensor l0CTensor; + + TileMmad tileMmad; + CopyL1ToL0A copyL1ToL0A; + CopyL1ToL0B copyL1ToL0B; +}; + +//////////////////////////////////////////////////////////////////// + +} // namespace Catlass::Gemm::Block + +#endif // CATLASS_GEMM_BLOCK_MMAD_QK_TAIL_TLA_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_qk_normal.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_qk_normal.hpp new file mode 100644 index 0000000000..d43f63ba13 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_qk_normal.hpp @@ -0,0 +1,337 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_GEMM_BLOCK_MMAD_QK_HPP +#define CATLASS_GEMM_BLOCK_MMAD_QK_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/kda_coord.hpp" +#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" +#include "catlass/gemm/kda_gemm_helper.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" +#include "catlass/gemm/tile/kda_tile_mmad.hpp" + +//////////////////////////////////////////////////////////////////// + +namespace Catlass::Gemm::Block { +//////////////////////////////////////////////////////////////////// + +template < + bool PAGED_CACHE_FLAG_, + bool ENABLE_UNIT_FLAG_, + class L1TileShape_, + class L0TileShape_, + class AType_, + class BType_, + class CType_, + class BiasType_, + class TileCopy_, + class TileMmad_> +struct BlockMmad< + MmadAtlasA2FAIQK, + L1TileShape_, + L0TileShape_, + AType_, + BType_, + CType_, + BiasType_, + TileCopy_, + TileMmad_> { +public: + // Type Aliases + using DispatchPolicy = MmadAtlasA2FAIQK; + using ArchTag = typename DispatchPolicy::ArchTag; + using L1TileShape = L1TileShape_; + using L0TileShape = L0TileShape_; + using ElementA = typename AType_::Element; + using LayoutA = typename AType_::Layout; + using ElementB = typename BType_::Element; + using LayoutB = typename BType_::Layout; + using ElementC = typename CType_::Element; + using LayoutC = typename CType_::Layout; + using TileMmad = TileMmad_; + using CopyGmToL1A = typename TileCopy_::CopyGmToL1A; + using CopyGmToL1B = typename TileCopy_::CopyGmToL1B; + using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; + using CopyL0CToGm = typename TileCopy_::CopyL0CToGm; + using ElementAccumulator = + typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; + using LayoutAInL1 = typename CopyL1ToL0A::LayoutSrc; + using LayoutBInL1 = typename CopyL1ToL0B::LayoutSrc; + using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst; + using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst; + using LayoutCInL0 = layout::zN; + + using L1AAlignHelper = Gemm::helper::L1AlignHelper; + using L1BAlignHelper = Gemm::helper::L1AlignHelper; + + static constexpr uint32_t STAGES = DispatchPolicy::STAGES; + static constexpr uint32_t L1A_SIZE = L1TileShape::M * L1TileShape::K * sizeof(ElementA); + static constexpr uint32_t L1B_SIZE = L1TileShape::N * L1TileShape::K * sizeof(ElementB); + static constexpr uint32_t L0A_SIZE = ArchTag::L0A_SIZE; + static constexpr uint32_t L0B_SIZE = ArchTag::L0B_SIZE; + static constexpr uint32_t L0C_SIZE = ArchTag::L0C_SIZE; + static constexpr uint32_t L0A_PINGPONG_BUF_SIZE = L0A_SIZE / STAGES; + static constexpr uint32_t L0B_PINGPONG_BUF_SIZE = L0B_SIZE / STAGES; + static constexpr uint32_t L0C_PINGPONG_BUF_SIZE = L0C_SIZE / STAGES; + static constexpr uint32_t BLOCK_SIZE = 16; + static constexpr uint32_t EMBED_SPLIT_SIZE = 128; + static constexpr uint32_t UNIT_BLOCK_STACK_NUM = 4; + static constexpr uint32_t KV_BASE_BLOCK = 512; + static constexpr uint32_t KV_SPLIT_SIZE = 128; + + static_assert(std::is_same_v, "LayoutC only support RowMajor yet!"); + + CATLASS_DEVICE + BlockMmad(Arch::Resource &resource, uint32_t l1BufAddrStart = 0) + { + // Allocate L1 memory space + l1ATensor = resource.l1Buf.template GetBufferByByte(l1BufAddrStart); + for (uint32_t i = 0; i < STAGES; i++) { + l1BTensor[i] = + resource.l1Buf.template GetBufferByByte(l1BufAddrStart + L1A_SIZE + L1B_SIZE * i); + l0ATensor[i] = resource.l0ABuf.template GetBufferByByte(L0A_PINGPONG_BUF_SIZE * i); + l0BTensor[i] = resource.l0BBuf.template GetBufferByByte(L0B_PINGPONG_BUF_SIZE * i); + } + l0CTensor = resource.l0CBuf.template GetBufferByByte(0); + } + + CATLASS_DEVICE + ~BlockMmad() {} + + CATLASS_DEVICE + void loadQGM(AscendC::GlobalTensor gA, LayoutA layoutA, uint32_t rowNum, uint32_t &singleGroupHeads, uint32_t &qHeads) + { + uint32_t embed = layoutA.shape(1); + uint32_t rowNumRound = RoundUp(rowNum); + uint32_t tokenNumPerGroup = rowNum / singleGroupHeads; + auto layoutSingleANd = layoutA.GetTileLayout(MakeCoord(singleGroupHeads, embed)); + LayoutAInL1 layoutAInL1 = LayoutAInL1::template MakeLayout(rowNum, embed); + copyGmToL1A( + l1ATensor, gA, + layoutAInL1, layoutSingleANd, + tokenNumPerGroup, qHeads * embed, tokenNumPerGroup, BLOCK_SIZE, rowNumRound); + // AscendC::Nd2NzParams intriParams; + // intriParams.nValue = singleGroupHeads; + // intriParams.dValue = embed; + // intriParams.srcDValue = embed; + // intriParams.dstNzNStride = tokenNumPerGroup; + // intriParams.dstNzC0Stride = rowNumRound; + // intriParams.ndNum = tokenNumPerGroup; + // intriParams.dstNzMatrixStride = 16; + // AscendC::DataCopy(l1ATensor, gA, intriParams); + AscendC::SetFlag(EVENT_ID3); + AscendC::WaitFlag(EVENT_ID3); + } + + CATLASS_DEVICE + void getBlockShape(GemmCoord &actualShape, uint32_t &nowNIdx, uint32_t &kIdx, + uint32_t &nLoop, uint32_t &kLoop, uint32_t &kvSeqlen, uint32_t &embed, bool firstBlock, uint32_t maskTailS = 0) + { + uint32_t nSplitSize = KV_SPLIT_SIZE; + uint32_t embedSplitSize = EMBED_SPLIT_SIZE; + if (nowNIdx == nLoop - 1) { + nSplitSize = kvSeqlen - nowNIdx * KV_SPLIT_SIZE; + } + if (firstBlock && maskTailS != 0) { + nSplitSize = nSplitSize - maskTailS; + } + // } + if (kIdx == kLoop - 1) { + embedSplitSize = embed - kIdx * EMBED_SPLIT_SIZE; + } + actualShape[1] = nSplitSize; + actualShape[2] = embedSplitSize; + } + + CATLASS_DEVICE + void getKVOffset(AscendC::GlobalTensor &gBlockTable, uint32_t &kOffset, uint32_t &nowNIdx, uint32_t &kIdx, + uint32_t &nLoop, uint32_t &kLoop, uint32_t &strideKV, uint32_t &blockSize, uint32_t maskTailS = 0) + { + if (nowNIdx >= nLoop || kIdx >= kLoop) { + kOffset = 0; + } + if constexpr (PAGED_CACHE_FLAG_) { + uint32_t blockTableId = gBlockTable.GetValue(nowNIdx); + kOffset = blockTableId * blockSize * strideKV + maskTailS * strideKV + kIdx * EMBED_SPLIT_SIZE; + } else { + kOffset = nowNIdx * KV_SPLIT_SIZE * strideKV + kIdx * EMBED_SPLIT_SIZE; + } + } + + CATLASS_DEVICE + void operator()(AscendC::GlobalTensor gA, + AscendC::GlobalTensor gB, + AscendC::GlobalTensor gC, + AscendC::GlobalTensor gBlockTable, + LayoutA layoutA, LayoutB layoutB, GemmCoord actualOriShape, + uint32_t &nIdx, uint32_t &nLoop, uint32_t &blockSize, uint32_t kvSeqlen, uint32_t strideKV) + { + uint32_t rowNum = actualOriShape[0]; + uint32_t embed = actualOriShape[2]; + uint32_t kLoop = CeilDiv(embed); + uint32_t nkBlockLoop = nLoop * kLoop; + GemmCoord actualShape{rowNum, 0, 0}; + GemmCoord actualNextShape{rowNum, 0, 0}; + uint32_t nkBlockNextIdx = nIdx * kLoop + 1; + uint32_t gBOffset = 0; + uint32_t gBNextOffset = 0; + uint32_t stackTile = 0; + for (uint32_t blockStackIdx = 0; (blockStackIdx < UNIT_BLOCK_STACK_NUM) && ((nIdx + blockStackIdx) < nLoop); + ++blockStackIdx) { + for (uint32_t kIdx = 0; kIdx < kLoop; kIdx++) { + uint32_t nowNIdx = nIdx + blockStackIdx; + uint32_t nLoopNextIdx = nkBlockNextIdx / kLoop; + uint32_t kLoopNextIdx = nkBlockNextIdx % kLoop; + uint32_t gCOffset = blockStackIdx / 2 * 2 * KV_SPLIT_SIZE; + getBlockShape(actualShape, nowNIdx, kIdx, nLoop, kLoop, kvSeqlen, embed, nowNIdx == nIdx); + getBlockShape(actualNextShape, nLoopNextIdx, kLoopNextIdx, nLoop, kLoop, kvSeqlen, embed, nLoopNextIdx == nIdx); + getKVOffset(gBlockTable, gBOffset, nowNIdx, kIdx, nLoop, kLoop, strideKV, blockSize); + getKVOffset(gBlockTable, gBNextOffset, nLoopNextIdx, kLoopNextIdx, nLoop, kLoop, strideKV, blockSize); + bool firstItr = ((blockStackIdx % 2) == 0) && (kIdx == 0); + bool endItr = (((blockStackIdx % 2) == 1) || (nowNIdx == nLoop - 1)) && (kIdx == kLoop - 1); + bool firstQtr = blockStackIdx == 0; + bool endQItr = ((nowNIdx == nLoop - 1) || (blockStackIdx == UNIT_BLOCK_STACK_NUM - 1)) && (kIdx == kLoop - 1); + int cc = 1; + bool initMmad = kIdx == 0; + stackTile += actualShape[1]; + LayoutC layOutSTemp(rowNum, stackTile, 512); + // AscendC::printf("firstItr:%d\n", firstItr); + // AscendC::printf("endItr:%d\n", endItr); + // AscendC::printf("initMmad:%d\n", initMmad); + // AscendC::printf("stackTile:%d\n", stackTile); + // AscendC::printf("blockStackIdx:%d\n", blockStackIdx); + // AscendC::printf("gBOffset:%d\n", gBOffset); + // AscendC::printf("gBNextOffset:%d\n", gBNextOffset); + // AscendC::printf("actualShape[0]:%d\n", actualShape[0]); + // AscendC::printf("actualShape[1]:%d\n", actualShape[1]); + // AscendC::printf("actualShape[2]:%d\n", actualShape[2]); + // AscendC::printf("actualShape.m():%d\n", actualShape.m()); + // AscendC::printf("actualShape.n():%d\n", actualShape.n()); + // AscendC::printf("actualShape.k():%d\n", actualShape.k()); + // AscendC::printf("actualNextShape[0]:%d\n", actualNextShape[0]); + // AscendC::printf("actualNextShape[1]:%d\n", actualNextShape[1]); + // AscendC::printf("actualNextShape[2]:%d\n", actualNextShape[2]); + computeQK(gA, gB[gBOffset], gC[gCOffset], gB[gBNextOffset], layoutA, layoutB, layOutSTemp, + actualShape, actualNextShape, blockStackIdx, nkBlockNextIdx, nkBlockLoop, firstItr, endItr, initMmad, firstQtr, endQItr); + ++nkBlockNextIdx; + if (endItr) { + stackTile = 0; + } + } + } + } + + CATLASS_DEVICE void computeQK( + AscendC::GlobalTensor const &gA, + AscendC::GlobalTensor const &gB, + AscendC::GlobalTensor const &gC, + AscendC::GlobalTensor const &gmNextBlockB, + LayoutA layoutA, LayoutB layoutB, LayoutC layoutC, + GemmCoord actualShape, GemmCoord actualNextShape, uint32_t nowIdx, uint32_t &nkblockIdx, + uint32_t &nkblockLoop, bool firstItr, bool endItr, bool initMmad, bool firstQItr, bool endQItr) + { + uint32_t mActual = actualShape.m(); + uint32_t kActual = actualShape.k(); + uint32_t nActual = actualShape.n(); + uint32_t mRound = RoundUp(mActual); + uint32_t kRound = RoundUp(kActual); + uint32_t nRound = RoundUp(nActual); + LayoutAInL1 layoutAInL1 = LayoutAInL1::template MakeLayout(mRound, kActual); // embed + LayoutAInL0 layoutAInL0 = LayoutAInL0::template MakeLayout(mRound, kActual); + LayoutBInL1 layoutBInL1 = LayoutBInL1::template MakeLayout(kActual, nActual); + LayoutBInL0 layoutBInL0 = LayoutBInL0::template MakeLayout(kActual, nRound); + uint32_t locPingPongFlag = nowIdx % 2; + uint32_t l1KvPingPongFlag = nkblockIdx % 2; + uint32_t l0ABPingPongFlag = nkblockIdx % 2; + if (nkblockIdx == 1) { + auto layoutBTile = layoutB.GetTileLayout(MakeCoord(kActual, nActual)); + AscendC::WaitFlag(l1KvPingPongFlag); + copyGmToL1B(l1BTensor[l1KvPingPongFlag], gB, layoutBInL1, layoutBTile); + AscendC::SetFlag(l1KvPingPongFlag); + } + if (nkblockIdx != nkblockLoop) { + uint32_t nNextActual = actualNextShape.n(); + uint32_t kNextActual = actualNextShape.k(); + LayoutBInL1 layoutBNextInL1 = LayoutBInL1::template MakeLayout(kNextActual, nNextActual); + auto layoutNextBTile = layoutB.GetTileLayout(MakeCoord(kNextActual, nNextActual)); + AscendC::WaitFlag(1 - l1KvPingPongFlag); + copyGmToL1B(l1BTensor[1 - l1KvPingPongFlag], gmNextBlockB, layoutBNextInL1, layoutNextBTile); + AscendC::SetFlag(1 - l1KvPingPongFlag); + } + if (firstQItr) { + AscendC::WaitFlag(0); + AscendC::WaitFlag(1); + copyL1ToL0A(l0ATensor[0], l1ATensor, layoutAInL0, layoutAInL1); + } + + AscendC::WaitFlag(l1KvPingPongFlag); + AscendC::WaitFlag(l0ABPingPongFlag + 2); + copyL1ToL0B(l0BTensor[l0ABPingPongFlag], l1BTensor[l1KvPingPongFlag], layoutBInL0, layoutBInL1); + AscendC::SetFlag(l1KvPingPongFlag); + + AscendC::SetFlag(l0ABPingPongFlag); + AscendC::WaitFlag(l0ABPingPongFlag); + uint8_t unitFlag = 0b00; + if constexpr (!ENABLE_UNIT_FLAG_) { + if (firstItr) { + AscendC::WaitFlag(0); + AscendC::WaitFlag(1); + } + } else { + unitFlag = 0b11; + } + tileMmad( + l0CTensor[locPingPongFlag * mRound * 128], l0ATensor[0], l0BTensor[l0ABPingPongFlag], + mRound, nActual, kActual, initMmad, unitFlag); + // AscendC::PipeBarrier(); + if (endQItr) { + AscendC::SetFlag(0); + AscendC::SetFlag(1); + } + AscendC::SetFlag(l0ABPingPongFlag + 2); + auto layoutInL0C = LayoutCInL0::MakeLayoutInL0C(MakeCoord(mActual, (uint32_t)256)); + if (endItr) { + if constexpr (!ENABLE_UNIT_FLAG_) { + AscendC::SetFlag(l1KvPingPongFlag); + AscendC::WaitFlag(l1KvPingPongFlag); + copyL0CToGm(gC, l0CTensor, layoutC, layoutInL0C); + AscendC::SetFlag(0); + AscendC::SetFlag(1); + } else { + copyL0CToGm(gC, l0CTensor, layoutC, layoutInL0C, unitFlag); + } + } + } + +protected: + /// Data members + AscendC::LocalTensor l1ATensor; + AscendC::LocalTensor l1BTensor[STAGES]; + AscendC::LocalTensor l0ATensor[STAGES]; + AscendC::LocalTensor l0BTensor[STAGES]; + AscendC::LocalTensor l0CTensor; + + TileMmad tileMmad; + CopyGmToL1A copyGmToL1A; + CopyGmToL1B copyGmToL1B; + CopyL1ToL0A copyL1ToL0A; + CopyL1ToL0B copyL1ToL0B; + CopyL0CToGm copyL0CToGm; +}; + +//////////////////////////////////////////////////////////////////// + +} // namespace Catlass::Gemm::Block + +#endif // CATLASS_GEMM_BLOCK_MMAD_QK_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_qk_normal_tla.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_qk_normal_tla.hpp new file mode 100644 index 0000000000..540cd1aedd --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_qk_normal_tla.hpp @@ -0,0 +1,347 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_GEMM_BLOCK_MMAD_QK_TLA_HPP +#define CATLASS_GEMM_BLOCK_MMAD_QK_TLA_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/kda_coord.hpp" +#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" +#include "catlass/gemm/kda_gemm_helper.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" +#include "catlass/gemm/tile/kda_tile_mmad.hpp" +#include "tla/kda_tla_layout.hpp" +#include "tla/kda_tensor.hpp" + +//////////////////////////////////////////////////////////////////// + +namespace Catlass::Gemm::Block { +//////////////////////////////////////////////////////////////////// + +template < + class ArchTag_, + bool PAGED_CACHE_FLAG_, + bool ENABLE_UNIT_FLAG_, + class L1TileShape_, + class L0TileShape_, + class ElementA_, + class ElementB_, + class ElementC_, + class ElementBias_, + class TileCopy_, + class TileMmad_> +struct BlockMmadTla< + MmadFAIQK, + L1TileShape_, + L0TileShape_, + ElementA_, + ElementB_, + ElementC_, + ElementBias_, + TileCopy_, + TileMmad_> { +public: + // Type Aliases + using DispatchPolicy = MmadFAIQK; + using ArchTag = typename DispatchPolicy::ArchTag; + using L1TileShape = L1TileShape_; + using L0TileShape = L0TileShape_; + using ElementA = ElementA_; + using LayoutA = typename TileCopy_::LayoutA; + using ElementB = ElementB_; + using LayoutB = typename TileCopy_::LayoutB; + using ElementC = ElementC_; + using LayoutC = typename TileCopy_::LayoutC; + + using TileMmad = TileMmad_; + + using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; + + using ElementAccumulator = typename TileCopy_::ElementAccumulator; + + using LayoutTagL1A = typename TileCopy_::LayoutTagL1A; + using LayoutTagL1B = typename TileCopy_::LayoutTagL1B; + using LayoutTagL0A = typename TileCopy_::LayoutTagL0A; + using LayoutTagL0B = typename TileCopy_::LayoutTagL0B; + + using L1AAlignHelper = typename TileCopy_::L1AAlignHelper; + using L1BAlignHelper = typename TileCopy_::L1BAlignHelper; + + static_assert(tla::is_tuple::value && tla::is_static::value, + "L1TileShape must be tla::tuple and static!"); + static_assert(tla::is_tuple::value && tla::is_static::value, + "L0TileShape must be tla::tuple and static!"); + + static constexpr bool ENABLE_UNIT_FLAG = DispatchPolicy::ENABLE_UNIT_FLAG; + static constexpr uint32_t STAGES = DispatchPolicy::STAGES; + static constexpr uint32_t L1_TILE_M = tla::get<0>(L1TileShape{}); + static constexpr uint32_t L1_TILE_N = tla::get<1>(L1TileShape{}); + static constexpr uint32_t L1_TILE_K = tla::get<2>(L1TileShape{}); + static constexpr uint32_t L0_TILE_M = tla::get<0>(L0TileShape{}); + static constexpr uint32_t L0_TILE_N = tla::get<1>(L0TileShape{}); + static constexpr uint32_t L0_TILE_K = tla::get<2>(L0TileShape{}); + + // L1 tile size + static constexpr uint32_t L1A_TILE_SIZE = L1_TILE_M * L1_TILE_K * sizeof(ElementA); + static constexpr uint32_t L1B_TILE_SIZE = L1_TILE_N * L1_TILE_K * sizeof(ElementB); + // L0 tile size + static constexpr uint32_t L0A_TILE_SIZE = L0_TILE_M * L0_TILE_K * sizeof(ElementA); + static constexpr uint32_t L0B_TILE_SIZE = L0_TILE_K * L0_TILE_N * sizeof(ElementB); + static constexpr uint32_t L0C_TILE_SIZE = L0_TILE_M * L0_TILE_N * sizeof(ElementAccumulator); + static constexpr uint32_t L0A_PINGPONG_BUF_SIZE = ArchTag::L0A_SIZE / STAGES; + static constexpr uint32_t L0B_PINGPONG_BUF_SIZE = ArchTag::L0B_SIZE / STAGES; + static constexpr uint32_t L0C_PINGPONG_BUF_SIZE = ArchTag::L0C_SIZE / STAGES; + + // Check LayoutC + static_assert(tla::detail::isRowMajor::value || + ((std::is_same_v || std::is_same_v || + std::is_same_v) && + tla::detail::iszN::value), + "LayoutC only supports zN in half or bfloat16 or float, RowMajor in all dtype yet!"); + + // Check L1TileShape + static_assert(L1A_TILE_SIZE + L1B_TILE_SIZE * STAGES <= ArchTag::L1_SIZE, + "L1TileShape exceeding the L1 space!"); + + // Check L0TileShape + static_assert(L0A_TILE_SIZE * STAGES <= ArchTag::L0A_SIZE, "L0TileShape exceeding the L0A space!"); + static_assert(L0B_TILE_SIZE * STAGES <= ArchTag::L0B_SIZE, "L0TileShape exceeding the L0B space!"); + static_assert(L0C_TILE_SIZE <= ArchTag::L0C_SIZE, "L0TileShape exceeding the L0C space!"); + + static constexpr uint32_t BLOCK_SIZE = 16; + static constexpr uint32_t EMBED_SPLIT_SIZE = 128; + static constexpr uint32_t UNIT_BLOCK_STACK_NUM = 4; + static constexpr uint32_t KV_BASE_BLOCK = 512; + static constexpr uint32_t KV_SPLIT_SIZE = 128; + + CATLASS_DEVICE + BlockMmadTla(Arch::Resource &resource, uint32_t l1BufAddrStart = 0) + { + // Allocate L1 memory space + l1ATensor = resource.l1Buf.template GetBufferByByte(l1BufAddrStart); + for (uint32_t i = 0; i < STAGES; i++) { + l1BTensor[i] = + resource.l1Buf.template GetBufferByByte(l1BufAddrStart + L1A_TILE_SIZE + L1B_TILE_SIZE * i); + l0ATensor[i] = resource.l0ABuf.template GetBufferByByte(L0A_PINGPONG_BUF_SIZE * i); + l0BTensor[i] = resource.l0BBuf.template GetBufferByByte(L0B_PINGPONG_BUF_SIZE * i); + } + l0CTensor = resource.l0CBuf.template GetBufferByByte(0); + } + + CATLASS_DEVICE + ~BlockMmadTla() {} + + template + CATLASS_DEVICE void loadQGM(TensorA &tensorA) + { + using CopyGmToL1A = Gemm::Tile::TileCopyFAQTla; + CopyGmToL1A copyGmToL1A; + + uint32_t rowNum = tla::get<0>(tensorA.shape()) * tla::get<1>(tensorA.shape()); + uint32_t embed = tla::get<2>(tensorA.shape()); + auto L1A_LAYOUT = tla::MakeLayout(rowNum, embed); + auto tensorL1A = tla::MakeTensor(l1ATensor, L1A_LAYOUT, Arch::PositionL1{}); + copyGmToL1A(tensorL1A, tensorA); + AscendC::SetFlag(EVENT_ID3); + AscendC::WaitFlag(EVENT_ID3); + } + + CATLASS_DEVICE + void getBlockShape(GemmCoord &actualShape, uint32_t &nowNIdx, uint32_t &kIdx, + uint32_t &nLoop, uint32_t &kLoop, uint32_t &kvSeqlen, uint32_t &embed, bool firstBlock, uint32_t maskTailS = 0) + { + uint32_t nSplitSize = KV_SPLIT_SIZE; + uint32_t embedSplitSize = EMBED_SPLIT_SIZE; + if (nowNIdx == nLoop - 1) { + nSplitSize = kvSeqlen - nowNIdx * KV_SPLIT_SIZE; + } + if (firstBlock && maskTailS != 0) { + nSplitSize = nSplitSize - maskTailS; + } + // } + if (kIdx == kLoop - 1) { + embedSplitSize = embed - kIdx * EMBED_SPLIT_SIZE; + } + actualShape[1] = nSplitSize; + actualShape[2] = embedSplitSize; + } + + CATLASS_DEVICE + void getKVOffset(AscendC::GlobalTensor &gBlockTable, uint32_t &kOffset, uint32_t &nowNIdx, uint32_t &kIdx, + uint32_t &nLoop, uint32_t &kLoop, uint32_t &strideKV, uint32_t &blockSize, uint32_t maskTailS = 0) + { + if (nowNIdx >= nLoop || kIdx >= kLoop) { + kOffset = 0; + } + if constexpr (PAGED_CACHE_FLAG_) { + uint32_t blockTableId = gBlockTable.GetValue(nowNIdx); + kOffset = blockTableId * blockSize * strideKV + maskTailS * strideKV + kIdx * EMBED_SPLIT_SIZE; + } else { + kOffset = nowNIdx * KV_SPLIT_SIZE * strideKV + kIdx * EMBED_SPLIT_SIZE; + } + } + + template + CATLASS_DEVICE void operator()(TensorA &tensorA, TensorB &tensorB, AscendC::GlobalTensor gC, + AscendC::GlobalTensor gBlockTable, GemmCoord actualOriShape, + uint32_t &nIdx, uint32_t &nLoop, uint32_t &blockSize, uint32_t kvSeqlen, uint32_t strideKV) + { + uint32_t rowNum = actualOriShape[0]; + uint32_t embed = actualOriShape[2]; + uint32_t kLoop = CeilDiv(embed); + uint32_t nkBlockLoop = nLoop * kLoop; + GemmCoord actualShape{rowNum, 0, 0}; + GemmCoord actualNextShape{rowNum, 0, 0}; + uint32_t nkBlockNextIdx = nIdx * kLoop + 1; + uint32_t gBOffset = 0; + uint32_t gBNextOffset = 0; + uint32_t stackTile = 0; + for (uint32_t blockStackIdx = 0; (blockStackIdx < UNIT_BLOCK_STACK_NUM) && ((nIdx + blockStackIdx) < nLoop); + ++blockStackIdx) { + for (uint32_t kIdx = 0; kIdx < kLoop; kIdx++) { + uint32_t nowNIdx = nIdx + blockStackIdx; + uint32_t nLoopNextIdx = nkBlockNextIdx / kLoop; + uint32_t kLoopNextIdx = nkBlockNextIdx % kLoop; + uint32_t gCOffset = blockStackIdx / 2 * 2 * KV_SPLIT_SIZE; + getBlockShape(actualShape, nowNIdx, kIdx, nLoop, kLoop, kvSeqlen, embed, nowNIdx == nIdx); + getBlockShape(actualNextShape, nLoopNextIdx, kLoopNextIdx, nLoop, kLoop, kvSeqlen, embed, nLoopNextIdx == nIdx); + getKVOffset(gBlockTable, gBOffset, nowNIdx, kIdx, nLoop, kLoop, strideKV, blockSize); + getKVOffset(gBlockTable, gBNextOffset, nLoopNextIdx, kLoopNextIdx, nLoop, kLoop, strideKV, blockSize); + bool firstItr = ((blockStackIdx % 2) == 0) && (kIdx == 0); + bool endItr = (((blockStackIdx % 2) == 1) || (nowNIdx == nLoop - 1)) && (kIdx == kLoop - 1); + bool firstQtr = blockStackIdx == 0; + bool endQItr = ((nowNIdx == nLoop - 1) || (blockStackIdx == UNIT_BLOCK_STACK_NUM - 1)) && (kIdx == kLoop - 1); + int cc = 1; + bool initMmad = kIdx == 0; + stackTile += actualShape[1]; + auto layoutS = tla::MakeLayout(tla::MakeShape(rowNum, stackTile), tla::MakeStride(Int<512>{}, Int<1>{})); + auto tensorCOffset = tla::MakeTensor(gC[gCOffset], layoutS, Arch::PositionGM{}); + computeQK(tensorA, tensorB, tensorCOffset, gBOffset, gBNextOffset, + actualShape, actualNextShape, blockStackIdx, nkBlockNextIdx, nkBlockLoop, strideKV, firstItr, endItr, initMmad, firstQtr, endQItr); + ++nkBlockNextIdx; + if (endItr) { + stackTile = 0; + } + } + } + } + + template + CATLASS_DEVICE void computeQK( + TensorA &tensorA, TensorB &tensorB, TensorC &tensorC, uint32_t gBOffset, uint32_t gBNextOffset, + GemmCoord actualShape, GemmCoord actualNextShape, uint32_t nowIdx, uint32_t &nkblockIdx, + uint32_t &nkblockLoop, uint32_t strideKV, bool firstItr, bool endItr, bool initMmad, bool firstQItr, bool endQItr) + { + using CopyGmToL1B = typename TileCopy_::template CopyGmToL1B; + using CopyL0CToGm = typename TileCopy_::template CopyL0CToGm; + CopyGmToL1B copyGmToL1B; + CopyL0CToGm copyL0CToGm; + + uint32_t mActual = actualShape.m(); + uint32_t kActual = actualShape.k(); + uint32_t nActual = actualShape.n(); + uint32_t mRound = RoundUp(mActual); + uint32_t kRound = RoundUp(kActual); + uint32_t nRound = RoundUp(nActual); + auto layoutAInL1 = tla::MakeLayout(mRound, kActual); + auto layoutAInL0 = tla::MakeLayout(mRound, kActual); + auto layoutBInL1 = tla::MakeLayout(kActual, nActual); + auto layoutBInL0 = tla::MakeLayout(kActual, nRound); + uint32_t locPingPongFlag = nowIdx % 2; + uint32_t l1KvPingPongFlag = nkblockIdx % 2; + uint32_t l0ABPingPongFlag = nkblockIdx % 2; + auto tensorL1A = tla::MakeTensor(l1ATensor, layoutAInL1, Arch::PositionL1{}); + auto tensorL0A = tla::MakeTensor(l0ATensor[0], layoutAInL0, Arch::PositionL0A{}); + auto tensorTileB = GetTile(tensorB, tla::MakeCoord(gBOffset % strideKV, gBOffset / strideKV), tla::MakeShape(kActual, nActual)); + auto tensorL1B = tla::MakeTensor(l1BTensor[l1KvPingPongFlag], layoutBInL1, Arch::PositionL1{}); + auto tensorL0B = tla::MakeTensor(l0BTensor[l0ABPingPongFlag], layoutBInL0, Arch::PositionL0B{}); + if (nkblockIdx == 1) { + AscendC::WaitFlag(l1KvPingPongFlag); + copyGmToL1B(tensorL1B, tensorTileB); + AscendC::SetFlag(l1KvPingPongFlag); + } + if (nkblockIdx != nkblockLoop) { + uint32_t nNextActual = actualNextShape.n(); + uint32_t kNextActual = actualNextShape.k(); + auto tensorTileBNext = GetTile(tensorB, tla::MakeCoord(gBNextOffset % strideKV, gBNextOffset / strideKV), tla::MakeShape(kNextActual, nNextActual)); + auto layoutBNextInL1 = tla::MakeLayout(kNextActual, nNextActual); + auto tensorL1BNext = tla::MakeTensor(l1BTensor[1 - l1KvPingPongFlag], layoutBNextInL1, Arch::PositionL1{}); + AscendC::WaitFlag(1 - l1KvPingPongFlag); + copyGmToL1B(tensorL1BNext, tensorTileBNext); + AscendC::SetFlag(1 - l1KvPingPongFlag); + } + if (firstQItr) { + AscendC::WaitFlag(0); + AscendC::WaitFlag(1); + copyL1ToL0A(tensorL0A, tensorL1A); + } + + AscendC::WaitFlag(l1KvPingPongFlag); + AscendC::WaitFlag(l0ABPingPongFlag + 2); + copyL1ToL0B(tensorL0B, tensorL1B); + AscendC::SetFlag(l1KvPingPongFlag); + + AscendC::SetFlag(l0ABPingPongFlag); + AscendC::WaitFlag(l0ABPingPongFlag); + uint8_t unitFlag = 0b00; + if constexpr (!ENABLE_UNIT_FLAG_) { + if (firstItr) { + AscendC::WaitFlag(0); + AscendC::WaitFlag(1); + } + } else { + unitFlag = 0b11; + } + auto layoutCTileInL0 = tla::MakeLayoutL0C(mRound, nActual); + auto tensorL0CTile = tla::MakeTensor(l0CTensor[locPingPongFlag * mRound * 128], layoutCTileInL0, Arch::PositionL0C{}); + tileMmad( + tensorL0CTile, tensorL0A, tensorL0B, + mRound, nActual, kActual, initMmad, unitFlag); + if (endQItr) { + AscendC::SetFlag(0); + AscendC::SetFlag(1); + } + AscendC::SetFlag(l0ABPingPongFlag + 2); + auto layoutCInL0 = tla::MakeLayoutL0C(mActual, (uint32_t)256); + auto tensorL0C = tla::MakeTensor(l0CTensor, layoutCInL0, Arch::PositionL0C{}); + if (endItr) { + if constexpr (!ENABLE_UNIT_FLAG_) { + AscendC::SetFlag(l1KvPingPongFlag); + AscendC::WaitFlag(l1KvPingPongFlag); + copyL0CToGm(tensorC, tensorL0C); + AscendC::SetFlag(0); + AscendC::SetFlag(1); + } else { + copyL0CToGm(tensorC, tensorL0C, unitFlag); + } + } + } + +protected: + /// Data members + AscendC::LocalTensor l1ATensor; + AscendC::LocalTensor l1BTensor[STAGES]; + AscendC::LocalTensor l0ATensor[STAGES]; + AscendC::LocalTensor l0BTensor[STAGES]; + AscendC::LocalTensor l0CTensor; + + TileMmad tileMmad; + CopyL1ToL0A copyL1ToL0A; + CopyL1ToL0B copyL1ToL0B; +}; + +//////////////////////////////////////////////////////////////////// + +} // namespace Catlass::Gemm::Block + +#endif // CATLASS_GEMM_BLOCK_MMAD_QK_TLA_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_qk_tla.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_qk_tla.hpp new file mode 100644 index 0000000000..9903b37153 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_qk_tla.hpp @@ -0,0 +1,349 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ +#ifndef CATLASS_GEMM_BLOCK_BLOCK_MMAD_FAI_QK_TLA_HPP +#define CATLASS_GEMM_BLOCK_BLOCK_MMAD_FAI_QK_TLA_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/kda_coord.hpp" +#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" +#include "catlass/gemm/kda_gemm_helper.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" +#include "catlass/gemm/tile/kda_tile_mmad.hpp" +#include "tla/kda_tensor.hpp" +#include "tla/kda_tla_layout.hpp" + +//////////////////////////////////////////////////////////////////// + +namespace Catlass::Gemm::Block { +//////////////////////////////////////////////////////////////////// + +template < + bool PAGED_CACHE_FLAG_, + bool ENABLE_UNIT_FLAG_, + class L1TileShape_, + class L0TileShape_, + class ElementA_, + class ElementB_, + class ElementC_, + class ElementBias_, + class TileCopy_, + class TileMmad_> +struct BlockMmadTla< + MmadFAIQK, + L1TileShape_, + L0TileShape_, + ElementA_, + ElementB_, + ElementC_, + ElementBias_, + TileCopy_, + TileMmad_> { +public: + using DispatchPolicy = MmadFAIQK; + using ArchTag = typename DispatchPolicy::ArchTag; + using L1TileShape = L1TileShape_; + using L0TileShape = L0TileShape_; + using TileCopy = TileCopy_; + using TileMmad = TileMmad_; + + using ElementA = ElementA_; + using ElementB = ElementB_; + using ElementC = ElementC_; + using ElementAccumulator = typename TileCopy::ElementAccumulator; + + using CopyL1ToL0A = typename TileCopy::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopy::CopyL1ToL0B; + + using LayoutTagA = typename TileCopy::LayoutTagA; + using LayoutTagB = typename TileCopy::LayoutTagB; + using LayoutTagC = typename TileCopy::LayoutTagC; + + using LayoutTagL1A = typename TileCopy::LayoutTagL1A; + using LayoutTagL1B = typename TileCopy::LayoutTagL1B; + using LayoutTagL0A = typename TileCopy::LayoutTagL0A; + using LayoutTagL0B = typename TileCopy::LayoutTagL0B; + + static constexpr uint32_t STAGES = DispatchPolicy::STAGES; + // L1 tile shape + static constexpr uint32_t L1_TILE_M = tla::get<0>(L1TileShape{}); // s1 + static constexpr uint32_t L1_TILE_N = tla::get<1>(L1TileShape{}); // s2 + static constexpr uint32_t L1_TILE_K = tla::get<2>(L1TileShape{}); // d + // L0 tile shape + static constexpr uint32_t L0_TILE_M = tla::get<0>(L0TileShape{}); // s1 + static constexpr uint32_t L0_TILE_N = tla::get<1>(L0TileShape{}); // s2 + static constexpr uint32_t L0_TILE_K = tla::get<2>(L0TileShape{}); // d + // L1 tile size + static constexpr uint32_t L1A_TILE_SIZE = L1_TILE_M * L1_TILE_K * sizeof(ElementA); + static constexpr uint32_t L1B_TILE_SIZE = L1_TILE_N * L1_TILE_K * sizeof(ElementB); + // L0 tile size + static constexpr uint32_t L0A_TILE_SIZE = L0_TILE_M * L0_TILE_K * sizeof(ElementA); + static constexpr uint32_t L0B_TILE_SIZE = L0_TILE_K * L0_TILE_N * sizeof(ElementB); + static constexpr uint32_t L0C_TILE_SIZE = L0_TILE_M * L0_TILE_N * sizeof(ElementAccumulator); + + static constexpr uint32_t BLOCK_L1_SIZE = (L1A_TILE_SIZE + L1B_TILE_SIZE) * STAGES; + static constexpr uint32_t BLOCK_L0C_SIZE = L0C_TILE_SIZE * STAGES; + + // Check L1/L0TileShape + static_assert(L1_TILE_M == L0_TILE_M && L1_TILE_N == L0_TILE_N && L1_TILE_K == L0_TILE_K, + "The situation where the basic blocks of L1 and L0 differ on the m, n, k axes is not supported yet"); + static_assert(L0A_TILE_SIZE * STAGES <= ArchTag::L0A_SIZE, "L0TileShape exceeding the L0A space!"); + static_assert(L0B_TILE_SIZE * STAGES <= ArchTag::L0B_SIZE, "L0TileShape exceeding the L0B space!"); + +public: + /// Construct + CATLASS_DEVICE + BlockMmadTla(Arch::Resource &resource, uint32_t &l1BufAddrStart, uint32_t &l0CBufAddrStart) + { + uint32_t l1AOffset = l1BufAddrStart; + uint32_t l1BOffset = l1AOffset + L1A_TILE_SIZE * STAGES; + uint32_t l0COffset = l0CBufAddrStart; + + for (uint32_t i = 0; i < STAGES; i++) { + // Assign L1/L0A/L0B space for each stages + l1ATensorList_[i] = resource.l1Buf.template GetBufferByByte(l1AOffset + L1A_TILE_SIZE * i); + l1BTensorList_[i] = resource.l1Buf.template GetBufferByByte(l1BOffset + L1B_TILE_SIZE * i); + l0ATensorList_[i] = resource.l0ABuf.template GetBufferByByte(L0A_TILE_SIZE * i); + l0BTensorList_[i] = resource.l0BBuf.template GetBufferByByte(L0B_TILE_SIZE * i); + l0CTensorList_[i] = resource.l0CBuf.template GetBufferByByte(l0COffset + L0C_TILE_SIZE * i); + + // Assign event ID for each stages + l1AEventList_[i] = BLOCK_EVENT_ID + i; + l1BEventList_[i] = BLOCK_EVENT_ID + i + STAGES; + l0AEventList_[i] = i; + l0BEventList_[i] = i + STAGES; + l0CEventList_[i] = BLOCK_EVENT_ID + i; + + // The event id that needs to be set before the loop + AscendC::SetFlag(l1AEventList_[i]); + AscendC::SetFlag(l1BEventList_[i]); + AscendC::SetFlag(l0AEventList_[i]); + AscendC::SetFlag(l0BEventList_[i]); + AscendC::SetFlag(l0CEventList_[i]); + } + + l1BufAddrStart += BLOCK_L1_SIZE; + l0CBufAddrStart += BLOCK_L0C_SIZE; + } + + /// Destructor + CATLASS_DEVICE + ~BlockMmadTla() + { + for (uint32_t i = 0; i < STAGES; i++) { + AscendC::WaitFlag(l1AEventList_[i]); + AscendC::WaitFlag(l1BEventList_[i]); + AscendC::WaitFlag(l0AEventList_[i]); + AscendC::WaitFlag(l0BEventList_[i]); + AscendC::WaitFlag(l0CEventList_[i]); + } + } + +public: + template + CATLASS_DEVICE void operator()( + TensorA &tensorA, + TensorB &tensorB, + TensorC &tensorC, + AscendC::GlobalTensor &blockTable, + Shape &actualShape, + int64_t taskId, int32_t blockSize, + bool isFirstLoop = false, bool isLastLoop = false) + { + int32_t blockM = tla::get<0>(actualShape); + int32_t blockN = tla::get<1>(actualShape); + int32_t blockK = tla::get<2>(actualShape); + + AscendC::WaitFlag(l0CEventList_[l0CListId_]); + auto layoutInL0C = tla::MakeLayoutL0C(blockM, blockN); + auto tensorL0C = tla::MakeTensor(l0CTensorList_[l0CListId_], layoutInL0C, Arch::PositionL0C{}); + + if (unlikely(isFirstLoop)) { + AscendC::WaitFlag(l1AEventList_[l1AListId_]); + } + + int32_t kLoops = (blockK + L0_TILE_K - 1) / L0_TILE_K; + for (int32_t kIdx = 0; kIdx < kLoops; ++kIdx) { + int32_t tileK; + if (kIdx == kLoops - 1) { + int32_t tailSize = blockK % L0_TILE_K; + tileK = tailSize ? tailSize : L0_TILE_K; + } else { + tileK = L0_TILE_K; + } + + // load matrix A tile from GM to L1 + auto layoutAInL1 = tla::MakeLayout(blockM, blockK); + auto tensorL1A = tla::MakeTensor(l1ATensorList_[l1AListId_], layoutAInL1, Arch::PositionL1{}); + if (unlikely(isFirstLoop)) { + CopyInL1A(tensorL1A, tensorA, blockM, tileK, kIdx * L0_TILE_K); + AscendC::SetFlag(l1AEventList_[l1AListId_]); + AscendC::WaitFlag(l1AEventList_[l1AListId_]); + } + + // load matrix B tile from GM to L1 + auto layoutBInL1 = tla::MakeLayout(tileK, blockN); + auto tensorL1B = tla::MakeTensor(l1BTensorList_[l1BListId_], layoutBInL1, Arch::PositionL1{}); + AscendC::WaitFlag(l1BEventList_[l1BListId_]); + CopyInL1B(tensorL1B, tensorB, blockTable, tileK, blockN, kIdx * L0_TILE_K, blockSize); + AscendC::SetFlag(l1BEventList_[l1BListId_]); + AscendC::WaitFlag(l1BEventList_[l1BListId_]); + + // load matrix A tile from L1 to L0A + auto layoutAInL0 = tla::MakeLayout(blockM, tileK); + auto tensorL0A = tla::MakeTensor(l0ATensorList_[l0ListId_], layoutAInL0, Arch::PositionL0A{}); + auto tensorL1TileA = GetTile( + tensorL1A, + tla::MakeCoord(0, kIdx * L0_TILE_K), + tla::MakeShape(blockM, tileK)); + AscendC::WaitFlag(l0AEventList_[l0ListId_]); + copyL1ToL0A(tensorL0A, tensorL1TileA); + + // load matrix B tile from L1 to L0B + auto layoutBInL0 = tla::MakeLayout(tileK, blockN); + auto tensorL0B = tla::MakeTensor(l0BTensorList_[l0ListId_], layoutBInL0, Arch::PositionL0B{}); + auto tensorL1TileB = GetTile( + tensorL1B, + tla::MakeCoord(0, 0), + tla::MakeShape(tileK, blockN)); + AscendC::WaitFlag(l0BEventList_[l0ListId_]); + copyL1ToL0B(tensorL0B, tensorL1TileB); + + AscendC::SetFlag(l0CEventList_[l0CListId_]); + AscendC::WaitFlag(l0CEventList_[l0CListId_]); + + AscendC::SetFlag(l1BEventList_[l1BListId_]); + l1BListId_ = (l1BListId_ + 1 < STAGES) ? (l1BListId_ + 1) : 0; + + // Mmad + bool initC = kIdx == 0; + uint32_t tileM = (blockM == 1 ? M_ALIGN : blockM); + uint32_t tileN = blockN; + tileMmad(tensorL0C, tensorL0A, tensorL0B, tileM, tileN, tileK, initC); + + AscendC::SetFlag(l0AEventList_[l0ListId_]); + AscendC::SetFlag(l0BEventList_[l0ListId_]); + l0ListId_ = (l0ListId_ + 1 < STAGES) ? (l0ListId_ + 1) : 0; + } + + if (unlikely(isLastLoop)) { + AscendC::SetFlag(l1AEventList_[l1AListId_]); + l1AListId_ = (l1AListId_ + 1 < STAGES) ? (l1AListId_ + 1) : 0; + } + + AscendC::SetFlag(l0CEventList_[l0CListId_]); + AscendC::WaitFlag(l0CEventList_[l0CListId_]); + + AscendC::CrossCoreWaitFlag(MM1_RES_INTRA_EVENT[taskId]); + AscendC::CrossCoreWaitFlag(16 + MM1_RES_INTRA_EVENT[taskId]); + + // copy block out + using CopyL0CToDst = typename TileCopy::template CopyL0CToDst; + CopyL0CToDst copyL0CToDst; + copyL0CToDst(tensorC, tensorL0C); + + AscendC::SetFlag(l0CEventList_[l0CListId_]); + l0CListId_ = (l0CListId_ + 1 < STAGES) ? (l0CListId_ + 1) : 0; + } + +private: + template + CATLASS_DEVICE void CopyInL1A(TensorL1A &tensorL1A, TensorA &tensorA, int32_t tileM, int32_t tileK, int32_t kOffset) + { + using CopyGmToL1A = typename TileCopy::template CopyGmToL1A; + CopyGmToL1A copyGmToL1A; + + auto tensorTileA = GetTile( + tensorA, + tla::MakeCoord(0, kOffset), + tla::MakeShape(tileM, tileK)); + + auto tensorL1TileA = GetTile( + tensorL1A, + tla::MakeCoord(0, kOffset), + tla::MakeShape(tileM, tileK)); + + copyGmToL1A(tensorL1TileA, tensorTileA); + } + + template + CATLASS_DEVICE void CopyInL1B( + TensorL1B &tensorL1B, + TensorB &tensorB, + AscendC::GlobalTensor &blockTable, + int32_t tileK, int32_t tileN, + int32_t kOffset, int32_t blockSize) + { + using CopyGmToL1B = typename TileCopy::template CopyGmToL1B; + CopyGmToL1B copyGmToL1B; + + if constexpr (PAGED_CACHE_FLAG_) { + int32_t copyColCnt = 0; + int32_t blockLoops = (tileN + blockSize - 1) / blockSize; + + for (int32_t blockIdx = 0; blockIdx < blockLoops; ++blockIdx) { + int32_t curCopyCol = blockIdx == blockLoops - 1 ? tileN - copyColCnt : blockSize; + int32_t idInBlockTable = blockTable.GetValue(blockIdx); + + auto tensorL1TileB = GetTile( + tensorL1B, + tla::MakeCoord(0, copyColCnt), + tla::MakeShape(tileK, curCopyCol)); + + auto tensorTileB = GetTile( + tensorB, + tla::MakeCoord(kOffset, idInBlockTable * blockSize), + tla::MakeShape(tileK, curCopyCol)); + + copyGmToL1B(tensorL1TileB, tensorTileB); + copyColCnt += curCopyCol; + } + } else { + auto tensorTileB = GetTile( + tensorB, + tla::MakeCoord(kOffset, 0), + tla::MakeShape(tileK, tileN)); + + copyGmToL1B(tensorL1B, tensorTileB); + } + } + +private: + static constexpr uint16_t M_ALIGN = 16; + static constexpr uint32_t SYNC_MODE = 4; + static constexpr uint32_t MM1_RES_INTRA_EVENT[2] = {9, 10}; + static constexpr uint32_t BLOCK_EVENT_ID = 0; + + // Multi-stage tensors list + AscendC::LocalTensor l1ATensorList_[STAGES]; + AscendC::LocalTensor l1BTensorList_[STAGES]; + AscendC::LocalTensor l0ATensorList_[STAGES]; + AscendC::LocalTensor l0BTensorList_[STAGES]; + AscendC::LocalTensor l0CTensorList_[STAGES]; + + // Multi-stage event id list + int32_t l1AEventList_[STAGES]; + int32_t l1BEventList_[STAGES]; + int32_t l0AEventList_[STAGES]; + int32_t l0BEventList_[STAGES]; + int32_t l0CEventList_[STAGES]; + + uint32_t l1AListId_{0}; + uint32_t l1BListId_{0}; + uint32_t l0ListId_{0}; + uint32_t l0CListId_{0}; + + CopyL1ToL0A copyL1ToL0A; + CopyL1ToL0B copyL1ToL0B; + TileMmad tileMmad; +}; +} // namespace Catlass::Gemm::Block +#endif diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_pingpong_dequant_tla.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_pingpong_dequant_tla.hpp new file mode 100644 index 0000000000..d286e68a0d --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_pingpong_dequant_tla.hpp @@ -0,0 +1,666 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_GEMM_BLOCK_BLOCK_MMAD_PINGPONG_DEQUANT_TLA_HPP +#define CATLASS_GEMM_BLOCK_BLOCK_MMAD_PINGPONG_DEQUANT_TLA_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/kda_coord.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" +#include "catlass/gemm/kda_gemm_helper.hpp" +#include "tla/kda_tla_layout.hpp" +#include "tla/kda_tensor.hpp" + +namespace Catlass::Gemm::Block { + +template < + class ArchTag_, + bool ENABLE_UNIT_FLAG_, + bool USE_HF32_MODE_, + uint32_t L0C_STAGES_, + bool ENABLE_L1_RESIDENT_, + uint32_t L1A_STAGES_, + uint32_t L1B_STAGES_, + uint32_t L0A_STAGES_, + uint32_t L0B_STAGES_, + class L1TileShape_, + class L0TileShape_, + class ElementA_, + class ElementB_, + class ElementC_, + class ElementBias_, + class TileCopy_, + class TileMmad_> +struct BlockMmadTla< + MmadDequant, + L1TileShape_, + L0TileShape_, + ElementA_, + ElementB_, + ElementC_, + ElementBias_, + TileCopy_, + TileMmad_> { +public: + // Type Aliases + using DispatchPolicy = MmadDequant; + using ArchTag = typename DispatchPolicy::ArchTag; + using TileCopy = TileCopy_; + using L1TileShape = L1TileShape_; + using L0TileShape = L0TileShape_; + using ElementA = ElementA_; + using LayoutA = typename TileCopy::LayoutA; + using ElementB = ElementB_; + using LayoutB = typename TileCopy::LayoutB; + using ElementC = ElementC_; + using LayoutC = typename TileCopy::LayoutC; + using ElementBias = ElementBias_; + + using TileMmad = TileMmad_; + + using CopyL1ToL0A = typename TileCopy::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopy::CopyL1ToL0B; + using CopyL1ToBT = typename TileCopy::CopyL1ToBT; + + using ElementAccumulator = typename TileCopy::ElementAccumulator; + + static constexpr bool HAS_BIAS = TileCopy::HAS_BIAS; + static constexpr Tile::ScaleGranularity DEQUANT_MODE = TileCopy::DEQUANT_GRANULARITY; + static constexpr bool HAS_QUANT_TENSOR = TileCopy::HAS_QUANT_TENSOR; + + using LayoutTagL1A = typename TileCopy::LayoutTagL1A; + using LayoutTagL1B = typename TileCopy::LayoutTagL1B; + using LayoutTagL0A = typename TileCopy::LayoutTagL0A; + using LayoutTagL0B = typename TileCopy::LayoutTagL0B; + + static_assert(tla::is_tuple::value && tla::is_static::value, + "L1TileShape must be tla::tuple and static!"); + static_assert(tla::is_tuple::value && tla::is_static::value, + "L0TileShape must be tla::tuple and static!"); + + static constexpr bool ENABLE_UNIT_FLAG = DispatchPolicy::ENABLE_UNIT_FLAG; + static constexpr bool USE_HF32_MODE = DispatchPolicy::USE_HF32_MODE; + static constexpr bool ENABLE_L1_RESIDENT = DispatchPolicy::ENABLE_L1_RESIDENT; + static constexpr uint32_t L1A_STAGES = DispatchPolicy::L1A_STAGES; + static constexpr uint32_t L1B_STAGES = DispatchPolicy::L1B_STAGES; + static constexpr uint32_t L0A_STAGES = DispatchPolicy::L0A_STAGES; + static constexpr uint32_t L0B_STAGES = DispatchPolicy::L0B_STAGES; + static constexpr uint32_t L0C_STAGES = DispatchPolicy::L0C_STAGES; + static constexpr uint32_t L1_TILE_M = tla::get<0>(L1TileShape{}); + static constexpr uint32_t L1_TILE_N = tla::get<1>(L1TileShape{}); + static constexpr uint32_t L1_TILE_K = tla::get<2>(L1TileShape{}); + static constexpr uint32_t L0_TILE_M = tla::get<0>(L0TileShape{}); + static constexpr uint32_t L0_TILE_N = tla::get<1>(L0TileShape{}); + static constexpr uint32_t L0_TILE_K = tla::get<2>(L0TileShape{}); + + // L1 tile size + static constexpr uint32_t L1A_TILE_SIZE = L1_TILE_M * L1_TILE_K * sizeof(ElementA); + static constexpr uint32_t L1B_TILE_SIZE = L1_TILE_N * L1_TILE_K * sizeof(ElementB); + static constexpr uint32_t L1AB_TILE_SIZE = L1A_TILE_SIZE * L1A_STAGES + L1B_TILE_SIZE * L1B_STAGES; + // L0 tile size + static constexpr uint32_t L0A_TILE_SIZE = L0_TILE_M * L0_TILE_K * sizeof(ElementA); + static constexpr uint32_t L0B_TILE_SIZE = L0_TILE_K * L0_TILE_N * sizeof(ElementB); + static constexpr uint32_t L0C_TILE_SIZE = L1_TILE_M * L1_TILE_N * sizeof(ElementAccumulator); + static constexpr uint32_t BIAS_BUF_SIZE = L0_TILE_N * sizeof(ElementAccumulator); + + // Check HF32_MODE + static_assert( + !USE_HF32_MODE || (USE_HF32_MODE && std::is_same_v && std::is_same_v), + "HF32 MODE only supports in float!"); + + // Check L0C_STAGES + static_assert(!(ENABLE_UNIT_FLAG && L0C_STAGES != 1), "L0C_STAGES must be 1 when UnitFlag is true!"); + + // Check LayoutC + static_assert(tla::detail::isRowMajor::value || + ((std::is_same_v || std::is_same_v || + std::is_same_v) && + tla::detail::iszN::value), + "LayoutC only supports zN in half or bfloat16 or float, RowMajor in all dtype yet!"); + + // Check L1TileShape + static_assert(L1AB_TILE_SIZE <= ArchTag::L1_SIZE, "L1TileShape exceeding the L1 space!"); + + // Check L0TileShape + static_assert(L0A_TILE_SIZE * L0A_STAGES <= ArchTag::L0A_SIZE, "L0TileShape exceeding the L0A space!"); + static_assert(L0B_TILE_SIZE * L0B_STAGES <= ArchTag::L0B_SIZE, "L0TileShape exceeding the L0B space!"); + static_assert(L0C_TILE_SIZE * L0C_STAGES <= ArchTag::L0C_SIZE, "L0TileShape exceeding the L0C space!"); + static_assert(BIAS_BUF_SIZE <= ArchTag::BIAS_SIZE, "BIAS_BUF_SIZE exceeding the BT space! Reduce L0_TILE_N"); + + static_assert(L1_TILE_M == L0_TILE_M && L1_TILE_N == L0_TILE_N, + "The situation where the basic blocks of L1 and L0 differ on the m and n axes is not supported yet"); + static_assert(L0_TILE_K <= L1_TILE_K, "L0TileShape::K cannot exceed L1TileShape::K"); + + static_assert((!HAS_BIAS && !HAS_QUANT_TENSOR && (L1A_STAGES + L1B_STAGES) <= 8) || + (HAS_BIAS && !HAS_QUANT_TENSOR && (L1A_STAGES + L1B_STAGES) <= 7) || + (HAS_QUANT_TENSOR && (L1A_STAGES + L1B_STAGES) <= 6), + "L1 Buffer overflow: Exceeds the supported range of EVENT(0~7)"); + + static_assert((!HAS_BIAS && (L0A_STAGES + L0B_STAGES) <= 8) || (HAS_BIAS && (L0A_STAGES + L0B_STAGES) <= 7), + "L0 Buffer overflow: Exceeds the supported range of EVENT_ID(0~7)"); + + static constexpr auto L1A_LAYOUT = + tla::MakeLayout(tla::Int{}, tla::Int{}); + static constexpr auto L1B_LAYOUT = + tla::MakeLayout(tla::Int{}, tla::Int{}); + static constexpr auto L1BIAS_LAYOUT = tla::MakeLayout(tla::Int{}); + static constexpr auto L0BIAS_LAYOUT = tla::MakeLayout(tla::Int{}); + static constexpr auto L1QUANT_LAYOUT = tla::MakeLayout(tla::Int{}); + + // When enableing L1 resident mode, restore the pointer and coordinates that record the last state + // to the initial state. if two blockmmad instances need to be consecutively invoked at the kernel layer, + // RestoreStatus() must be inserted between them. + CATLASS_DEVICE + void RestoreStatus() + { + for (int i = 0; i < L1A_STAGES; ++i) { + lastAddrA[i] = nullptr; + lastCoordA[i] = MatrixCoord{0U, 0U}; + } + for (int i = 0; i < L1B_STAGES; ++i) { + lastAddrB[i] = nullptr; + lastCoordB[i] = MatrixCoord{0U, 0U}; + } + } + + /// Construct + CATLASS_DEVICE + BlockMmadTla(Arch::Resource &resource, uint32_t l1BufAddrStart = 0) + { + if ASCEND_IS_AIC { + // use HF32 when USE_HF32_MODE is true + if constexpr (USE_HF32_MODE) { + AscendC::SetHF32Mode(true); + } else { + AscendC::SetHF32Mode(false); + } + if constexpr (ENABLE_UNIT_FLAG && tla::detail::isRowMajor::value) { + AscendC::SetMMLayoutTransform(true); + } + uint32_t l1AOffset = l1BufAddrStart; + uint32_t l1BOffset = l1BufAddrStart + L1A_TILE_SIZE * L1A_STAGES; + uint32_t l1BiasOffset = l1BOffset + L1B_TILE_SIZE * L1B_STAGES; + uint32_t l1QuantOffset = l1BiasOffset; + if constexpr (HAS_BIAS) { + l1QuantOffset += L1_TILE_N * sizeof(ElementBias); + } + // Init buffers + for (uint32_t i = 0; i < L1A_STAGES; i++) { + // Assign L1/L0A/L0B space for each stages + l1ATensorList[i] = resource.l1Buf.template GetBufferByByte(l1AOffset + L1A_TILE_SIZE * i); + // Assign event ID for each stages + l1AEventList[i] = i; + // The event id that needs to be set before the loop + AscendC::SetFlag(l1AEventList[i]); + } + for (uint32_t i = 0; i < L1B_STAGES; i++) { + // Assign L1/L0A/L0B space for each stages + l1BTensorList[i] = resource.l1Buf.template GetBufferByByte(l1BOffset + L1B_TILE_SIZE * i); + // Assign event ID for each stages + l1BEventList[i] = i + L1A_STAGES; + // The event id that needs to be set before the loop + AscendC::SetFlag(l1BEventList[i]); + } + for (uint32_t i = 0; i < L0A_STAGES; i++) { + // Assign L1/L0A/L0B space for each stages + l0ATensorList[i] = resource.l0ABuf.template GetBufferByByte(L0A_TILE_SIZE * i); + // Assign event ID for each stages + l0AEventList[i] = i; + // The event id that needs to be set before the loop + AscendC::SetFlag(l0AEventList[i]); + } + for (uint32_t i = 0; i < L0B_STAGES; i++) { + // Assign L1/L0A/L0B space for each stages + l0BTensorList[i] = resource.l0BBuf.template GetBufferByByte(L0B_TILE_SIZE * i); + // Assign event ID for each stages + l0BEventList[i] = i + L0A_STAGES; + // The event id that needs to be set before the loop + AscendC::SetFlag(l0BEventList[i]); + } + if constexpr (!ENABLE_UNIT_FLAG) { + for (uint32_t i = 0; i < L0C_STAGES; i++) { + l0CTensorList[i] = resource.l0CBuf.template GetBufferByByte(L0C_TILE_SIZE * i); + l0CEventList[i] = i; + AscendC::SetFlag(l0CEventList[i]); + } + } else { + l0CTensorList[0] = resource.l0CBuf.template GetBufferByByte(0); + } + if constexpr (HAS_BIAS) { + l1BiasTensor = resource.l1Buf.template GetBufferByByte(l1BiasOffset); + l0BiasTensor = resource.btBuf.template GetBufferByByte(0); + AscendC::SetFlag(L1A_STAGES + L1B_STAGES); + AscendC::SetFlag(L0A_STAGES + L0B_STAGES); + } + if constexpr (HAS_QUANT_TENSOR) { + l1QuantTensor = resource.l1Buf.template GetBufferByByte(l1QuantOffset); + AscendC::SetFlag(QUANT_EVENT_ID); + } + + if constexpr (ENABLE_L1_RESIDENT) { + RestoreStatus(); + } + } + } + + /// Destructor + CATLASS_DEVICE + ~BlockMmadTla() + { + if ASCEND_IS_AIC { + if constexpr (USE_HF32_MODE) { + AscendC::SetHF32Mode(false); + } + if constexpr (ENABLE_UNIT_FLAG && tla::detail::isRowMajor::value) { + AscendC::SetMMLayoutTransform(false); + } + for (uint32_t i = 0; i < L1A_STAGES; i++) { + AscendC::WaitFlag(l1AEventList[i]); + } + for (uint32_t i = 0; i < L1B_STAGES; i++) { + AscendC::WaitFlag(l1BEventList[i]); + } + for (uint32_t i = 0; i < L0A_STAGES; i++) { + AscendC::WaitFlag(l0AEventList[i]); + } + for (uint32_t i = 0; i < L0B_STAGES; i++) { + AscendC::WaitFlag(l0BEventList[i]); + } + if constexpr (!ENABLE_UNIT_FLAG) { + for (uint32_t i = 0; i < L0C_STAGES; i++) { + AscendC::WaitFlag(l0CEventList[i]); + } + } + if constexpr (HAS_BIAS) { + AscendC::WaitFlag(L1A_STAGES + L1B_STAGES); + AscendC::WaitFlag(L0A_STAGES + L0B_STAGES); + } + if constexpr (HAS_QUANT_TENSOR) { + AscendC::WaitFlag(QUANT_EVENT_ID); + } + } + } + + /// Perform a block-scoped matrix multiply-accumulate + template + CATLASS_DEVICE void operator()(TensorA &tensorA, TensorB &tensorB, TensorC &tensorC, GemmCoord const &actualShape, + TensorBias const &tensorBias = {}, TensorQuant const &tensorQuant = {}, float scaleParam = 1.0) + { + // Check TileShape + if constexpr (HAS_BIAS) { + static constexpr uint32_t BIAS_BUF_SIZE = L0_TILE_N * sizeof(ElementAccumulator); + static constexpr uint32_t L1BIAS_SIZE = L1_TILE_N * sizeof(ElementBias); + static_assert(BIAS_BUF_SIZE <= ArchTag::BIAS_SIZE, + "BIAS_BUF_SIZE exceeding the BT space! Reduce L0_TILE_N"); + static_assert(L1AB_TILE_SIZE + L1BIAS_SIZE <= ArchTag::L1_SIZE, + "L1TileShape exceeding the L1 space!"); + } + if constexpr (HAS_QUANT_TENSOR) { + static constexpr uint32_t L1QUANT_SIZE = L1_TILE_N * sizeof(uint64_t); + if constexpr (HAS_BIAS) { + static constexpr uint32_t L1BIAS_SIZE = L1_TILE_N * sizeof(ElementBias); + static_assert(L1AB_TILE_SIZE + L1BIAS_SIZE + L1QUANT_SIZE <= ArchTag::L1_SIZE, + "L1TileShape exceeding the L1 space!"); + } else { + static_assert(L1AB_TILE_SIZE + L1QUANT_SIZE <= ArchTag::L1_SIZE, + "L1TileShape exceeding the L1 space!"); + } + } + + using CopyGmToL1A = typename TileCopy_::template CopyGmToL1A; + using CopyGmToL1B = typename TileCopy_::template CopyGmToL1B; + using CopyL0CToDst = typename TileCopy_::template CopyL0CToDst; + CopyGmToL1A copyGmToL1A; + CopyGmToL1B copyGmToL1B; + CopyL0CToDst copyL0CToDst; + + uint32_t mBlockActual = actualShape.m(); + uint32_t kBlockActual = actualShape.k(); + uint32_t nBlockActual = actualShape.n(); + + uint32_t mL1Actual = mBlockActual; + if constexpr (std::is_same_v) { + // Avoid using the gemv mode in mmad + if (mL1Actual == 1) { + mL1Actual = 16; + } + } + uint32_t nL1Actual = nBlockActual; + + auto layoutInL0C = tla::MakeLayoutL0C(mL1Actual, nL1Actual); + auto tensorL0C = tla::MakeTensor(l0CTensorList[l0CListId], layoutInL0C, Arch::PositionL0C{}); + auto tensorL0Bias = tla::MakeTensor(l0BiasTensor, L0BIAS_LAYOUT, Arch::PositionBias{}); + + uint32_t kL1Actual = min(kBlockActual, L1_TILE_K); + // load first matrix A tile from GM to L1 + AscendC::WaitFlag(l1AEventList[l1AListId]); + auto tensorL1A = tla::MakeTensor(l1ATensorList[l1AListId], L1A_LAYOUT, Arch::PositionL1{}); + auto tensorTileA = GetTileA(tensorA, 0, 0, mBlockActual, kL1Actual); + if constexpr (ENABLE_L1_RESIDENT) { + // If the currently loaded GM pointer and block coordinates are the same as the last loaded ones, + // skip this loading. + if (lastAddrA[l1AListId] != tensorTileA.data().GetPhyAddr() || tla::get<0>(tensorTileA.coord()) != lastCoordA[l1AListId].row() || tla::get<1>(tensorTileA.coord()) != lastCoordA[l1AListId].column()) { + copyGmToL1A(tensorL1A, tensorTileA); + lastCoordA[l1AListId] = MatrixCoord{tla::get<0>(tensorTileA.coord()), tla::get<1>(tensorTileA.coord())}; + lastAddrA[l1AListId] = const_cast<__gm__ typename AscendC::GlobalTensor::PrimType *>( + tensorTileA.data().GetPhyAddr()); + } + } else { + copyGmToL1A(tensorL1A, tensorTileA); + } + AscendC::SetFlag(l1AEventList[l1AListId]); + + // load first matrix B tile from GM to L1 + AscendC::WaitFlag(l1BEventList[l1BListId]); + auto tensorL1B = tla::MakeTensor(l1BTensorList[l1BListId], L1B_LAYOUT, Arch::PositionL1{}); + auto tensorTileB = GetTile(tensorB, tla::MakeCoord(0, 0), tla::MakeShape(kL1Actual, nBlockActual)); + if constexpr (ENABLE_L1_RESIDENT) { + if (lastAddrB[l1BListId] != tensorTileB.data().GetPhyAddr() || tla::get<0>(tensorTileB.coord()) != lastCoordB[l1BListId].row() || tla::get<1>(tensorTileB.coord()) != lastCoordB[l1BListId].column()) { + copyGmToL1B(tensorL1B, tensorTileB); + lastCoordB[l1BListId] = MatrixCoord{tla::get<0>(tensorTileB.coord()), tla::get<1>(tensorTileB.coord())}; + lastAddrB[l1BListId] = const_cast<__gm__ typename AscendC::GlobalTensor::PrimType *>( + tensorTileB.data().GetPhyAddr()); + } + } else { + copyGmToL1B(tensorL1B, tensorTileB); + } + AscendC::SetFlag(l1BEventList[l1BListId]); + + if constexpr (HAS_BIAS && !std::is_same_v) { + using CopyGmToL1Bias = typename TileCopy::template CopyGmToL1Bias; + CopyGmToL1Bias copyGmToL1Bias; + AscendC::WaitFlag(L1A_STAGES + L1B_STAGES); + auto l1Bias = l1BiasTensor.template ReinterpretCast(); + auto tensorL1Bias = tla::MakeTensor(l1Bias, L1BIAS_LAYOUT, Arch::PositionL1{}); + copyGmToL1Bias(tensorL1Bias, tensorBias); + AscendC::SetFlag(L1A_STAGES + L1B_STAGES); + } + + if constexpr (!ENABLE_UNIT_FLAG) { + AscendC::WaitFlag(l0CEventList[l0CListId]); + } + + uint32_t mL0Loop = CeilDiv(mL1Actual); + uint32_t nL0Loop = CeilDiv(nL1Actual); + + // main loop + uint32_t kL1Loop = CeilDiv(kBlockActual); + for (uint32_t kL1Idx = 0; kL1Idx < kL1Loop; kL1Idx++) { + uint32_t l1AListIdNext = (l1AListId + 1 < L1A_STAGES) ? (l1AListId + 1) : 0; + uint32_t l1BListIdNext = (l1BListId + 1 < L1B_STAGES) ? (l1BListId + 1) : 0; + uint32_t kL1ActualNext{0}; + // preload next tile from GM to L1 + if (kL1Idx < kL1Loop - 1) { + uint32_t kL1IdxNext = kL1Idx + 1; + kL1ActualNext = (kL1IdxNext < kL1Loop - 1) ? L1_TILE_K : (kBlockActual - kL1IdxNext * L1_TILE_K); + + // Get L1 tensor for next stage + auto l1ATensor = l1ATensorList[l1AListIdNext]; + auto l1BTensor = l1BTensorList[l1BListIdNext]; + auto tensorL1A = tla::MakeTensor(l1ATensor, L1A_LAYOUT, Arch::PositionL1{}); + auto tensorL1B = tla::MakeTensor(l1BTensor, L1B_LAYOUT, Arch::PositionL1{}); + // Get GM tile for next stage + auto tensorTileA = GetTileA(tensorA, 0, kL1IdxNext * L1_TILE_K, mBlockActual, kL1ActualNext); + auto tensorTileB = GetTile(tensorB, tla::MakeCoord(kL1IdxNext * L1_TILE_K, 0), + tla::MakeShape(kL1ActualNext, nBlockActual)); + + // load next matrix A tile from GM to L1 + AscendC::WaitFlag(l1AEventList[l1AListIdNext]); + if constexpr (ENABLE_L1_RESIDENT) { + if (lastAddrA[l1AListIdNext] != tensorTileA.data().GetPhyAddr() || tla::get<0>(tensorTileA.coord()) != lastCoordA[l1AListIdNext].row() || tla::get<1>(tensorTileA.coord()) != lastCoordA[l1AListIdNext].column()) { + copyGmToL1A(tensorL1A, tensorTileA); + lastCoordA[l1AListIdNext] = + MatrixCoord{tla::get<0>(tensorTileA.coord()), tla::get<1>(tensorTileA.coord())}; + lastAddrA[l1AListIdNext] = + const_cast<__gm__ typename AscendC::GlobalTensor::PrimType *>( + tensorTileA.data().GetPhyAddr()); + } + } else { + copyGmToL1A(tensorL1A, tensorTileA); + } + AscendC::SetFlag(l1AEventList[l1AListIdNext]); + + // load next matrix B tile from GM to L1 + AscendC::WaitFlag(l1BEventList[l1BListIdNext]); + if constexpr (ENABLE_L1_RESIDENT) { + if (lastAddrB[l1BListIdNext] != tensorTileB.data().GetPhyAddr() || tla::get<0>(tensorTileB.coord()) != lastCoordB[l1BListIdNext].row() || tla::get<1>(tensorTileB.coord()) != lastCoordB[l1BListIdNext].column()) { + copyGmToL1B(tensorL1B, tensorTileB); + lastCoordB[l1BListIdNext] = + MatrixCoord{tla::get<0>(tensorTileB.coord()), tla::get<1>(tensorTileB.coord())}; + lastAddrB[l1BListIdNext] = + const_cast<__gm__ typename AscendC::GlobalTensor::PrimType *>( + tensorTileB.data().GetPhyAddr()); + } + } else { + copyGmToL1B(tensorL1B, tensorTileB); + } + AscendC::SetFlag(l1BEventList[l1BListIdNext]); + } + + // Get L1 tensor for current stage + auto l1ATensor = l1ATensorList[l1AListId]; + auto l1BTensor = l1BTensorList[l1BListId]; + tensorL1A = tla::MakeTensor(l1ATensor, L1A_LAYOUT, Arch::PositionL1{}); + tensorL1B = tla::MakeTensor(l1BTensor, L1B_LAYOUT, Arch::PositionL1{}); + // Get the loop nums on L0 + uint32_t kL0Loop = CeilDiv(kL1Actual); + + for (int mL0Idx = 0; mL0Idx < mL0Loop; mL0Idx++) { + uint32_t mL0Actual = (mL0Idx < mL0Loop - 1) ? L0_TILE_M : (mL1Actual - mL0Idx * L0_TILE_M); + + for (int kL0Idx = 0; kL0Idx < kL0Loop; kL0Idx++) { + uint32_t kL0Actual = (kL0Idx < kL0Loop - 1) ? L0_TILE_K : (kL1Actual - kL0Idx * L0_TILE_K); + + // Locate the current tile on L0A + auto l0ATile = l0ATensorList[l0AListId]; + auto layoutAInL0 = tla::MakeLayout(mL0Actual, kL0Actual); + auto tensorL0A = tla::MakeTensor(l0ATile, layoutAInL0, Arch::PositionL0A{}); + // Locate the current tile of matrix A on L1 + auto tensorTileL1A = GetTileA(tensorL1A, mL0Idx * L0_TILE_M, kL0Idx * L0_TILE_K, mL0Actual, kL0Actual); + + AscendC::WaitFlag(l0AEventList[l0AListId]); + if ((mL0Idx == 0) && (kL0Idx == 0)) { + AscendC::WaitFlag(l1AEventList[l1AListId]); + } + + // Load current tile from L1 to L0A + copyL1ToL0A(tensorL0A, tensorTileL1A); + + if ((mL0Idx == mL0Loop - 1) && (kL0Idx == kL0Loop - 1)) { + AscendC::SetFlag(l1AEventList[l1AListId]); + } + + bool initC = ((kL1Idx == 0) && (kL0Idx == 0)); + for (int nL0Idx = 0; nL0Idx < nL0Loop; nL0Idx++) { + uint32_t nL0Actual = (nL0Idx < nL0Loop - 1) ? L0_TILE_N : (nL1Actual - nL0Idx * L0_TILE_N); + + // Locate the current tile on L0B + auto l0BTile = l0BTensorList[l0BListId]; + auto layoutBInL0 = tla::MakeLayout(kL0Actual, nL0Actual); + auto tensorL0B = tla::MakeTensor(l0BTile, layoutBInL0, Arch::PositionL0B{}); + // Locate the current tile of matrix B on L1 + auto tensorTileL1B = GetTile(tensorL1B, + tla::MakeCoord(kL0Idx * L0_TILE_K, nL0Idx * L0_TILE_N), + tla::MakeShape(kL0Actual, nL0Actual)); + + // Wait for mmad finished + AscendC::WaitFlag(l0BEventList[l0BListId]); + // If the current tile is the first one on the k&n axis, wait for loading matrix B from GM to L1 + if ((mL0Idx == 0) && (kL0Idx == 0) && (nL0Idx == 0)) { + AscendC::WaitFlag(l1BEventList[l1BListId]); + } + + // Load current tile from L1 to L0B + copyL1ToL0B(tensorL0B, tensorTileL1B); + + // If the current tile is the last one on the k&n axis, notify to load matrix B from GM to L1 + if ((mL0Idx == mL0Loop - 1) && (kL0Idx == kL0Loop - 1) && (nL0Idx == nL0Loop - 1)) { + AscendC::SetFlag(l1BEventList[l1BListId]); + } + + if constexpr (HAS_BIAS && !std::is_same_v) { + if (initC) { + if (nL0Idx == 0) { + AscendC::WaitFlag(L1A_STAGES + L1B_STAGES); + } + AscendC::WaitFlag(L0A_STAGES + L0B_STAGES); + auto l1Bias = l1BiasTensor.template ReinterpretCast(); + auto tensorL1Bias = tla::MakeTensor(l1Bias, L1BIAS_LAYOUT, Arch::PositionL1{}); + auto tensorTileL1Bias = GetTile(tensorL1Bias, + tla::MakeCoord(nL0Idx * L0_TILE_N), + tla::MakeShape(nL0Actual)); + // Load bias to l0 biastable + copyL1ToBT(tensorL0Bias, tensorTileL1Bias); + if (nL0Idx == nL0Loop - 1) { + AscendC::SetFlag(L1A_STAGES + L1B_STAGES); + } + } + } + + // Notify to do mmad + AscendC::SetFlag(l0CEventList[l0CListId]); + + // Locate the current tile on L0C + auto tensorTileL0C = GetTile(tensorL0C, + tla::MakeCoord(mL0Idx * L0_TILE_M, nL0Idx * L0_TILE_N), + tla::MakeShape(mL0Actual, nL0Actual)); + + // Compute the matrix multiplication on L0A and L0B and write the result to the accumulator + // Wait for loading L0B + AscendC::WaitFlag(l0CEventList[l0CListId]); + + // If the unit flag is enabled, the unit flag is set according to the calculation progress + uint8_t unitFlag = 0b00; + if constexpr (ENABLE_UNIT_FLAG) { + if ((kL1Idx == kL1Loop - 1) && (mL0Idx == mL0Loop - 1) && + (kL0Idx == kL0Loop - 1) && (nL0Idx == nL0Loop - 1)) { + unitFlag = 0b11; + } else { + unitFlag = 0b10; + } + } + + if constexpr (HAS_BIAS && !std::is_same_v) { + if (initC) { + tileMmad(tensorTileL0C, tensorL0A, tensorL0B, tensorL0Bias, + mL0Actual, nL0Actual, kL0Actual, initC, unitFlag); + AscendC::SetFlag(L0A_STAGES + L0B_STAGES); + } else { + tileMmad(tensorTileL0C, tensorL0A, tensorL0B, + mL0Actual, nL0Actual, kL0Actual, initC, unitFlag); + } + } else { + tileMmad(tensorTileL0C, tensorL0A, tensorL0B, + mL0Actual, nL0Actual, kL0Actual, initC, unitFlag); + } + + // Notify to move the next L0B tile + AscendC::SetFlag(l0BEventList[l0BListId]); + l0BListId = (l0BListId + 1 < L0B_STAGES) ? (l0BListId + 1) : 0; + } + AscendC::SetFlag(l0AEventList[l0AListId]); + l0AListId = (l0AListId + 1 < L0A_STAGES) ? (l0AListId + 1) : 0; + } + } + l1AListId = l1AListIdNext; + l1BListId = l1BListIdNext; + kL1Actual = kL1ActualNext; + } + + // copy block out + if constexpr (DEQUANT_MODE == Tile::ScaleGranularity::PER_TENSOR) { + copyL0CToDst.params.scale = scaleParam; + } + + if constexpr (HAS_QUANT_TENSOR && !std::is_same_v) { + using CopyGmToL1Scale = typename TileCopy::template CopyGmToL1Scale; + CopyGmToL1Scale copyGmToL1Scale; + auto tensorL1Quant = tla::MakeTensor(l1QuantTensor, L1QUANT_LAYOUT, Arch::PositionL1{}); + AscendC::WaitFlag(QUANT_EVENT_ID); + copyGmToL1Scale(tensorL1Quant, tensorQuant); + AscendC::SetFlag(QUANT_EVENT_ID); + AscendC::WaitFlag(QUANT_EVENT_ID); + if constexpr (!ENABLE_UNIT_FLAG) { + AscendC::SetFlag(l0CEventList[l0CListId]); + AscendC::WaitFlag(l0CEventList[l0CListId]); + copyL0CToDst(tensorC, tensorL0C, tensorL1Quant); + AscendC::SetFlag(l0CEventList[l0CListId]); + l0CListId = (l0CListId + 1 < L0C_STAGES) ? (l0CListId + 1) : 0; + } else { + copyL0CToDst(tensorC, tensorL0C, tensorL1Quant, 0b11); + } + AscendC::SetFlag(QUANT_EVENT_ID); + } else { + if constexpr (!ENABLE_UNIT_FLAG) { + AscendC::SetFlag(l0CEventList[l0CListId]); + AscendC::WaitFlag(l0CEventList[l0CListId]); + copyL0CToDst(tensorC, tensorL0C); + AscendC::SetFlag(l0CEventList[l0CListId]); + l0CListId = (l0CListId + 1 < L0C_STAGES) ? (l0CListId + 1) : 0; + } else { + copyL0CToDst(tensorC, tensorL0C, 0b11); + } + } + } + +protected: + template + CATLASS_DEVICE auto GetTileA(TensorA &tensorA, uint32_t mIndex, uint32_t kIndex, uint32_t mSize, uint32_t kSize) + { + if constexpr (tla::detail::isVector::value) { + return GetTile(tensorA, tla::MakeCoord(kIndex), tla::MakeShape(kSize)); + } else { + return GetTile(tensorA, tla::MakeCoord(mIndex, kIndex), tla::MakeShape(mSize, kSize)); + } + } + + // Multi-stage tensors list + AscendC::LocalTensor l1ATensorList[L1A_STAGES]; + AscendC::LocalTensor l1BTensorList[L1B_STAGES]; + AscendC::LocalTensor l0ATensorList[L0A_STAGES]; + AscendC::LocalTensor l0BTensorList[L0B_STAGES]; + AscendC::LocalTensor l0CTensorList[L0C_STAGES]; + AscendC::LocalTensor l1BiasTensor; + AscendC::LocalTensor l0BiasTensor; + AscendC::LocalTensor l1QuantTensor; + + // Multi-stage event id list + int32_t l1AEventList[L1A_STAGES]; + int32_t l1BEventList[L1B_STAGES]; + int32_t l0AEventList[L0A_STAGES]; + int32_t l0BEventList[L0B_STAGES]; + int32_t l0CEventList[L0C_STAGES]; + static constexpr int32_t QUANT_EVENT_ID = L1A_STAGES + L1B_STAGES + 1; + + __gm__ typename AscendC::GlobalTensor::PrimType *lastAddrA[L1A_STAGES]; + __gm__ typename AscendC::GlobalTensor::PrimType *lastAddrB[L1B_STAGES]; + MatrixCoord lastCoordA[L1A_STAGES]; + MatrixCoord lastCoordB[L1B_STAGES]; + + // The id of current stage + uint32_t l1AListId{0}; + uint32_t l1BListId{0}; + uint32_t l0AListId{0}; + uint32_t l0BListId{0}; + uint32_t l0CListId{0}; + + TileMmad tileMmad; + CopyL1ToL0A copyL1ToL0A; + CopyL1ToL0B copyL1ToL0B; + CopyL1ToBT copyL1ToBT; +}; + +} // namespace Catlass::Gemm::Block + +#endif // CATLASS_GEMM_BLOCK_BLOCK_MMAD_PINGPONG_DEQUANT_TLA_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_pingpong_full_loadA.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_pingpong_full_loadA.hpp new file mode 100644 index 0000000000..802dc7a5bc --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_pingpong_full_loadA.hpp @@ -0,0 +1,347 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_GEMM_BLOCK_BLOCK_MMAD_PINGPONG_FULL_LOADA_HPP +#define CATLASS_GEMM_BLOCK_BLOCK_MMAD_PINGPONG_FULL_LOADA_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/kda_coord.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" +#include "catlass/gemm/kda_gemm_helper.hpp" +#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" +#include "catlass/gemm/tile/kda_tile_mmad.hpp" + +namespace Catlass::Gemm::Block { + +template < + bool ENABLE_UNIT_FLAG_, + class L1TileShape_, + class L0TileShape_, + class AType_, + class BType_, + class CType_, + class BiasType_, + class TileCopy_, + class TileMmad_> +struct BlockMmad< + MmadAtlasA2FullLoadA, + L1TileShape_, + L0TileShape_, + AType_, + BType_, + CType_, + BiasType_, + TileCopy_, + TileMmad_> { +public: + // Type Aliases + using DispatchPolicy = MmadAtlasA2FullLoadA; + using ArchTag = typename DispatchPolicy::ArchTag; + using L1TileShape = L1TileShape_; + using L0TileShape = L0TileShape_; + using ElementA = typename AType_::Element; + using LayoutA = typename AType_::Layout; + using ElementB = typename BType_::Element; + using LayoutB = typename BType_::Layout; + using ElementC = typename CType_::Element; + using LayoutC = typename CType_::Layout; + using TileMmad = TileMmad_; + using CopyGmToL1A = typename TileCopy_::CopyGmToL1A; + using CopyGmToL1B = typename TileCopy_::CopyGmToL1B; + using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; + using CopyL0CToGm = typename TileCopy_::CopyL0CToGm; + using ElementAccumulator = + typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; + using LayoutAInL1 = typename CopyL1ToL0A::LayoutSrc; + using LayoutBInL1 = typename CopyL1ToL0B::LayoutSrc; + using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst; + using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst; + using LayoutCInL0 = layout::zN; + + using L1AAlignHelper = Gemm::helper::L1AlignHelper; + using L1BAlignHelper = Gemm::helper::L1AlignHelper; + + static constexpr bool ENABLE_UNIT_FLAG = DispatchPolicy::ENABLE_UNIT_FLAG; + static constexpr uint32_t STAGES = DispatchPolicy::STAGES; + static constexpr uint32_t L1A_SIZE = ArchTag::L1_SIZE / 2; + static constexpr uint32_t L1B_SIZE = L1TileShape::N * L1TileShape::K * sizeof(ElementB); + static constexpr uint32_t L0A_SIZE = ArchTag::L0A_SIZE; + static constexpr uint32_t L0B_SIZE = ArchTag::L0B_SIZE; + static constexpr uint32_t L0C_SIZE = ArchTag::L0C_SIZE; + static constexpr uint32_t L0A_PINGPONG_BUF_SIZE = L0A_SIZE / STAGES; + static constexpr uint32_t L0B_PINGPONG_BUF_SIZE = L0B_SIZE / STAGES; + + // Check LayoutC + static_assert(std::is_same_v, "LayoutC only support RowMajor yet!"); + + // Check L1TileShape in example, because problemShape.k() cannot get here. + + // Check L0TileShape + static constexpr uint32_t L0A_TILE_SIZE = L0TileShape::M * L0TileShape::K * sizeof(ElementA); + static constexpr uint32_t L0B_TILE_SIZE = L0TileShape::K * L0TileShape::N * sizeof(ElementB); + static constexpr uint32_t L0C_TILE_SIZE = L0TileShape::M * L0TileShape::N * sizeof(ElementAccumulator); + static_assert((L0A_TILE_SIZE * STAGES) <= L0A_SIZE, "L0TileShape exceeding the L0A space!"); + static_assert((L0B_TILE_SIZE * STAGES) <= L0B_SIZE, "L0TileShape exceeding the L0B space!"); + static_assert(L0C_TILE_SIZE <= L0C_SIZE, "L0TileShape exceeding the L0C space!"); + + static_assert(L1TileShape::M == L0TileShape::M && L1TileShape::N == L0TileShape::N, + "The situation where the basic blocks of L1 and L0 differ on the m and n axes is not supported yet"); + // 32B (256b) aligned + static_assert(Gemm::helper::TileShapeAlignChecker::_ALIGN == 256, + "Tile shape must be 32B aligned."); + + /// Construct + CATLASS_DEVICE + BlockMmad(Arch::Resource &resource, uint32_t l1BufAddrStart = 0) + { + uint32_t l1AOffset = l1BufAddrStart; + uint32_t l1BOffset = l1BufAddrStart + L1A_SIZE; + // Init L1A related + l1ATensorList[0] = resource.l1Buf.template GetBufferByByte(l1AOffset); + l1AEventList[0] = 0; + + // Init buffers + for (uint32_t i = 0; i < STAGES; i++) { + // Assign L1/L0A/L0B space for each stages + l1BTensorList[i] = resource.l1Buf.template GetBufferByByte(l1BOffset + L1B_SIZE * i); + l0ATensorList[i] = resource.l0ABuf.template GetBufferByByte(L0A_PINGPONG_BUF_SIZE * i); + l0BTensorList[i] = resource.l0BBuf.template GetBufferByByte(L0B_PINGPONG_BUF_SIZE * i); + + // Assign event ID for each stages + l1BEventList[i] = i + STAGES; + l0AEventList[i] = i; + l0BEventList[i] = i + STAGES; + + // The event id that needs to be set before the loop + AscendC::SetFlag(l1BEventList[i]); + AscendC::SetFlag(l0AEventList[i]); + AscendC::SetFlag(l0BEventList[i]); + } + l0CTensor = resource.l0CBuf.template GetBufferByByte(0); + AscendC::SetFlag(EVENT_ID0); + } + + /// Destructor + CATLASS_DEVICE + ~BlockMmad() + { + for (uint32_t i = 0; i < STAGES; i++) { + AscendC::WaitFlag(l1BEventList[i]); + AscendC::WaitFlag(l0AEventList[i]); + AscendC::WaitFlag(l0BEventList[i]); + } + AscendC::WaitFlag(EVENT_ID0); + } + + /// Perform a block-scoped matrix multiply-accumulate + CATLASS_DEVICE + void operator()( + AscendC::GlobalTensor const &gmA, LayoutA const &layoutA, + AscendC::GlobalTensor const &gmB, LayoutB const &layoutB, + AscendC::GlobalTensor const &gmC, LayoutC const &layoutC, + GemmCoord const &actualShape, bool needLoadL1) + { + uint32_t mRound = RoundUp(actualShape.m()); + uint32_t nRound = RoundUp(actualShape.n()); + + auto layoutAInL1 = LayoutAInL1::template MakeLayout(L1TileShape::M, actualShape.k()); + auto layoutBInL1 = LayoutBInL1::template MakeLayout(L1TileShape::K, L1TileShape::N); + auto layoutInL0C = LayoutCInL0::MakeLayoutInL0C(MakeCoord(mRound, nRound)); + + uint32_t kActual = min(actualShape.k(), L1TileShape::K); + + if (needLoadL1) { + // load first matrix A tile from GM to L1 + auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShape.m(), actualShape.k())); + AscendC::SetFlag(l1AEventList[0]); + AscendC::WaitFlag(l1AEventList[0]); + copyGmToL1A(l1ATensorList[0], gmA, layoutAInL1, layoutTileA); + AscendC::SetFlag(l1AEventList[0]); + AscendC::WaitFlag(l1AEventList[0]); + } + + // load first matrix B tile from GM to L1 + AscendC::WaitFlag(l1BEventList[l1ListId]); + auto layoutTileB = layoutB.GetTileLayout(MakeCoord(kActual, actualShape.n())); + copyGmToL1B(l1BTensorList[l1ListId], gmB, layoutBInL1, layoutTileB); + AscendC::SetFlag(l1BEventList[l1ListId]); + + if constexpr (!ENABLE_UNIT_FLAG) { + AscendC::WaitFlag(EVENT_ID0); + } + + uint32_t mPartLoop = CeilDiv(mRound); + uint32_t nPartLoop = CeilDiv(nRound); + + // main loop + uint32_t kTileCount = CeilDiv(actualShape.k()); + for (uint32_t kLoopIdx = 0; kLoopIdx < kTileCount; kLoopIdx++) { + uint32_t l1ListIdNext = (l1ListId + 1 < STAGES) ? (l1ListId + 1) : 0; + uint32_t kActualNext{0}; + // preload next tile from GM to L1 + if (kLoopIdx < kTileCount - 1) { + uint32_t kLoopIdxNext = kLoopIdx + 1; + kActualNext = (kLoopIdxNext < kTileCount - 1) ? + L1TileShape::K : + (actualShape.k() - kLoopIdxNext * L1TileShape::K); + + // Get L1 tensor for next stage + auto l1BTensor = l1BTensorList[l1ListIdNext]; + // Get GM tile for next stage + MatrixCoord gmTileBOffset{kLoopIdxNext * L1TileShape::K, 0}; + auto gmTileB = gmB[layoutB.GetOffset(gmTileBOffset)]; + + // load next matrix B tile from GM to L1 + AscendC::WaitFlag(l1BEventList[l1ListIdNext]); + layoutTileB = layoutB.GetTileLayout(MakeCoord(kActualNext, actualShape.n())); + copyGmToL1B(l1BTensor, gmTileB, layoutBInL1, layoutTileB); + AscendC::SetFlag(l1BEventList[l1ListIdNext]); + } + + // Get L1 tensor for current stage + auto l1ATensor = l1ATensorList[0]; + auto l1BTensor = l1BTensorList[l1ListId]; + + // Get the loop nums on L0 + uint32_t kPartLoop = CeilDiv(kActual); + + for (int mPartIdx = 0; mPartIdx < mPartLoop; mPartIdx++) { + uint32_t mPartActual = (mPartIdx < mPartLoop - 1) ? + L0TileShape::M : + (mRound - mPartIdx * L0TileShape::M); + + for (int kPartIdx = 0; kPartIdx < kPartLoop; kPartIdx++) { + uint32_t kPartActual = (kPartIdx < kPartLoop - 1) ? + L0TileShape::K : + (kActual - kPartIdx * L0TileShape::K); + + // Locate the current tile on L0A + auto l0ATile = l0ATensorList[l0AListId]; + LayoutAInL0 layoutAInL0 = LayoutAInL0::template MakeLayout(mPartActual, kPartActual); + // Locate the current tile of matrix A on L1 + MatrixCoord l1AOffset{mPartIdx * L0TileShape::M, + kPartIdx * L0TileShape::K + kLoopIdx * L1TileShape::K}; + auto l1ATile = l1ATensor[layoutAInL1.GetOffset(l1AOffset)]; + + AscendC::WaitFlag(l0AEventList[l0AListId]); + + // Load current tile from L1 to L0A + copyL1ToL0A(l0ATile, l1ATile, layoutAInL0, layoutAInL1); + + for (int nPartIdx = 0; nPartIdx < nPartLoop; nPartIdx++) { + uint32_t nPartActual = (nPartIdx < nPartLoop - 1) ? + L0TileShape::N : + (nRound - nPartIdx * L0TileShape::N); + + // Locate the current tile on L0B + auto l0BTile = l0BTensorList[l0BListId]; + LayoutBInL0 layoutBInL0 = LayoutBInL0::template MakeLayout(kPartActual, nPartActual); + // Locate the current tile of matrix B on L1 + MatrixCoord l1BOffset{kPartIdx * L0TileShape::K, nPartIdx * L0TileShape::N}; + auto l1BTile = l1BTensor[layoutBInL1.GetOffset(l1BOffset)]; + + // Wait for mmad finished + AscendC::WaitFlag(l0BEventList[l0BListId]); + // If the current tile is the first one on the k&n axis, wait for loading matrix B from GM to L1 + if ((kPartIdx == 0) && (nPartIdx == 0)) { + AscendC::WaitFlag(l1BEventList[l1ListId]); + } + + // Load current tile from L1 to L0B + copyL1ToL0B(l0BTile, l1BTile, layoutBInL0, layoutBInL1); + + // If the current tile is the last one on the k&n axis, notify to load matrix B from GM to L1 + if ((kPartIdx == kPartLoop - 1) && (nPartIdx == nPartLoop - 1)) { + AscendC::SetFlag(l1BEventList[l1ListId]); + } + // Notify to do mmad + AscendC::SetFlag(EVENT_ID0); + + // Locate the current tile on L0C + MatrixCoord l0COffset{mPartIdx * L0TileShape::M, nPartIdx * L0TileShape::N}; + auto l0CTile = l0CTensor[layoutInL0C.GetOffset(l0COffset)]; + + // Compute the matrix multiplication on L0A and L0B and write the result to the accumulator + // Wait for loading L0B + AscendC::WaitFlag(EVENT_ID0); + + // If the current tile is the first tile on the k axis, the accumulator needs to be reset to 0 + bool initC = ((kLoopIdx == 0) && (kPartIdx == 0)); + // If the unit flag is enabled, the unit flag is set according to the calculation progress + uint8_t unitFlag = 0b00; + if constexpr (ENABLE_UNIT_FLAG) { + if ((kLoopIdx == kTileCount - 1) && (mPartIdx == mPartLoop - 1) && + (kPartIdx == kPartLoop - 1) && (nPartIdx == nPartLoop - 1)) { + unitFlag = 0b11; + } else { + unitFlag = 0b10; + } + } + // Perform calculation operations + tileMmad(l0CTile, l0ATile, l0BTile, mPartActual, nPartActual, kPartActual, initC, unitFlag); + + // Notify to move the next L0B tile + AscendC::SetFlag(l0BEventList[l0BListId]); + l0BListId = (l0BListId + 1 < STAGES) ? (l0BListId + 1) : 0; + } + AscendC::SetFlag(l0AEventList[l0AListId]); + l0AListId = (l0AListId + 1 < STAGES) ? (l0AListId + 1) : 0; + } + } + l1ListId = l1ListIdNext; + kActual = kActualNext; + } + + // copy block out + LayoutC layoutBlock = layoutC.GetTileLayout(actualShape.GetCoordMN()); + + if constexpr (!ENABLE_UNIT_FLAG) { + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + copyL0CToGm(gmC, l0CTensor, layoutBlock, layoutInL0C); + AscendC::SetFlag(EVENT_ID0); + } else { + copyL0CToGm(gmC, l0CTensor, layoutBlock, layoutInL0C, 0b11); + } + } + +protected: + // Multi-stage tensors list + AscendC::LocalTensor l1ATensorList[STAGES]; + AscendC::LocalTensor l1BTensorList[STAGES]; + AscendC::LocalTensor l0ATensorList[STAGES]; + AscendC::LocalTensor l0BTensorList[STAGES]; + AscendC::LocalTensor l0CTensor; + + // Multi-stage event id list + int32_t l1AEventList[STAGES]; + int32_t l1BEventList[STAGES]; + int32_t l0AEventList[STAGES]; + int32_t l0BEventList[STAGES]; + + // The id of current stage + uint32_t l1ListId{0}; + uint32_t l0AListId{0}; + uint32_t l0BListId{0}; + + TileMmad tileMmad; + CopyGmToL1A copyGmToL1A; + CopyGmToL1B copyGmToL1B; + CopyL1ToL0A copyL1ToL0A; + CopyL1ToL0B copyL1ToL0B; + CopyL0CToGm copyL0CToGm; +}; + +} // namespace Catlass::Gemm::Block + +#endif // CATLASS_GEMM_BLOCK_BLOCK_MMAD_PINGPONG_FULL_LOADA_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_pingpong_full_loadA_tla.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_pingpong_full_loadA_tla.hpp new file mode 100644 index 0000000000..fd323844ad --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_pingpong_full_loadA_tla.hpp @@ -0,0 +1,582 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ +#ifndef CATLASS_GEMM_BLOCK_BLOCK_MMAD_PINGPONG_FULL_LOADA_TLA_HPP +#define CATLASS_GEMM_BLOCK_BLOCK_MMAD_PINGPONG_FULL_LOADA_TLA_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/kda_coord.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" +#include "catlass/gemm/kda_gemm_helper.hpp" +#include "tla/kda_tla_layout.hpp" +#include "tla/kda_tensor.hpp" + +namespace Catlass::Gemm::Block { + +template < + class ArchTag_, + bool ENABLE_UNIT_FLAG_, + bool ENABLE_SHUFFLE_K_, + bool USE_HF32_MODE_, + uint32_t L0C_STAGES_, + bool ENABLE_L1_RESIDENT_, + uint32_t L1A_STAGES_, + uint32_t L1B_STAGES_, + uint32_t L0A_STAGES_, + uint32_t L0B_STAGES_, + class L1TileShape_, + class L0TileShape_, + class ElementA_, + class ElementB_, + class ElementC_, + class ElementBias_, + class TileCopy_, + class TileMmad_> +struct BlockMmadTla< + MmadFullLoadA, + L1TileShape_, + L0TileShape_, + ElementA_, + ElementB_, + ElementC_, + ElementBias_, + TileCopy_, + TileMmad_> { +public: + // Type Aliases + using DispatchPolicy = MmadFullLoadA; + using ArchTag = typename DispatchPolicy::ArchTag; + using TileCopy = TileCopy_; + using L1TileShape = L1TileShape_; + using L0TileShape = L0TileShape_; + using ElementA = ElementA_; + using LayoutA = typename TileCopy::LayoutA; + using ElementB = ElementB_; + using LayoutB = typename TileCopy::LayoutB; + using ElementC = ElementC_; + using LayoutC = typename TileCopy::LayoutC; + using ElementBias = ElementBias_; + + using TileMmad = TileMmad_; + + using CopyL1ToL0A = typename TileCopy::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopy::CopyL1ToL0B; + using CopyL1ToBT = typename TileCopy::CopyL1ToBT; + + using ElementAccumulator = typename TileCopy::ElementAccumulator; + + static constexpr bool HAS_BIAS = TileCopy::HAS_BIAS; + + using LayoutTagL1A = typename TileCopy::LayoutTagL1A; + using LayoutTagL1B = typename TileCopy::LayoutTagL1B; + using LayoutTagL0A = typename TileCopy::LayoutTagL0A; + using LayoutTagL0B = typename TileCopy::LayoutTagL0B; + + static_assert(tla::is_tuple::value && tla::is_static::value, + "L1TileShape must be tla::tuple and static!"); + static_assert(tla::is_tuple::value && tla::is_static::value, + "L0TileShape must be tla::tuple and static!"); + + static constexpr bool ENABLE_UNIT_FLAG = DispatchPolicy::ENABLE_UNIT_FLAG; + static constexpr bool ENABLE_SHUFFLE_K = DispatchPolicy::ENABLE_SHUFFLE_K; + static constexpr bool USE_HF32_MODE = DispatchPolicy::USE_HF32_MODE; + static constexpr bool ENABLE_L1_RESIDENT = DispatchPolicy::ENABLE_L1_RESIDENT; + static constexpr uint32_t L1A_STAGES = DispatchPolicy::L1A_STAGES; + static constexpr uint32_t L1B_STAGES = DispatchPolicy::L1B_STAGES; + static constexpr uint32_t L0A_STAGES = DispatchPolicy::L0A_STAGES; + static constexpr uint32_t L0B_STAGES = DispatchPolicy::L0B_STAGES; + static constexpr uint32_t L0C_STAGES = DispatchPolicy::L0C_STAGES; + static constexpr uint32_t L1_TILE_M = tla::get<0>(L1TileShape{}); + static constexpr uint32_t L1_TILE_N = tla::get<1>(L1TileShape{}); + static constexpr uint32_t L1_TILE_K = tla::get<2>(L1TileShape{}); + static constexpr uint32_t L0_TILE_M = tla::get<0>(L0TileShape{}); + static constexpr uint32_t L0_TILE_N = tla::get<1>(L0TileShape{}); + static constexpr uint32_t L0_TILE_K = tla::get<2>(L0TileShape{}); + + // L1 tile size + static constexpr uint32_t L1A_TILE_SIZE = ArchTag::L1_SIZE / 2; + static constexpr uint32_t L1B_TILE_SIZE = L1_TILE_N * L1_TILE_K * sizeof(ElementB); + // L0 tile size + static constexpr uint32_t L0A_TILE_SIZE = L0_TILE_M * L0_TILE_K * sizeof(ElementA); + static constexpr uint32_t L0B_TILE_SIZE = L0_TILE_K * L0_TILE_N * sizeof(ElementB); + static constexpr uint32_t L0C_TILE_SIZE = L1_TILE_M * L1_TILE_N * sizeof(ElementAccumulator); + + // Check HF32_MODE + static_assert( + !USE_HF32_MODE || (USE_HF32_MODE && std::is_same_v && std::is_same_v), + "HF32 MODE only supports in float!"); + + // Check L0C_STAGES + static_assert(!(ENABLE_UNIT_FLAG && L0C_STAGES != 1), "L0C_STAGES must be 1 when UnitFlag is true!"); + + // Check LayoutC + static_assert(tla::detail::isRowMajor::value || + ((std::is_same_v || std::is_same_v || + std::is_same_v) && + tla::detail::iszN::value), + "LayoutC only supports zN in half or bfloat16 or float, RowMajor in all dtype yet!"); + + // Check L1TileShape + static_assert(L1A_TILE_SIZE * L1A_STAGES + L1B_TILE_SIZE * L1B_STAGES <= ArchTag::L1_SIZE, + "L1TileShape exceeding the L1 space!"); + + // Check L0TileShape + static_assert(L0A_TILE_SIZE * L0A_STAGES <= ArchTag::L0A_SIZE, "L0TileShape exceeding the L0A space!"); + static_assert(L0B_TILE_SIZE * L0B_STAGES <= ArchTag::L0B_SIZE, "L0TileShape exceeding the L0B space!"); + static_assert(L0C_TILE_SIZE * L0C_STAGES <= ArchTag::L0C_SIZE, "L0TileShape exceeding the L0C space!"); + + static_assert(L1_TILE_M == L0_TILE_M && L1_TILE_N == L0_TILE_N, + "The situation where the basic blocks of L1 and L0 differ on the m and n axes is not supported yet"); + static_assert(L0_TILE_K <= L1_TILE_K, "L0TileShape::K cannot exceed L1TileShape::K"); + + static_assert((!HAS_BIAS && (L1A_STAGES + L1B_STAGES) <= 8) || (HAS_BIAS && (L1A_STAGES + L1B_STAGES) <= 7), + "L1 Buffer overflow: Exceeds the supported range of EVENT(0~7)"); + + static_assert((!HAS_BIAS && (L0A_STAGES + L0B_STAGES) <= 8) || (HAS_BIAS && (L0A_STAGES + L0B_STAGES) <= 7), + "L0 Buffer overflow: Exceeds the supported range of EVENT_ID(0~7)"); + + static constexpr auto L1B_LAYOUT = + tla::MakeLayout(tla::Int{}, tla::Int{}); + static constexpr auto L1BIAS_LAYOUT = tla::MakeLayout(tla::Int{}); + static constexpr auto L0BIAS_LAYOUT = tla::MakeLayout(tla::Int{}); + + // When enableing L1 resident mode, restore the pointer and coordinates that record the last state + // to the initial state. if tow blockmmad instances need to be consecutively invoked at the kernel layer, + // RestoreStatus() must be inserted between them. + CATLASS_DEVICE + void RestoreStatus() + { + for (int i = 0; i < L1A_STAGES; ++i) { + lastAddrA[i] = nullptr; + lastCoordA[i] = MatrixCoord{0U, 0U}; + } + for (int i = 0; i < L1B_STAGES; ++i) { + lastAddrB[i] = nullptr; + lastCoordB[i] = MatrixCoord{0U, 0U}; + } + } + + /// Construct + CATLASS_DEVICE + BlockMmadTla(Arch::Resource &resource, uint32_t l1BufAddrStart = 0) + { + if ASCEND_IS_AIC { + // use HF32 when USE_HF32_MODE is true + if constexpr (USE_HF32_MODE) { + AscendC::SetHF32Mode(true); + } else { + AscendC::SetHF32Mode(false); + } + if constexpr (ENABLE_UNIT_FLAG && tla::detail::isRowMajor::value) { + AscendC::SetMMLayoutTransform(true); + } + uint32_t l1AOffset = l1BufAddrStart; + uint32_t l1BOffset = l1BufAddrStart + L1A_TILE_SIZE * L1A_STAGES; + // Init buffers + for (uint32_t i = 0; i < L1A_STAGES; i++) { + // Assign L1/L0A/L0B space for each stages + l1ATensorList[i] = resource.l1Buf.template GetBufferByByte(l1AOffset + L1A_TILE_SIZE * i); + // Assign event ID for each stages + l1AEventList[i] = i; + } + for (uint32_t i = 0; i < L1B_STAGES; i++) { + // Assign L1/L0A/L0B space for each stages + l1BTensorList[i] = resource.l1Buf.template GetBufferByByte(l1BOffset + L1B_TILE_SIZE * i); + // Assign event ID for each stages + l1BEventList[i] = i + L1A_STAGES; + // The event id that needs to be set before the loop + AscendC::SetFlag(l1BEventList[i]); + } + for (uint32_t i = 0; i < L0A_STAGES; i++) { + // Assign L1/L0A/L0B space for each stages + l0ATensorList[i] = resource.l0ABuf.template GetBufferByByte(L0A_TILE_SIZE * i); + // Assign event ID for each stages + l0AEventList[i] = i; + // The event id that needs to be set before the loop + AscendC::SetFlag(l0AEventList[i]); + } + for (uint32_t i = 0; i < L0B_STAGES; i++) { + // Assign L1/L0A/L0B space for each stages + l0BTensorList[i] = resource.l0BBuf.template GetBufferByByte(L0B_TILE_SIZE * i); + // Assign event ID for each stages + l0BEventList[i] = i + L0A_STAGES; + // The event id that needs to be set before the loop + AscendC::SetFlag(l0BEventList[i]); + } + if constexpr (!ENABLE_UNIT_FLAG) { + for (uint32_t i = 0; i < L0C_STAGES; i++) { + l0CTensorList[i] = resource.l0CBuf.template GetBufferByByte(L0C_TILE_SIZE * i); + l0CEventList[i] = i; + AscendC::SetFlag(l0CEventList[i]); + } + } else { + l0CTensorList[0] = resource.l0CBuf.template GetBufferByByte(0); + } + if constexpr (HAS_BIAS) { + uint32_t l1BiasOffset = l1BOffset + L1B_TILE_SIZE * L1B_STAGES; + l1BiasTensor = resource.l1Buf.template GetBufferByByte(l1BiasOffset); + l0BiasTensor = resource.btBuf.template GetBufferByByte(0); + AscendC::SetFlag(L1A_STAGES + L1B_STAGES); + AscendC::SetFlag(L0A_STAGES + L0B_STAGES); + } + + if constexpr (ENABLE_L1_RESIDENT) { + RestoreStatus(); + } + } + } + + /// Destructor + CATLASS_DEVICE + ~BlockMmadTla() + { + if ASCEND_IS_AIC { + if constexpr (USE_HF32_MODE) { + AscendC::SetHF32Mode(false); + } + if constexpr (ENABLE_UNIT_FLAG && tla::detail::isRowMajor::value) { + AscendC::SetMMLayoutTransform(false); + } + for (uint32_t i = 0; i < L1B_STAGES; i++) { + AscendC::WaitFlag(l1BEventList[i]); + } + for (uint32_t i = 0; i < L0A_STAGES; i++) { + AscendC::WaitFlag(l0AEventList[i]); + } + for (uint32_t i = 0; i < L0B_STAGES; i++) { + AscendC::WaitFlag(l0BEventList[i]); + } + if constexpr (!ENABLE_UNIT_FLAG) { + for (uint32_t i = 0; i < L0C_STAGES; i++) { + AscendC::WaitFlag(l0CEventList[i]); + } + } + if constexpr (HAS_BIAS) { + AscendC::WaitFlag(L1A_STAGES + L1B_STAGES); + AscendC::WaitFlag(L0A_STAGES + L0B_STAGES); + } + } + } + + /// Perform a block-scoped matrix multiply-accumulate + template + CATLASS_DEVICE void operator()(TensorA &tensorA, TensorB &tensorB, TensorC &tensorC, GemmCoord const &actualShape, + bool needLoadL1, TensorBias const &tensorBias = {}) + { + // Check L1TileShape + if constexpr (HAS_BIAS) { + static constexpr uint32_t BIAS_BUF_SIZE = L0_TILE_N * sizeof(ElementAccumulator); + static constexpr uint32_t L1_BIAS_SIZE = L1_TILE_N * sizeof(ElementBias); + static_assert(BIAS_BUF_SIZE <= ArchTag::BIAS_SIZE, + "BIAS_BUF_SIZE exceeding the BT space! Reduce L0_TILE_N"); + static_assert(L1A_TILE_SIZE * L1A_STAGES + L1B_TILE_SIZE * L1B_STAGES + L1_BIAS_SIZE <= ArchTag::L1_SIZE, + "L1TileShape exceeding the L1 space!"); + } + + using CopyGmToL1A = typename TileCopy_::template CopyGmToL1A; + using CopyGmToL1B = typename TileCopy_::template CopyGmToL1B; + using CopyL0CToGm = typename TileCopy_::template CopyL0CToGm; + CopyGmToL1A copyGmToL1A; + CopyGmToL1B copyGmToL1B; + CopyL0CToGm copyL0CToGm; + + // load first matrix A tile from GM to L1 + auto L1A_LAYOUT = tla::MakeLayout(tla::Int{}, actualShape.k()); + if (needLoadL1) { + AscendC::SetFlag(l1AEventList[0]); + AscendC::WaitFlag(l1AEventList[0]); + auto tensorL1A = tla::MakeTensor(l1ATensorList[0], L1A_LAYOUT, Arch::PositionL1{}); + auto tensorTileA = GetTileA(tensorA, 0, 0, actualShape.m(), actualShape.k()); + copyGmToL1A(tensorL1A, tensorTileA); + AscendC::SetFlag(l1AEventList[0]); + AscendC::WaitFlag(l1AEventList[0]); + } + + uint32_t mBlockActual = actualShape.m(); + uint32_t kBlockActual = actualShape.k(); + uint32_t nBlockActual = actualShape.n(); + + uint32_t mL1Actual = mBlockActual; + if constexpr (std::is_same_v) { + // Avoid using the gemv mode in mmad + if (mL1Actual == 1) { + mL1Actual = 16; + } + } + uint32_t nL1Actual = nBlockActual; + + auto layoutInL0C = tla::MakeLayoutL0C(mL1Actual, nL1Actual); + auto tensorL0C = tla::MakeTensor(l0CTensorList[l0CListId], layoutInL0C, Arch::PositionL0C{}); + auto tensorL0Bias = tla::MakeTensor(l0BiasTensor, L0BIAS_LAYOUT, Arch::PositionBias{}); + + // uint32_t kL1Actual = min(kBlockActual, L1_TILE_K); + uint32_t kTileCount = CeilDiv(kBlockActual); + uint32_t startTileIdx = 0; + if constexpr (ENABLE_SHUFFLE_K) { + startTileIdx = AscendC::GetBlockIdx(); + } + uint32_t firstTileIdx = startTileIdx % kTileCount; + uint32_t lastTileIdx = (startTileIdx + kTileCount - 1) % kTileCount; + uint32_t kL1Actual = + (firstTileIdx < kTileCount - 1) ? L1_TILE_K : (kBlockActual - firstTileIdx * L1_TILE_K); + + // load first matrix B tile from GM to L1 + AscendC::WaitFlag(l1BEventList[l1BListId]); + auto tensorL1B = tla::MakeTensor(l1BTensorList[l1BListId], L1B_LAYOUT, Arch::PositionL1{}); + auto tensorTileB = GetTile(tensorB, tla::MakeCoord(firstTileIdx * L1_TILE_K, 0), tla::MakeShape(kL1Actual, nBlockActual)); + if constexpr (ENABLE_L1_RESIDENT) { + if (lastAddrB[l1BListId] != tensorTileB.data().GetPhyAddr() || tla::get<0>(tensorTileB.coord()) != lastCoordB[l1BListId].row() || tla::get<1>(tensorTileB.coord()) != lastCoordB[l1BListId].column()) { + copyGmToL1B(tensorL1B, tensorTileB); + lastCoordB[l1BListId] = MatrixCoord{tla::get<0>(tensorTileB.coord()), tla::get<1>(tensorTileB.coord())}; + lastAddrB[l1BListId] = const_cast<__gm__ typename AscendC::GlobalTensor::PrimType *>( + tensorTileB.data().GetPhyAddr()); + } + } else { + copyGmToL1B(tensorL1B, tensorTileB); + } + AscendC::SetFlag(l1BEventList[l1BListId]); + + if constexpr (HAS_BIAS && !std::is_same_v) { + using CopyGmToL1Bias = typename TileCopy::template CopyGmToL1Bias; + CopyGmToL1Bias copyGmToL1Bias; + AscendC::WaitFlag(L1A_STAGES + L1B_STAGES); + auto l1Bias = l1BiasTensor.template ReinterpretCast(); + auto tensorL1Bias = tla::MakeTensor(l1Bias, L1BIAS_LAYOUT, Arch::PositionL1{}); + copyGmToL1Bias(tensorL1Bias, tensorBias); + AscendC::SetFlag(L1A_STAGES + L1B_STAGES); + } + + if constexpr (!ENABLE_UNIT_FLAG) { + AscendC::WaitFlag(l0CEventList[l0CListId]); + } + + uint32_t mL0Loop = CeilDiv(mL1Actual); + uint32_t nL0Loop = CeilDiv(nL1Actual); + + // main loop + for (uint32_t kL1Idx = 0; kL1Idx < kTileCount; kL1Idx++) { + uint32_t shuffleKIdx = (startTileIdx + kL1Idx) % kTileCount; + uint32_t l1BListIdNext = (l1BListId + 1 < L1B_STAGES) ? (l1BListId + 1) : 0; + uint32_t kL1ActualNext{0}; + // preload next tile from GM to L1 + if (kL1Idx < kTileCount - 1) { + uint32_t kL1IdxNext = (startTileIdx + kL1Idx + 1) % kTileCount; + kL1ActualNext = (kL1IdxNext < kTileCount - 1) ? L1_TILE_K : (kBlockActual - kL1IdxNext * L1_TILE_K); + + // Get L1 tensor for next stage + auto l1BTensor = l1BTensorList[l1BListIdNext]; + auto tensorL1B = tla::MakeTensor(l1BTensor, L1B_LAYOUT, Arch::PositionL1{}); + // Get GM tile for next stage + auto tensorTileB = GetTile(tensorB, tla::MakeCoord(kL1IdxNext * L1_TILE_K, 0), + tla::MakeShape(kL1ActualNext, nBlockActual)); + + // load next matrix B tile from GM to L1 + AscendC::WaitFlag(l1BEventList[l1BListIdNext]); + if constexpr (ENABLE_L1_RESIDENT) { + if (lastAddrB[l1BListIdNext] != tensorTileB.data().GetPhyAddr() || tla::get<0>(tensorTileB.coord()) != lastCoordB[l1BListIdNext].row() || tla::get<1>(tensorTileB.coord()) != lastCoordB[l1BListIdNext].column()) { + copyGmToL1B(tensorL1B, tensorTileB); + lastCoordB[l1BListIdNext] = + MatrixCoord{tla::get<0>(tensorTileB.coord()), tla::get<1>(tensorTileB.coord())}; + lastAddrB[l1BListIdNext] = + const_cast<__gm__ typename AscendC::GlobalTensor::PrimType *>( + tensorTileB.data().GetPhyAddr()); + } + } else { + copyGmToL1B(tensorL1B, tensorTileB); + } + AscendC::SetFlag(l1BEventList[l1BListIdNext]); + } + + // Get L1 tensor for current stage + auto l1ATensor = l1ATensorList[0]; + auto l1BTensor = l1BTensorList[l1BListId]; + auto tensorL1A = tla::MakeTensor(l1ATensor, L1A_LAYOUT, Arch::PositionL1{}); + tensorL1B = tla::MakeTensor(l1BTensor, L1B_LAYOUT, Arch::PositionL1{}); + // Get the loop nums on L0 + uint32_t kL0Loop = CeilDiv(kL1Actual); + + for (int mL0Idx = 0; mL0Idx < mL0Loop; mL0Idx++) { + uint32_t mL0Actual = (mL0Idx < mL0Loop - 1) ? L0_TILE_M : (mL1Actual - mL0Idx * L0_TILE_M); + + for (int kL0Idx = 0; kL0Idx < kL0Loop; kL0Idx++) { + uint32_t kL0Actual = (kL0Idx < kL0Loop - 1) ? L0_TILE_K : (kL1Actual - kL0Idx * L0_TILE_K); + + // Locate the current tile on L0A + auto l0ATile = l0ATensorList[l0AListId]; + auto layoutAInL0 = tla::MakeLayout(mL0Actual, kL0Actual); + auto tensorL0A = tla::MakeTensor(l0ATile, layoutAInL0, Arch::PositionL0A{}); + // Locate the current tile of matrix A on L1 + auto tensorTileL1A = GetTileA(tensorL1A, mL0Idx * L0_TILE_M, kL0Idx * L0_TILE_K + shuffleKIdx * L1_TILE_K, + mL0Actual, kL0Actual); + + AscendC::WaitFlag(l0AEventList[l0AListId]); + + // Load current tile from L1 to L0A + copyL1ToL0A(tensorL0A, tensorTileL1A); + + bool initC = ((kL1Idx == 0) && (kL0Idx == 0)); + for (int nL0Idx = 0; nL0Idx < nL0Loop; nL0Idx++) { + uint32_t nL0Actual = (nL0Idx < nL0Loop - 1) ? L0_TILE_N : (nL1Actual - nL0Idx * L0_TILE_N); + + // Locate the current tile on L0B + auto l0BTile = l0BTensorList[l0BListId]; + auto layoutBInL0 = tla::MakeLayout(kL0Actual, nL0Actual); + auto tensorL0B = tla::MakeTensor(l0BTile, layoutBInL0, Arch::PositionL0B{}); + // Locate the current tile of matrix B on L1 + auto tensorTileL1B = GetTile(tensorL1B, + tla::MakeCoord(kL0Idx * L0_TILE_K, nL0Idx * L0_TILE_N), + tla::MakeShape(kL0Actual, nL0Actual)); + + // Wait for mmad finished + AscendC::WaitFlag(l0BEventList[l0BListId]); + // If the current tile is the first one on the k&n axis, wait for loading matrix B from GM to L1 + if ((mL0Idx == 0) && (kL0Idx == 0) && (nL0Idx == 0)) { + AscendC::WaitFlag(l1BEventList[l1BListId]); + } + + // Load current tile from L1 to L0B + copyL1ToL0B(tensorL0B, tensorTileL1B); + + // If the current tile is the last one on the k&n axis, notify to load matrix B from GM to L1 + if ((mL0Idx == mL0Loop - 1) && (kL0Idx == kL0Loop - 1) && (nL0Idx == nL0Loop - 1)) { + AscendC::SetFlag(l1BEventList[l1BListId]); + } + + if constexpr (HAS_BIAS && !std::is_same_v) { + if (initC) { + if (nL0Idx == 0) { + AscendC::WaitFlag(L1A_STAGES + L1B_STAGES); + } + AscendC::WaitFlag(L0A_STAGES + L0B_STAGES); + auto l1Bias = l1BiasTensor.template ReinterpretCast(); + auto tensorL1Bias = tla::MakeTensor(l1Bias, L1BIAS_LAYOUT, Arch::PositionL1{}); + auto tensorTileL1Bias = GetTile(tensorL1Bias, + tla::MakeCoord(nL0Idx * L0_TILE_N), + tla::MakeShape(nL0Actual)); + // Load bias to l0 biastable + copyL1ToBT(tensorL0Bias, tensorTileL1Bias); + if (nL0Idx == nL0Loop - 1) { + AscendC::SetFlag(L1A_STAGES + L1B_STAGES); + } + } + } + + // Notify to do mmad + AscendC::SetFlag(l0CEventList[l0CListId]); + + // Locate the current tile on L0C + auto tensorTileL0C = GetTile(tensorL0C, + tla::MakeCoord(mL0Idx * L0_TILE_M, nL0Idx * L0_TILE_N), + tla::MakeShape(mL0Actual, nL0Actual)); + + // Compute the matrix multiplication on L0A and L0B and write the result to the accumulator + // Wait for loading L0B + AscendC::WaitFlag(l0CEventList[l0CListId]); + + // If the unit flag is enabled, the unit flag is set according to the calculation progress + uint8_t unitFlag = 0b00; + if constexpr (ENABLE_UNIT_FLAG) { + if ((kL1Idx == kTileCount - 1) && (mL0Idx == mL0Loop - 1) && + (kL0Idx == kL0Loop - 1) && (nL0Idx == nL0Loop - 1)) { + unitFlag = 0b11; + } else { + unitFlag = 0b10; + } + } + + if constexpr (HAS_BIAS && !std::is_same_v) { + if (initC) { + tileMmad(tensorTileL0C, tensorL0A, tensorL0B, tensorL0Bias, + mL0Actual, nL0Actual, kL0Actual, initC, unitFlag); + AscendC::SetFlag(L0A_STAGES + L0B_STAGES); + } else { + tileMmad(tensorTileL0C, tensorL0A, tensorL0B, + mL0Actual, nL0Actual, kL0Actual, initC, unitFlag); + } + } else { + tileMmad(tensorTileL0C, tensorL0A, tensorL0B, + mL0Actual, nL0Actual, kL0Actual, initC, unitFlag); + } + + // Notify to move the next L0B tile + AscendC::SetFlag(l0BEventList[l0BListId]); + l0BListId = (l0BListId + 1 < L0B_STAGES) ? (l0BListId + 1) : 0; + } + AscendC::SetFlag(l0AEventList[l0AListId]); + l0AListId = (l0AListId + 1 < L0A_STAGES) ? (l0AListId + 1) : 0; + } + } + l1BListId = l1BListIdNext; + kL1Actual = kL1ActualNext; + } + + // copy block out + if constexpr (!ENABLE_UNIT_FLAG) { + AscendC::SetFlag(l0CEventList[l0CListId]); + AscendC::WaitFlag(l0CEventList[l0CListId]); + copyL0CToGm(tensorC, tensorL0C); + AscendC::SetFlag(l0CEventList[l0CListId]); + l0CListId = (l0CListId + 1 < L0C_STAGES) ? (l0CListId + 1) : 0; + } else { + copyL0CToGm(tensorC, tensorL0C, 0b11); + } + } + +protected: + template + CATLASS_DEVICE auto GetTileA(TensorA &tensorA, uint32_t mIndex, uint32_t kIndex, uint32_t mSize, uint32_t kSize) + { + if constexpr (tla::detail::isVector::value) { + return GetTile(tensorA, tla::MakeCoord(kIndex), tla::MakeShape(kSize)); + } else { + return GetTile(tensorA, tla::MakeCoord(mIndex, kIndex), tla::MakeShape(mSize, kSize)); + } + } + + // Multi-stage tensors list + AscendC::LocalTensor l1ATensorList[L1A_STAGES]; + AscendC::LocalTensor l1BTensorList[L1B_STAGES]; + AscendC::LocalTensor l0ATensorList[L0A_STAGES]; + AscendC::LocalTensor l0BTensorList[L0B_STAGES]; + AscendC::LocalTensor l0CTensorList[L0C_STAGES]; + AscendC::LocalTensor l1BiasTensor; + AscendC::LocalTensor l0BiasTensor; + + // Multi-stage event id list + int32_t l1AEventList[L1A_STAGES]; + int32_t l1BEventList[L1B_STAGES]; + int32_t l0AEventList[L0A_STAGES]; + int32_t l0BEventList[L0B_STAGES]; + int32_t l0CEventList[L0C_STAGES]; + + __gm__ typename AscendC::GlobalTensor::PrimType *lastAddrA[L1A_STAGES]; + __gm__ typename AscendC::GlobalTensor::PrimType *lastAddrB[L1B_STAGES]; + MatrixCoord lastCoordA[L1A_STAGES]; + MatrixCoord lastCoordB[L1B_STAGES]; + + // The id of current stage + uint32_t l1AListId{0}; + uint32_t l1BListId{0}; + uint32_t l0AListId{0}; + uint32_t l0BListId{0}; + uint32_t l0CListId{0}; + + TileMmad tileMmad; + CopyL1ToL0A copyL1ToL0A; + CopyL1ToL0B copyL1ToL0B; + CopyL1ToBT copyL1ToBT; +}; + +} // namespace Catlass::Gemm::Block + +#endif // CATLASS_GEMM_BLOCK_BLOCK_MMAD_PINGPONG_FULL_LOADA_TLA_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_pingpong_per_group_per_block_tla.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_pingpong_per_group_per_block_tla.hpp new file mode 100644 index 0000000000..56b51d3c41 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_pingpong_per_group_per_block_tla.hpp @@ -0,0 +1,320 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_GEMM_BLOCK_BLOCK_MMAD_PINGPONG_PER_GROUP_PER_BLOCK_TLA_HPP +#define CATLASS_GEMM_BLOCK_BLOCK_MMAD_PINGPONG_PER_GROUP_PER_BLOCK_TLA_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/kda_coord.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" +#include "catlass/gemm/kda_gemm_helper.hpp" +#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" +#include "catlass/gemm/tile/kda_tile_mmad.hpp" +#include "tla/kda_tla_layout.hpp" +#include "tla/kda_tensor.hpp" + +namespace Catlass::Gemm::Block { + +template < + class ArchTag_, + bool ENABLE_UNIT_FLAG_, + class L1TileShape_, + class L0TileShape_, + class ElementA_, + class ElementB_, + class ElementC_, + class ElementBias_, + class TileCopy_, + class TileMmad_> +struct BlockMmadTla< + MmadPingpongPertile, + L1TileShape_, + L0TileShape_, + ElementA_, + ElementB_, + ElementC_, + ElementBias_, + TileCopy_, + TileMmad_> { +public: + struct Arguments { + GM_ADDR aGmAddr{nullptr}; + GM_ADDR bGmAddr{nullptr}; + GM_ADDR cGmAddr{nullptr}; + GM_ADDR workspaceGmAddr{nullptr}; + }; + + using Params = Arguments; + + using DispatchPolicy = MmadPingpongPertile; + using ArchTag = typename DispatchPolicy::ArchTag; + + using TileCopy = TileCopy_; + using L1TileShape = L1TileShape_; + using L0TileShape = L0TileShape_; + using ElementA = ElementA_; + using LayoutA = typename TileCopy::LayoutA; + using ElementB = ElementB_; + using LayoutB = typename TileCopy::LayoutB; + using ElementC = ElementC_; + using LayoutC = typename TileCopy::LayoutC; + using ElementBias = ElementBias_; + + using TileMmad = TileMmad_; + + using CopyL1ToL0A = typename TileCopy::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopy::CopyL1ToL0B; + using CopyL1ToBT = typename TileCopy::CopyL1ToBT; + + using ElementAccumulator = typename TileCopy::ElementAccumulator; + + static constexpr bool HAS_BIAS = TileCopy::HAS_BIAS; + + using LayoutTagA = typename TileCopy::LayoutTagA; + using LayoutTagB = typename TileCopy::LayoutTagB; + using LayoutTagC = typename TileCopy::LayoutTagC; + using LayoutTagL1A = typename TileCopy::LayoutTagL1A; + using LayoutTagL1B = typename TileCopy::LayoutTagL1B; + using LayoutTagL0A = typename TileCopy::LayoutTagL0A; + using LayoutTagL0B = typename TileCopy::LayoutTagL0B; + + uint64_t m_{1}; + uint64_t n_{1}; + uint64_t k_{1}; + + constexpr static uint64_t STAGES = DispatchPolicy::STAGES; + + constexpr static uint64_t ML1_ = tla::get<0>(L1TileShape{}); + constexpr static uint64_t NL1_ = tla::get<1>(L1TileShape{}); + constexpr static uint64_t KL1_ = tla::get<2>(L1TileShape{}); + + constexpr static uint64_t ML0_ = tla::get<0>(L0TileShape{}); + constexpr static uint64_t NL0_ = tla::get<1>(L0TileShape{}); + constexpr static uint64_t KL0_ = tla::get<2>(L0TileShape{}); + + constexpr static uint64_t L1A_TILE_SIZE = ML1_ * KL1_ * sizeof(ElementA); + constexpr static uint64_t L1B_TILE_SIZE = NL1_ * KL1_ * sizeof(ElementB); + + constexpr static uint64_t L0A_TILE_SIZE = ML0_ * KL0_ * sizeof(ElementA); + constexpr static uint64_t L0B_TILE_SIZE = NL0_ * KL0_ * sizeof(ElementB); + constexpr static uint64_t L0C_TILE_SIZE = ML0_ * NL0_ * sizeof(ElementAccumulator); + + CATLASS_DEVICE + BlockMmadTla(const GemmCoord &shape) + { + uint32_t l1AOffset = 0; + uint32_t l1BOffset = l1AOffset + L1A_TILE_SIZE * STAGES; + // Init buffers + for (uint32_t i = 0; i < STAGES; i++) { + l1ATensorList_[i] = l1Buf_[l1AOffset + L1A_TILE_SIZE * i].template ReinterpretCast(); + l1BTensorList_[i] = l1Buf_[l1BOffset + L1B_TILE_SIZE * i].template ReinterpretCast(); + l0ATensorList_[i] = l0ABuf_[L0A_TILE_SIZE * i].template ReinterpretCast(); + l0BTensorList_[i] = l0BBuf_[L0B_TILE_SIZE * i].template ReinterpretCast(); + l0CTensorList_[i] = l0CBuf_[L0C_TILE_SIZE * i].template ReinterpretCast(); + + l1AEventList_[i] = i; + l1BEventList_[i] = i + STAGES; + l0AEventList_[i] = i; + l0BEventList_[i] = i + STAGES; + l0CEventList_[i] = i; + + AscendC::SetFlag(l1AEventList_[i]); + AscendC::SetFlag(l1BEventList_[i]); + AscendC::SetFlag(l0AEventList_[i]); + AscendC::SetFlag(l0BEventList_[i]); + AscendC::SetFlag(l0CEventList_[i]); + } + m_ = shape.m(); + n_ = shape.n(); + k_ = shape.k(); + } + + CATLASS_DEVICE + ~BlockMmadTla() + { + for (uint32_t i = 0; i < STAGES; i++) { + AscendC::WaitFlag(l1AEventList_[i]); + AscendC::WaitFlag(l1BEventList_[i]); + AscendC::WaitFlag(l0AEventList_[i]); + AscendC::WaitFlag(l0BEventList_[i]); + AscendC::WaitFlag(l0CEventList_[i]); + } + } + + CATLASS_DEVICE void WaitForVector(uint32_t crossPingPongID) + { + AscendC::CrossCoreWaitFlag(AIV_SYNC_AIC_FLAG + (crossPingPongID_)); + AscendC::CrossCoreWaitFlag(AIV_SYNC_AIC_FLAG + FLAG_ID_MAX + (crossPingPongID_)); + } + + CATLASS_DEVICE void NotifyVector(uint32_t crossPingPongID) + { + AscendC::CrossCoreSetFlag(AIC_SYNC_AIV_FLAG + (crossPingPongID_)); + AscendC::CrossCoreSetFlag(AIC_SYNC_AIV_FLAG + FLAG_ID_MAX + (crossPingPongID_)); + } + +public: + template + CATLASS_DEVICE void operator()( + DstTensor &cTensor, + SrcATensor &aTensor, + SrcBTensor &bTensor, + GemmCoord tileShape) + { + using CopyL0CToOut = typename TileCopy::template CopyL0CToDst; + CopyL0CToOut copyL0CToOut; + + uint32_t curML0 = tileShape.m(); + uint32_t curNL0 = tileShape.n(); + auto l0CLayout = tla::MakeLayoutL0C(curML0, curNL0); + + uint32_t kL1Loops = CeilDiv(k_, KL1_); + + // GM -> L1 + uint32_t kL1OffsetLength = 0; + for (uint32_t kL1Idx = 0; kL1Idx < kL1Loops; ++kL1Idx) { + uint32_t curKL1 = (kL1Idx + 1 == kL1Loops) ? (k_ - kL1OffsetLength) : KL1_; + + AscendC::WaitFlag(l1AEventList_[l1ListId_]); + auto l1ALayout = tla::MakeLayout(curML0, curKL1); + auto l1ATensor = tla::MakeTensor(l1ATensorList_[l1ListId_], l1ALayout, Arch::PositionL1{}); + + auto aTileTensor = GetTile( + aTensor, + tla::MakeCoord(0, kL1OffsetLength), + tla::MakeShape(curML0, curKL1)); + + using CopyGmToL1A = typename TileCopy::template CopyGmToL1A; + CopyGmToL1A copyGmToL1A; + copyGmToL1A(l1ATensor, aTileTensor); + + AscendC::SetFlag(l1AEventList_[l1ListId_]); + AscendC::WaitFlag(l1AEventList_[l1ListId_]); + + AscendC::WaitFlag(l1BEventList_[l1ListId_]); + + auto l1BLayout = tla::MakeLayout(curKL1, curNL0); + auto l1BTensor = tla::MakeTensor(l1BTensorList_[l1ListId_], l1BLayout, Arch::PositionL1{}); + auto bTileTensor = GetTile( + bTensor, + tla::MakeCoord(kL1OffsetLength, 0), + tla::MakeShape(curKL1, curNL0)); + + using CopyGmToL1B = typename TileCopy::template CopyGmToL1B; + CopyGmToL1B copyGmToL1B; + copyGmToL1B(l1BTensor, bTileTensor); + + AscendC::SetFlag(l1BEventList_[l1ListId_]); + AscendC::WaitFlag(l1BEventList_[l1ListId_]); + + kL1OffsetLength += curKL1; + + // L1 -> L0 + uint32_t kL0Loops = CeilDiv(curKL1, KL0_); + for (uint32_t kL0Idx = 0; kL0Idx < kL0Loops; ++kL0Idx) { + uint32_t curKL0 = (kL0Idx + 1 == kL0Loops) ? (curKL1 - kL0Idx * KL0_) : KL0_; + + AscendC::WaitFlag(l0AEventList_[l0ListId_]); + + auto l0ALayout = tla::MakeLayout(curML0, curKL0); + auto l0ATensor = tla::MakeTensor(l0ATensorList_[l1ListId_], l0ALayout, Arch::PositionL0A{}); + auto l1ATileTensor = GetTile( + l1ATensor, + tla::MakeCoord(0, kL0Idx * KL0_), + tla::MakeShape(curML0, curKL0)); + + copyL1ToL0A(l0ATensor, l1ATileTensor); + + AscendC::WaitFlag(l0BEventList_[l0ListId_]); + + auto l0BLayout = tla::MakeLayout(curKL0, curNL0); + auto l0BTensor = tla::MakeTensor(l0BTensorList_[l0ListId_], l0BLayout, Arch::PositionL0B{}); + auto l1BTileTensor = GetTile( + l1BTensor, + tla::MakeCoord(kL0Idx * KL0_, 0), + tla::MakeShape(curKL0, curNL0)); + + copyL1ToL0B(l0BTensor, l1BTileTensor); + + AscendC::WaitFlag(l0CEventList_[l0CListId_]); + auto l0CTensor = tla::MakeTensor(l0CTensorList_[l0CListId_], l0CLayout, Arch::PositionL0C{}); + + AscendC::SetFlag(l0CEventList_[l0CListId_]); + AscendC::WaitFlag(l0CEventList_[l0CListId_]); + + // Mmad + bool initC = (kL1Idx == 0) && (kL0Idx == 0); + curML0 = (curML0 == 1 ? M_ALIGN : curML0); + tileMmad(l0CTensor, l0ATensor, l0BTensor, curML0, curNL0, curKL0, true); + + AscendC::SetFlag(l0AEventList_[l0ListId_]); + AscendC::SetFlag(l0BEventList_[l0ListId_]); + l0ListId_ = (l0ListId_ + 1 < STAGES) ? (l0ListId_ + 1) : 0; + + AscendC::SetFlag(l0CEventList_[l0CListId_]); + AscendC::WaitFlag(l0CEventList_[l0CListId_]); + + WaitForVector(crossPingPongID_); + auto outTensor = tla::MakeTensor(crossPingPongID_ == 0 ? cTensor.data() : cTensor.data()[UB_OFFSET], + cTensor.layout(), Arch::PositionUB{}); + copyL0CToOut(outTensor, l0CTensor); + NotifyVector(crossPingPongID_); + + AscendC::SetFlag(l0CEventList_[l0CListId_]); + l0CListId_ = (l0CListId_ + 1 < STAGES) ? (l0CListId_ + 1) : 0; + crossPingPongID_ = (crossPingPongID_ + 1 < STAGES) ? (crossPingPongID_ + 1) : 0; + } + AscendC::SetFlag(l1AEventList_[l1ListId_]); + AscendC::SetFlag(l1BEventList_[l1ListId_]); + l1ListId_ = (l1ListId_ + 1 < STAGES) ? (l1ListId_ + 1) : 0; + } + } + +private: + constexpr static uint16_t M_ALIGN = 16; + constexpr static uint16_t AIC_SYNC_AIV_MODE_4 = 4; + constexpr static uint16_t AIV_SYNC_AIC_FLAG = 8; + constexpr static uint16_t AIC_SYNC_AIV_FLAG = 6; + constexpr static uint16_t FLAG_ID_MAX = 16; + constexpr static uint32_t UB_TWO_BANK_ELEMS_B32 = 128U; + constexpr static int64_t PER_BLOCK_SIZE = 128LL; + constexpr static uint64_t UB_OFFSET = UB_TWO_BANK_ELEMS_B32 * PER_BLOCK_SIZE; + + CopyL1ToL0A copyL1ToL0A; + CopyL1ToL0B copyL1ToL0B; + TileMmad tileMmad; + + AscendC::LocalTensor l1Buf_{AscendC::LocalTensor(AscendC::TPosition::A1, 0, ArchTag::L1_SIZE)}; + AscendC::LocalTensor l0ABuf_{AscendC::LocalTensor(AscendC::TPosition::A2, 0, ArchTag::L0A_SIZE)}; + AscendC::LocalTensor l0BBuf_{AscendC::LocalTensor(AscendC::TPosition::B2, 0, ArchTag::L0B_SIZE)}; + AscendC::LocalTensor l0CBuf_{AscendC::LocalTensor(AscendC::TPosition::CO1, 0, ArchTag::L0C_SIZE)}; + + AscendC::LocalTensor l1ATensorList_[STAGES]; + AscendC::LocalTensor l1BTensorList_[STAGES]; + AscendC::LocalTensor l0ATensorList_[STAGES]; + AscendC::LocalTensor l0BTensorList_[STAGES]; + AscendC::LocalTensor l0CTensorList_[STAGES]; + + int32_t l1AEventList_[STAGES]; + int32_t l1BEventList_[STAGES]; + int32_t l0AEventList_[STAGES]; + int32_t l0BEventList_[STAGES]; + int32_t l0CEventList_[STAGES]; + + uint32_t l1ListId_{0}; + uint32_t l0ListId_{0}; + uint32_t l0CListId_{0}; + uint32_t crossPingPongID_{0}; +}; +} // namespace Catlass::Gemm::Block + +#endif // CATLASS_GEMM_BLOCK_BLOCK_MMAD_PINGPONG_PER_GROUP_PER_BLOCK_TLA_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_pingpong_slice_k_with_prologue.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_pingpong_slice_k_with_prologue.hpp new file mode 100644 index 0000000000..ca4f77289f --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_pingpong_slice_k_with_prologue.hpp @@ -0,0 +1,806 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_GEMM_BLOCK_BLOCK_MMAD_PINGPONG_SLICE_K_WITH_PROLOGUE_HPP +#define CATLASS_GEMM_BLOCK_BLOCK_MMAD_PINGPONG_SLICE_K_WITH_PROLOGUE_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/kda_coord.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" +#include "catlass/gemm/kda_gemm_helper.hpp" +#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" +#include "catlass/gemm/tile/kda_tile_traits.hpp" + +namespace Catlass::Gemm::Block { + +template < + bool ENABLE_UNIT_FLAG_, + class L1TileShape_, + class L0TileShape_, + class AType_, + class BType_, + class CType_, + class BiasType_, + class TileCopy_, + class TileMmad_> +struct BlockMmad< + MmadAtlasA2PingpongSliceKWithPrologue, + L1TileShape_, + L0TileShape_, + AType_, + BType_, + CType_, + BiasType_, + TileCopy_, + TileMmad_> { +public: + // Type Aliases + using DispatchPolicy = MmadAtlasA2PingpongSliceKWithPrologue; + using ArchTag = typename DispatchPolicy::ArchTag; + using L1TileShape = L1TileShape_; + using L0TileShape = L0TileShape_; + using ElementA = typename AType_::Element; + using LayoutA = typename AType_::Layout; + using ElementB = typename BType_::Element; + using LayoutB = typename BType_::Layout; + using ElementC = typename CType_::Element; + using LayoutC = typename CType_::Layout; + using TileMmad = TileMmad_; + using CopyGmToL1A = typename TileCopy_::CopyGmToL1A; + using CopyGmToL1B = typename TileCopy_::CopyGmToL1B; + using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; + using CopyL0CToGm = typename TileCopy_::CopyL0CToGm; + + using PrologueA = typename TileCopy_::PrologueA; + using PrologueB = typename TileCopy_::PrologueB; + using PrologueC = typename TileCopy_::PrologueB; // NOTE: SHOULD BE Epilogue + + using ElementAccumulator = + typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; + using LayoutAInL1 = typename CopyL1ToL0A::LayoutSrc; + using LayoutBInL1 = typename CopyL1ToL0B::LayoutSrc; + using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst; + using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst; + using LayoutCInL0 = layout::zN; + + using L1AAlignHelper = Gemm::helper::L1AlignHelper; + using L1BAlignHelper = Gemm::helper::L1AlignHelper; + + static constexpr bool ENABLE_UNIT_FLAG = DispatchPolicy::ENABLE_UNIT_FLAG; + static constexpr uint32_t STAGES = DispatchPolicy::STAGES; + static constexpr uint32_t L1A_SIZE = L1TileShape::M * L1TileShape::K * sizeof(ElementA); + static constexpr uint32_t L1B_SIZE = L1TileShape::N * L1TileShape::K * sizeof(ElementB); + static constexpr uint32_t L0A_SIZE = ArchTag::L0A_SIZE; + static constexpr uint32_t L0B_SIZE = ArchTag::L0B_SIZE; + static constexpr uint32_t L0A_PINGPONG_BUF_SIZE = L0A_SIZE / STAGES; + static constexpr uint32_t L0B_PINGPONG_BUF_SIZE = L0B_SIZE / STAGES; + + static constexpr bool HAS_PROLOGUE_A = !std::is_same_v; + static constexpr bool HAS_PROLOGUE_B = !std::is_same_v; + + // Check L1TileShape + static_assert((L1A_SIZE * STAGES + L1B_SIZE * STAGES) <= ArchTag::L1_SIZE, "L1TileShape exceeding the L1 space!"); + + // Check L0TileShape + static constexpr uint32_t L0A_TILE_SIZE = L0TileShape::M * L0TileShape::K * sizeof(ElementA); + static constexpr uint32_t L0B_TILE_SIZE = L0TileShape::K * L0TileShape::N * sizeof(ElementB); + static_assert((L0A_TILE_SIZE * STAGES) <= L0A_SIZE, "L0TileShape exceeding the L0A space!"); + static_assert((L0B_TILE_SIZE * STAGES) <= L0B_SIZE, "L0TileShape exceeding the L0B space!"); + + // 32B (256b) aligned + static_assert(Gemm::helper::TileShapeAlignChecker::_ALIGN == 256, + "Tile shape must be 32B aligned."); + + struct Params { + typename Tile::PrologueTraits::Params prologueAParams{}; + typename Tile::PrologueTraits::Params prologueBParams{}; + uint32_t mScalar; + uint32_t nScalar; + uint32_t splitkLength; + }; + + /// Construct + CATLASS_DEVICE + BlockMmad(Arch::Resource &resource, Params const ¶ms_ = {}, + uint32_t l1BufAddrStart = 0) + : params(params_) + { + // PREVIOUSLY: + // Initilize PrologueA & PrologueB instantly + // prologueA(resource, params_.prologueAParams), prologueB(resource, params_.prologueBParams) + uint32_t l1AOffset = l1BufAddrStart; + uint32_t l1BOffset = l1BufAddrStart + L1A_SIZE * STAGES; + // Init set/wait flags + flag0[0].id = flagID0; + flag0[1].id = flagID1; + flag1[0].id = flagID2; + flag1[1].id = flagID3; + + // AIC Type + if constexpr (g_coreType == AscendC::AIC) { + // Init buffers + for (uint32_t i = 0; i < STAGES; i++) { + // Assign L1/L0A/L0B space for each stages + l1ATensorList[i] = resource.l1Buf.template GetBufferByByte(l1AOffset + L1A_SIZE * i); + l1BTensorList[i] = resource.l1Buf.template GetBufferByByte(l1BOffset + L1B_SIZE * i); + l0ATensorList[i] = resource.l0ABuf.template GetBufferByByte(L0A_PINGPONG_BUF_SIZE * i); + l0BTensorList[i] = resource.l0BBuf.template GetBufferByByte(L0B_PINGPONG_BUF_SIZE * i); + + // Assign event ID for each stages + l1AEventList[i] = i; + l1BEventList[i] = i + STAGES; + l0AEventList[i] = i; + l0BEventList[i] = i + STAGES; + + // The event id that needs to be set before the loop + AscendC::SetFlag(l1AEventList[i]); + AscendC::SetFlag(l1BEventList[i]); + AscendC::SetFlag(l0AEventList[i]); + AscendC::SetFlag(l0BEventList[i]); + } + l0CTensor = resource.l0CBuf.template GetBufferByByte(0); + AscendC::SetFlag(EVENT_ID0); + + // CV 同步 + Catlass::Arch::CrossCoreSetFlag<0x2, PIPE_FIX>(flag0[0]); + Catlass::Arch::CrossCoreSetFlag<0x2, PIPE_FIX>(flag0[1]); + } + // AIV Coretype + if constexpr (g_coreType == AscendC::AIV) { + AscendC::SetFlag(EVENT_ID0); + AscendC::SetFlag(EVENT_ID1); + AscendC::SetFlag(EVENT_ID2); + AscendC::SetFlag(EVENT_ID3); + AscendC::SetFlag(EVENT_ID0); + AscendC::SetFlag(EVENT_ID1); + AscendC::SetFlag(EVENT_ID2); + AscendC::SetFlag(EVENT_ID3); + } + } + + /// Destructor + CATLASS_DEVICE + ~BlockMmad() + { + if constexpr (g_coreType == AscendC::AIC) { + for (uint32_t i = 0; i < STAGES; i++) { + AscendC::WaitFlag(l1AEventList[i]); + AscendC::WaitFlag(l1BEventList[i]); + AscendC::WaitFlag(l0AEventList[i]); + AscendC::WaitFlag(l0BEventList[i]); + } + AscendC::WaitFlag(EVENT_ID0); + } + + if constexpr (g_coreType == AscendC::AIV) { + AscendC::WaitFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID1); + AscendC::WaitFlag(EVENT_ID2); + AscendC::WaitFlag(EVENT_ID3); + AscendC::WaitFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID1); + AscendC::WaitFlag(EVENT_ID2); + AscendC::WaitFlag(EVENT_ID3); + + Catlass::Arch::CrossCoreWaitFlag(flag0[0]); + Catlass::Arch::CrossCoreWaitFlag(flag0[1]); + } + } + + /// Perform a block-scoped matrix multiply-accumulate + CATLASS_DEVICE + void operator()( + AscendC::GlobalTensor const &gmA, + AscendC::GlobalTensor const &gmB, + AscendC::GlobalTensor const &gmC, + GemmCoord actualBlockShape, + GemmCoord problemShape) + { + MNScalerMatmulImpl(gmA, gmB, gmC, {}, actualBlockShape, problemShape); + } + + /// Prologue + // template + CATLASS_DEVICE + void Prologue( + AscendC::GlobalTensor const &gmDstA, + AscendC::GlobalTensor const &gmSrcA, typename PrologueA::LayoutSrc const &layoutSrcA, + AscendC::GlobalTensor const &gmDstB, + AscendC::GlobalTensor const &gmSrcB, typename PrologueB::LayoutSrc const &layoutSrcB, + AscendC::GlobalTensor const &gmDstC, + AscendC::GlobalTensor const &gmSrcC, LayoutC const &layoutSrcC, + GemmCoord const &blockCoord, GemmCoord const &nextBlockCoord, + GemmCoord const &actualBlockShape, GemmCoord const &nextActualBlockShape, + GemmCoord const &problemShape, + bool isFirstBlock, bool hasNextBlock, + uint32_t &bufferIndex) + { + PrologueImpl( + gmDstA, gmSrcA, layoutSrcA, + gmDstB, gmSrcB, layoutSrcB, + gmDstC, gmSrcC, layoutSrcC, + blockCoord, nextBlockCoord, + actualBlockShape, nextActualBlockShape, problemShape, + isFirstBlock, hasNextBlock, bufferIndex); + } + +protected: + CATLASS_DEVICE + void PrologueImpl( + AscendC::GlobalTensor const &gmDstA, + AscendC::GlobalTensor const &gmSrcA, + typename PrologueA::LayoutSrc const &layoutSrcA, + AscendC::GlobalTensor const &gmDstB, + AscendC::GlobalTensor const &gmSrcB, + typename PrologueB::LayoutSrc const &layoutSrcB, + AscendC::GlobalTensor const &gmDstC, + AscendC::GlobalTensor const &gmSrcC, + LayoutC const &layoutSrcC, + GemmCoord blockCoord, + GemmCoord nextBlockCoord, + GemmCoord actualBlockShape, + GemmCoord nextActualBlockShape, + GemmCoord problemShape, + bool const isFirstBlock, + bool const hasNextBlock, + uint32_t &bufferIndex) + { + // 当前任务块 Current block + // 计算基础偏移量 + // layoutSrcA: params.layoutA, layoutSrcB: params.layoutB, layoutSrcC: params.layoutC + int64_t gmOffsetA, gmOffsetB, gmOffsetC, gmOffsetWA, gmOffsetWB, gmOffsetWC; + int64_t gmOffsetNextA, gmOffsetNextB, gmOffsetWADelta, gmOffsetWBDelta; + + gmOffsetA = layoutSrcA.GetOffset(MakeCoord(blockCoord.m() * (L1TileShape::M * params.mScalar), 0U)); + gmOffsetNextA = layoutSrcA.GetOffset(MakeCoord(nextBlockCoord.m() * (L1TileShape::M * params.mScalar), 0U)); + gmOffsetWA = (AscendC::GetBlockIdx() / AIVPERCORE) * + (params.mScalar * L1TileShape::M) * params.splitkLength * STAGES; + gmOffsetWADelta = (params.mScalar * L1TileShape::M) * params.splitkLength; + + gmOffsetB = layoutSrcB.GetOffset(MakeCoord(0U, blockCoord.n() * (L1TileShape::N * params.nScalar))); + gmOffsetNextB = layoutSrcB.GetOffset(MakeCoord(0U, nextBlockCoord.n() * (L1TileShape::N * params.nScalar))); + gmOffsetWB = (AscendC::GetBlockIdx() / AIVPERCORE) * + (params.nScalar * L1TileShape::N) * params.splitkLength * STAGES; + gmOffsetWBDelta = params.splitkLength * (params.nScalar * L1TileShape::N); + + MatrixCoord offsetC{blockCoord.m() * (L1TileShape::M * params.mScalar), + blockCoord.n() * (L1TileShape::N * params.nScalar)}; + gmOffsetC = layoutSrcC.GetOffset(offsetC); + gmOffsetWC = (AscendC::GetBlockIdx() / AIVPERCORE) * + (params.mScalar * L1TileShape::M) * (params.nScalar * L1TileShape::N); + + uint32_t srcAStride = layoutSrcA.stride( + std::is_same_v ? 0 : 1); + uint32_t srcBStride = layoutSrcB.stride( + std::is_same_v ? 0 : 1); + uint32_t mShape_ = actualBlockShape.m(); + uint32_t nShape_ = actualBlockShape.n(); + + // Begin resource + Arch::Resource resource; + + // SPLIT-K MainLoop + uint32_t kLoop = (problemShape.k() + params.splitkLength - 1) / params.splitkLength; + uint32_t kResidueLenght = problemShape.k() % params.splitkLength; + for (uint32_t ldk = 0; ldk < kLoop; ldk++) { + uint32_t kActual = (problemShape.k() < (ldk + 1) * params.splitkLength) ? kResidueLenght : params.splitkLength; + uint32_t kActual_ = kActual; + + if (ldk == 0 && isFirstBlock) { // 首块 + Catlass::Arch::CrossCoreWaitFlag(flag0[crossCoreBufferIndexAIV]); + + uint32_t kActualAligned_ = RoundUp<256, uint32_t>(kActual_); + LayoutA layoutWA(actualBlockShape.m(), kActual_, srcAStride); + LayoutA layoutDstA(actualBlockShape.m(), kActualAligned_); + LayoutB layoutWB(kActual_, actualBlockShape.n(), srcBStride); + LayoutB layoutDstB(kActualAligned_, actualBlockShape.n()); + + PrologueA prologueA(resource, params.prologueAParams); + PrologueB prologueB(resource, params.prologueBParams); + prologueA( + gmDstA[gmOffsetWA + crossCoreBufferIndexAIV * gmOffsetWADelta], + layoutDstA, + gmSrcA[gmOffsetA], + layoutWA, + bufferIndex); + prologueB( + gmDstB[gmOffsetWB + crossCoreBufferIndexAIV * gmOffsetWBDelta], + layoutDstB, + gmSrcB[gmOffsetB], + layoutWB, + bufferIndex); + } + if (ldk < kLoop - 1) { // 后续块 + kActual_ = (problemShape.k() < (ldk + 2) * params.splitkLength) ? kResidueLenght : params.splitkLength; + + Catlass::Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(flag1[crossCoreBufferIndexAIV]); + Catlass::Arch::CrossCoreWaitFlag(flag0[1 - crossCoreBufferIndexAIV]); + gmOffsetA += layoutSrcA.GetOffset(MakeCoord(0U, kActual)); + gmOffsetB += layoutSrcB.GetOffset(MakeCoord(kActual, 0U)); + + uint32_t kActualAligned_ = RoundUp<256, uint32_t>(kActual_); + LayoutA layoutWA(actualBlockShape.m(), kActual_, srcAStride); + LayoutA layoutDstA(actualBlockShape.m(), kActualAligned_); + LayoutB layoutWB(kActual_, actualBlockShape.n(), srcBStride); + LayoutB layoutDstB(kActualAligned_, actualBlockShape.n()); + + PrologueA prologueA(resource, params.prologueAParams); + PrologueB prologueB(resource, params.prologueBParams); + prologueA( + gmDstA[gmOffsetWA + (1 - crossCoreBufferIndexAIV) * gmOffsetWADelta], + layoutDstA, + gmSrcA[gmOffsetA], + layoutWA, + bufferIndex); + prologueB( + gmDstB[gmOffsetWB + (1 - crossCoreBufferIndexAIV) * gmOffsetWBDelta], + layoutDstB, + gmSrcB[gmOffsetB], + layoutWB, + bufferIndex); + } + if ((ldk == kLoop - 1) && hasNextBlock) { // SliceK尾部块 + kActual_ = (problemShape.k() < params.splitkLength) ? kResidueLenght : params.splitkLength; + + Catlass::Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(flag1[crossCoreBufferIndexAIV]); + Catlass::Arch::CrossCoreWaitFlag(flag0[1 - crossCoreBufferIndexAIV]); + + uint32_t kActualAligned_ = RoundUp<256, uint32_t>(kActual_); + LayoutA layoutWA(nextActualBlockShape.m(), kActual_, srcAStride); + LayoutA layoutDstA(nextActualBlockShape.m(), kActualAligned_); + LayoutB layoutWB(kActual_, nextActualBlockShape.n(), srcBStride); + LayoutB layoutDstB(kActualAligned_, nextActualBlockShape.n()); + + PrologueA prologueA(resource, params.prologueAParams); + PrologueB prologueB(resource, params.prologueBParams); + prologueA( + gmDstA[gmOffsetWA + (1 - crossCoreBufferIndexAIV) * gmOffsetWADelta], + layoutDstA, + gmSrcA[gmOffsetNextA], + layoutWA, + bufferIndex); + prologueB( + gmDstB[gmOffsetWB + (1 - crossCoreBufferIndexAIV) * gmOffsetWBDelta], + layoutDstB, + gmSrcB[gmOffsetNextB], + layoutWB, + bufferIndex); + } + + if (ldk == kLoop - 1) { // 尾部块 + if (!hasNextBlock) { + Catlass::Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(flag1[crossCoreBufferIndexAIV]); + } + Catlass::Arch::CrossCoreWaitFlag(flag4); + + Catlass::layout::RowMajor layoutBlockC( + actualBlockShape.m(), actualBlockShape.n(), params.nScalar * L1TileShape::N); + Catlass::layout::RowMajor layoutBlockDstC( + problemShape.m(), problemShape.n()); + + PrologueC epCast; + epCast.EpCastFp32ToFp16( + gmDstC[gmOffsetC], + layoutBlockDstC, + gmSrcC[gmOffsetWC], + layoutBlockC); + } + crossCoreBufferIndexAIV = 1 - crossCoreBufferIndexAIV; + } // end: for-loop (ldk) + } + + CATLASS_DEVICE + void MNScalerMatmulImpl( + AscendC::GlobalTensor const &gmWA, + AscendC::GlobalTensor const &gmWB, + AscendC::GlobalTensor const &gmWC, + [[maybe_unused]] GemmCoord blockCoord, + GemmCoord actualBlockShape, + GemmCoord problemShape) + { + constexpr bool isLayoutAColumnMajor = std::is_same_v; + constexpr bool isLayoutBColumnMajor = std::is_same_v; + + // uint32_t splitKLen = params.splitkLength; + uint32_t kLoop = CeilDiv(problemShape.k(), params.splitkLength); + uint32_t mLoop = CeilDiv(actualBlockShape.m()); + uint32_t nLoop = CeilDiv(actualBlockShape.n()); + + uint32_t nextMIdx, nextNIdx; + // SPLIT-K MainLoop + for (uint32_t ldk = 0; ldk < kLoop; ldk++) { // K-loop + bool isFirstKSlice = (ldk == 0); + uint32_t kActual = (problemShape.k() < (ldk + 1) * params.splitkLength) ? + problemShape.k() % params.splitkLength : + params.splitkLength; + uint32_t kActualAligned = RoundUp<256>(kActual); + + uint32_t srcStrideA = isLayoutAColumnMajor ? actualBlockShape.m() : kActualAligned; + uint32_t srcStrideB = isLayoutBColumnMajor ? kActualAligned : actualBlockShape.n(); + + // CV 同步 + Catlass::Arch::CrossCoreWaitFlag(flag1[crossCoreBufferIndexAIC]); + + // M-N-loop + for (uint32_t mIdx = 0; mIdx < mLoop; mIdx++) { // M-loop + for (uint32_t nIdx = 0; nIdx < nLoop; nIdx++) { // N-loop + bool isFirstBlock = (mIdx == 0 && nIdx == 0); + bool hasNextBlock = !(mIdx == mLoop - 1 && nIdx == nLoop - 1); + + // 当前偏移及实际尺寸计算(No padding) + int64_t gmOffsetWA_ = AscendC::GetBlockIdx() * (L1TileShape::M * params.mScalar) * + params.splitkLength * STAGES; + int64_t gmOffsetWB_ = AscendC::GetBlockIdx() * (L1TileShape::N * params.nScalar) * + params.splitkLength * STAGES; + gmOffsetWA_ += crossCoreBufferIndexAIC ? params.splitkLength * + (L1TileShape::M * params.mScalar) : + 0; + gmOffsetWB_ += crossCoreBufferIndexAIC ? params.splitkLength * + (L1TileShape::N * params.nScalar) : + 0; + int64_t gmOffsetWA = isLayoutAColumnMajor ? gmOffsetWA_ + mIdx * L1TileShape::M : + gmOffsetWA_ + mIdx * L1TileShape::M * kActualAligned; + int64_t gmOffsetWB = isLayoutBColumnMajor ? gmOffsetWB_ + nIdx * L1TileShape::N * + kActualAligned : + gmOffsetWB_ + nIdx * L1TileShape::N; + int64_t gmOffsetWC = AscendC::GetBlockIdx() * (L1TileShape::M * params.mScalar) * + (L1TileShape::N * params.nScalar) + + mIdx * L1TileShape::M * + (L1TileShape::N * params.nScalar) + + nIdx * L1TileShape::N; + + uint32_t mActual = (mIdx == mLoop - 1 && actualBlockShape.m() % L1TileShape::M != 0) ? + actualBlockShape.m() % L1TileShape::M : + L1TileShape::M; + uint32_t nActual = (nIdx == nLoop - 1 && actualBlockShape.n() % L1TileShape::N != 0) ? + actualBlockShape.n() % L1TileShape::N : + L1TileShape::N; + LayoutA layoutWA(mActual, kActual, srcStrideA); + LayoutB layoutWB(kActual, nActual, srcStrideB); + LayoutC layoutWC(mActual, nActual, params.nScalar * L1TileShape::N); + GemmCoord actualSmallBlockShape(mActual, nActual, kActual); + + // Next实际尺寸(No padding) + int64_t gmOffsetWANext, gmOffsetWBNext, gmOffsetWCNext; + LayoutA layoutWANext; + LayoutB layoutWBNext; + LayoutC layoutWCNext; + GemmCoord nextSmallBlockShape; + if (hasNextBlock) { + uint32_t nextNIdx = (nIdx + 1) % nLoop; + uint32_t nextMIdx = mIdx + static_cast(nextNIdx == 0); + + gmOffsetWANext = isLayoutAColumnMajor ? gmOffsetWA_ + nextMIdx * L1TileShape::M : + gmOffsetWA_ + nextMIdx * L1TileShape::M * kActualAligned; + gmOffsetWBNext = isLayoutBColumnMajor ? gmOffsetWB_ + nextNIdx * L1TileShape::N * + kActualAligned : + gmOffsetWB_ + nextNIdx * L1TileShape::N; + gmOffsetWCNext = AscendC::GetBlockIdx() * (L1TileShape::M * params.mScalar) * + (L1TileShape::N * params.nScalar) + + nextMIdx * L1TileShape::M * + (L1TileShape::N * params.nScalar) + + nextNIdx * L1TileShape::N; + + uint32_t mActualNext = (nextMIdx == mLoop - 1 && actualBlockShape.m() % L1TileShape::M != 0) ? + actualBlockShape.m() % L1TileShape::M : + L1TileShape::M; + uint32_t nActualNext = (nextNIdx == nLoop - 1 && actualBlockShape.n() % L1TileShape::N != 0) ? + actualBlockShape.n() % L1TileShape::N : + L1TileShape::N; + layoutWANext = LayoutA(mActualNext, kActual, srcStrideA); + layoutWBNext = LayoutB(kActual, nActualNext, srcStrideB); + layoutWCNext = LayoutC(mActualNext, nActualNext, params.nScalar * L1TileShape::N); + nextSmallBlockShape = GemmCoord(mActualNext, nActualNext, kActual); + } + + // 完成一个128 * 256的小结果矩阵基本块的运算 + smallBlockMmadImpl(gmWA[gmOffsetWA], + layoutWA, + gmWB[gmOffsetWB], + layoutWB, + gmWC[gmOffsetWC], + layoutWC, + gmWA[gmOffsetWANext], + layoutWANext, + gmWB[gmOffsetWBNext], + layoutWBNext, + actualSmallBlockShape, + nextSmallBlockShape, + isFirstKSlice, + isFirstBlock, + hasNextBlock); + } // End of 'N-loop' + } // End of 'M-loop' + + Catlass::Arch::CrossCoreSetFlag<0x2, PIPE_FIX>(flag0[crossCoreBufferIndexAIC]); + crossCoreBufferIndexAIC = 1 - crossCoreBufferIndexAIC; + if (ldk == kLoop - 1) { + // cast 256 * 512的fp32大结果基本块为fp16 + Catlass::Arch::CrossCoreSetFlag<0x2, PIPE_FIX>(flag4); + } + } // End of 'K-loop' + } + + CATLASS_DEVICE + void smallBlockMmadImpl( + AscendC::GlobalTensor const &gmWA, LayoutA const &layoutA, + AscendC::GlobalTensor const &gmWB, LayoutB const &layoutB, + AscendC::GlobalTensor const &gmC, LayoutC const &layoutC, + AscendC::GlobalTensor const &gmNextWA, LayoutA const &layoutNextA, + AscendC::GlobalTensor const &gmNextWB, LayoutB const &layoutNextB, + GemmCoord const &actualShape, GemmCoord const &nextActualShape, + bool isFirstKSlice, bool isFirstBlock, bool hasNextBlock) + { + uint32_t mRound = RoundUp(actualShape.m()); + uint32_t nRound = RoundUp(actualShape.n()); + + auto layoutAInL1 = LayoutAInL1::template MakeLayout(L1TileShape::M, L1TileShape::K); + auto layoutBInL1 = LayoutBInL1::template MakeLayout(L1TileShape::K, L1TileShape::N); + auto layoutInL0C = LayoutCInL0::MakeLayoutInL0C(MakeCoord(mRound, nRound)); + + uint32_t kActual = min(actualShape.k(), L1TileShape::K); + + // load first matrix A tile from GM to L1 + if (isFirstBlock) { + AscendC::WaitFlag(l1AEventList[l1ListId]); + auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShape.m(), kActual)); + copyGmToL1A(l1ATensorList[l1ListId], gmWA, layoutAInL1, layoutTileA); + AscendC::SetFlag(l1AEventList[l1ListId]); + + // load first matrix B tile from GM to L1 + AscendC::WaitFlag(l1BEventList[l1ListId]); + auto layoutTileB = layoutB.GetTileLayout(MakeCoord(kActual, actualShape.n())); + copyGmToL1B(l1BTensorList[l1ListId], gmWB, layoutBInL1, layoutTileB); + AscendC::SetFlag(l1BEventList[l1ListId]); + } + + if constexpr (!ENABLE_UNIT_FLAG) { + AscendC::WaitFlag(EVENT_ID0); + } + + uint32_t mPartLoop = CeilDiv(mRound); + uint32_t nPartLoop = CeilDiv(nRound); + + // main loop + uint32_t kTileCount = CeilDiv(actualShape.k()); + for (uint32_t kLoopIdx = 0; kLoopIdx < kTileCount; kLoopIdx++) { + uint32_t l1ListIdNext = (l1ListId + 1 < STAGES) ? (l1ListId + 1) : 0; + uint32_t kActualNext{0}; + // preload next tile from GM to L1 + if (kLoopIdx < kTileCount - 1) { + uint32_t kLoopIdxNext = kLoopIdx + 1; + kActualNext = (kLoopIdxNext < kTileCount - 1) ? + L1TileShape::K : + (actualShape.k() - kLoopIdxNext * L1TileShape::K); + + // Get L1 tensor for next stage + auto l1ATensor = l1ATensorList[l1ListIdNext]; + auto l1BTensor = l1BTensorList[l1ListIdNext]; + // Get GM tile for next stage + MatrixCoord gmTileAOffset{0, kLoopIdxNext * L1TileShape::K}; + MatrixCoord gmTileBOffset{kLoopIdxNext * L1TileShape::K, 0}; + auto gmTileA = gmWA[layoutA.GetOffset(gmTileAOffset)]; + auto gmTileB = gmWB[layoutB.GetOffset(gmTileBOffset)]; + + // load next matrix A tile from GM to L1 + AscendC::WaitFlag(l1AEventList[l1ListIdNext]); + auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShape.m(), kActualNext)); + copyGmToL1A(l1ATensor, gmTileA, layoutAInL1, layoutTileA); + AscendC::SetFlag(l1AEventList[l1ListIdNext]); + + // load next matrix B tile from GM to L1 + AscendC::WaitFlag(l1BEventList[l1ListIdNext]); + auto layoutTileB = layoutB.GetTileLayout(MakeCoord(kActualNext, actualShape.n())); + copyGmToL1B(l1BTensor, gmTileB, layoutBInL1, layoutTileB); + AscendC::SetFlag(l1BEventList[l1ListIdNext]); + } else if (hasNextBlock) { + uint32_t kLoopIdxNext = 0; + kActualNext = (kLoopIdxNext < kTileCount - 1) ? L1TileShape::K : (nextActualShape.k()); + + // Get L1 tensor for next stage + auto l1ATensor = l1ATensorList[l1ListIdNext]; + auto l1BTensor = l1BTensorList[l1ListIdNext]; + // Get GM tile for next stage + MatrixCoord gmTileAOffset{0, kLoopIdxNext * L1TileShape::K}; + MatrixCoord gmTileBOffset{kLoopIdxNext * L1TileShape::K, 0}; + auto gmTileA = gmNextWA[layoutNextA.GetOffset(gmTileAOffset)]; + auto gmTileB = gmNextWB[layoutNextB.GetOffset(gmTileBOffset)]; + + // load next matrix A tile from GM to L1 + AscendC::WaitFlag(l1AEventList[l1ListIdNext]); + auto layoutTileA = layoutNextA.GetTileLayout(MakeCoord(nextActualShape.m(), kActualNext)); + copyGmToL1A(l1ATensor, gmTileA, layoutAInL1, layoutTileA); + AscendC::SetFlag(l1AEventList[l1ListIdNext]); + + // load next matrix B tile from GM to L1 + AscendC::WaitFlag(l1BEventList[l1ListIdNext]); + auto layoutTileB = layoutNextB.GetTileLayout(MakeCoord(kActualNext, nextActualShape.n())); + copyGmToL1B(l1BTensor, gmTileB, layoutBInL1, layoutTileB); + AscendC::SetFlag(l1BEventList[l1ListIdNext]); + } + + // Get L1 tensor for current stage + auto l1ATensor = l1ATensorList[l1ListId]; + auto l1BTensor = l1BTensorList[l1ListId]; + + // Get the loop nums on L0 + uint32_t kPartLoop = CeilDiv(kActual); + + for (int mPartIdx = 0; mPartIdx < mPartLoop; mPartIdx++) { + uint32_t mPartActual = (mPartIdx < mPartLoop - 1) ? + L0TileShape::M : + (mRound - mPartIdx * L0TileShape::M); + + for (int kPartIdx = 0; kPartIdx < kPartLoop; kPartIdx++) { + uint32_t kPartActual = (kPartIdx < kPartLoop - 1) ? + L0TileShape::K : + (kActual - kPartIdx * L0TileShape::K); + + // Locate the current tile on L0A + auto l0ATile = l0ATensorList[l0AListId]; + LayoutAInL0 layoutAInL0 = LayoutAInL0::template MakeLayout(mPartActual, kPartActual); + // Locate the current tile of matrix A on L1 + MatrixCoord l1AOffset{mPartIdx * L0TileShape::M, kPartIdx * L0TileShape::K}; + auto l1ATile = l1ATensor[layoutAInL1.GetOffset(l1AOffset)]; + + AscendC::WaitFlag(l0AEventList[l0AListId]); + if ((mPartIdx == 0) && (kPartIdx == 0)) { + AscendC::WaitFlag(l1AEventList[l1ListId]); + } + + // Load current tile from L1 to L0A + copyL1ToL0A(l0ATile, l1ATile, layoutAInL0, layoutAInL1); + + if ((mPartIdx == mPartLoop - 1) && (kPartIdx == kPartLoop - 1)) { + AscendC::SetFlag(l1AEventList[l1ListId]); + } + + for (int nPartIdx = 0; nPartIdx < nPartLoop; nPartIdx++) { + uint32_t nPartActual = (nPartIdx < nPartLoop - 1) ? + L0TileShape::N : + (nRound - nPartIdx * L0TileShape::N); + + // Locate the current tile on L0B + auto l0BTile = l0BTensorList[l0BListId]; + LayoutBInL0 layoutBInL0 = LayoutBInL0::template MakeLayout(kPartActual, nPartActual); + // Locate the current tile of matrix B on L1 + MatrixCoord l1BOffset{kPartIdx * L0TileShape::K, nPartIdx * L0TileShape::N}; + auto l1BTile = l1BTensor[layoutBInL1.GetOffset(l1BOffset)]; + + // Wait for mmad finished + AscendC::WaitFlag(l0BEventList[l0BListId]); + // If the current tile is the first one on the k&n axis, wait for loading matrix B from GM to L1 + if ((kPartIdx == 0) && (nPartIdx == 0)) { + AscendC::WaitFlag(l1BEventList[l1ListId]); + } + + // Load current tile from L1 to L0B + copyL1ToL0B(l0BTile, l1BTile, layoutBInL0, layoutBInL1); + + // If the current tile is the last one on the k&n axis, notify to load matrix B from GM to L1 + if ((kPartIdx == kPartLoop - 1) && (nPartIdx == nPartLoop - 1)) { + AscendC::SetFlag(l1BEventList[l1ListId]); + } + // Notify to do mmad + AscendC::SetFlag(EVENT_ID0); + + // Locate the current tile on L0C + MatrixCoord l0COffset{mPartIdx * L0TileShape::M, nPartIdx * L0TileShape::N}; + auto l0CTile = l0CTensor[layoutInL0C.GetOffset(l0COffset)]; + + // Compute the matrix multiplication on L0A and L0B and write the result to the accumulator + // Wait for loading L0B + AscendC::WaitFlag(EVENT_ID0); + + // If the current tile is the first tile on the k axis, the accumulator needs to be reset to 0 + bool initC = ((kLoopIdx == 0) && (kPartIdx == 0)); + // If the unit flag is enabled, the unit flag is set according to the calculation progress + uint8_t unitFlag = 0b00; + if constexpr (ENABLE_UNIT_FLAG) { + if ((kLoopIdx == kTileCount - 1) && (mPartIdx == mPartLoop - 1) && + (kPartIdx == kPartLoop - 1) && (nPartIdx == nPartLoop - 1)) { + unitFlag = 0b11; + } else { + unitFlag = 0b10; + } + } + // Perform calculation operations + tileMmad(l0CTile, l0ATile, l0BTile, mPartActual, nPartActual, kPartActual, initC, unitFlag); + + // Notify to move the next L0B tile + AscendC::SetFlag(l0BEventList[l0BListId]); + l0BListId = (l0BListId + 1 < STAGES) ? (l0BListId + 1) : 0; + } + AscendC::SetFlag(l0AEventList[l0AListId]); + l0AListId = (l0AListId + 1 < STAGES) ? (l0AListId + 1) : 0; + } + } + l1ListId = l1ListIdNext; + kActual = kActualNext; + } + + // copy block out + LayoutC layoutBlock = layoutC.GetTileLayout(actualShape.GetCoordMN()); + + if (!isFirstKSlice) { // 切K后非第一块 开启原子加 + AscendC::SetAtomicAdd(); + } + if constexpr (!ENABLE_UNIT_FLAG) { + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + copyL0CToGm(gmC, l0CTensor, layoutBlock, layoutInL0C); + AscendC::SetFlag(EVENT_ID0); + } else { + copyL0CToGm(gmC, l0CTensor, layoutBlock, layoutInL0C, 0b11); + } + if (!isFirstKSlice) { + AscendC::SetAtomicNone(); + } + } + +protected: + /// Data members + Params params; + + // static constexpr uint32_t STAGES = 2; + static constexpr uint32_t AIVPERCORE = 2; + + // Multi-stage tensors list + AscendC::LocalTensor l1ATensorList[STAGES]; + AscendC::LocalTensor l1BTensorList[STAGES]; + AscendC::LocalTensor l0ATensorList[STAGES]; + AscendC::LocalTensor l0BTensorList[STAGES]; + AscendC::LocalTensor l0CTensor; + + // Multi-stage event id list + int32_t l1AEventList[STAGES]; + int32_t l1BEventList[STAGES]; + int32_t l0AEventList[STAGES]; + int32_t l0BEventList[STAGES]; + + // The id of current stage + uint32_t l1ListId{0}; + uint32_t l0AListId{0}; + uint32_t l0BListId{0}; + + // CV Flag + static constexpr Arch::FlagID flagID0 = 0; + static constexpr Arch::FlagID flagID1 = 1; + static constexpr Arch::FlagID flagID2 = 2; + static constexpr Arch::FlagID flagID3 = 3; + static constexpr Arch::FlagID flagID4 = 4; + + Arch::CrossCoreFlag flag0[STAGES]; + Arch::CrossCoreFlag flag1[STAGES]; + Arch::CrossCoreFlag flag4{flagID4}; + + uint32_t crossCoreBufferIndexAIC{0}; + uint32_t crossCoreBufferIndexAIV{0}; + + // Tile action + TileMmad tileMmad; + CopyGmToL1A copyGmToL1A; + CopyGmToL1B copyGmToL1B; + CopyL1ToL0A copyL1ToL0A; + CopyL1ToL0B copyL1ToL0B; + CopyL0CToGm copyL0CToGm; + + // Buffer index (for prologue) + uint32_t bufferIndex{0}; + + Tile::PrologueTraits prologueA; + Tile::PrologueTraits prologueB; +}; + +} // namespace Catlass::Gemm::Block +#endif // CATLASS_GEMM_BLOCK_BLOCK_MMAD_PINGPONG_SLICE_K_WITH_PROLOGUE_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_pingpong_tla.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_pingpong_tla.hpp new file mode 100644 index 0000000000..f8049a3da5 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_pingpong_tla.hpp @@ -0,0 +1,630 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_GEMM_BLOCK_BLOCK_MMAD_PINGPONG_TLA_HPP +#define CATLASS_GEMM_BLOCK_BLOCK_MMAD_PINGPONG_TLA_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/kda_coord.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" +#include "catlass/gemm/kda_gemm_helper.hpp" +#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" +#include "catlass/gemm/tile/kda_tile_mmad.hpp" +#include "tla/kda_tla_layout.hpp" +#include "tla/kda_tensor.hpp" + +namespace Catlass::Gemm::Block { + +template < + class ArchTag_, + bool ENABLE_UNIT_FLAG_, + bool USE_HF32_MODE_, + uint32_t L0C_STAGES_, + bool ENABLE_L1_RESIDENT_, + uint32_t L1A_STAGES_, + uint32_t L1B_STAGES_, + uint32_t L0A_STAGES_, + uint32_t L0B_STAGES_, + class L1TileShape_, + class L0TileShape_, + class ElementA_, + class ElementB_, + class ElementC_, + class ElementBias_, + class TileCopy_, + class TileMmad_> +struct BlockMmadTla< + MmadPingpong, + L1TileShape_, + L0TileShape_, + ElementA_, + ElementB_, + ElementC_, + ElementBias_, + TileCopy_, + TileMmad_> { +public: + // Type Aliases + using DispatchPolicy = MmadPingpong; + using ArchTag = typename DispatchPolicy::ArchTag; + using TileCopy = TileCopy_; + using L1TileShape = L1TileShape_; + using L0TileShape = L0TileShape_; + using ElementA = ElementA_; + using LayoutA = typename TileCopy::LayoutA; + using ElementB = ElementB_; + using LayoutB = typename TileCopy::LayoutB; + using ElementC = ElementC_; + using LayoutC = typename TileCopy::LayoutC; + using ElementBias = ElementBias_; + + using TileMmad = TileMmad_; + + using CopyL1ToL0A = typename TileCopy::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopy::CopyL1ToL0B; + using CopyL1ToBT = typename TileCopy::CopyL1ToBT; + + using ElementAccumulator = typename TileCopy::ElementAccumulator; + + static constexpr bool HAS_BIAS = TileCopy::HAS_BIAS; + + using LayoutTagL1A = typename TileCopy::LayoutTagL1A; + using LayoutTagL1B = typename TileCopy::LayoutTagL1B; + using LayoutTagL0A = typename TileCopy::LayoutTagL0A; + using LayoutTagL0B = typename TileCopy::LayoutTagL0B; + + using L1AAlignHelper = typename TileCopy_::L1AAlignHelper; + using L1BAlignHelper = typename TileCopy_::L1BAlignHelper; + + static_assert(tla::is_tuple::value && tla::is_static::value, + "L1TileShape must be tla::tuple and static!"); + static_assert(tla::is_tuple::value && tla::is_static::value, + "L0TileShape must be tla::tuple and static!"); + + static constexpr bool ENABLE_UNIT_FLAG = DispatchPolicy::ENABLE_UNIT_FLAG; + static constexpr bool USE_HF32_MODE = DispatchPolicy::USE_HF32_MODE; + static constexpr bool ENABLE_L1_RESIDENT = DispatchPolicy::ENABLE_L1_RESIDENT; + static constexpr uint32_t L1A_STAGES = DispatchPolicy::L1A_STAGES; + static constexpr uint32_t L1B_STAGES = DispatchPolicy::L1B_STAGES; + static constexpr uint32_t L0A_STAGES = DispatchPolicy::L0A_STAGES; + static constexpr uint32_t L0B_STAGES = DispatchPolicy::L0B_STAGES; + static constexpr uint32_t L0C_STAGES = DispatchPolicy::L0C_STAGES; + static constexpr uint32_t L1_TILE_M = tla::get<0>(L1TileShape{}); + static constexpr uint32_t L1_TILE_N = tla::get<1>(L1TileShape{}); + static constexpr uint32_t L1_TILE_K = tla::get<2>(L1TileShape{}); + static constexpr uint32_t L0_TILE_M = tla::get<0>(L0TileShape{}); + static constexpr uint32_t L0_TILE_N = tla::get<1>(L0TileShape{}); + static constexpr uint32_t L0_TILE_K = tla::get<2>(L0TileShape{}); + + // L1 tile size + static constexpr uint32_t L1A_TILE_SIZE = L1_TILE_M * L1_TILE_K * sizeof(ElementA); + static constexpr uint32_t L1B_TILE_SIZE = L1_TILE_N * L1_TILE_K * sizeof(ElementB); + // L0 tile size + static constexpr uint32_t L0A_TILE_SIZE = L0_TILE_M * L0_TILE_K * sizeof(ElementA); + static constexpr uint32_t L0B_TILE_SIZE = L0_TILE_K * L0_TILE_N * sizeof(ElementB); + static constexpr uint32_t L0C_TILE_SIZE = L1_TILE_M * L1_TILE_N * sizeof(ElementAccumulator); + + // Check HF32_MODE + static_assert( + !USE_HF32_MODE || (USE_HF32_MODE && std::is_same_v && std::is_same_v), + "HF32 MODE only supports in float!"); + + // Check L0C_STAGES + static_assert(!(ENABLE_UNIT_FLAG && L0C_STAGES != 1), "L0C_STAGES must be 1 when UnitFlag is true!"); + + // Check LayoutC + static_assert(tla::detail::isRowMajor::value || + ((std::is_same_v || std::is_same_v || + std::is_same_v) && + tla::detail::iszN::value), + "LayoutC only supports zN in half or bfloat16 or float, RowMajor in all dtype yet!"); + + // Check L1TileShape + static_assert(L1A_TILE_SIZE * L1A_STAGES + L1B_TILE_SIZE * L1B_STAGES <= ArchTag::L1_SIZE, + "L1TileShape exceeding the L1 space!"); + + // Check L0TileShape + static_assert(L0A_TILE_SIZE * L0A_STAGES <= ArchTag::L0A_SIZE, "L0TileShape exceeding the L0A space!"); + static_assert(L0B_TILE_SIZE * L0B_STAGES <= ArchTag::L0B_SIZE, "L0TileShape exceeding the L0B space!"); + static_assert(L0C_TILE_SIZE * L0C_STAGES <= ArchTag::L0C_SIZE, "L0TileShape exceeding the L0C space!"); + + static constexpr uint32_t _32B = 32 * 8; // in bits + static_assert(L1_TILE_M == L0_TILE_M && L1_TILE_N == L0_TILE_N, + "The situation where the basic blocks of L1 and L0 differ on the m and n axes is not supported yet"); + static_assert(L0_TILE_K <= L1_TILE_K, "L0TileShape::K cannot exceed L1TileShape::K"); +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 2201) + static_assert(L1_TILE_M * SizeOfBits::value % _32B == 0, "L1TileShape::M must be 32B aligned."); + static_assert(L1_TILE_K * SizeOfBits::value % _32B == 0, "L1TileShape::K must be 32B aligned."); + static_assert(L1_TILE_K * SizeOfBits::value % _32B == 0, "L1TileShape::K must be 32B aligned."); + static_assert(L1_TILE_N * SizeOfBits::value % _32B == 0, "L1TileShape::N must be 32B aligned."); + static_assert(L0_TILE_K * SizeOfBits::value % _32B == 0, "L0TileShape::K must be 32B aligned."); +#endif + + static_assert((!HAS_BIAS && (L1A_STAGES + L1B_STAGES) <= 8) || (HAS_BIAS && (L1A_STAGES + L1B_STAGES) <= 7), + "L1 Buffer overflow: Exceeds the supported range of EVENT(0~7)"); + + static_assert((!HAS_BIAS && (L0A_STAGES + L0B_STAGES) <= 8) || (HAS_BIAS && (L0A_STAGES + L0B_STAGES) <= 7), + "L0 Buffer overflow: Exceeds the supported range of EVENT_ID(0~7)"); + + static constexpr auto L1A_LAYOUT = + tla::MakeLayout(tla::Int{}, tla::Int{}); + static constexpr auto L1B_LAYOUT = + tla::MakeLayout(tla::Int{}, tla::Int{}); + static constexpr auto L1BIAS_LAYOUT = tla::MakeLayout(tla::Int{}); + static constexpr auto L0BIAS_LAYOUT = tla::MakeLayout(tla::Int{}); + + // When enabling L1 resident mode, restore the pointer and coordinates that record the last state + // to the initial state. if two blockmmad instances need to be consecutively invoked at the kernel layer, + // RestoreStatus() must be inserted between them. + CATLASS_DEVICE + void RestoreStatus() + { + for (int i = 0; i < L1A_STAGES; ++i) { + lastAddrA[i] = nullptr; + lastCoordA[i] = MatrixCoord{0U, 0U}; + } + for (int i = 0; i < L1B_STAGES; ++i) { + lastAddrB[i] = nullptr; + lastCoordB[i] = MatrixCoord{0U, 0U}; + } + } + + /// Construct + CATLASS_DEVICE + BlockMmadTla(Arch::Resource &resource, uint32_t l1BufAddrStart = 0) + { + if ASCEND_IS_AIC { + // use HF32 when USE_HF32_MODE is true + if constexpr (USE_HF32_MODE) { + AscendC::SetHF32Mode(true); + } else { + AscendC::SetHF32Mode(false); + } + if constexpr (ENABLE_UNIT_FLAG && tla::detail::isRowMajor::value) { + AscendC::SetMMLayoutTransform(true); + } + uint32_t l1AOffset = l1BufAddrStart; + uint32_t l1BOffset = l1BufAddrStart + L1A_TILE_SIZE * L1A_STAGES; + // Init buffers + for (uint32_t i = 0; i < L1A_STAGES; i++) { + // Assign L1/L0A/L0B space for each stages + l1ATensorList[i] = resource.l1Buf.template GetBufferByByte(l1AOffset + L1A_TILE_SIZE * i); + // Assign event ID for each stages + l1AEventList[i] = i; + // The event id that needs to be set before the loop + AscendC::SetFlag(l1AEventList[i]); + } + for (uint32_t i = 0; i < L1B_STAGES; i++) { + // Assign L1/L0A/L0B space for each stages + l1BTensorList[i] = resource.l1Buf.template GetBufferByByte(l1BOffset + L1B_TILE_SIZE * i); + // Assign event ID for each stages + l1BEventList[i] = i + L1A_STAGES; + // The event id that needs to be set before the loop + AscendC::SetFlag(l1BEventList[i]); + } + for (uint32_t i = 0; i < L0A_STAGES; i++) { + // Assign L1/L0A/L0B space for each stages + l0ATensorList[i] = resource.l0ABuf.template GetBufferByByte(L0A_TILE_SIZE * i); + // Assign event ID for each stages + l0AEventList[i] = i; + // The event id that needs to be set before the loop + AscendC::SetFlag(l0AEventList[i]); + } + for (uint32_t i = 0; i < L0B_STAGES; i++) { + // Assign L1/L0A/L0B space for each stages + l0BTensorList[i] = resource.l0BBuf.template GetBufferByByte(L0B_TILE_SIZE * i); + // Assign event ID for each stages + l0BEventList[i] = i + L0A_STAGES; + // The event id that needs to be set before the loop + AscendC::SetFlag(l0BEventList[i]); + } + if constexpr (!ENABLE_UNIT_FLAG) { + for (uint32_t i = 0; i < L0C_STAGES; i++) { + l0CTensorList[i] = resource.l0CBuf.template GetBufferByByte(L0C_TILE_SIZE * i); + l0CEventList[i] = i; + AscendC::SetFlag(l0CEventList[i]); + } + } else { + l0CTensorList[0] = resource.l0CBuf.template GetBufferByByte(0); + } + if constexpr (HAS_BIAS) { + uint32_t l1BiasOffset = l1BOffset + L1B_TILE_SIZE * L1B_STAGES; + l1BiasTensor = resource.l1Buf.template GetBufferByByte(l1BiasOffset); + l0BiasTensor = resource.btBuf.template GetBufferByByte(0); + AscendC::SetFlag(L1A_STAGES + L1B_STAGES); + AscendC::SetFlag(L0A_STAGES + L0B_STAGES); + } + + if constexpr (ENABLE_L1_RESIDENT) { + RestoreStatus(); + } + } + } + + /// Destructor + CATLASS_DEVICE + ~BlockMmadTla() + { + if ASCEND_IS_AIC { + if constexpr (USE_HF32_MODE) { + AscendC::SetHF32Mode(false); + } + if constexpr (ENABLE_UNIT_FLAG && tla::detail::isRowMajor::value) { + AscendC::SetMMLayoutTransform(false); + } + for (uint32_t i = 0; i < L1A_STAGES; i++) { + AscendC::WaitFlag(l1AEventList[i]); + } + for (uint32_t i = 0; i < L1B_STAGES; i++) { + AscendC::WaitFlag(l1BEventList[i]); + } + for (uint32_t i = 0; i < L0A_STAGES; i++) { + AscendC::WaitFlag(l0AEventList[i]); + } + for (uint32_t i = 0; i < L0B_STAGES; i++) { + AscendC::WaitFlag(l0BEventList[i]); + } + if constexpr (!ENABLE_UNIT_FLAG) { + for (uint32_t i = 0; i < L0C_STAGES; i++) { + AscendC::WaitFlag(l0CEventList[i]); + } + } + if constexpr (HAS_BIAS) { + AscendC::WaitFlag(L1A_STAGES + L1B_STAGES); + AscendC::WaitFlag(L0A_STAGES + L0B_STAGES); + } + } + } + + /// Perform a block-scoped matrix multiply-accumulate + template + CATLASS_DEVICE void operator()(TensorA &tensorA, TensorB &tensorB, TensorC &tensorC, GemmCoord const &actualShape, + TensorBias const &tensorBias = {}) + { + // Check L1TileShape + if constexpr (HAS_BIAS) { + static constexpr uint32_t BIAS_BUF_SIZE = L0_TILE_N * sizeof(ElementAccumulator); + static constexpr uint32_t L1BIAS_SIZE = L1_TILE_N * sizeof(ElementBias); + static_assert(BIAS_BUF_SIZE <= ArchTag::BIAS_SIZE, + "BIAS_BUF_SIZE exceeding the BT space! Reduce L0_TILE_N"); + static_assert(L1A_TILE_SIZE * L1A_STAGES + L1B_TILE_SIZE * L1B_STAGES + L1BIAS_SIZE <= ArchTag::L1_SIZE, + "L1TileShape exceeding the L1 space!"); + } + + using CopyGmToL1A = typename TileCopy_::template CopyGmToL1A; + using CopyGmToL1B = typename TileCopy_::template CopyGmToL1B; + CopyGmToL1A copyGmToL1A; + CopyGmToL1B copyGmToL1B; +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 2201) + using CopyL0CToGm = typename TileCopy_::template CopyL0CToGm; + CopyL0CToGm copyL0CToDst; +#endif +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 3510) + using CopyL0CToDst = typename TileCopy_::template CopyL0CToDst; + CopyL0CToDst copyL0CToDst; +#endif + + uint32_t mBlockActual = actualShape.m(); + uint32_t kBlockActual = actualShape.k(); + uint32_t nBlockActual = actualShape.n(); + + uint32_t mL1Actual = mBlockActual; + if constexpr (std::is_same_v) { + // Avoid using the gemv mode in mmad + if (mL1Actual == 1) { + mL1Actual = 16; + } + } + uint32_t nL1Actual = nBlockActual; + + auto layoutInL0C = tla::MakeLayoutL0C(mL1Actual, nL1Actual); + auto tensorL0C = tla::MakeTensor(l0CTensorList[l0CListId], layoutInL0C, Arch::PositionL0C{}); + auto tensorL0Bias = tla::MakeTensor(l0BiasTensor, L0BIAS_LAYOUT, Arch::PositionBias{}); + + uint32_t kL1Actual = min(kBlockActual, L1_TILE_K); + // load first matrix A tile from GM to L1 + AscendC::WaitFlag(l1AEventList[l1AListId]); + auto tensorL1A = tla::MakeTensor(l1ATensorList[l1AListId], L1A_LAYOUT, Arch::PositionL1{}); + auto tensorTileA = GetTileA(tensorA, 0, 0, mBlockActual, kL1Actual); + if constexpr (ENABLE_L1_RESIDENT) { + // If the currently loaded GM pointer and block coordinates are the same as the last loaded ones, + // skip this loading. + if (lastAddrA[l1AListId] != tensorTileA.data().GetPhyAddr() || tla::get<0>(tensorTileA.coord()) != lastCoordA[l1AListId].row() || tla::get<1>(tensorTileA.coord()) != lastCoordA[l1AListId].column()) { + copyGmToL1A(tensorL1A, tensorTileA); + lastCoordA[l1AListId] = MatrixCoord{tla::get<0>(tensorTileA.coord()), tla::get<1>(tensorTileA.coord())}; + lastAddrA[l1AListId] = const_cast<__gm__ typename AscendC::GlobalTensor::PrimType *>( + tensorTileA.data().GetPhyAddr()); + } + } else { + copyGmToL1A(tensorL1A, tensorTileA); + } + AscendC::SetFlag(l1AEventList[l1AListId]); + + // load first matrix B tile from GM to L1 + AscendC::WaitFlag(l1BEventList[l1BListId]); + auto tensorL1B = tla::MakeTensor(l1BTensorList[l1BListId], L1B_LAYOUT, Arch::PositionL1{}); + auto tensorTileB = GetTile(tensorB, tla::MakeCoord(0, 0), tla::MakeShape(kL1Actual, nBlockActual)); + if constexpr (ENABLE_L1_RESIDENT) { + if (lastAddrB[l1BListId] != tensorTileB.data().GetPhyAddr() || tla::get<0>(tensorTileB.coord()) != lastCoordB[l1BListId].row() || tla::get<1>(tensorTileB.coord()) != lastCoordB[l1BListId].column()) { + copyGmToL1B(tensorL1B, tensorTileB); + lastCoordB[l1BListId] = MatrixCoord{tla::get<0>(tensorTileB.coord()), tla::get<1>(tensorTileB.coord())}; + lastAddrB[l1BListId] = const_cast<__gm__ typename AscendC::GlobalTensor::PrimType *>( + tensorTileB.data().GetPhyAddr()); + } + } else { + copyGmToL1B(tensorL1B, tensorTileB); + } + AscendC::SetFlag(l1BEventList[l1BListId]); + + if constexpr (HAS_BIAS && !std::is_same_v) { + using CopyGmToL1Bias = typename TileCopy::template CopyGmToL1Bias; + CopyGmToL1Bias copyGmToL1Bias; + AscendC::WaitFlag(L1A_STAGES + L1B_STAGES); + auto l1Bias = l1BiasTensor.template ReinterpretCast(); + auto tensorL1Bias = tla::MakeTensor(l1Bias, L1BIAS_LAYOUT, Arch::PositionL1{}); + copyGmToL1Bias(tensorL1Bias, tensorBias); + AscendC::SetFlag(L1A_STAGES + L1B_STAGES); + } + + if constexpr (!ENABLE_UNIT_FLAG) { + AscendC::WaitFlag(l0CEventList[l0CListId]); + } + + uint32_t mL0Loop = CeilDiv(mL1Actual); + uint32_t nL0Loop = CeilDiv(nL1Actual); + + // main loop + uint32_t kL1Loop = CeilDiv(kBlockActual); + for (uint32_t kL1Idx = 0; kL1Idx < kL1Loop; kL1Idx++) { + uint32_t l1AListIdNext = (l1AListId + 1 < L1A_STAGES) ? (l1AListId + 1) : 0; + uint32_t l1BListIdNext = (l1BListId + 1 < L1B_STAGES) ? (l1BListId + 1) : 0; + uint32_t kL1ActualNext{0}; + // preload next tile from GM to L1 + if (kL1Idx < kL1Loop - 1) { + uint32_t kL1IdxNext = kL1Idx + 1; + kL1ActualNext = (kL1IdxNext < kL1Loop - 1) ? L1_TILE_K : (kBlockActual - kL1IdxNext * L1_TILE_K); + + // Get L1 tensor for next stage + auto l1ATensor = l1ATensorList[l1AListIdNext]; + auto l1BTensor = l1BTensorList[l1BListIdNext]; + auto tensorL1A = tla::MakeTensor(l1ATensor, L1A_LAYOUT, Arch::PositionL1{}); + auto tensorL1B = tla::MakeTensor(l1BTensor, L1B_LAYOUT, Arch::PositionL1{}); + // Get GM tile for next stage + auto tensorTileA = GetTileA(tensorA, 0, kL1IdxNext * L1_TILE_K, mBlockActual, kL1ActualNext); + auto tensorTileB = GetTile(tensorB, tla::MakeCoord(kL1IdxNext * L1_TILE_K, 0), + tla::MakeShape(kL1ActualNext, nBlockActual)); + + // load next matrix A tile from GM to L1 + AscendC::WaitFlag(l1AEventList[l1AListIdNext]); + if constexpr (ENABLE_L1_RESIDENT) { + if (lastAddrA[l1AListIdNext] != tensorTileA.data().GetPhyAddr() || tla::get<0>(tensorTileA.coord()) != lastCoordA[l1AListIdNext].row() || tla::get<1>(tensorTileA.coord()) != lastCoordA[l1AListIdNext].column()) { + copyGmToL1A(tensorL1A, tensorTileA); + lastCoordA[l1AListIdNext] = + MatrixCoord{tla::get<0>(tensorTileA.coord()), tla::get<1>(tensorTileA.coord())}; + lastAddrA[l1AListIdNext] = + const_cast<__gm__ typename AscendC::GlobalTensor::PrimType *>( + tensorTileA.data().GetPhyAddr()); + } + } else { + copyGmToL1A(tensorL1A, tensorTileA); + } + AscendC::SetFlag(l1AEventList[l1AListIdNext]); + + // load next matrix B tile from GM to L1 + AscendC::WaitFlag(l1BEventList[l1BListIdNext]); + if constexpr (ENABLE_L1_RESIDENT) { + if (lastAddrB[l1BListIdNext] != tensorTileB.data().GetPhyAddr() || tla::get<0>(tensorTileB.coord()) != lastCoordB[l1BListIdNext].row() || tla::get<1>(tensorTileB.coord()) != lastCoordB[l1BListIdNext].column()) { + copyGmToL1B(tensorL1B, tensorTileB); + lastCoordB[l1BListIdNext] = + MatrixCoord{tla::get<0>(tensorTileB.coord()), tla::get<1>(tensorTileB.coord())}; + lastAddrB[l1BListIdNext] = + const_cast<__gm__ typename AscendC::GlobalTensor::PrimType *>( + tensorTileB.data().GetPhyAddr()); + } + } else { + copyGmToL1B(tensorL1B, tensorTileB); + } + AscendC::SetFlag(l1BEventList[l1BListIdNext]); + } + + // Get L1 tensor for current stage + auto l1ATensor = l1ATensorList[l1AListId]; + auto l1BTensor = l1BTensorList[l1BListId]; + tensorL1A = tla::MakeTensor(l1ATensor, L1A_LAYOUT, Arch::PositionL1{}); + tensorL1B = tla::MakeTensor(l1BTensor, L1B_LAYOUT, Arch::PositionL1{}); + // Get the loop nums on L0 + uint32_t kL0Loop = CeilDiv(kL1Actual); + + for (int mL0Idx = 0; mL0Idx < mL0Loop; mL0Idx++) { + uint32_t mL0Actual = (mL0Idx < mL0Loop - 1) ? L0_TILE_M : (mL1Actual - mL0Idx * L0_TILE_M); + + for (int kL0Idx = 0; kL0Idx < kL0Loop; kL0Idx++) { + uint32_t kL0Actual = (kL0Idx < kL0Loop - 1) ? L0_TILE_K : (kL1Actual - kL0Idx * L0_TILE_K); + + // Locate the current tile on L0A + auto l0ATile = l0ATensorList[l0AListId]; + auto layoutAInL0 = tla::MakeLayout(mL0Actual, kL0Actual); + auto tensorL0A = tla::MakeTensor(l0ATile, layoutAInL0, Arch::PositionL0A{}); + // Locate the current tile of matrix A on L1 + auto tensorTileL1A = GetTileA(tensorL1A, mL0Idx * L0_TILE_M, kL0Idx * L0_TILE_K, mL0Actual, kL0Actual); + + AscendC::WaitFlag(l0AEventList[l0AListId]); + if ((mL0Idx == 0) && (kL0Idx == 0)) { + AscendC::WaitFlag(l1AEventList[l1AListId]); + } + + // Load current tile from L1 to L0A + copyL1ToL0A(tensorL0A, tensorTileL1A); + + if ((mL0Idx == mL0Loop - 1) && (kL0Idx == kL0Loop - 1)) { + AscendC::SetFlag(l1AEventList[l1AListId]); + } + + bool initC = ((kL1Idx == 0) && (kL0Idx == 0)); + for (int nL0Idx = 0; nL0Idx < nL0Loop; nL0Idx++) { + uint32_t nL0Actual = (nL0Idx < nL0Loop - 1) ? L0_TILE_N : (nL1Actual - nL0Idx * L0_TILE_N); + + // Locate the current tile on L0B + auto l0BTile = l0BTensorList[l0BListId]; + auto layoutBInL0 = tla::MakeLayout(kL0Actual, nL0Actual); + auto tensorL0B = tla::MakeTensor(l0BTile, layoutBInL0, Arch::PositionL0B{}); + // Locate the current tile of matrix B on L1 + auto tensorTileL1B = GetTile(tensorL1B, + tla::MakeCoord(kL0Idx * L0_TILE_K, nL0Idx * L0_TILE_N), + tla::MakeShape(kL0Actual, nL0Actual)); + + // Wait for mmad finished + AscendC::WaitFlag(l0BEventList[l0BListId]); + // If the current tile is the first one on the k&n axis, wait for loading matrix B from GM to L1 + if ((mL0Idx == 0) && (kL0Idx == 0) && (nL0Idx == 0)) { + AscendC::WaitFlag(l1BEventList[l1BListId]); + } + + // Load current tile from L1 to L0B + copyL1ToL0B(tensorL0B, tensorTileL1B); + + // If the current tile is the last one on the k&n axis, notify to load matrix B from GM to L1 + if ((mL0Idx == mL0Loop - 1) && (kL0Idx == kL0Loop - 1) && (nL0Idx == nL0Loop - 1)) { + AscendC::SetFlag(l1BEventList[l1BListId]); + } + + if constexpr (HAS_BIAS && !std::is_same_v) { + if (initC) { + if (nL0Idx == 0) { + AscendC::WaitFlag(L1A_STAGES + L1B_STAGES); + } + AscendC::WaitFlag(L0A_STAGES + L0B_STAGES); + auto l1Bias = l1BiasTensor.template ReinterpretCast(); + auto tensorL1Bias = tla::MakeTensor(l1Bias, L1BIAS_LAYOUT, Arch::PositionL1{}); + auto tensorTileL1Bias = GetTile(tensorL1Bias, + tla::MakeCoord(nL0Idx * L0_TILE_N), + tla::MakeShape(nL0Actual)); + // Load bias to l0 biasTable + copyL1ToBT(tensorL0Bias, tensorTileL1Bias); + if (nL0Idx == nL0Loop - 1) { + AscendC::SetFlag(L1A_STAGES + L1B_STAGES); + } + } + } + + // Notify to do mmad + AscendC::SetFlag(l0CEventList[l0CListId]); + + // Locate the current tile on L0C + auto tensorTileL0C = GetTile(tensorL0C, + tla::MakeCoord(mL0Idx * L0_TILE_M, nL0Idx * L0_TILE_N), + tla::MakeShape(mL0Actual, nL0Actual)); + + // Compute the matrix multiplication on L0A and L0B and write the result to the accumulator + // Wait for loading L0B + AscendC::WaitFlag(l0CEventList[l0CListId]); + + // If the unit flag is enabled, the unit flag is set according to the calculation progress + uint8_t unitFlag = 0b00; + if constexpr (ENABLE_UNIT_FLAG) { + if ((kL1Idx == kL1Loop - 1) && (mL0Idx == mL0Loop - 1) && + (kL0Idx == kL0Loop - 1) && (nL0Idx == nL0Loop - 1)) { + unitFlag = 0b11; + } else { + unitFlag = 0b10; + } + } + + if constexpr (HAS_BIAS && !std::is_same_v) { + if (initC) { + tileMmad(tensorTileL0C, tensorL0A, tensorL0B, tensorL0Bias, + mL0Actual, nL0Actual, kL0Actual, initC, unitFlag); + AscendC::SetFlag(L0A_STAGES + L0B_STAGES); + } else { + tileMmad(tensorTileL0C, tensorL0A, tensorL0B, + mL0Actual, nL0Actual, kL0Actual, initC, unitFlag); + } + } else { + tileMmad(tensorTileL0C, tensorL0A, tensorL0B, + mL0Actual, nL0Actual, kL0Actual, initC, unitFlag); + } + + // Notify to move the next L0B tile + AscendC::SetFlag(l0BEventList[l0BListId]); + l0BListId = (l0BListId + 1 < L0B_STAGES) ? (l0BListId + 1) : 0; + } + AscendC::SetFlag(l0AEventList[l0AListId]); + l0AListId = (l0AListId + 1 < L0A_STAGES) ? (l0AListId + 1) : 0; + } + } + l1AListId = l1AListIdNext; + l1BListId = l1BListIdNext; + kL1Actual = kL1ActualNext; + } + + // copy block out + if constexpr (!ENABLE_UNIT_FLAG) { + AscendC::SetFlag(l0CEventList[l0CListId]); + AscendC::WaitFlag(l0CEventList[l0CListId]); + copyL0CToDst(tensorC, tensorL0C); + AscendC::SetFlag(l0CEventList[l0CListId]); + l0CListId = (l0CListId + 1 < L0C_STAGES) ? (l0CListId + 1) : 0; + } else { + copyL0CToDst(tensorC, tensorL0C, 0b11); + } + } + +protected: + template + CATLASS_DEVICE auto GetTileA(TensorA &tensorA, uint32_t mIndex, uint32_t kIndex, uint32_t mSize, uint32_t kSize) + { + if constexpr (tla::detail::isVector::value) { + return GetTile(tensorA, tla::MakeCoord(kIndex), tla::MakeShape(kSize)); + } else { + return GetTile(tensorA, tla::MakeCoord(mIndex, kIndex), tla::MakeShape(mSize, kSize)); + } + } + + // Multi-stage tensors list + AscendC::LocalTensor l1ATensorList[L1A_STAGES]; + AscendC::LocalTensor l1BTensorList[L1B_STAGES]; + AscendC::LocalTensor l0ATensorList[L0A_STAGES]; + AscendC::LocalTensor l0BTensorList[L0B_STAGES]; + AscendC::LocalTensor l0CTensorList[L0C_STAGES]; + AscendC::LocalTensor l1BiasTensor; + AscendC::LocalTensor l0BiasTensor; + + // Multi-stage event id list + int32_t l1AEventList[L1A_STAGES]; + int32_t l1BEventList[L1B_STAGES]; + int32_t l0AEventList[L0A_STAGES]; + int32_t l0BEventList[L0B_STAGES]; + int32_t l0CEventList[L0C_STAGES]; + + __gm__ typename AscendC::GlobalTensor::PrimType *lastAddrA[L1A_STAGES]; + __gm__ typename AscendC::GlobalTensor::PrimType *lastAddrB[L1B_STAGES]; + MatrixCoord lastCoordA[L1A_STAGES]; + MatrixCoord lastCoordB[L1B_STAGES]; + + // The id of current stage + uint32_t l1AListId{0}; + uint32_t l1BListId{0}; + uint32_t l0AListId{0}; + uint32_t l0BListId{0}; + uint32_t l0CListId{0}; + + TileMmad tileMmad; + CopyL1ToL0A copyL1ToL0A; + CopyL1ToL0B copyL1ToL0B; + CopyL1ToBT copyL1ToBT; +}; + +} // namespace Catlass::Gemm::Block + +#endif // CATLASS_GEMM_BLOCK_BLOCK_MMAD_PINGPONG_TLA_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_pingpong_tla_v2.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_pingpong_tla_v2.hpp new file mode 100644 index 0000000000..de6198e50e --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_pingpong_tla_v2.hpp @@ -0,0 +1,394 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_GEMM_BLOCK_BLOCK_MMAD_PINGPONG_TLA_V2_HPP +#define CATLASS_GEMM_BLOCK_BLOCK_MMAD_PINGPONG_TLA_V2_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/kda_coord.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" +#include "catlass/gemm/kda_gemm_helper.hpp" +#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" +#include "catlass/gemm/tile/kda_tile_mmad.hpp" +#include "tla/kda_tla_layout.hpp" +#include "tla/kda_tensor.hpp" + +namespace Catlass::Gemm::Block { + +template < + class ArchTag_, + bool ENABLE_UNIT_FLAG_, + class L1TileShape_, + class L0TileShape_, + class ElementA_, + class ElementB_, + class ElementC_, + class ElementBias_, + class TileCopy_, + class TileMmad_> +struct BlockMmadTla< + MmadPingpongTlaV2, + L1TileShape_, + L0TileShape_, + ElementA_, + ElementB_, + ElementC_, + ElementBias_, + TileCopy_, + TileMmad_> { +public: + // Type Aliases + using DispatchPolicy = MmadPingpongTlaV2; + using ArchTag = typename DispatchPolicy::ArchTag; + using L1TileShape = L1TileShape_; + using L0TileShape = L0TileShape_; + using ElementA = ElementA_; + using LayoutA = typename TileCopy_::LayoutA; + using ElementB = ElementB_; + using LayoutB = typename TileCopy_::LayoutB; + using ElementC = ElementC_; + using LayoutC = typename TileCopy_::LayoutC; + + using TileMmad = TileMmad_; + + using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; + + using ElementAccumulator = typename TileCopy_::ElementAccumulator; + + using LayoutTagL1A = typename TileCopy_::LayoutTagL1A; + using LayoutTagL1B = typename TileCopy_::LayoutTagL1B; + using LayoutTagL0A = typename TileCopy_::LayoutTagL0A; + using LayoutTagL0B = typename TileCopy_::LayoutTagL0B; + using LayoutTagL0C = typename TileCopy_::LayoutTagL0C; + + using L1AAlignHelper = typename TileCopy_::L1AAlignHelper; + using L1BAlignHelper = typename TileCopy_::L1BAlignHelper; + + static_assert(tla::is_tuple::value && tla::is_static::value, + "L1TileShape must be tla::tuple and static!"); + static_assert(tla::is_tuple::value && tla::is_static::value, + "L0TileShape must be tla::tuple and static!"); + + static constexpr bool ENABLE_UNIT_FLAG = DispatchPolicy::ENABLE_UNIT_FLAG; + static constexpr uint32_t STAGES = DispatchPolicy::STAGES; + static constexpr uint32_t L1_TILE_M = tla::get<0>(L1TileShape{}); + static constexpr uint32_t L1_TILE_N = tla::get<1>(L1TileShape{}); + static constexpr uint32_t L1_TILE_K = tla::get<2>(L1TileShape{}); + static constexpr uint32_t L0_TILE_M = tla::get<0>(L0TileShape{}); + static constexpr uint32_t L0_TILE_N = tla::get<1>(L0TileShape{}); + static constexpr uint32_t L0_TILE_K = tla::get<2>(L0TileShape{}); + + // L1 tile size + static constexpr uint32_t L1A_TILE_SIZE = L1_TILE_M * L1_TILE_K * sizeof(ElementA); + static constexpr uint32_t L1B_TILE_SIZE = L1_TILE_N * L1_TILE_K * sizeof(ElementB); + // L0 tile size + static constexpr uint32_t L0A_TILE_SIZE = L0_TILE_M * L0_TILE_K * sizeof(ElementA); + static constexpr uint32_t L0B_TILE_SIZE = L0_TILE_K * L0_TILE_N * sizeof(ElementB); + static constexpr uint32_t L0C_TILE_SIZE = L1_TILE_M * L1_TILE_N * sizeof(ElementAccumulator); + + // Check LayoutC + static_assert(tla::detail::isRowMajor::value || + ((std::is_same_v || std::is_same_v || + std::is_same_v) && + tla::detail::iszN::value), + "LayoutC only supports zN in half or bfloat16 or float, RowMajor in all dtype yet!"); + + // Check L1TileShape + static_assert((L1A_TILE_SIZE + L1B_TILE_SIZE) * STAGES <= ArchTag::L1_SIZE, + "L1TileShape exceeding the L1 space!"); + + // Check L0TileShape + static_assert(L0A_TILE_SIZE * STAGES <= ArchTag::L0A_SIZE, "L0TileShape exceeding the L0A space!"); + static_assert(L0B_TILE_SIZE * STAGES <= ArchTag::L0B_SIZE, "L0TileShape exceeding the L0B space!"); + static_assert(L0C_TILE_SIZE <= ArchTag::L0C_SIZE, "L0TileShape exceeding the L0C space!"); + + static constexpr uint32_t _32B = 32 * 8; // in bits + static_assert(L1_TILE_M == L0_TILE_M && L1_TILE_N == L0_TILE_N, + "The situation where the basic blocks of L1 and L0 differ on the m and n axes is not supported yet"); + static_assert(L0_TILE_K <= L1_TILE_K, "L0TileShape::K cannot exceed L1TileShape::K"); + static_assert(L1_TILE_M * SizeOfBits::value % _32B == 0, "L1TileShape::M must be 32B aligned."); + static_assert(L1_TILE_K * SizeOfBits::value % _32B == 0, "L1TileShape::K must be 32B aligned."); + static_assert(L1_TILE_K * SizeOfBits::value % _32B == 0, "L1TileShape::K must be 32B aligned."); + static_assert(L1_TILE_N * SizeOfBits::value % _32B == 0, "L1TileShape::N must be 32B aligned."); + static_assert(L0_TILE_K * SizeOfBits::value % _32B == 0, "L0TileShape::K must be 32B aligned."); + + static constexpr auto L1A_LAYOUT = tla::MakeLayout(tla::Int{}, tla::Int{}); + static constexpr auto L1B_LAYOUT = tla::MakeLayout(tla::Int{}, tla::Int{}); + + /// Construct + CATLASS_DEVICE + BlockMmadTla(Arch::Resource &resource, uint32_t l1BufAddrStart = 0) + { + uint32_t l1AOffset = l1BufAddrStart; + uint32_t l1BOffset = l1BufAddrStart + L1A_TILE_SIZE * STAGES; + // Init buffers + for (uint32_t i = 0; i < STAGES; i++) { + // Assign L1/L0A/L0B space for each stages + l1ATensorList[i] = resource.l1Buf.template GetBufferByByte(l1AOffset + L1A_TILE_SIZE * i); + l1BTensorList[i] = resource.l1Buf.template GetBufferByByte(l1BOffset + L1B_TILE_SIZE * i); + l0ATensorList[i] = resource.l0ABuf.template GetBufferByByte(L0A_TILE_SIZE * i); + l0BTensorList[i] = resource.l0BBuf.template GetBufferByByte(L0B_TILE_SIZE * i); + + // Assign event ID for each stages + l1AEventList[i] = i; + l1BEventList[i] = i + STAGES; + l0AEventList[i] = i; + l0BEventList[i] = i + STAGES; + + // The event id that needs to be set before the loop + AscendC::SetFlag(l1AEventList[i]); + AscendC::SetFlag(l1BEventList[i]); + AscendC::SetFlag(l0AEventList[i]); + AscendC::SetFlag(l0BEventList[i]); + } + l0CTensor = resource.l0CBuf.template GetBufferByByte(0); + AscendC::SetFlag(EVENT_ID0); + } + + /// Destructor + CATLASS_DEVICE + ~BlockMmadTla() + { + for (uint32_t i = 0; i < STAGES; i++) { + AscendC::WaitFlag(l1AEventList[i]); + AscendC::WaitFlag(l1BEventList[i]); + AscendC::WaitFlag(l0AEventList[i]); + AscendC::WaitFlag(l0BEventList[i]); + } + AscendC::WaitFlag(EVENT_ID0); + } + + /// Perform a block-scoped matrix multiply-accumulate + template + CATLASS_DEVICE void operator()(TensorA &tensorA, TensorB &tensorB, TensorC &tensorC) + { + using CopyGmToL1A = typename TileCopy_::template CopyGmToL1A; + using CopyGmToL1B = typename TileCopy_::template CopyGmToL1B; + CopyGmToL1A copyGmToL1A; + CopyGmToL1B copyGmToL1B; +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 2201) + using CopyL0CToGm = typename TileCopy_::template CopyL0CToGm; + CopyL0CToGm copyL0CToDst; +#endif +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 3510) + using CopyL0CToDst = typename TileCopy_::template CopyL0CToDst; + CopyL0CToDst copyL0CToDst; +#endif + + // Create an accumulator tensor view on L0C buffer. + // - Logical size comes from tensorC.layout().originShape() (tail-aware) + // - Layout is constructed from LayoutTagL0C + // - coord is initialized to (0, 0) (new buffer view) + auto tensorL0C = tla::MakeTensorLike(l0CTensor, tensorC, Arch::PositionL0C{}); + + // load first matrix A tile from GM to L1 + AscendC::WaitFlag(l1AEventList[l1ListId]); + // TileView: tileCoord is in tile units (not element units). + // It internally converts tileCoord to elementOffset = tileCoord ⊙ tileShape and handles tail tiles via originShape. + auto tensorTileA = tla::TileView(tensorA, + tla::MakeCoord(0u, 0u), // (m_tile, k_tile) + tla::MakeShape(tla::Int{}, tla::Int{})); + auto tensorL1A = tla::MakeTensorLike(l1ATensorList[l1ListId], tensorTileA, Arch::PositionL1{}, L1A_LAYOUT); + copyGmToL1A(tensorL1A, tensorTileA); + AscendC::SetFlag(l1AEventList[l1ListId]); + + // load first matrix B tile from GM to L1 + AscendC::WaitFlag(l1BEventList[l1ListId]); + auto tensorTileB = tla::TileView(tensorB, + tla::MakeCoord(0u, 0u), // (k_tile, n_tile) + tla::MakeShape(tla::Int{}, tla::Int{})); + auto tensorL1B = tla::MakeTensorLike(l1BTensorList[l1ListId], tensorTileB, Arch::PositionL1{}, L1B_LAYOUT); + copyGmToL1B(tensorL1B, tensorTileB); + AscendC::SetFlag(l1BEventList[l1ListId]); + + if constexpr (!ENABLE_UNIT_FLAG) { + AscendC::WaitFlag(EVENT_ID0); + } + + uint32_t mPartLoop = CeilDiv(tla::get<0>(tensorL0C.originShape())); // dim 0 = M + uint32_t nPartLoop = CeilDiv(tla::get<1>(tensorL0C.originShape())); // dim 1 = N + + // main loop + uint32_t kTileCount = CeilDiv(tla::get<1>(tensorA.originShape())); // dim 1 = K + for (uint32_t kLoopIdx = 0; kLoopIdx < kTileCount; kLoopIdx++) { + uint32_t l1ListIdNext = (l1ListId + 1 < STAGES) ? (l1ListId + 1) : 0; + // preload next tile from GM to L1 + if (kLoopIdx < kTileCount - 1) { + uint32_t kLoopIdxNext = kLoopIdx + 1; + + // Get L1 tensor for next stage + auto l1ATensor = l1ATensorList[l1ListIdNext]; + auto l1BTensor = l1BTensorList[l1ListIdNext]; + // Get GM tile for next stage + auto tensorTileA = tla::TileView(tensorA, + tla::MakeCoord(0u, kLoopIdxNext), // (m_tile, k_tile) + tla::MakeShape(tla::Int{}, tla::Int{})); + auto tensorTileB = tla::TileView(tensorB, + tla::MakeCoord(kLoopIdxNext, 0u), // (k_tile, n_tile) + tla::MakeShape(tla::Int{}, tla::Int{})); + auto tensorL1A = tla::MakeTensorLike(l1ATensor, tensorTileA, Arch::PositionL1{}, L1A_LAYOUT); + auto tensorL1B = tla::MakeTensorLike(l1BTensor, tensorTileB, Arch::PositionL1{}, L1B_LAYOUT); + + // load next matrix A tile from GM to L1 + AscendC::WaitFlag(l1AEventList[l1ListIdNext]); + copyGmToL1A(tensorL1A, tensorTileA); + AscendC::SetFlag(l1AEventList[l1ListIdNext]); + + // load next matrix B tile from GM to L1 + AscendC::WaitFlag(l1BEventList[l1ListIdNext]); + copyGmToL1B(tensorL1B, tensorTileB); + AscendC::SetFlag(l1BEventList[l1ListIdNext]); + } + + // Get L1 tensor for current stage + auto l1ATensor = l1ATensorList[l1ListId]; + auto l1BTensor = l1BTensorList[l1ListId]; + // Create tile view for current K iteration + auto tensorTileA = tla::TileView(tensorA, + tla::MakeCoord(0u, kLoopIdx), // (m_tile, k_tile) + tla::MakeShape(tla::Int{}, tla::Int{})); + auto tensorTileB = tla::TileView(tensorB, + tla::MakeCoord(kLoopIdx, 0u), // (k_tile, n_tile) + tla::MakeShape(tla::Int{}, tla::Int{})); + auto tensorL1A = tla::MakeTensorLike(l1ATensor, tensorTileA, Arch::PositionL1{}, L1A_LAYOUT); + auto tensorL1B = tla::MakeTensorLike(l1BTensor, tensorTileB, Arch::PositionL1{}, L1B_LAYOUT); + // Get the loop nums on L0 based on current L1 tile's actual K size + uint32_t kPartLoop = CeilDiv(tla::get<1>(tensorL1A.originShape())); // dim 1 = K + + for (int mPartIdx = 0; mPartIdx < mPartLoop; mPartIdx++) { + for (int kPartIdx = 0; kPartIdx < kPartLoop; kPartIdx++) { + // Locate the current tile on L0A + auto l0ATile = l0ATensorList[l0AListId]; + // Locate the current tile of matrix A on L1 + // Take a (L0_TILE_M, L0_TILE_K) tile from the current L1A tile (tile coordinates within L1 tile). + auto tensorTileL1A = tla::TileView( + tensorL1A, + tla::MakeCoord(mPartIdx, kPartIdx), + tla::MakeShape(tla::Int{}, tla::Int{})); + auto tensorL0A = tla::MakeTensorLike(l0ATile, tensorTileL1A, Arch::PositionL0A{}); + + AscendC::WaitFlag(l0AEventList[l0AListId]); + if ((mPartIdx == 0) && (kPartIdx == 0)) { + AscendC::WaitFlag(l1AEventList[l1ListId]); + } + + // Load current tile from L1 to L0A + copyL1ToL0A(tensorL0A, tensorTileL1A); + + if ((mPartIdx == mPartLoop - 1) && (kPartIdx == kPartLoop - 1)) { + AscendC::SetFlag(l1AEventList[l1ListId]); + } + + for (int nPartIdx = 0; nPartIdx < nPartLoop; nPartIdx++) { + // Locate the current tile on L0B + auto l0BTile = l0BTensorList[l0BListId]; + // Locate the current tile of matrix B on L1 + // Take a (L0_TILE_K, L0_TILE_N) tile from the current L1B tile (tile coordinates within L1 tile). + auto tensorTileL1B = tla::TileView( + tensorL1B, + tla::MakeCoord(kPartIdx, nPartIdx), + tla::MakeShape(tla::Int{}, tla::Int{})); + auto tensorL0B = tla::MakeTensorLike(l0BTile, tensorTileL1B, Arch::PositionL0B{}); + + // Wait for mmad finished + AscendC::WaitFlag(l0BEventList[l0BListId]); + // If the current tile is the first one on the k&n axis, wait for loading matrix B from GM to L1 + if ((kPartIdx == 0) && (nPartIdx == 0)) { + AscendC::WaitFlag(l1BEventList[l1ListId]); + } + + // Load current tile from L1 to L0B + copyL1ToL0B(tensorL0B, tensorTileL1B); + + // If the current tile is the last one on the k&n axis, notify to load matrix B from GM to L1 + if ((kPartIdx == kPartLoop - 1) && (nPartIdx == nPartLoop - 1)) { + AscendC::SetFlag(l1BEventList[l1ListId]); + } + // Notify to do mmad + AscendC::SetFlag(EVENT_ID0); + + // Locate the current tile on L0C + // View into L0C accumulator tile (tile coordinates in (m_part, n_part)). + auto tensorTileL0C = tla::TileView( + tensorL0C, + tla::MakeCoord(mPartIdx, nPartIdx), + tla::MakeShape(tla::Int{}, tla::Int{})); + + // Compute the matrix multiplication on L0A and L0B and write the result to the accumulator + // Wait for loading L0B + AscendC::WaitFlag(EVENT_ID0); + + // If the current tile is the first tile on the k axis, the accumulator needs to be reset to 0 + bool initC = ((kLoopIdx == 0) && (kPartIdx == 0)); + // If the unit flag is enabled, the unit flag is set according to the calculation progress + uint8_t unitFlag = 0b00; + if constexpr (ENABLE_UNIT_FLAG) { + if ((kLoopIdx == kTileCount - 1) && (mPartIdx == mPartLoop - 1) && + (kPartIdx == kPartLoop - 1) && (nPartIdx == nPartLoop - 1)) { + unitFlag = 0b11; + } else { + unitFlag = 0b10; + } + } + // Perform calculation operations + tileMmad(tensorTileL0C, tensorL0A, tensorL0B, initC, unitFlag); + + // Notify to move the next L0B tile + AscendC::SetFlag(l0BEventList[l0BListId]); + l0BListId = (l0BListId + 1 < STAGES) ? (l0BListId + 1) : 0; + } + AscendC::SetFlag(l0AEventList[l0AListId]); + l0AListId = (l0AListId + 1 < STAGES) ? (l0AListId + 1) : 0; + } + } + l1ListId = l1ListIdNext; + } + + // copy block out + if constexpr (!ENABLE_UNIT_FLAG) { + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + copyL0CToDst(tensorC, tensorL0C); + AscendC::SetFlag(EVENT_ID0); + } else { + copyL0CToDst(tensorC, tensorL0C, 0b11); + } + } + +protected: + // Multi-stage tensors list + AscendC::LocalTensor l1ATensorList[STAGES]; + AscendC::LocalTensor l1BTensorList[STAGES]; + AscendC::LocalTensor l0ATensorList[STAGES]; + AscendC::LocalTensor l0BTensorList[STAGES]; + AscendC::LocalTensor l0CTensor; + + // Multi-stage event id list + int32_t l1AEventList[STAGES]; + int32_t l1BEventList[STAGES]; + int32_t l0AEventList[STAGES]; + int32_t l0BEventList[STAGES]; + + // The id of current stage + uint32_t l1ListId{0}; + uint32_t l0AListId{0}; + uint32_t l0BListId{0}; + + TileMmad tileMmad; + CopyL1ToL0A copyL1ToL0A; + CopyL1ToL0B copyL1ToL0B; +}; + +} // namespace Catlass::Gemm::Block + +#endif // CATLASS_GEMM_BLOCK_BLOCK_MMAD_PINGPONG_TLA_V2_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_pingpong_with_prologue.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_pingpong_with_prologue.hpp new file mode 100644 index 0000000000..dc1c234b83 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_pingpong_with_prologue.hpp @@ -0,0 +1,535 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_GEMM_BLOCK_BLOCK_MMAD_PINGPONG_WITH_PROLOGUE +#define CATLASS_GEMM_BLOCK_BLOCK_MMAD_PINGPONG_WITH_PROLOGUE + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/kda_coord.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" +#include "catlass/gemm/kda_gemm_helper.hpp" +#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" +#include "catlass/gemm/tile/kda_tile_traits.hpp" + +namespace Catlass::Gemm::Block { + +template < + bool ENABLE_UNIT_FLAG_, + class L1TileShape_, + class L0TileShape_, + class AType_, + class BType_, + class CType_, + class BiasType_, + class TileCopy_, + class TileMmad_> +struct BlockMmad< + MmadAtlasA2PingPongWithPrologue, + L1TileShape_, + L0TileShape_, + AType_, + BType_, + CType_, + BiasType_, + TileCopy_, + TileMmad_> { +public: + // Type Aliases + using DispatchPolicy = MmadAtlasA2PingPongWithPrologue; + using ArchTag = typename DispatchPolicy::ArchTag; + using L1TileShape = L1TileShape_; + using L0TileShape = L0TileShape_; + using ElementA = typename AType_::Element; + using LayoutA = typename AType_::Layout; + using ElementB = typename BType_::Element; + using LayoutB = typename BType_::Layout; + using ElementC = typename CType_::Element; + using LayoutC = typename CType_::Layout; + using TileMmad = TileMmad_; + + using CopyGmToL1A = typename TileCopy_::CopyGmToL1A; + using CopyGmToL1B = typename TileCopy_::CopyGmToL1B; + using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; + + using PrologueA = typename TileCopy_::PrologueA; + using PrologueB = typename TileCopy_::PrologueB; + + using ElementAccumulator = + typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; + using CopyL0CToGm = typename TileCopy_::CopyL0CToGm; + using LayoutAInL1 = typename CopyL1ToL0A::LayoutSrc; + using LayoutBInL1 = typename CopyL1ToL0B::LayoutSrc; + using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst; + using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst; + using LayoutCInL0 = layout::zN; + + using L1AAlignHelper = Gemm::helper::L1AlignHelper; + using L1BAlignHelper = Gemm::helper::L1AlignHelper; + + static constexpr bool ENABLE_UNIT_FLAG = DispatchPolicy::ENABLE_UNIT_FLAG; + static constexpr uint32_t STAGES = DispatchPolicy::STAGES; + static constexpr uint32_t L1A_SIZE = L1TileShape::M * L1TileShape::K * sizeof(ElementA); + static constexpr uint32_t L1B_SIZE = L1TileShape::N * L1TileShape::K * sizeof(ElementB); + static constexpr uint32_t L0A_SIZE = ArchTag::L0A_SIZE; + static constexpr uint32_t L0B_SIZE = ArchTag::L0B_SIZE; + static constexpr uint32_t L0C_SIZE = ArchTag::L0C_SIZE; + static constexpr uint32_t L0A_PINGPONG_BUF_SIZE = L0A_SIZE / STAGES; + static constexpr uint32_t L0B_PINGPONG_BUF_SIZE = L0B_SIZE / STAGES; + + static constexpr bool HAS_PROLOGUE_A = !std::is_same_v; + static constexpr bool HAS_PROLOGUE_B = !std::is_same_v; + + // Check L1TileShape + static_assert((L1A_SIZE * STAGES + L1B_SIZE * STAGES) <= ArchTag::L1_SIZE, + "L1TileShape exceeding the L1 space!"); + + // Check L0TileShape + static constexpr uint32_t L0A_TILE_SIZE = L0TileShape::M * L0TileShape::K * sizeof(ElementA); + static constexpr uint32_t L0B_TILE_SIZE = L0TileShape::K * L0TileShape::N * sizeof(ElementB); + static constexpr uint32_t L0C_TILE_SIZE = L0TileShape::M * L0TileShape::N * sizeof(ElementAccumulator); + static_assert((L0A_TILE_SIZE * STAGES) <= L0A_SIZE, + "L0TileShape exceeding the L0A space!"); + static_assert((L0B_TILE_SIZE * STAGES) <= L0B_SIZE, + "L0TileShape exceeding the L0B space!"); + static_assert(L0C_TILE_SIZE <= L0C_SIZE, "L0TileShape exceeding the L0C space!"); + + static_assert(L1TileShape::M == L0TileShape::M && L1TileShape::N == L0TileShape::N, + "The situation where the basic blocks of L1 and L0 differ on the m and n axes is not supported yet"); + static_assert(L0TileShape::K <= L1TileShape::K, "L0TileShape::K cannot exceed L1TileShape::K"); + + // 32B (256b) aligned + static_assert(Gemm::helper::TileShapeAlignChecker::_ALIGN == 256, + "Tile shape must be 32B aligned."); + + struct Params { + typename Tile::PrologueTraits::Params prologueA{}; + typename Tile::PrologueTraits::Params prologueB{}; + typename CopyL0CToGm::Params copyL0CToGm{}; + }; + + /// Construct + CATLASS_DEVICE + BlockMmad(Arch::Resource const &resource, Params const ¶ms_ = {}) + : params(params_), + prologueA(resource, params_.prologueA), + prologueB(resource, params_.prologueB), + copyL0CToGm(params_.copyL0CToGm) + { + Arch::FlagID flagId = 0; + for (uint32_t i = 0; i < STAGES; ++i) { + if constexpr (HAS_PROLOGUE_A) { + flagCopyAFinish[i] = Arch::CrossCoreFlag(flagId++); + flagPrologueAFinish[i] = Arch::CrossCoreFlag(flagId++); + } + + if constexpr (HAS_PROLOGUE_B) { + flagCopyBFinish[i] = Arch::CrossCoreFlag(flagId++); + flagPrologueBFinish[i] = Arch::CrossCoreFlag(flagId++); + } + } + + if constexpr (g_coreType == AscendC::AIC) { + uint32_t l1AOffset = 0; + uint32_t l1BOffset = L1A_SIZE * STAGES; + // Init buffers + for (uint32_t i = 0; i < STAGES; i++) { + if constexpr (HAS_PROLOGUE_A) { + Catlass::Arch::CrossCoreSetFlag<0x02, PIPE_MTE2>(flagCopyAFinish[i]); + } + + if constexpr (HAS_PROLOGUE_B) { + Catlass::Arch::CrossCoreSetFlag<0x02, PIPE_MTE2>(flagCopyBFinish[i]); + } + + l1AEventList[i] = i; + l1BEventList[i] = i + STAGES; + l0AEventList[i] = i; + l0BEventList[i] = i + STAGES; + + AscendC::SetFlag(l1AEventList[i]); + AscendC::SetFlag(l1BEventList[i]); + AscendC::SetFlag(l0AEventList[i]); + AscendC::SetFlag(l0BEventList[i]); + + l1ATensorList[i] = resource.l1Buf.template GetBufferByByte(l1AOffset + L1A_SIZE * i); + l1BTensorList[i] = resource.l1Buf.template GetBufferByByte(l1BOffset + L1B_SIZE * i); + l0ATensorList[i] = resource.l0ABuf.template GetBufferByByte(L0A_PINGPONG_BUF_SIZE * i); + l0BTensorList[i] = resource.l0BBuf.template GetBufferByByte(L0B_PINGPONG_BUF_SIZE * i); + } + AscendC::SetFlag(EVENT_ID0); + l0CTensor = resource.l0CBuf.template GetBufferByByte(0); + } + } + + /// Destructor + CATLASS_DEVICE + ~BlockMmad() + { + if constexpr (g_coreType == AscendC::AIC) { + for (uint32_t i = 0; i < STAGES; i++) { + AscendC::WaitFlag(l1AEventList[i]); + AscendC::WaitFlag(l1BEventList[i]); + AscendC::WaitFlag(l0AEventList[i]); + AscendC::WaitFlag(l0BEventList[i]); + } + AscendC::WaitFlag(EVENT_ID0); + } else { + for (uint32_t i = 0; i < STAGES; i++) { + if constexpr (HAS_PROLOGUE_A) { + Catlass::Arch::CrossCoreWaitFlag(flagCopyAFinish[i]); + } + + if constexpr (HAS_PROLOGUE_B) { + Catlass::Arch::CrossCoreWaitFlag(flagCopyBFinish[i]); + } + } + } + } + + template + CATLASS_DEVICE + std::enable_if_t && !std::is_void_v, void> + Prologue( + AscendC::GlobalTensor const &gmSrcA, typename T::LayoutSrc const &layoutSrcA, + AscendC::GlobalTensor const &gmDstA, typename T::LayoutDst const &layoutDstA, + AscendC::GlobalTensor const &gmSrcB, typename U::LayoutSrc const &layoutSrcB, + AscendC::GlobalTensor const &gmDstB, typename U::LayoutDst const &layoutDstB, + GemmCoord const &actualBlockShape) + { + PrologueImpl( + gmSrcA, layoutSrcA, gmDstA, layoutDstA, + gmSrcB, layoutSrcB, gmDstB, layoutDstB, + actualBlockShape); + } + + template + CATLASS_DEVICE + std::enable_if_t && std::is_void_v, void> + Prologue( + AscendC::GlobalTensor const &gmSrcA, typename T::LayoutSrc const &layoutSrcA, + AscendC::GlobalTensor const &gmDstA, typename T::LayoutDst const &layoutDstA, + GemmCoord const &actualBlockShape) + { + PrologueImpl( + gmSrcA, layoutSrcA, gmDstA, layoutDstA, + {}, {}, {}, {}, + actualBlockShape); + } + + template + CATLASS_DEVICE + std::enable_if_t && !std::is_void_v, void> + Prologue( + AscendC::GlobalTensor const &gmSrcB, typename U::LayoutSrc const &layoutSrcB, + AscendC::GlobalTensor const &gmDstB, typename U::LayoutDst const &layoutDstB, + GemmCoord const &actualBlockShape) + { + PrologueImpl( + {}, {}, {}, {}, + gmSrcB, layoutSrcB, gmDstB, layoutDstB, + actualBlockShape); + } + + /// Perform a block-scoped matrix multiply-accumulate + CATLASS_DEVICE + void operator()( + AscendC::GlobalTensor const &gmA, LayoutA const &layoutA, + AscendC::GlobalTensor const &gmB, LayoutB const &layoutB, + AscendC::GlobalTensor const &gmC, LayoutC const &layoutC, + GemmCoord const &actualShape) + { + uint32_t mRound = RoundUp(actualShape.m()); + uint32_t nRound = RoundUp(actualShape.n()); + + auto layoutAInL1 = LayoutAInL1::template MakeLayout(L1TileShape::M, L1TileShape::K); + auto layoutBInL1 = LayoutBInL1::template MakeLayout(L1TileShape::K, L1TileShape::N); + auto layoutInL0C = LayoutCInL0::MakeLayoutInL0C(MakeCoord(mRound, nRound)); + + uint32_t kActual = min(actualShape.k(), L1TileShape::K); + + // load first matrix A tile from GM to L1 + AscendC::WaitFlag(l1AEventList[l1ListId]); + auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShape.m(), kActual)); + if constexpr (HAS_PROLOGUE_A) { + auto gmTileA = gmA[l1ListId * layoutA.Capacity()]; + Catlass::Arch::CrossCoreWaitFlag(flagPrologueAFinish[l1ListId]); + copyGmToL1A(l1ATensorList[l1ListId], gmTileA, layoutAInL1, layoutTileA); + Catlass::Arch::CrossCoreSetFlag<0x02, PIPE_MTE2>(flagCopyAFinish[l1ListId]); + } else { + copyGmToL1A(l1ATensorList[l1ListId], gmA, layoutAInL1, layoutTileA); + } + AscendC::SetFlag(l1AEventList[l1ListId]); + + // load first matrix B tile from GM to L1 + AscendC::WaitFlag(l1BEventList[l1ListId]); + auto layoutTileB = layoutB.GetTileLayout(MakeCoord(kActual, actualShape.n())); + if constexpr (HAS_PROLOGUE_B) { + auto gmTileB = gmB[l1ListId * layoutB.Capacity()]; + Catlass::Arch::CrossCoreWaitFlag(flagPrologueBFinish[l1ListId]); + copyGmToL1B(l1BTensorList[l1ListId], gmTileB, layoutBInL1, layoutTileB); + Catlass::Arch::CrossCoreSetFlag<0x02, PIPE_MTE2>(flagCopyBFinish[l1ListId]); + } else { + copyGmToL1B(l1BTensorList[l1ListId], gmB, layoutBInL1, layoutTileB); + } + AscendC::SetFlag(l1BEventList[l1ListId]); + + if constexpr (!ENABLE_UNIT_FLAG) { + AscendC::WaitFlag(EVENT_ID0); + } + + uint32_t mPartLoop = CeilDiv(mRound); + uint32_t nPartLoop = CeilDiv(nRound); + + // main loop + uint32_t kTileCount = CeilDiv(actualShape.k()); + for (uint32_t kLoopIdx = 0; kLoopIdx < kTileCount; kLoopIdx++) { + uint32_t l1ListIdNext = (l1ListId + 1 < STAGES) ? (l1ListId + 1) : 0; + uint32_t kActualNext{0}; + // preload next tile from GM to L1 + if (kLoopIdx < kTileCount - 1) { + uint32_t kLoopIdxNext = kLoopIdx + 1; + kActualNext = (kLoopIdxNext < kTileCount - 1) ? + L1TileShape::K : + (actualShape.k() - kLoopIdxNext * L1TileShape::K); + + // Get L1 tensor for next stage + auto l1ATensor = l1ATensorList[l1ListIdNext]; + auto l1BTensor = l1BTensorList[l1ListIdNext]; + + // load next matrix A tile from GM to L1 + AscendC::WaitFlag(l1AEventList[l1ListIdNext]); + layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShape.m(), kActualNext)); + if constexpr (HAS_PROLOGUE_A) { + auto gmTileA = gmA[l1ListIdNext * layoutA.Capacity()]; + Catlass::Arch::CrossCoreWaitFlag(flagPrologueAFinish[l1ListIdNext]); + copyGmToL1A(l1ATensor, gmTileA, layoutAInL1, layoutTileA); + Catlass::Arch::CrossCoreSetFlag<0x02, PIPE_MTE2>(flagCopyAFinish[l1ListIdNext]); + } else { + MatrixCoord gmTileAOffset{0, kLoopIdxNext * L1TileShape::K}; + auto gmTileA = gmA[layoutA.GetOffset(gmTileAOffset)]; + copyGmToL1A(l1ATensor, gmTileA, layoutAInL1, layoutTileA); + } + AscendC::SetFlag(l1AEventList[l1ListIdNext]); + + // load next matrix B tile from GM to L1 + AscendC::WaitFlag(l1BEventList[l1ListIdNext]); + layoutTileB = layoutB.GetTileLayout(MakeCoord(kActualNext, actualShape.n())); + if constexpr (HAS_PROLOGUE_B) { + auto gmTileB = gmB[l1ListIdNext * layoutB.Capacity()]; + Catlass::Arch::CrossCoreWaitFlag(flagPrologueBFinish[l1ListIdNext]); + copyGmToL1B(l1BTensor, gmTileB, layoutBInL1, layoutTileB); + Catlass::Arch::CrossCoreSetFlag<0x02, PIPE_MTE2>(flagCopyBFinish[l1ListIdNext]); + } else { + MatrixCoord gmTileBOffset{kLoopIdxNext * L1TileShape::K, 0}; + auto gmTileB = gmB[layoutB.GetOffset(gmTileBOffset)]; + copyGmToL1B(l1BTensor, gmTileB, layoutBInL1, layoutTileB); + } + AscendC::SetFlag(l1BEventList[l1ListIdNext]); + } + + // Get L1 tensor for current stage + auto l1ATensor = l1ATensorList[l1ListId]; + auto l1BTensor = l1BTensorList[l1ListId]; + + // Get the loop nums on L0 + uint32_t kPartLoop = CeilDiv(kActual); + + for (int mPartIdx = 0; mPartIdx < mPartLoop; mPartIdx++) { + uint32_t mPartActual = (mPartIdx < mPartLoop - 1) ? + L0TileShape::M : + (mRound - mPartIdx * L0TileShape::M); + + for (int kPartIdx = 0; kPartIdx < kPartLoop; kPartIdx++) { + uint32_t kPartActual = (kPartIdx < kPartLoop - 1) ? + L0TileShape::K : + (kActual - kPartIdx * L0TileShape::K); + + // Locate the current tile on L0A + auto l0ATile = l0ATensorList[l0AListId]; + LayoutAInL0 layoutAInL0 = LayoutAInL0::template MakeLayout(mPartActual, kPartActual); + // Locate the current tile of matrix A on L1 + MatrixCoord l1AOffset{mPartIdx * L0TileShape::M, kPartIdx * L0TileShape::K}; + auto l1ATile = l1ATensor[layoutAInL1.GetOffset(l1AOffset)]; + + AscendC::WaitFlag(l0AEventList[l0AListId]); + if ((mPartIdx == 0) && (kPartIdx == 0)) { + AscendC::WaitFlag(l1AEventList[l1ListId]); + } + + // Load current tile from L1 to L0A + copyL1ToL0A(l0ATile, l1ATile, layoutAInL0, layoutAInL1); + + if ((mPartIdx == mPartLoop - 1) && (kPartIdx == kPartLoop - 1)) { + AscendC::SetFlag(l1AEventList[l1ListId]); + } + + for (int nPartIdx = 0; nPartIdx < nPartLoop; nPartIdx++) { + uint32_t nPartActual = (nPartIdx < nPartLoop - 1) ? + L0TileShape::N : + (nRound - nPartIdx * L0TileShape::N); + + // Locate the current tile on L0B + auto l0BTile = l0BTensorList[l0BListId]; + LayoutBInL0 layoutBInL0 = LayoutBInL0::template MakeLayout(kPartActual, nPartActual); + // Locate the current tile of matrix B on L1 + MatrixCoord l1BOffset{kPartIdx * L0TileShape::K, nPartIdx * L0TileShape::N}; + auto l1BTile = l1BTensor[layoutBInL1.GetOffset(l1BOffset)]; + + // Wait for mmad finished + AscendC::WaitFlag(l0BEventList[l0BListId]); + // If the current tile is the first one on the k&n axis, wait for loading matrix B from GM to L1 + if ((kPartIdx == 0) && (nPartIdx == 0)) { + AscendC::WaitFlag(l1BEventList[l1ListId]); + } + + // Load current tile from L1 to L0B + copyL1ToL0B(l0BTile, l1BTile, layoutBInL0, layoutBInL1); + + // If the current tile is the last one on the k&n axis, notify to load matrix B from GM to L1 + if ((kPartIdx == kPartLoop - 1) && (nPartIdx == nPartLoop - 1)) { + AscendC::SetFlag(l1BEventList[l1ListId]); + } + // Notify to do mmad + AscendC::SetFlag(EVENT_ID0); + + // Locate the current tile on L0C + MatrixCoord l0COffset{mPartIdx * L0TileShape::M, nPartIdx * L0TileShape::N}; + auto l0CTile = l0CTensor[layoutInL0C.GetOffset(l0COffset)]; + + // Compute the matrix multiplication on L0A and L0B and write the result to the accumulator + // Wait for loading L0B + AscendC::WaitFlag(EVENT_ID0); + + // If the current tile is the first tile on the k axis, the accumulator needs to be reset to 0 + bool initC = ((kLoopIdx == 0) && (kPartIdx == 0)); + // If the unit flag is enabled, the unit flag is set according to the calculation progress + uint8_t unitFlag = 0b00; + if constexpr (ENABLE_UNIT_FLAG) { + if ((kLoopIdx == kTileCount - 1) && (mPartIdx == mPartLoop - 1) && + (kPartIdx == kPartLoop - 1) && (nPartIdx == nPartLoop - 1)) { + unitFlag = 0b11; + } else { + unitFlag = 0b10; + } + } + // Perform calculation operations + tileMmad(l0CTile, l0ATile, l0BTile, mPartActual, nPartActual, kPartActual, initC, unitFlag); + + // Notify to move the next L0B tile + AscendC::SetFlag(l0BEventList[l0BListId]); + l0BListId = (l0BListId + 1 < STAGES) ? (l0BListId + 1) : 0; + } + AscendC::SetFlag(l0AEventList[l0AListId]); + l0AListId = (l0AListId + 1 < STAGES) ? (l0AListId + 1) : 0; + } + } + l1ListId = l1ListIdNext; + kActual = kActualNext; + } + + // copy block out + LayoutC layoutBlock = layoutC.GetTileLayout(actualShape.GetCoordMN()); + + if constexpr (!ENABLE_UNIT_FLAG) { + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + copyL0CToGm(gmC, l0CTensor, layoutBlock, layoutInL0C); + AscendC::SetFlag(EVENT_ID0); + } else { + copyL0CToGm(gmC, l0CTensor, layoutBlock, layoutInL0C, 0b11); + } + } + +protected: + CATLASS_DEVICE + void PrologueImpl( + typename Tile::PrologueTraits::TensorSrc const &gmSrcA, + typename Tile::PrologueTraits::LayoutSrc const &layoutSrcA, + typename Tile::PrologueTraits::TensorDst const &gmDstA, + typename Tile::PrologueTraits::LayoutDst const &layoutDstA, + typename Tile::PrologueTraits::TensorSrc const &gmSrcB, + typename Tile::PrologueTraits::LayoutSrc const &layoutSrcB, + typename Tile::PrologueTraits::TensorDst const &gmDstB, + typename Tile::PrologueTraits::LayoutDst const &layoutDstB, + GemmCoord const &actualShape) + { + uint32_t kTileCount = CeilDiv(actualShape.k()); + for (uint32_t kLoopIdx = 0; kLoopIdx < kTileCount; kLoopIdx++) { + uint32_t kOffset = kLoopIdx * L1TileShape::K; + uint32_t kActual = (kLoopIdx == kTileCount - 1) ? (actualShape.k() - kOffset) : L1TileShape::K; + + if constexpr (HAS_PROLOGUE_A) { + MatrixCoord offsetCoordA{0, kOffset}; + MatrixCoord actualTileShapeA{actualShape.m(), kActual}; + auto gmTileSrcA = gmSrcA[layoutSrcA.GetOffset(offsetCoordA)]; + auto layoutTileSrcA = layoutSrcA.GetTileLayout(actualTileShapeA); + auto gmTileDstA = gmDstA[l1ListId * layoutDstA.Capacity()]; + auto layoutTileDstA = layoutDstA.GetTileLayout(actualTileShapeA); + Catlass::Arch::CrossCoreWaitFlag(flagCopyAFinish[l1ListId]); + prologueA(gmTileDstA, layoutTileDstA, gmTileSrcA, layoutTileSrcA); + Catlass::Arch::CrossCoreSetFlag<0x02, PIPE_MTE3>(flagPrologueAFinish[l1ListId]); + } + + if constexpr (HAS_PROLOGUE_B) { + MatrixCoord offsetCoordB{kOffset, 0}; + MatrixCoord actualTileShapeB{kActual, actualShape.n()}; + auto gmTileSrcB = gmSrcB[layoutSrcB.GetOffset(offsetCoordB)]; + auto layoutTileSrcB = layoutSrcB.GetTileLayout(actualTileShapeB); + auto gmTileDstB = gmDstB[l1ListId * layoutDstB.Capacity()]; + auto layoutTileDstB = layoutDstB.GetTileLayout(actualTileShapeB); + Catlass::Arch::CrossCoreWaitFlag(flagCopyBFinish[l1ListId]); + prologueB(gmTileDstB, layoutTileDstB, gmTileSrcB, layoutTileSrcB); + Catlass::Arch::CrossCoreSetFlag<0x02, PIPE_MTE3>(flagPrologueBFinish[l1ListId]); + } + + l1ListId = (l1ListId + 1 == STAGES) ? 0 : (l1ListId + 1); + } + } + + /// Data members + Params params; + + AscendC::LocalTensor l1ATensorList[STAGES]; + AscendC::LocalTensor l1BTensorList[STAGES]; + AscendC::LocalTensor l0ATensorList[STAGES]; + AscendC::LocalTensor l0BTensorList[STAGES]; + AscendC::LocalTensor l0CTensor; + + int32_t l1AEventList[STAGES]; + int32_t l1BEventList[STAGES]; + int32_t l0AEventList[STAGES]; + int32_t l0BEventList[STAGES]; + + Arch::CrossCoreFlag flagPrologueAFinish[STAGES]; + Arch::CrossCoreFlag flagCopyAFinish[STAGES]; + Arch::CrossCoreFlag flagPrologueBFinish[STAGES]; + Arch::CrossCoreFlag flagCopyBFinish[STAGES]; + + uint32_t l1ListId{0}; + uint32_t l0AListId{0}; + uint32_t l0BListId{0}; + + TileMmad tileMmad; + CopyGmToL1A copyGmToL1A; + CopyGmToL1B copyGmToL1B; + CopyL1ToL0A copyL1ToL0A; + CopyL1ToL0B copyL1ToL0B; + CopyL0CToGm copyL0CToGm; + + Tile::PrologueTraits prologueA; + Tile::PrologueTraits prologueB; +}; + +} // namespace Catlass::Gemm::Block + +#endif // CATLASS_GEMM_BLOCK_BLOCK_MMAD_PINGPONG_WITH_PROLOGUE diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_preload_async_with_callback_tla.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_preload_async_with_callback_tla.hpp new file mode 100644 index 0000000000..b17a26df5c --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_preload_async_with_callback_tla.hpp @@ -0,0 +1,650 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ +#ifndef CATLASS_GEMM_BLOCK_BLOCK_MMAD_PRELOAD_ASYNC_WITH_CALLBACK_TLA_HPP +#define CATLASS_GEMM_BLOCK_BLOCK_MMAD_PRELOAD_ASYNC_WITH_CALLBACK_TLA_HPP + +#include "catlass/arch/kda_resource.hpp" +#include "catlass/kda_catlass.hpp" +#include "catlass/kda_coord.hpp" +#include "catlass/detail/kda_callback.hpp" +#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" +#include "catlass/gemm/kda_gemm_helper.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "tla/kda_tla_layout.hpp" +#include "tla/kda_tensor.hpp" + +namespace Catlass::Gemm::Block { + +template < + class ArchTag_, + uint32_t PRELOAD_STAGES_, + uint32_t L1A_STAGES_, + uint32_t L1B_STAGES_, + uint32_t L0A_STAGES_, + uint32_t L0B_STAGES_, + uint32_t L0C_STAGES_, + bool ENABLE_UNIT_FLAG_, + bool ENABLE_SHUFFLE_K_, + bool ENABLE_L1_RESIDENT_, + bool USE_HF32_MODE_, + class L1TileShape_, + class L0TileShape_, + class ElementA_, + class ElementB_, + class ElementC_, + class ElementBias_, + class TileCopy_, + class TileMmad_> +struct BlockMmadTla< + MmadPreloadAsyncWithCallback< + ArchTag_, + PRELOAD_STAGES_, + L1A_STAGES_, + L1B_STAGES_, + L0A_STAGES_, + L0B_STAGES_, + L0C_STAGES_, + ENABLE_UNIT_FLAG_, + ENABLE_SHUFFLE_K_, + USE_HF32_MODE_, + ENABLE_L1_RESIDENT_>, + L1TileShape_, + L0TileShape_, + ElementA_, + ElementB_, + ElementC_, + ElementBias_, + TileCopy_, + TileMmad_> { +public: + // Type Aliases + using DispatchPolicy = MmadPreloadAsyncWithCallback< + ArchTag_, + PRELOAD_STAGES_, + L1A_STAGES_, + L1B_STAGES_, + L0A_STAGES_, + L0B_STAGES_, + L0C_STAGES_, + ENABLE_UNIT_FLAG_, + ENABLE_SHUFFLE_K_, + USE_HF32_MODE_, + ENABLE_L1_RESIDENT_>; + using ArchTag = typename DispatchPolicy::ArchTag; + using TileCopy = TileCopy_; + using L1TileShape = L1TileShape_; + using L0TileShape = L0TileShape_; + using ElementA = ElementA_; + using LayoutA = typename TileCopy::LayoutA; + using ElementB = ElementB_; + using LayoutB = typename TileCopy::LayoutB; + using ElementC = ElementC_; + using LayoutC = typename TileCopy::LayoutC; + using ElementBias = ElementBias_; + + using TileMmad = TileMmad_; + + using CopyL1ToL0A = typename TileCopy::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopy::CopyL1ToL0B; + using CopyL1ToBT = typename TileCopy::CopyL1ToBT; + + using ElementAccumulator = typename TileCopy::ElementAccumulator; + + static constexpr bool HAS_BIAS = TileCopy::HAS_BIAS; + + using LayoutTagL1A = typename TileCopy::LayoutTagL1A; + using LayoutTagL1B = typename TileCopy::LayoutTagL1B; + using LayoutTagL0A = typename TileCopy::LayoutTagL0A; + using LayoutTagL0B = typename TileCopy::LayoutTagL0B; + + static_assert( + tla::is_tuple::value && tla::is_static::value, + "L1TileShape must be tla::tuple and static!"); + static_assert( + tla::is_tuple::value && tla::is_static::value, + "L0TileShape must be tla::tuple and static!"); + + static constexpr uint32_t PRELOAD_STAGES = DispatchPolicy::PRELOAD_STAGES; + static constexpr uint32_t L1A_STAGES = DispatchPolicy::L1A_STAGES; + static constexpr uint32_t L1B_STAGES = DispatchPolicy::L1B_STAGES; + static constexpr uint32_t L0A_STAGES = DispatchPolicy::L0A_STAGES; + static constexpr uint32_t L0B_STAGES = DispatchPolicy::L0B_STAGES; + static constexpr uint32_t L0C_STAGES = DispatchPolicy::L0C_STAGES; + static constexpr uint32_t MIN_L1_STAGES = L1A_STAGES > L1B_STAGES ? L1B_STAGES : L1A_STAGES; + + static constexpr bool ENABLE_UNIT_FLAG = DispatchPolicy::ENABLE_UNIT_FLAG; + static constexpr bool ENABLE_SHUFFLE_K = DispatchPolicy::ENABLE_SHUFFLE_K; + static constexpr bool ENABLE_L1_RESIDENT = DispatchPolicy::ENABLE_L1_RESIDENT; + static constexpr bool USE_HF32_MODE = DispatchPolicy::USE_HF32_MODE; + + static constexpr uint32_t L1_TILE_M = tla::get<0>(L1TileShape{}); + static constexpr uint32_t L1_TILE_N = tla::get<1>(L1TileShape{}); + static constexpr uint32_t L1_TILE_K = tla::get<2>(L1TileShape{}); + static constexpr uint32_t L0_TILE_M = tla::get<0>(L0TileShape{}); + static constexpr uint32_t L0_TILE_N = tla::get<1>(L0TileShape{}); + static constexpr uint32_t L0_TILE_K = tla::get<2>(L0TileShape{}); + + // L1 tile size + static constexpr uint32_t L1A_TILE_SIZE = L1_TILE_M * L1_TILE_K * sizeof(ElementA); + static constexpr uint32_t L1B_TILE_SIZE = L1_TILE_N * L1_TILE_K * sizeof(ElementB); + // L0 tile size + static constexpr uint32_t L0A_TILE_SIZE = L0_TILE_M * L0_TILE_K * sizeof(ElementA); + static constexpr uint32_t L0B_TILE_SIZE = L0_TILE_K * L0_TILE_N * sizeof(ElementB); + static constexpr uint32_t L0C_TILE_SIZE = L1_TILE_M * L1_TILE_N * sizeof(ElementAccumulator); + + // Check HF32_MODE + static_assert( + !USE_HF32_MODE || (USE_HF32_MODE && std::is_same_v && std::is_same_v), + "HF32 MODE only supports in float!"); + + // Check LayoutC + static_assert( + tla::detail::isRowMajor::value || ((std::is_same_v || std::is_same_v || std::is_same_v) && tla::detail::iszN::value), + "LayoutC only supports zN in half or bfloat16 or float, RowMajor in all dtype yet!"); + + // Check L1TileShape + static_assert( + L1A_TILE_SIZE * L1A_STAGES + L1B_TILE_SIZE * L1B_STAGES <= ArchTag::L1_SIZE, + "L1TileShape exceeding the L1 space!"); + + // Check L0TileShape + static_assert(L0A_TILE_SIZE * L0A_STAGES <= ArchTag::L0A_SIZE, "L0TileShape exceeding the L0A space!"); + static_assert(L0B_TILE_SIZE * L0B_STAGES <= ArchTag::L0B_SIZE, "L0TileShape exceeding the L0B space!"); + static_assert(L0C_TILE_SIZE * L0C_STAGES <= ArchTag::L0C_SIZE, "L0TileShape exceeding the L0C space!"); + + static_assert( + L1_TILE_M == L0_TILE_M && L1_TILE_N == L0_TILE_N, + "The situation where the basic blocks of L1 and L0 differ on the m and n axes is not supported yet"); + static_assert(L0_TILE_K <= L1_TILE_K, "L0TileShape::K cannot exceed L1TileShape::K"); + static_assert(PRELOAD_STAGES <= MIN_L1_STAGES, + "PRELOAD_STAGES must not exceed the smaller of L1A_STAGES and L1B_STAGES!"); + + static_assert( + (!HAS_BIAS && (L1A_STAGES + L1B_STAGES) <= 8) || (HAS_BIAS && (L1A_STAGES + L1B_STAGES + MIN_L1_STAGES) <= 8), + "L1 Buffer overflow: Exceeds the supported range of EVENT(0~7)"); + static_assert( + (!HAS_BIAS && (L0A_STAGES + L0B_STAGES) <= 8) || (HAS_BIAS && (L0A_STAGES + L0B_STAGES) <= 7), + "L0 Buffer overflow: Exceeds the supported range of EVENT_ID(0~7)"); + + static constexpr auto L1A_LAYOUT = + tla::MakeLayout(tla::Int{}, tla::Int{}); + static constexpr auto L1B_LAYOUT = + tla::MakeLayout(tla::Int{}, tla::Int{}); + static constexpr auto L1BIAS_LAYOUT = tla::MakeLayout(tla::Int{}); + static constexpr auto L0BIAS_LAYOUT = tla::MakeLayout(tla::Int{}); + + CATLASS_DEVICE + BlockMmadTla(Arch::Resource &resource, uint32_t l1BufAddrStart = 0) + { + // use HF32 when USE_HF32_MODE is true + if constexpr (USE_HF32_MODE) { + AscendC::SetHF32Mode(true); + } else { + AscendC::SetHF32Mode(false); + } + if constexpr (ENABLE_UNIT_FLAG && tla::detail::isRowMajor::value) { + AscendC::SetMMLayoutTransform(true); + } + InitL1(resource, l1BufAddrStart); + InitL0A(resource); + InitL0B(resource); + InitL0C(resource); + if constexpr (HAS_BIAS) { + l0BiasTensor = resource.btBuf.template GetBufferByByte(0); + AscendC::SetFlag(L0A_STAGES + L0B_STAGES); + } + if constexpr (ENABLE_L1_RESIDENT) { + RestoreStatus(); + } + } + + CATLASS_DEVICE + ~BlockMmadTla() + { + if constexpr (USE_HF32_MODE) { + AscendC::SetHF32Mode(false); + } + if constexpr (ENABLE_UNIT_FLAG && tla::detail::isRowMajor::value) { + AscendC::SetMMLayoutTransform(false); + } + for (uint32_t i = 0; i < L1A_STAGES; ++i) { + AscendC::WaitFlag(l1AEventList[i]); + } + for (uint32_t i = 0; i < L1B_STAGES; ++i) { + AscendC::WaitFlag(l1BEventList[i]); + } + for (uint32_t i = 0; i < L0A_STAGES; ++i) { + AscendC::WaitFlag(l0AEventList[i]); + } + for (uint32_t i = 0; i < L0B_STAGES; ++i) { + AscendC::WaitFlag(l0BEventList[i]); + } + for (uint32_t i = 0; i < L0C_STAGES; ++i) { + AscendC::WaitFlag(l0CEventList[i]); + } + if constexpr (HAS_BIAS) { + for (uint32_t i = 0; i < MIN_L1_STAGES; ++i) { + AscendC::WaitFlag(l1BiasEventList[i]); + } + AscendC::WaitFlag(L0A_STAGES + L0B_STAGES); + } + } + + template + CATLASS_DEVICE void operator()( + TensorA &tensorA, + TensorB &tensorB, + TensorC &tensorC, + GemmCoord const &actualShape, + TensorBias const &tensorBias = {}, + Callback const &callbackBeforeFixpipe = {}, + Callback const &callbackAfterFixpipe = {}) + { + // Check L1TileShape + if constexpr (HAS_BIAS) { + static constexpr uint32_t BIAS_BUF_SIZE = L0_TILE_N * sizeof(ElementAccumulator); + static constexpr uint32_t L1BIAS_SIZE = L1_TILE_N * MIN_L1_STAGES * sizeof(ElementBias); + static_assert(BIAS_BUF_SIZE <= ArchTag::BIAS_SIZE, "BIAS_BUF_SIZE exceeding the BT space! Reduce L0_TILE_N"); + static_assert( + L1A_TILE_SIZE * L1A_STAGES + L1B_TILE_SIZE * L1B_STAGES + L1BIAS_SIZE <= ArchTag::L1_SIZE, + "L1TileShape exceeding the L1 space!"); + } + + using CopyGmToL1A = typename TileCopy::template CopyGmToL1A; + using CopyGmToL1B = typename TileCopy::template CopyGmToL1B; + CopyGmToL1A copyGmToL1A; + CopyGmToL1B copyGmToL1B; + + uint32_t mBlockActual = actualShape.m(); + uint32_t kBlockActual = actualShape.k(); + uint32_t nBlockActual = actualShape.n(); + + uint32_t kL1Loop = CeilDiv(kBlockActual); + + uint32_t mL1Actual = mBlockActual; + if constexpr (std::is_same_v) { + // Avoid using the gemv mode in mmad + if (mL1Actual == 1) { + mL1Actual = 16; + } + } + uint32_t nL1Actual = nBlockActual; + + uint32_t startTileIdx = 0; + if constexpr (ENABLE_SHUFFLE_K) { + startTileIdx = AscendC::GetBlockIdx() % kL1Loop; + } + + for (uint32_t kL1Idx = 0; kL1Idx < kL1Loop; ++kL1Idx) { + uint32_t kL1TileIdx = (startTileIdx + kL1Idx < kL1Loop) ? (startTileIdx + kL1Idx) : (startTileIdx + kL1Idx - kL1Loop); + + uint32_t kL1Actual = (kL1TileIdx < kL1Loop - 1) ? L1_TILE_K : (kBlockActual - kL1TileIdx * L1_TILE_K); + + // Load matrix A tile from GM to L1 + AscendC::WaitFlag(l1AEventList[l1AListId]); + auto tensorL1A = tla::MakeTensor(l1ATensorList[l1AListId], L1A_LAYOUT, Arch::PositionL1{}); + auto tensorTileA = + GetTile(tensorA, tla::MakeCoord(0, kL1TileIdx * L1_TILE_K), tla::MakeShape(mBlockActual, kL1Actual)); + if constexpr (ENABLE_L1_RESIDENT) { + // If the currently loaded GM pointer and block coordinates are the same as the last loaded ones, + // skip this loading. + if (lastAddrA[l1AListId] != tensorTileA.data().GetPhyAddr() || tla::get<0>(tensorTileA.coord()) != lastCoordA[l1AListId].row() || tla::get<1>(tensorTileA.coord()) != lastCoordA[l1AListId].column()) { + copyGmToL1A(tensorL1A, tensorTileA); + lastCoordA[l1AListId] = MatrixCoord{ + tla::get<0>(tensorTileA.coord()), tla::get<1>(tensorTileA.coord())}; + lastAddrA[l1AListId] = const_cast<__gm__ typename AscendC::GlobalTensor::PrimType *>( + tensorTileA.data().GetPhyAddr()); + } + } else { + copyGmToL1A(tensorL1A, tensorTileA); + } + AscendC::SetFlag(l1AEventList[l1AListId]); + // Load matrix B tile from GM to L1 + AscendC::WaitFlag(l1BEventList[l1BListId]); + auto tensorL1B = tla::MakeTensor(l1BTensorList[l1BListId], L1B_LAYOUT, Arch::PositionL1{}); + auto tensorTileB = + GetTile(tensorB, tla::MakeCoord(kL1TileIdx * L1_TILE_K, 0), tla::MakeShape(kL1Actual, nBlockActual)); + if constexpr (ENABLE_L1_RESIDENT) { + if (lastAddrB[l1BListId] != tensorTileB.data().GetPhyAddr() || tla::get<0>(tensorTileB.coord()) != lastCoordB[l1BListId].row() || tla::get<1>(tensorTileB.coord()) != lastCoordB[l1BListId].column()) { + copyGmToL1B(tensorL1B, tensorTileB); + lastCoordB[l1BListId] = MatrixCoord{ + tla::get<0>(tensorTileB.coord()), tla::get<1>(tensorTileB.coord())}; + lastAddrB[l1BListId] = const_cast<__gm__ typename AscendC::GlobalTensor::PrimType *>( + tensorTileB.data().GetPhyAddr()); + } + } else { + copyGmToL1B(tensorL1B, tensorTileB); + } + AscendC::SetFlag(l1BEventList[l1BListId]); + + if constexpr (HAS_BIAS && !std::is_same_v) { + if (kL1Idx == 0) { + using CopyGmToL1Bias = typename TileCopy::template CopyGmToL1Bias; + CopyGmToL1Bias copyGmToL1Bias; + AscendC::WaitFlag(l1BiasEventList[l1BiasListId]); + auto l1Bias = l1BiasTensor[l1BiasListId].template ReinterpretCast(); + auto tensorL1Bias = tla::MakeTensor(l1Bias, L1BIAS_LAYOUT, Arch::PositionL1{}); + copyGmToL1Bias(tensorL1Bias, tensorBias); + AscendC::SetFlag(l1BiasEventList[l1BiasListId]); + } + } + + // If the number of preload instructions reaches the upper limit, perform an mmad calculation on L1 tile + if (preloadCount == PRELOAD_STAGES) { + L1TileMmad(l1TileMmadParamsList[l1TileMmadParamsId]); + } + + // Store the current load status + uint32_t preloadL1TileMmadParamsId = (l1TileMmadParamsId + preloadCount < PRELOAD_STAGES) ? (l1TileMmadParamsId + preloadCount) : (l1TileMmadParamsId + preloadCount - PRELOAD_STAGES); + auto &l1TileMmadParams = l1TileMmadParamsList[preloadL1TileMmadParamsId]; + l1TileMmadParams.l1AListId = l1AListId; + l1TileMmadParams.l1BListId = l1BListId; + l1TileMmadParams.l1BiasListId = l1BiasListId; + l1TileMmadParams.mL1Actual = mL1Actual; + l1TileMmadParams.nL1Actual = nL1Actual; + l1TileMmadParams.kL1Actual = kL1Actual; + l1TileMmadParams.isKLoopFirst = (kL1Idx == 0); + l1TileMmadParams.isKLoopLast = (kL1Idx == kL1Loop - 1); + l1TileMmadParams.isNeedAddBias = !std::is_same_v; + if (kL1Idx == kL1Loop - 1) { + l1TileMmadParams.gmBlockC = tensorC.data(); + l1TileMmadParams.layoutCInGm = tensorC.layout(); + l1TileMmadParams.coord = tensorC.coord(); + l1TileMmadParams.callbackBeforeFixpipe = callbackBeforeFixpipe; + l1TileMmadParams.callbackAfterFixpipe = callbackAfterFixpipe; + } + + if (preloadCount < PRELOAD_STAGES) { + ++preloadCount; + } else { + l1TileMmadParamsId = (l1TileMmadParamsId + 1 < PRELOAD_STAGES) ? (l1TileMmadParamsId + 1) : 0; + } + l1AListId = (l1AListId + 1 < L1A_STAGES) ? (l1AListId + 1) : 0; + l1BListId = (l1BListId + 1 < L1B_STAGES) ? (l1BListId + 1) : 0; + if constexpr (HAS_BIAS && !std::is_same_v) { + if (kL1Idx == 0) { + l1BiasListId = (l1BiasListId + 1 < MIN_L1_STAGES) ? (l1BiasListId + 1) : 0; + } + } + } + } + + template + CATLASS_DEVICE void SynchronizeBlock() + { + while (preloadCount > 0) { + L1TileMmad(l1TileMmadParamsList[l1TileMmadParamsId]); + l1TileMmadParamsId = (l1TileMmadParamsId + 1 < PRELOAD_STAGES) ? (l1TileMmadParamsId + 1) : 0; + --preloadCount; + } + } + + // When enabling L1 resident mode, restore the pointer and coordinates that record the last state + // to the initial state. if two blockmmad instances need to be consecutively invoked at the kernel layer, + // RestoreStatus() must be inserted between them. + CATLASS_DEVICE + void RestoreStatus() + { + for (int i = 0; i < L1A_STAGES; ++i) { + lastAddrA[i] = nullptr; + lastCoordA[i] = MatrixCoord{0U, 0U}; + } + for (int i = 0; i < L1B_STAGES; ++i) { + lastAddrB[i] = nullptr; + lastCoordB[i] = MatrixCoord{0U, 0U}; + } + } + +private: + struct L1TileMmadParams { + uint32_t l1AListId; + uint32_t l1BListId; + uint32_t l1BiasListId; + uint32_t mL1Actual; + uint32_t nL1Actual; + uint32_t kL1Actual; + bool isKLoopFirst; + bool isKLoopLast; + bool isNeedAddBias; + AscendC::GlobalTensor gmBlockC; + LayoutC layoutCInGm; + tla::tuple coord; + Callback callbackBeforeFixpipe; + Callback callbackAfterFixpipe; + + CATLASS_DEVICE + L1TileMmadParams() = default; + }; + + CATLASS_DEVICE + void InitL1(Arch::Resource &resource, uint32_t l1BufAddrStart) + { + uint32_t l1AOffset = l1BufAddrStart; + uint32_t l1BOffset = l1BufAddrStart + L1A_TILE_SIZE * L1A_STAGES; + uint32_t l1BiasOffset = l1BOffset + L1B_TILE_SIZE * L1B_STAGES; + for (uint32_t i = 0; i < L1A_STAGES; ++i) { + l1ATensorList[i] = resource.l1Buf.template GetBufferByByte(l1AOffset + L1A_TILE_SIZE * i); + l1AEventList[i] = i; + AscendC::SetFlag(l1AEventList[i]); + } + for (uint32_t i = 0; i < L1B_STAGES; ++i) { + l1BTensorList[i] = resource.l1Buf.template GetBufferByByte(l1BOffset + L1B_TILE_SIZE * i); + l1BEventList[i] = i + L1A_STAGES; + AscendC::SetFlag(l1BEventList[i]); + } + for (uint32_t i = 0; i < MIN_L1_STAGES; ++i) { + if constexpr (HAS_BIAS) { + l1BiasTensor[i] = resource.l1Buf.template GetBufferByByte( + l1BiasOffset + L1_TILE_N * sizeof(ElementBias) * i); + l1BiasEventList[i] = i + L1A_STAGES + L1B_STAGES; + AscendC::SetFlag(l1BiasEventList[i]); + } + } + } + + CATLASS_DEVICE + void InitL0A(Arch::Resource &resource) + { + for (uint32_t i = 0; i < L0A_STAGES; ++i) { + l0ATensorList[i] = resource.l0ABuf.template GetBufferByByte(L0A_TILE_SIZE * i); + l0AEventList[i] = i; + AscendC::SetFlag(l0AEventList[i]); + } + } + + CATLASS_DEVICE + void InitL0B(Arch::Resource &resource) + { + for (uint32_t i = 0; i < L0B_STAGES; ++i) { + l0BTensorList[i] = resource.l0BBuf.template GetBufferByByte(L0B_TILE_SIZE * i); + l0BEventList[i] = i + L0A_STAGES; + AscendC::SetFlag(l0BEventList[i]); + } + } + + CATLASS_DEVICE + void InitL0C(Arch::Resource &resource) + { + for (uint32_t i = 0; i < L0C_STAGES; ++i) { + l0CTensorList[i] = resource.l0CBuf.template GetBufferByByte(L0C_TILE_SIZE * i); + l0CEventList[i] = i; + AscendC::SetFlag(l0CEventList[i]); + } + } + + template + CATLASS_DEVICE void L1TileMmad(L1TileMmadParams const ¶ms) + { + using CopyL0CToDst = typename TileCopy_::template CopyL0CToDst; + CopyL0CToDst copyL0CToDst; + + uint32_t kL0Loop = CeilDiv(params.kL1Actual); + auto &l1ATensor = l1ATensorList[params.l1AListId]; + auto &l1BTensor = l1BTensorList[params.l1BListId]; + auto tensorL1A = tla::MakeTensor(l1ATensor, L1A_LAYOUT, Arch::PositionL1{}); + auto tensorL1B = tla::MakeTensor(l1BTensor, L1B_LAYOUT, Arch::PositionL1{}); + + auto &l0CTensor = l0CTensorList[l0CListId]; + auto layoutInL0C = tla::MakeLayoutL0C(params.mL1Actual, params.nL1Actual); + auto tensorL0C = tla::MakeTensor(l0CTensor, layoutInL0C, Arch::PositionL0C{}); + auto tensorL0Bias = tla::MakeTensor(l0BiasTensor, L0BIAS_LAYOUT, Arch::PositionBias{}); + + if constexpr (!ENABLE_UNIT_FLAG) { + if (params.isKLoopFirst) { + AscendC::WaitFlag(l0CEventList[l0CListId]); + } + } + + for (uint32_t kL0Idx = 0; kL0Idx < kL0Loop; ++kL0Idx) { + uint32_t kL0Actual = (kL0Idx < kL0Loop - 1) ? L0_TILE_K : (params.kL1Actual - kL0Idx * L0_TILE_K); + + auto &l0ATile = l0ATensorList[l0AListId]; + auto layoutAInL0 = tla::MakeLayout(params.mL1Actual, kL0Actual); + auto tensorL0A = tla::MakeTensor(l0ATile, layoutAInL0, Arch::PositionL0A{}); + auto tensorTileL1A = + GetTile(tensorL1A, tla::MakeCoord(0, kL0Idx * L0_TILE_K), tla::MakeShape(params.mL1Actual, kL0Actual)); + + AscendC::WaitFlag(l0AEventList[l0AListId]); + if (kL0Idx == 0) { + AscendC::WaitFlag(l1AEventList[params.l1AListId]); + } + copyL1ToL0A(tensorL0A, tensorTileL1A); + if (kL0Idx == kL0Loop - 1) { + AscendC::SetFlag(l1AEventList[params.l1AListId]); + } + + auto &l0BTile = l0BTensorList[l0BListId]; + auto layoutBInL0 = tla::MakeLayout(kL0Actual, params.nL1Actual); + auto tensorL0B = tla::MakeTensor(l0BTile, layoutBInL0, Arch::PositionL0B{}); + // Locate the current tile of matrix B on L1 + auto tensorTileL1B = + GetTile(tensorL1B, tla::MakeCoord(kL0Idx * L0_TILE_K, 0), tla::MakeShape(kL0Actual, params.nL1Actual)); + + AscendC::WaitFlag(l0BEventList[l0BListId]); + if (kL0Idx == 0) { + AscendC::WaitFlag(l1BEventList[params.l1BListId]); + } + copyL1ToL0B(tensorL0B, tensorTileL1B); + if (kL0Idx == kL0Loop - 1) { + AscendC::SetFlag(l1BEventList[params.l1BListId]); + } + + bool initC = (params.isKLoopFirst && (kL0Idx == 0)); + if constexpr (HAS_BIAS) { + if (params.isNeedAddBias && initC) { + AscendC::WaitFlag(l1BiasEventList[params.l1BiasListId]); + AscendC::WaitFlag(L0A_STAGES + L0B_STAGES); + auto l1Bias = l1BiasTensor[params.l1BiasListId].template ReinterpretCast(); + auto tensorL1Bias = tla::MakeTensor(l1Bias, L1BIAS_LAYOUT, Arch::PositionL1{}); + auto tensorTileL1Bias = GetTile(tensorL1Bias, tla::MakeCoord(0), tla::MakeShape(params.nL1Actual)); + // Load bias to l0 biasTable + copyL1ToBT(tensorL0Bias, tensorTileL1Bias); + AscendC::SetFlag(l1BiasEventList[params.l1BiasListId]); + } + } + + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + + // If the unit flag is enabled, the unit flag is set according to the calculation progress + uint8_t unitFlag = 0b00; + if constexpr (ENABLE_UNIT_FLAG) { + if (params.isKLoopLast && (kL0Idx == kL0Loop - 1)) { + unitFlag = 0b11; + } else { + unitFlag = 0b10; + } + } + + if constexpr (HAS_BIAS) { + if (params.isNeedAddBias && initC) { + tileMmad( + tensorL0C, tensorL0A, tensorL0B, tensorL0Bias, params.mL1Actual, params.nL1Actual, kL0Actual, + initC, unitFlag); + AscendC::SetFlag(L0A_STAGES + L0B_STAGES); + } else { + tileMmad( + tensorL0C, tensorL0A, tensorL0B, params.mL1Actual, params.nL1Actual, kL0Actual, initC, unitFlag); + } + } else { + tileMmad( + tensorL0C, tensorL0A, tensorL0B, params.mL1Actual, params.nL1Actual, kL0Actual, initC, unitFlag); + } + + AscendC::SetFlag(l0BEventList[l0BListId]); + l0BListId = (l0BListId + 1 < L0B_STAGES) ? (l0BListId + 1) : 0; + + AscendC::SetFlag(l0AEventList[l0AListId]); + l0AListId = (l0AListId + 1 < L0A_STAGES) ? (l0AListId + 1) : 0; + } + + if (params.isKLoopLast) { + auto layoutCInGm = params.layoutCInGm; + auto tensorTileC = + tla::MakeTensor(params.gmBlockC, layoutCInGm, params.coord, Arch::PositionType{}); + + params.callbackBeforeFixpipe(); + + if constexpr (!ENABLE_UNIT_FLAG) { + AscendC::SetFlag(l0CEventList[l0CListId]); + AscendC::WaitFlag(l0CEventList[l0CListId]); + copyL0CToDst(tensorTileC, tensorL0C); + AscendC::SetFlag(l0CEventList[l0CListId]); + } else { + copyL0CToDst(tensorTileC, tensorL0C, 0b11); + } + l0CListId = (l0CListId + 1 < L0C_STAGES) ? (l0CListId + 1) : 0; + + params.callbackAfterFixpipe(); + } + } + + AscendC::LocalTensor l1ATensorList[L1A_STAGES]; + AscendC::LocalTensor l1BTensorList[L1B_STAGES]; + int32_t l1AEventList[L1A_STAGES]; + int32_t l1BEventList[L1B_STAGES]; + uint32_t l1AListId{0}; + uint32_t l1BListId{0}; + uint32_t l1BiasListId{0}; + + AscendC::LocalTensor l0ATensorList[L0A_STAGES]; + int32_t l0AEventList[L0A_STAGES]; + uint32_t l0AListId{0}; + + AscendC::LocalTensor l0BTensorList[L0B_STAGES]; + int32_t l0BEventList[L0B_STAGES]; + uint32_t l0BListId{0}; + + AscendC::LocalTensor l0CTensorList[L0C_STAGES_]; + int32_t l0CEventList[L0C_STAGES_]; + uint32_t l0CListId{0}; + + AscendC::LocalTensor l1BiasTensor[MIN_L1_STAGES]; + int32_t l1BiasEventList[MIN_L1_STAGES]; + AscendC::LocalTensor l0BiasTensor; + + L1TileMmadParams l1TileMmadParamsList[PRELOAD_STAGES]; + uint32_t l1TileMmadParamsId{0}; + uint32_t preloadCount{0}; + + __gm__ typename AscendC::GlobalTensor::PrimType *lastAddrA[L1A_STAGES]; + __gm__ typename AscendC::GlobalTensor::PrimType *lastAddrB[L1B_STAGES]; + MatrixCoord lastCoordA[L1A_STAGES]; + MatrixCoord lastCoordB[L1B_STAGES]; + + TileMmad tileMmad; + CopyL1ToL0A copyL1ToL0A; + CopyL1ToL0B copyL1ToL0B; + CopyL1ToBT copyL1ToBT; +}; + +} // namespace Catlass::Gemm::Block + +#endif // CATLASS_GEMM_BLOCK_BLOCK_MMAD_PRELOAD_ASYNC_WITH_CALLBACK_TLA_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_preload_tla.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_preload_tla.hpp new file mode 100644 index 0000000000..7bcdad82aa --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_preload_tla.hpp @@ -0,0 +1,445 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_GEMM_BLOCK_BLOCK_MMAD_PRELOAD_TLA_HPP +#define CATLASS_GEMM_BLOCK_BLOCK_MMAD_PRELOAD_TLA_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/kda_coord.hpp" +#include "catlass/detail/kda_callback.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" +#include "catlass/gemm/kda_gemm_helper.hpp" +#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" +#include "catlass/gemm/tile/kda_tile_mmad.hpp" +#include "tla/kda_tla_layout.hpp" +#include "tla/kda_tensor.hpp" + +namespace Catlass::Gemm::Block { + +template < + bool ENABLE_UNIT_FLAG_, + bool ENABLE_SHUFFLE_K_, + class L1TileShape_, + class L0TileShape_, + class TensorA_, + class TensorB_, + class TensorC_, + class TensorBias_, + class TileCopy_, + class TileMmad_> +struct BlockMmadTla< + MmadAtlasA2Preload, + L1TileShape_, + L0TileShape_, + TensorA_, + TensorB_, + TensorC_, + TensorBias_, + TileCopy_, + TileMmad_> { +public: + // Type Aliases + using DispatchPolicy = MmadAtlasA2Preload; + using ArchTag = typename DispatchPolicy::ArchTag; + using L1TileShape = L1TileShape_; + using L0TileShape = L0TileShape_; + using TensorA = TensorA_; + using TensorB = TensorB_; + using TensorC = TensorC_; + using ElementA = typename TensorA::Element; + using LayoutA = typename TensorA::Layout; + using ElementB = typename TensorB::Element; + using LayoutB = typename TensorB::Layout; + using ElementC = typename TensorC::Element; + using LayoutC = typename TensorC::Layout; + + using TileMmad = TileMmad_; + using CopyGmToL1A = typename TileCopy_::CopyGmToL1A; + using CopyGmToL1B = typename TileCopy_::CopyGmToL1B; + using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; + using CopyL0CToGm = typename TileCopy_::CopyL0CToGm; + using ElementAccumulator = typename CopyL0CToGm::ElementSrc; + + using LayoutTagL1A = typename TileCopy_::LayoutTagL1A; + using LayoutTagL1B = typename TileCopy_::LayoutTagL1B; + using LayoutTagL0A = typename TileCopy_::LayoutTagL0A; + using LayoutTagL0B = typename TileCopy_::LayoutTagL0B; + + using L1AAlignHelper = typename TileCopy_::L1AAlignHelper; + using L1BAlignHelper = typename TileCopy_::L1BAlignHelper; + + static_assert(tla::is_tuple::value && tla::is_static::value, + "L1TileShape must be tla::tuple and static!"); + static_assert(tla::is_tuple::value && tla::is_static::value, + "L0TileShape must be tla::tuple and static!"); + + static constexpr bool ENABLE_UNIT_FLAG = DispatchPolicy::ENABLE_UNIT_FLAG; + static constexpr bool ENABLE_SHUFFLE_K = DispatchPolicy::ENABLE_SHUFFLE_K; + static constexpr uint32_t STAGES = DispatchPolicy::STAGES; + static constexpr uint32_t L1_TILE_M = tla::get<0>(L1TileShape{}); + static constexpr uint32_t L1_TILE_N = tla::get<1>(L1TileShape{}); + static constexpr uint32_t L1_TILE_K = tla::get<2>(L1TileShape{}); + static constexpr uint32_t L0_TILE_M = tla::get<0>(L0TileShape{}); + static constexpr uint32_t L0_TILE_N = tla::get<1>(L0TileShape{}); + static constexpr uint32_t L0_TILE_K = tla::get<2>(L0TileShape{}); + + // L1 tile size + static constexpr uint32_t L1A_TILE_SIZE = L1_TILE_M * L1_TILE_K * sizeof(ElementA); + static constexpr uint32_t L1B_TILE_SIZE = L1_TILE_N * L1_TILE_K * sizeof(ElementB); + // L0 tile size + static constexpr uint32_t L0A_TILE_SIZE = L0_TILE_M * L0_TILE_K * sizeof(ElementA); + static constexpr uint32_t L0B_TILE_SIZE = L0_TILE_K * L0_TILE_N * sizeof(ElementB); + static constexpr uint32_t L0C_TILE_SIZE = L1_TILE_M * L1_TILE_N * sizeof(ElementAccumulator); + + // Check LayoutC + static_assert(tla::detail::isRowMajor::value, "LayoutC only support RowMajor yet!"); + + // Check L1TileShape + static_assert((L1A_TILE_SIZE + L1B_TILE_SIZE) * STAGES <= ArchTag::L1_SIZE, + "L1TileShape exceeding the L1 space!"); + + // Check L0TileShape + static_assert(L0A_TILE_SIZE * STAGES <= ArchTag::L0A_SIZE, "L0TileShape exceeding the L0A space!"); + static_assert(L0B_TILE_SIZE * STAGES <= ArchTag::L0B_SIZE, "L0TileShape exceeding the L0B space!"); + static_assert(L0C_TILE_SIZE <= ArchTag::L0C_SIZE, "L0TileShape exceeding the L0C space!"); + + static constexpr uint32_t _32B = 32 * 8; // in bits + static_assert(L1_TILE_M == L0_TILE_M && L1_TILE_N == L0_TILE_N, + "The situation where the basic blocks of L1 and L0 differ on the m and n axes is not supported yet"); + static_assert(L0_TILE_K <= L1_TILE_K, "L0TileShape::K cannot exceed L1TileShape::K"); + static_assert(L1_TILE_M * SizeOfBits::value % _32B == 0, "L1TileShape::M must be 32B aligned."); + static_assert(L1_TILE_K * SizeOfBits::value % _32B == 0, "L1TileShape::K must be 32B aligned."); + static_assert(L1_TILE_K * SizeOfBits::value % _32B == 0, "L1TileShape::K must be 32B aligned."); + static_assert(L1_TILE_N * SizeOfBits::value % _32B == 0, "L1TileShape::N must be 32B aligned."); + static_assert(L0_TILE_K * SizeOfBits::value % _32B == 0, "L0TileShape::K must be 32B aligned."); + + static constexpr auto L1A_LAYOUT = tla::MakeLayout(L1_TILE_M, L1_TILE_K); + static constexpr auto L1B_LAYOUT = tla::MakeLayout(L1_TILE_K, L1_TILE_N); + + /// Construct + CATLASS_DEVICE + BlockMmadTla(Arch::Resource &resource, uint32_t l1BufAddrStart = 0) + { + uint32_t l1AOffset = l1BufAddrStart; + uint32_t l1BOffset = l1BufAddrStart + L1A_TILE_SIZE * STAGES; + // Init buffers + for (uint32_t i = 0; i < STAGES; i++) { + // Assign L1/L0A/L0B space for each stages + l1ATensorList[i] = resource.l1Buf.template GetBufferByByte(l1AOffset + L1A_TILE_SIZE * i); + l1BTensorList[i] = resource.l1Buf.template GetBufferByByte(l1BOffset + L1B_TILE_SIZE * i); + l0ATensorList[i] = resource.l0ABuf.template GetBufferByByte(L0A_TILE_SIZE * i); + l0BTensorList[i] = resource.l0BBuf.template GetBufferByByte(L0B_TILE_SIZE * i); + + // Assign event ID for each stages + l1AEventList[i] = i; + l1BEventList[i] = i + STAGES; + l0AEventList[i] = i; + l0BEventList[i] = i + STAGES; + + // The event id that needs to be set before the loop + AscendC::SetFlag(l1AEventList[i]); + AscendC::SetFlag(l1BEventList[i]); + AscendC::SetFlag(l0AEventList[i]); + AscendC::SetFlag(l0BEventList[i]); + } + l0CTensor = resource.l0CBuf.template GetBufferByByte(0); + AscendC::SetFlag(EVENT_ID0); + } + + /// Destructor + CATLASS_DEVICE + ~BlockMmadTla() + { + for (uint32_t i = 0; i < STAGES; i++) { + AscendC::WaitFlag(l1AEventList[i]); + AscendC::WaitFlag(l1BEventList[i]); + AscendC::WaitFlag(l0AEventList[i]); + AscendC::WaitFlag(l0BEventList[i]); + } + AscendC::WaitFlag(EVENT_ID0); + } + + /// Perform a block-scoped matrix multiply-accumulate + template + CATLASS_DEVICE void operator()( + TensorA &tensorA, TensorB &tensorB, TensorC &tensorC, TensorA &tensorNextA, TensorB &tensorNextB, + GemmCoord const &actualShape, GemmCoord const &actualShapeNext, bool isFirstBlock, bool hasNextBlock, + Callback const &callbackBeforeFixpipe = {}, Callback const &callbackAfterFixpipe = {}) + { + uint32_t mBlockActual = actualShape.m(); + uint32_t kBlockActual = actualShape.k(); + uint32_t nBlockActual = actualShape.n(); + uint32_t mNextBlockActual = actualShapeNext.m(); + uint32_t kNextBlockActual = actualShapeNext.k(); + uint32_t nNextBlockActual = actualShapeNext.n(); + + uint32_t mRound = RoundUp(mBlockActual); + uint32_t nRound = RoundUp(nBlockActual); + + auto layoutInL0C = tla::MakeLayoutL0C(mRound, nRound); + auto tensorL0C = tla::MakeTensor(l0CTensor, layoutInL0C, Arch::PositionL0C{}); + + if constexpr (!ENABLE_UNIT_FLAG) { + AscendC::WaitFlag(EVENT_ID0); + } + uint32_t startTileIdx = 0; + if constexpr (ENABLE_SHUFFLE_K) { + startTileIdx = AscendC::GetBlockIdx(); + } + + uint32_t kTileCount = CeilDiv(kBlockActual); + uint32_t firstTileIdx = startTileIdx % kTileCount; + uint32_t lastTileIdx = (startTileIdx + kTileCount - 1) % kTileCount; + uint32_t kActual = + (firstTileIdx < kTileCount - 1) ? L1_TILE_K : (kBlockActual - firstTileIdx * L1_TILE_K); + uint32_t kTileCountNext = CeilDiv(kNextBlockActual); + uint32_t firstTileIdxNext = startTileIdx % kTileCountNext; + + if (isFirstBlock) { + // load first matrix A tile from GM to L1 + AscendC::WaitFlag(l1AEventList[l1ListId]); + auto tensorL1A = tla::MakeTensor(l1ATensorList[l1ListId], L1A_LAYOUT, Arch::PositionL1{}); + auto tensorTileA = GetTile( + tensorA, + tla::MakeCoord(0, firstTileIdx * L1_TILE_K), + tla::MakeShape(mBlockActual, kActual)); + copyGmToL1A(tensorL1A, tensorTileA, tla::MakeShape(mBlockActual, kActual)); + AscendC::SetFlag(l1AEventList[l1ListId]); + + // load first matrix B tile from GM to L1 + AscendC::WaitFlag(l1BEventList[l1ListId]); + auto tensorL1B = tla::MakeTensor(l1BTensorList[l1ListId], L1B_LAYOUT, Arch::PositionL1{}); + auto tensorTileB = GetTile( + tensorB, + tla::MakeCoord(firstTileIdx * L1_TILE_K, 0), + tla::MakeShape(kActual, nBlockActual)); + copyGmToL1B(tensorL1B, tensorTileB, tla::MakeShape(kActual, nBlockActual)); + AscendC::SetFlag(l1BEventList[l1ListId]); + } + + uint32_t mPartLoop = CeilDiv(mRound); + uint32_t nPartLoop = CeilDiv(nRound); + + // main loop + for (uint32_t kLoopIdx = 0; kLoopIdx < kTileCount; kLoopIdx++) { + uint32_t shuffleKIdx = (startTileIdx + kLoopIdx) % kTileCount; + uint32_t l1ListIdNext = (l1ListId + 1 < STAGES) ? (l1ListId + 1) : 0; + uint32_t kActualNext{0}; + + if (shuffleKIdx != lastTileIdx) { + uint32_t shuffleKIdxNext = (startTileIdx + kLoopIdx + 1) % kTileCount; + kActualNext = (shuffleKIdxNext < kTileCount - 1) ? + L1_TILE_K : + (kBlockActual - shuffleKIdxNext * L1_TILE_K); + + auto tensorL1A = tla::MakeTensor(l1ATensorList[l1ListIdNext], L1A_LAYOUT, Arch::PositionL1{}); + auto tensorL1B = tla::MakeTensor(l1BTensorList[l1ListIdNext], L1B_LAYOUT, Arch::PositionL1{}); + auto tensorTileA = GetTile( + tensorA, + tla::MakeCoord(0, shuffleKIdxNext * L1_TILE_K), + tla::MakeShape(mBlockActual, kActualNext)); + auto tensorTileB = GetTile( + tensorB, + tla::MakeCoord(shuffleKIdxNext * L1_TILE_K, 0), + tla::MakeShape(kActualNext, nBlockActual)); + + // load next matrix A tile from GM to L1 + AscendC::WaitFlag(l1AEventList[l1ListIdNext]); + copyGmToL1A(tensorL1A, tensorTileA, tla::MakeShape(mBlockActual, kActualNext)); + AscendC::SetFlag(l1AEventList[l1ListIdNext]); + + // load next matrix B tile from GM to L1 + AscendC::WaitFlag(l1BEventList[l1ListIdNext]); + copyGmToL1B(tensorL1B, tensorTileB, tla::MakeShape(kActualNext, nBlockActual)); + AscendC::SetFlag(l1BEventList[l1ListIdNext]); + } + + // preload next tile from GM to L1 + if (shuffleKIdx == lastTileIdx && hasNextBlock) { + kActualNext = (firstTileIdxNext < kTileCountNext - 1) ? + L1_TILE_K : + (kNextBlockActual - firstTileIdxNext * L1_TILE_K); + + // Get L1 tensor for next stage + auto tensorL1A = tla::MakeTensor(l1ATensorList[l1ListIdNext], L1A_LAYOUT, Arch::PositionL1{}); + auto tensorL1B = tla::MakeTensor(l1BTensorList[l1ListIdNext], L1B_LAYOUT, Arch::PositionL1{}); + // Get GM tile for next stage + auto tensorTileA = GetTile( + tensorNextA, + tla::MakeCoord(0, firstTileIdxNext * L1_TILE_K), + tla::MakeShape(mNextBlockActual, kActualNext)); + auto tensorTileB = GetTile( + tensorNextB, + tla::MakeCoord(firstTileIdxNext * L1_TILE_K, 0), + tla::MakeShape(kActualNext, nNextBlockActual)); + + // load next matrix A tile from GM to L1 + AscendC::WaitFlag(l1AEventList[l1ListIdNext]); + copyGmToL1A(tensorL1A, tensorTileA, tla::MakeShape(mNextBlockActual, kActualNext)); + AscendC::SetFlag(l1AEventList[l1ListIdNext]); + + // load next matrix B tile from GM to L1 + AscendC::WaitFlag(l1BEventList[l1ListIdNext]); + copyGmToL1B(tensorL1B, tensorTileB, tla::MakeShape(kActualNext, nNextBlockActual)); + AscendC::SetFlag(l1BEventList[l1ListIdNext]); + } + + // Get L1 tensor for current stage + auto l1ATensor = l1ATensorList[l1ListId]; + auto l1BTensor = l1BTensorList[l1ListId]; + auto tensorL1A = tla::MakeTensor(l1ATensor, L1A_LAYOUT, Arch::PositionL1{}); + auto tensorL1B = tla::MakeTensor(l1BTensor, L1B_LAYOUT, Arch::PositionL1{}); + // Get the loop nums on L0 + uint32_t kPartLoop = CeilDiv(kActual); + + for (int mPartIdx = 0; mPartIdx < mPartLoop; mPartIdx++) { + uint32_t mPartActual = (mPartIdx < mPartLoop - 1) ? + L0_TILE_M : + (mRound - mPartIdx * L0_TILE_M); + + for (int kPartIdx = 0; kPartIdx < kPartLoop; kPartIdx++) { + uint32_t kPartActual = (kPartIdx < kPartLoop - 1) ? + L0_TILE_K : + (kActual - kPartIdx * L0_TILE_K); + + // Locate the current tile on L0A + auto l0ATile = l0ATensorList[l0AListId]; + auto layoutAInL0 = tla::MakeLayout(mPartActual, kPartActual); + auto tensorL0A = tla::MakeTensor(l0ATile, layoutAInL0, Arch::PositionL0A{}); + // Locate the current tile of matrix A on L1 + auto tensorTileL1A = GetTile(tensorL1A, tla::MakeCoord(mPartIdx * L0_TILE_M, kPartIdx * L0_TILE_K), + tla::MakeShape(mPartActual, kPartActual)); + + AscendC::WaitFlag(l0AEventList[l0AListId]); + if ((mPartIdx == 0) && (kPartIdx == 0)) { + AscendC::WaitFlag(l1AEventList[l1ListId]); + } + + // Load current tile from L1 to L0A + copyL1ToL0A(tensorL0A, tensorTileL1A); + + if ((mPartIdx == mPartLoop - 1) && (kPartIdx == kPartLoop - 1)) { + AscendC::SetFlag(l1AEventList[l1ListId]); + } + + for (int nPartIdx = 0; nPartIdx < nPartLoop; nPartIdx++) { + uint32_t nPartActual = (nPartIdx < nPartLoop - 1) ? + L0_TILE_N : + (nRound - nPartIdx * L0_TILE_N); + + // Locate the current tile on L0B + auto l0BTile = l0BTensorList[l0BListId]; + auto layoutBInL0 = tla::MakeLayout(kPartActual, nPartActual); + auto tensorL0B = tla::MakeTensor(l0BTile, layoutBInL0, Arch::PositionL0B{}); + // Locate the current tile of matrix B on L1 + auto tensorTileL1B = GetTile(tensorL1B, + tla::MakeCoord(kPartIdx * L0_TILE_K, nPartIdx * L0_TILE_N), + tla::MakeShape(kPartActual, nPartActual)); + + // Wait for mmad finished + AscendC::WaitFlag(l0BEventList[l0BListId]); + // If the current tile is the first one on the k&n axis, wait for loading matrix B from GM to L1 + if ((kPartIdx == 0) && (nPartIdx == 0)) { + AscendC::WaitFlag(l1BEventList[l1ListId]); + } + + // Load current tile from L1 to L0B + copyL1ToL0B(tensorL0B, tensorTileL1B); + + // If the current tile is the last one on the k&n axis, notify to load matrix B from GM to L1 + if ((kPartIdx == kPartLoop - 1) && (nPartIdx == nPartLoop - 1)) { + AscendC::SetFlag(l1BEventList[l1ListId]); + } + // Notify to do mmad + AscendC::SetFlag(EVENT_ID0); + + // Locate the current tile on L0C + auto tensorTileL0C = GetTile(tensorL0C, + tla::MakeCoord(mPartIdx * L0_TILE_M, nPartIdx * L0_TILE_N), + tla::MakeShape(mPartActual, nPartActual)); + + // Compute the matrix multiplication on L0A and L0B and write the result to the accumulator + // Wait for loading L0B + AscendC::WaitFlag(EVENT_ID0); + + // If the current tile is the first tile on the k axis, the accumulator needs to be reset to 0 + bool initC = ((kLoopIdx == 0) && (kPartIdx == 0)); + // If the unit flag is enabled, the unit flag is set according to the calculation progress + uint8_t unitFlag = 0b00; + if constexpr (ENABLE_UNIT_FLAG) { + if ((kLoopIdx == kTileCount - 1) && (mPartIdx == mPartLoop - 1) && + (kPartIdx == kPartLoop - 1) && (nPartIdx == nPartLoop - 1)) { + unitFlag = 0b11; + } else { + unitFlag = 0b10; + } + } + // Perform calculation operations + tileMmad(tensorTileL0C, tensorL0A, tensorL0B, + mPartActual, nPartActual, kPartActual, initC, unitFlag); + + // Notify to move the next L0B tile + AscendC::SetFlag(l0BEventList[l0BListId]); + l0BListId = (l0BListId + 1 < STAGES) ? (l0BListId + 1) : 0; + } + AscendC::SetFlag(l0AEventList[l0AListId]); + l0AListId = (l0AListId + 1 < STAGES) ? (l0AListId + 1) : 0; + } + } + l1ListId = l1ListIdNext; + kActual = kActualNext; + } + + callbackBeforeFixpipe(); + // copy block out + if constexpr (!ENABLE_UNIT_FLAG) { + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + copyL0CToGm(tensorC, tensorL0C); + AscendC::SetFlag(EVENT_ID0); + } else { + copyL0CToGm(tensorC, tensorL0C, 0b11); + } + callbackAfterFixpipe(); + } + +protected: + // Multi-stage tensors list + AscendC::LocalTensor l1ATensorList[STAGES]; + AscendC::LocalTensor l1BTensorList[STAGES]; + AscendC::LocalTensor l0ATensorList[STAGES]; + AscendC::LocalTensor l0BTensorList[STAGES]; + AscendC::LocalTensor l0CTensor; + + // Multi-stage event id list + int32_t l1AEventList[STAGES]; + int32_t l1BEventList[STAGES]; + int32_t l0AEventList[STAGES]; + int32_t l0BEventList[STAGES]; + + // The id of current stage + uint32_t l1ListId{0}; + uint32_t l0AListId{0}; + uint32_t l0BListId{0}; + + TileMmad tileMmad; + CopyGmToL1A copyGmToL1A; + CopyGmToL1B copyGmToL1B; + CopyL1ToL0A copyL1ToL0A; + CopyL1ToL0B copyL1ToL0B; + CopyL0CToGm copyL0CToGm; +}; + +} // namespace Catlass::Gemm::Block + +#endif // CATLASS_GEMM_BLOCK_BLOCK_MMAD_PRELOAD_TLA_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_single_core_splitk.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_single_core_splitk.hpp new file mode 100644 index 0000000000..9101ffee2d --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_single_core_splitk.hpp @@ -0,0 +1,391 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_GEMM_BLOCK_BLOCK_MMAD_SINGLE_CORE_SPLITK_HPP +#define CATLASS_GEMM_BLOCK_BLOCK_MMAD_SINGLE_CORE_SPLITK_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/kda_coord.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" +#include "catlass/gemm/kda_gemm_helper.hpp" +#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" +#include "catlass/gemm/tile/kda_tile_mmad.hpp" + +namespace Catlass::Gemm::Block { + +template < + uint32_t L1A_STAGES_, + uint32_t L1B_STAGES_, + uint32_t L0C_STAGES_, + bool ENABLE_UNIT_FLAG_, + class L1TileShape_, + class L0TileShape_, + class AType_, + class BType_, + class CType_, + class BiasType_, + class TileCopy_, + class TileMmad_> +struct BlockMmad< + MmadAtlasA2SingleCoreSplitk, + L1TileShape_, + L0TileShape_, + AType_, + BType_, + CType_, + BiasType_, + TileCopy_, + TileMmad_> { +public: + // Type Aliases + using DispatchPolicy = MmadAtlasA2SingleCoreSplitk; + using ArchTag = typename DispatchPolicy::ArchTag; + using ElementA = typename AType_::Element; + using L1TileShape = L1TileShape_; + using L0TileShape = L0TileShape_; + using LayoutA = typename AType_::Layout; + using ElementB = typename BType_::Element; + using LayoutB = typename BType_::Layout; + using ElementC = typename CType_::Element; + using LayoutC = typename CType_::Layout; + using TileMmad = TileMmad_; + using CopyGmToL1A = typename TileCopy_::CopyGmToL1A; + using CopyGmToL1B = typename TileCopy_::CopyGmToL1B; + using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; + using CopyL0CToGm = typename TileCopy_::CopyL0CToGm; + using ElementAccumulator = + typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; + using LayoutAInL1 = typename CopyL1ToL0A::LayoutSrc; + using LayoutBInL1 = typename CopyL1ToL0B::LayoutSrc; + using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst; + using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst; + using LayoutCInL0 = layout::zN; + + using L1AAlignHelper = Gemm::helper::L1AlignHelper; + using L1BAlignHelper = Gemm::helper::L1AlignHelper; + + static constexpr bool ENABLE_UNIT_FLAG = DispatchPolicy::ENABLE_UNIT_FLAG; + static constexpr uint32_t L1A_STAGES = DispatchPolicy::L1A_STAGES; + static constexpr uint32_t L1B_STAGES = DispatchPolicy::L1B_STAGES; + static constexpr uint32_t L0AB_STAGES = DispatchPolicy::L0AB_STAGES; + static constexpr uint32_t L0C_STAGES = DispatchPolicy::L0C_STAGES; + + static constexpr uint32_t L1A_SIZE = L1TileShape::M * L1TileShape::K * sizeof(ElementA); + static constexpr uint32_t L1B_SIZE = L1TileShape::K * L1TileShape::N * sizeof(ElementB); + static constexpr uint32_t L0C_SIZE = L0TileShape::M * L0TileShape::N * sizeof(ElementC); + + // Check L1A & L1B + static_assert((L1A_SIZE * L1A_STAGES + L1B_SIZE * L1B_STAGES) <= ArchTag::L1_SIZE, + "L1TileShape exceeding the L1 space!"); + + static constexpr uint32_t L0A_TILE_SIZE = L0TileShape::M * L0TileShape::K * sizeof(ElementA); + static constexpr uint32_t L0B_TILE_SIZE = L0TileShape::K * L0TileShape::N * sizeof(ElementB); + static constexpr uint32_t L0C_TILE_SIZE = L0TileShape::M * L0TileShape::N * sizeof(ElementAccumulator); + + static_assert((L0C_TILE_SIZE * L0C_STAGES <= ArchTag::L0C_SIZE), "L0TileShape exceeding the L0C space!"); + static_assert(L0A_TILE_SIZE * L0AB_STAGES <= ArchTag::L0A_SIZE, "L0TileShape exceeding the L0A space!"); + static_assert(L0B_TILE_SIZE * L0AB_STAGES <= ArchTag::L0B_SIZE, "L0TileShape exceeding the L0B space!"); + + static_assert(L0TileShape::K <= L1TileShape::K, "L0TileShape::K cannot exceed L1TileShape::K"); + static_assert(L0TileShape::M <= L1TileShape::M, "L0TileShape::M cannot exceed L1TileShape::M"); + static_assert(L0TileShape::N <= L1TileShape::N, "L0TileShape::N cannot exceed L1TileShape::N"); + // 32B (256b) aligned + static_assert(Gemm::helper::TileShapeAlignChecker::_ALIGN == 256, + "Tile shape must be 32B aligned."); + + static constexpr uint32_t L0A_PINGPONG_BUF_SIZE = ArchTag::L0A_SIZE / L0AB_STAGES; + static constexpr uint32_t L0B_PINGPONG_BUF_SIZE = ArchTag::L0B_SIZE / L0AB_STAGES; + static constexpr uint32_t L0C_PINGPONG_BUF_SIZE = ArchTag::L0C_SIZE / L0C_STAGES; + + // Check LayoutC + static_assert(std::is_same_v, "LayoutC only support RowMajor yet!"); + + /// Construct + CATLASS_DEVICE + BlockMmad(Arch::Resource &resource, uint32_t l1BufAddrStart = 0) + { + uint32_t l1AOffset = l1BufAddrStart; + uint32_t l1BOffset = l1BufAddrStart + L1A_SIZE * L1A_STAGES; + // Init buffers + for (uint32_t i = 0; i < L1A_STAGES; i++) { + l1ATensorList[i] = resource.l1Buf.template GetBufferByByte(l1AOffset + L1A_SIZE * i); + l1AEventList[i] = i; + AscendC::SetFlag(l1AEventList[i]); + } + for (uint32_t i = 0; i < L1B_STAGES; i++) { + l1BTensorList[i] = resource.l1Buf.template GetBufferByByte(l1BOffset + L1B_SIZE * i); + l1BEventList[i] = i + L1A_STAGES; + AscendC::SetFlag(l1BEventList[i]); + } + + for (uint32_t i = 0; i < L0AB_STAGES; i++) { + l0ATensorList[i] = resource.l0ABuf.template GetBufferByByte(L0A_PINGPONG_BUF_SIZE * i); + l0BTensorList[i] = resource.l0BBuf.template GetBufferByByte(L0B_PINGPONG_BUF_SIZE * i); + l0AEventList[i] = i; + l0BEventList[i] = i + L0AB_STAGES; + AscendC::SetFlag(l0AEventList[i]); + AscendC::SetFlag(l0BEventList[i]); + } + for (uint32_t i = 0; i < L0C_STAGES; i++) { + l0CTensorList[i] = resource.l0CBuf.template GetBufferByByte(L0C_PINGPONG_BUF_SIZE * i); + l0CEventList[i] = i; + AscendC::SetFlag(l0CEventList[i]); + } + } + + /// Destructor + CATLASS_DEVICE + ~BlockMmad() + { + for (uint32_t i = 0; i < L1A_STAGES; i++) { + AscendC::WaitFlag(l1AEventList[i]); + } + for (uint32_t i = 0; i < L1B_STAGES; i++) { + AscendC::WaitFlag(l1BEventList[i]); + } + + for (uint32_t i = 0; i < L0AB_STAGES; i++) { + AscendC::WaitFlag(l0AEventList[i]); + AscendC::WaitFlag(l0BEventList[i]); + } + for (uint32_t i = 0; i < L0C_STAGES; i++) { + AscendC::WaitFlag(l0CEventList[i]); + } + } + + /// Perform a block-scoped matrix multiply-accumulate + CATLASS_DEVICE + void operator()(AscendC::GlobalTensor const &gmBlockA, LayoutA const &layoutA, + AscendC::GlobalTensor const &gmBlockB, LayoutB const &layoutB, + AscendC::GlobalTensor const &gmBlockC, LayoutC const &layoutC, + AscendC::GlobalTensor const &gmNextBlockA, AscendC::GlobalTensor const &gmNextBlockB, + GemmCoord const &actualShape, GemmCoord const &actualShapeNext, bool needLoadNextA, bool needLoadNextB, + bool atomicAdd) + { + auto layoutAInL1 = LayoutAInL1::template MakeLayout(L1TileShape::M, L1TileShape::K); + auto layoutBInL1 = LayoutBInL1::template MakeLayout(L1TileShape::K, L1TileShape::N); + + bool isFirstBlock = !preLoadedA && !preLoadedB; + bool isLastBlock = !needLoadNextA && !needLoadNextB; + uint32_t l1AListIdNext = (l1AListId + 1 < L1A_STAGES) ? (l1AListId + 1) : 0; + uint32_t l1BListIdNext = (l1BListId + 1 < L1B_STAGES) ? (l1BListId + 1) : 0; + + if (isFirstBlock) { + // Load matrix A tile from GM to L1 + AscendC::WaitFlag(l1AEventList[l1AListId]); + auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShape.m(), actualShape.k())); + copyGmToL1A(l1ATensorList[l1AListId], gmBlockA, layoutAInL1, layoutTileA); + AscendC::SetFlag(l1AEventList[l1AListId]); + + // Load matrix B tile from GM to L1 + AscendC::WaitFlag(l1BEventList[l1BListId]); + auto layoutTileB = layoutB.GetTileLayout(MakeCoord(actualShape.k(), actualShape.n())); + copyGmToL1B(l1BTensorList[l1BListId], gmBlockB, layoutBInL1, layoutTileB); + AscendC::SetFlag(l1BEventList[l1BListId]); + } + + if (needLoadNextA && (L1A_STAGES > 1)) { + AscendC::WaitFlag(l1AEventList[l1AListIdNext]); + auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShapeNext.m(), actualShapeNext.k())); + copyGmToL1A(l1ATensorList[l1AListIdNext], gmNextBlockA, layoutAInL1, layoutTileA); + AscendC::SetFlag(l1AEventList[l1AListIdNext]); + } + if (needLoadNextB && (L1B_STAGES > 1)) { + AscendC::WaitFlag(l1BEventList[l1BListIdNext]); + auto layoutTileB = layoutB.GetTileLayout(MakeCoord(actualShapeNext.k(), actualShapeNext.n())); + copyGmToL1B(l1BTensorList[l1BListIdNext], gmNextBlockB, layoutBInL1, layoutTileB); + AscendC::SetFlag(l1BEventList[l1BListIdNext]); + } + + // Get L1 tensor for current stage + auto l1ATensor = l1ATensorList[l1AListId]; + auto l1BTensor = l1BTensorList[l1BListId]; + + // Atomic set once in mmad level + if (atomicAdd) { + AscendC::SetAtomicAdd(); + } else { + AscendC::SetAtomicNone(); + } + + uint32_t l1mLoops = CeilDiv(actualShape.m(), L0TileShape::M); + uint32_t l1nLoops = CeilDiv(actualShape.n(), L0TileShape::N); + for (uint32_t l1mLoopIdx = 0; l1mLoopIdx < l1mLoops; l1mLoopIdx++) { + for (uint32_t l1nLoopIdx = 0; l1nLoopIdx < l1nLoops; l1nLoopIdx++) { + // layoutAInL0 + uint32_t mActual = (l1mLoopIdx < l1mLoops - 1) ? L0TileShape::M : (actualShape.m() - l1mLoopIdx * L0TileShape::M); + uint32_t nActual = (l1nLoopIdx < l1nLoops - 1) ? L0TileShape::N : (actualShape.n() - l1nLoopIdx * L0TileShape::N); + uint32_t mRound = RoundUp(mActual); + uint32_t nRound = RoundUp(nActual); + auto layoutInL0C = LayoutCInL0::MakeLayoutInL0C(MakeCoord(mRound, nRound)); + + uint32_t l0ABufId = 0; + uint32_t l0BBufId = 0; + + // wait/set conditions in block/l1mLoops/l1nLoops level + bool mnLoopStart = (l1mLoopIdx == 0 && l1nLoopIdx == 0); + bool mnLoopEnd = (l1mLoopIdx == l1mLoops - 1) && (l1nLoopIdx == l1nLoops - 1); + bool waitFlagA = (isFirstBlock || preLoadedA) && mnLoopStart; + bool setFlagA = (needLoadNextA || isLastBlock) && mnLoopEnd; + bool waitFlagB = (isFirstBlock || preLoadedB) && mnLoopStart; + bool setFlagB = (needLoadNextB || isLastBlock) && mnLoopEnd; + + if constexpr (!ENABLE_UNIT_FLAG) { + AscendC::WaitFlag(l0CEventList[l0CListId]); + } + + // Get the loop nums on L0 + uint32_t l1kLoops = CeilDiv(actualShape.k(), L0TileShape::K); + for (uint32_t l1kLoopIdx = 0; l1kLoopIdx < l1kLoops; l1kLoopIdx++) { + uint32_t kActual = + (l1kLoopIdx < l1kLoops - 1) ? L0TileShape::K : (actualShape.k() - l1kLoopIdx * L0TileShape::K); + + // Locate the current tile on L0A + auto l0ATile = l0ATensorList[l0ABufId]; + LayoutAInL0 layoutAInL0 = LayoutAInL0::template MakeLayout(mRound, kActual); + // Locate the current tile of matrix A on L1 + MatrixCoord l1ACoord{l1mLoopIdx * L0TileShape::M, l1kLoopIdx * L0TileShape::K}; + auto l1ATile = l1ATensor[layoutAInL1.GetOffset(l1ACoord)]; + + AscendC::WaitFlag(l0AEventList[l0ABufId]); + if (waitFlagA && (l1kLoopIdx == 0)) { + AscendC::WaitFlag(l1AEventList[l1AListId]); + } + + copyL1ToL0A(l0ATile, l1ATile, layoutAInL0, layoutAInL1); + + if (setFlagA && (l1kLoopIdx == l1kLoops - 1)) { + AscendC::SetFlag(l1AEventList[l1AListId]); + } + + // Locate the current tile on L0B + auto l0BTile = l0BTensorList[l0BBufId]; + LayoutBInL0 layoutBInL0 = LayoutBInL0::template MakeLayout(kActual, nRound); + // Locate the current tile of matrix B on L1 + MatrixCoord l1BCoord{l1kLoopIdx * L0TileShape::K, l1nLoopIdx * L0TileShape::N}; + auto l1BTile = l1BTensor[layoutBInL1.GetOffset(l1BCoord)]; + + // Wait for mmad finished + AscendC::WaitFlag(l0BEventList[l0BBufId]); + // If the current tile is the first one on the k&n axis, wait for loading matrix B from GM to L1 + if (waitFlagB && (l1kLoopIdx == 0)) { + AscendC::WaitFlag(l1BEventList[l1BListId]); + } + + // Load current tile from L1 to L0B + copyL1ToL0B(l0BTile, l1BTile, layoutBInL0, layoutBInL1); + + // If the current tile is the last one on the k&n axis, notify to load matrix B from GM to L1 + if (setFlagB && (l1kLoopIdx == l1kLoops - 1)) { + AscendC::SetFlag(l1BEventList[l1BListId]); + } + // Notify to do mmad + AscendC::SetFlag(l0CEventList[l0CListId]); + + // Locate the current tile on L0C + auto l0CTile = l0CTensorList[l0CListId]; + + // Compute the matrix multiplication on L0A and L0B and write the result to the accumulator + // Wait for loading L0B + AscendC::WaitFlag(l0CEventList[l0CListId]); + + // If the current tile is the first tile on the k axis, the accumulator needs to be reset to 0 + bool initC = (l1kLoopIdx == 0); + // If the unit flag is enabled, the unit flag is set according to the calculation progress + uint8_t unitFlag = 0b00; + if constexpr (ENABLE_UNIT_FLAG) { + if (l1kLoopIdx == l1kLoops - 1) { + unitFlag = 0b11; + } else { + unitFlag = 0b10; + } + } + // Perform calculation operations + tileMmad(l0CTile, l0ATile, l0BTile, mRound, nRound, kActual, initC, unitFlag); + + // Notify to move the next L0B tile + AscendC::SetFlag(l0BEventList[l0BBufId]); + l0BBufId = (l0BBufId + 1 < L0AB_STAGES) ? (l0BBufId + 1) : 0; + + AscendC::SetFlag(l0AEventList[l0ABufId]); + l0ABufId = (l0ABufId + 1 < L0AB_STAGES) ? (l0ABufId + 1) : 0; + } + // copy block out + LayoutC layoutBlock = layoutC.GetTileLayout(MakeCoord(mActual, nActual)); + MatrixCoord gmCoordC{l1mLoopIdx * L0TileShape::M, l1nLoopIdx * L0TileShape::N}; + auto gmTileC = gmBlockC[layoutC.GetOffset(gmCoordC)]; + if constexpr (!ENABLE_UNIT_FLAG) { + AscendC::SetFlag(l0CEventList[l0CListId]); + AscendC::WaitFlag(l0CEventList[l0CListId]); + copyL0CToGm(gmTileC, l0CTensorList[l0CListId], layoutBlock, layoutInL0C); + AscendC::SetFlag(l0CEventList[l0CListId]); + } else { + copyL0CToGm(gmTileC, l0CTensorList[l0CListId], layoutBlock, layoutInL0C, 0b11); + } + l0CListId = (l0CListId + 1 < L0C_STAGES) ? (l0CListId + 1) : 0; + } + } + + // if L1A_STAGES/L1B_STAGES == 1, load next block after compute + if (needLoadNextA && (L1A_STAGES == 1)) { + // Load first matrix A tile from GM to L1 + AscendC::WaitFlag(l1AEventList[l1AListIdNext]); + auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShapeNext.m(), actualShapeNext.k())); + copyGmToL1A(l1ATensorList[l1AListIdNext], gmNextBlockA, layoutAInL1, layoutTileA); + AscendC::SetFlag(l1AEventList[l1AListIdNext]); + } + if (needLoadNextB && (L1B_STAGES == 1)) { + // Load first matrix B tile from GM to L1 + AscendC::WaitFlag(l1BEventList[l1BListIdNext]); + auto layoutTileB = layoutB.GetTileLayout(MakeCoord(actualShapeNext.k(), actualShapeNext.n())); + copyGmToL1B(l1BTensorList[l1BListIdNext], gmNextBlockB, layoutBInL1, layoutTileB); + AscendC::SetFlag(l1BEventList[l1BListIdNext]); + } + + preLoadedA = needLoadNextA; + preLoadedB = needLoadNextB; + l1AListId = needLoadNextA ? l1AListIdNext : l1AListId; + l1BListId = needLoadNextB ? l1BListIdNext : l1BListId; + } + +protected: + /// Data members + AscendC::LocalTensor l1ATensorList[L1A_STAGES]; + AscendC::LocalTensor l1BTensorList[L1B_STAGES]; + AscendC::LocalTensor l0ATensorList[L0AB_STAGES]; + AscendC::LocalTensor l0BTensorList[L0AB_STAGES]; + AscendC::LocalTensor l0CTensorList[L0C_STAGES]; + + int32_t l1AEventList[L1A_STAGES]; + int32_t l1BEventList[L1B_STAGES]; + int32_t l0AEventList[L0AB_STAGES]; + int32_t l0BEventList[L0AB_STAGES]; + int32_t l0CEventList[L0C_STAGES]; + + bool preLoadedA{false}; + bool preLoadedB{false}; + uint32_t l1AListId{0}; + uint32_t l1BListId{0}; + uint32_t l0CListId{0}; + + TileMmad tileMmad; + CopyGmToL1A copyGmToL1A; + CopyGmToL1B copyGmToL1B; + CopyL1ToL0A copyL1ToL0A; + CopyL1ToL0B copyL1ToL0B; + CopyL0CToGm copyL0CToGm; +}; + +} // namespace Catlass::Gemm::Block + +#endif // CATLASS_GEMM_BLOCK_BLOCK_MMAD_SINGLE_CORE_SPLITK_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_small.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_small.hpp new file mode 100644 index 0000000000..d74e23afff --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_small.hpp @@ -0,0 +1,319 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_GEMM_BLOCK_BLOCK_MMAD_SMALL_HPP +#define CATLASS_GEMM_BLOCK_BLOCK_MMAD_SMALL_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/kda_coord.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" +#include "catlass/gemm/kda_gemm_helper.hpp" +#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" +#include "catlass/gemm/tile/kda_tile_mmad.hpp" + +namespace Catlass::Gemm::Block { + +template < + uint32_t STAGES_, + bool ENABLE_UNIT_FLAG_, + bool ENABLE_SHUFFLE_K_, + class L1TileShape_, + class L0TileShape_, + class AType_, + class BType_, + class CType_, + class BiasType_, + class TileCopy_, + class TileMmad_> +struct BlockMmad< + MmadAtlasA2Small, + L1TileShape_, + L0TileShape_, + AType_, + BType_, + CType_, + BiasType_, + TileCopy_, + TileMmad_> { +public: + // Type Aliases + using DispatchPolicy = MmadAtlasA2Small; + using ArchTag = typename DispatchPolicy::ArchTag; + using L1TileShape = L1TileShape_; + using L0TileShape = L0TileShape_; + using ElementA = typename AType_::Element; + using LayoutA = typename AType_::Layout; + using ElementB = typename BType_::Element; + using LayoutB = typename BType_::Layout; + using ElementC = typename CType_::Element; + using LayoutC = typename CType_::Layout; + using TileMmad = TileMmad_; + using CopyGmToL1A = typename TileCopy_::CopyGmToL1A; + using CopyGmToL1B = typename TileCopy_::CopyGmToL1B; + using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; + using CopyL0CToGm = typename TileCopy_::CopyL0CToGm; + using ElementAccumulator = + typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; + using LayoutAInL1 = typename CopyL1ToL0A::LayoutSrc; + using LayoutBInL1 = typename CopyL1ToL0B::LayoutSrc; + using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst; + using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst; + using LayoutCInL0 = layout::zN; + + using L1AAlignHelper = Gemm::helper::L1AlignHelper; + using L1BAlignHelper = Gemm::helper::L1AlignHelper; + + static constexpr bool ENABLE_UNIT_FLAG = DispatchPolicy::ENABLE_UNIT_FLAG; + static constexpr bool ENABLE_SHUFFLE_K = DispatchPolicy::ENABLE_SHUFFLE_K; + static constexpr uint32_t STAGES = DispatchPolicy::STAGES; + static constexpr uint32_t L1A_SIZE = L1TileShape::M * L1TileShape::K * sizeof(ElementA); + static constexpr uint32_t L1B_SIZE = L1TileShape::N * L1TileShape::K * sizeof(ElementB); + static constexpr uint32_t L0A_SIZE = ArchTag::L0A_SIZE; + static constexpr uint32_t L0B_SIZE = ArchTag::L0B_SIZE; + static constexpr uint32_t L0C_SIZE = ArchTag::L0C_SIZE; + static constexpr uint32_t L0A_PINGPONG_BUF_SIZE = L0A_SIZE / STAGES; + static constexpr uint32_t L0B_PINGPONG_BUF_SIZE = L0B_SIZE / STAGES; + + // Check LayoutC + static_assert(std::is_same_v, "LayoutC only support RowMajor yet!"); + + // Check L1TileShape + static_assert((L1A_SIZE * STAGES + L1B_SIZE * STAGES) <= ArchTag::L1_SIZE, "L1TileShape exceeding the L1 space!"); + + // Check L0TileShape + static constexpr uint32_t L0A_TILE_SIZE = L0TileShape::M * L0TileShape::K * sizeof(ElementA); + static constexpr uint32_t L0B_TILE_SIZE = L0TileShape::K * L0TileShape::N * sizeof(ElementB); + static constexpr uint32_t L0C_TILE_SIZE = L0TileShape::M * L0TileShape::N * sizeof(ElementAccumulator); + static_assert((L0A_TILE_SIZE * STAGES) <= L0A_SIZE, "L0TileShape exceeding the L0A space!"); + static_assert((L0B_TILE_SIZE * STAGES) <= L0B_SIZE, "L0TileShape exceeding the L0B space!"); + static_assert(L0C_TILE_SIZE <= L0C_SIZE, "L0TileShape exceeding the L0C space!"); + + static_assert(L1TileShape::M == L0TileShape::M && L1TileShape::N == L0TileShape::N, + "The situation where the basic blocks of L1 and L0 differ on the m and n axes is not supported yet"); + static_assert(L0TileShape::K <= L1TileShape::K, "L0TileShape::K cannot exceed L1TileShape::K"); + // 32B (256b) aligned + static_assert(Gemm::helper::TileShapeAlignChecker::_ALIGN == 256, + "Tile shape must be 32B aligned."); + + /// Construct + CATLASS_DEVICE + BlockMmad(Arch::Resource &resource, uint32_t l1BufAddrStart = 0) + { + uint32_t l1AOffset = l1BufAddrStart; + uint32_t l1BOffset = l1BufAddrStart + L1A_SIZE * STAGES; + // Init buffers + for (uint32_t i = 0; i < STAGES; i++) { + // Assign L1/L0A/L0B space for each stages + l1ATensorList[i] = resource.l1Buf.template GetBufferByByte(l1AOffset + L1A_SIZE * i); + l1BTensorList[i] = resource.l1Buf.template GetBufferByByte(l1BOffset + L1B_SIZE * i); + l0ATensorList[i] = resource.l0ABuf.template GetBufferByByte(L0A_PINGPONG_BUF_SIZE * i); + l0BTensorList[i] = resource.l0BBuf.template GetBufferByByte(L0B_PINGPONG_BUF_SIZE * i); + + // Assign event ID for each stages + l1AEventList[i] = i; + l1BEventList[i] = i + STAGES; + l0AEventList[i] = i; + l0BEventList[i] = i + STAGES; + + // The event id that needs to be set before the loop + AscendC::SetFlag(l1AEventList[i]); + AscendC::SetFlag(l1BEventList[i]); + AscendC::SetFlag(l0AEventList[i]); + AscendC::SetFlag(l0BEventList[i]); + } + l0CTensor = resource.l0CBuf.template GetBufferByByte(0); + if constexpr (!ENABLE_UNIT_FLAG) { + AscendC::SetFlag(EVENT_ID0); + } + } + + /// Destructor + CATLASS_DEVICE + ~BlockMmad() + { + for (uint32_t i = 0; i < STAGES; i++) { + AscendC::WaitFlag(l1AEventList[i]); + AscendC::WaitFlag(l1BEventList[i]); + AscendC::WaitFlag(l0AEventList[i]); + AscendC::WaitFlag(l0BEventList[i]); + } + if constexpr (!ENABLE_UNIT_FLAG) { + AscendC::WaitFlag(EVENT_ID0); + } + } + + /// Perform a block-scoped matrix multiply-accumulate + CATLASS_DEVICE + void operator()( + AscendC::GlobalTensor const &gmA, LayoutA const &layoutA, + AscendC::GlobalTensor const &gmB, LayoutB const &layoutB, + AscendC::GlobalTensor const &gmC, LayoutC const &layoutC, + GemmCoord const &actualShape) + { + uint32_t mRound = RoundUp(actualShape.m()); + uint32_t nRound = RoundUp(actualShape.n()); + + auto layoutAInL1 = LayoutAInL1::template MakeLayout(L1TileShape::M, L1TileShape::K); + auto layoutBInL1 = LayoutBInL1::template MakeLayout(L1TileShape::K, L1TileShape::N); + auto layoutInL0C = LayoutCInL0::MakeLayoutInL0C(MakeCoord(mRound, nRound)); + + if constexpr (!ENABLE_UNIT_FLAG) { + AscendC::WaitFlag(EVENT_ID0); + } + + uint32_t kActual = actualShape.k(); + + // k loop + { + // For small-shape matmul, the k-loop executes only once. + + // load first matrix A tile from GM to L1 + AscendC::WaitFlag(l1AEventList[l1ListId]); + auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShape.m(), kActual)); + copyGmToL1A(l1ATensorList[l1ListId], gmA, layoutAInL1, layoutTileA); + AscendC::SetFlag(l1AEventList[l1ListId]); + + // load first matrix B tile from GM to L1 + AscendC::WaitFlag(l1BEventList[l1ListId]); + auto layoutTileB = layoutB.GetTileLayout(MakeCoord(kActual, actualShape.n())); + copyGmToL1B(l1BTensorList[l1ListId], gmB, layoutBInL1, layoutTileB); + AscendC::SetFlag(l1BEventList[l1ListId]); + + // Get L1 tensor for current stage + auto l1ATensor = l1ATensorList[l1ListId]; + auto l1BTensor = l1BTensorList[l1ListId]; + + // Get the loop nums on L0 + uint32_t kPartLoop = CeilDiv(kActual); + + for (int kPartIdx = 0; kPartIdx < kPartLoop; kPartIdx++) { + uint32_t kPartActual = (kPartIdx < kPartLoop - 1) ? + L0TileShape::K : + (kActual - kPartIdx * L0TileShape::K); + + // Locate the current tile on L0A + auto l0ATile = l0ATensorList[l0AListId]; + LayoutAInL0 layoutAInL0 = LayoutAInL0::template MakeLayout(mRound, kPartActual); + // Locate the current tile of matrix A on L1 + MatrixCoord l1AOffset{0, kPartIdx * L0TileShape::K}; + auto l1ATile = l1ATensor[layoutAInL1.GetOffset(l1AOffset)]; + + AscendC::WaitFlag(l0AEventList[l0AListId]); + if (kPartIdx == 0) { + AscendC::WaitFlag(l1AEventList[l1ListId]); + } + + // Load current tile from L1 to L0A + copyL1ToL0A(l0ATile, l1ATile, layoutAInL0, layoutAInL1); + + if (kPartIdx == kPartLoop - 1) { + AscendC::SetFlag(l1AEventList[l1ListId]); + } + + // Locate the current tile on L0B + auto l0BTile = l0BTensorList[l0BListId]; + LayoutBInL0 layoutBInL0 = LayoutBInL0::template MakeLayout(kPartActual, nRound); + // Locate the current tile of matrix B on L1 + MatrixCoord l1BOffset{kPartIdx * L0TileShape::K, 0}; + auto l1BTile = l1BTensor[layoutBInL1.GetOffset(l1BOffset)]; + + // Wait for mmad finished + AscendC::WaitFlag(l0BEventList[l0BListId]); + // If the current tile is the first one on the k&n axis, wait for loading matrix B from GM to L1 + if (kPartIdx == 0) { + AscendC::WaitFlag(l1BEventList[l1ListId]); + } + + // Load current tile from L1 to L0B + copyL1ToL0B(l0BTile, l1BTile, layoutBInL0, layoutBInL1); + + // If the current tile is the last one on the k&n axis, notify to load matrix B from GM to L1 + if (kPartIdx == kPartLoop - 1) { + AscendC::SetFlag(l1BEventList[l1ListId]); + } + // Notify to do mmad + AscendC::SetFlag(EVENT_ID0); + + // Locate the current tile on L0C + auto l0CTile = l0CTensor; + + // Compute the matrix multiplication on L0A and L0B and write the result to the accumulator + // Wait for loading L0B + AscendC::WaitFlag(EVENT_ID0); + + // If the current tile is the first tile on the k axis, the accumulator needs to be reset to 0 + bool initC = (kPartIdx == 0); + // If the unit flag is enabled, the unit flag is set according to the calculation progress + uint8_t unitFlag = 0b00; + if constexpr (ENABLE_UNIT_FLAG) { + if (kPartIdx == kPartLoop - 1) { + unitFlag = 0b11; + } else { + unitFlag = 0b10; + } + } + // Perform calculation operations + tileMmad(l0CTile, l0ATile, l0BTile, mRound, nRound, kPartActual, initC, unitFlag); + + // Notify to move the next L0B tile + AscendC::SetFlag(l0BEventList[l0BListId]); + l0BListId = (l0BListId + 1 < STAGES) ? (l0BListId + 1) : 0; + + AscendC::SetFlag(l0AEventList[l0AListId]); + l0AListId = (l0AListId + 1 < STAGES) ? (l0AListId + 1) : 0; + } + } + + // copy block out + LayoutC layoutBlock = layoutC.GetTileLayout(actualShape.GetCoordMN()); + + if constexpr (!ENABLE_UNIT_FLAG) { + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + copyL0CToGm(gmC, l0CTensor, layoutBlock, layoutInL0C); + AscendC::SetFlag(EVENT_ID0); + } else { + copyL0CToGm(gmC, l0CTensor, layoutBlock, layoutInL0C, 0b11); + } + } + +protected: + // Multi-stage tensors list + AscendC::LocalTensor l1ATensorList[STAGES]; + AscendC::LocalTensor l1BTensorList[STAGES]; + AscendC::LocalTensor l0ATensorList[STAGES]; + AscendC::LocalTensor l0BTensorList[STAGES]; + AscendC::LocalTensor l0CTensor; + + // Multi-stage event id list + int32_t l1AEventList[STAGES]; + int32_t l1BEventList[STAGES]; + int32_t l0AEventList[STAGES]; + int32_t l0BEventList[STAGES]; + + // The id of current stage + uint32_t l1ListId{0}; + uint32_t l0AListId{0}; + uint32_t l0BListId{0}; + + TileMmad tileMmad; + CopyGmToL1A copyGmToL1A; + CopyGmToL1B copyGmToL1B; + CopyL1ToL0A copyL1ToL0A; + CopyL1ToL0B copyL1ToL0B; + CopyL0CToGm copyL0CToGm; +}; + +} // namespace Catlass::Gemm::Block + +#endif // CATLASS_GEMM_BLOCK_BLOCK_MMAD_SMALL_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_sparse_tla.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_sparse_tla.hpp new file mode 100644 index 0000000000..99433b669d --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_sparse_tla.hpp @@ -0,0 +1,397 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ +#ifndef CATLASS_GEMM_BLOCK_BLOCK_MMAD_SPARSE_TLA_HPP +#define CATLASS_GEMM_BLOCK_BLOCK_MMAD_SPARSE_TLA_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/kda_coord.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" +#include "catlass/gemm/kda_gemm_helper.hpp" +#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" +#include "catlass/gemm/tile/kda_tile_mmad.hpp" +#include "tla/kda_tla_layout.hpp" +#include "tla/kda_tensor.hpp" + +namespace Catlass::Gemm::Block { + +template < + class ArchTag_, + bool ENABLE_UNIT_FLAG_, + class L1TileShape_, + class L0TileShape_, + class ElementA_, + class ElementB_, + class ElementC_, + class ElementBias_, + class TileCopy_> +struct BlockMmadSparseTla< + SparseMatmulMultiBlockOnKAxis, + L1TileShape_, + L0TileShape_, + ElementA_, + ElementB_, + ElementC_, + ElementBias_, + TileCopy_> { +public: + using DispatchPolicy = SparseMatmulMultiBlockOnKAxis; + using ArchTag = typename DispatchPolicy::ArchTag; + using L1Shape = L1TileShape_; + using L0Shape = L0TileShape_; + using ElementA = ElementA_; + using LayoutA = typename TileCopy_::LayoutA; + using ElementB = ElementB_; + using LayoutB = typename TileCopy_::LayoutB; + using ElementC = ElementC_; + using LayoutC = typename TileCopy_::LayoutC; + using ElementSparseIndex = uint8_t; + using ElementAccumulator = typename TileCopy_::ElementAccumulator; + + using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; + + using L1ALayout = typename TileCopy_::LayoutTagL1A; + using L1BLayout = typename TileCopy_::LayoutTagL1B; + using L1IdxLayout = typename TileCopy_::LayoutTagL1BIdx; // sparse index fractal 16*8 + using L0ALayout = detail::TagToLayout_t; + using L0BLayout = detail::TagToLayout_t; + using L0IdxLayout = detail::TagToLayout_t; + + using L1ATlaTensor = tla::Tensor, L1ALayout, + tla::Coord, AscendC::TPosition::GM>; + using L1BTlaTensor = tla::Tensor, L1BLayout, + tla::Coord, AscendC::TPosition::GM>; + + static constexpr uint16_t LIMIT_MNSIZE = 10; + static constexpr uint16_t ALIGN_NUM = 16; + static constexpr int32_t DENSE_MATRIX_B_OFFSET = 2; + static constexpr int32_t INDEX_MATRIX_OFFSET = 8; + static constexpr int32_t SHARE_LENS_COUNT = 3; + static constexpr auto L1_M = tla::get<0>(L1Shape{}); + static constexpr auto L1_N = tla::get<1>(L1Shape{}); + static constexpr auto L1_K = tla::get<2>(L1Shape{}); + static constexpr auto L0_M = tla::get<0>(L0Shape{}); + static constexpr auto L0_N = tla::get<1>(L0Shape{}); + static constexpr auto L0_K = tla::get<2>(L0Shape{}); + static constexpr uint32_t STAGES = DispatchPolicy::STAGES; + + // L1 tile size + static constexpr uint32_t L1A_TILE_SIZE = L1_M * L1_K * sizeof(ElementA); + static constexpr uint32_t L1B_TILE_SIZE = L1_N * L1_K * sizeof(ElementB) / DENSE_MATRIX_B_OFFSET; + static constexpr uint32_t L1B_IDX_TILE_SIZE = L1_N * L1_K * sizeof(ElementB) / INDEX_MATRIX_OFFSET; + // L0 tile size + static constexpr uint32_t L0A_TILE_SIZE = L0_M * L0_K * sizeof(ElementA); + static constexpr uint32_t L0B_TILE_SIZE = L0_K * L0_N * sizeof(ElementB); + static constexpr uint32_t L0C_TILE_SIZE = L0_M * L0_N * sizeof(ElementAccumulator); + + static constexpr uint32_t L0A_SIZE = ArchTag::L0A_SIZE; + static constexpr uint32_t L0B_SIZE = ArchTag::L0B_SIZE; + static constexpr uint32_t L0C_SIZE = ArchTag::L0C_SIZE; + static constexpr uint32_t L0A_PINGPONG_BUF_SIZE = L0A_SIZE / STAGES; + static constexpr uint32_t L0B_PINGPONG_BUF_SIZE = L0B_SIZE / STAGES; + + static constexpr bool TRANS_A = tla::detail::isColumnMajor::value; + static constexpr bool TRANS_B = tla::detail::isColumnMajor::value; + + // Check L1TileShape + static_assert((L1A_TILE_SIZE + L1B_TILE_SIZE + L1B_IDX_TILE_SIZE) * STAGES <= ArchTag::L1_SIZE, + "L1TileShape exceeding the L1 space!"); + + // Check L0TileShape + static_assert(L0A_TILE_SIZE <= ArchTag::L0A_SIZE, "L0TileShape exceeding the L0A space!"); + static_assert(L0B_TILE_SIZE <= ArchTag::L0B_SIZE, "L0TileShape exceeding the L0B space!"); + static_assert(L0C_TILE_SIZE <= ArchTag::L0C_SIZE, "L0TileShape exceeding the L0C space!"); + + static_assert(L1_M == L0_M && L1_N == L0_N, + "The situation where the basic blocks of L1 and L0 differ on the m and n axes is not supported yet"); + static_assert(L0_K <= L1_K, "L0TileShape::K cannot exceed L1TileShape::K"); + + static_assert((std::is_same_v && std::is_same_v && + std::is_same_v), + "Unsupported dtype"); + static_assert(tla::detail::isColumnMajor::value, "Only support MatrixB transpose"); + + CATLASS_DEVICE + BlockMmadSparseTla() + { + Init(); + } + + CATLASS_DEVICE + ~BlockMmadSparseTla() + { + End(); + } + + /** + * @brief Initialization function + * + * This function initializes the buffers and parameters required for matrix multiplication operations + */ + CATLASS_DEVICE + void Init() + { + constexpr static int32_t aL0MatrixByteSize = + L0_M * L0_K * sizeof(ElementA); + constexpr static int32_t bL0MatrixByteSize = + L0_N * L0_K * sizeof(ElementB); + constexpr static int32_t cMatrixByteSize = + L0_M * L0_N * sizeof(ElementAccumulator); + + constexpr static int dbL0AFlag = (aL0MatrixByteSize * STAGES > ArchTag::L0A_SIZE) ? 1 : STAGES; + constexpr static int dbL0BFlag = (bL0MatrixByteSize * STAGES > ArchTag::L0B_SIZE) ? 1 : STAGES; + + uint32_t l1AOffset = 0; + uint32_t l1BOffset = L1A_TILE_SIZE * STAGES; + uint32_t l1BIdxOffset = l1BOffset + L1B_TILE_SIZE * STAGES; + // Init buffers + for (uint32_t i = 0; i < STAGES; i++) { + // Assign L1/L0A/L0B space for each stages + l1ATensorList_[i] = l1Buf_[l1AOffset + L1A_TILE_SIZE * i].template ReinterpretCast(); + l1BTensorList_[i] = l1Buf_[l1BOffset + L1B_TILE_SIZE * i].template ReinterpretCast(); + l1BIdxTensorList_[i] = l1Buf_[l1BIdxOffset + L1B_IDX_TILE_SIZE * i].template ReinterpretCast(); + l0ATensorList_[i] = l0ABuf_[L0A_PINGPONG_BUF_SIZE * i].template ReinterpretCast(); + l0BTensorList_[i] = l0BBuf_[L0B_PINGPONG_BUF_SIZE * i].template ReinterpretCast(); + + // Assign event ID for each stages + l1AEventList_[i] = i; + l1BEventList_[i] = i + STAGES; + l1BIdxEventList_[i] = i + STAGES * 2; + l0AEventList_[i] = i; + l0BEventList_[i] = i + STAGES; + + // The event id that needs to be set before the loop + AscendC::SetFlag(l1AEventList_[i]); + AscendC::SetFlag(l1BEventList_[i]); + AscendC::SetFlag(l1BIdxEventList_[i]); + AscendC::SetFlag(l0AEventList_[i]); + AscendC::SetFlag(l0BEventList_[i]); + } + l0CTensor_ = l0CBuf_.template ReinterpretCast(); + AscendC::SetFlag(EVENT_ID0); + } + + /** + * @brief Perform the matrix multiplication operations + * @param [in] DstTensor: destination tensor type + * @param [in] SrcATensor: matrix A source tensor type + * @param [in] SrcBTensor: matrix B source tensor type + * @param [in] SparseIdxTensor: the type of the sparse index tensor + * @param [in] Shape: shape type + * @param [in] c: destination tensor + * @param [in] a: matrix A source tensor + * @param [in] b: matrix B source tensor + * @param [in] sparseIdx: the sparse index tensor + * @param [in] actualShape: actual shape + */ + template + CATLASS_DEVICE void operator()(DstTensor &c, const SrcATensor &a, const SrcBTensor &b, + const SparseIdxTensor &sparseIdx, const Shape &actualShape) + { + using CopyL0CToGm = typename TileCopy_::template CopyL0CToGm; + CopyL0CToGm copyL0CToGm; + + int32_t mIter = CeilDiv(tla::get<0>(actualShape), L0_M); + int32_t nIter = CeilDiv(tla::get<1>(actualShape), L0_N); + int32_t kOuterIter = CeilDiv(tla::get<2>(actualShape), L1_K); + int32_t tailBaseM = (tla::get<0>(actualShape) % L0_M) == 0 ? L0_M : (tla::get<0>(actualShape) % L0_M); + int32_t tailBaseN = (tla::get<1>(actualShape) % L0_N) == 0 ? L0_N : (tla::get<1>(actualShape) % L0_N); + int32_t tailBaseK = (tla::get<2>(actualShape) % L0_K) == 0 ? L0_K : (tla::get<2>(actualShape) % L0_K); + int32_t tailL1K = (tla::get<2>(actualShape) % L1_K) == 0 ? L1_K : (tla::get<2>(actualShape) % L1_K); + + for (auto mIndex = 0; mIndex < mIter; ++mIndex) { + for (auto nIndex = 0; nIndex < nIter; ++nIndex) { + int mPartActual = (mIndex + 1 == mIter) ? tailBaseM : L0_M; + int nPartActual = (nIndex + 1 == nIter) ? tailBaseN : L0_N; + auto l0CLayout = tla::MakeLayoutL0C(mPartActual, nPartActual); + AscendC::WaitFlag(EVENT_ID0); + auto l0CTlaTensor = tla::MakeTensor(l0CTensor_, l0CLayout, Arch::PositionL0C{}); + + for (auto kOuterIndex = 0; kOuterIndex < kOuterIter; ++kOuterIndex) { + uint32_t l1ListIdNext = (l1ListId + 1 < STAGES) ? (l1ListId + 1) : 0; + int l1K = (kOuterIndex + 1 == kOuterIter) ? tailL1K : L1_K; + // -----------------Step1: GM -> L1 ----------------- + auto aTileHeight = mPartActual; + auto aTileWidth = l1K; + auto aL1Row = mIndex; + auto aL1Col = kOuterIndex; + + auto l1ALayout = tla::MakeLayout(aTileHeight, aTileWidth); + AscendC::WaitFlag(l1AEventList_[l1ListIdNext]); + using CopyGmToL1A = typename TileCopy_::template CopyGmToL1A; + + auto l1ATensor = tla::MakeTensor(l1ATensorList_[l1ListIdNext], l1ALayout, Arch::PositionL1{}); + auto SrcATlaTensor = tla::MakeTensor(a.data(), a.layout(), + tla::MakeCoord(aL1Row * L0_M, + aL1Col * L1_K), + Arch::PositionGM{}); + CopyGmToL1A copyGmToL1A; + copyGmToL1A(l1ATensor, SrcATlaTensor); + + AscendC::SetFlag(l1AEventList_[l1ListIdNext]); + + auto bTileHeight = l1K; + auto bTileWidth = nPartActual; + auto bL1Row = kOuterIndex; + auto bL1Col = nIndex; + using CopyGmToL1B = typename TileCopy_::template CopyGmToL1B; + using CopyGmToL1BIdx = typename TileCopy_::template CopyGmToL1BIdx; + CopyGmToL1B copyGmToL1B; + CopyGmToL1BIdx copyGmToL1BIdx; + + auto l1BLayout = tla::MakeLayout(bTileHeight / DENSE_MATRIX_B_OFFSET, bTileWidth); + + AscendC::WaitFlag(l1BEventList_[l1ListIdNext]); + + auto l1BTensor = tla::MakeTensor(l1BTensorList_[l1ListIdNext], l1BLayout, Arch::PositionL1{}); + auto srcBTlaTensor = GetTile(b, tla::MakeCoord(bL1Row * L1_K / DENSE_MATRIX_B_OFFSET, bL1Col * L0_N), + tla::MakeShape(l1K / DENSE_MATRIX_B_OFFSET, nPartActual)); + + auto l1BIndexLayout = + tla::MakeLayout(bTileHeight / INDEX_MATRIX_OFFSET, bTileWidth); + AscendC::WaitFlag(l1BIdxEventList_[l1ListIdNext]); + + auto l1BIdxTensor = tla::MakeTensor(l1BIdxTensorList_[l1ListIdNext], l1BIndexLayout, Arch::PositionL1{}); + auto sparseIdxTlaTensor = tla::MakeTensor(sparseIdx.data(), sparseIdx.layout(), + tla::MakeCoord(bL1Row * L1_K / INDEX_MATRIX_OFFSET, bL1Col * L0_N), + Arch::PositionGM{}); + + copyGmToL1B(l1BTensor, srcBTlaTensor); + + copyGmToL1BIdx(l1BIdxTensor, sparseIdxTlaTensor); + + AscendC::SetFlag(l1BEventList_[l1ListIdNext]); + AscendC::SetFlag(l1BIdxEventList_[l1ListIdNext]); + AscendC::WaitFlag(l1AEventList_[l1ListIdNext]); + AscendC::WaitFlag(l1BEventList_[l1ListIdNext]); + AscendC::WaitFlag(l1BIdxEventList_[l1ListIdNext]); + + int32_t kInnerIter = + (kOuterIndex + 1 == kOuterIter) ? CeilDiv(tailL1K, L0_K) : CeilDiv(L1_K, L0_K); + for (auto kInnerIndex = 0; kInnerIndex < kInnerIter; ++kInnerIndex) { + int baseK = + (kOuterIndex + 1 == kOuterIter) && (kInnerIndex + 1 == kInnerIter) ? tailBaseK : L0_K; + + // -----------------Step2: L1 -> L0----------------- + AscendC::WaitFlag(l0AEventList_[l0AListId]); + + auto aBaseHeight = mPartActual; + auto aBaseWidth = baseK; + auto aL0Row = mIndex; + auto aL0Col = kInnerIndex; + auto l0ALayout = tla::MakeLayout(aBaseHeight, aBaseWidth); + auto l0ATensor = tla::MakeTensor(l0ATensorList_[l0AListId], l0ALayout, Arch::PositionL0A{}); + + auto l1ATileTensor = tla::MakeTensor(l1ATensor.data(), l1ATensor.layout(), + tla::MakeCoord(aL0Row * 0, + aL0Col * L0_K), + Arch::PositionL1{}); + copyL1ToL0A_(l0ATensor, l1ATileTensor); + + auto bBaseHeight = baseK; + auto bBaseWidth = nPartActual; + auto bL0Row = kInnerIndex; + auto bL0Col = nIndex; + auto l0BLayout = tla::MakeLayout( + bBaseHeight / DENSE_MATRIX_B_OFFSET, bBaseWidth); + + auto l0BTensor = tla::MakeTensor(l0BTensorList_[l0BListId], l0BLayout, Arch::PositionL0B{}); + auto l1BTileTensor = tla::MakeTensor(l1BTensor.data(), l1BTensor.layout(), + tla::MakeCoord(bL0Row * L0_K / DENSE_MATRIX_B_OFFSET, 0), + Arch::PositionL1{}); + + copyL1ToL0B_(l0BTensor, l1BTileTensor, l1BIdxTensor); + + AscendC::SetFlag(l0AEventList_[l0AListId]); + AscendC::WaitFlag(l0AEventList_[l0AListId]); + + // -----------------Step3: compute----------------- + AscendC::MmadParams mmadParams; + // GEMV is automatically enabled when setting M=1 in normal mode + mmadParams.m = (mPartActual == 1 ? ALIGN_NUM : mPartActual); + mmadParams.k = baseK; + mmadParams.n = nPartActual; + mmadParams.unitFlag = 0; + mmadParams.cmatrixInitVal = (kOuterIndex == 0) && (kInnerIndex == 0) ? true : false; + MmadWithSparse(l0CTensor_, l0ATensorList_[l0AListId], l0BTensorList_[l0BListId], mmadParams); + if ((mPartActual / ALIGN_NUM) * (nPartActual / ALIGN_NUM) < LIMIT_MNSIZE) { + AscendC::PipeBarrier(); + } + AscendC::SetFlag(l0AEventList_[l0AListId]); + l0AListId = (l0AListId + 1 < STAGES) ? (l0AListId + 1) : 0; + l0BListId = (l0BListId + 1 < STAGES) ? (l0BListId + 1) : 0; + } + AscendC::SetFlag(l1AEventList_[l1ListIdNext]); + AscendC::SetFlag(l1BEventList_[l1ListIdNext]); + AscendC::SetFlag(l1BIdxEventList_[l1ListIdNext]); + l1ListId = l1ListIdNext; + } + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + // -----------------Step4: L0C -> GM----------------- + auto l0CSrcTensor = tla::MakeTensor(l0CTlaTensor.data(), l0CTlaTensor.layout(), + tla::MakeCoord(mIndex * L0_M, nIndex * L0_N), + Arch::PositionL0C{}); + copyL0CToGm(c, l0CSrcTensor); + + AscendC::SetFlag(EVENT_ID0); + } + } + } + +private: + /** + * @brief End function, release all events + */ + CATLASS_DEVICE + void End() + { + for (uint32_t i = 0; i < STAGES; i++) { + AscendC::WaitFlag(l1AEventList_[i]); + AscendC::WaitFlag(l1BEventList_[i]); + AscendC::WaitFlag(l1BIdxEventList_[i]); + AscendC::WaitFlag(l0AEventList_[i]); + AscendC::WaitFlag(l0BEventList_[i]); + } + AscendC::WaitFlag(EVENT_ID0); + } + +private: + CopyL1ToL0A copyL1ToL0A_; + CopyL1ToL0B copyL1ToL0B_; + + AscendC::LocalTensor l1Buf_{AscendC::LocalTensor(AscendC::TPosition::A1, 0, ArchTag::L1_SIZE)}; + AscendC::LocalTensor l0ABuf_{AscendC::LocalTensor(AscendC::TPosition::A2, 0, ArchTag::L0A_SIZE)}; + AscendC::LocalTensor l0BBuf_{AscendC::LocalTensor(AscendC::TPosition::B2, 0, ArchTag::L0B_SIZE)}; + AscendC::LocalTensor l0CBuf_{AscendC::LocalTensor(AscendC::TPosition::CO1, 0, ArchTag::L0C_SIZE)}; + AscendC::LocalTensor l1ATensorList_[STAGES]; + AscendC::LocalTensor l1BTensorList_[STAGES]; + AscendC::LocalTensor l1BIdxTensorList_[STAGES]; + AscendC::LocalTensor l0ATensorList_[STAGES]; + AscendC::LocalTensor l0BTensorList_[STAGES]; + AscendC::LocalTensor l0CTensor_; + + // Multi-stage event id list + int32_t l1AEventList_[STAGES]; + int32_t l1BEventList_[STAGES]; + int32_t l1BIdxEventList_[STAGES]; + int32_t l0AEventList_[STAGES]; + int32_t l0BEventList_[STAGES]; + + // The id of current stage + uint32_t l1ListId{0}; + uint32_t l0AListId{0}; + uint32_t l0BListId{0}; +}; +} // namespace Catlass::Gemm::Block + +#endif // CATLASS_GEMM_BLOCK_BLOCK_MMAD_SPARSE_TLA_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_streamk.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_streamk.hpp new file mode 100644 index 0000000000..73f760d116 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_streamk.hpp @@ -0,0 +1,406 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_GEMM_BLOCK_BLOCK_MMAD_STREAMK_HPP +#define CATLASS_GEMM_BLOCK_BLOCK_MMAD_STREAMK_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/kda_coord.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" +#include "catlass/gemm/kda_gemm_helper.hpp" +#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" +#include "catlass/gemm/tile/kda_tile_mmad.hpp" + +namespace Catlass::Gemm::Block { + +template < + bool ENABLE_UNIT_FLAG_, + bool ENABLE_SHUFFLE_K_, + class L1TileShape_, + class L0TileShape_, + class AType_, + class BType_, + class CType_, + class BiasType_, + class TileCopy_, + class TileMmad_> +struct BlockMmad< + MmadAtlasA2Streamk, + L1TileShape_, + L0TileShape_, + AType_, + BType_, + CType_, + BiasType_, + TileCopy_, + TileMmad_> { +public: + // Type Aliases + using DispatchPolicy = MmadAtlasA2Streamk; + using ArchTag = typename DispatchPolicy::ArchTag; + using L1TileShape = L1TileShape_; + using L0TileShape = L0TileShape_; + using ElementA = typename AType_::Element; + using LayoutA = typename AType_::Layout; + using ElementB = typename BType_::Element; + using LayoutB = typename BType_::Layout; + using ElementC = typename CType_::Element; + using LayoutC = typename CType_::Layout; + using TileMmad = TileMmad_; + using CopyGmToL1A = typename TileCopy_::CopyGmToL1A; + using CopyGmToL1B = typename TileCopy_::CopyGmToL1B; + using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; + using CopyL0CToGmNormalBlock = typename TileCopy_::CopyL0CToGm; + using CopyL0CToGmStreamkBlock = Gemm::Tile::CopyL0CToGm>; + using ElementAccumulator = + typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; + using LayoutAInL1 = typename CopyL1ToL0A::LayoutSrc; + using LayoutBInL1 = typename CopyL1ToL0B::LayoutSrc; + using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst; + using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst; + using LayoutCInL0 = layout::zN; + + using L1AAlignHelper = Gemm::helper::L1AlignHelper; + using L1BAlignHelper = Gemm::helper::L1AlignHelper; + + static constexpr bool ENABLE_UNIT_FLAG = DispatchPolicy::ENABLE_UNIT_FLAG; + static constexpr bool ENABLE_SHUFFLE_K = DispatchPolicy::ENABLE_SHUFFLE_K; + static constexpr uint32_t STAGES = DispatchPolicy::STAGES; + static constexpr uint32_t L1A_SIZE = L1TileShape::M * L1TileShape::K * sizeof(ElementA); + static constexpr uint32_t L1B_SIZE = L1TileShape::N * L1TileShape::K * sizeof(ElementB); + static constexpr uint32_t L0A_SIZE = ArchTag::L0A_SIZE; + static constexpr uint32_t L0B_SIZE = ArchTag::L0B_SIZE; + static constexpr uint32_t L0C_SIZE = ArchTag::L0C_SIZE; + static constexpr uint32_t L0A_PINGPONG_BUF_SIZE = L0A_SIZE / STAGES; + static constexpr uint32_t L0B_PINGPONG_BUF_SIZE = L0B_SIZE / STAGES; + + // Check LayoutC + static_assert(std::is_same_v, "LayoutC only support RowMajor yet!"); + + // Check L1TileShape + static_assert((L1A_SIZE * STAGES + L1B_SIZE * STAGES) <= ArchTag::L1_SIZE, "L1TileShape exceeding the L1 space!"); + + // Check L0TileShape + static constexpr uint32_t L0A_TILE_SIZE = L0TileShape::M * L0TileShape::K * sizeof(ElementA); + static constexpr uint32_t L0B_TILE_SIZE = L0TileShape::K * L0TileShape::N * sizeof(ElementB); + static constexpr uint32_t L0C_TILE_SIZE = L0TileShape::M * L0TileShape::N * sizeof(ElementAccumulator); + static_assert((L0A_TILE_SIZE * STAGES) <= L0A_SIZE, "L0TileShape exceeding the L0A space!"); + static_assert((L0B_TILE_SIZE * STAGES) <= L0B_SIZE, "L0TileShape exceeding the L0B space!"); + static_assert(L0C_TILE_SIZE <= L0C_SIZE, "L0TileShape exceeding the L0C space!"); + + static_assert(L1TileShape::M == L0TileShape::M && L1TileShape::N == L0TileShape::N, + "The situation where the basic blocks of L1 and L0 differ on the m and n axes is not supported yet"); + static_assert(L0TileShape::K <= L1TileShape::K, "L0TileShape::K cannot exceed L1TileShape::K"); + // 32B (256b) aligned + static_assert(Gemm::helper::TileShapeAlignChecker::_ALIGN == 256, + "Tile shape must be 32B aligned."); + + /// Construct + CATLASS_DEVICE + BlockMmad(Arch::Resource &resource, uint32_t l1BufAddrStart = 0) + { + kPartLenMax = min(L0A_PINGPONG_BUF_SIZE / sizeof(ElementA) / L1TileShape::M / L1AAlignHelper::ELE_NUM_PER_C0 * + L1AAlignHelper::ELE_NUM_PER_C0, + L0B_PINGPONG_BUF_SIZE / sizeof(ElementB) / L1TileShape::N / L1BAlignHelper::ELE_NUM_PER_C0 * + L1BAlignHelper::ELE_NUM_PER_C0); + + if constexpr (std::is_same_v && std::is_same_v) { + kPartLenMax = RoundDown(kPartLenMax); + } + + uint32_t l1AOffset = l1BufAddrStart; + uint32_t l1BOffset = l1BufAddrStart + L1A_SIZE * STAGES; + // Init buffers + for (uint32_t i = 0; i < STAGES; i++) { + l1ATensorList[i] = resource.l1Buf.template GetBufferByByte(l1AOffset + L1A_SIZE * i); + l1BTensorList[i] = resource.l1Buf.template GetBufferByByte(l1BOffset + L1B_SIZE * i); + l0ATensorList[i] = resource.l0ABuf.template GetBufferByByte(L0A_PINGPONG_BUF_SIZE * i); + l0BTensorList[i] = resource.l0BBuf.template GetBufferByByte(L0B_PINGPONG_BUF_SIZE * i); + + l1AEventList[i] = i; + l1BEventList[i] = i + STAGES; + l0AEventList[i] = i; + l0BEventList[i] = i + STAGES; + AscendC::SetFlag(l1AEventList[i]); + AscendC::SetFlag(l1BEventList[i]); + AscendC::SetFlag(l0AEventList[i]); + AscendC::SetFlag(l0BEventList[i]); + } + l0CTensor = resource.l0CBuf.template GetBufferByByte(0); + if constexpr (!ENABLE_UNIT_FLAG) { + AscendC::SetFlag(EVENT_ID0); + } + } + + /// Destructor + CATLASS_DEVICE + ~BlockMmad() + { + for (uint32_t i = 0; i < STAGES; i++) { + AscendC::WaitFlag(l1AEventList[i]); + AscendC::WaitFlag(l1BEventList[i]); + AscendC::WaitFlag(l0AEventList[i]); + AscendC::WaitFlag(l0BEventList[i]); + } + if constexpr (!ENABLE_UNIT_FLAG) { + AscendC::WaitFlag(EVENT_ID0); + } + } + + /// Perform a block-scoped matrix multiply-accumulate + CATLASS_DEVICE + void operator()( + AscendC::GlobalTensor const &gmBlockA, LayoutA const &layoutA, + AscendC::GlobalTensor const &gmBlockB, LayoutB const &layoutB, + AscendC::GlobalTensor const &gmBlockC, LayoutC const &layoutC, + AscendC::GlobalTensor const &gmW, LayoutC const &layoutW, + AscendC::GlobalTensor const &gmNextBlockA, AscendC::GlobalTensor const &gmNextBlockB, + GemmCoord const &actualShape, GemmCoord const &actualShapeNext, + bool isFirstBlock, bool hasNextBlock, bool isSkBlock) + { + uint32_t mRound = RoundUp(actualShape.m()); + uint32_t nRound = RoundUp(actualShape.n()); + + auto layoutAInL1 = LayoutAInL1::template MakeLayout(L1TileShape::M, L1TileShape::K); + auto layoutBInL1 = LayoutBInL1::template MakeLayout(L1TileShape::K, L1TileShape::N); + auto layoutInL0C = LayoutCInL0::MakeLayoutInL0C(MakeCoord(mRound, nRound)); + + uint32_t kTileCount = CeilDiv(actualShape.k()); + + if constexpr (!ENABLE_UNIT_FLAG) { + AscendC::WaitFlag(EVENT_ID0); + } + uint32_t startTileIdx = 0; + if constexpr (ENABLE_SHUFFLE_K) { + startTileIdx = AscendC::GetBlockIdx(); + } + uint32_t firstTileIdx = startTileIdx % kTileCount; + uint32_t lastTileIdx = (startTileIdx + kTileCount - 1) % kTileCount; + uint32_t kActual = + (firstTileIdx < kTileCount - 1) ? L1TileShape::K : (actualShape.k() - firstTileIdx * L1TileShape::K); + + // k loop + for (uint32_t kLoopIdx = 0; kLoopIdx < kTileCount; kLoopIdx++) { + uint32_t shuffleKIdx = (startTileIdx + kLoopIdx) % kTileCount; + // Load first matrix A tile in total kernel loop from GM to L1 + if (shuffleKIdx == firstTileIdx && isFirstBlock) { + MatrixCoord gmTileACoord{0, shuffleKIdx * L1TileShape::K}; + MatrixCoord gmTileBCoord{shuffleKIdx * L1TileShape::K, 0}; + auto gmTileA = gmBlockA[layoutA.GetOffset(gmTileACoord)]; + auto gmTileB = gmBlockB[layoutB.GetOffset(gmTileBCoord)]; + // Load first matrix A tile from GM to L1 + AscendC::WaitFlag(l1AEventList[l1ListId]); + auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShape.m(), kActual)); + copyGmToL1A(l1ATensorList[l1ListId], gmTileA, layoutAInL1, layoutTileA); + AscendC::SetFlag(l1AEventList[l1ListId]); + + // Load first matrix B tile from GM to L1 + AscendC::WaitFlag(l1BEventList[l1ListId]); + auto layoutTileB = layoutB.GetTileLayout(MakeCoord(kActual, actualShape.n())); + copyGmToL1B(l1BTensorList[l1ListId], gmTileB, layoutBInL1, layoutTileB); + AscendC::SetFlag(l1BEventList[l1ListId]); + } + + uint32_t l1ListIdNext = (l1ListId + 1 < STAGES) ? (l1ListId + 1) : 0; + uint32_t kActualNext{0}; + + // preload next tile from GM to L1 + if (shuffleKIdx != lastTileIdx) { + uint32_t shuffleKIdxNext = (startTileIdx + kLoopIdx + 1) % kTileCount; + // Get L1 tensor for next stage + auto l1ATensor = l1ATensorList[l1ListIdNext]; + auto l1BTensor = l1BTensorList[l1ListIdNext]; + // Get GM tensor for next stage + kActualNext = (shuffleKIdxNext < kTileCount - 1) ? L1TileShape::K : (actualShape.k() - shuffleKIdxNext * L1TileShape::K); + MatrixCoord gmTileACoord{0, shuffleKIdxNext * L1TileShape::K}; + MatrixCoord gmTileBCoord{shuffleKIdxNext * L1TileShape::K, 0}; + auto gmTileA = gmBlockA[layoutA.GetOffset(gmTileACoord)]; + auto gmTileB = gmBlockB[layoutB.GetOffset(gmTileBCoord)]; + + // load next matrix A tile from GM to L1 + AscendC::WaitFlag(l1AEventList[l1ListIdNext]); + auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShape.m(), kActualNext)); + copyGmToL1A(l1ATensor, gmTileA, layoutAInL1, layoutTileA); + AscendC::SetFlag(l1AEventList[l1ListIdNext]); + + // load next matrix B tile from GM to L1 + AscendC::WaitFlag(l1BEventList[l1ListIdNext]); + auto layoutTileB = layoutB.GetTileLayout(MakeCoord(kActualNext, actualShape.n())); + copyGmToL1B(l1BTensor, gmTileB, layoutBInL1, layoutTileB); + AscendC::SetFlag(l1BEventList[l1ListIdNext]); + } + if (shuffleKIdx == lastTileIdx && hasNextBlock) { + uint32_t kTileCount = CeilDiv(actualShapeNext.k()); + uint32_t firstTileIdx = startTileIdx % kTileCount; + // Get L1 tensor for next stage + auto l1ATensor = l1ATensorList[l1ListIdNext]; + auto l1BTensor = l1BTensorList[l1ListIdNext]; + // Get GM tensor for next stage + kActualNext = (firstTileIdx < kTileCount - 1) ? L1TileShape::K : (actualShapeNext.k() - firstTileIdx * L1TileShape::K); + MatrixCoord gmTileACoord{0, firstTileIdx * L1TileShape::K}; + MatrixCoord gmTileBCoord{firstTileIdx * L1TileShape::K, 0}; + auto gmTileA = gmNextBlockA[layoutA.GetOffset(gmTileACoord)]; + auto gmTileB = gmNextBlockB[layoutB.GetOffset(gmTileBCoord)]; + // load next matrix A tile from GM to L1 + AscendC::WaitFlag(l1AEventList[l1ListIdNext]); + auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShapeNext.m(), kActualNext)); + copyGmToL1A(l1ATensor, gmTileA, layoutAInL1, layoutTileA); + AscendC::SetFlag(l1AEventList[l1ListIdNext]); + + // load next matrix B tile from GM to L1 + AscendC::WaitFlag(l1BEventList[l1ListIdNext]); + auto layoutTileB = layoutB.GetTileLayout(MakeCoord(kActualNext, actualShapeNext.n())); + copyGmToL1B(l1BTensor, gmTileB, layoutBInL1, layoutTileB); + AscendC::SetFlag(l1BEventList[l1ListIdNext]); + } + + // Get L1 tensor for current stage + auto l1ATensor = l1ATensorList[l1ListId]; + auto l1BTensor = l1BTensorList[l1ListId]; + + // Get the loop nums on L0 + uint32_t kPartLoop = CeilDiv(kActual, kPartLenMax); + + uint32_t l0ABufId = 0; + uint32_t l0BBufId = 0; + + for (int kPartIdx = 0; kPartIdx < kPartLoop; kPartIdx++) { + uint32_t kPartActual = + (kPartIdx < kPartLoop - 1) ? kPartLenMax : (kActual - kPartIdx * kPartLenMax); + + // Locate the current tile on L0A + auto l0ATile = l0ATensorList[l0ABufId]; + LayoutAInL0 layoutAInL0 = LayoutAInL0::template MakeLayout(mRound, kPartActual); + // Locate the current tile of matrix A on L1 + MatrixCoord l1ACoord{0, kPartIdx * kPartLenMax}; + auto l1ATile = l1ATensor[layoutAInL1.GetOffset(l1ACoord)]; + + AscendC::WaitFlag(l0AEventList[l0ABufId]); + if (kPartIdx == 0) { + AscendC::WaitFlag(l1AEventList[l1ListId]); + } + + copyL1ToL0A(l0ATile, l1ATile, layoutAInL0, layoutAInL1); + + if (kPartIdx == kPartLoop - 1) { + AscendC::SetFlag(l1AEventList[l1ListId]); + } + + // Locate the current tile on L0B + auto l0BTile = l0BTensorList[l0BBufId]; + LayoutBInL0 layoutBInL0 = LayoutBInL0::template MakeLayout(kPartActual, nRound); + // Locate the current tile of matrix B on L1 + MatrixCoord l1BCoord{kPartIdx * kPartLenMax, 0}; + auto l1BTile = l1BTensor[layoutBInL1.GetOffset(l1BCoord)]; + + // Wait for mmad finished + AscendC::WaitFlag(l0BEventList[l0BBufId]); + // If the current tile is the first one on the k&n axis, wait for loading matrix B from GM to L1 + if (kPartIdx == 0) { + AscendC::WaitFlag(l1BEventList[l1ListId]); + } + + // Load current tile from L1 to L0B + copyL1ToL0B(l0BTile, l1BTile, layoutBInL0, layoutBInL1); + + // If the current tile is the last one on the k&n axis, notify to load matrix B from GM to L1 + if (kPartIdx == kPartLoop - 1) { + AscendC::SetFlag(l1BEventList[l1ListId]); + } + // Notify to do mmad + AscendC::SetFlag(EVENT_ID0); + + // Locate the current tile on L0C + MatrixCoord l0CCoord{static_cast(0U), static_cast(0U)}; + auto l0CTile = l0CTensor[layoutInL0C.GetOffset(l0CCoord)]; + + // Compute the matrix multiplication on L0A and L0B and write the result to the accumulator + // Wait for loading L0B + AscendC::WaitFlag(EVENT_ID0); + + // If the current tile is the first tile on the k axis, the accumulator needs to be reset to 0 + bool initC = ((kLoopIdx == 0) && (kPartIdx == 0)); + // If the unit flag is enabled, the unit flag is set according to the calculation progress + uint8_t unitFlag = 0b00; + if constexpr (ENABLE_UNIT_FLAG) { + if ((kLoopIdx == kTileCount - 1) && (kPartIdx == kPartLoop - 1)) { + unitFlag = 0b11; + } else { + unitFlag = 0b10; + } + } + // Perform calculation operations + tileMmad(l0CTile, l0ATile, l0BTile, mRound, nRound, kPartActual, initC, unitFlag); + + // Notify to move the next L0B tile + AscendC::SetFlag(l0BEventList[l0BBufId]); + + l0BBufId = (l0BBufId + 1 < STAGES) ? (l0BBufId + 1) : 0; + + AscendC::SetFlag(l0AEventList[l0ABufId]); + l0ABufId = (l0ABufId + 1 < STAGES) ? (l0ABufId + 1) : 0; + } + + l1ListId = l1ListIdNext; + kActual = kActualNext; + } + + // copy block out + LayoutC layoutBlock = layoutC.GetTileLayout(actualShape.GetCoordMN()); + + if constexpr (!ENABLE_UNIT_FLAG) { + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + if (isSkBlock) { + copyL0CToGmStreamkBlock(gmW, l0CTensor, layoutW, layoutInL0C); + } else { + copyL0CToGmNormalBlock(gmBlockC, l0CTensor, layoutBlock, layoutInL0C); + } + AscendC::SetFlag(EVENT_ID0); + } else { + if (isSkBlock) { + copyL0CToGmStreamkBlock(gmW, l0CTensor, layoutW, layoutInL0C, 0b11); + } else { + copyL0CToGmNormalBlock(gmBlockC, l0CTensor, layoutBlock, layoutInL0C, 0b11); + } + } + } + +protected: + /// Data members + AscendC::LocalTensor l1ATensorList[STAGES]; + AscendC::LocalTensor l1BTensorList[STAGES]; + AscendC::LocalTensor l0ATensorList[STAGES]; + AscendC::LocalTensor l0BTensorList[STAGES]; + AscendC::LocalTensor l0CTensor; + + int32_t l1AEventList[STAGES]; + int32_t l1BEventList[STAGES]; + int32_t l0AEventList[STAGES]; + int32_t l0BEventList[STAGES]; + + uint32_t l1ListId{0}; + uint32_t kPartLenMax; + + TileMmad tileMmad; + CopyGmToL1A copyGmToL1A; + CopyGmToL1B copyGmToL1B; + CopyL1ToL0A copyL1ToL0A; + CopyL1ToL0B copyL1ToL0B; + CopyL0CToGmNormalBlock copyL0CToGmNormalBlock; + CopyL0CToGmStreamkBlock copyL0CToGmStreamkBlock; +}; + +} // namespace Catlass::Gemm::Block + +#endif // CATLASS_GEMM_BLOCK_BLOCK_MMAD_STREAMK_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_w4a4_per_token_per_channel_dequant.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_w4a4_per_token_per_channel_dequant.hpp new file mode 100644 index 0000000000..fc30e82382 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_w4a4_per_token_per_channel_dequant.hpp @@ -0,0 +1,494 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ +#ifndef CATLASS_GEMM_BLOCK_BLOCK_MMAD_W4A4_HPP +#define CATLASS_GEMM_BLOCK_BLOCK_MMAD_W4A4_HPP + +#include "catlass/arch/kda_resource.hpp" +#include "catlass/kda_catlass.hpp" +#include "catlass/kda_coord.hpp" +#include "catlass/detail/kda_callback.hpp" +#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" +#include "catlass/gemm/kda_gemm_helper.hpp" +#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" +#include "catlass/gemm/tile/kda_tile_mmad.hpp" +#include "catlass/kda_gemm_coord.hpp" + +namespace Catlass::Gemm::Block { + +template < + uint32_t PRELOAD_STAGES_, + uint32_t L1_STAGES_, + uint32_t L0A_STAGES_, + uint32_t L0B_STAGES_, + uint32_t L0C_STAGES_, + bool ENABLE_UNIT_FLAG_, + bool ENABLE_SHUFFLE_K_, + class L1TileShape_, + class L0TileShape_, + class AType_, + class BType_, + class CType_, + class BiasType_, + class TileCopy_, + class TileMmad_> +struct BlockMmad< + MmadAtlasA2W4A4MatmulPerTokenPerChannelDequant< + PRELOAD_STAGES_, + L1_STAGES_, + L0A_STAGES_, + L0B_STAGES_, + L0C_STAGES_, + ENABLE_UNIT_FLAG_, + ENABLE_SHUFFLE_K_>, + L1TileShape_, + L0TileShape_, + AType_, + BType_, + CType_, + BiasType_, + TileCopy_, + TileMmad_> { +public: + // Type Aliases + using DispatchPolicy = MmadAtlasA2W4A4MatmulPerTokenPerChannelDequant< + PRELOAD_STAGES_, + L1_STAGES_, + L0A_STAGES_, + L0B_STAGES_, + L0C_STAGES_, + ENABLE_UNIT_FLAG_, + ENABLE_SHUFFLE_K_>; + using ArchTag = typename DispatchPolicy::ArchTag; + using L1TileShape = L1TileShape_; + using L0TileShape = L0TileShape_; + using ElementA = typename AType_::Element; + using LayoutA = typename AType_::Layout; + using ElementB = typename BType_::Element; + using LayoutB = typename BType_::Layout; + using ElementC = typename CType_::Element; + using LayoutC = typename CType_::Layout; + using TileMmad = TileMmad_; + using CopyGmToL1A = typename TileCopy_::CopyGmToL1A; + using CopyGmToL1B = typename TileCopy_::CopyGmToL1B; + using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; + using CopyL0CToGm = typename TileCopy_::CopyL0CToGm; + using CopyGmToL1Scale = typename TileCopy_::CopyGmToL1Scale; + using CopyL1ToFP = typename TileCopy_::CopyL1ToFP; + using ElementAccumulator = typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; + using LayoutAInL1 = typename CopyL1ToL0A::LayoutSrc; + using LayoutBInL1 = typename CopyL1ToL0B::LayoutSrc; + using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst; + using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst; + using LayoutCInL0 = layout::zN; + using LayoutScale = layout::VectorLayout; + using LayoutL1Scale = layout::VectorLayout; + using LayoutFpBuf = layout::VectorLayout; + + using L1AAlignHelper = Gemm::helper::L1AlignHelper; + using L1BAlignHelper = Gemm::helper::L1AlignHelper; + + static constexpr uint32_t PRELOAD_STAGES = DispatchPolicy::PRELOAD_STAGES; + static constexpr uint32_t L1_STAGES = DispatchPolicy::L1_STAGES; + static constexpr uint32_t L0A_STAGES = DispatchPolicy::L0A_STAGES; + static constexpr uint32_t L0B_STAGES = DispatchPolicy::L0B_STAGES; + static constexpr uint32_t L0C_STAGES = DispatchPolicy::L0C_STAGES; + + static constexpr bool ENABLE_UNIT_FLAG = DispatchPolicy::ENABLE_UNIT_FLAG; + static constexpr bool ENABLE_SHUFFLE_K = DispatchPolicy::ENABLE_SHUFFLE_K; + + // L1 tile size + static constexpr uint32_t L1A_TILE_SIZE = L1TileShape::M * L1TileShape::K * sizeof(int8_t) / 2; + static constexpr uint32_t L1B_TILE_SIZE = L1TileShape::N * L1TileShape::K * sizeof(int8_t) / 2; + static constexpr uint32_t SCALE_KTILE_SIZE = L1TileShape::N * sizeof(uint64_t); + // L0 tile size + static constexpr uint32_t L0A_TILE_SIZE = L0TileShape::M * L0TileShape::K * sizeof(int8_t) / 2; + static constexpr uint32_t L0B_TILE_SIZE = L0TileShape::K * L0TileShape::N * sizeof(int8_t) / 2; + static constexpr uint32_t L0C_TILE_SIZE = L1TileShape::M * L1TileShape::N * sizeof(ElementAccumulator); + + // Check LayoutC + static_assert(std::is_same_v, "LayoutC only support RowMajor yet!"); + + // Check L1TileShape + static_assert( + (L1A_TILE_SIZE + L1B_TILE_SIZE + SCALE_KTILE_SIZE) * L1_STAGES <= ArchTag::L1_SIZE, + "L1TileShape exceeding the L1 space!"); + + // Check L0TileShape + static_assert(L0A_TILE_SIZE * L0A_STAGES <= ArchTag::L0A_SIZE, "L0TileShape exceeding the L0A space!"); + static_assert(L0B_TILE_SIZE * L0B_STAGES <= ArchTag::L0B_SIZE, "L0TileShape exceeding the L0B space!"); + static_assert(L0C_TILE_SIZE * L0C_STAGES <= ArchTag::L0C_SIZE, "L0TileShape exceeding the L0C space!"); + + static_assert( + L1TileShape::M == L0TileShape::M && L1TileShape::N == L0TileShape::N, + "The situation where the basic blocks of L1 and L0 differ on the m and n axes is not supported yet"); + static_assert(L0TileShape::K <= L1TileShape::K, "L0TileShape::K cannot exceed L1TileShape::K"); + + static constexpr auto L1A_LAYOUT = LayoutAInL1::template MakeLayout(L1TileShape::M, L1TileShape::K); + static constexpr auto L1B_LAYOUT = LayoutBInL1::template MakeLayout(L1TileShape::K, L1TileShape::N); + + CATLASS_DEVICE + BlockMmad(Arch::Resource &resource, uint32_t l1BufAddrStart = 0, uint32_t FpAddrStart = 0) + { + InitL1(resource, l1BufAddrStart); + InitFpBuf(resource, FpAddrStart); + InitL0A(resource); + InitL0B(resource); + InitL0C(resource); + } + + CATLASS_DEVICE + ~BlockMmad() + { + for (uint32_t i = 0; i < L1_STAGES; ++i) { + AscendC::WaitFlag(l1AEventList[i]); + AscendC::WaitFlag(l1BEventList[i]); + AscendC::WaitFlag(l1ScaleEventList[i]); + } + for (uint32_t i = 0; i < L0A_STAGES; ++i) { + AscendC::WaitFlag(l0AEventList[i]); + } + for (uint32_t i = 0; i < L0B_STAGES; ++i) { + AscendC::WaitFlag(l0BEventList[i]); + } + for (uint32_t i = 0; i < L0C_STAGES; ++i) { + AscendC::WaitFlag(l0CEventList[i]); + } + } + + CATLASS_DEVICE + void operator()( + AscendC::GlobalTensor const &gmScale, + LayoutScale const &layoutScale, + AscendC::GlobalTensor const &gmBlockA, + LayoutA const &layoutA, + AscendC::GlobalTensor const &gmBlockB, + LayoutB const &layoutB, + AscendC::GlobalTensor const &gmBlockC, + LayoutC const &layoutC, + GemmCoord const &actualShape, + Callback const &callbackBeforeFixpipe, + Callback const &callbackAfterFixpipe) + { + uint32_t kTileCount = CeilDiv(actualShape.k()); + + uint32_t mRound = RoundUp(actualShape.m()); + uint32_t nRound = RoundUp(actualShape.n()); + + uint32_t startTileIdx = 0; + if constexpr (ENABLE_SHUFFLE_K) { + startTileIdx = AscendC::GetBlockIdx() % kTileCount; + } + + // Copy scale from GM to L1 + auto layoutTileScale = layoutScale.GetTileLayout(MakeCoord(L1TileShape::N)); + LayoutL1Scale layoutL1KTileScale{L1TileShape::N}; + AscendC::WaitFlag(l1ScaleEventList[l1ScaleListId]); + copyGmToL1Scale(l1ScaleBufList[l1ScaleListId], gmScale, layoutL1KTileScale, layoutTileScale); + AscendC::SetFlag(l1ScaleEventList[l1ScaleListId]); + + for (uint32_t kLoopIdx = 0; kLoopIdx < kTileCount; ++kLoopIdx) { + uint32_t kTileIdx = (startTileIdx + kLoopIdx < kTileCount) ? (startTileIdx + kLoopIdx) : (startTileIdx + kLoopIdx - kTileCount); + + uint32_t kActual = (kTileIdx < kTileCount - 1) ? L1TileShape::K : (actualShape.k() - kTileIdx * L1TileShape::K); + + // Emission load instruction from GM to L1 + MatrixCoord gmTileAOffset{0, kTileIdx * L1TileShape::K}; + MatrixCoord gmTileBOffset{kTileIdx * L1TileShape::K, 0}; + auto gmTileA = gmBlockA[layoutA.GetOffset(gmTileAOffset)]; + auto gmTileB = gmBlockB[layoutB.GetOffset(gmTileBOffset)]; + // Load first matrix A tile from GM to L1 + AscendC::WaitFlag(l1AEventList[l1ListId]); + auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShape.m(), kActual)); + copyGmToL1A(l1ATensorList[l1ListId], gmTileA, L1A_LAYOUT, layoutTileA); + AscendC::SetFlag(l1AEventList[l1ListId]); + // Load first matrix B tile from GM to L1 + AscendC::WaitFlag(l1BEventList[l1ListId]); + auto layoutTileB = layoutB.GetTileLayout(MakeCoord(kActual, actualShape.n())); + copyGmToL1B(l1BTensorList[l1ListId], gmTileB, L1B_LAYOUT, layoutTileB); + AscendC::SetFlag(l1BEventList[l1ListId]); + + // If the number of preload instructions reaches the upper limit, perform an mmad calculation on L1 tile + if (preloadCount == PRELOAD_STAGES) { + L1TileMmad(l1TileMmadParamsList[l1TileMmadParamsId]); + } + + // Store the current load status + uint32_t preloadL1TileMmadParamsId = (l1TileMmadParamsId + preloadCount < PRELOAD_STAGES) ? (l1TileMmadParamsId + preloadCount) : (l1TileMmadParamsId + preloadCount - PRELOAD_STAGES); + auto &l1TileMmadParams = l1TileMmadParamsList[preloadL1TileMmadParamsId]; + l1TileMmadParams.l1ListId = l1ListId; + l1TileMmadParams.mRound = mRound; + l1TileMmadParams.nRound = nRound; + l1TileMmadParams.kActual = kActual; + l1TileMmadParams.isKLoopFirst = (kLoopIdx == 0); + l1TileMmadParams.isKLoopLast = (kLoopIdx == kTileCount - 1); + if (kLoopIdx == kTileCount - 1) { + l1TileMmadParams.l1ScaleListId = l1ScaleListId; + l1TileMmadParams.gmBlockC = gmBlockC; + l1TileMmadParams.layoutCInGm = layoutC.GetTileLayout(actualShape.GetCoordMN()); + l1TileMmadParams.callbackBeforeFixpipe = callbackBeforeFixpipe; + l1TileMmadParams.callbackAfterFixpipe = callbackAfterFixpipe; + } + + if (preloadCount < PRELOAD_STAGES) { + ++preloadCount; + } else { + l1TileMmadParamsId = (l1TileMmadParamsId + 1 < PRELOAD_STAGES) ? (l1TileMmadParamsId + 1) : 0; + } + l1ListId = (l1ListId + 1 < L1_STAGES) ? (l1ListId + 1) : 0; + } + l1ScaleListId = (l1ScaleListId + 1 < L1_STAGES) ? (l1ScaleListId + 1) : 0; + } + + CATLASS_DEVICE + void SynchronizeBlock() + { + while (preloadCount > 0) { + L1TileMmad(l1TileMmadParamsList[l1TileMmadParamsId]); + l1TileMmadParamsId = (l1TileMmadParamsId + 1 < PRELOAD_STAGES) ? (l1TileMmadParamsId + 1) : 0; + --preloadCount; + } + } + +private: + struct L1TileMmadParams { + uint32_t l1ListId; + uint32_t l1ScaleListId; + uint32_t mRound; + uint32_t nRound; + uint32_t kActual; + bool isKLoopFirst; + bool isKLoopLast; + AscendC::GlobalTensor gmBlockC; + LayoutC layoutCInGm; + Callback callbackBeforeFixpipe; + Callback callbackAfterFixpipe; + + CATLASS_DEVICE + L1TileMmadParams() = default; + }; + + CATLASS_DEVICE + void InitL1(Arch::Resource &resource, uint32_t l1BufAddrStart) + { + uint32_t l1AOffset = l1BufAddrStart; + uint32_t l1BOffset = l1BufAddrStart + L1A_TILE_SIZE * L1_STAGES; + uint32_t l1ScaleOffset = l1BOffset + L1B_TILE_SIZE * L1_STAGES; + for (uint32_t i = 0; i < L1_STAGES; ++i) { + l1ATensorList[i] = resource.l1Buf.template GetBufferByByte(l1AOffset + L1A_TILE_SIZE * i); + l1BTensorList[i] = resource.l1Buf.template GetBufferByByte(l1BOffset + L1B_TILE_SIZE * i); + l1ScaleBufList[i] = resource.l1Buf.template GetBufferByByte(l1ScaleOffset + SCALE_KTILE_SIZE * i); + l1AEventList[i] = i; + l1BEventList[i] = i + L1_STAGES; + l1ScaleEventList[i] = i; + AscendC::SetFlag(l1AEventList[i]); + AscendC::SetFlag(l1BEventList[i]); + AscendC::SetFlag(l1ScaleEventList[i]); + } + } + + CATLASS_DEVICE + void InitFpBuf(Arch::Resource &resource, uint32_t FpAddrStart) + { + uint32_t FpOffset = FpAddrStart; + fixpipeBuf = resource.fpBuf.template GetBufferByByte(FpOffset); + } + + CATLASS_DEVICE + void InitL0A(Arch::Resource &resource) + { + for (uint32_t i = 0; i < L0A_STAGES; ++i) { + l0ATensorList[i] = resource.l0ABuf.template GetBufferByByte(L0A_TILE_SIZE * i); + l0AEventList[i] = i; + AscendC::SetFlag(l0AEventList[i]); + } + } + + CATLASS_DEVICE + void InitL0B(Arch::Resource &resource) + { + for (uint32_t i = 0; i < L0B_STAGES; ++i) { + l0BTensorList[i] = resource.l0BBuf.template GetBufferByByte(L0B_TILE_SIZE * i); + l0BEventList[i] = i + L0A_STAGES; + AscendC::SetFlag(l0BEventList[i]); + } + } + + CATLASS_DEVICE + void InitL0C(Arch::Resource &resource) + { + for (uint32_t i = 0; i < L0C_STAGES; ++i) { + l0CTensorList[i] = resource.l0CBuf.template GetBufferByByte(L0C_TILE_SIZE * i); + l0CEventList[i] = i; + AscendC::SetFlag(l0CEventList[i]); + } + } + + CATLASS_DEVICE + void L1TileMmad(L1TileMmadParams const ¶ms) + { + uint32_t mPartLoop = CeilDiv(params.mRound); + uint32_t nPartLoop = CeilDiv(params.nRound); + uint32_t kPartLoop = CeilDiv(params.kActual); + auto &l1ATensor = l1ATensorList[params.l1ListId]; + auto &l1BTensor = l1BTensorList[params.l1ListId]; + + LayoutFpBuf layoutFpBuf{L1TileShape::N}; + LayoutL1Scale layoutL1Scale{L1TileShape::N}; + + auto &l0CTensor = l0CTensorList[l0CListId]; + LayoutCInL0 layoutCInL0 = LayoutCInL0::MakeLayoutInL0C(MakeCoord(params.mRound, params.nRound)); + + if constexpr (!ENABLE_UNIT_FLAG) { + if (params.isKLoopFirst) { + AscendC::WaitFlag(l0CEventList[l0CListId]); + } + } + + for (uint32_t mPartIdx = 0; mPartIdx < mPartLoop; ++mPartIdx) { + uint32_t mPartActual = (mPartIdx < mPartLoop - 1) ? L0TileShape::M : (params.mRound - mPartIdx * L0TileShape::M); + + for (uint32_t kPartIdx = 0; kPartIdx < kPartLoop; ++kPartIdx) { + uint32_t kPartActual = (kPartIdx < kPartLoop - 1) ? L0TileShape::K : (params.kActual - kPartIdx * L0TileShape::K); + + auto &l0ATile = l0ATensorList[l0AListId]; + auto layoutAInL0 = LayoutAInL0::template MakeLayout(mPartActual, kPartActual); + auto l1AOffset = MakeCoord(mPartIdx, kPartIdx) * L0TileShape::ToCoordMK(); + auto l1ATile = l1ATensor[L1A_LAYOUT.GetOffset(l1AOffset)]; + + AscendC::WaitFlag(l0AEventList[l0AListId]); + if ((mPartIdx == 0) && (kPartIdx == 0)) { + AscendC::WaitFlag(l1AEventList[params.l1ListId]); + } + copyL1ToL0A(l0ATile, l1ATile, layoutAInL0, L1A_LAYOUT); + if ((mPartIdx == mPartLoop - 1) && (kPartIdx == kPartLoop - 1)) { + AscendC::SetFlag(l1AEventList[params.l1ListId]); + } + + for (uint32_t nPartIdx = 0; nPartIdx < nPartLoop; ++nPartIdx) { + uint32_t nPartActual = (nPartIdx < nPartLoop - 1) ? L0TileShape::N : (params.nRound - nPartIdx * L0TileShape::N); + + auto &l0BTile = l0BTensorList[l0BListId]; + auto layoutBInL0 = LayoutBInL0::template MakeLayout(kPartActual, nPartActual); + auto l1BOffset = MakeCoord(kPartIdx, nPartIdx) * L0TileShape::ToCoordKN(); + auto l1BTile = l1BTensor[L1B_LAYOUT.GetOffset(l1BOffset)]; + + AscendC::WaitFlag(l0BEventList[l0BListId]); + if ((kPartIdx == 0) && (nPartIdx == 0)) { + AscendC::WaitFlag(l1BEventList[params.l1ListId]); + } + copyL1ToL0B(l0BTile, l1BTile, layoutBInL0, L1B_LAYOUT); + if ((kPartIdx == kPartLoop - 1) && (nPartIdx == nPartLoop - 1)) { + AscendC::SetFlag(l1BEventList[params.l1ListId]); + } + + AscendC::SetFlag(EVENT_ID0); + + auto l0COffset = MakeCoord(mPartIdx, nPartIdx) * L0TileShape::ToCoordMN(); + auto l0CTile = l0CTensor[layoutCInL0.GetOffset(l0COffset)]; + + AscendC::WaitFlag(EVENT_ID0); + // If the current tile is the first tile on the k axis, the accumulator needs to be reset to 0 + bool initC = (params.isKLoopFirst && (kPartIdx == 0)); + // If the unit flag is enabled, the unit flag is set according to the calculation progress + uint8_t unitFlag = 0b00; + if constexpr (ENABLE_UNIT_FLAG) { + if (params.isKLoopLast && (mPartIdx == mPartLoop - 1) && (kPartIdx == kPartLoop - 1) && (nPartIdx == nPartLoop - 1)) { + unitFlag = 0b11; + } else { + unitFlag = 0b10; + } + } + tileMmad(l0CTile, l0ATile, l0BTile, mPartActual, nPartActual, kPartActual, initC, unitFlag); + + AscendC::SetFlag(l0BEventList[l0BListId]); + l0BListId = (l0BListId + 1 < L0B_STAGES) ? (l0BListId + 1) : 0; + } + AscendC::SetFlag(l0AEventList[l0AListId]); + l0AListId = (l0AListId + 1 < L0A_STAGES) ? (l0AListId + 1) : 0; + } + } + + if (params.isKLoopLast) { + auto layoutCInGm = params.layoutCInGm; + params.callbackBeforeFixpipe(); + + if constexpr (!ENABLE_UNIT_FLAG) { + AscendC::SetFlag(l0CEventList[l0CListId]); + AscendC::WaitFlag(l0CEventList[l0CListId]); + + // Copy scale from L1 to FixpipeBuffer + auto &l1ScaleBuf = l1ScaleBufList[params.l1ScaleListId]; + AscendC::WaitFlag(params.l1ScaleListId); + copyL1ToFP(fixpipeBuf, l1ScaleBuf, layoutFpBuf, layoutL1Scale); + AscendC::SetFlag(params.l1ScaleListId); + + AscendC::PipeBarrier(); + + copyL0CToGm(params.gmBlockC, l0CTensor, fixpipeBuf, layoutCInGm, layoutCInL0); + AscendC::SetFlag(l0CEventList[l0CListId]); + } else { + // Copy scale from L1 to FixpipeBuffer + auto &l1ScaleBuf = l1ScaleBufList[params.l1ScaleListId]; + AscendC::WaitFlag(params.l1ScaleListId); + copyL1ToFP(fixpipeBuf, l1ScaleBuf, layoutFpBuf, layoutL1Scale); + AscendC::SetFlag(params.l1ScaleListId); + + AscendC::PipeBarrier(); + + copyL0CToGm(params.gmBlockC, l0CTensor, fixpipeBuf, layoutCInGm, layoutCInL0, 0b11); + } + + l0CListId = (l0CListId + 1 < L0C_STAGES) ? (l0CListId + 1) : 0; + + params.callbackAfterFixpipe(); + } + } + + AscendC::LocalTensor fixpipeBuf; + + AscendC::LocalTensor l1ATensorList[L1_STAGES]; + AscendC::LocalTensor l1BTensorList[L1_STAGES]; + AscendC::LocalTensor l1ScaleBufList[L1_STAGES]; + int32_t l1AEventList[L1_STAGES]; + int32_t l1BEventList[L1_STAGES]; + int32_t l1ScaleEventList[L1_STAGES]; + uint32_t l1ListId{0}; + uint32_t l1ScaleListId{0}; + + AscendC::LocalTensor l0ATensorList[L0A_STAGES]; + int32_t l0AEventList[L0A_STAGES]; + uint32_t l0AListId{0}; + + AscendC::LocalTensor l0BTensorList[L0B_STAGES]; + int32_t l0BEventList[L0B_STAGES]; + uint32_t l0BListId{0}; + + AscendC::LocalTensor l0CTensorList[L0C_STAGES_]; + int32_t l0CEventList[L0C_STAGES_]; + uint32_t l0CListId{0}; + + L1TileMmadParams l1TileMmadParamsList[PRELOAD_STAGES]; + uint32_t l1TileMmadParamsId{0}; + uint32_t preloadCount{0}; + + TileMmad tileMmad; + CopyGmToL1A copyGmToL1A; + CopyGmToL1B copyGmToL1B; + CopyL1ToL0A copyL1ToL0A; + CopyL1ToL0B copyL1ToL0B; + CopyL0CToGm copyL0CToGm; + CopyGmToL1Scale copyGmToL1Scale; + CopyL1ToFP copyL1ToFP; +}; + +} // namespace Catlass::Gemm::Block + +#endif // CATLASS_GEMM_BLOCK_BLOCK_MMAD_W4A4_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad.hpp new file mode 100644 index 0000000000..d832b08517 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad.hpp @@ -0,0 +1,143 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_BLOCK_MMAD_HPP +#define CATLASS_KDA_BLOCK_MMAD_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" +#include "catlass/gemm/tile/kda_tile_mmad.hpp" + +namespace Catlass::Gemm::Block { + +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 2201) +template < + class DispatchPolicy, + class L1TileShape, + class L0TileShape, + class AType, + class BType, + class CType, + class BiasType = void, + class TileCopy = Gemm::Tile::TileCopy, + class TileMmad = Gemm::Tile::TileMmad > +struct BlockMmad { + static_assert(DEPENDENT_FALSE, "BlockMmad is not implemented for this DispatchPolicy"); +}; + +/// new add for the reason that i am using the dispatchpolicy which is same as the policy of the optimized_matmul +// so i add a new one class to avoid the conflict +template < + class DispatchPolicy, + class L1TileShape, + class L0TileShape, + class AType, + class BType, + class CType, + class BiasType = void, + class TileCopy = Gemm::Tile::TileCopyGemm, // change the name + class TileMmad = Gemm::Tile::TileMmad > +struct BlockGemm { + static_assert(DEPENDENT_FALSE, "BlockMmad is not implemented for this DispatchPolicy"); +}; + +template < + class DispatchPolicy, + class AType, + class BType, + class CType, + class BiasType, + class TileCopy, + class TileMmad> +struct BlockMmadAiv { + static_assert(DEPENDENT_FALSE, "BlockMmadAiv is not implemented for this DispatchPolicy"); +}; + +template < + class DispatchPolicy, + class L1TileShape, + class L0TileShape, + class ElementA, + class ElementB, + class ElementC, + class ElementBias = void, + class TileCopy = Gemm::Tile::SparseTileCopyTla > +struct BlockMmadSparseTla { + static_assert(DEPENDENT_FALSE, "BlockMmadSparseTla is not implemented for this DispatchPolicy"); +}; + +#endif + +template < + class DispatchPolicy, + class L1TileShape, + class L0TileShape, + class ElementA, + class ElementB, + class ElementC, + class ElementBias = void, + class TileCopy = Gemm::Tile::PackedTileCopyTla, + class TileMmad = + Gemm::Tile::TileMmadTla > +struct BlockMmadTla { + static_assert(DEPENDENT_FALSE, "BlockMmadTla is not implemented for this DispatchPolicy"); +}; + +} // namespace Catlass::Gemm::Block + +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 2201) +#include "catlass/gemm/block/kda_block_mmad_pingpong.hpp" +#include "catlass/gemm/block/kda_block_mmad_fa_qk.hpp" +#include "catlass/gemm/block/kda_block_mmad_fa_pv.hpp" +#include "catlass/gemm/block/kda_block_mmad_mla_qk.hpp" +#include "catlass/gemm/block/kda_block_mmad_mla_pv.hpp" +#include "catlass/gemm/block/kda_block_mmad_mla_qk_tp1_spec.hpp" +#include "catlass/gemm/block/kda_block_mmad_mla_pv_tp1_spec.hpp" +#include "catlass/gemm/block/kda_block_mmad_preload.hpp" +#include "catlass/gemm/block/kda_block_mmad_preload_async.hpp" +#include "catlass/gemm/block/kda_block_mmad_preload_async_with_callback.hpp" +#include "catlass/gemm/block/kda_block_mmad_gemm.hpp" +#include "catlass/gemm/block/kda_block_mmad_pingpong_bias.hpp" +#include "catlass/gemm/block/block_mmad_fai_qk_head_tail.hpp" +#include "catlass/gemm/block/block_mmad_fai_qk_normal.hpp" +#include "catlass/gemm/block/block_mmad_fai_pv_head_tail.hpp" +#include "catlass/gemm/block/block_mmad_fai_pv_normal.hpp" +#include "catlass/gemm/block/block_mmad_pingpong_full_loadA.hpp" +#include "catlass/gemm/block/block_mmad_pingpong_full_loadA_tla.hpp" +#include "catlass/gemm/block/block_mmad_pingpong_with_prologue.hpp" +#include "catlass/gemm/block/block_mmad_pingpong_slice_k_with_prologue.hpp" +#include "catlass/gemm/block/block_mmad_dynamic_common.hpp" +#include "catlass/gemm/block/block_mmad_dynamic_small.hpp" +#include "catlass/gemm/block/block_mmad_dynamic_streamk.hpp" +#include "catlass/gemm/block/block_mmad_dynamic_single_core_splitk.hpp" +#include "catlass/gemm/block/block_mmad_dynamic_preload_async_with_callback.hpp" +#include "catlass/gemm/block/block_mmad_small.hpp" +#include "catlass/gemm/block/block_mmad_single_core_splitk.hpp" +#include "catlass/gemm/block/block_mmad_dynamic_aiv.hpp" +#include "catlass/gemm/block/block_mmad_streamk.hpp" +#include "catlass/gemm/block/block_mmad_w4a4_per_token_per_channel_dequant.hpp" +#include "catlass/gemm/block/block_mmad_sparse_tla.hpp" +#include "catlass/gemm/block/block_mmad_fai_qk_head_tail_tla.hpp" +#include "catlass/gemm/block/block_mmad_fai_qk_normal_tla.hpp" +#include "catlass/gemm/block/block_mmad_fai_pv_head_tail_tla.hpp" +#include "catlass/gemm/block/block_mmad_fai_pv_normal_tla.hpp" +#endif + +#include "catlass/gemm/block/block_mmad_pingpong_tla.hpp" +#include "catlass/gemm/block/block_mmad_pingpong_dequant_tla.hpp" +#include "catlass/gemm/block/block_mmad_pingpong_tla_v2.hpp" +#include "catlass/gemm/block/block_mmad_preload_tla.hpp" +#include "catlass/gemm/block/block_mmad_preload_async_with_callback_tla.hpp" +#include "catlass/gemm/block/block_mmad_fai_pv_tla.hpp" +#include "catlass/gemm/block/block_mmad_fai_qk_tla.hpp" +#include "catlass/gemm/block/block_mmad_pingpong_per_group_per_block_tla.hpp" +#endif // CATLASS_KDA_BLOCK_MMAD_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_fa_pv.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_fa_pv.hpp new file mode 100644 index 0000000000..dfbe1ede13 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_fa_pv.hpp @@ -0,0 +1,193 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_BLOCK_MMAD_FA_PV_HPP +#define CATLASS_KDA_BLOCK_MMAD_FA_PV_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_cross_core_sync.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/kda_coord.hpp" +#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" +#include "catlass/gemm/kda_gemm_helper.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" +#include "catlass/gemm/tile/kda_tile_mmad.hpp" + +//////////////////////////////////////////////////////////////////// + +namespace Catlass::Gemm::Block { +//////////////////////////////////////////////////////////////////// + +template < + class L1TileShape_, + class L0TileShape_, + class AType_, + class BType_, + class CType_, + class BiasType_, + class TileCopy_, + class TileMmad_> +struct BlockMmad< + MmadAtlasA2FAPV, + L1TileShape_, + L0TileShape_, + AType_, + BType_, + CType_, + BiasType_, + TileCopy_, + TileMmad_> { +public: + // Type Aliases + using DispatchPolicy = MmadAtlasA2FAPV; + using ArchTag = typename DispatchPolicy::ArchTag; + using L1TileShape = L1TileShape_; + using L0TileShape = L0TileShape_; + using ElementA = typename AType_::Element; + using LayoutA = typename AType_::Layout; + using ElementB = typename BType_::Element; + using LayoutB = typename BType_::Layout; + using ElementC = typename CType_::Element; + using LayoutC = typename CType_::Layout; + using TileMmad = TileMmad_; + using CopyGmToL1A = typename TileCopy_::CopyGmToL1A; + using CopyGmToL1B = typename TileCopy_::CopyGmToL1B; + using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; + using CopyL0CToGm = typename TileCopy_::CopyL0CToGm; + using ElementAccumulator = + typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; + using LayoutAInL1 = typename CopyL1ToL0A::LayoutSrc; + using LayoutBInL1 = typename CopyL1ToL0B::LayoutSrc; + using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst; + using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst; + using LayoutCInL0 = layout::zN; + + using L1AAlignHelper = Gemm::helper::L1AlignHelper; + using L1BAlignHelper = Gemm::helper::L1AlignHelper; + + static constexpr uint32_t STAGES = DispatchPolicy::STAGES; + static constexpr uint32_t L1A_SIZE = L1TileShape::M * L1TileShape::K * sizeof(ElementA); + static constexpr uint32_t L1B_SIZE = L1TileShape::N * L1TileShape::K * sizeof(ElementB); + static constexpr uint32_t L0A_SIZE = ArchTag::L0A_SIZE; + static constexpr uint32_t L0B_SIZE = ArchTag::L0B_SIZE; + static constexpr uint32_t L0C_SIZE = ArchTag::L0C_SIZE; + static constexpr uint32_t L0A_PINGPONG_BUF_SIZE = L0A_SIZE / STAGES; + static constexpr uint32_t L0B_PINGPONG_BUF_SIZE = L0B_SIZE / STAGES; + static constexpr uint32_t L0C_PINGPONG_BUF_SIZE = L0C_SIZE / STAGES; + + // Check LayoutC + static_assert(std::is_same_v, "LayoutC only support RowMajor yet!"); + + /// Construct + CATLASS_DEVICE + BlockMmad(Arch::Resource &resource, uint32_t l1BufAddrStart = 0) + { + for (uint32_t i = 0; i < STAGES; i++) { + l1ATensor[i] = resource.l1Buf.template GetBufferByByte(l1BufAddrStart + L1A_SIZE * i); + l1BTensor[i] = + resource.l1Buf.template GetBufferByByte(l1BufAddrStart + L1A_SIZE * STAGES + L1B_SIZE * i); + l0ATensor[i] = resource.l0ABuf.template GetBufferByByte(L0A_PINGPONG_BUF_SIZE * i); + l0BTensor[i] = resource.l0BBuf.template GetBufferByByte(L0B_PINGPONG_BUF_SIZE * i); + l0CTensor[i] = resource.l0CBuf.template GetBufferByByte(L0C_PINGPONG_BUF_SIZE * i); + } + + AscendC::SetFlag(EVENT_ID2); + AscendC::SetFlag(EVENT_ID3); + } + + /// Destructor + CATLASS_DEVICE + ~BlockMmad() + { + AscendC::WaitFlag(EVENT_ID2); + AscendC::WaitFlag(EVENT_ID3); + } + + /// Perform a block-scoped matrix multiply-accumulate + CATLASS_DEVICE + void operator()( + AscendC::GlobalTensor gA, + AscendC::GlobalTensor gB, + AscendC::GlobalTensor gC, + LayoutA layoutA, LayoutB layoutB, LayoutC layoutC, + GemmCoord actualShape, uint32_t &pingpongFlag, + Arch::CrossCoreFlag softmaxReady) + { + LayoutAInL1 layoutAInL1 = LayoutAInL1::template MakeLayout(L1TileShape::M, L1TileShape::K); + LayoutBInL1 layoutBInL1 = LayoutBInL1::template MakeLayout(L1TileShape::K, L1TileShape::N); + + uint32_t mRound = RoundUp(actualShape.m()); + uint32_t nRound = RoundUp(actualShape.n()); + uint32_t kRound = RoundUp(actualShape.k()); + + constexpr uint32_t PINGPONG_FLAG_OFFSET = 2; + AscendC::WaitFlag(pingpongFlag + PINGPONG_FLAG_OFFSET); + auto layoutTileB = layoutB.GetTileLayout(MakeCoord(actualShape.k(), actualShape.n())); + copyGmToL1B(l1BTensor[pingpongFlag], gB, layoutBInL1, layoutTileB); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + + AscendC::WaitFlag(pingpongFlag); + LayoutBInL0 layoutBInL0 = LayoutBInL0::template MakeLayout(kRound, nRound); + copyL1ToL0B(l0BTensor[pingpongFlag], l1BTensor[pingpongFlag], layoutBInL0, layoutBInL1); + + Arch::CrossCoreWaitFlag(softmaxReady); + auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShape.m(), actualShape.k())); + copyGmToL1A(l1ATensor[pingpongFlag], gA, layoutAInL1, layoutTileA); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + LayoutAInL0 layoutAInL0 = LayoutAInL0::template MakeLayout(mRound, kRound); + copyL1ToL0A(l0ATensor[pingpongFlag], l1ATensor[pingpongFlag], layoutAInL0, layoutAInL1); + AscendC::SetFlag(pingpongFlag + PINGPONG_FLAG_OFFSET); + + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + + AscendC::WaitFlag(pingpongFlag); + tileMmad( + l0CTensor[pingpongFlag], l0ATensor[pingpongFlag], l0BTensor[pingpongFlag], mRound, nRound, actualShape.k()); + AscendC::SetFlag(pingpongFlag); + + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + + // copy block out + auto blockShape = MakeCoord(actualShape.m(), actualShape.n()); + auto layoutInL0C = LayoutCInL0::MakeLayoutInL0C(blockShape); + + copyL0CToGm(gC, l0CTensor[pingpongFlag], layoutC, layoutInL0C); + AscendC::SetFlag(pingpongFlag); + + pingpongFlag = 1 - pingpongFlag; + } + +protected: + /// Data members + AscendC::LocalTensor l1ATensor[STAGES]; + AscendC::LocalTensor l1BTensor[STAGES]; + AscendC::LocalTensor l0ATensor[STAGES]; + AscendC::LocalTensor l0BTensor[STAGES]; + AscendC::LocalTensor l0CTensor[STAGES]; + + TileMmad tileMmad; + CopyGmToL1A copyGmToL1A; + CopyGmToL1B copyGmToL1B; + CopyL1ToL0A copyL1ToL0A; + CopyL1ToL0B copyL1ToL0B; + CopyL0CToGm copyL0CToGm; +}; + +//////////////////////////////////////////////////////////////////// + +} // namespace Catlass::Gemm::Block + +#endif // CATLASS_KDA_BLOCK_MMAD_FA_PV_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_fa_qk.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_fa_qk.hpp new file mode 100644 index 0000000000..e58ceee9ff --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_fa_qk.hpp @@ -0,0 +1,199 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_BLOCK_MMAD_FA_QK_HPP +#define CATLASS_KDA_BLOCK_MMAD_FA_QK_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/kda_coord.hpp" +#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" +#include "catlass/gemm/kda_gemm_helper.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" +#include "catlass/gemm/tile/kda_tile_mmad.hpp" + +//////////////////////////////////////////////////////////////////// + +namespace Catlass::Gemm::Block { +//////////////////////////////////////////////////////////////////// + +template < + class L1TileShape_, + class L0TileShape_, + class AType_, + class BType_, + class CType_, + class BiasType_, + class TileCopy_, + class TileMmad_> +struct BlockMmad< + MmadAtlasA2FAQK, + L1TileShape_, + L0TileShape_, + AType_, + BType_, + CType_, + BiasType_, + TileCopy_, + TileMmad_> { +public: + // Type Aliases + using DispatchPolicy = MmadAtlasA2FAQK; + using ArchTag = typename DispatchPolicy::ArchTag; + using L1TileShape = L1TileShape_; + using L0TileShape = L0TileShape_; + using ElementA = typename AType_::Element; + using LayoutA = typename AType_::Layout; + using ElementB = typename BType_::Element; + using LayoutB = typename BType_::Layout; + using ElementC = typename CType_::Element; + using LayoutC = typename CType_::Layout; + using TileMmad = TileMmad_; + using CopyGmToL1A = typename TileCopy_::CopyGmToL1A; + using CopyGmToL1B = typename TileCopy_::CopyGmToL1B; + using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; + using CopyL0CToGm = typename TileCopy_::CopyL0CToGm; + using ElementAccumulator = + typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; + using LayoutAInL1 = typename CopyL1ToL0A::LayoutSrc; + using LayoutBInL1 = typename CopyL1ToL0B::LayoutSrc; + using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst; + using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst; + using LayoutCInL0 = layout::zN; + + using L1AAlignHelper = Gemm::helper::L1AlignHelper; + using L1BAlignHelper = Gemm::helper::L1AlignHelper; + + static constexpr uint32_t STAGES = DispatchPolicy::STAGES; + static constexpr uint32_t L1A_SIZE = L1TileShape::M * L1TileShape::K * sizeof(ElementA); + static constexpr uint32_t L1B_SIZE = L1TileShape::N * L1TileShape::K * sizeof(ElementB); + static constexpr uint32_t L0A_SIZE = ArchTag::L0A_SIZE; + static constexpr uint32_t L0B_SIZE = ArchTag::L0B_SIZE; + static constexpr uint32_t L0C_SIZE = ArchTag::L0C_SIZE; + static constexpr uint32_t L0A_PINGPONG_BUF_SIZE = L0A_SIZE / STAGES; + static constexpr uint32_t L0B_PINGPONG_BUF_SIZE = L0B_SIZE / STAGES; + static constexpr uint32_t L0C_PINGPONG_BUF_SIZE = L0C_SIZE / STAGES; + + // Check LayoutC + static_assert(std::is_same_v, "LayoutC only support RowMajor yet!"); + + /// Construct + CATLASS_DEVICE + BlockMmad(Arch::Resource &resource, uint32_t l1BufAddrStart = 0) + { + for (uint32_t i = 0; i < STAGES; i++) { + l1ATensor[i] = resource.l1Buf.template GetBufferByByte(l1BufAddrStart + L1A_SIZE * i); + l1BTensor[i] = + resource.l1Buf.template GetBufferByByte(l1BufAddrStart + L1A_SIZE * STAGES + L1B_SIZE * i); + l0ATensor[i] = resource.l0ABuf.template GetBufferByByte(L0A_PINGPONG_BUF_SIZE * i); + l0BTensor[i] = resource.l0BBuf.template GetBufferByByte(L0B_PINGPONG_BUF_SIZE * i); + l0CTensor[i] = resource.l0CBuf.template GetBufferByByte(L0C_PINGPONG_BUF_SIZE * i); + } + + AscendC::SetFlag(EVENT_ID0); + AscendC::SetFlag(EVENT_ID1); + AscendC::SetFlag(EVENT_ID0); + AscendC::SetFlag(EVENT_ID1); + AscendC::SetFlag(EVENT_ID0); + AscendC::SetFlag(EVENT_ID1); + } + + /// Destructor + CATLASS_DEVICE + ~BlockMmad() + { + AscendC::WaitFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID1); + AscendC::WaitFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID1); + AscendC::WaitFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID1); + } + + /// Perform a block-scoped matrix multiply-accumulate + CATLASS_DEVICE + void operator()( + AscendC::GlobalTensor gA, + AscendC::GlobalTensor gB, + AscendC::GlobalTensor gC, + LayoutA layoutA, LayoutB layoutB, LayoutC layoutC, + GemmCoord actualShape, uint32_t &pingpongFlag, bool isFirst) + { + LayoutAInL1 layoutAInL1 = LayoutAInL1::template MakeLayout(L1TileShape::M, L1TileShape::K); + LayoutBInL1 layoutBInL1 = LayoutBInL1::template MakeLayout(L1TileShape::K, L1TileShape::N); + + uint32_t mRound = RoundUp(actualShape.m()); + uint32_t nRound = RoundUp(actualShape.n()); + uint32_t kRound = RoundUp(actualShape.k()); + + if (isFirst) { + auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShape.m(), actualShape.k())); + copyGmToL1A(l1ATensor[0], gA, layoutAInL1, layoutTileA); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + } + + AscendC::WaitFlag(pingpongFlag); + LayoutAInL0 layoutAInL0 = LayoutAInL0::template MakeLayout(mRound, kRound); + copyL1ToL0A(l0ATensor[pingpongFlag], l1ATensor[0], layoutAInL0, layoutAInL1); + + AscendC::WaitFlag(pingpongFlag); + auto layoutTileB = layoutB.GetTileLayout(MakeCoord(actualShape.k(), actualShape.n())); + copyGmToL1B(l1BTensor[pingpongFlag], gB, layoutBInL1, layoutTileB); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + LayoutBInL0 layoutBInL0 = LayoutBInL0::template MakeLayout(kRound, nRound); + copyL1ToL0B(l0BTensor[pingpongFlag], l1BTensor[pingpongFlag], layoutBInL0, layoutBInL1); + AscendC::SetFlag(pingpongFlag); + + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + + AscendC::WaitFlag(pingpongFlag); + tileMmad( + l0CTensor[pingpongFlag], l0ATensor[pingpongFlag], l0BTensor[pingpongFlag], mRound, nRound, actualShape.k()); + AscendC::SetFlag(pingpongFlag); + + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + + // copy block out + auto blockShape = MakeCoord(actualShape.m(), actualShape.n()); + auto layoutInL0C = LayoutCInL0::MakeLayoutInL0C(blockShape); + + copyL0CToGm(gC, l0CTensor[pingpongFlag], layoutC, layoutInL0C); + AscendC::SetFlag(pingpongFlag); + + pingpongFlag = 1 - pingpongFlag; + } + +protected: + /// Data members + AscendC::LocalTensor l1ATensor[STAGES]; + AscendC::LocalTensor l1BTensor[STAGES]; + AscendC::LocalTensor l0ATensor[STAGES]; + AscendC::LocalTensor l0BTensor[STAGES]; + AscendC::LocalTensor l0CTensor[STAGES]; + + TileMmad tileMmad; + CopyGmToL1A copyGmToL1A; + CopyGmToL1B copyGmToL1B; + CopyL1ToL0A copyL1ToL0A; + CopyL1ToL0B copyL1ToL0B; + CopyL0CToGm copyL0CToGm; +}; + +//////////////////////////////////////////////////////////////////// + +} // namespace Catlass::Gemm::Block + +#endif // CATLASS_KDA_BLOCK_MMAD_FA_QK_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_gemm.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_gemm.hpp new file mode 100644 index 0000000000..8f97aa18ea --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_gemm.hpp @@ -0,0 +1,348 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_BLOCK_MMAD_GEMM_HPP +#define CATLASS_KDA_BLOCK_MMAD_GEMM_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/gemm/kda_gemm_helper.hpp" +#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" +#include "catlass/gemm/tile/kda_tile_mmad.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" +#include "catlass/arch/kda_resource.hpp" + +namespace Catlass::Gemm::Block { + +template < + bool ENABLE_UNIT_FLAG_, + bool ENABLE_SHUFFLE_K_, + bool ENABLE_ABBA_, + class L1TileShape_, + class L0TileShape_, + class AType_, + class BType_, + class CType_, + class BiasType_, + class TileCopy_, + class TileMmad_> +struct BlockGemm< + Gemm::GemmAtlasA2, + L1TileShape_, + L0TileShape_, + AType_, + BType_, + CType_, + BiasType_, + TileCopy_, + TileMmad_> { +public: + using DispatchPolicy = Gemm::GemmAtlasA2; + using ArchTag = typename DispatchPolicy::ArchTag; + using L1TileShape = L1TileShape_; + using L0TileShape = L0TileShape_; + using ElementA = typename AType_::Element; + using LayoutA = typename AType_::Layout; + using ElementB = typename BType_::Element; + using LayoutB = typename BType_::Layout; + using ElementC = typename CType_::Element; + using LayoutC = typename CType_::Layout; + + using TileMmad = TileMmad_; + using CopyGmToL1A = typename TileCopy_::CopyGmToL1A; + using CopyGmToL1B = typename TileCopy_::CopyGmToL1B; + using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; + using CopyL0CToGm = typename TileCopy_::CopyL0CToGm; + using ElementAccumulator = + typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; + using LayoutAInL1 = typename CopyGmToL1A::LayoutDst; + using LayoutBInL1 = typename CopyGmToL1B::LayoutDst; + using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst; + using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst; + using LayoutCInL0 = layout::zN; + + static constexpr uint32_t STAGES = DispatchPolicy::STAGES; + static constexpr bool ENABLE_UNIT_FLAG = DispatchPolicy::ENABLE_UNIT_FLAG; + static constexpr bool ENABLE_SHUFFLE_K = DispatchPolicy::ENABLE_SHUFFLE_K; + static constexpr bool ENABLE_ABBA = DispatchPolicy::ENABLE_ABBA; + const uint32_t L1Size = ArchTag::L1_SIZE; + const uint32_t L1ASize = L1TileShape::M * L1TileShape::K * sizeof(ElementA); + const uint32_t L1BSize = L1TileShape::K * L1TileShape::N * sizeof(ElementB); + const uint32_t cSize = L1TileShape::M * L1TileShape::N * sizeof(ElementAccumulator); + const uint32_t BlockCnt = L1TileShape::M * L1TileShape::N; + const uint32_t L0ASize = ArchTag::L0A_SIZE; + const uint32_t L0BSize = ArchTag::L0B_SIZE; + const uint32_t L0CSize = ArchTag::L0C_SIZE; + const uint32_t L0A_PINGPONG_BUF_LEN = (L0ASize / STAGES); + const uint32_t L0B_PINGPONG_BUF_LEN = (L0BSize / STAGES); + const uint32_t l0CBlockNum = ArchTag::L0C_SIZE / cSize; + + static_assert(std::is_same_v, "LayoutC only support RowMajor yet!"); + + // 32B (256b) aligned + static_assert(Gemm::helper::TileShapeAlignChecker::_ALIGN == 256, + "Tile shape must be 32B aligned."); + + CATLASS_DEVICE + BlockGemm(Arch::Resource &resource, uint32_t l1BufAddrStart = 0) + { + uint32_t l1AOffset = l1BufAddrStart; + uint32_t l1BOffset = l1BufAddrStart + L1ASize * STAGES; + for (uint32_t i = 0; i < STAGES; i++) { + l1ATensor[i] = resource.l1Buf.template GetBufferByByte(l1AOffset + L1ASize * i); + l1BTensor[i] = resource.l1Buf.template GetBufferByByte(l1BOffset + L1BSize * i); + l0ATensor[i] = resource.l0ABuf.template GetBufferByByte(L0A_PINGPONG_BUF_LEN * i); + l0BTensor[i] = resource.l0BBuf.template GetBufferByByte(L0B_PINGPONG_BUF_LEN * i); + // Assign event ID for each stages + l1AEventList[i] = i; + l1BEventList[i] = i + STAGES; + l0AEventList[i] = i; + l0BEventList[i] = i + STAGES; + // The event id that needs to be set before the loop + AscendC::SetFlag(l1AEventList[i]); + AscendC::SetFlag(l1BEventList[i]); + AscendC::SetFlag(l0AEventList[i]); + AscendC::SetFlag(l0BEventList[i]); + } + l0CTensor = resource.l0CBuf.template GetBufferByByte(0); + } + // destroy function + CATLASS_DEVICE + ~BlockGemm() + { + for (uint32_t i = 0; i < STAGES; i++) { + AscendC::WaitFlag(l1AEventList[i]); + AscendC::WaitFlag(l1BEventList[i]); + AscendC::WaitFlag(l0AEventList[i]); + AscendC::WaitFlag(l0BEventList[i]); + } + } + + CATLASS_DEVICE + void operator()( + AscendC::GlobalTensor const &gmA, LayoutA const &layoutA, + AscendC::GlobalTensor const &gmB, LayoutB const &layoutB, + AscendC::GlobalTensor const &gmC, LayoutC const &layoutC, + AscendC::GlobalTensor const &gmNextBlockA, + AscendC::GlobalTensor const &gmNextBlockB, + GemmCoord const &actualShape, GemmCoord const &actualShapeNext, + bool isFirstBlock, bool hasNextBlock, uint32_t singleIdx) + { + uint32_t K = actualShape.k(); + uint32_t maxKPerBlock = L1TileShape::K; + uint32_t kLoops = CeilDiv(K, maxKPerBlock); + uint32_t kLoopsNext = CeilDiv(actualShapeNext.k(), maxKPerBlock); + uint32_t startTileIdx{0}; + if (ENABLE_SHUFFLE_K) { + startTileIdx = AscendC::GetBlockIdx(); + } + uint32_t firstTileIdx = startTileIdx % kLoops; + uint32_t firstTileIdxNext = startTileIdx % kLoopsNext; + uint32_t lastTileIdx = (startTileIdx + kLoops - 1) % kLoops; + uint32_t kGmActual = (firstTileIdx == kLoops - 1) ? (K - firstTileIdx * maxKPerBlock) : maxKPerBlock; + auto layoutAInL1 = LayoutAInL1::template MakeLayout(L1TileShape::M, L1TileShape::K); + auto layoutBInL1 = LayoutBInL1::template MakeLayout(L1TileShape::K, L1TileShape::N); + for (uint32_t kIdx = 0; kIdx < kLoops; kIdx++) { + uint32_t shuffleKIdx = (startTileIdx + kIdx) % kLoops; + if (shuffleKIdx == firstTileIdx && isFirstBlock) { + auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShape.m(), kGmActual)); + auto layoutTileB = layoutB.GetTileLayout(MakeCoord(kGmActual, actualShape.n())); + MatrixCoord gmTileAOffset{0, shuffleKIdx * maxKPerBlock}; + auto gmTileA = gmA[layoutA.GetOffset(gmTileAOffset)]; + MatrixCoord gmTileBOffset{shuffleKIdx * maxKPerBlock, 0}; + auto gmTileB = gmB[layoutB.GetOffset(gmTileBOffset)]; + AscendC::WaitFlag(l1AEventList[l1ListId]); + copyGmToL1A(l1ATensor[l1ListId], gmTileA, layoutAInL1, layoutTileA); + AscendC::SetFlag(l1AEventList[l1ListId]); + AscendC::WaitFlag(l1BEventList[l1ListId]); + copyGmToL1B(l1BTensor[l1ListId], gmTileB, layoutBInL1, layoutTileB); + AscendC::SetFlag(l1BEventList[l1ListId]); + } + l1ListIdNext = 1 - l1ListId; + uint32_t kGmActualNext = 0; + if (shuffleKIdx != lastTileIdx) { + uint32_t shuffleKIdxNext = (startTileIdx + kIdx + 1) % kLoops; + kGmActualNext = (shuffleKIdxNext == kLoops - 1) ? (K - shuffleKIdxNext * maxKPerBlock) : maxKPerBlock; + auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShape.m(), kGmActualNext)); + auto layoutTileB = layoutB.GetTileLayout(MakeCoord(kGmActualNext, actualShape.n())); + MatrixCoord gmTileAOffset{0, shuffleKIdxNext * maxKPerBlock}; + auto gmTileA = gmA[layoutA.GetOffset(gmTileAOffset)]; + MatrixCoord gmTileBOffset{shuffleKIdxNext * maxKPerBlock, 0}; + auto gmTileB = gmB[layoutB.GetOffset(gmTileBOffset)]; + if (ENABLE_ABBA) { + if (shuffleKIdxNext % 2 == 1) { + AscendC::WaitFlag(l1BEventList[l1ListIdNext]); + copyGmToL1B(l1BTensor[l1ListIdNext], gmTileB, layoutBInL1, layoutTileB); + AscendC::SetFlag(l1BEventList[l1ListIdNext]); + AscendC::WaitFlag(l1AEventList[l1ListIdNext]); + copyGmToL1A(l1ATensor[l1ListIdNext], gmTileA, layoutAInL1, layoutTileA); + AscendC::SetFlag(l1AEventList[l1ListIdNext]); + } else { + AscendC::WaitFlag(l1AEventList[l1ListIdNext]); + copyGmToL1A(l1ATensor[l1ListIdNext], gmTileA, layoutAInL1, layoutTileA); + AscendC::SetFlag(l1AEventList[l1ListIdNext]); + AscendC::WaitFlag(l1BEventList[l1ListIdNext]); + copyGmToL1B(l1BTensor[l1ListIdNext], gmTileB, layoutBInL1, layoutTileB); + AscendC::SetFlag(l1BEventList[l1ListIdNext]); + } + } else { + AscendC::WaitFlag(l1AEventList[l1ListIdNext]); + copyGmToL1A(l1ATensor[l1ListIdNext], gmTileA, layoutAInL1, layoutTileA); + AscendC::SetFlag(l1AEventList[l1ListIdNext]); + AscendC::WaitFlag(l1BEventList[l1ListIdNext]); + copyGmToL1B(l1BTensor[l1ListIdNext], gmTileB, layoutBInL1, layoutTileB); + AscendC::SetFlag(l1BEventList[l1ListIdNext]); + } + } + if (shuffleKIdx == lastTileIdx && hasNextBlock) { + kGmActualNext = (firstTileIdxNext == kLoopsNext - 1) ? (actualShapeNext.k() - firstTileIdxNext * maxKPerBlock) : maxKPerBlock; + auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShapeNext.m(), kGmActualNext)); + auto layoutTileB = layoutB.GetTileLayout(MakeCoord(kGmActualNext, actualShapeNext.n())); + MatrixCoord gmTileAOffset{0, firstTileIdxNext * maxKPerBlock}; + auto gmNextTileA = gmNextBlockA[layoutA.GetOffset(gmTileAOffset)]; + MatrixCoord gmTileBOffset{firstTileIdxNext * maxKPerBlock, 0}; + auto gmNextTileB = gmNextBlockB[layoutB.GetOffset(gmTileBOffset)]; + if (ENABLE_ABBA) { + if (shuffleKIdx % 2 == 0) { + AscendC::WaitFlag(l1BEventList[l1ListIdNext]); + copyGmToL1B(l1BTensor[l1ListIdNext], gmNextTileB, layoutBInL1, layoutTileB); + AscendC::SetFlag(l1BEventList[l1ListIdNext]); + AscendC::WaitFlag(l1AEventList[l1ListIdNext]); + copyGmToL1A(l1ATensor[l1ListIdNext], gmNextTileA, layoutAInL1, layoutTileA); + AscendC::SetFlag(l1AEventList[l1ListIdNext]); + } else { + AscendC::WaitFlag(l1AEventList[l1ListIdNext]); + copyGmToL1A(l1ATensor[l1ListIdNext], gmNextTileA, layoutAInL1, layoutTileA); + AscendC::SetFlag(l1AEventList[l1ListIdNext]); + AscendC::WaitFlag(l1BEventList[l1ListIdNext]); + copyGmToL1B(l1BTensor[l1ListIdNext], gmNextTileB, layoutBInL1, layoutTileB); + AscendC::SetFlag(l1BEventList[l1ListIdNext]); + } + } else { + AscendC::WaitFlag(l1AEventList[l1ListIdNext]); + copyGmToL1A(l1ATensor[l1ListIdNext], gmNextTileA, layoutAInL1, layoutTileA); + AscendC::SetFlag(l1AEventList[l1ListIdNext]); + AscendC::WaitFlag(l1BEventList[l1ListIdNext]); + copyGmToL1B(l1BTensor[l1ListIdNext], gmNextTileB, layoutBInL1, layoutTileB); + AscendC::SetFlag(l1BEventList[l1ListIdNext]); + } + } + + uint32_t kL0TileSize = L0TileShape::K; + uint32_t kL0Loops = CeilDiv(kGmActual, kL0TileSize); + AscendC::WaitFlag(l1AEventList[l1ListId]); + AscendC::WaitFlag(l1BEventList[l1ListId]); + auto l1ATile = l1ATensor[l1ListId]; + auto l1BTile = l1BTensor[l1ListId]; + uint32_t mActual{0}; + uint32_t nActual{0}; + for (uint32_t kL0Idx = 0; kL0Idx < kL0Loops; kL0Idx++) { + uint32_t kL0Actual = (kL0Idx == kL0Loops - 1) ? (kGmActual - kL0Idx * kL0TileSize) : kL0TileSize; + LayoutAInL0 layoutAInL0 = LayoutAInL0::template MakeLayout(L1TileShape::M, kL0Actual); + LayoutBInL0 layoutBInL0 = LayoutBInL0::template MakeLayout(kL0Actual, L1TileShape::N); + uint32_t l1TileAOffset = layoutAInL1.GetOffset(MatrixCoord(0, kL0Idx * kL0TileSize)); + uint32_t l1TileBOffset = layoutBInL1.GetOffset(MatrixCoord(kL0Idx * kL0TileSize, 0)); + auto l1TileA = l1ATile[l1TileAOffset]; + auto l1TileB = l1BTile[l1TileBOffset]; + auto l0TileA = l0ATensor[l0ListId]; + auto l0TileB = l0BTensor[l0ListId]; + mActual = L1TileShape::M; + nActual = L1TileShape::N; + if (ENABLE_ABBA) { + if (shuffleKIdx % 2 == 0) { + if (kL0Idx % 2 == 0) { + AscendC::WaitFlag(l0BEventList[l0ListId]); + copyL1ToL0B(l0TileB, l1TileB, layoutBInL0, layoutBInL1); + AscendC::SetFlag(l0BEventList[l0ListId]); + AscendC::WaitFlag(l0AEventList[l0ListId]); + copyL1ToL0A(l0TileA, l1TileA, layoutAInL0, layoutAInL1); + AscendC::SetFlag(l0AEventList[l0ListId]); + } else { + AscendC::WaitFlag(l0AEventList[l0ListId]); + copyL1ToL0A(l0TileA, l1TileA, layoutAInL0, layoutAInL1); + AscendC::SetFlag(l0AEventList[l0ListId]); + AscendC::WaitFlag(l0BEventList[l0ListId]); + copyL1ToL0B(l0TileB, l1TileB, layoutBInL0, layoutBInL1); + AscendC::SetFlag(l0BEventList[l0ListId]); + } + } else { + if (kL0Idx % 2 == 0) { + AscendC::WaitFlag(l0AEventList[l0ListId]); + copyL1ToL0A(l0TileA, l1TileA, layoutAInL0, layoutAInL1); + AscendC::SetFlag(l0AEventList[l0ListId]); + AscendC::WaitFlag(l0BEventList[l0ListId]); + copyL1ToL0B(l0TileB, l1TileB, layoutBInL0, layoutBInL1); + AscendC::SetFlag(l0BEventList[l0ListId]); + } else { + AscendC::WaitFlag(l0BEventList[l0ListId]); + copyL1ToL0B(l0TileB, l1TileB, layoutBInL0, layoutBInL1); + AscendC::SetFlag(l0BEventList[l0ListId]); + AscendC::WaitFlag(l0AEventList[l0ListId]); + copyL1ToL0A(l0TileA, l1TileA, layoutAInL0, layoutAInL1); + AscendC::SetFlag(l0AEventList[l0ListId]); + } + } + } else { + AscendC::WaitFlag(l0AEventList[l0ListId]); + copyL1ToL0A(l0TileA, l1TileA, layoutAInL0, layoutAInL1); + AscendC::SetFlag(l0AEventList[l0ListId]); + AscendC::WaitFlag(l0BEventList[l0ListId]); + copyL1ToL0B(l0TileB, l1TileB, layoutBInL0, layoutBInL1); + AscendC::SetFlag(l0BEventList[l0ListId]); + } + if (kL0Idx == kL0Loops - 1) { + AscendC::SetFlag(l1AEventList[l1ListId]); + AscendC::SetFlag(l1BEventList[l1ListId]); + l1ListId = l1ListIdNext; + kGmActual = kGmActualNext; + } + AscendC::WaitFlag(l0BEventList[l0ListId]); + AscendC::WaitFlag(l0AEventList[l0ListId]); + tileMmad(l0CTensor[(singleIdx % l0CBlockNum) * BlockCnt], l0TileA, l0TileB, mActual, nActual, kL0Actual, (kIdx == 0) && (kL0Idx == 0)); + AscendC::SetFlag(l0AEventList[l0ListId]); + AscendC::SetFlag(l0BEventList[l0ListId]); + l0ListId = 1 - l0ListId; + } + } + AscendC::SetFlag((int32_t)(singleIdx % l0CBlockNum)); + AscendC::WaitFlag((int32_t)(singleIdx % l0CBlockNum)); + auto layoutInL0X = LayoutCInL0::MakeLayoutInL0C(MakeCoord(L1TileShape::M, L1TileShape::N)); + LayoutC layoutBlock = layoutC.GetTileLayout(MakeCoord(actualShape.m(), actualShape.n())); + copyL0CToGm(gmC, l0CTensor[(singleIdx % l0CBlockNum) * BlockCnt], layoutBlock, layoutInL0X); + } + +private: + AscendC::LocalTensor l1ATensor[STAGES]; + AscendC::LocalTensor l1BTensor[STAGES]; + AscendC::LocalTensor l0ATensor[STAGES]; + AscendC::LocalTensor l0BTensor[STAGES]; + AscendC::LocalTensor l0CTensor; + // Multi-stage event id list + int32_t l1AEventList[STAGES]; + int32_t l1BEventList[STAGES]; + int32_t l0AEventList[STAGES]; + int32_t l0BEventList[STAGES]; + + // The id of current stage + uint32_t l1ListId{0}; + uint32_t l0ListId{0}; + uint32_t l1ListIdNext{0}; + + TileMmad tileMmad; + CopyGmToL1A copyGmToL1A; + CopyGmToL1B copyGmToL1B; + CopyL1ToL0A copyL1ToL0A; + CopyL1ToL0B copyL1ToL0B; + CopyL0CToGm copyL0CToGm; +}; +} // namespace Catlass::Gemm::Block + +#endif // CATLASS_KDA_BLOCK_MMAD_GEMM_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_mla_pv.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_mla_pv.hpp new file mode 100644 index 0000000000..3d37211bc2 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_mla_pv.hpp @@ -0,0 +1,207 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_BLOCK_MMAD_MLA_PV_HPP +#define CATLASS_KDA_BLOCK_MMAD_MLA_PV_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_cross_core_sync.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/kda_coord.hpp" +#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" +#include "catlass/gemm/kda_gemm_helper.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" +#include "catlass/gemm/tile/kda_tile_mmad.hpp" + +//////////////////////////////////////////////////////////////////// + +namespace Catlass::Gemm::Block { +//////////////////////////////////////////////////////////////////// + +template < + class L1TileShape_, + class L0TileShape_, + class AType_, + class BType_, + class CType_, + class BiasType_, + class TileCopy_, + class TileMmad_> +struct BlockMmad< + MmadAtlasA2MLAPV, + L1TileShape_, + L0TileShape_, + AType_, + BType_, + CType_, + BiasType_, + TileCopy_, + TileMmad_> { +public: + // Type Aliases + using DispatchPolicy = MmadAtlasA2MLAPV; + using ArchTag = typename DispatchPolicy::ArchTag; + using L1TileShape = L1TileShape_; + using L0TileShape = L0TileShape_; + using ElementA = typename AType_::Element; + using LayoutA = typename AType_::Layout; + using ElementB = typename BType_::Element; + using LayoutB = typename BType_::Layout; + using ElementC = typename CType_::Element; + using LayoutC = typename CType_::Layout; + using TileMmad = TileMmad_; + using CopyGmToL1A = typename TileCopy_::CopyGmToL1A; + using CopyGmToL1B = typename TileCopy_::CopyGmToL1B; + using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; + using CopyL0CToGm = typename TileCopy_::CopyL0CToGm; + using ElementAccumulator = + typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; + using LayoutAInL1 = typename CopyL1ToL0A::LayoutSrc; + using LayoutBInL1 = typename CopyL1ToL0B::LayoutSrc; + using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst; + using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst; + using LayoutCInL0 = layout::zN; + + using L1AAlignHelper = Gemm::helper::L1AlignHelper; + using L1BAlignHelper = Gemm::helper::L1AlignHelper; + + static constexpr uint32_t STAGES = DispatchPolicy::STAGES; + static constexpr uint32_t L1A_SIZE = L1TileShape::M * L1TileShape::K * sizeof(ElementA); + static constexpr uint32_t L1B_SIZE = L1TileShape::N * L1TileShape::K * sizeof(ElementB); + static constexpr uint32_t L0A_SIZE = ArchTag::L0A_SIZE; + static constexpr uint32_t L0B_SIZE = ArchTag::L0B_SIZE; + static constexpr uint32_t L0C_SIZE = ArchTag::L0C_SIZE; + static constexpr uint32_t L0A_PINGPONG_BUF_SIZE = L0A_SIZE / STAGES; + static constexpr uint32_t L0B_PINGPONG_BUF_SIZE = L0B_SIZE / STAGES; + static constexpr uint32_t L0C_PINGPONG_BUF_SIZE = L0C_SIZE / STAGES; + static constexpr uint32_t BLOCK_SIZE = 16; + static constexpr uint32_t EMBED_SPLIT_SIZE = 128; + static constexpr uint32_t EMBED_SPLIT_LOOP = 4; + + // Check LayoutC + static_assert(std::is_same_v, "LayoutC only support RowMajor yet!"); + + /// Construct + CATLASS_DEVICE + BlockMmad(Arch::Resource &resource, uint32_t l1BufAddrStart = 0) + { + // Allocate L1 memory space + l1ATensor = + resource.l1Buf.template GetBufferByByte(l1BufAddrStart + L1A_SIZE + L1B_SIZE * STAGES); + for (uint32_t i = 0; i < STAGES; i++) { + l1BTensor[i] = + resource.l1Buf.template GetBufferByByte(l1BufAddrStart + L1A_SIZE + L1B_SIZE * i); + l0ATensor[i] = resource.l0ABuf.template GetBufferByByte(L0A_PINGPONG_BUF_SIZE * i); + l0BTensor[i] = resource.l0BBuf.template GetBufferByByte(L0B_PINGPONG_BUF_SIZE * i); + l0CTensor[i] = resource.l0CBuf.template GetBufferByByte(L0C_PINGPONG_BUF_SIZE * i); + } + } + + /// Destructor + CATLASS_DEVICE + ~BlockMmad() {} + + /// Perform a block-scoped matrix multiply-accumulate + CATLASS_DEVICE + void operator()( + AscendC::GlobalTensor gA, + AscendC::GlobalTensor gC, + LayoutA layoutA, LayoutB layoutB, LayoutC layoutC, + GemmCoord actualShape, uint32_t &nIdx, Arch::CrossCoreFlag softmaxReady) + { + uint32_t rowNum = actualShape.m(); + uint32_t vSeqTile = actualShape.k(); + uint32_t embed = actualShape.n(); + uint32_t embedSplitSize = EMBED_SPLIT_SIZE; + uint32_t embedSplitLoopV = EMBED_SPLIT_LOOP; + uint32_t rowNumRound = RoundUp(rowNum); + uint32_t embedSplitSizeRound = RoundUp(embedSplitSize); + uint32_t vSeqTileRound = RoundUp(vSeqTile); + uint32_t L1BPingPongFlag = (nIdx - 1) % 2; + uint32_t L0APingPongFlag = (nIdx - 1) % 2; + + for (uint32_t embedSplitIdx = 0; embedSplitIdx < embedSplitLoopV; embedSplitIdx++) { + uint32_t L0CPingPongFlag = (nIdx + embedSplitIdx) % 2; + uint32_t L0BPingPongFlag = (embedSplitIdx + 1) % 2; + AscendC::WaitFlag(L0BPingPongFlag + 2); + LayoutBInL1 layoutBInL1 = LayoutBInL1::template MakeLayout(vSeqTile, embed); + LayoutBInL0 layoutBInL0 = LayoutBInL0::template MakeLayout(vSeqTile, embedSplitSize); + // copy V from L1 to L0B + copyL1ToL0B( + l0BTensor[L0BPingPongFlag], + l1BTensor[L1BPingPongFlag][embedSplitIdx * vSeqTileRound * EMBED_SPLIT_SIZE], + layoutBInL0, layoutBInL1); + if (embedSplitIdx == embedSplitLoopV - 1) { + AscendC::SetFlag(L1BPingPongFlag); + } + + if (embedSplitIdx == 0) { + LayoutAInL1 layoutAInL1 = LayoutAInL1::template MakeLayout(rowNum, vSeqTile); + Arch::CrossCoreWaitFlag(softmaxReady); + AscendC::WaitFlag(EVENT_ID7); + // copy P to L1 + copyGmToL1A(l1ATensor, gA, layoutAInL1, layoutA); + AscendC::SetFlag(EVENT_ID7); + AscendC::WaitFlag(EVENT_ID7); + // move p from l1 to l0a + AscendC::WaitFlag(L0APingPongFlag); + LayoutAInL0 layoutAInL0 = LayoutAInL0::template MakeLayout(rowNum, vSeqTile); + copyL1ToL0A(l0ATensor[L0APingPongFlag], l1ATensor, layoutAInL0, layoutAInL1); + AscendC::SetFlag(EVENT_ID7); + } + + AscendC::SetFlag(L0BPingPongFlag); + AscendC::WaitFlag(L0BPingPongFlag); + AscendC::WaitFlag(L0CPingPongFlag); + // mmad + tileMmad( + l0CTensor[L0CPingPongFlag], l0ATensor[L0APingPongFlag], l0BTensor[L0BPingPongFlag], + rowNumRound, embedSplitSize, vSeqTile); + AscendC::SetFlag(L0BPingPongFlag + 2); + if (embedSplitIdx == embedSplitLoopV - 1) { + AscendC::SetFlag(L0APingPongFlag); + } + AscendC::SetFlag(L0CPingPongFlag); + + AscendC::WaitFlag(L0CPingPongFlag); + auto blockShape = MakeCoord(rowNum, embedSplitSize); + auto layoutInL0C = LayoutCInL0::MakeLayoutInL0C(blockShape); + auto layoutCSplitK = layoutC.GetTileLayout(MakeCoord(rowNumRound, embedSplitSizeRound)); + // copy Otmp to gm + copyL0CToGm( + gC[embedSplitIdx * embedSplitSizeRound], l0CTensor[L0CPingPongFlag], + layoutCSplitK, layoutInL0C); + AscendC::SetFlag(L0CPingPongFlag); + } + } + +protected: + /// Data members + AscendC::LocalTensor l1ATensor; + AscendC::LocalTensor l1BTensor[STAGES]; + AscendC::LocalTensor l0ATensor[STAGES]; + AscendC::LocalTensor l0BTensor[STAGES]; + AscendC::LocalTensor l0CTensor[STAGES]; + + TileMmad tileMmad; + CopyGmToL1A copyGmToL1A; + CopyGmToL1B copyGmToL1B; + CopyL1ToL0A copyL1ToL0A; + CopyL1ToL0B copyL1ToL0B; + CopyL0CToGm copyL0CToGm; +}; + +//////////////////////////////////////////////////////////////////// + +} // namespace Catlass::Gemm::Block + +#endif // CATLASS_KDA_BLOCK_MMAD_MLA_PV_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_mla_pv_tp1_spec.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_mla_pv_tp1_spec.hpp new file mode 100644 index 0000000000..36a8ba0188 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_mla_pv_tp1_spec.hpp @@ -0,0 +1,221 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_BLOCK_MMAD_MLA_PV_TP1_SPEC_HPP +#define CATLASS_KDA_BLOCK_MMAD_MLA_PV_TP1_SPEC_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_cross_core_sync.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/kda_coord.hpp" +#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" +#include "catlass/gemm/kda_gemm_helper.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" +#include "catlass/gemm/tile/kda_tile_mmad.hpp" + +//////////////////////////////////////////////////////////////////// + +namespace Catlass::Gemm::Block { +//////////////////////////////////////////////////////////////////// + +template < + class L1TileShape_, + class L0TileShape_, + class AType_, + class BType_, + class CType_, + class BiasType_, + class TileCopy_, + class TileMmad_> +struct BlockMmad< + MmadAtlasA2MLAPVTp1Spec, + L1TileShape_, + L0TileShape_, + AType_, + BType_, + CType_, + BiasType_, + TileCopy_, + TileMmad_> { +public: + // Type Aliases + using DispatchPolicy = MmadAtlasA2MLAPV; + using ArchTag = typename DispatchPolicy::ArchTag; + using L1TileShape = L1TileShape_; + using L0TileShape = L0TileShape_; + using ElementA = typename AType_::Element; + using LayoutA = typename AType_::Layout; + using ElementB = typename BType_::Element; + using LayoutB = typename BType_::Layout; + using ElementC = typename CType_::Element; + using LayoutC = typename CType_::Layout; + using TileMmad = TileMmad_; + using CopyGmToL1A = typename TileCopy_::CopyGmToL1A; + using CopyGmToL1B = typename TileCopy_::CopyGmToL1B; + using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; + using CopyL0CToGm = typename TileCopy_::CopyL0CToGm; + using ElementAccumulator = + typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; + using LayoutAInL1 = typename CopyL1ToL0A::LayoutSrc; + using LayoutBInL1 = typename CopyL1ToL0B::LayoutSrc; + using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst; + using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst; + using LayoutCInL0 = layout::zN; + + using L1AAlignHelper = Gemm::helper::L1AlignHelper; + using L1BAlignHelper = Gemm::helper::L1AlignHelper; + + static constexpr uint32_t STAGES = DispatchPolicy::STAGES; + static constexpr uint32_t L0A_SIZE = ArchTag::L0A_SIZE; + static constexpr uint32_t L0B_SIZE = ArchTag::L0B_SIZE; + static constexpr uint32_t L0C_SIZE = ArchTag::L0C_SIZE; + static constexpr uint32_t L0A_PINGPONG_BUF_SIZE = L0A_SIZE / STAGES; + static constexpr uint32_t L0B_PINGPONG_BUF_SIZE = L0B_SIZE / STAGES; + static constexpr uint32_t L0C_PINGPONG_BUF_SIZE = L0C_SIZE / STAGES; + static constexpr uint32_t UNIT_BLOCK_STACK_NUM = 4; + static constexpr uint32_t EMBED_SPLIT_LOOP = 4; + static constexpr uint32_t BLOCK_SIZE = 16; + static constexpr uint32_t EMBED_SPLIT_SIZE = 128; + static constexpr uint32_t L1_PV_ADDR_START = 311296; // reserved for Q(no db) and K(db) + static constexpr uint32_t L1A_SIZE = L1TileShape::M * L1TileShape::N * sizeof(ElementA); + static constexpr uint32_t L1B_SIZE = L1TileShape::N * EMBED_SPLIT_SIZE * sizeof(ElementB); + + /// Construct + CATLASS_DEVICE + BlockMmad(Arch::Resource &resource, uint32_t l1BufAddrStart = L1_PV_ADDR_START) + { + // Allocate L1 memory space + for (uint32_t i = 0; i < STAGES; i++) { + l1ATensor[i] = + resource.l1Buf.template GetBufferByByte(l1BufAddrStart + L1B_SIZE * STAGES + L1A_SIZE * i); + l1BTensor[i] = resource.l1Buf.template GetBufferByByte(l1BufAddrStart + L1B_SIZE * i); + l0ATensor[i] = resource.l0ABuf.template GetBufferByByte(L0A_PINGPONG_BUF_SIZE * i); + l0BTensor[i] = resource.l0BBuf.template GetBufferByByte(L0B_PINGPONG_BUF_SIZE * i); + l0CTensor[i] = resource.l0CBuf.template GetBufferByByte(L0C_PINGPONG_BUF_SIZE * i); + } + } + + /// Destructor + CATLASS_DEVICE + ~BlockMmad() {} + + /// Perform a block-scoped matrix multiply-accumulate + CATLASS_DEVICE + void operator()(AscendC::GlobalTensor gA, AscendC::GlobalTensor gB, + AscendC::GlobalTensor gblockTable, AscendC::GlobalTensor gC, LayoutA layoutA, + LayoutB layoutB, LayoutC layoutC, GemmCoord actualShape, uint32_t &nIdx, uint32_t &nLoop, + uint32_t &blockSize, uint32_t kvSeqlen, Arch::CrossCoreFlag softmaxReady) + { + uint32_t rowNum = actualShape.m(); + uint32_t stackSeqTile = actualShape.k(); + uint32_t seqTile = blockSize; + uint32_t embed = actualShape.n(); + uint32_t embedSplitSize = EMBED_SPLIT_SIZE; + uint32_t embedSplitLoopV = EMBED_SPLIT_LOOP; + uint32_t rowNumRound = RoundUp(rowNum); + uint32_t embedSplitSizeRound = RoundUp(embedSplitSize); + uint32_t stackSeqTileRound = layoutA.stride(0); + uint32_t seqTileRound = RoundUp(seqTile); + + for (uint32_t embedSplitIdx = 0; embedSplitIdx < embedSplitLoopV; embedSplitIdx++) { + uint32_t L0CPingPongFlag = embedSplitIdx % 2; + for (uint32_t blockStackIdx = 0; + (blockStackIdx < UNIT_BLOCK_STACK_NUM) && ((nIdx + blockStackIdx) < (nLoop + UNIT_BLOCK_STACK_NUM)); + blockStackIdx++) { + uint32_t nIdxActual = nIdx + blockStackIdx; + if (nIdxActual == (nLoop + UNIT_BLOCK_STACK_NUM - 1)) { + seqTile = (kvSeqlen - (nIdxActual - UNIT_BLOCK_STACK_NUM) * blockSize); + } else { + seqTile = blockSize; + } + seqTileRound = RoundUp(seqTile); + uint32_t L1ABPingPongFlag = nIdxActual % 2; + uint32_t L0ABPingPongFlag = nIdxActual % 2; + uint32_t blockTableId = gblockTable.GetValue(nIdxActual - UNIT_BLOCK_STACK_NUM); + uint64_t kvOffset = (uint64_t)blockTableId * blockSize * embed + embedSplitIdx * embedSplitSize; + + AscendC::WaitFlag(L1ABPingPongFlag + 4); + auto layoutUnitBSplitN = layoutB.GetTileLayout(MakeCoord(seqTile, embedSplitSize)); + // copy V to L1 + LayoutBInL1 layoutUnitBSplitNInL1 = LayoutBInL1::template MakeLayout(seqTile, embedSplitSize); + copyGmToL1B(l1BTensor[L1ABPingPongFlag], gB[kvOffset], layoutUnitBSplitNInL1, layoutUnitBSplitN); + AscendC::SetFlag(L1ABPingPongFlag); + + AscendC::WaitFlag(L1ABPingPongFlag); + AscendC::WaitFlag(L0ABPingPongFlag + 2); + LayoutBInL0 layoutUnitBSplitNInL0 = LayoutBInL0::template MakeLayout(seqTile, embedSplitSize); + // copy V from L1 to L0B + copyL1ToL0B(l0BTensor[L0ABPingPongFlag], l1BTensor[L1ABPingPongFlag], layoutUnitBSplitNInL0, + layoutUnitBSplitNInL1); + AscendC::SetFlag(L1ABPingPongFlag + 4); + + if (embedSplitIdx == 0 && blockStackIdx == 0) { + Arch::CrossCoreWaitFlag(softmaxReady); + } + AscendC::WaitFlag(L1ABPingPongFlag + 6); + auto layoutASplitK = layoutA.GetTileLayout(MakeCoord(rowNum, seqTile)); + LayoutAInL1 layoutASplitKInL1 = LayoutAInL1::template MakeLayout(rowNum, seqTile); + // copy P to L1 + copyGmToL1A(l1ATensor[L1ABPingPongFlag], gA[blockStackIdx * blockSize], layoutASplitKInL1, + layoutASplitK); + AscendC::SetFlag(EVENT_ID7); + // copy P to l0a + AscendC::WaitFlag(EVENT_ID7); + AscendC::WaitFlag(L0ABPingPongFlag); + LayoutAInL0 layoutASplitKInL0 = LayoutAInL0::template MakeLayout(rowNum, seqTile); + copyL1ToL0A(l0ATensor[L0ABPingPongFlag], l1ATensor[L1ABPingPongFlag], layoutASplitKInL0, + layoutASplitKInL1); + AscendC::SetFlag(L1ABPingPongFlag + 6); + AscendC::SetFlag(L0ABPingPongFlag); + // mmad + AscendC::WaitFlag(L0ABPingPongFlag); + if (blockStackIdx == 0) { + AscendC::WaitFlag(L0CPingPongFlag); + } + tileMmad(l0CTensor[L0CPingPongFlag], l0ATensor[L0ABPingPongFlag], l0BTensor[L0ABPingPongFlag], + rowNumRound, embedSplitSize, seqTile, blockStackIdx == 0); + AscendC::SetFlag(L0ABPingPongFlag + 2); + AscendC::SetFlag(L0ABPingPongFlag); + } + // copy Otmp to gm + AscendC::SetFlag(L0CPingPongFlag); + AscendC::WaitFlag(L0CPingPongFlag); + auto blockShape = MakeCoord(rowNum, embedSplitSize); + auto layoutInL0C = LayoutCInL0::MakeLayoutInL0C(blockShape); + auto layoutCSplitN = layoutC.GetTileLayout(MakeCoord(rowNumRound, embedSplitSizeRound)); + copyL0CToGm(gC[embedSplitIdx * embedSplitSizeRound], l0CTensor[L0CPingPongFlag], layoutCSplitN, + layoutInL0C); + AscendC::SetFlag(L0CPingPongFlag); + } + } + +protected: + /// Data members + AscendC::LocalTensor l1ATensor[STAGES]; + AscendC::LocalTensor l1BTensor[STAGES]; + AscendC::LocalTensor l0ATensor[STAGES]; + AscendC::LocalTensor l0BTensor[STAGES]; + AscendC::LocalTensor l0CTensor[STAGES]; + + TileMmad tileMmad; + CopyGmToL1A copyGmToL1A; + CopyGmToL1B copyGmToL1B; + CopyL1ToL0A copyL1ToL0A; + CopyL1ToL0B copyL1ToL0B; + CopyL0CToGm copyL0CToGm; +}; + +//////////////////////////////////////////////////////////////////// + +} // namespace Catlass::Gemm::Block + +#endif // CATLASS_KDA_BLOCK_MMAD_MLA_PV_TP1_SPEC_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_mla_qk.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_mla_qk.hpp new file mode 100644 index 0000000000..3b1fedb199 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_mla_qk.hpp @@ -0,0 +1,241 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_BLOCK_MMAD_MLA_QK_HPP +#define CATLASS_KDA_BLOCK_MMAD_MLA_QK_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/kda_coord.hpp" +#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" +#include "catlass/gemm/kda_gemm_helper.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" +#include "catlass/gemm/tile/kda_tile_mmad.hpp" + +//////////////////////////////////////////////////////////////////// + +namespace Catlass::Gemm::Block { +//////////////////////////////////////////////////////////////////// + +template < + class L1TileShape_, + class L0TileShape_, + class AType_, + class BType_, + class CType_, + class BiasType_, + class TileCopy_, + class TileMmad_> +struct BlockMmad< + MmadAtlasA2MLAQK, + L1TileShape_, + L0TileShape_, + AType_, + BType_, + CType_, + BiasType_, + TileCopy_, + TileMmad_> { +public: + // Type Aliases + using DispatchPolicy = MmadAtlasA2MLAQK; + using ArchTag = typename DispatchPolicy::ArchTag; + using L1TileShape = L1TileShape_; + using L0TileShape = L0TileShape_; + using ElementA = typename AType_::Element; + using LayoutA = typename AType_::Layout; + using ElementB = typename BType_::Element; + using LayoutB = typename BType_::Layout; + using ElementC = typename CType_::Element; + using LayoutC = typename CType_::Layout; + using TileMmad = TileMmad_; + using CopyGmToL1A = typename TileCopy_::CopyGmToL1A; + using CopyGmToL1B = typename TileCopy_::CopyGmToL1B; + using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; + using CopyL0CToGm = typename TileCopy_::CopyL0CToGm; + using ElementAccumulator = + typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; + using LayoutAInL1 = typename CopyL1ToL0A::LayoutSrc; + using LayoutBInL1 = typename CopyL1ToL0B::LayoutSrc; + using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst; + using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst; + using LayoutCInL0 = layout::zN; + + using L1AAlignHelper = Gemm::helper::L1AlignHelper; + using L1BAlignHelper = Gemm::helper::L1AlignHelper; + + static constexpr uint32_t STAGES = DispatchPolicy::STAGES; + static constexpr uint32_t L1A_SIZE = L1TileShape::M * L1TileShape::K * sizeof(ElementA); + static constexpr uint32_t L1B_SIZE = L1TileShape::N * L1TileShape::K * sizeof(ElementB); + static constexpr uint32_t L0A_SIZE = ArchTag::L0A_SIZE; + static constexpr uint32_t L0B_SIZE = ArchTag::L0B_SIZE; + static constexpr uint32_t L0C_SIZE = ArchTag::L0C_SIZE; + static constexpr uint32_t L0A_PINGPONG_BUF_SIZE = L0A_SIZE / STAGES; + static constexpr uint32_t L0B_PINGPONG_BUF_SIZE = L0B_SIZE / STAGES; + static constexpr uint32_t L0C_PINGPONG_BUF_SIZE = L0C_SIZE / STAGES; + static constexpr uint32_t BLOCK_SIZE = 16; + static constexpr uint32_t EMBED_SPLIT_SIZE = 128; + static constexpr uint32_t EMBED_SPLIT_LOOP = 5; + + // Check LayoutC + static_assert(std::is_same_v, "LayoutC only support RowMajor yet!"); + + /// Construct + CATLASS_DEVICE + BlockMmad(Arch::Resource &resource, uint32_t l1BufAddrStart = 0) + { + // Allocate L1 memory space + l1ATensor = resource.l1Buf.template GetBufferByByte(l1BufAddrStart); + for (uint32_t i = 0; i < STAGES; i++) { + l1BTensor[i] = + resource.l1Buf.template GetBufferByByte(l1BufAddrStart + L1A_SIZE + L1B_SIZE * i); + l0ATensor[i] = resource.l0ABuf.template GetBufferByByte(L0A_PINGPONG_BUF_SIZE * i); + l0BTensor[i] = resource.l0BBuf.template GetBufferByByte(L0B_PINGPONG_BUF_SIZE * i); + l0CTensor[i] = resource.l0CBuf.template GetBufferByByte(L0C_PINGPONG_BUF_SIZE * i); + } + } + + /// Destructor + CATLASS_DEVICE + ~BlockMmad() {} + + /// Perform a block-scoped matrix multiply-accumulate + CATLASS_DEVICE + void operator()( + AscendC::GlobalTensor gA, + AscendC::GlobalTensor gARope, + AscendC::GlobalTensor gB, + AscendC::GlobalTensor gBRope, + AscendC::GlobalTensor gC, + LayoutA layoutA, LayoutA layoutARope, LayoutB layoutB, LayoutB layoutBRope, LayoutC layoutC, + GemmCoord actualShape, MatrixCoord qShapeSingleNd, + uint32_t &qHeads, uint32_t &nIdx) + { + uint32_t rowNum = actualShape.m(); + uint32_t kSeqTile = actualShape.n(); + uint32_t embed = layoutB.shape(0); + uint32_t embedRope = layoutBRope.shape(0); + uint32_t embedCat = actualShape.k(); + uint32_t embedSplitSize = EMBED_SPLIT_SIZE; + uint32_t embedSplitLoopK = EMBED_SPLIT_LOOP; + uint32_t curHeadNum = qShapeSingleNd.row(); + uint32_t tokenNumPerHead = rowNum / curHeadNum; + uint32_t kSeqTileRound = RoundUp(kSeqTile); + uint32_t rowNumRound = RoundUp(rowNum); + uint32_t l1KvPingPongFlag = nIdx % 2; + + if (nIdx == 0) { + // copy Q to L1 + auto layoutASingleNd = layoutA.GetTileLayout(MakeCoord(curHeadNum, embed)); + LayoutAInL1 layoutAInL1 = LayoutAInL1::template MakeLayout(rowNum, embed); + copyGmToL1A( + l1ATensor, gA, + layoutAInL1, layoutASingleNd, + tokenNumPerHead, qHeads * embed, tokenNumPerHead, BLOCK_SIZE, rowNumRound); + + // copy QRope to L1 + auto layoutARopeSingleNd = layoutARope.GetTileLayout(MakeCoord(curHeadNum, embedRope)); + LayoutAInL1 layoutARopeInL1 = LayoutAInL1::template MakeLayout(rowNum, embedRope); + copyGmToL1A( + l1ATensor[rowNumRound * embed], gARope, + layoutARopeInL1, layoutARopeSingleNd, + tokenNumPerHead, qHeads * embedRope, tokenNumPerHead, BLOCK_SIZE, rowNumRound); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + } + + AscendC::WaitFlag(l1KvPingPongFlag); + // copy K to L1 + LayoutBInL1 layoutBInL1 = LayoutBInL1::template MakeLayout(embed, kSeqTile); + copyGmToL1B(l1BTensor[l1KvPingPongFlag], gB, layoutBInL1, layoutB); + AscendC::SetFlag(l1KvPingPongFlag); + + AscendC::WaitFlag(l1KvPingPongFlag + 2); + // copy KRope to L1 + LayoutBInL1 layoutBRopeInL1 = LayoutBInL1::template MakeLayout(embedRope, kSeqTile); + copyGmToL1B(l1BTensor[l1KvPingPongFlag][kSeqTileRound * embed], gBRope, layoutBRopeInL1, layoutBRope); + AscendC::SetFlag(l1KvPingPongFlag + 2); + + for (uint32_t embedSplitIdx = 0; embedSplitIdx < embedSplitLoopK; embedSplitIdx++) { + uint32_t l0ABPingPongFlag = embedSplitIdx % 2; + if (embedSplitIdx == embedSplitLoopK - 1) { + embedSplitSize = embedRope; + } + // copy Q from L1 to l0a + AscendC::WaitFlag(l0ABPingPongFlag); + LayoutAInL1 layoutACatInL1 = LayoutAInL1::template MakeLayout(rowNum, embedCat); + LayoutAInL0 layoutACatInL0 = LayoutAInL0::template MakeLayout(rowNum, embedSplitSize); + copyL1ToL0A( + l0ATensor[l0ABPingPongFlag], l1ATensor[embedSplitIdx * rowNumRound * EMBED_SPLIT_SIZE], + layoutACatInL0, layoutACatInL1); + AscendC::SetFlag(l0ABPingPongFlag); + + if (embedSplitIdx == 0) { + AscendC::WaitFlag(l1KvPingPongFlag); + } else if (embedSplitIdx == embedSplitLoopK - 1) { + AscendC::WaitFlag(l1KvPingPongFlag + 2); + } + AscendC::WaitFlag(l0ABPingPongFlag + 2); + // copy K from L1 to l0b + LayoutBInL1 layoutBCatInL1 = LayoutBInL1::template MakeLayout(embedCat, kSeqTile); + LayoutBInL0 layoutBCatInL0 = LayoutBInL0::template MakeLayout(embedSplitSize, kSeqTile); + copyL1ToL0B(l0BTensor[l0ABPingPongFlag], + l1BTensor[l1KvPingPongFlag][embedSplitIdx * kSeqTileRound * EMBED_SPLIT_SIZE], + layoutBCatInL0, layoutBCatInL1); + if (embedSplitIdx == embedSplitLoopK - 1) { + AscendC::SetFlag(l1KvPingPongFlag + 2); + } + AscendC::SetFlag(l0ABPingPongFlag + 2); + // mmad + AscendC::WaitFlag(l0ABPingPongFlag); + AscendC::WaitFlag(l0ABPingPongFlag + 2); + if (embedSplitIdx == 0) { + AscendC::WaitFlag(l1KvPingPongFlag); + } + // mmad + tileMmad( + l0CTensor[l1KvPingPongFlag], l0ATensor[l0ABPingPongFlag], l0BTensor[l0ABPingPongFlag], + rowNumRound, kSeqTile, embedSplitSize, embedSplitIdx == 0); + AscendC::SetFlag(l0ABPingPongFlag); + AscendC::SetFlag(l0ABPingPongFlag + 2); + } + // copy S to gm + AscendC::SetFlag(l1KvPingPongFlag); + AscendC::WaitFlag(l1KvPingPongFlag); + auto blockShape = MakeCoord(rowNum, kSeqTile); + auto layoutInL0C = LayoutCInL0::MakeLayoutInL0C(blockShape); + // copy L0C to gm + copyL0CToGm(gC, l0CTensor[l1KvPingPongFlag], layoutC, layoutInL0C); + AscendC::SetFlag(l1KvPingPongFlag); + } + +protected: + /// Data members + AscendC::LocalTensor l1ATensor; + AscendC::LocalTensor l1BTensor[STAGES]; + AscendC::LocalTensor l0ATensor[STAGES]; + AscendC::LocalTensor l0BTensor[STAGES]; + AscendC::LocalTensor l0CTensor[STAGES]; + + TileMmad tileMmad; + CopyGmToL1A copyGmToL1A; + CopyGmToL1B copyGmToL1B; + CopyL1ToL0A copyL1ToL0A; + CopyL1ToL0B copyL1ToL0B; + CopyL0CToGm copyL0CToGm; +}; + +//////////////////////////////////////////////////////////////////// + +} // namespace Catlass::Gemm::Block + +#endif // CATLASS_KDA_BLOCK_MMAD_MLA_QK_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_mla_qk_tp1_spec.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_mla_qk_tp1_spec.hpp new file mode 100644 index 0000000000..999cfa0bf0 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_mla_qk_tp1_spec.hpp @@ -0,0 +1,259 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_BLOCK_MMAD_MLA_QK_TP1_SPEC_HPP +#define CATLASS_KDA_BLOCK_MMAD_MLA_QK_TP1_SPEC_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/kda_coord.hpp" +#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" +#include "catlass/gemm/kda_gemm_helper.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" +#include "catlass/gemm/tile/kda_tile_mmad.hpp" + +//////////////////////////////////////////////////////////////////// + +namespace Catlass::Gemm::Block { +//////////////////////////////////////////////////////////////////// + +template < + class L1TileShape_, + class L0TileShape_, + class AType_, + class BType_, + class CType_, + class BiasType_, + class TileCopy_, + class TileMmad_> +struct BlockMmad< + MmadAtlasA2MLAQKTp1Spec, + L1TileShape_, + L0TileShape_, + AType_, + BType_, + CType_, + BiasType_, + TileCopy_, + TileMmad_> { +public: + // Type Aliases + using DispatchPolicy = MmadAtlasA2MLAQKTp1Spec; + using ArchTag = typename DispatchPolicy::ArchTag; + using L1TileShape = L1TileShape_; + using L0TileShape = L0TileShape_; + using ElementA = typename AType_::Element; + using LayoutA = typename AType_::Layout; + using ElementB = typename BType_::Element; + using LayoutB = typename BType_::Layout; + using ElementC = typename CType_::Element; + using LayoutC = typename CType_::Layout; + using TileMmad = TileMmad_; + using CopyGmToL1A = typename TileCopy_::CopyGmToL1A; + using CopyGmToL1B = typename TileCopy_::CopyGmToL1B; + using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; + using CopyL0CToGm = typename TileCopy_::CopyL0CToGm; + using ElementAccumulator = + typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; + using LayoutAInL1 = typename CopyL1ToL0A::LayoutSrc; + using LayoutBInL1 = typename CopyL1ToL0B::LayoutSrc; + using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst; + using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst; + using LayoutCInL0 = layout::zN; + + using L1AAlignHelper = Gemm::helper::L1AlignHelper; + using L1BAlignHelper = Gemm::helper::L1AlignHelper; + + static constexpr uint32_t STAGES = DispatchPolicy::STAGES; + static constexpr uint32_t L1A_SIZE = L1TileShape::M * L1TileShape::K * sizeof(ElementA); + static constexpr uint32_t L0A_SIZE = ArchTag::L0A_SIZE; + static constexpr uint32_t L0B_SIZE = ArchTag::L0B_SIZE; + static constexpr uint32_t L0C_SIZE = ArchTag::L0C_SIZE; + static constexpr uint32_t L0A_PINGPONG_BUF_SIZE = L0A_SIZE / STAGES; + static constexpr uint32_t L0B_PINGPONG_BUF_SIZE = L0B_SIZE / STAGES; + static constexpr uint32_t L0C_PINGPONG_BUF_SIZE = L0C_SIZE / STAGES; + static constexpr uint32_t UNIT_BLOCK_STACK_NUM = 4; + static constexpr uint32_t BLOCK_SIZE = 16; + static constexpr uint32_t EMBED_SPLIT_SIZE = 128; + static constexpr uint32_t EMBED_ROPE = 64; + static constexpr uint32_t GM_L1_EMBED_SPLIT_SIZE = 256; + static constexpr uint32_t EMBED_SPLIT_LOOP = 5; + static constexpr uint32_t L1B_SIZE = L1TileShape::N * GM_L1_EMBED_SPLIT_SIZE * sizeof(ElementB); + static constexpr uint32_t L1BROPE_SIZE = L1TileShape::N * EMBED_ROPE * sizeof(ElementB); + static constexpr uint32_t L1B_ROPE_START = L1TileShape::N * GM_L1_EMBED_SPLIT_SIZE; + + /// Construct + CATLASS_DEVICE + BlockMmad(Arch::Resource &resource, uint32_t l1BufAddrStart = 0) + { + // Allocate L1 memory space + l1ATensor = resource.l1Buf.template GetBufferByByte(l1BufAddrStart); + for (uint32_t i = 0; i < STAGES; i++) { + l1BTensor[i] = resource.l1Buf.template GetBufferByByte(l1BufAddrStart + L1A_SIZE + L1B_SIZE * i); + l1BRopeTensor[i] = resource.l1Buf.template GetBufferByByte(l1BufAddrStart + L1A_SIZE + + L1B_SIZE * STAGES + L1BROPE_SIZE * i); + l0ATensor[i] = resource.l0ABuf.template GetBufferByByte(L0A_PINGPONG_BUF_SIZE * i); + l0BTensor[i] = resource.l0BBuf.template GetBufferByByte(L0B_PINGPONG_BUF_SIZE * i); + l0CTensor[i] = resource.l0CBuf.template GetBufferByByte(L0C_PINGPONG_BUF_SIZE * i); + } + } + + /// Destructor + CATLASS_DEVICE + ~BlockMmad() {} + + /// Perform a block-scoped matrix multiply-accumulate + CATLASS_DEVICE + void operator()(AscendC::GlobalTensor gA, AscendC::GlobalTensor gARope, + AscendC::GlobalTensor gB, AscendC::GlobalTensor gBRope, + AscendC::GlobalTensor gblockTable, AscendC::GlobalTensor gC, LayoutA layoutA, + LayoutA layoutARope, LayoutB layoutB, LayoutB layoutBRope, LayoutC layoutC, GemmCoord actualShape, + uint32_t &nIdx, uint32_t &nLoop, uint32_t &blockSize, uint32_t kvSeqlen) + { + uint32_t rowNum = actualShape.m(); + uint32_t stackSeqTile = actualShape.n(); + uint32_t seqTile = blockSize; + uint32_t embed = layoutA.shape(1); + uint32_t embedRope = layoutARope.shape(1); + uint32_t embedCat = actualShape.k(); + uint32_t embedSplitGm2L1 = GM_L1_EMBED_SPLIT_SIZE; + uint32_t stackSeqTileRound = layoutC.shape(1); + uint32_t rowNumRound = layoutC.shape(0); + uint32_t seqTileRound = RoundUp(seqTile); + + if (nIdx == 0) { + // copy Q to L1 + LayoutAInL1 layoutAInL1 = LayoutAInL1::template MakeLayout(rowNum, embed); + copyGmToL1A(l1ATensor, gA, layoutAInL1, layoutA); + + // copy QRope to L1 + LayoutAInL1 layoutARopeInL1 = LayoutAInL1::template MakeLayout(rowNum, embedRope); + copyGmToL1A(l1ATensor[rowNumRound * embed], gARope, layoutARopeInL1, layoutARope); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + } + + for (uint32_t blockStackIdx = 0; (blockStackIdx < UNIT_BLOCK_STACK_NUM) && ((nIdx + blockStackIdx) < nLoop); + blockStackIdx++) { + uint32_t nIdxActual = nIdx + blockStackIdx; + uint32_t L0CPingPongFlag = nIdxActual % 2; + uint32_t L1BRopePingPongFlag = nIdxActual % 2; + if (nIdxActual == (nLoop - 1)) { + seqTile = (kvSeqlen - nIdxActual * blockSize); + seqTileRound = RoundUp(seqTile); + } + uint32_t blockTableId = gblockTable.GetValue(nIdxActual); + uint64_t kvOffset = (uint64_t)blockTableId * blockSize * embed; + uint64_t kvOffsetRope = (uint64_t)blockTableId * blockSize * embedRope; + uint64_t l1bSplitOffset = 0; + uint32_t embedSplitSize = EMBED_SPLIT_SIZE; + uint32_t embedSplitLoopK = EMBED_SPLIT_LOOP; + for (uint32_t embedSplitIdx = 0; embedSplitIdx < embedSplitLoopK; embedSplitIdx++) { + uint32_t L0ABPingPongFlag = (blockStackIdx + embedSplitIdx) % 2; + uint32_t innerSplitIdxL12L0 = embedSplitIdx % 2; + uint32_t L1BPingPongFlag = embedSplitIdx / 2; + if (embedSplitIdx == 4) { + embedSplitSize = embedRope; + } + // copy Q to l0a + AscendC::WaitFlag(L0ABPingPongFlag); + LayoutAInL1 layoutACatSplitKInL1 = LayoutAInL1::template MakeLayout(rowNum, embedSplitSize); + LayoutAInL0 layoutACatSplitKInL0 = LayoutAInL0::template MakeLayout(rowNum, embedSplitSize); + copyL1ToL0A(l0ATensor[L0ABPingPongFlag], l1ATensor[embedSplitIdx * rowNumRound * EMBED_SPLIT_SIZE], + layoutACatSplitKInL0, layoutACatSplitKInL1); + AscendC::SetFlag(L0ABPingPongFlag); + + if (embedSplitIdx == 0 || embedSplitIdx == 2) { + // copy K to l1b + AscendC::WaitFlag(L1BPingPongFlag); + auto layoutUnitBSplitK = layoutB.GetTileLayout(MakeCoord(embedSplitGm2L1, seqTile)); + LayoutBInL1 layoutUnitBSplitKInL1 = + LayoutBInL1::template MakeLayout(embedSplitGm2L1, seqTile); + copyGmToL1B(l1BTensor[L1BPingPongFlag], gB[kvOffset + embedSplitIdx * embedSplitSize], + layoutUnitBSplitKInL1, layoutUnitBSplitK); + AscendC::SetFlag(L1BPingPongFlag); + AscendC::WaitFlag(L1BPingPongFlag); + } else if (embedSplitIdx == 4) { + // copy KRope to L1 + AscendC::WaitFlag(L1BRopePingPongFlag + 2); + auto layoutUnitBRope = layoutBRope.GetTileLayout(MakeCoord(embedRope, seqTile)); + LayoutBInL1 layoutBRopeInL1 = LayoutBInL1::template MakeLayout(embedRope, seqTile); + copyGmToL1B(l1BRopeTensor[L1BRopePingPongFlag], gBRope[kvOffsetRope], layoutBRopeInL1, + layoutUnitBRope); + AscendC::SetFlag(L1BRopePingPongFlag); + AscendC::WaitFlag(L1BRopePingPongFlag); + } + // copy K to l0b + AscendC::WaitFlag(L0ABPingPongFlag + 2); + LayoutBInL1 layoutBCatSplitKInL1 = LayoutBInL1::template MakeLayout(embedSplitSize, seqTile); + LayoutBInL0 layoutBCatSplitKInL0 = LayoutBInL0::template MakeLayout(embedSplitSize, seqTile); + if (embedSplitIdx != 4) { + copyL1ToL0B(l0BTensor[L0ABPingPongFlag], + l1BTensor[L1BPingPongFlag][innerSplitIdxL12L0 * EMBED_SPLIT_SIZE * seqTileRound], + layoutBCatSplitKInL0, layoutBCatSplitKInL1); + } else { + copyL1ToL0B(l0BTensor[L0ABPingPongFlag], l1BRopeTensor[L1BRopePingPongFlag], layoutBCatSplitKInL0, + layoutBCatSplitKInL1); + } + + if (embedSplitIdx == 1 || embedSplitIdx == 3) { + AscendC::SetFlag(L1BPingPongFlag); + } else if (embedSplitIdx == 4) { + AscendC::SetFlag(L1BRopePingPongFlag + 2); + } + AscendC::SetFlag(L0ABPingPongFlag + 2); + + AscendC::WaitFlag(L0ABPingPongFlag); + AscendC::WaitFlag(L0ABPingPongFlag + 2); + if (embedSplitIdx == 0) { + AscendC::WaitFlag(L0CPingPongFlag); + } + // mmad + tileMmad(l0CTensor[L0CPingPongFlag], l0ATensor[L0ABPingPongFlag], l0BTensor[L0ABPingPongFlag], + rowNumRound, seqTile, embedSplitSize, embedSplitIdx == 0); + AscendC::SetFlag(L0ABPingPongFlag); + AscendC::SetFlag(L0ABPingPongFlag + 2); + } + // copy S to gm + AscendC::SetFlag(L0CPingPongFlag); + AscendC::WaitFlag(L0CPingPongFlag); + auto blockShape = MakeCoord(rowNum, seqTile); + auto layoutInL0C = LayoutCInL0::MakeLayoutInL0C(blockShape); + auto layoutCSplitN = layoutC.GetTileLayout(MakeCoord(rowNumRound, seqTileRound)); + // copy L0C to gm + copyL0CToGm(gC[blockStackIdx * blockSize], l0CTensor[L0CPingPongFlag], layoutCSplitN, layoutInL0C); + AscendC::SetFlag(L0CPingPongFlag); + } + } + +protected: + /// Data members + AscendC::LocalTensor l1ATensor; + AscendC::LocalTensor l1BTensor[STAGES]; + AscendC::LocalTensor l1BRopeTensor[STAGES]; + AscendC::LocalTensor l0ATensor[STAGES]; + AscendC::LocalTensor l0BTensor[STAGES]; + AscendC::LocalTensor l0CTensor[STAGES]; + + TileMmad tileMmad; + CopyGmToL1A copyGmToL1A; + CopyGmToL1B copyGmToL1B; + CopyL1ToL0A copyL1ToL0A; + CopyL1ToL0B copyL1ToL0B; + CopyL0CToGm copyL0CToGm; +}; + +//////////////////////////////////////////////////////////////////// + +} // namespace Catlass::Gemm::Block + +#endif // CATLASS_KDA_BLOCK_MMAD_MLA_QK_TP1_SPEC_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_pingpong.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_pingpong.hpp new file mode 100644 index 0000000000..685821c0e4 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_pingpong.hpp @@ -0,0 +1,362 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_BLOCK_MMAD_PINGPONG_HPP +#define CATLASS_KDA_BLOCK_MMAD_PINGPONG_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/kda_coord.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" +#include "catlass/gemm/kda_gemm_helper.hpp" +#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" +#include "catlass/gemm/tile/kda_tile_mmad.hpp" + +namespace Catlass::Gemm::Block { + +template < + bool ENABLE_UNIT_FLAG_, + class L1TileShape_, + class L0TileShape_, + class AType_, + class BType_, + class CType_, + class BiasType_, + class TileCopy_, + class TileMmad_> +struct BlockMmad< + MmadAtlasA2Pingpong, + L1TileShape_, + L0TileShape_, + AType_, + BType_, + CType_, + BiasType_, + TileCopy_, + TileMmad_> { +public: + // Type Aliases + using DispatchPolicy = MmadAtlasA2Pingpong; + using ArchTag = typename DispatchPolicy::ArchTag; + using L1TileShape = L1TileShape_; + using L0TileShape = L0TileShape_; + using ElementA = typename AType_::Element; + using LayoutA = typename AType_::Layout; + using ElementB = typename BType_::Element; + using LayoutB = typename BType_::Layout; + using ElementC = typename CType_::Element; + using LayoutC = typename CType_::Layout; + using TileMmad = TileMmad_; + using CopyGmToL1A = typename TileCopy_::CopyGmToL1A; + using CopyGmToL1B = typename TileCopy_::CopyGmToL1B; + using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; + using CopyL0CToGm = typename TileCopy_::CopyL0CToGm; + using ElementAccumulator = + typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; + using LayoutAInL1 = typename CopyL1ToL0A::LayoutSrc; + using LayoutBInL1 = typename CopyL1ToL0B::LayoutSrc; + using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst; + using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst; + using LayoutCInL0 = layout::zN; + + using L1AAlignHelper = Gemm::helper::L1AlignHelper; + using L1BAlignHelper = Gemm::helper::L1AlignHelper; + + static constexpr bool ENABLE_UNIT_FLAG = DispatchPolicy::ENABLE_UNIT_FLAG; + static constexpr uint32_t STAGES = DispatchPolicy::STAGES; + static constexpr uint32_t L1A_SIZE = L1TileShape::M * L1TileShape::K * sizeof(ElementA); + static constexpr uint32_t L1B_SIZE = L1TileShape::N * L1TileShape::K * sizeof(ElementB); + static constexpr uint32_t L0A_SIZE = ArchTag::L0A_SIZE; + static constexpr uint32_t L0B_SIZE = ArchTag::L0B_SIZE; + static constexpr uint32_t L0C_SIZE = ArchTag::L0C_SIZE; + static constexpr uint32_t L0A_PINGPONG_BUF_SIZE = L0A_SIZE / STAGES; + static constexpr uint32_t L0B_PINGPONG_BUF_SIZE = L0B_SIZE / STAGES; + + // Check LayoutC + static_assert(std::is_same_v, "LayoutC only support RowMajor yet!"); + + // Check L1TileShape + static_assert((L1A_SIZE * STAGES + L1B_SIZE * STAGES) <= ArchTag::L1_SIZE, "L1TileShape exceeding the L1 space!"); + + // Check L0TileShape + static constexpr uint32_t L0A_TILE_SIZE = L0TileShape::M * L0TileShape::K * sizeof(ElementA); + static constexpr uint32_t L0B_TILE_SIZE = L0TileShape::K * L0TileShape::N * sizeof(ElementB); + static constexpr uint32_t L0C_TILE_SIZE = L0TileShape::M * L0TileShape::N * sizeof(ElementAccumulator); + static_assert((L0A_TILE_SIZE * STAGES) <= L0A_SIZE, "L0TileShape exceeding the L0A space!"); + static_assert((L0B_TILE_SIZE * STAGES) <= L0B_SIZE, "L0TileShape exceeding the L0B space!"); + static_assert(L0C_TILE_SIZE <= L0C_SIZE, "L0TileShape exceeding the L0C space!"); + + // Check tileshape + static_assert(L1TileShape::M == L0TileShape::M && L1TileShape::N == L0TileShape::N, + "The situation where the basic blocks of L1 and L0 differ on the m and n axes is not supported yet"); + static_assert(L0TileShape::K <= L1TileShape::K, "L0TileShape::K cannot exceed L1TileShape::K"); + + // 32B (256b) aligned + static_assert(Gemm::helper::TileShapeAlignChecker::_ALIGN == 256, + "Tile shape must be 32B aligned."); + + /// Construct + CATLASS_DEVICE + BlockMmad(Arch::Resource &resource, uint32_t l1BufAddrStart = 0) + { + uint32_t l1AOffset = l1BufAddrStart; + uint32_t l1BOffset = l1BufAddrStart + L1A_SIZE * STAGES; + // Init buffers + for (uint32_t i = 0; i < STAGES; i++) { + // Assign L1/L0A/L0B space for each stages + l1ATensorList[i] = resource.l1Buf.template GetBufferByByte(l1AOffset + L1A_SIZE * i); + l1BTensorList[i] = resource.l1Buf.template GetBufferByByte(l1BOffset + L1B_SIZE * i); + l0ATensorList[i] = resource.l0ABuf.template GetBufferByByte(L0A_PINGPONG_BUF_SIZE * i); + l0BTensorList[i] = resource.l0BBuf.template GetBufferByByte(L0B_PINGPONG_BUF_SIZE * i); + + // Assign event ID for each stages + l1AEventList[i] = i; + l1BEventList[i] = i + STAGES; + l0AEventList[i] = i; + l0BEventList[i] = i + STAGES; + + // The event id that needs to be set before the loop + AscendC::SetFlag(l1AEventList[i]); + AscendC::SetFlag(l1BEventList[i]); + AscendC::SetFlag(l0AEventList[i]); + AscendC::SetFlag(l0BEventList[i]); + } + l0CTensor = resource.l0CBuf.template GetBufferByByte(0); + AscendC::SetFlag(EVENT_ID0); + } + + /// Destructor + CATLASS_DEVICE + ~BlockMmad() + { + for (uint32_t i = 0; i < STAGES; i++) { + AscendC::WaitFlag(l1AEventList[i]); + AscendC::WaitFlag(l1BEventList[i]); + AscendC::WaitFlag(l0AEventList[i]); + AscendC::WaitFlag(l0BEventList[i]); + } + AscendC::WaitFlag(EVENT_ID0); + } + + /// Perform a block-scoped matrix multiply-accumulate + CATLASS_DEVICE + void operator()( + AscendC::GlobalTensor const &gmA, LayoutA const &layoutA, + AscendC::GlobalTensor const &gmB, LayoutB const &layoutB, + AscendC::GlobalTensor const &gmC, LayoutC const &layoutC, + GemmCoord const &actualShape) + { + uint32_t mRound = RoundUp(actualShape.m()); + uint32_t nRound = RoundUp(actualShape.n()); + + auto layoutAInL1 = LayoutAInL1::template MakeLayout(L1TileShape::M, L1TileShape::K); + auto layoutBInL1 = LayoutBInL1::template MakeLayout(L1TileShape::K, L1TileShape::N); + auto layoutInL0C = LayoutCInL0::MakeLayoutInL0C(MakeCoord(mRound, nRound)); + + uint32_t kActual = min(actualShape.k(), L1TileShape::K); + + // load first matrix A tile from GM to L1 + AscendC::WaitFlag(l1AEventList[l1ListId]); + auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShape.m(), kActual)); + copyGmToL1A(l1ATensorList[l1ListId], gmA, layoutAInL1, layoutTileA); + AscendC::SetFlag(l1AEventList[l1ListId]); + + // load first matrix B tile from GM to L1 + AscendC::WaitFlag(l1BEventList[l1ListId]); + auto layoutTileB = layoutB.GetTileLayout(MakeCoord(kActual, actualShape.n())); + copyGmToL1B(l1BTensorList[l1ListId], gmB, layoutBInL1, layoutTileB); + AscendC::SetFlag(l1BEventList[l1ListId]); + + if constexpr (!ENABLE_UNIT_FLAG) { + AscendC::WaitFlag(EVENT_ID0); + } + + uint32_t mPartLoop = CeilDiv(mRound); + uint32_t nPartLoop = CeilDiv(nRound); + + // main loop + uint32_t kTileCount = CeilDiv(actualShape.k()); + for (uint32_t kLoopIdx = 0; kLoopIdx < kTileCount; kLoopIdx++) { + uint32_t l1ListIdNext = (l1ListId + 1 < STAGES) ? (l1ListId + 1) : 0; + uint32_t kActualNext{0}; + // preload next tile from GM to L1 + if (kLoopIdx < kTileCount - 1) { + uint32_t kLoopIdxNext = kLoopIdx + 1; + kActualNext = (kLoopIdxNext < kTileCount - 1) ? + L1TileShape::K : + (actualShape.k() - kLoopIdxNext * L1TileShape::K); + + // Get L1 tensor for next stage + auto l1ATensor = l1ATensorList[l1ListIdNext]; + auto l1BTensor = l1BTensorList[l1ListIdNext]; + // Get GM tile for next stage + MatrixCoord gmTileAOffset{0, kLoopIdxNext * L1TileShape::K}; + MatrixCoord gmTileBOffset{kLoopIdxNext * L1TileShape::K, 0}; + auto gmTileA = gmA[layoutA.GetOffset(gmTileAOffset)]; + auto gmTileB = gmB[layoutB.GetOffset(gmTileBOffset)]; + + // load next matrix A tile from GM to L1 + AscendC::WaitFlag(l1AEventList[l1ListIdNext]); + layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShape.m(), kActualNext)); + copyGmToL1A(l1ATensor, gmTileA, layoutAInL1, layoutTileA); + AscendC::SetFlag(l1AEventList[l1ListIdNext]); + + // load next matrix B tile from GM to L1 + AscendC::WaitFlag(l1BEventList[l1ListIdNext]); + layoutTileB = layoutB.GetTileLayout(MakeCoord(kActualNext, actualShape.n())); + copyGmToL1B(l1BTensor, gmTileB, layoutBInL1, layoutTileB); + AscendC::SetFlag(l1BEventList[l1ListIdNext]); + } + + // Get L1 tensor for current stage + auto l1ATensor = l1ATensorList[l1ListId]; + auto l1BTensor = l1BTensorList[l1ListId]; + + // Get the loop nums on L0 + uint32_t kPartLoop = CeilDiv(kActual); + + for (int mPartIdx = 0; mPartIdx < mPartLoop; mPartIdx++) { + uint32_t mPartActual = (mPartIdx < mPartLoop - 1) ? + L0TileShape::M : + (mRound - mPartIdx * L0TileShape::M); + + for (int kPartIdx = 0; kPartIdx < kPartLoop; kPartIdx++) { + uint32_t kPartActual = (kPartIdx < kPartLoop - 1) ? + L0TileShape::K : + (kActual - kPartIdx * L0TileShape::K); + + // Locate the current tile on L0A + auto l0ATile = l0ATensorList[l0AListId]; + LayoutAInL0 layoutAInL0 = LayoutAInL0::template MakeLayout(mPartActual, kPartActual); + // Locate the current tile of matrix A on L1 + MatrixCoord l1AOffset{mPartIdx * L0TileShape::M, kPartIdx * L0TileShape::K}; + auto l1ATile = l1ATensor[layoutAInL1.GetOffset(l1AOffset)]; + + AscendC::WaitFlag(l0AEventList[l0AListId]); + if ((mPartIdx == 0) && (kPartIdx == 0)) { + AscendC::WaitFlag(l1AEventList[l1ListId]); + } + + // Load current tile from L1 to L0A + copyL1ToL0A(l0ATile, l1ATile, layoutAInL0, layoutAInL1); + + if ((mPartIdx == mPartLoop - 1) && (kPartIdx == kPartLoop - 1)) { + AscendC::SetFlag(l1AEventList[l1ListId]); + } + + for (int nPartIdx = 0; nPartIdx < nPartLoop; nPartIdx++) { + uint32_t nPartActual = (nPartIdx < nPartLoop - 1) ? + L0TileShape::N : + (nRound - nPartIdx * L0TileShape::N); + + // Locate the current tile on L0B + auto l0BTile = l0BTensorList[l0BListId]; + LayoutBInL0 layoutBInL0 = LayoutBInL0::template MakeLayout(kPartActual, nPartActual); + // Locate the current tile of matrix B on L1 + MatrixCoord l1BOffset{kPartIdx * L0TileShape::K, nPartIdx * L0TileShape::N}; + auto l1BTile = l1BTensor[layoutBInL1.GetOffset(l1BOffset)]; + + // Wait for mmad finished + AscendC::WaitFlag(l0BEventList[l0BListId]); + // If the current tile is the first one on the k&n axis, wait for loading matrix B from GM to L1 + if ((kPartIdx == 0) && (nPartIdx == 0)) { + AscendC::WaitFlag(l1BEventList[l1ListId]); + } + + // Load current tile from L1 to L0B + copyL1ToL0B(l0BTile, l1BTile, layoutBInL0, layoutBInL1); + + // If the current tile is the last one on the k&n axis, notify to load matrix B from GM to L1 + if ((kPartIdx == kPartLoop - 1) && (nPartIdx == nPartLoop - 1)) { + AscendC::SetFlag(l1BEventList[l1ListId]); + } + // Notify to do mmad + AscendC::SetFlag(EVENT_ID0); + + // Locate the current tile on L0C + MatrixCoord l0COffset{mPartIdx * L0TileShape::M, nPartIdx * L0TileShape::N}; + auto l0CTile = l0CTensor[layoutInL0C.GetOffset(l0COffset)]; + + // Compute the matrix multiplication on L0A and L0B and write the result to the accumulator + // Wait for loading L0B + AscendC::WaitFlag(EVENT_ID0); + + // If the current tile is the first tile on the k axis, the accumulator needs to be reset to 0 + bool initC = ((kLoopIdx == 0) && (kPartIdx == 0)); + // If the unit flag is enabled, the unit flag is set according to the calculation progress + uint8_t unitFlag = 0b00; + if constexpr (ENABLE_UNIT_FLAG) { + if ((kLoopIdx == kTileCount - 1) && (mPartIdx == mPartLoop - 1) && + (kPartIdx == kPartLoop - 1) && (nPartIdx == nPartLoop - 1)) { + unitFlag = 0b11; + } else { + unitFlag = 0b10; + } + } + // Perform calculation operations + tileMmad(l0CTile, l0ATile, l0BTile, mPartActual, nPartActual, kPartActual, initC, unitFlag); + + // Notify to move the next L0B tile + AscendC::SetFlag(l0BEventList[l0BListId]); + l0BListId = (l0BListId + 1 < STAGES) ? (l0BListId + 1) : 0; + } + AscendC::SetFlag(l0AEventList[l0AListId]); + l0AListId = (l0AListId + 1 < STAGES) ? (l0AListId + 1) : 0; + } + } + l1ListId = l1ListIdNext; + kActual = kActualNext; + } + + // copy block out + LayoutC layoutBlock = layoutC.GetTileLayout(actualShape.GetCoordMN()); + + if constexpr (!ENABLE_UNIT_FLAG) { + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + copyL0CToGm(gmC, l0CTensor, layoutBlock, layoutInL0C); + AscendC::SetFlag(EVENT_ID0); + } else { + copyL0CToGm(gmC, l0CTensor, layoutBlock, layoutInL0C, 0b11); + } + } + +protected: + // Multi-stage tensors list + AscendC::LocalTensor l1ATensorList[STAGES]; + AscendC::LocalTensor l1BTensorList[STAGES]; + AscendC::LocalTensor l0ATensorList[STAGES]; + AscendC::LocalTensor l0BTensorList[STAGES]; + AscendC::LocalTensor l0CTensor; + + // Multi-stage event id list + int32_t l1AEventList[STAGES]; + int32_t l1BEventList[STAGES]; + int32_t l0AEventList[STAGES]; + int32_t l0BEventList[STAGES]; + + // The id of current stage + uint32_t l1ListId{0}; + uint32_t l0AListId{0}; + uint32_t l0BListId{0}; + + TileMmad tileMmad; + CopyGmToL1A copyGmToL1A; + CopyGmToL1B copyGmToL1B; + CopyL1ToL0A copyL1ToL0A; + CopyL1ToL0B copyL1ToL0B; + CopyL0CToGm copyL0CToGm; +}; + +} // namespace Catlass::Gemm::Block + +#endif // CATLASS_KDA_BLOCK_MMAD_PINGPONG_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_pingpong_bias.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_pingpong_bias.hpp new file mode 100644 index 0000000000..f03d3a75d8 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_pingpong_bias.hpp @@ -0,0 +1,421 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_BLOCK_MMAD_PINGPONG_BIAS_HPP +#define CATLASS_KDA_BLOCK_MMAD_PINGPONG_BIAS_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/kda_coord.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" +#include "catlass/gemm/kda_gemm_helper.hpp" +#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" +#include "catlass/gemm/tile/kda_tile_mmad.hpp" + +namespace Catlass::Gemm::Block { + +template < + bool ENABLE_UNIT_FLAG_, + class L1TileShape_, + class L0TileShape_, + class AType_, + class BType_, + class CType_, + class BiasType_, + class TileCopy_, + class TileMmad_> +struct BlockMmad< + MmadAtlasA2PingpongBias, + L1TileShape_, + L0TileShape_, + AType_, + BType_, + CType_, + BiasType_, + TileCopy_, + TileMmad_> { +public: + // Type Aliases + using DispatchPolicy = MmadAtlasA2PingpongBias; + using ArchTag = typename DispatchPolicy::ArchTag; + using L1TileShape = L1TileShape_; + using L0TileShape = L0TileShape_; + using ElementA = typename AType_::Element; + using LayoutA = typename AType_::Layout; + using ElementB = typename BType_::Element; + using LayoutB = typename BType_::Layout; + using ElementC = typename CType_::Element; + using LayoutC = typename CType_::Layout; + using ElementBias = typename BiasType_::Element; + using LayoutBias = typename BiasType_::Layout; + using TileMmad = TileMmad_; + using CopyGmToL1A = typename TileCopy_::CopyGmToL1A; + using CopyGmToL1B = typename TileCopy_::CopyGmToL1B; + using CopyGmToL1Bias = typename TileCopy_::CopyGmToL1Bias; + using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; + using CopyL0CToGm = typename TileCopy_::CopyL0CToGm; + using ElementAccumulator = + typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; + using CopyL1ToBT = typename TileCopy_::CopyL1ToBT; + using LayoutAInL1 = typename CopyL1ToL0A::LayoutSrc; + using LayoutBInL1 = typename CopyL1ToL0B::LayoutSrc; + using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst; + using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst; + using LayoutCInL0 = layout::zN; + + using L1AAlignHelper = Gemm::helper::L1AlignHelper; + using L1BAlignHelper = Gemm::helper::L1AlignHelper; + + static constexpr bool ENABLE_UNIT_FLAG = DispatchPolicy::ENABLE_UNIT_FLAG; + static constexpr uint32_t STAGES = DispatchPolicy::STAGES; + static constexpr uint32_t L1A_SIZE = L1TileShape::M * L1TileShape::K * sizeof(ElementA); + static constexpr uint32_t L1B_SIZE = L1TileShape::N * L1TileShape::K * sizeof(ElementB); + static constexpr uint32_t L1BIAS_SIZE = L1TileShape::N * sizeof(ElementBias); + static constexpr uint32_t L0A_SIZE = ArchTag::L0A_SIZE; + static constexpr uint32_t L0B_SIZE = ArchTag::L0B_SIZE; + static constexpr uint32_t L0C_SIZE = ArchTag::L0C_SIZE; + static constexpr uint32_t BT_SIZE = ArchTag::BIAS_SIZE; + static constexpr uint32_t L0A_PINGPONG_BUF_SIZE = L0A_SIZE / STAGES; + static constexpr uint32_t L0B_PINGPONG_BUF_SIZE = L0B_SIZE / STAGES; + static constexpr uint32_t BIAS_BUF_SIZE = L0TileShape::N * sizeof(ElementAccumulator); + + // Check LayoutC + static_assert(std::is_same_v, "LayoutC only support RowMajor yet!"); + + // Check L1TileShape + static_assert((L1A_SIZE * STAGES + L1B_SIZE * STAGES + L1BIAS_SIZE) <= ArchTag::L1_SIZE, + "L1TileShape exceeding the L1 space!"); + + // Check L0TileShape + static constexpr uint32_t L0A_TILE_SIZE = L0TileShape::M * L0TileShape::K * sizeof(ElementA); + static constexpr uint32_t L0B_TILE_SIZE = L0TileShape::K * L0TileShape::N * sizeof(ElementB); + static constexpr uint32_t L0C_TILE_SIZE = L0TileShape::M * L0TileShape::N * sizeof(ElementAccumulator); + static_assert((L0A_TILE_SIZE * STAGES) <= L0A_SIZE, "L0TileShape exceeding the L0A space!"); + static_assert((L0B_TILE_SIZE * STAGES) <= L0B_SIZE, "L0TileShape exceeding the L0B space!"); + static_assert(L0C_TILE_SIZE <= L0C_SIZE, "L0TileShape exceeding the L0C space!"); + static_assert(BIAS_BUF_SIZE <= BT_SIZE, "BIAS_BUF_SIZE exceeding the BT space! Reduce L0TileShape::N"); + + static_assert(L1TileShape::M == L0TileShape::M && L1TileShape::N == L0TileShape::N, + "The situation where the basic blocks of L1 and L0 differ on the m and n axes is not supported yet"); + static_assert(L0TileShape::K <= L1TileShape::K, "L0TileShape::K cannot exceed L1TileShape::K"); + static_assert(!std::is_same_v, "Bias does not support bfloat16 yet"); + static_assert(Gemm::helper::TileShapeAlignChecker::_ALIGN == 256, + "Tile shape must be 32B aligned."); + + /// Construct + CATLASS_DEVICE + BlockMmad(Arch::Resource &resource, uint32_t l1BufAddrStart = 0) + { + uint32_t l1AOffset = l1BufAddrStart; + uint32_t l1BOffset = l1BufAddrStart + L1A_SIZE * STAGES; + uint32_t l1BiasOffset = l1BufAddrStart + L1A_SIZE * STAGES + L1B_SIZE * STAGES; + // Init buffers + for (uint32_t i = 0; i < STAGES; i++) { + // Assign L1/L0A/L0B space for each stages + l1ATensorList[i] = resource.l1Buf.template GetBufferByByte(l1AOffset + L1A_SIZE * i); + l1BTensorList[i] = resource.l1Buf.template GetBufferByByte(l1BOffset + L1B_SIZE * i); + l0ATensorList[i] = resource.l0ABuf.template GetBufferByByte(L0A_PINGPONG_BUF_SIZE * i); + l0BTensorList[i] = resource.l0BBuf.template GetBufferByByte(L0B_PINGPONG_BUF_SIZE * i); + + // Assign event ID for each stages + l1AEventList[i] = i; + l1BEventList[i] = i + STAGES; + l0AEventList[i] = i; + l0BEventList[i] = i + STAGES; + + // The event id that needs to be set before the loop + AscendC::SetFlag(l1AEventList[i]); + AscendC::SetFlag(l1BEventList[i]); + AscendC::SetFlag(l0AEventList[i]); + AscendC::SetFlag(l0BEventList[i]); + } + EVENT_ID_BIAS = 2 * STAGES; + l0CTensor = resource.l0CBuf.template GetBufferByByte(0); + l1BiasTensor = resource.l1Buf.template GetBufferByByte(l1BiasOffset); + l0BiasTensor = resource.btBuf.template GetBufferByByte(0); + AscendC::SetFlag(EVENT_ID0); + // Initialize bias independent event + AscendC::SetFlag(EVENT_ID_BIAS); + AscendC::SetFlag(EVENT_ID_BIAS); + } + + /// Destructor + CATLASS_DEVICE + ~BlockMmad() + { + for (uint32_t i = 0; i < STAGES; i++) { + AscendC::WaitFlag(l1AEventList[i]); + AscendC::WaitFlag(l1BEventList[i]); + AscendC::WaitFlag(l0AEventList[i]); + AscendC::WaitFlag(l0BEventList[i]); + } + AscendC::WaitFlag(EVENT_ID0); + // Wait for bias independent event + AscendC::WaitFlag(EVENT_ID_BIAS); + AscendC::WaitFlag(EVENT_ID_BIAS); + } + + /// Perform a block-scoped matrix multiply-accumulate + CATLASS_DEVICE + void operator()( + AscendC::GlobalTensor const &gmA, LayoutA const &layoutA, + AscendC::GlobalTensor const &gmB, LayoutB const &layoutB, + AscendC::GlobalTensor const &gmC, LayoutC const &layoutC, + AscendC::GlobalTensor const &gmBias, + GemmCoord const &actualShape) + { + uint32_t mRound = RoundUp(actualShape.m()); + uint32_t nRound = RoundUp(actualShape.n()); + + auto layoutAInL1 = LayoutAInL1::template MakeLayout(L1TileShape::M, L1TileShape::K); + auto layoutBInL1 = LayoutBInL1::template MakeLayout(L1TileShape::K, L1TileShape::N); + auto layoutBiasInL1 = layout::VectorLayout(L1TileShape::N); + auto layoutBiasInL0 = layout::VectorLayout(L0TileShape::N); + auto layoutInL0C = LayoutCInL0::MakeLayoutInL0C(MakeCoord(mRound, nRound)); + + uint32_t kActual = min(actualShape.k(), L1TileShape::K); + + // load first matrix A tile from GM to L1 + AscendC::WaitFlag(l1AEventList[l1ListId]); + auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShape.m(), kActual)); + copyGmToL1A(l1ATensorList[l1ListId], gmA, layoutAInL1, layoutTileA); + AscendC::SetFlag(l1AEventList[l1ListId]); + + // load first matrix B tile from GM to L1 + AscendC::WaitFlag(l1BEventList[l1ListId]); + auto layoutTileB = layoutB.GetTileLayout(MakeCoord(kActual, actualShape.n())); + copyGmToL1B(l1BTensorList[l1ListId], gmB, layoutBInL1, layoutTileB); + AscendC::SetFlag(l1BEventList[l1ListId]); + + // load bias from GM to L1 independently + AscendC::WaitFlag(EVENT_ID_BIAS); + auto layoutTileBias = layout::VectorLayout(actualShape.n()); + copyGmToL1Bias(l1BiasTensor, gmBias, layoutBiasInL1, layoutTileBias); + AscendC::SetFlag(EVENT_ID_BIAS); + + if constexpr (!ENABLE_UNIT_FLAG) { + AscendC::WaitFlag(EVENT_ID0); + } + + uint32_t mPartLoop = CeilDiv(mRound); + uint32_t nPartLoop = CeilDiv(nRound); + + // main loop + uint32_t kTileCount = CeilDiv(actualShape.k()); + for (uint32_t kLoopIdx = 0; kLoopIdx < kTileCount; kLoopIdx++) { + uint32_t l1ListIdNext = (l1ListId + 1 < STAGES) ? (l1ListId + 1) : 0; + uint32_t kActualNext{0}; + // preload next tile from GM to L1 + if (kLoopIdx < kTileCount - 1) { + uint32_t kLoopIdxNext = kLoopIdx + 1; + kActualNext = (kLoopIdxNext < kTileCount - 1) ? + L1TileShape::K : + (actualShape.k() - kLoopIdxNext * L1TileShape::K); + + // Get L1 tensor for next stage + auto l1ATensor = l1ATensorList[l1ListIdNext]; + auto l1BTensor = l1BTensorList[l1ListIdNext]; + // Get GM tile for next stage + MatrixCoord gmTileAOffset{0, kLoopIdxNext * L1TileShape::K}; + MatrixCoord gmTileBOffset{kLoopIdxNext * L1TileShape::K, 0}; + auto gmTileA = gmA[layoutA.GetOffset(gmTileAOffset)]; + auto gmTileB = gmB[layoutB.GetOffset(gmTileBOffset)]; + + // load next matrix A tile from GM to L1 + AscendC::WaitFlag(l1AEventList[l1ListIdNext]); + layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShape.m(), kActualNext)); + copyGmToL1A(l1ATensor, gmTileA, layoutAInL1, layoutTileA); + AscendC::SetFlag(l1AEventList[l1ListIdNext]); + + // load next matrix B tile from GM to L1 + AscendC::WaitFlag(l1BEventList[l1ListIdNext]); + layoutTileB = layoutB.GetTileLayout(MakeCoord(kActualNext, actualShape.n())); + copyGmToL1B(l1BTensor, gmTileB, layoutBInL1, layoutTileB); + AscendC::SetFlag(l1BEventList[l1ListIdNext]); + } + + // Get L1 tensor for current stage + auto l1ATensor = l1ATensorList[l1ListId]; + auto l1BTensor = l1BTensorList[l1ListId]; + + // Get the loop nums on L0 + uint32_t kPartLoop = CeilDiv(kActual); + + for (int mPartIdx = 0; mPartIdx < mPartLoop; mPartIdx++) { + uint32_t mPartActual = (mPartIdx < mPartLoop - 1) ? + L0TileShape::M : + (mRound - mPartIdx * L0TileShape::M); + + for (int kPartIdx = 0; kPartIdx < kPartLoop; kPartIdx++) { + uint32_t kPartActual = (kPartIdx < kPartLoop - 1) ? + L0TileShape::K : + (kActual - kPartIdx * L0TileShape::K); + + // Locate the current tile on L0A + auto l0ATile = l0ATensorList[l0AListId]; + LayoutAInL0 layoutAInL0 = LayoutAInL0::template MakeLayout(mPartActual, kPartActual); + // Locate the current tile of matrix A on L1 + MatrixCoord l1AOffset{mPartIdx * L0TileShape::M, kPartIdx * L0TileShape::K}; + auto l1ATile = l1ATensor[layoutAInL1.GetOffset(l1AOffset)]; + + AscendC::WaitFlag(l0AEventList[l0AListId]); + if ((mPartIdx == 0) && (kPartIdx == 0)) { + AscendC::WaitFlag(l1AEventList[l1ListId]); + } + + // Load current tile from L1 to L0A + copyL1ToL0A(l0ATile, l1ATile, layoutAInL0, layoutAInL1); + + if ((mPartIdx == mPartLoop - 1) && (kPartIdx == kPartLoop - 1)) { + AscendC::SetFlag(l1AEventList[l1ListId]); + } + + for (int nPartIdx = 0; nPartIdx < nPartLoop; nPartIdx++) { + uint32_t nPartActual = (nPartIdx < nPartLoop - 1) ? + L0TileShape::N : + (nRound - nPartIdx * L0TileShape::N); + + // Locate the current tile on L0B + auto l0BTile = l0BTensorList[l0BListId]; + LayoutBInL0 layoutBInL0 = LayoutBInL0::template MakeLayout(kPartActual, nPartActual); + // Locate the current tile of matrix B on L1 + MatrixCoord l1BOffset{kPartIdx * L0TileShape::K, nPartIdx * L0TileShape::N}; + auto l1BTile = l1BTensor[layoutBInL1.GetOffset(l1BOffset)]; + + // Wait for mmad finished + AscendC::WaitFlag(l0BEventList[l0BListId]); + // If the current tile is the first one on the k&n axis, wait for loading matrix B from GM to L1 + if ((kPartIdx == 0) && (nPartIdx == 0)) { + AscendC::WaitFlag(l1BEventList[l1ListId]); + } + + // Load current tile from L1 to L0B + copyL1ToL0B(l0BTile, l1BTile, layoutBInL0, layoutBInL1); + + // If the current tile is the last one on the k&n axis, notify to load matrix B from GM to L1 + if ((kPartIdx == kPartLoop - 1) && (nPartIdx == nPartLoop - 1)) { + AscendC::SetFlag(l1BEventList[l1ListId]); + } + + // If the current tile is the first tile on the k axis, the accumulator needs to be reset to 0 + bool initC = ((kLoopIdx == 0) && (kPartIdx == 0)); + + // Load bias to l0 biastable only when initC is true (kLoopIdx==0 && kPartIdx==0) + // Load L0_TILE_N length data in nPartIdx loop, starting from nPartIdx * L0_TILE_N + if (initC) { + if (nPartIdx == 0) { + AscendC::WaitFlag(EVENT_ID_BIAS); + } + AscendC::WaitFlag(EVENT_ID_BIAS); + // Create layout for actual nPartActual size + auto layoutBiasInL0Actual = layout::VectorLayout(nPartActual); + // Get L1 bias tile starting from nPartIdx * L0_TILE_N + // For VectorLayout, directly use index offset + auto l1BiasTile = l1BiasTensor[nPartIdx * L0TileShape::N]; + // Create layout for L1 source with offset + auto layoutBiasInL1Offset = layout::VectorLayout(L1TileShape::N); + copyL1ToBT(l0BiasTensor, l1BiasTile, layoutBiasInL0Actual, layoutBiasInL1Offset); + if (nPartIdx == nPartLoop - 1) { + AscendC::SetFlag(EVENT_ID_BIAS); + } + } + // Notify to do mmad + AscendC::SetFlag(EVENT_ID0); + + // Locate the current tile on L0C + MatrixCoord l0COffset{mPartIdx * L0TileShape::M, nPartIdx * L0TileShape::N}; + auto l0CTile = l0CTensor[layoutInL0C.GetOffset(l0COffset)]; + + // Compute the matrix multiplication on L0A and L0B and write the result to the accumulator + // Wait for loading L0B + AscendC::WaitFlag(EVENT_ID0); + // If the unit flag is enabled, the unit flag is set according to the calculation progress + uint8_t unitFlag = 0b00; + if constexpr (ENABLE_UNIT_FLAG) { + if ((kLoopIdx == kTileCount - 1) && (mPartIdx == mPartLoop - 1) && + (kPartIdx == kPartLoop - 1) && (nPartIdx == nPartLoop - 1)) { + unitFlag = 0b11; + } else { + unitFlag = 0b10; + } + } + // Perform calculation operations + if (initC) { + tileMmad(l0CTile, l0ATile, l0BTile, l0BiasTensor, mPartActual, nPartActual, kPartActual, + initC, unitFlag); + AscendC::SetFlag(EVENT_ID_BIAS); + } else { + tileMmad(l0CTile, l0ATile, l0BTile, mPartActual, nPartActual, kPartActual, initC, unitFlag); + } + + // Notify to move the next L0B tile + AscendC::SetFlag(l0BEventList[l0BListId]); + l0BListId = (l0BListId + 1 < STAGES) ? (l0BListId + 1) : 0; + } + AscendC::SetFlag(l0AEventList[l0AListId]); + l0AListId = (l0AListId + 1 < STAGES) ? (l0AListId + 1) : 0; + } + } + l1ListId = l1ListIdNext; + kActual = kActualNext; + } + + // copy block out + LayoutC layoutBlock = layoutC.GetTileLayout(actualShape.GetCoordMN()); + + if constexpr (!ENABLE_UNIT_FLAG) { + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + copyL0CToGm(gmC, l0CTensor, layoutBlock, layoutInL0C); + AscendC::SetFlag(EVENT_ID0); + } else { + copyL0CToGm(gmC, l0CTensor, layoutBlock, layoutInL0C, 0b11); + } + } + +protected: + // Multi-stage tensors list + AscendC::LocalTensor l1ATensorList[STAGES]; + AscendC::LocalTensor l1BTensorList[STAGES]; + AscendC::LocalTensor l1BiasTensor; + AscendC::LocalTensor l0ATensorList[STAGES]; + AscendC::LocalTensor l0BTensorList[STAGES]; + AscendC::LocalTensor l0CTensor; + AscendC::LocalTensor l0BiasTensor; + + // Multi-stage event id list + int32_t l1AEventList[STAGES]; + int32_t l1BEventList[STAGES]; + int32_t l0AEventList[STAGES]; + int32_t l0BEventList[STAGES]; + + // Event ID for Bias + int32_t EVENT_ID_BIAS; + + // The id of current stage + uint32_t l1ListId{0}; + uint32_t l0AListId{0}; + uint32_t l0BListId{0}; + + TileMmad tileMmad; + CopyGmToL1A copyGmToL1A; + CopyGmToL1B copyGmToL1B; + CopyGmToL1Bias copyGmToL1Bias; + CopyL1ToL0A copyL1ToL0A; + CopyL1ToL0B copyL1ToL0B; + CopyL1ToBT copyL1ToBT; + CopyL0CToGm copyL0CToGm; +}; + +} // namespace Catlass::Gemm::Block + +#endif // CATLASS_KDA_BLOCK_MMAD_PINGPONG_BIAS_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_preload.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_preload.hpp new file mode 100644 index 0000000000..b720d161a6 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_preload.hpp @@ -0,0 +1,402 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_BLOCK_MMAD_PRELOAD_HPP +#define CATLASS_KDA_BLOCK_MMAD_PRELOAD_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/kda_coord.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" +#include "catlass/gemm/kda_gemm_helper.hpp" +#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" +#include "catlass/gemm/tile/kda_tile_mmad.hpp" + +namespace Catlass::Gemm::Block { + +template < + bool ENABLE_UNIT_FLAG_, + bool ENABLE_SHUFFLE_K_, + class L1TileShape_, + class L0TileShape_, + class AType_, + class BType_, + class CType_, + class BiasType_, + class TileCopy_, + class TileMmad_> +struct BlockMmad< + MmadAtlasA2Preload, + L1TileShape_, + L0TileShape_, + AType_, + BType_, + CType_, + BiasType_, + TileCopy_, + TileMmad_> { +public: + // Type Aliases + using DispatchPolicy = MmadAtlasA2Preload; + using ArchTag = typename DispatchPolicy::ArchTag; + using L1TileShape = L1TileShape_; + using L0TileShape = L0TileShape_; + using ElementA = typename AType_::Element; + using LayoutA = typename AType_::Layout; + using ElementB = typename BType_::Element; + using LayoutB = typename BType_::Layout; + using ElementC = typename CType_::Element; + using LayoutC = typename CType_::Layout; + using TileMmad = TileMmad_; + using CopyGmToL1A = typename TileCopy_::CopyGmToL1A; + using CopyGmToL1B = typename TileCopy_::CopyGmToL1B; + using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; + using CopyL0CToGm = typename TileCopy_::CopyL0CToGm; + using ElementAccumulator = + typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; + using LayoutAInL1 = typename CopyL1ToL0A::LayoutSrc; + using LayoutBInL1 = typename CopyL1ToL0B::LayoutSrc; + using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst; + using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst; + using LayoutCInL0 = layout::zN; + + using L1AAlignHelper = Gemm::helper::L1AlignHelper; + using L1BAlignHelper = Gemm::helper::L1AlignHelper; + + static constexpr bool ENABLE_UNIT_FLAG = DispatchPolicy::ENABLE_UNIT_FLAG; + static constexpr bool ENABLE_SHUFFLE_K = DispatchPolicy::ENABLE_SHUFFLE_K; + static constexpr uint32_t STAGES = DispatchPolicy::STAGES; + static constexpr uint32_t L1A_SIZE = L1TileShape::M * L1TileShape::K * sizeof(ElementA); + static constexpr uint32_t L1B_SIZE = L1TileShape::N * L1TileShape::K * sizeof(ElementB); + static constexpr uint32_t L0A_SIZE = ArchTag::L0A_SIZE; + static constexpr uint32_t L0B_SIZE = ArchTag::L0B_SIZE; + static constexpr uint32_t L0C_SIZE = ArchTag::L0C_SIZE; + static constexpr uint32_t L0A_PINGPONG_BUF_SIZE = L0A_SIZE / STAGES; + static constexpr uint32_t L0B_PINGPONG_BUF_SIZE = L0B_SIZE / STAGES; + + // Check LayoutC + static_assert(std::is_same_v, "LayoutC only support RowMajor yet!"); + + // Check L1TileShape + static_assert((L1A_SIZE * STAGES + L1B_SIZE * STAGES) <= ArchTag::L1_SIZE, "L1TileShape exceeding the L1 space!"); + + // Check L0TileShape + static constexpr uint32_t L0A_TILE_SIZE = L0TileShape::M * L0TileShape::K * sizeof(ElementA); + static constexpr uint32_t L0B_TILE_SIZE = L0TileShape::K * L0TileShape::N * sizeof(ElementB); + static constexpr uint32_t L0C_TILE_SIZE = L0TileShape::M * L0TileShape::N * sizeof(ElementAccumulator); + static_assert((L0A_TILE_SIZE * STAGES) <= L0A_SIZE, "L0TileShape exceeding the L0A space!"); + static_assert((L0B_TILE_SIZE * STAGES) <= L0B_SIZE, "L0TileShape exceeding the L0B space!"); + static_assert(L0C_TILE_SIZE <= L0C_SIZE, "L0TileShape exceeding the L0C space!"); + + static_assert(L1TileShape::M == L0TileShape::M && L1TileShape::N == L0TileShape::N, + "The situation where the basic blocks of L1 and L0 differ on the m and n axes is not supported yet"); + static_assert(L0TileShape::K <= L1TileShape::K, "L0TileShape::K cannot exceed L1TileShape::K"); + // 32B (256b) aligned + static_assert(Gemm::helper::TileShapeAlignChecker::_ALIGN == 256, + "Tile shape must be 32B aligned."); + + /// Construct + CATLASS_DEVICE + BlockMmad(Arch::Resource &resource, uint32_t l1BufAddrStart = 0) + { + uint32_t l1AOffset = l1BufAddrStart; + uint32_t l1BOffset = l1BufAddrStart + L1A_SIZE * STAGES; + // Init buffers + for (uint32_t i = 0; i < STAGES; i++) { + l1ATensorList[i] = resource.l1Buf.template GetBufferByByte(l1AOffset + L1A_SIZE * i); + l1BTensorList[i] = resource.l1Buf.template GetBufferByByte(l1BOffset + L1B_SIZE * i); + l0ATensorList[i] = resource.l0ABuf.template GetBufferByByte(L0A_PINGPONG_BUF_SIZE * i); + l0BTensorList[i] = resource.l0BBuf.template GetBufferByByte(L0B_PINGPONG_BUF_SIZE * i); + + l1AEventList[i] = i; + l1BEventList[i] = i + STAGES; + l0AEventList[i] = i; + l0BEventList[i] = i + STAGES; + AscendC::SetFlag(l1AEventList[i]); + AscendC::SetFlag(l1BEventList[i]); + AscendC::SetFlag(l0AEventList[i]); + AscendC::SetFlag(l0BEventList[i]); + } + l0CTensor = resource.l0CBuf.template GetBufferByByte(0); + AscendC::SetFlag(EVENT_ID0); + } + + /// Destructor + CATLASS_DEVICE + ~BlockMmad() + { + for (uint32_t i = 0; i < STAGES; i++) { + AscendC::WaitFlag(l1AEventList[i]); + AscendC::WaitFlag(l1BEventList[i]); + AscendC::WaitFlag(l0AEventList[i]); + AscendC::WaitFlag(l0BEventList[i]); + } + AscendC::WaitFlag(EVENT_ID0); + } + + /// Perform a block-scoped matrix multiply-accumulate + CATLASS_DEVICE + void operator()( + AscendC::GlobalTensor const &gmBlockA, LayoutA const &layoutA, + AscendC::GlobalTensor const &gmBlockB, LayoutB const &layoutB, + AscendC::GlobalTensor const &gmBlockC, LayoutC const &layoutC, + AscendC::GlobalTensor const &gmNextBlockA, + AscendC::GlobalTensor const &gmNextBlockB, + GemmCoord const &actualShape, GemmCoord const &actualShapeNext, + bool isFirstBlock, bool hasNextBlock) + { + uint32_t mRound = RoundUp(actualShape.m()); + uint32_t nRound = RoundUp(actualShape.n()); + + auto layoutAInL1 = LayoutAInL1::template MakeLayout(L1TileShape::M, L1TileShape::K); + auto layoutBInL1 = LayoutBInL1::template MakeLayout(L1TileShape::K, L1TileShape::N); + auto layoutInL0C = LayoutCInL0::MakeLayoutInL0C(MakeCoord(mRound, nRound)); + + uint32_t kTileCount = CeilDiv(actualShape.k()); + uint32_t kTileCountNext = CeilDiv(actualShapeNext.k()); + + if constexpr (!ENABLE_UNIT_FLAG) { + AscendC::WaitFlag(EVENT_ID0); + } + uint32_t startTileIdx = 0; + if constexpr (ENABLE_SHUFFLE_K) { + startTileIdx = AscendC::GetBlockIdx(); + } + uint32_t firstTileIdx = startTileIdx % kTileCount; + uint32_t lastTileIdx = (startTileIdx + kTileCount - 1) % kTileCount; + uint32_t kActual = + (firstTileIdx < kTileCount - 1) ? L1TileShape::K : (actualShape.k() - firstTileIdx * L1TileShape::K); + uint32_t firstTileIdxNext = startTileIdx % kTileCountNext; + + uint32_t mPartLoop = CeilDiv(mRound); + uint32_t nPartLoop = CeilDiv(nRound); + + // k loop + for (uint32_t kLoopIdx = 0; kLoopIdx < kTileCount; kLoopIdx++) { + uint32_t shuffleKIdx = (startTileIdx + kLoopIdx) % kTileCount; + // Load first matrix A tile in total kernel loop from GM to L1 + if (shuffleKIdx == firstTileIdx && isFirstBlock) { + MatrixCoord gmTileAOffset{0, shuffleKIdx * L1TileShape::K}; + MatrixCoord gmTileBOffset{shuffleKIdx * L1TileShape::K, 0}; + auto gmTileA = gmBlockA[layoutA.GetOffset(gmTileAOffset)]; + auto gmTileB = gmBlockB[layoutB.GetOffset(gmTileBOffset)]; + // Load first matrix A tile from GM to L1 + AscendC::WaitFlag(l1AEventList[l1ListId]); + auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShape.m(), kActual)); + copyGmToL1A(l1ATensorList[l1ListId], gmTileA, layoutAInL1, layoutTileA); + AscendC::SetFlag(l1AEventList[l1ListId]); + + // Load first matrix B tile from GM to L1 + AscendC::WaitFlag(l1BEventList[l1ListId]); + auto layoutTileB = layoutB.GetTileLayout(MakeCoord(kActual, actualShape.n())); + copyGmToL1B(l1BTensorList[l1ListId], gmTileB, layoutBInL1, layoutTileB); + AscendC::SetFlag(l1BEventList[l1ListId]); + } + + uint32_t l1ListIdNext = (l1ListId + 1 < STAGES) ? (l1ListId + 1) : 0; + uint32_t kActualNext{0}; + + // preload next tile from GM to L1 + if (shuffleKIdx != lastTileIdx) { + uint32_t shuffleKIdxNext = (startTileIdx + kLoopIdx + 1) % kTileCount; + // Get L1 tensor for next stage + auto l1ATensor = l1ATensorList[l1ListIdNext]; + auto l1BTensor = l1BTensorList[l1ListIdNext]; + // Get GM tensor for next stage + kActualNext = (shuffleKIdxNext < kTileCount - 1) ? + L1TileShape::K : + (actualShape.k() - shuffleKIdxNext * L1TileShape::K); + MatrixCoord gmTileAOffset{0, shuffleKIdxNext * L1TileShape::K}; + MatrixCoord gmTileBOffset{shuffleKIdxNext * L1TileShape::K, 0}; + auto gmTileA = gmBlockA[layoutA.GetOffset(gmTileAOffset)]; + auto gmTileB = gmBlockB[layoutB.GetOffset(gmTileBOffset)]; + + // load next matrix A tile from GM to L1 + AscendC::WaitFlag(l1AEventList[l1ListIdNext]); + auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShape.m(), kActualNext)); + copyGmToL1A(l1ATensor, gmTileA, layoutAInL1, layoutTileA); + AscendC::SetFlag(l1AEventList[l1ListIdNext]); + + // load next matrix B tile from GM to L1 + AscendC::WaitFlag(l1BEventList[l1ListIdNext]); + auto layoutTileB = layoutB.GetTileLayout(MakeCoord(kActualNext, actualShape.n())); + copyGmToL1B(l1BTensor, gmTileB, layoutBInL1, layoutTileB); + AscendC::SetFlag(l1BEventList[l1ListIdNext]); + } + if (shuffleKIdx == lastTileIdx && hasNextBlock) { + // Get L1 tensor for next stage + auto l1ATensor = l1ATensorList[l1ListIdNext]; + auto l1BTensor = l1BTensorList[l1ListIdNext]; + // Get GM tensor for next stage + kActualNext = (firstTileIdxNext < kTileCountNext - 1) ? + L1TileShape::K : + (actualShapeNext.k() - firstTileIdxNext * L1TileShape::K); + MatrixCoord gmTileAOffset{0, firstTileIdxNext * L1TileShape::K}; + MatrixCoord gmTileBOffset{firstTileIdxNext * L1TileShape::K, 0}; + auto gmTileA = gmNextBlockA[layoutA.GetOffset(gmTileAOffset)]; + auto gmTileB = gmNextBlockB[layoutB.GetOffset(gmTileBOffset)]; + // load next matrix A tile from GM to L1 + AscendC::WaitFlag(l1AEventList[l1ListIdNext]); + auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShapeNext.m(), kActualNext)); + copyGmToL1A(l1ATensor, gmTileA, layoutAInL1, layoutTileA); + AscendC::SetFlag(l1AEventList[l1ListIdNext]); + + // load next matrix B tile from GM to L1 + AscendC::WaitFlag(l1BEventList[l1ListIdNext]); + auto layoutTileB = layoutB.GetTileLayout(MakeCoord(kActualNext, actualShapeNext.n())); + copyGmToL1B(l1BTensor, gmTileB, layoutBInL1, layoutTileB); + AscendC::SetFlag(l1BEventList[l1ListIdNext]); + } + + // Get L1 tensor for current stage + auto l1ATensor = l1ATensorList[l1ListId]; + auto l1BTensor = l1BTensorList[l1ListId]; + + // Get the loop nums on L0 + uint32_t kPartLoop = CeilDiv(kActual); + + uint32_t l0ABufId = 0; + uint32_t l0BBufId = 0; + + for (int mPartIdx = 0; mPartIdx < mPartLoop; mPartIdx++) { + uint32_t mPartActual = (mPartIdx < mPartLoop - 1) ? + L0TileShape::M : + (mRound - mPartIdx * L0TileShape::M); + + for (int kPartIdx = 0; kPartIdx < kPartLoop; kPartIdx++) { + uint32_t kPartActual = (kPartIdx < kPartLoop - 1) ? + L0TileShape::K : + (kActual - kPartIdx * L0TileShape::K); + + // Locate the current tile on L0A + auto l0ATile = l0ATensorList[l0ABufId]; + LayoutAInL0 layoutAInL0 = LayoutAInL0::template MakeLayout(mPartActual, kPartActual); + // Locate the current tile of matrix A on L1 + MatrixCoord l1AOffset{mPartIdx * L0TileShape::M, kPartIdx * L0TileShape::K}; + auto l1ATile = l1ATensor[layoutAInL1.GetOffset(l1AOffset)]; + + AscendC::WaitFlag(l0AEventList[l0ABufId]); + if ((mPartIdx == 0) && (kPartIdx == 0)) { + AscendC::WaitFlag(l1AEventList[l1ListId]); + } + + copyL1ToL0A(l0ATile, l1ATile, layoutAInL0, layoutAInL1); + + if ((mPartIdx == mPartLoop - 1) && (kPartIdx == kPartLoop - 1)) { + AscendC::SetFlag(l1AEventList[l1ListId]); + } + + for (int nPartIdx = 0; nPartIdx < nPartLoop; nPartIdx++) { + uint32_t nPartActual = (nPartIdx < nPartLoop - 1) ? + L0TileShape::N : + (nRound - nPartIdx * L0TileShape::N); + + // Locate the current tile on L0B + auto l0BTile = l0BTensorList[l0BBufId]; + LayoutBInL0 layoutBInL0 = LayoutBInL0::template MakeLayout(kPartActual, nPartActual); + // Locate the current tile of matrix B on L1 + MatrixCoord l1BOffset{kPartIdx * L0TileShape::K, nPartIdx * L0TileShape::N}; + auto l1BTile = l1BTensor[layoutBInL1.GetOffset(l1BOffset)]; + + // Wait for mmad finished + AscendC::WaitFlag(l0BEventList[l0BBufId]); + // If the current tile is the first one on the k&n axis, wait for loading matrix B from GM to L1 + if ((kPartIdx == 0) && (nPartIdx == 0)) { + AscendC::WaitFlag(l1BEventList[l1ListId]); + } + + // Load current tile from L1 to L0B + copyL1ToL0B(l0BTile, l1BTile, layoutBInL0, layoutBInL1); + + // If the current tile is the last one on the k&n axis, notify to load matrix B from GM to L1 + if ((kPartIdx == kPartLoop - 1) && (nPartIdx == nPartLoop - 1)) { + AscendC::SetFlag(l1BEventList[l1ListId]); + } + // Notify to do mmad + AscendC::SetFlag(EVENT_ID0); + + // Locate the current tile on L0C + MatrixCoord l0COffset{mPartIdx * L0TileShape::M, nPartIdx * L0TileShape::N}; + auto l0CTile = l0CTensor[layoutInL0C.GetOffset(l0COffset)]; + + // Compute the matrix multiplication on L0A and L0B and write the result to the accumulator + // Wait for loading L0B + AscendC::WaitFlag(EVENT_ID0); + + // If the current tile is the first tile on the k axis, the accumulator needs to be reset to 0 + bool initC = ((kLoopIdx == 0) && (kPartIdx == 0)); + // If the unit flag is enabled, the unit flag is set according to the calculation progress + uint8_t unitFlag = 0b00; + if constexpr (ENABLE_UNIT_FLAG) { + if ((kLoopIdx == kTileCount - 1) && (mPartIdx == mPartLoop - 1) && + (kPartIdx == kPartLoop - 1) && (nPartIdx == nPartLoop - 1)) { + unitFlag = 0b11; + } else { + unitFlag = 0b10; + } + } + // Perform calculation operations + tileMmad(l0CTile, l0ATile, l0BTile, mPartActual, nPartActual, kPartActual, initC, unitFlag); + + // Notify to move the next L0B tile + AscendC::SetFlag(l0BEventList[l0BBufId]); + + l0BBufId = (l0BBufId + 1 < STAGES) ? (l0BBufId + 1) : 0; + } + AscendC::SetFlag(l0AEventList[l0ABufId]); + l0ABufId = (l0ABufId + 1 < STAGES) ? (l0ABufId + 1) : 0; + } + } + + l1ListId = l1ListIdNext; + kActual = kActualNext; + } + + // copy block out + LayoutC layoutBlock = layoutC.GetTileLayout(actualShape.GetCoordMN()); + + if constexpr (!ENABLE_UNIT_FLAG) { + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + copyL0CToGm(gmBlockC, l0CTensor, layoutBlock, layoutInL0C); + AscendC::SetFlag(EVENT_ID0); + } else { + copyL0CToGm(gmBlockC, l0CTensor, layoutBlock, layoutInL0C, 0b11); + } + } + +protected: + /// Data members + AscendC::LocalTensor l1ATensorList[STAGES]; + AscendC::LocalTensor l1BTensorList[STAGES]; + AscendC::LocalTensor l0ATensorList[STAGES]; + AscendC::LocalTensor l0BTensorList[STAGES]; + AscendC::LocalTensor l0CTensor; + + int32_t l1AEventList[STAGES]; + int32_t l1BEventList[STAGES]; + int32_t l0AEventList[STAGES]; + int32_t l0BEventList[STAGES]; + + uint32_t l1ListId{0}; + + TileMmad tileMmad; + CopyGmToL1A copyGmToL1A; + CopyGmToL1B copyGmToL1B; + CopyL1ToL0A copyL1ToL0A; + CopyL1ToL0B copyL1ToL0B; + CopyL0CToGm copyL0CToGm; +}; + +} // namespace Catlass::Gemm::Block + +#endif // CATLASS_KDA_BLOCK_MMAD_PRELOAD_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_preload_async.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_preload_async.hpp new file mode 100644 index 0000000000..490cb3b11b --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_preload_async.hpp @@ -0,0 +1,448 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_BLOCK_MMAD_PRELOAD_ASYNC_HPP +#define CATLASS_KDA_BLOCK_MMAD_PRELOAD_ASYNC_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/kda_coord.hpp" +#include "catlass/detail/kda_callback.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" +#include "catlass/gemm/kda_gemm_helper.hpp" +#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" +#include "catlass/gemm/tile/kda_tile_mmad.hpp" + +namespace Catlass::Gemm::Block { + +template < + uint32_t PRELOAD_STAGES_, + uint32_t L1_STAGES_, + uint32_t L0A_STAGES_, + uint32_t L0B_STAGES_, + uint32_t L0C_STAGES_, + bool ENABLE_UNIT_FLAG_, + bool ENABLE_SHUFFLE_K_, + class L1TileShape_, + class L0TileShape_, + class AType_, + class BType_, + class CType_, + class BiasType_, + class TileCopy_, + class TileMmad_> +struct BlockMmad< + MmadAtlasA2PreloadAsync< + PRELOAD_STAGES_, + L1_STAGES_, + L0A_STAGES_, + L0B_STAGES_, + L0C_STAGES_, + ENABLE_UNIT_FLAG_, + ENABLE_SHUFFLE_K_>, + L1TileShape_, + L0TileShape_, + AType_, + BType_, + CType_, + BiasType_, + TileCopy_, + TileMmad_> { +public: + // Type Aliases + using DispatchPolicy = MmadAtlasA2PreloadAsync< + PRELOAD_STAGES_, + L1_STAGES_, + L0A_STAGES_, + L0B_STAGES_, + L0C_STAGES_, + ENABLE_UNIT_FLAG_, + ENABLE_SHUFFLE_K_>; + using ArchTag = typename DispatchPolicy::ArchTag; + using L1TileShape = L1TileShape_; + using L0TileShape = L0TileShape_; + using ElementA = typename AType_::Element; + using LayoutA = typename AType_::Layout; + using ElementB = typename BType_::Element; + using LayoutB = typename BType_::Layout; + using ElementC = typename CType_::Element; + using LayoutC = typename CType_::Layout; + using TileMmad = TileMmad_; + using CopyGmToL1A = typename TileCopy_::CopyGmToL1A; + using CopyGmToL1B = typename TileCopy_::CopyGmToL1B; + using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; + using CopyL0CToGm = typename TileCopy_::CopyL0CToGm; + using ElementAccumulator = + typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; + using LayoutAInL1 = typename CopyL1ToL0A::LayoutSrc; + using LayoutBInL1 = typename CopyL1ToL0B::LayoutSrc; + using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst; + using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst; + using LayoutCInL0 = layout::zN; + + using L1AAlignHelper = Gemm::helper::L1AlignHelper; + using L1BAlignHelper = Gemm::helper::L1AlignHelper; + + static constexpr uint32_t PRELOAD_STAGES = DispatchPolicy::PRELOAD_STAGES; + static constexpr uint32_t L1_STAGES = DispatchPolicy::L1_STAGES; + static constexpr uint32_t L0A_STAGES = DispatchPolicy::L0A_STAGES; + static constexpr uint32_t L0B_STAGES = DispatchPolicy::L0B_STAGES; + static constexpr uint32_t L0C_STAGES = DispatchPolicy::L0C_STAGES; + + static constexpr bool ENABLE_UNIT_FLAG = DispatchPolicy::ENABLE_UNIT_FLAG; + static constexpr bool ENABLE_SHUFFLE_K = DispatchPolicy::ENABLE_SHUFFLE_K; + + // L1 tile size + static constexpr uint32_t L1A_TILE_SIZE = L1TileShape::M * L1TileShape::K * sizeof(ElementA); + static constexpr uint32_t L1B_TILE_SIZE = L1TileShape::N * L1TileShape::K * sizeof(ElementB); + // L0 tile size + static constexpr uint32_t L0A_TILE_SIZE = L0TileShape::M * L0TileShape::K * sizeof(ElementA); + static constexpr uint32_t L0B_TILE_SIZE = L0TileShape::K * L0TileShape::N * sizeof(ElementB); + static constexpr uint32_t L0C_TILE_SIZE = L1TileShape::M * L1TileShape::N * sizeof(ElementAccumulator); + + // Check LayoutC + static_assert(std::is_same_v, "LayoutC only support RowMajor yet!"); + + // Check L1TileShape + static_assert((L1A_TILE_SIZE + L1B_TILE_SIZE) * L1_STAGES <= ArchTag::L1_SIZE, + "L1TileShape exceeding the L1 space!"); + + // Check L0TileShape + static_assert(L0A_TILE_SIZE * L0A_STAGES <= ArchTag::L0A_SIZE, "L0TileShape exceeding the L0A space!"); + static_assert(L0B_TILE_SIZE * L0B_STAGES <= ArchTag::L0B_SIZE, "L0TileShape exceeding the L0B space!"); + static_assert(L0C_TILE_SIZE * L0C_STAGES <= ArchTag::L0C_SIZE, "L0TileShape exceeding the L0C space!"); + + static_assert(L1TileShape::M == L0TileShape::M && L1TileShape::N == L0TileShape::N, + "The situation where the basic blocks of L1 and L0 differ on the m and n axes is not supported yet"); + static_assert(L0TileShape::K <= L1TileShape::K, "L0TileShape::K cannot exceed L1TileShape::K"); + // 32B (256b) aligned + static_assert(Gemm::helper::TileShapeAlignChecker::_ALIGN == 256, + "Tile shape must be 32B aligned."); + + static constexpr auto L1A_LAYOUT = LayoutAInL1::template MakeLayout( + L1TileShape::M, L1TileShape::K); + static constexpr auto L1B_LAYOUT = LayoutBInL1::template MakeLayout( + L1TileShape::K, L1TileShape::N); + + CATLASS_DEVICE + BlockMmad(Arch::Resource &resource, uint32_t l1BufAddrStart = 0) + { + InitL1(resource, l1BufAddrStart); + InitL0A(resource); + InitL0B(resource); + InitL0C(resource); + } + + CATLASS_DEVICE + ~BlockMmad() + { + SynchronizeBlock(); + for (uint32_t i = 0; i < L1_STAGES; ++i) { + AscendC::WaitFlag(l1AEventList[i]); + AscendC::WaitFlag(l1BEventList[i]); + } + for (uint32_t i = 0; i < L0A_STAGES; ++i) { + AscendC::WaitFlag(l0AEventList[i]); + } + for (uint32_t i = 0; i < L0B_STAGES; ++i) { + AscendC::WaitFlag(l0BEventList[i]); + } + for (uint32_t i = 0; i < L0C_STAGES; ++i) { + AscendC::WaitFlag(l0CEventList[i]); + } + } + + CATLASS_DEVICE + void operator()( + AscendC::GlobalTensor const &gmBlockA, LayoutA const &layoutA, + AscendC::GlobalTensor const &gmBlockB, LayoutB const &layoutB, + AscendC::GlobalTensor const &gmBlockC, LayoutC const &layoutC, + GemmCoord const &actualShape, Callback &&callback = Callback{}) + { + uint32_t kTileCount = CeilDiv(actualShape.k()); + + uint32_t mRound = RoundUp(actualShape.m()); + uint32_t nRound = RoundUp(actualShape.n()); + + uint32_t startTileIdx = 0; + if constexpr (ENABLE_SHUFFLE_K) { + startTileIdx = AscendC::GetBlockIdx() % kTileCount; + } + + for (uint32_t kLoopIdx = 0; kLoopIdx < kTileCount; ++kLoopIdx) { + uint32_t kTileIdx = (startTileIdx + kLoopIdx < kTileCount) ? + (startTileIdx + kLoopIdx) : + (startTileIdx + kLoopIdx - kTileCount); + + uint32_t kActual = (kTileIdx < kTileCount - 1) ? + L1TileShape::K : + (actualShape.k() - kTileIdx * L1TileShape::K); + + // Emission load instruction from GM to L1 + MatrixCoord gmTileAOffset{0, kTileIdx * L1TileShape::K}; + MatrixCoord gmTileBOffset{kTileIdx * L1TileShape::K, 0}; + auto gmTileA = gmBlockA[layoutA.GetOffset(gmTileAOffset)]; + auto gmTileB = gmBlockB[layoutB.GetOffset(gmTileBOffset)]; + // Load first matrix A tile from GM to L1 + AscendC::WaitFlag(l1AEventList[l1ListId]); + auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShape.m(), kActual)); + copyGmToL1A(l1ATensorList[l1ListId], gmTileA, L1A_LAYOUT, layoutTileA); + AscendC::SetFlag(l1AEventList[l1ListId]); + // Load first matrix B tile from GM to L1 + AscendC::WaitFlag(l1BEventList[l1ListId]); + auto layoutTileB = layoutB.GetTileLayout(MakeCoord(kActual, actualShape.n())); + copyGmToL1B(l1BTensorList[l1ListId], gmTileB, L1B_LAYOUT, layoutTileB); + AscendC::SetFlag(l1BEventList[l1ListId]); + + // If the number of preload instructions reaches the upper limit, perform an mmad calculation on L1 tile + if (preloadCount == PRELOAD_STAGES) { + L1TileMmad(l1TileMmadParamsList[l1TileMmadParamsId]); + } + + // Store the current load status + uint32_t preloadL1TileMmadParamsId = (l1TileMmadParamsId + preloadCount < PRELOAD_STAGES) ? + (l1TileMmadParamsId + preloadCount) : + (l1TileMmadParamsId + preloadCount - PRELOAD_STAGES); + auto &l1TileMmadParams = l1TileMmadParamsList[preloadL1TileMmadParamsId]; + l1TileMmadParams.l1ListId = l1ListId; + l1TileMmadParams.mRound = mRound; + l1TileMmadParams.nRound = nRound; + l1TileMmadParams.kActual = kActual; + l1TileMmadParams.isKLoopFirst = (kLoopIdx == 0); + l1TileMmadParams.isKLoopLast = (kLoopIdx == kTileCount - 1); + if (kLoopIdx == kTileCount - 1) { + l1TileMmadParams.gmBlockC = gmBlockC; + l1TileMmadParams.layoutCInGm = layoutC.GetTileLayout(actualShape.GetCoordMN()); + l1TileMmadParams.callback = callback; + } + + if (preloadCount < PRELOAD_STAGES) { + ++preloadCount; + } else { + l1TileMmadParamsId = (l1TileMmadParamsId + 1 < PRELOAD_STAGES) ? (l1TileMmadParamsId + 1) : 0; + } + l1ListId = (l1ListId + 1 < L1_STAGES) ? (l1ListId + 1) : 0; + } + } + + CATLASS_DEVICE + void SynchronizeBlock() + { + while (preloadCount > 0) { + L1TileMmad(l1TileMmadParamsList[l1TileMmadParamsId]); + l1TileMmadParamsId = (l1TileMmadParamsId + 1 < PRELOAD_STAGES) ? (l1TileMmadParamsId + 1) : 0; + --preloadCount; + } + } + +private: + struct L1TileMmadParams { + uint32_t l1ListId; + uint32_t mRound; + uint32_t nRound; + uint32_t kActual; + bool isKLoopFirst; + bool isKLoopLast; + AscendC::GlobalTensor gmBlockC; + LayoutC layoutCInGm; + Callback callback; + + CATLASS_DEVICE + L1TileMmadParams() = default; + }; + + CATLASS_DEVICE + void InitL1(Arch::Resource &resource, uint32_t l1BufAddrStart) + { + uint32_t l1AOffset = l1BufAddrStart; + uint32_t l1BOffset = l1BufAddrStart + L1A_TILE_SIZE * L1_STAGES; + for (uint32_t i = 0; i < L1_STAGES; ++i) { + l1ATensorList[i] = resource.l1Buf.template GetBufferByByte(l1AOffset + L1A_TILE_SIZE * i); + l1BTensorList[i] = resource.l1Buf.template GetBufferByByte(l1BOffset + L1B_TILE_SIZE * i); + l1AEventList[i] = i; + l1BEventList[i] = i + L1_STAGES; + AscendC::SetFlag(l1AEventList[i]); + AscendC::SetFlag(l1BEventList[i]); + } + } + + CATLASS_DEVICE + void InitL0A(Arch::Resource &resource) + { + for (uint32_t i = 0; i < L0A_STAGES; ++i) { + l0ATensorList[i] = resource.l0ABuf.template GetBufferByByte(L0A_TILE_SIZE * i); + l0AEventList[i] = i; + AscendC::SetFlag(l0AEventList[i]); + } + } + + CATLASS_DEVICE + void InitL0B(Arch::Resource &resource) + { + for (uint32_t i = 0; i < L0B_STAGES; ++i) { + l0BTensorList[i] = resource.l0BBuf.template GetBufferByByte(L0B_TILE_SIZE * i); + l0BEventList[i] = i + L0A_STAGES; + AscendC::SetFlag(l0BEventList[i]); + } + } + + CATLASS_DEVICE + void InitL0C(Arch::Resource &resource) + { + for (uint32_t i = 0; i < L0C_STAGES; ++i) { + l0CTensorList[i] = resource.l0CBuf.template GetBufferByByte(L0C_TILE_SIZE * i); + l0CEventList[i] = i; + AscendC::SetFlag(l0CEventList[i]); + } + } + + CATLASS_DEVICE + void L1TileMmad(L1TileMmadParams const ¶ms) + { + uint32_t mPartLoop = CeilDiv(params.mRound); + uint32_t nPartLoop = CeilDiv(params.nRound); + uint32_t kPartLoop = CeilDiv(params.kActual); + auto &l1ATensor = l1ATensorList[params.l1ListId]; + auto &l1BTensor = l1BTensorList[params.l1ListId]; + + auto &l0CTensor = l0CTensorList[l0CListId]; + LayoutCInL0 layoutCInL0 = LayoutCInL0::MakeLayoutInL0C(MakeCoord(params.mRound, params.nRound)); + + if constexpr (!ENABLE_UNIT_FLAG) { + if (params.isKLoopFirst) { + AscendC::WaitFlag(l0CEventList[l0CListId]); + } + } + + for (uint32_t mPartIdx = 0; mPartIdx < mPartLoop; ++mPartIdx) { + uint32_t mPartActual = (mPartIdx < mPartLoop - 1) ? + L0TileShape::M : + (params.mRound - mPartIdx * L0TileShape::M); + + for (uint32_t kPartIdx = 0; kPartIdx < kPartLoop; ++kPartIdx) { + uint32_t kPartActual = (kPartIdx < kPartLoop - 1) ? + L0TileShape::K : + (params.kActual - kPartIdx * L0TileShape::K); + + auto &l0ATile = l0ATensorList[l0AListId]; + auto layoutAInL0 = LayoutAInL0::template MakeLayout(mPartActual, kPartActual); + auto l1AOffset = MakeCoord(mPartIdx, kPartIdx) * L0TileShape::ToCoordMK(); + auto l1ATile = l1ATensor[L1A_LAYOUT.GetOffset(l1AOffset)]; + + AscendC::WaitFlag(l0AEventList[l0AListId]); + if ((mPartIdx == 0) && (kPartIdx == 0)) { + AscendC::WaitFlag(l1AEventList[params.l1ListId]); + } + copyL1ToL0A(l0ATile, l1ATile, layoutAInL0, L1A_LAYOUT); + if ((mPartIdx == mPartLoop - 1) && (kPartIdx == kPartLoop - 1)) { + AscendC::SetFlag(l1AEventList[params.l1ListId]); + } + + for (uint32_t nPartIdx = 0; nPartIdx < nPartLoop; ++nPartIdx) { + uint32_t nPartActual = (nPartIdx < nPartLoop - 1) ? + L0TileShape::N : + (params.nRound - nPartIdx * L0TileShape::N); + + auto &l0BTile = l0BTensorList[l0BListId]; + auto layoutBInL0 = LayoutBInL0::template MakeLayout(kPartActual, nPartActual); + auto l1BOffset = MakeCoord(kPartIdx, nPartIdx) * L0TileShape::ToCoordKN(); + auto l1BTile = l1BTensor[L1B_LAYOUT.GetOffset(l1BOffset)]; + + AscendC::WaitFlag(l0BEventList[l0BListId]); + if ((kPartIdx == 0) && (nPartIdx == 0)) { + AscendC::WaitFlag(l1BEventList[params.l1ListId]); + } + copyL1ToL0B(l0BTile, l1BTile, layoutBInL0, L1B_LAYOUT); + if ((kPartIdx == kPartLoop - 1) && (nPartIdx == nPartLoop - 1)) { + AscendC::SetFlag(l1BEventList[params.l1ListId]); + } + + AscendC::SetFlag(EVENT_ID0); + + auto l0COffset = MakeCoord(mPartIdx, nPartIdx) * L0TileShape::ToCoordMN(); + auto l0CTile = l0CTensor[layoutCInL0.GetOffset(l0COffset)]; + + AscendC::WaitFlag(EVENT_ID0); + // If the current tile is the first tile on the k axis, the accumulator needs to be reset to 0 + bool initC = (params.isKLoopFirst && (kPartIdx == 0)); + // If the unit flag is enabled, the unit flag is set according to the calculation progress + uint8_t unitFlag = 0b00; + if constexpr (ENABLE_UNIT_FLAG) { + if (params.isKLoopLast && + (mPartIdx == mPartLoop - 1) && (kPartIdx == kPartLoop - 1) && (nPartIdx == nPartLoop - 1)) { + unitFlag = 0b11; + } else { + unitFlag = 0b10; + } + } + tileMmad(l0CTile, l0ATile, l0BTile, mPartActual, nPartActual, kPartActual, initC, unitFlag); + + AscendC::SetFlag(l0BEventList[l0BListId]); + l0BListId = (l0BListId + 1 < L0B_STAGES) ? (l0BListId + 1) : 0; + } + AscendC::SetFlag(l0AEventList[l0AListId]); + l0AListId = (l0AListId + 1 < L0A_STAGES) ? (l0AListId + 1) : 0; + } + } + + if (params.isKLoopLast) { + auto layoutCInGm = params.layoutCInGm; + + if constexpr (!ENABLE_UNIT_FLAG) { + AscendC::SetFlag(l0CEventList[l0CListId]); + AscendC::WaitFlag(l0CEventList[l0CListId]); + copyL0CToGm(params.gmBlockC, l0CTensor, layoutCInGm, layoutCInL0); + AscendC::SetFlag(l0CEventList[l0CListId]); + } else { + copyL0CToGm(params.gmBlockC, l0CTensor, layoutCInGm, layoutCInL0, 0b11); + } + l0CListId = (l0CListId + 1 < L0C_STAGES) ? (l0CListId + 1) : 0; + + if (params.callback) { + params.callback(); + } + } + } + + AscendC::LocalTensor l1ATensorList[L1_STAGES]; + AscendC::LocalTensor l1BTensorList[L1_STAGES]; + int32_t l1AEventList[L1_STAGES]; + int32_t l1BEventList[L1_STAGES]; + uint32_t l1ListId{0}; + + AscendC::LocalTensor l0ATensorList[L0A_STAGES]; + int32_t l0AEventList[L0A_STAGES]; + uint32_t l0AListId{0}; + + AscendC::LocalTensor l0BTensorList[L0B_STAGES]; + int32_t l0BEventList[L0B_STAGES]; + uint32_t l0BListId{0}; + + AscendC::LocalTensor l0CTensorList[L0C_STAGES_]; + int32_t l0CEventList[L0C_STAGES_]; + uint32_t l0CListId{0}; + + L1TileMmadParams l1TileMmadParamsList[PRELOAD_STAGES]; + uint32_t l1TileMmadParamsId{0}; + uint32_t preloadCount{0}; + + TileMmad tileMmad; + CopyGmToL1A copyGmToL1A; + CopyGmToL1B copyGmToL1B; + CopyL1ToL0A copyL1ToL0A; + CopyL1ToL0B copyL1ToL0B; + CopyL0CToGm copyL0CToGm; +}; + +} // namespace Catlass::Gemm::Block + +#endif // CATLASS_KDA_BLOCK_MMAD_PRELOAD_ASYNC_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_preload_async_with_callback.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_preload_async_with_callback.hpp new file mode 100644 index 0000000000..df26fa7548 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_preload_async_with_callback.hpp @@ -0,0 +1,449 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_BLOCK_MMAD_PRELOAD_ASYNC_WITH_CALLBACK_HPP +#define CATLASS_KDA_BLOCK_MMAD_PRELOAD_ASYNC_WITH_CALLBACK_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/kda_coord.hpp" +#include "catlass/detail/kda_callback.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" +#include "catlass/gemm/kda_gemm_helper.hpp" +#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" +#include "catlass/gemm/tile/kda_tile_mmad.hpp" + +namespace Catlass::Gemm::Block { + +template < + uint32_t PRELOAD_STAGES_, + uint32_t L1_STAGES_, + uint32_t L0A_STAGES_, + uint32_t L0B_STAGES_, + uint32_t L0C_STAGES_, + bool ENABLE_UNIT_FLAG_, + bool ENABLE_SHUFFLE_K_, + class L1TileShape_, + class L0TileShape_, + class AType_, + class BType_, + class CType_, + class BiasType_, + class TileCopy_, + class TileMmad_> +struct BlockMmad< + MmadAtlasA2PreloadAsyncWithCallback< + PRELOAD_STAGES_, + L1_STAGES_, + L0A_STAGES_, + L0B_STAGES_, + L0C_STAGES_, + ENABLE_UNIT_FLAG_, + ENABLE_SHUFFLE_K_>, + L1TileShape_, + L0TileShape_, + AType_, + BType_, + CType_, + BiasType_, + TileCopy_, + TileMmad_> { +public: + // Type Aliases + using DispatchPolicy = MmadAtlasA2PreloadAsyncWithCallback< + PRELOAD_STAGES_, + L1_STAGES_, + L0A_STAGES_, + L0B_STAGES_, + L0C_STAGES_, + ENABLE_UNIT_FLAG_, + ENABLE_SHUFFLE_K_>; + using ArchTag = typename DispatchPolicy::ArchTag; + using L1TileShape = L1TileShape_; + using L0TileShape = L0TileShape_; + using ElementA = typename AType_::Element; + using LayoutA = typename AType_::Layout; + using ElementB = typename BType_::Element; + using LayoutB = typename BType_::Layout; + using ElementC = typename CType_::Element; + using LayoutC = typename CType_::Layout; + using TileMmad = TileMmad_; + using CopyGmToL1A = typename TileCopy_::CopyGmToL1A; + using CopyGmToL1B = typename TileCopy_::CopyGmToL1B; + using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; + using CopyL0CToGm = typename TileCopy_::CopyL0CToGm; + using ElementAccumulator = + typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; + using LayoutAInL1 = typename CopyL1ToL0A::LayoutSrc; + using LayoutBInL1 = typename CopyL1ToL0B::LayoutSrc; + using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst; + using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst; + using LayoutCInL0 = layout::zN; + + using L1AAlignHelper = Gemm::helper::L1AlignHelper; + using L1BAlignHelper = Gemm::helper::L1AlignHelper; + + static constexpr uint32_t PRELOAD_STAGES = DispatchPolicy::PRELOAD_STAGES; + static constexpr uint32_t L1_STAGES = DispatchPolicy::L1_STAGES; + static constexpr uint32_t L0A_STAGES = DispatchPolicy::L0A_STAGES; + static constexpr uint32_t L0B_STAGES = DispatchPolicy::L0B_STAGES; + static constexpr uint32_t L0C_STAGES = DispatchPolicy::L0C_STAGES; + + static constexpr bool ENABLE_UNIT_FLAG = DispatchPolicy::ENABLE_UNIT_FLAG; + static constexpr bool ENABLE_SHUFFLE_K = DispatchPolicy::ENABLE_SHUFFLE_K; + + // L1 tile size + static constexpr uint32_t L1A_TILE_SIZE = L1TileShape::M * L1TileShape::K * sizeof(ElementA); + static constexpr uint32_t L1B_TILE_SIZE = L1TileShape::N * L1TileShape::K * sizeof(ElementB); + // L0 tile size + static constexpr uint32_t L0A_TILE_SIZE = L0TileShape::M * L0TileShape::K * sizeof(ElementA); + static constexpr uint32_t L0B_TILE_SIZE = L0TileShape::K * L0TileShape::N * sizeof(ElementB); + static constexpr uint32_t L0C_TILE_SIZE = L1TileShape::M * L1TileShape::N * sizeof(ElementAccumulator); + + // Check LayoutC + static_assert(std::is_same_v, "LayoutC only support RowMajor yet!"); + + // Check L1TileShape + static_assert((L1A_TILE_SIZE + L1B_TILE_SIZE) * L1_STAGES <= ArchTag::L1_SIZE, + "L1TileShape exceeding the L1 space!"); + + // Check L0TileShape + static_assert(L0A_TILE_SIZE * L0A_STAGES <= ArchTag::L0A_SIZE, "L0TileShape exceeding the L0A space!"); + static_assert(L0B_TILE_SIZE * L0B_STAGES <= ArchTag::L0B_SIZE, "L0TileShape exceeding the L0B space!"); + static_assert(L0C_TILE_SIZE * L0C_STAGES <= ArchTag::L0C_SIZE, "L0TileShape exceeding the L0C space!"); + + static_assert(L1TileShape::M == L0TileShape::M && L1TileShape::N == L0TileShape::N, + "The situation where the basic blocks of L1 and L0 differ on the m and n axes is not supported yet"); + static_assert(L0TileShape::K <= L1TileShape::K, "L0TileShape::K cannot exceed L1TileShape::K"); + // 32B (256b) aligned + static_assert(Gemm::helper::TileShapeAlignChecker::_ALIGN == 256, + "Tile shape must be 32B aligned."); + + static constexpr auto L1A_LAYOUT = LayoutAInL1::template MakeLayout(L1TileShape::M, L1TileShape::K); + static constexpr auto L1B_LAYOUT = LayoutBInL1::template MakeLayout(L1TileShape::K, L1TileShape::N); + + CATLASS_DEVICE + BlockMmad(Arch::Resource &resource, uint32_t l1BufAddrStart = 0) + { + InitL1(resource, l1BufAddrStart); + InitL0A(resource); + InitL0B(resource); + InitL0C(resource); + } + + CATLASS_DEVICE + ~BlockMmad() + { + SynchronizeBlock(); + for (uint32_t i = 0; i < L1_STAGES; ++i) { + AscendC::WaitFlag(l1AEventList[i]); + AscendC::WaitFlag(l1BEventList[i]); + } + for (uint32_t i = 0; i < L0A_STAGES; ++i) { + AscendC::WaitFlag(l0AEventList[i]); + } + for (uint32_t i = 0; i < L0B_STAGES; ++i) { + AscendC::WaitFlag(l0BEventList[i]); + } + for (uint32_t i = 0; i < L0C_STAGES; ++i) { + AscendC::WaitFlag(l0CEventList[i]); + } + } + + CATLASS_DEVICE + void operator()( + AscendC::GlobalTensor const &gmBlockA, LayoutA const &layoutA, + AscendC::GlobalTensor const &gmBlockB, LayoutB const &layoutB, + AscendC::GlobalTensor const &gmBlockC, LayoutC const &layoutC, + GemmCoord const &actualShape, + Callback const &callbackBeforeFixpipe, Callback const &callbackAfterFixpipe) + { + uint32_t kTileCount = CeilDiv(actualShape.k()); + + uint32_t mRound = RoundUp(actualShape.m()); + uint32_t nRound = RoundUp(actualShape.n()); + + uint32_t startTileIdx = 0; + if constexpr (ENABLE_SHUFFLE_K) { + startTileIdx = AscendC::GetBlockIdx() % kTileCount; + } + + for (uint32_t kLoopIdx = 0; kLoopIdx < kTileCount; ++kLoopIdx) { + uint32_t kTileIdx = (startTileIdx + kLoopIdx < kTileCount) ? + (startTileIdx + kLoopIdx) : + (startTileIdx + kLoopIdx - kTileCount); + + uint32_t kActual = (kTileIdx < kTileCount - 1) ? + L1TileShape::K : + (actualShape.k() - kTileIdx * L1TileShape::K); + + // Emission load instruction from GM to L1 + MatrixCoord gmTileAOffset{0, kTileIdx * L1TileShape::K}; + MatrixCoord gmTileBOffset{kTileIdx * L1TileShape::K, 0}; + auto gmTileA = gmBlockA[layoutA.GetOffset(gmTileAOffset)]; + auto gmTileB = gmBlockB[layoutB.GetOffset(gmTileBOffset)]; + // Load first matrix A tile from GM to L1 + AscendC::WaitFlag(l1AEventList[l1ListId]); + auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShape.m(), kActual)); + copyGmToL1A(l1ATensorList[l1ListId], gmTileA, L1A_LAYOUT, layoutTileA); + AscendC::SetFlag(l1AEventList[l1ListId]); + // Load first matrix B tile from GM to L1 + AscendC::WaitFlag(l1BEventList[l1ListId]); + auto layoutTileB = layoutB.GetTileLayout(MakeCoord(kActual, actualShape.n())); + copyGmToL1B(l1BTensorList[l1ListId], gmTileB, L1B_LAYOUT, layoutTileB); + AscendC::SetFlag(l1BEventList[l1ListId]); + + // If the number of preload instructions reaches the upper limit, perform an mmad calculation on L1 tile + if (preloadCount == PRELOAD_STAGES) { + L1TileMmad(l1TileMmadParamsList[l1TileMmadParamsId]); + } + + // Store the current load status + uint32_t preloadL1TileMmadParamsId = (l1TileMmadParamsId + preloadCount < PRELOAD_STAGES) ? + (l1TileMmadParamsId + preloadCount) : + (l1TileMmadParamsId + preloadCount - PRELOAD_STAGES); + auto &l1TileMmadParams = l1TileMmadParamsList[preloadL1TileMmadParamsId]; + l1TileMmadParams.l1ListId = l1ListId; + l1TileMmadParams.mRound = mRound; + l1TileMmadParams.nRound = nRound; + l1TileMmadParams.kActual = kActual; + l1TileMmadParams.isKLoopFirst = (kLoopIdx == 0); + l1TileMmadParams.isKLoopLast = (kLoopIdx == kTileCount - 1); + if (kLoopIdx == kTileCount - 1) { + l1TileMmadParams.gmBlockC = gmBlockC; + l1TileMmadParams.layoutCInGm = layoutC.GetTileLayout(actualShape.GetCoordMN()); + l1TileMmadParams.callbackBeforeFixpipe = callbackBeforeFixpipe; + l1TileMmadParams.callbackAfterFixpipe = callbackAfterFixpipe; + } + + if (preloadCount < PRELOAD_STAGES) { + ++preloadCount; + } else { + l1TileMmadParamsId = (l1TileMmadParamsId + 1 < PRELOAD_STAGES) ? (l1TileMmadParamsId + 1) : 0; + } + l1ListId = (l1ListId + 1 < L1_STAGES) ? (l1ListId + 1) : 0; + } + } + + CATLASS_DEVICE + void SynchronizeBlock() + { + while (preloadCount > 0) { + L1TileMmad(l1TileMmadParamsList[l1TileMmadParamsId]); + l1TileMmadParamsId = (l1TileMmadParamsId + 1 < PRELOAD_STAGES) ? (l1TileMmadParamsId + 1) : 0; + --preloadCount; + } + } + +private: + struct L1TileMmadParams { + uint32_t l1ListId; + uint32_t mRound; + uint32_t nRound; + uint32_t kActual; + bool isKLoopFirst; + bool isKLoopLast; + AscendC::GlobalTensor gmBlockC; + LayoutC layoutCInGm; + Callback callbackBeforeFixpipe; + Callback callbackAfterFixpipe; + + CATLASS_DEVICE + L1TileMmadParams() = default; + }; + + CATLASS_DEVICE + void InitL1(Arch::Resource &resource, uint32_t l1BufAddrStart) + { + uint32_t l1AOffset = l1BufAddrStart; + uint32_t l1BOffset = l1BufAddrStart + L1A_TILE_SIZE * L1_STAGES; + for (uint32_t i = 0; i < L1_STAGES; ++i) { + l1ATensorList[i] = resource.l1Buf.template GetBufferByByte(l1AOffset + L1A_TILE_SIZE * i); + l1BTensorList[i] = resource.l1Buf.template GetBufferByByte(l1BOffset + L1B_TILE_SIZE * i); + l1AEventList[i] = i; + l1BEventList[i] = i + L1_STAGES; + AscendC::SetFlag(l1AEventList[i]); + AscendC::SetFlag(l1BEventList[i]); + } + } + + CATLASS_DEVICE + void InitL0A(Arch::Resource &resource) + { + for (uint32_t i = 0; i < L0A_STAGES; ++i) { + l0ATensorList[i] = resource.l0ABuf.template GetBufferByByte(L0A_TILE_SIZE * i); + l0AEventList[i] = i; + AscendC::SetFlag(l0AEventList[i]); + } + } + + CATLASS_DEVICE + void InitL0B(Arch::Resource &resource) + { + for (uint32_t i = 0; i < L0B_STAGES; ++i) { + l0BTensorList[i] = resource.l0BBuf.template GetBufferByByte(L0B_TILE_SIZE * i); + l0BEventList[i] = i + L0A_STAGES; + AscendC::SetFlag(l0BEventList[i]); + } + } + + CATLASS_DEVICE + void InitL0C(Arch::Resource &resource) + { + for (uint32_t i = 0; i < L0C_STAGES; ++i) { + l0CTensorList[i] = resource.l0CBuf.template GetBufferByByte(L0C_TILE_SIZE * i); + l0CEventList[i] = i; + AscendC::SetFlag(l0CEventList[i]); + } + } + + CATLASS_DEVICE + void L1TileMmad(L1TileMmadParams const ¶ms) + { + uint32_t mPartLoop = CeilDiv(params.mRound); + uint32_t nPartLoop = CeilDiv(params.nRound); + uint32_t kPartLoop = CeilDiv(params.kActual); + auto &l1ATensor = l1ATensorList[params.l1ListId]; + auto &l1BTensor = l1BTensorList[params.l1ListId]; + + auto &l0CTensor = l0CTensorList[l0CListId]; + LayoutCInL0 layoutCInL0 = LayoutCInL0::MakeLayoutInL0C(MakeCoord(params.mRound, params.nRound)); + + if constexpr (!ENABLE_UNIT_FLAG) { + if (params.isKLoopFirst) { + AscendC::WaitFlag(l0CEventList[l0CListId]); + } + } + + for (uint32_t mPartIdx = 0; mPartIdx < mPartLoop; ++mPartIdx) { + uint32_t mPartActual = (mPartIdx < mPartLoop - 1) ? + L0TileShape::M : + (params.mRound - mPartIdx * L0TileShape::M); + + for (uint32_t kPartIdx = 0; kPartIdx < kPartLoop; ++kPartIdx) { + uint32_t kPartActual = (kPartIdx < kPartLoop - 1) ? + L0TileShape::K : + (params.kActual - kPartIdx * L0TileShape::K); + + auto &l0ATile = l0ATensorList[l0AListId]; + auto layoutAInL0 = LayoutAInL0::template MakeLayout(mPartActual, kPartActual); + auto l1AOffset = MakeCoord(mPartIdx, kPartIdx) * L0TileShape::ToCoordMK(); + auto l1ATile = l1ATensor[L1A_LAYOUT.GetOffset(l1AOffset)]; + + AscendC::WaitFlag(l0AEventList[l0AListId]); + if ((mPartIdx == 0) && (kPartIdx == 0)) { + AscendC::WaitFlag(l1AEventList[params.l1ListId]); + } + copyL1ToL0A(l0ATile, l1ATile, layoutAInL0, L1A_LAYOUT); + if ((mPartIdx == mPartLoop - 1) && (kPartIdx == kPartLoop - 1)) { + AscendC::SetFlag(l1AEventList[params.l1ListId]); + } + + for (uint32_t nPartIdx = 0; nPartIdx < nPartLoop; ++nPartIdx) { + uint32_t nPartActual = (nPartIdx < nPartLoop - 1) ? + L0TileShape::N : + (params.nRound - nPartIdx * L0TileShape::N); + + auto &l0BTile = l0BTensorList[l0BListId]; + auto layoutBInL0 = LayoutBInL0::template MakeLayout(kPartActual, nPartActual); + auto l1BOffset = MakeCoord(kPartIdx, nPartIdx) * L0TileShape::ToCoordKN(); + auto l1BTile = l1BTensor[L1B_LAYOUT.GetOffset(l1BOffset)]; + + AscendC::WaitFlag(l0BEventList[l0BListId]); + if ((kPartIdx == 0) && (nPartIdx == 0)) { + AscendC::WaitFlag(l1BEventList[params.l1ListId]); + } + copyL1ToL0B(l0BTile, l1BTile, layoutBInL0, L1B_LAYOUT); + if ((kPartIdx == kPartLoop - 1) && (nPartIdx == nPartLoop - 1)) { + AscendC::SetFlag(l1BEventList[params.l1ListId]); + } + + AscendC::SetFlag(EVENT_ID0); + + auto l0COffset = MakeCoord(mPartIdx, nPartIdx) * L0TileShape::ToCoordMN(); + auto l0CTile = l0CTensor[layoutCInL0.GetOffset(l0COffset)]; + + AscendC::WaitFlag(EVENT_ID0); + // If the current tile is the first tile on the k axis, the accumulator needs to be reset to 0 + bool initC = (params.isKLoopFirst && (kPartIdx == 0)); + // If the unit flag is enabled, the unit flag is set according to the calculation progress + uint8_t unitFlag = 0b00; + if constexpr (ENABLE_UNIT_FLAG) { + if (params.isKLoopLast && + (mPartIdx == mPartLoop - 1) && (kPartIdx == kPartLoop - 1) && (nPartIdx == nPartLoop - 1)) { + unitFlag = 0b11; + } else { + unitFlag = 0b10; + } + } + tileMmad(l0CTile, l0ATile, l0BTile, mPartActual, nPartActual, kPartActual, initC, unitFlag); + + AscendC::SetFlag(l0BEventList[l0BListId]); + l0BListId = (l0BListId + 1 < L0B_STAGES) ? (l0BListId + 1) : 0; + } + AscendC::SetFlag(l0AEventList[l0AListId]); + l0AListId = (l0AListId + 1 < L0A_STAGES) ? (l0AListId + 1) : 0; + } + } + + if (params.isKLoopLast) { + auto layoutCInGm = params.layoutCInGm; + + params.callbackBeforeFixpipe(); + + if constexpr (!ENABLE_UNIT_FLAG) { + AscendC::SetFlag(l0CEventList[l0CListId]); + AscendC::WaitFlag(l0CEventList[l0CListId]); + copyL0CToGm(params.gmBlockC, l0CTensor, layoutCInGm, layoutCInL0); + AscendC::SetFlag(l0CEventList[l0CListId]); + } else { + copyL0CToGm(params.gmBlockC, l0CTensor, layoutCInGm, layoutCInL0, 0b11); + } + l0CListId = (l0CListId + 1 < L0C_STAGES) ? (l0CListId + 1) : 0; + + params.callbackAfterFixpipe(); + } + } + + AscendC::LocalTensor l1ATensorList[L1_STAGES]; + AscendC::LocalTensor l1BTensorList[L1_STAGES]; + int32_t l1AEventList[L1_STAGES]; + int32_t l1BEventList[L1_STAGES]; + uint32_t l1ListId{0}; + + AscendC::LocalTensor l0ATensorList[L0A_STAGES]; + int32_t l0AEventList[L0A_STAGES]; + uint32_t l0AListId{0}; + + AscendC::LocalTensor l0BTensorList[L0B_STAGES]; + int32_t l0BEventList[L0B_STAGES]; + uint32_t l0BListId{0}; + + AscendC::LocalTensor l0CTensorList[L0C_STAGES_]; + int32_t l0CEventList[L0C_STAGES_]; + uint32_t l0CListId{0}; + + L1TileMmadParams l1TileMmadParamsList[PRELOAD_STAGES]; + uint32_t l1TileMmadParamsId{0}; + uint32_t preloadCount{0}; + + TileMmad tileMmad; + CopyGmToL1A copyGmToL1A; + CopyGmToL1B copyGmToL1B; + CopyL1ToL0A copyL1ToL0A; + CopyL1ToL0B copyL1ToL0B; + CopyL0CToGm copyL0CToGm; +}; + +} // namespace Catlass::Gemm::Block + +#endif // CATLASS_KDA_BLOCK_MMAD_PRELOAD_ASYNC_WITH_CALLBACK_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_swizzle.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_swizzle.hpp new file mode 100644 index 0000000000..db6bb00841 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_swizzle.hpp @@ -0,0 +1,1018 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_BLOCK_SWIZZLE_HPP +#define CATLASS_KDA_BLOCK_SWIZZLE_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/detail/kda_alignment.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/kda_matrix_coord.hpp" + +namespace Catlass::Gemm::Block { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Block swizzling function for Gemms +template +struct GemmIdentityBlockSwizzle { + /// Data members + + GemmCoord problemShape; + MatrixCoord tileMN; + MatrixCoord loopsMN; + + /// Methods + + CATLASS_DEVICE + GemmIdentityBlockSwizzle() {} + + CATLASS_DEVICE + GemmIdentityBlockSwizzle(GemmCoord const &problemShape_, MatrixCoord const &tileMN_) + : problemShape(problemShape_), + tileMN(tileMN_) + { + loopsMN = CeilDiv(MatrixCoord(problemShape.GetCoordMN()), tileMN); + } + + CATLASS_DEVICE + GemmIdentityBlockSwizzle(GemmCoord const &problemShape_, MatrixCoord const &tileMN_, + MatrixCoord const &loopsMN_) + : problemShape(problemShape_), + tileMN(tileMN_), + loopsMN(loopsMN_) + {} + + CATLASS_DEVICE + void Update(GemmCoord const &problemShape_, MatrixCoord const &tileMN_) + { + problemShape = problemShape_; + tileMN = tileMN_; + + loopsMN = CeilDiv(MatrixCoord(problemShape.GetCoordMN()), tileMN); + } + + CATLASS_DEVICE + void Update(GemmCoord const &problemShape_, MatrixCoord const &tileMN_, MatrixCoord const &loopsMN_) + { + problemShape = problemShape_; + tileMN = tileMN_; + loopsMN = loopsMN_; + } + + CATLASS_DEVICE + uint32_t GetCoreLoops() const + { + return loopsMN.row() * loopsMN.column(); + } + + CATLASS_DEVICE + uint32_t GetBatchIdx(uint32_t taskIdx) + { + return taskIdx / (GetCoreLoops()); + } + + CATLASS_DEVICE + GemmCoord GetBlockCoord(uint32_t taskIdx) + { + uint32_t innerIdx = taskIdx % GetCoreLoops(); + if constexpr (SwizzleDirection == 0) { // Zn + uint32_t tileBlockLoop = CeilDiv(loopsMN.row(), SwizzleOffset); + uint32_t tileBlockIdx = innerIdx / (SwizzleOffset * loopsMN.column()); + uint32_t inTileBlockIdx = innerIdx % (SwizzleOffset * loopsMN.column()); + + uint32_t nRow = SwizzleOffset; + if (tileBlockIdx == tileBlockLoop - 1) { + nRow = loopsMN.row() - SwizzleOffset * tileBlockIdx; + } + uint32_t mIdx = tileBlockIdx * SwizzleOffset + inTileBlockIdx % nRow; + uint32_t nIdx = inTileBlockIdx / nRow; + if (tileBlockIdx % 2 == 1) { + nIdx = loopsMN.column() - nIdx - 1; + } + return GemmCoord{mIdx, nIdx, 0}; + } else if constexpr (SwizzleDirection == 1) { // Nz + uint32_t tileBlockLoop = CeilDiv(loopsMN.column(), SwizzleOffset); + uint32_t tileBlockIdx = innerIdx / (SwizzleOffset * loopsMN.row()); + uint32_t inTileBlockIdx = innerIdx % (SwizzleOffset * loopsMN.row()); + + uint32_t nCol = SwizzleOffset; + if (tileBlockIdx == tileBlockLoop - 1) { + nCol = loopsMN.column() - SwizzleOffset * tileBlockIdx; + } + uint32_t mIdx = inTileBlockIdx / nCol; + uint32_t nIdx = tileBlockIdx * SwizzleOffset + inTileBlockIdx % nCol; + if (tileBlockIdx % 2 == 1) { + mIdx = loopsMN.row() - mIdx - 1; + } + return GemmCoord{mIdx, nIdx, 0}; + } + } + + CATLASS_DEVICE + GemmCoord GetActualBlockShape(GemmCoord blockCoord) + { + uint32_t mActual = (blockCoord.m() == (loopsMN.row() - 1)) ? + (problemShape.m() - blockCoord.m() * tileMN.row()) : + tileMN.row(); + uint32_t nActual = (blockCoord.n() == (loopsMN.column() - 1)) ? + (problemShape.n() - blockCoord.n() * tileMN.column()) : + tileMN.column(); + uint32_t kActual = problemShape.k(); + return GemmCoord{mActual, nActual, kActual}; + } +}; + +struct DynamicGemmIdentityBlockSwizzle : public GemmIdentityBlockSwizzle<> { + uint32_t swizzleOffset{1}; + uint32_t swizzleDirection{0}; + + CATLASS_DEVICE + DynamicGemmIdentityBlockSwizzle(GemmCoord const &problemShape_, MatrixCoord const &tileMN_, + uint32_t swizzleOffset_, uint32_t swizzleDirection_) + : swizzleOffset(swizzleOffset_), + swizzleDirection(swizzleDirection_), + GemmIdentityBlockSwizzle<>(problemShape_, tileMN_) + {} + + CATLASS_DEVICE + DynamicGemmIdentityBlockSwizzle(GemmCoord const &problemShape_, MatrixCoord const &tileMN_) + : GemmIdentityBlockSwizzle<>(problemShape_, tileMN_) + {} + + CATLASS_DEVICE + DynamicGemmIdentityBlockSwizzle() {} + + CATLASS_DEVICE + void SetSwizzleParams(uint32_t swizzleOffset_, uint32_t swizzleDirection_) + { + swizzleOffset = swizzleOffset_; + swizzleDirection = swizzleDirection_; + } + + CATLASS_DEVICE + GemmCoord GetBlockCoord(uint32_t taskIdx) + { + uint32_t innerIdx = taskIdx % GetCoreLoops(); + if (swizzleDirection == 0) { // Zn + uint32_t tileBlockLoop = CeilDiv(loopsMN.row(), swizzleOffset); + uint32_t tileBlockIdx = innerIdx / (swizzleOffset * loopsMN.column()); + uint32_t inTileBlockIdx = innerIdx % (swizzleOffset * loopsMN.column()); + + uint32_t nRow = swizzleOffset; + if (tileBlockIdx == tileBlockLoop - 1) { + nRow = loopsMN.row() - swizzleOffset * tileBlockIdx; + } + uint32_t mIdx = tileBlockIdx * swizzleOffset + inTileBlockIdx % nRow; + uint32_t nIdx = inTileBlockIdx / nRow; + if (tileBlockIdx % 2 == 1) { + nIdx = loopsMN.column() - nIdx - 1; + } + return GemmCoord{mIdx, nIdx, 0}; + } else { // Nz + uint32_t tileBlockLoop = CeilDiv(loopsMN.column(), swizzleOffset); + uint32_t tileBlockIdx = innerIdx / (swizzleOffset * loopsMN.row()); + uint32_t inTileBlockIdx = innerIdx % (swizzleOffset * loopsMN.row()); + + uint32_t nCol = swizzleOffset; + if (tileBlockIdx == tileBlockLoop - 1) { + nCol = loopsMN.column() - swizzleOffset * tileBlockIdx; + } + uint32_t mIdx = inTileBlockIdx / nCol; + uint32_t nIdx = tileBlockIdx * swizzleOffset + inTileBlockIdx % nCol; + if (tileBlockIdx % 2 == 1) { + mIdx = loopsMN.row() - mIdx - 1; + } + return GemmCoord{mIdx, nIdx, 0}; + } + } +}; + +/// Block swizzling function for Splitk Gemms +template +struct SplitkGemmIdentityBlockSwizzle { + /// Data members + + GemmCoord problemShape; + GemmCoord tileShape; + GemmCoord loopsMNK; + uint32_t splitkFactor = 1; // splite k dim into virtual cores + + /// Methods + + CATLASS_DEVICE + SplitkGemmIdentityBlockSwizzle() {} + + CATLASS_DEVICE + SplitkGemmIdentityBlockSwizzle( + GemmCoord const &problemShape_, GemmCoord const &tileShape_, uint32_t splitkFactor_ = 1) + : problemShape(problemShape_), + tileShape(tileShape_), + splitkFactor(splitkFactor_) + { + loopsMNK = CeilDiv(problemShape, tileShape); + } + + CATLASS_DEVICE + uint32_t GetKIdxBySplitkSliceIdx(uint32_t splitkSliceIdx) const + { + if (splitkSliceIdx < loopsMNK.k() % splitkFactor) { + return (loopsMNK.k() / splitkFactor + 1) * splitkSliceIdx; + } else { + return splitkSliceIdx * (loopsMNK.k() / splitkFactor) + loopsMNK.k() % splitkFactor; + } + } + + CATLASS_DEVICE + uint32_t GetSplitkSliceIdx(uint32_t taskIdx) const + { + uint32_t mnLoops = loopsMNK.m() * loopsMNK.n(); + return taskIdx % GetCoreLoops() / mnLoops; + } + + CATLASS_DEVICE + uint32_t GetCoreLoops() const + { + return loopsMNK.m() * loopsMNK.n() * splitkFactor; + } + + CATLASS_DEVICE + uint32_t GetBatchIdx(uint32_t taskIdx) + { + return taskIdx / GetCoreLoops(); + } + + CATLASS_DEVICE + GemmCoord GetBlockCoord(uint32_t taskIdx) + { + uint32_t splitkSliceIdx = GetSplitkSliceIdx(taskIdx); + uint32_t kIdx = GetKIdxBySplitkSliceIdx(splitkSliceIdx); + + uint32_t innerIdx = taskIdx % (loopsMNK.m() * loopsMNK.n()); + if constexpr (SwizzleDirection == 0) { // Zn + uint32_t tileBlockLoop = CeilDiv(loopsMNK.m(), SwizzleOffset); + uint32_t tileBlockIdx = innerIdx / (SwizzleOffset * loopsMNK.n()); + uint32_t inTileBlockIdx = innerIdx % (SwizzleOffset * loopsMNK.n()); + + uint32_t nRow = SwizzleOffset; + if (tileBlockIdx == tileBlockLoop - 1) { + nRow = loopsMNK.m() - SwizzleOffset * tileBlockIdx; + } + uint32_t mIdx = tileBlockIdx * SwizzleOffset + inTileBlockIdx % nRow; + uint32_t nIdx = inTileBlockIdx / nRow; + if (tileBlockIdx % 2 == 1) { + nIdx = loopsMNK.n() - nIdx - 1; + } + return GemmCoord{mIdx, nIdx, kIdx}; + } else if constexpr (SwizzleDirection == 1) { // Nz + uint32_t tileBlockLoop = CeilDiv(loopsMNK.n(), SwizzleOffset); + uint32_t tileBlockIdx = innerIdx / (SwizzleOffset * loopsMNK.m()); + uint32_t inTileBlockIdx = innerIdx % (SwizzleOffset * loopsMNK.m()); + + uint32_t nCol = SwizzleOffset; + if (tileBlockIdx == tileBlockLoop - 1) { + nCol = loopsMNK.n() - SwizzleOffset * tileBlockIdx; + } + uint32_t mIdx = inTileBlockIdx / nCol; + uint32_t nIdx = tileBlockIdx * SwizzleOffset + inTileBlockIdx % nCol; + if (tileBlockIdx % 2 == 1) { + mIdx = loopsMNK.m() - mIdx - 1; + } + return GemmCoord{mIdx, nIdx, kIdx}; + } + } + + CATLASS_DEVICE + GemmCoord GetActualBlockShape(GemmCoord blockCoord, uint32_t splitkSliceIdx) + { + uint32_t splitkSliceLen; + if (splitkSliceIdx < loopsMNK.k() % splitkFactor) { + splitkSliceLen = (loopsMNK.k() / splitkFactor + 1) * tileShape.k(); + } else { + splitkSliceLen = (loopsMNK.k() / splitkFactor) * tileShape.k(); + } + uint32_t mActual = (blockCoord.m() == (loopsMNK.m() - 1)) ? + (problemShape.m() - blockCoord.m() * tileShape.m()) : + tileShape.m(); + uint32_t nActual = (blockCoord.n() == (loopsMNK.n() - 1)) ? + (problemShape.n() - blockCoord.n() * tileShape.n()) : + tileShape.n(); + uint32_t kActual = (splitkSliceIdx == (splitkFactor - 1)) ? + (problemShape.k() - blockCoord.k() * tileShape.k()) : + splitkSliceLen; + return GemmCoord{mActual, nActual, kActual}; + } +}; + +/// Block swizzleing function for single core splitk matmul +template +struct SingleCoreSplitkGemmIdentityBlockSwizzle { + /// Data members + static_assert(SwizzleDirection == 0 || SwizzleDirection == 1, "SwizzleDirection only support 0/1"); + + GemmCoord problemShape; + GemmCoord tileMNK; + GemmCoord loopsMNK; + uint32_t mnLoopRangeStart; + uint32_t mnLoopRangeLen; + + /// Methods + CATLASS_DEVICE + SingleCoreSplitkGemmIdentityBlockSwizzle() {} + + /// Methods + CATLASS_DEVICE + SingleCoreSplitkGemmIdentityBlockSwizzle(GemmCoord const &problemShape_, GemmCoord const &l1TileShape) + : problemShape(problemShape_), + tileMNK(l1TileShape) + { + loopsMNK = CeilDiv(problemShape, tileMNK); + GetMNLoopRange(AscendC::GetBlockNum(), AscendC::GetBlockIdx()); + } + + CATLASS_DEVICE + uint32_t GetSingleCoreLoops() const + { + return loopsMNK.k() * mnLoopRangeLen; + } + + CATLASS_DEVICE + void GetMNLoopRange(uint32_t coreNum, uint32_t coreIdx) + { + uint32_t mnLoops = loopsMNK.m() * loopsMNK.n(); + uint32_t bPerCore = CeilDiv(mnLoops, coreNum); + // the first tmpNum cores process bPerCore blocks + // and the last (coreNum - tmpNum) cores process (bPerCore - 1) blocks + // total mnLoops = bPerCore * tmpNum + (bPerCore - 1) * (coreNum - tmpNum) + uint32_t tmpNum = mnLoops - coreNum * (bPerCore - 1); + + mnLoopRangeStart = (coreIdx < tmpNum) ? (coreIdx * bPerCore) : (tmpNum * bPerCore + (coreIdx - tmpNum) * (bPerCore - 1)); + mnLoopRangeLen = (coreIdx < tmpNum) ? bPerCore : (bPerCore - 1); + } + + CATLASS_DEVICE + bool IsAtomicAdd(uint32_t loopIdx) + { + uint32_t kIdx = (loopIdx / mnLoopRangeLen); + return kIdx != 0; + } + + CATLASS_DEVICE + GemmCoord GetBlockCoord(uint32_t loopIdx) + { + uint32_t kIdx = (loopIdx / mnLoopRangeLen); + // mnLoopIdx + uint32_t innerIdx = mnLoopRangeStart + loopIdx % mnLoopRangeLen; + if constexpr (SwizzleDirection == 0) { // Zn + uint32_t tileBlockLoop = CeilDiv(loopsMNK.m(), SwizzleOffset); + uint32_t tileBlockIdx = innerIdx / (SwizzleOffset * loopsMNK.n()); + uint32_t inTileBlockIdx = innerIdx % (SwizzleOffset * loopsMNK.n()); + uint32_t tailOffset = loopsMNK.m() - (tileBlockLoop - 1) * SwizzleOffset; + uint32_t nRow = SwizzleOffset; + + if (tileBlockLoop >= 2 && tileBlockIdx >= tileBlockLoop - 2 && tailOffset <= SwizzleOffset / 2) { + tileBlockLoop = (loopsMNK.m() - tailOffset) / SwizzleOffset; + tileBlockIdx = tileBlockLoop - 1; + nRow = SwizzleOffset + tailOffset; + inTileBlockIdx = innerIdx - SwizzleOffset * loopsMNK.n() * tileBlockIdx; + } else if (tileBlockIdx == tileBlockLoop - 1) { + nRow = loopsMNK.m() - SwizzleOffset * tileBlockIdx; + } + + uint32_t nIdx = inTileBlockIdx / nRow; + if (tileBlockIdx % 2 == 1) { + nIdx = loopsMNK.n() - nIdx - 1; + } + + uint32_t mIdx = tileBlockIdx * SwizzleOffset + inTileBlockIdx % nRow; + if (nIdx % 2 == 1) { + mIdx = tileBlockIdx * SwizzleOffset + (nRow - inTileBlockIdx % nRow - 1); + } + return GemmCoord{mIdx, nIdx, kIdx}; + } else if constexpr (SwizzleDirection == 1) { // Nz + uint32_t tileBlockLoop = CeilDiv(loopsMNK.n(), SwizzleOffset); + uint32_t tileBlockIdx = innerIdx / (SwizzleOffset * loopsMNK.m()); + uint32_t inTileBlockIdx = innerIdx % (SwizzleOffset * loopsMNK.m()); + uint32_t tailOffset = loopsMNK.n() - (tileBlockLoop - 1) * SwizzleOffset; + uint32_t nCol = SwizzleOffset; + + if (tileBlockLoop >= 2 && tileBlockIdx >= tileBlockLoop - 2 && tailOffset < SwizzleOffset / 2) { + tileBlockLoop = (loopsMNK.n() - tailOffset) / SwizzleOffset; + tileBlockIdx = tileBlockLoop - 1; + nCol = SwizzleOffset + tailOffset; + inTileBlockIdx = innerIdx - SwizzleOffset * loopsMNK.m() * tileBlockIdx; + } else if (tileBlockIdx == tileBlockLoop - 1) { + nCol = loopsMNK.n() - SwizzleOffset * tileBlockIdx; + } + + uint32_t mIdx = inTileBlockIdx / nCol; + if (tileBlockIdx % 2 == 1) { + mIdx = loopsMNK.m() - mIdx - 1; + } + + uint32_t nIdx = tileBlockIdx * SwizzleOffset + inTileBlockIdx % nCol; + if (mIdx % 2 == 1) { + nIdx = tileBlockIdx * SwizzleOffset + (nCol - inTileBlockIdx % nCol - 1); + } + return GemmCoord{mIdx, nIdx, kIdx}; + } + } + + CATLASS_DEVICE + GemmCoord GetActualBlockShape(GemmCoord &blockCoord) + { + uint32_t mActual = (blockCoord.m() == (loopsMNK.m() - 1)) ? + (problemShape.m() - blockCoord.m() * tileMNK.m()) : + tileMNK.m(); + uint32_t nActual = (blockCoord.n() == (loopsMNK.n() - 1)) ? + (problemShape.n() - blockCoord.n() * tileMNK.n()) : + tileMNK.n(); + uint32_t kActual = (blockCoord.k() == (loopsMNK.k() - 1)) ? + (problemShape.k() - (loopsMNK.k() - 1) * tileMNK.k()) : + tileMNK.k(); + return GemmCoord{mActual, nActual, kActual}; + } +}; + +struct DynamicSingleCoreSplitkGemmIdentityBlockSwizzle : public SingleCoreSplitkGemmIdentityBlockSwizzle<> { + uint32_t swizzleOffset; + uint32_t swizzleDirection; + + CATLASS_DEVICE + DynamicSingleCoreSplitkGemmIdentityBlockSwizzle(GemmCoord const &problemShape_, GemmCoord const &l1TileShape, + uint32_t swizzleOffset_, uint32_t swizzleDirection_) + : swizzleOffset(swizzleOffset_), + swizzleDirection(swizzleDirection_), + SingleCoreSplitkGemmIdentityBlockSwizzle<>(problemShape_, l1TileShape) + {} + + CATLASS_DEVICE + GemmCoord GetBlockCoord(uint32_t loopIdx) + { + uint32_t kIdx = (loopIdx / mnLoopRangeLen); + // mnLoopIdx + uint32_t innerIdx = mnLoopRangeStart + loopIdx % mnLoopRangeLen; + if (swizzleDirection == 0) { // Zn + uint32_t tileBlockLoop = CeilDiv(loopsMNK.m(), swizzleOffset); + uint32_t tileBlockIdx = innerIdx / (swizzleOffset * loopsMNK.n()); + uint32_t inTileBlockIdx = innerIdx % (swizzleOffset * loopsMNK.n()); + uint32_t tailOffset = loopsMNK.m() - (tileBlockLoop - 1) * swizzleOffset; + uint32_t nRow = swizzleOffset; + + if (tileBlockLoop >= 2 && tileBlockIdx >= tileBlockLoop - 2 && tailOffset <= swizzleOffset / 2) { + tileBlockLoop = (loopsMNK.m() - tailOffset) / swizzleOffset; + tileBlockIdx = tileBlockLoop - 1; + nRow = swizzleOffset + tailOffset; + inTileBlockIdx = innerIdx - swizzleOffset * loopsMNK.n() * tileBlockIdx; + } else if (tileBlockIdx == tileBlockLoop - 1) { + nRow = loopsMNK.m() - swizzleOffset * tileBlockIdx; + } + + uint32_t nIdx = inTileBlockIdx / nRow; + if (tileBlockIdx % 2 == 1) { + nIdx = loopsMNK.n() - nIdx - 1; + } + + uint32_t mIdx = tileBlockIdx * swizzleOffset + inTileBlockIdx % nRow; + if (nIdx % 2 == 1) { + mIdx = tileBlockIdx * swizzleOffset + (nRow - inTileBlockIdx % nRow - 1); + } + return GemmCoord{mIdx, nIdx, kIdx}; + } else { // Nz + uint32_t tileBlockLoop = CeilDiv(loopsMNK.n(), swizzleOffset); + uint32_t tileBlockIdx = innerIdx / (swizzleOffset * loopsMNK.m()); + uint32_t inTileBlockIdx = innerIdx % (swizzleOffset * loopsMNK.m()); + uint32_t tailOffset = loopsMNK.n() - (tileBlockLoop - 1) * swizzleOffset; + uint32_t nCol = swizzleOffset; + + if (tileBlockLoop >= 2 && tileBlockIdx >= tileBlockLoop - 2 && tailOffset < swizzleOffset / 2) { + tileBlockLoop = (loopsMNK.n() - tailOffset) / swizzleOffset; + tileBlockIdx = tileBlockLoop - 1; + nCol = swizzleOffset + tailOffset; + inTileBlockIdx = innerIdx - swizzleOffset * loopsMNK.m() * tileBlockIdx; + } else if (tileBlockIdx == tileBlockLoop - 1) { + nCol = loopsMNK.n() - swizzleOffset * tileBlockIdx; + } + + uint32_t mIdx = inTileBlockIdx / nCol; + if (tileBlockIdx % 2 == 1) { + mIdx = loopsMNK.m() - mIdx - 1; + } + + uint32_t nIdx = tileBlockIdx * swizzleOffset + inTileBlockIdx % nCol; + if (mIdx % 2 == 1) { + nIdx = tileBlockIdx * swizzleOffset + (nCol - inTileBlockIdx % nCol - 1); + } + return GemmCoord{mIdx, nIdx, kIdx}; + } + } +}; + +/// Block swizzling function for Gemms +template +struct GemmIdentityBlockSwizzleL1FullLoad { + /// Data members + + GemmCoord problemShape; + MatrixCoord tileMN; + MatrixCoord loopsMN; + + uint32_t loopsPerCore; + uint32_t loopsTail; + uint32_t aicCoreNum; + + /// Methods + + CATLASS_DEVICE + GemmIdentityBlockSwizzleL1FullLoad() {} + + CATLASS_DEVICE + GemmIdentityBlockSwizzleL1FullLoad(GemmCoord const &problemShape_, MatrixCoord const &tileMN_) + : problemShape(problemShape_), + tileMN(tileMN_) + { + loopsMN = CeilDiv(MatrixCoord(problemShape.GetCoordMN()), tileMN); + uint32_t loopsTotalNum = GetCoreLoops(); + aicCoreNum = AscendC::GetBlockNum(); + loopsPerCore = loopsTotalNum / aicCoreNum; + loopsTail = loopsTotalNum % aicCoreNum; + } + + CATLASS_DEVICE + GemmIdentityBlockSwizzleL1FullLoad(GemmCoord const &problemShape_, MatrixCoord const &tileMN_, + MatrixCoord const &loopsMN_) + : problemShape(problemShape_), + tileMN(tileMN_), + loopsMN(loopsMN_) + { + uint32_t loopsTotalNum = GetCoreLoops(); + aicCoreNum = AscendC::GetBlockNum(); + loopsPerCore = loopsTotalNum / aicCoreNum; + loopsTail = loopsTotalNum % aicCoreNum; + } + + CATLASS_DEVICE + void Update(GemmCoord const &problemShape_, MatrixCoord const &tileMN_) + { + problemShape = problemShape_; + tileMN = tileMN_; + loopsMN = CeilDiv(MatrixCoord(problemShape.GetCoordMN()), tileMN); + + uint32_t loopsTotalNum = GetCoreLoops(); + aicCoreNum = AscendC::GetBlockNum(); + loopsPerCore = loopsTotalNum / aicCoreNum; + loopsTail = loopsTotalNum % aicCoreNum; + } + + CATLASS_DEVICE + void Update(GemmCoord const &problemShape_, MatrixCoord const &tileMN_, MatrixCoord const &loopsMN_) + { + problemShape = problemShape_; + tileMN = tileMN_; + loopsMN = loopsMN_; + + uint32_t loopsTotalNum = GetCoreLoops(); + aicCoreNum = AscendC::GetBlockNum(); + loopsPerCore = loopsTotalNum / aicCoreNum; + loopsTail = loopsTotalNum % aicCoreNum; + } + + CATLASS_DEVICE + uint32_t GetCoreLoops() const + { + return loopsMN.row() * loopsMN.column(); + } + + ////// WARNING: current strategy not support GetBatchIdx() + + CATLASS_DEVICE + GemmCoord GetBlockCoord(uint32_t taskIdx) + { + // calculate innerIdx from taskIdx + uint32_t CoreIdx = taskIdx % aicCoreNum; + uint32_t innerCoreIdx = taskIdx / aicCoreNum; + uint32_t innerIdx = CoreIdx * loopsPerCore + innerCoreIdx; + if (CoreIdx < loopsTail) { + innerIdx += CoreIdx; + } else { + innerIdx += loopsTail; + } + // calculate block location in swizzle, using innerIdx + if constexpr (SwizzleDirection == 0) { // Zn + uint32_t tileBlockLoop = CeilDiv(loopsMN.row(), SwizzleOffset); + uint32_t tileBlockIdx = innerIdx / (SwizzleOffset * loopsMN.column()); + uint32_t inTileBlockIdx = innerIdx % (SwizzleOffset * loopsMN.column()); + + uint32_t nRow = SwizzleOffset; + if (tileBlockIdx == tileBlockLoop - 1) { + nRow = loopsMN.row() - SwizzleOffset * tileBlockIdx; + } + uint32_t mIdx = tileBlockIdx * SwizzleOffset + inTileBlockIdx % nRow; + uint32_t nIdx = inTileBlockIdx / nRow; + if (tileBlockIdx % 2 == 1) { + nIdx = loopsMN.column() - nIdx - 1; + } + return GemmCoord{mIdx, nIdx, 0}; + } else if constexpr (SwizzleDirection == 1) { // Nz + uint32_t tileBlockLoop = CeilDiv(loopsMN.column(), SwizzleOffset); + uint32_t tileBlockIdx = innerIdx / (SwizzleOffset * loopsMN.row()); + uint32_t inTileBlockIdx = innerIdx % (SwizzleOffset * loopsMN.row()); + + uint32_t nCol = SwizzleOffset; + if (tileBlockIdx == tileBlockLoop - 1) { + nCol = loopsMN.column() - SwizzleOffset * tileBlockIdx; + } + uint32_t mIdx = inTileBlockIdx / nCol; + uint32_t nIdx = tileBlockIdx * SwizzleOffset + inTileBlockIdx % nCol; + if (tileBlockIdx % 2 == 1) { + mIdx = loopsMN.row() - mIdx - 1; + } + return GemmCoord{mIdx, nIdx, 0}; + } + } + + CATLASS_DEVICE + GemmCoord GetActualBlockShape(GemmCoord blockCoord) + { + uint32_t mActual = (blockCoord.m() == (loopsMN.row() - 1)) ? + (problemShape.m() - blockCoord.m() * tileMN.row()) : + tileMN.row(); + uint32_t nActual = (blockCoord.n() == (loopsMN.column() - 1)) ? + (problemShape.n() - blockCoord.n() * tileMN.column()) : + tileMN.column(); + uint32_t kActual = problemShape.k(); + return GemmCoord{mActual, nActual, kActual}; + } +}; + +struct DynamicSplitkGemmIdentityBlockSwizzle : public SplitkGemmIdentityBlockSwizzle<> { + uint32_t swizzleOffset{1}; + uint32_t swizzleDirection{0}; + + CATLASS_DEVICE + DynamicSplitkGemmIdentityBlockSwizzle(GemmCoord const &problemShape_, GemmCoord const &tileMNK_, + uint32_t splitkFactor_, uint32_t swizzleOffset_, uint32_t swizzleDirection_) + : swizzleOffset(swizzleOffset_), + swizzleDirection(swizzleDirection_), + SplitkGemmIdentityBlockSwizzle<>(problemShape_, tileMNK_, splitkFactor_) + {} + + CATLASS_DEVICE + DynamicSplitkGemmIdentityBlockSwizzle(GemmCoord const &problemShape_, GemmCoord const &tileMNK_, + uint32_t splitkFactor_) + : SplitkGemmIdentityBlockSwizzle<>(problemShape_, tileMNK_, splitkFactor_) + {} + + CATLASS_DEVICE + DynamicSplitkGemmIdentityBlockSwizzle() {} + + CATLASS_DEVICE + void SetSwizzleParams(uint32_t swizzleOffset_, uint32_t swizzleDirection_) + { + swizzleOffset = swizzleOffset_; + swizzleDirection = swizzleDirection_; + } + + CATLASS_DEVICE + GemmCoord GetBlockCoord(uint32_t taskIdx) + { + uint32_t splitkSliceIdx = GetSplitkSliceIdx(taskIdx); + uint32_t kIdx = GetKIdxBySplitkSliceIdx(splitkSliceIdx); + + uint32_t innerIdx = taskIdx % (loopsMNK.m() * loopsMNK.n()); + if (swizzleDirection == 0) { // Zn + uint32_t tileBlockLoop = CeilDiv(loopsMNK.m(), swizzleOffset); + uint32_t tileBlockIdx = innerIdx / (swizzleOffset * loopsMNK.n()); + uint32_t inTileBlockIdx = innerIdx % (swizzleOffset * loopsMNK.n()); + + uint32_t nRow = swizzleOffset; + if (tileBlockIdx == tileBlockLoop - 1) { + nRow = loopsMNK.m() - swizzleOffset * tileBlockIdx; + } + uint32_t mIdx = tileBlockIdx * swizzleOffset + inTileBlockIdx % nRow; + uint32_t nIdx = inTileBlockIdx / nRow; + if (tileBlockIdx % 2 == 1) { + nIdx = loopsMNK.n() - nIdx - 1; + } + return GemmCoord{mIdx, nIdx, kIdx}; + } else { // Nz + uint32_t tileBlockLoop = CeilDiv(loopsMNK.n(), swizzleOffset); + uint32_t tileBlockIdx = innerIdx / (swizzleOffset * loopsMNK.m()); + uint32_t inTileBlockIdx = innerIdx % (swizzleOffset * loopsMNK.m()); + + uint32_t nCol = swizzleOffset; + if (tileBlockIdx == tileBlockLoop - 1) { + nCol = loopsMNK.n() - swizzleOffset * tileBlockIdx; + } + uint32_t mIdx = inTileBlockIdx / nCol; + uint32_t nIdx = tileBlockIdx * swizzleOffset + inTileBlockIdx % nCol; + if (tileBlockIdx % 2 == 1) { + mIdx = loopsMNK.m() - mIdx - 1; + } + return GemmCoord{mIdx, nIdx, kIdx}; + } + } +}; + +/// Block swizzling function for streamk Gemms +template +struct StreamkGemmIdentityBlockSwizzle { + /// Data members + + GemmCoord problemShape; + GemmCoord tileMNK; + GemmCoord loopsMNK; + uint32_t blockDim; + // Represents the number of ordinary blocks. + uint32_t normalBlocks; + // Represents the number of blocks that need to be split according to the Streamk algorithm. + uint32_t streamkBlocks; + // If computing streamk blocks, how many kTile should each core process + uint32_t kTileNumPerCore; + // Since the number of kTiles may not be evenly divisible by the number of cores, + // some cores need to compute one additoinal kTile compared to others + uint32_t kTileRemain; + + struct StreamkBlockDec { + // If the current block is a ordinary block, set it to false. + // If it is a streamk block, set it to true. + bool isStreamkBlock{false}; + // If the current block is a cross block-meaning it consists partly of the tail k portion of the + // current basic block and partly of the head k portion of the next basic block, set it to true. + bool isCrossBlock{false}; + + GemmCoord blockCoord; + GemmCoord actualBlockShape; + GemmCoord streamkBlockCoord; + GemmCoord streamkActualBlockShape; + }; + + /// Methods + + CATLASS_DEVICE + StreamkGemmIdentityBlockSwizzle() {} + + CATLASS_DEVICE + StreamkGemmIdentityBlockSwizzle( + GemmCoord const &problemShape_, GemmCoord const &tileMNK_, uint32_t blockDim_) + : problemShape(problemShape_), + tileMNK(tileMNK_), + blockDim(blockDim_) + { + loopsMNK = CeilDiv(problemShape, tileMNK); + streamkBlocks = loopsMNK.m() * loopsMNK.n() % blockDim; + normalBlocks = loopsMNK.m() * loopsMNK.n() - streamkBlocks; + kTileNumPerCore = streamkBlocks * loopsMNK.k() / blockDim; + kTileRemain = streamkBlocks * loopsMNK.k() % blockDim; + } + + CATLASS_DEVICE + uint32_t GetCoreLoops() const + { + uint32_t blocksNum = loopsMNK.m() * loopsMNK.n(); + uint32_t coreLoops = blocksNum / blockDim * blockDim + min(streamkBlocks * loopsMNK.k(), blockDim); + return coreLoops; + } + + CATLASS_DEVICE + GemmCoord GetBlockCoord(uint32_t taskIdx) + { + uint32_t innerIdx = taskIdx % GetCoreLoops(); + if constexpr (SwizzleDirection == 0) { // Zn + uint32_t tileBlockLoop = CeilDiv(loopsMNK.m(), SwizzleOffset); + uint32_t tileBlockIdx = innerIdx / (SwizzleOffset * loopsMNK.n()); + uint32_t inTileBlockIdx = innerIdx % (SwizzleOffset * loopsMNK.n()); + + uint32_t nRow = SwizzleOffset; + if (tileBlockIdx == tileBlockLoop - 1) { + nRow = loopsMNK.m() - SwizzleOffset * tileBlockIdx; + } + uint32_t mIdx = tileBlockIdx * SwizzleOffset + inTileBlockIdx % nRow; + uint32_t nIdx = inTileBlockIdx / nRow; + if (tileBlockIdx % 2 == 1) { + nIdx = loopsMNK.n() - nIdx - 1; + } + return GemmCoord{mIdx, nIdx, 0}; + } else if constexpr (SwizzleDirection == 1) { // Nz + uint32_t tileBlockLoop = CeilDiv(loopsMNK.n(), SwizzleOffset); + uint32_t tileBlockIdx = innerIdx / (SwizzleOffset * loopsMNK.m()); + uint32_t inTileBlockIdx = innerIdx % (SwizzleOffset * loopsMNK.m()); + + uint32_t nCol = SwizzleOffset; + if (tileBlockIdx == tileBlockLoop - 1) { + nCol = loopsMNK.n() - SwizzleOffset * tileBlockIdx; + } + uint32_t mIdx = inTileBlockIdx / nCol; + uint32_t nIdx = tileBlockIdx * SwizzleOffset + inTileBlockIdx % nCol; + if (tileBlockIdx % 2 == 1) { + mIdx = loopsMNK.m() - mIdx - 1; + } + return GemmCoord{mIdx, nIdx, 0}; + } + } + + CATLASS_DEVICE + GemmCoord GetActualBlockShape(GemmCoord blockCoord) + { + uint32_t mActual = + (blockCoord.m() == (loopsMNK.m() - 1)) ? (problemShape.m() - blockCoord.m() * tileMNK.m()) : tileMNK.m(); + uint32_t nActual = + (blockCoord.n() == (loopsMNK.n() - 1)) ? (problemShape.n() - blockCoord.n() * tileMNK.n()) : tileMNK.n(); + uint32_t kActual = problemShape.k(); + return GemmCoord{mActual, nActual, kActual}; + } + + CATLASS_DEVICE + uint32_t GetStreamkBlockNum() { return streamkBlocks; } + + CATLASS_DEVICE + uint32_t GetNormalBlockNum() { return normalBlocks; } + + CATLASS_DEVICE + void GetStreamkBlockDec(uint32_t loopIdx, StreamkBlockDec &streamkBlockDec) + { + if (loopIdx < normalBlocks) { + streamkBlockDec.isStreamkBlock = false; + streamkBlockDec.isCrossBlock = false; + streamkBlockDec.blockCoord = GetBlockCoord(loopIdx); + streamkBlockDec.actualBlockShape = GetActualBlockShape(streamkBlockDec.blockCoord); + } else { + streamkBlockDec.isStreamkBlock = true; + uint32_t curKTileNum = kTileNumPerCore; + uint32_t kTileIdx = ((loopIdx - normalBlocks) * kTileNumPerCore + kTileRemain); + if (loopIdx - normalBlocks < kTileRemain) { + curKTileNum += 1; + kTileIdx = ((loopIdx - normalBlocks) * curKTileNum); + } + uint32_t streamkBlockIdx = kTileIdx / loopsMNK.k(); + GemmCoord blockCoordTmp = GetBlockCoord(normalBlocks + streamkBlockIdx); + streamkBlockDec.blockCoord = GemmCoord{blockCoordTmp.m(), blockCoordTmp.n(), kTileIdx % loopsMNK.k()}; + + uint32_t nextStreamkBlockIdx = (kTileIdx + curKTileNum) / loopsMNK.k(); + streamkBlockDec.isCrossBlock = false; + GemmCoord actualBlockShapeTmp = GetActualBlockShape(streamkBlockDec.blockCoord); + if ((kTileIdx % loopsMNK.k() + curKTileNum) * tileMNK.k() > problemShape.k()) { + streamkBlockDec.actualBlockShape = GemmCoord{actualBlockShapeTmp.m(), actualBlockShapeTmp.n(), + problemShape.k() - (kTileIdx % loopsMNK.k()) * tileMNK.k()}; + } else { + streamkBlockDec.actualBlockShape = + GemmCoord{actualBlockShapeTmp.m(), actualBlockShapeTmp.n(), curKTileNum * tileMNK.k()}; + } + if (kTileIdx % loopsMNK.k() + curKTileNum > loopsMNK.k()) { + streamkBlockDec.isCrossBlock = true; + streamkBlockDec.streamkBlockCoord = GetBlockCoord(normalBlocks + nextStreamkBlockIdx); + GemmCoord streamkActualBlockShapeTmp = GetActualBlockShape(streamkBlockDec.streamkBlockCoord); + streamkBlockDec.streamkActualBlockShape = GemmCoord{streamkActualBlockShapeTmp.m(), + streamkActualBlockShapeTmp.n(), ((kTileIdx + curKTileNum) % loopsMNK.k()) * tileMNK.k()}; + } + } + } + + CATLASS_DEVICE + uint32_t GetCoreIdx(uint32_t streamkBlockIdx) + { + if (streamkBlockIdx * loopsMNK.k() > kTileRemain * (kTileNumPerCore + 1)) { + return kTileRemain + + (streamkBlockIdx * loopsMNK.k() - kTileRemain * (kTileNumPerCore + 1)) / kTileNumPerCore; + } else { + return streamkBlockIdx * loopsMNK.k() / (kTileNumPerCore + 1); + } + } + + CATLASS_DEVICE + bool IsCross(uint32_t streamkBlockIdx) + { + // If the head part of the current block and the tail part of the previous block are computed by + // the same core, return true. + if (streamkBlockIdx * loopsMNK.k() > kTileRemain * (kTileNumPerCore + 1)) { + if ((streamkBlockIdx * loopsMNK.k() - kTileRemain * (kTileNumPerCore + 1)) % kTileNumPerCore == 0) { + return false; + } else { + return true; + } + } else { + if ((streamkBlockIdx * loopsMNK.k()) % (kTileNumPerCore + 1) == 0) { + return false; + } else { + return true; + } + } + } +}; + +struct DynamicStreamkGemmIdentityBlockSwizzle : public StreamkGemmIdentityBlockSwizzle<> { + uint32_t swizzleOffset{1}; + uint32_t swizzleDirection{0}; + + CATLASS_DEVICE + DynamicStreamkGemmIdentityBlockSwizzle(GemmCoord const &problemShape_, GemmCoord const &tileShape_, + uint32_t coreNum_, uint32_t swizzleOffset_, uint32_t swizzleDirection_) + : swizzleOffset(swizzleOffset_), + swizzleDirection(swizzleDirection_), + StreamkGemmIdentityBlockSwizzle<>(problemShape_, tileShape_, coreNum_) + {} + + CATLASS_DEVICE + DynamicStreamkGemmIdentityBlockSwizzle(GemmCoord const &problemShape_, GemmCoord const &tileShape_, + uint32_t coreNum_) + : StreamkGemmIdentityBlockSwizzle<>(problemShape_, tileShape_, coreNum_) + {} + + CATLASS_DEVICE + DynamicStreamkGemmIdentityBlockSwizzle() {} + + CATLASS_DEVICE + void SetSwizzleParams(uint32_t swizzleOffset_, uint32_t swizzleDirection_) + { + swizzleOffset = swizzleOffset_; + swizzleDirection = swizzleDirection_; + } + + CATLASS_DEVICE + void GetStreamkBlockDec(uint32_t loopIdx, StreamkBlockDec &streamkBlockDec) + { + if (loopIdx < normalBlocks) { + streamkBlockDec.isStreamkBlock = false; + streamkBlockDec.isCrossBlock = false; + streamkBlockDec.blockCoord = GetBlockCoord(loopIdx); + streamkBlockDec.actualBlockShape = GetActualBlockShape(streamkBlockDec.blockCoord); + } else { + streamkBlockDec.isStreamkBlock = true; + uint32_t curKTileNum = kTileNumPerCore; + uint32_t kTileIdx = ((loopIdx - normalBlocks) * kTileNumPerCore + kTileRemain); + if (loopIdx - normalBlocks < kTileRemain) { + curKTileNum += 1; + kTileIdx = ((loopIdx - normalBlocks) * curKTileNum); + } + uint32_t streamkBlockIdx = kTileIdx / loopsMNK.k(); + GemmCoord blockCoordTmp = GetBlockCoord(normalBlocks + streamkBlockIdx); + streamkBlockDec.blockCoord = GemmCoord{blockCoordTmp.m(), blockCoordTmp.n(), kTileIdx % loopsMNK.k()}; + + uint32_t nextStreamkBlockIdx = (kTileIdx + curKTileNum) / loopsMNK.k(); + streamkBlockDec.isCrossBlock = false; + GemmCoord actualBlockShapeTmp = GetActualBlockShape(streamkBlockDec.blockCoord); + if ((kTileIdx % loopsMNK.k() + curKTileNum) * tileMNK.k() > problemShape.k()) { + streamkBlockDec.actualBlockShape = GemmCoord{actualBlockShapeTmp.m(), actualBlockShapeTmp.n(), + problemShape.k() - (kTileIdx % loopsMNK.k()) * tileMNK.k()}; + } else { + streamkBlockDec.actualBlockShape = + GemmCoord{actualBlockShapeTmp.m(), actualBlockShapeTmp.n(), curKTileNum * tileMNK.k()}; + } + if (kTileIdx % loopsMNK.k() + curKTileNum > loopsMNK.k()) { + streamkBlockDec.isCrossBlock = true; + streamkBlockDec.streamkBlockCoord = GetBlockCoord(normalBlocks + nextStreamkBlockIdx); + GemmCoord streamkActualBlockShapeTmp = GetActualBlockShape(streamkBlockDec.streamkBlockCoord); + streamkBlockDec.streamkActualBlockShape = GemmCoord{streamkActualBlockShapeTmp.m(), + streamkActualBlockShapeTmp.n(), ((kTileIdx + curKTileNum) % loopsMNK.k()) * tileMNK.k()}; + } + } + } + + CATLASS_DEVICE + GemmCoord GetBlockCoord(uint32_t taskIdx) + { + uint32_t innerIdx = taskIdx % GetCoreLoops(); + if (swizzleDirection == 0) { // Zn + uint32_t tileBlockLoop = CeilDiv(loopsMNK.m(), swizzleOffset); + uint32_t tileBlockIdx = innerIdx / (swizzleOffset * loopsMNK.n()); + uint32_t inTileBlockIdx = innerIdx % (swizzleOffset * loopsMNK.n()); + + uint32_t nRow = swizzleOffset; + if (tileBlockIdx == tileBlockLoop - 1) { + nRow = loopsMNK.m() - swizzleOffset * tileBlockIdx; + } + uint32_t mIdx = tileBlockIdx * swizzleOffset + inTileBlockIdx % nRow; + uint32_t nIdx = inTileBlockIdx / nRow; + if (tileBlockIdx % 2 == 1) { + nIdx = loopsMNK.n() - nIdx - 1; + } + return GemmCoord{mIdx, nIdx, 0}; + } else { // Nz + uint32_t tileBlockLoop = CeilDiv(loopsMNK.n(), swizzleOffset); + uint32_t tileBlockIdx = innerIdx / (swizzleOffset * loopsMNK.m()); + uint32_t inTileBlockIdx = innerIdx % (swizzleOffset * loopsMNK.m()); + + uint32_t nCol = swizzleOffset; + if (tileBlockIdx == tileBlockLoop - 1) { + nCol = loopsMNK.n() - swizzleOffset * tileBlockIdx; + } + uint32_t mIdx = inTileBlockIdx / nCol; + uint32_t nIdx = tileBlockIdx * swizzleOffset + inTileBlockIdx % nCol; + if (tileBlockIdx % 2 == 1) { + mIdx = loopsMNK.m() - mIdx - 1; + } + return GemmCoord{mIdx, nIdx, 0}; + } + } +}; + +} // namespace Catlass::Gemm::Block + +#endif // CATLASS_KDA_BLOCK_SWIZZLE_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/kda_gemm_dispatch_policy.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/kda_gemm_dispatch_policy.hpp new file mode 100644 index 0000000000..8c30bb8837 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/kda_gemm_dispatch_policy.hpp @@ -0,0 +1,373 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_GEMM_DISPATCH_POLICY_HPP +#define CATLASS_KDA_GEMM_DISPATCH_POLICY_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_arch.hpp" + +namespace Catlass::Gemm { + +// Block Mmad Policies + +template +struct MmadBase { + using ArchTag = ArchTag_; + static constexpr uint32_t ASYNC = ASYNC_; +}; + +using MmadAtlasA2 = MmadBase; +using MmadAtlasA2Async = MmadBase; + +// Now ENABLE_UNIT_FLAG_ must be false when intput element is int8 +template +struct MmadAtlasA2Pingpong : public MmadAtlasA2 { + static constexpr uint32_t STAGES = 2; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; +}; + +template +struct MmadAtlasA2PingPongWithPrologue : public MmadAtlasA2 { + static constexpr uint32_t STAGES = 2; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; +}; + +template +struct MmadAtlasA2SingleCoreSplitk : public MmadAtlasA2 { + static constexpr uint32_t L1A_STAGES = L1A_STAGES_; + static constexpr uint32_t L1B_STAGES = L1B_STAGES_; + static constexpr uint32_t L0AB_STAGES = 2; + static constexpr uint32_t L0C_STAGES = L0C_STAGES_; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; + static_assert(!(ENABLE_UNIT_FLAG && (L0C_STAGES > 1)), "When L0C_STAGES > 1, can not enable unitflag"); +}; + +template +struct MmadAtlasA2PingpongSliceKWithPrologue : public MmadAtlasA2 { + static constexpr uint32_t STAGES = 2; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; +}; + +template +struct MmadAtlasA2Preload : public MmadAtlasA2 { + static constexpr uint32_t STAGES = 2; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; + static constexpr bool ENABLE_SHUFFLE_K = ENABLE_SHUFFLE_K_; +}; + +struct MmadAtlasA2FAQK : public MmadAtlasA2 { + static constexpr uint32_t STAGES = 2; +}; + +struct MmadAtlasA2FAPV : public MmadAtlasA2 { + static constexpr uint32_t STAGES = 2; +}; + +struct MmadAtlasA2MLAQK : public MmadAtlasA2 { + static constexpr uint32_t STAGES = 2; +}; + +struct MmadAtlasA2MLAPV : public MmadAtlasA2 { + static constexpr uint32_t STAGES = 2; +}; + +struct MmadAtlasA2MLAQKTp1Spec : public MmadAtlasA2 { + static constexpr uint32_t STAGES = 2; +}; + +struct MmadAtlasA2MLAPVTp1Spec : public MmadAtlasA2 { + static constexpr uint32_t STAGES = 2; +}; + +template +struct MmadAtlasA2PreloadAsync : public MmadAtlasA2Async { + static constexpr uint32_t PRELOAD_STAGES = PRELOAD_STAGES_; // Stages of emitting load instruction in advance + static constexpr uint32_t L1_STAGES = L1_STAGES_; + static constexpr uint32_t L0A_STAGES = L0A_STAGES_; + static constexpr uint32_t L0B_STAGES = L0B_STAGES_; + static constexpr uint32_t L0C_STAGES = L0C_STAGES_; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; + static constexpr bool ENABLE_SHUFFLE_K = ENABLE_SHUFFLE_K_; +}; + +template +struct MmadAtlasA2PreloadAsyncWithCallback : public MmadAtlasA2PreloadAsync< + PRELOAD_STAGES_, + L1_STAGES_, + L0A_STAGES_, + L0B_STAGES_, + L0C_STAGES_, + ENABLE_UNIT_FLAG_, + ENABLE_SHUFFLE_K_> { +}; + +template +struct MmadAtlasA2DynamicPreloadAsyncWithCallback : public MmadAtlasA2PreloadAsync< + PRELOAD_STAGES_, + L1_STAGES_, + L0A_STAGES_, + L0B_STAGES_, + L0C_STAGES_, + ENABLE_UNIT_FLAG_, + ENABLE_SHUFFLE_K_> { +}; + +template +struct MmadAtlasA2W4A4MatmulPerTokenPerChannelDequant : public MmadAtlasA2PreloadAsyncWithCallback< + PRELOAD_STAGES_, + L1_STAGES_, + L0A_STAGES_, + L0B_STAGES_, + L0C_STAGES_, + ENABLE_UNIT_FLAG_, + ENABLE_SHUFFLE_K_> { +}; +//////////////////// +// new add +template +struct GemmAtlasA2 : public MmadAtlasA2 { + static constexpr uint32_t STAGES = 2; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; + static constexpr bool ENABLE_SHUFFLE_K = ENABLE_SHUFFLE_K_; + static constexpr bool ENABLE_ABBA = ENABLE_ABBA_; +}; + +struct GemvAtlasA2 : public MmadAtlasA2 { + static constexpr uint32_t STAGES = 2; +}; +//////////////////// + +template +struct MmadAtlasA2PingpongBias : public MmadAtlasA2 { + static constexpr uint32_t STAGES = 2; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; +}; + +template +struct MmadAtlasA2FAIQK : public MmadAtlasA2 { + static constexpr uint32_t STAGES = 2; + static constexpr bool PAGED_CACHE_FLAG = PAGED_CACHE_FLAG_; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; +}; + +template +struct MmadAtlasA2FAIPV : public MmadAtlasA2 { + static constexpr uint32_t STAGES = 2; + static constexpr bool PAGED_CACHE_FLAG = PAGED_CACHE_FLAG_; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; +}; + +template +struct MmadAtlasA2FAITailQK : public MmadAtlasA2 { + static constexpr uint32_t STAGES = 2; + static constexpr bool PAGED_CACHE_FLAG = PAGED_CACHE_FLAG_; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; +}; + +template +struct MmadAtlasA2FAITailPV : public MmadAtlasA2 { + static constexpr uint32_t STAGES = 2; + static constexpr bool PAGED_CACHE_FLAG = PAGED_CACHE_FLAG_; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; +}; + +template +struct MmadAtlasA2FullLoadA : public MmadAtlasA2 { + static constexpr uint32_t STAGES = 2; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; +}; + +template +struct MmadFullLoadA : public MmadBase { + static constexpr uint32_t L1A_STAGES = L1A_STAGES_; + static constexpr uint32_t L1B_STAGES = L1B_STAGES_; + static constexpr uint32_t L0A_STAGES = L0A_STAGES_; + static constexpr uint32_t L0B_STAGES = L0B_STAGES_; + static constexpr uint32_t L0C_STAGES = L0C_STAGES_; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; + static constexpr bool ENABLE_SHUFFLE_K = ENABLE_SHUFFLE_K_; + static constexpr bool USE_HF32_MODE = USE_HF32_MODE_; + static constexpr bool ENABLE_L1_RESIDENT = ENABLE_L1_RESIDENT_; +}; + +template +struct MmadAtlasA2W8A16 : public MmadAtlasA2 { + static constexpr uint32_t STAGES = 2; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; + static constexpr bool ENABLE_SHUFFLE_K = ENABLE_SHUFFLE_K_; +}; + +template +struct MmadAtlasA2DynamicCommon : public MmadAtlasA2 { + static constexpr uint32_t STAGES = 2; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; + static constexpr bool ENABLE_SHUFFLE_K = ENABLE_SHUFFLE_K_; +}; + +template +struct MmadAtlasA2DynamicSmall : public MmadAtlasA2 { + static constexpr uint32_t STAGES = STAGES_; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; +}; + +template +struct MmadAtlasA2Streamk : public MmadAtlasA2 { + static constexpr uint32_t STAGES = 2; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; + static constexpr bool ENABLE_SHUFFLE_K = ENABLE_SHUFFLE_K_; +}; + +template +struct MmadAtlasA2DynamicStreamk : public MmadAtlasA2 { + static constexpr uint32_t STAGES = 2; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; + static constexpr bool ENABLE_SHUFFLE_K = ENABLE_SHUFFLE_K_; +}; + +template +struct MmadAtlasA2DynamicSingleCoreSplitk : public MmadAtlasA2 { + static constexpr uint32_t L1A_STAGES = L1A_STAGES_; + static constexpr uint32_t L1B_STAGES = L1B_STAGES_; + static constexpr uint32_t L0AB_STAGES = 2; + static constexpr uint32_t L0C_STAGES = L0C_STAGES_; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; + static_assert(!(ENABLE_UNIT_FLAG && (L0C_STAGES > 1)), "When L0C_STAGES > 1, can not enable unitflag"); +}; + +template +struct MmadAtlasA2DynamicAiv : public MmadAtlasA2 { + static constexpr uint32_t STAGES = STAGES_; + static constexpr uint32_t SCALAR_BUFFER_ELE_NUM = SCALAR_BUFFER_ELE_NUM_; +}; + +// UB must have a capacity of UBA(mTile) + UBB(nTile) + UBC(mTile, nTile) size. +// MmadAtlasA2AivSimple is suitable for cases: mTile <= mTileUbLimits && nTile >= params.n +template +struct MmadAtlasA2DynamicAivSimple : public MmadAtlasA2 { + static constexpr uint32_t SCALAR_BUFFER_ELE_NUM = SCALAR_BUFFER_ELE_NUM_; + static constexpr bool IS_TILE_M = IS_TILE_M_; +}; + +// mTile and nTile must be aligned to 16. +// UB must have a capacity of UBA(mTile) + UBB(nTile) + 2 * UBC(mTile, nTile) size. +// MmadAtlasA2AivTrans is suitable for cases: M > N +template +struct MmadAtlasA2DynamicAivTrans : public MmadAtlasA2 { + static constexpr uint32_t SCALAR_BUFFER_ELE_NUM = SCALAR_BUFFER_ELE_NUM_; +}; + +template +struct MmadAtlasA2Small : public MmadAtlasA2 { + static constexpr uint32_t STAGES = STAGES_; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; + static constexpr bool ENABLE_SHUFFLE_K = ENABLE_SHUFFLE_K_; +}; +// Now ENABLE_UNIT_FLAG_ must be false when intput element is int8 +template +struct MmadPingpong : public MmadBase { + static constexpr uint32_t L1A_STAGES = L1A_STAGES_; + static constexpr uint32_t L1B_STAGES = L1B_STAGES_; + static constexpr uint32_t L0A_STAGES = L0A_STAGES_; + static constexpr uint32_t L0B_STAGES = L0B_STAGES_; + static constexpr uint32_t L0C_STAGES = L0C_STAGES_; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; + static constexpr bool USE_HF32_MODE = USE_HF32_MODE_; + static constexpr bool ENABLE_L1_RESIDENT = ENABLE_L1_RESIDENT_; +}; + +template +struct MmadPreloadAsyncWithCallback : public MmadBase { + static constexpr uint32_t PRELOAD_STAGES = PRELOAD_STAGES_; + static constexpr uint32_t L1A_STAGES = L1A_STAGES_; + static constexpr uint32_t L1B_STAGES = L1B_STAGES_; + static constexpr uint32_t L0A_STAGES = L0A_STAGES_; + static constexpr uint32_t L0B_STAGES = L0B_STAGES_; + static constexpr uint32_t L0C_STAGES = L0C_STAGES_; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; + static constexpr bool ENABLE_SHUFFLE_K = ENABLE_SHUFFLE_K_; + static constexpr bool USE_HF32_MODE = USE_HF32_MODE_; + static constexpr bool ENABLE_L1_RESIDENT = ENABLE_L1_RESIDENT_; +}; + +// 基于TLA提供block层不感知尾块逻辑，tile层感知originShape的流程。使用TileView与MakeTensorLike。 +template +struct MmadPingpongTlaV2 : public MmadBase { + static constexpr uint32_t STAGES = 2; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; +}; + +template +struct SparseMatmulMultiBlockOnKAxis : public MmadBase { + static constexpr uint32_t STAGES = 2; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; +}; + +template +struct MmadFAIQK : public MmadBase { + static constexpr uint32_t STAGES = 2; + static constexpr bool PAGED_CACHE_FLAG = PAGED_CACHE_FLAG_; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; +}; + +template +struct MmadFAITailQK : public MmadBase { + static constexpr uint32_t STAGES = 2; + static constexpr bool PAGED_CACHE_FLAG = PAGED_CACHE_FLAG_; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; +}; + +template +struct MmadFAIPV : public MmadBase { + static constexpr uint32_t STAGES = 2; + static constexpr bool PAGED_CACHE_FLAG = PAGED_CACHE_FLAG_; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; +}; + +template +struct MmadFAITailPV : public MmadBase { + static constexpr uint32_t STAGES = 2; + static constexpr bool PAGED_CACHE_FLAG = PAGED_CACHE_FLAG_; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; +}; + +template +struct MmadDequant : public MmadBase { + static constexpr uint32_t L1A_STAGES = L1A_STAGES_; + static constexpr uint32_t L1B_STAGES = L1B_STAGES_; + static constexpr uint32_t L0A_STAGES = L0A_STAGES_; + static constexpr uint32_t L0B_STAGES = L0B_STAGES_; + static constexpr uint32_t L0C_STAGES = L0C_STAGES_; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; + static constexpr bool USE_HF32_MODE = USE_HF32_MODE_; + static constexpr bool ENABLE_L1_RESIDENT = ENABLE_L1_RESIDENT_; +}; + +template +struct MmadPingpongPertile : public MmadBase { + static constexpr uint32_t STAGES = 2; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; +}; + +} // namespace Catlass::Gemm + +#endif // CATLASS_KDA_GEMM_DISPATCH_POLICY_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/kda_gemm_helper.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/kda_gemm_helper.hpp new file mode 100644 index 0000000000..4ae8778601 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/kda_gemm_helper.hpp @@ -0,0 +1,385 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_GEMM_HELPER_HPP +#define CATLASS_KDA_GEMM_HELPER_HPP + +#include "catlass/arch/kda_arch.hpp" +#include "catlass/kda_catlass.hpp" +#include "catlass/kda_numeric_size.hpp" +#include "catlass/layout/kda_layout.hpp" +#include "catlass/gemm/kda_gemm_type.hpp" +#include "tla/kda_tla_layout.hpp" + +namespace Catlass::Gemm::helper { + +template +struct L1AlignHelper { + static_assert(DEPENDENT_FALSE, "Unsupported align helper, can not find the specialization."); +}; + +template +struct L1AlignHelper { + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t M_ALIGNED = C0_NUM_PER_FRACTAL; + static constexpr uint32_t K_ALIGNED = ELE_NUM_PER_C0; + static constexpr uint32_t N_ALIGNED = ELE_NUM_PER_C0; +}; + +template +struct L1AlignHelper { + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t M_ALIGNED = ELE_NUM_PER_C0; + static constexpr uint32_t K_ALIGNED = ELE_NUM_PER_C0; + static constexpr uint32_t N_ALIGNED = C0_NUM_PER_FRACTAL; +}; + +template +struct L1AlignHelper { + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t M_ALIGNED = C0_NUM_PER_FRACTAL; + static constexpr uint32_t K_ALIGNED = ELE_NUM_PER_C0; + static constexpr uint32_t N_ALIGNED = ELE_NUM_PER_C0; +}; + +template +struct L1AlignHelper { + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t M_ALIGNED = ELE_NUM_PER_C0; + static constexpr uint32_t K_ALIGNED = ELE_NUM_PER_C0; + static constexpr uint32_t N_ALIGNED = C0_NUM_PER_FRACTAL; +}; + +template +struct L1AlignHelper { + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t M_ALIGNED = C0_NUM_PER_FRACTAL; + static constexpr uint32_t K_ALIGNED = ELE_NUM_PER_C0; + static constexpr uint32_t N_ALIGNED = ELE_NUM_PER_C0; +}; + +template +struct L1AlignHelper { + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t M_ALIGNED = ELE_NUM_PER_C0; + static constexpr uint32_t K_ALIGNED = ELE_NUM_PER_C0; + static constexpr uint32_t N_ALIGNED = C0_NUM_PER_FRACTAL; +}; + +template +struct L1AlignHelper { + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t HOWO_ALIGNED = C0_NUM_PER_FRACTAL; +}; + +template +struct L1AlignHelper { + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t COUT_ALIGNED = C0_NUM_PER_FRACTAL; +}; + +template +struct ElementAccumulatorSelector { + static_assert(DEPENDENT_FALSE, + "Unsupported element accumulator selector, can not find the specialization."); +}; + +template <> +struct ElementAccumulatorSelector { + using ElementAccumulator = float; +}; + +template <> +struct ElementAccumulatorSelector { + using ElementAccumulator = float; +}; + +template <> +struct ElementAccumulatorSelector { + using ElementAccumulator = int32_t; +}; + +template <> +struct ElementAccumulatorSelector { + using ElementAccumulator = float; +}; + +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 3510) +template <> +struct ElementAccumulatorSelector { + using ElementAccumulator = float; +}; + +template <> +struct ElementAccumulatorSelector { + using ElementAccumulator = float; +}; + +template <> +struct ElementAccumulatorSelector { + using ElementAccumulator = float; +}; + +template <> +struct ElementAccumulatorSelector { + using ElementAccumulator = float; +}; +#endif + +template <> +struct ElementAccumulatorSelector { + using ElementAccumulator = int32_t; +}; + +template +struct L1ATypeSelector { + static_assert(DEPENDENT_FALSE, + "Unsupported layout selector, can not find the specialization."); +}; + +template +struct L1ATypeSelector> { + using L1AType = Gemm::GemmType; +}; + +template +struct L1ATypeSelector> { + using L1AType = Gemm::GemmType; +}; + +template +struct L1ATypeSelector> { + using L1AType = Gemm::GemmType; +}; + +template +struct L1ATypeSelector> { + using L1AType = Gemm::GemmType; +}; + +template +struct L1ATypeSelector> { + using L1AType = Gemm::GemmType; +}; + +template +struct L1ATypeSelector> { + using L1AType = Gemm::GemmType; +}; + +template +struct L1ATypeSelector> { + using L1AType = Gemm::GemmType; +}; + +template +struct L1ATypeSelector> { + using L1AType = Gemm::GemmType; +}; + +template +struct L1ATypeSelector> { + using L1AType = Gemm::GemmType; +}; + +template +struct L1BTypeSelector { + static_assert(DEPENDENT_FALSE, + "Unsupported layout selector, can not find the specialization."); +}; + +template +struct L1BTypeSelector> { + using L1BType = Gemm::GemmType; +}; + +template +struct L1BTypeSelector> { + using L1BType = Gemm::GemmType; +}; + +template +struct L1BTypeSelector> { + using L1BType = Gemm::GemmType; +}; + +template +struct L1BTypeSelector> { + using L1BType = Gemm::GemmType; +}; + +template +struct L1BTypeSelector> { + using L1BType = Gemm::GemmType; +}; + +template +struct L1BTypeSelector> { + using L1BType = Gemm::GemmType; +}; + +template +struct L1BTypeSelector> { + using L1BType = Gemm::GemmType; +}; + +template +struct L1BTypeSelector> { + using L1BType = Gemm::GemmType; +}; + +template +struct L1BiasTypeSelector { + static_assert(DEPENDENT_FALSE, + "Unsupported layout selector, can not find the specialization."); +}; + +template +struct L1BiasTypeSelector { + using GMBiasType = void; + using L1BiasType = void; + using L0BiasType = void; +}; + +template +struct L1BiasTypeSelector, ElementAccumulator> { + using GMBiasType = Gemm::GemmType; + using L1BiasType = Gemm::GemmType; + using L0BiasType = Gemm::GemmType; +}; + +template +struct L0ALayoutSelector { + static_assert(DEPENDENT_FALSE, + "Unsupported layout selector, can not find the specialization."); +}; + +template <> +struct L0ALayoutSelector { + using Layout = layout::zZ; +}; + +template <> +struct L0ALayoutSelector { + using Layout = layout::zN; +}; + +template +struct L1AlignHelperTla { + static_assert(DEPENDENT_FALSE, "Unsupported align helper tla, can not find the specialization."); +}; + +template +struct L1AlignHelperTla::value>> { + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t M_ALIGNED = C0_NUM_PER_FRACTAL; + static constexpr uint32_t K_ALIGNED = ELE_NUM_PER_C0; + static constexpr uint32_t N_ALIGNED = ELE_NUM_PER_C0; +}; + +template +struct L1AlignHelperTla::value>> { + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t M_ALIGNED = ELE_NUM_PER_C0; + static constexpr uint32_t K_ALIGNED = ELE_NUM_PER_C0; + static constexpr uint32_t N_ALIGNED = C0_NUM_PER_FRACTAL; +}; + +/////////////////////////////////////// +// new add +template <> +struct ElementAccumulatorSelector { + using ElementAccumulator = int32_t; +}; + +template +struct L1AndL0TypeSelectorGemm { + static_assert(DEPENDENT_FALSE, + "Unsupported layout selector, can not find the specialization."); + static_assert(DEPENDENT_FALSE, + "Unsupported layout selector, can not find the specialization."); +}; + +template +struct L1AndL0TypeSelectorGemm, Gemm::GemmType> { + using L1AType = Gemm::GemmType; + using L1BType = Gemm::GemmType; + using L0AType = Gemm::GemmType; + using L0BType = Gemm::GemmType; +}; + +template <> +struct L1AndL0TypeSelectorGemm, Gemm::GemmType> { + using L1AType = Gemm::GemmType; + using L1BType = Gemm::GemmType; + using L0AType = Gemm::GemmType; + using L0BType = Gemm::GemmType; +}; + +template +struct L1AndL0TypeSelectorGemm, Gemm::GemmType> { + using L1AType = Gemm::GemmType; + using L1BType = Gemm::GemmType; + using L0AType = Gemm::GemmType; + using L0BType = Gemm::GemmType; +}; + +template <> +struct L1AndL0TypeSelectorGemm, Gemm::GemmType> { + using L1AType = Gemm::GemmType; + using L1BType = Gemm::GemmType; + using L0AType = Gemm::GemmType; + using L0BType = Gemm::GemmType; +}; + +template +struct L1AndL0TypeSelectorGemm, Gemm::GemmType> { + using L1AType = Gemm::GemmType; + using L1BType = Gemm::GemmType; + using L0AType = Gemm::GemmType; + using L0BType = Gemm::GemmType; +}; + +template +struct L1AndL0TypeSelectorGemm, Gemm::GemmType> { + using L1AType = Gemm::GemmType; + using L1BType = Gemm::GemmType; + using L0AType = Gemm::GemmType; + using L0BType = Gemm::GemmType; +}; + +template <> +struct L1AndL0TypeSelectorGemm, Gemm::GemmType> { + using L1AType = Gemm::GemmType; + using L1BType = Gemm::GemmType; + using L0AType = Gemm::GemmType; + using L0BType = Gemm::GemmType; +}; + +/////////////////////////////////////// +// Check for whether the tile shape is aligned (Default: 32Bytes). +template +struct TileShapeAlignChecker { + static constexpr uint32_t _ALIGN = _aligned * 8; + + static_assert(L1TileShape::M * SizeOfBits::value % _ALIGN == 0, "L1TileShape::M is not aligned."); + static_assert(L0TileShape::M * SizeOfBits::value % _ALIGN == 0, "L0TileShape::M is not aligned."); + static_assert(L1TileShape::K * SizeOfBits::value % _ALIGN == 0, "L1TileShape::K is not aligned."); + static_assert(L1TileShape::K * SizeOfBits::value % _ALIGN == 0, "L1TileShape::K is not aligned."); + static_assert(L0TileShape::K * SizeOfBits::value % _ALIGN == 0, "L0TileShape::K is not aligned."); + static_assert(L0TileShape::K * SizeOfBits::value % _ALIGN == 0, "L0TileShape::K is not aligned."); + static_assert(L1TileShape::N * SizeOfBits::value % _ALIGN == 0, "L1TileShape::N is not aligned."); + static_assert(L0TileShape::N * SizeOfBits::value % _ALIGN == 0, "L0TileShape::N is not aligned."); +}; + +/////////////////////////////////////// +} // namespace Catlass::Gemm::helper + +#endif // CATLASS_KDA_GEMM_HELPER_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/kda_gemm_type.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/kda_gemm_type.hpp new file mode 100644 index 0000000000..9d680d795b --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/kda_gemm_type.hpp @@ -0,0 +1,29 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_GEMM_TYPE_HPP +#define CATLASS_KDA_GEMM_TYPE_HPP + +#include "catlass/kda_catlass.hpp" + +namespace Catlass::Gemm { + +//////////////////////////////////////////////////////////////////// + +template +struct GemmType { + using Element = Element_; + using Layout = Layout_; + static constexpr AscendC::TPosition POSITION = POSITION_; +}; + +} // namespace Catlass::Gemm + +#endif // CATLASS_KDA_GEMM_TYPE_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/ascend950/copy_l0c_to_ub.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/ascend950/copy_l0c_to_ub.hpp new file mode 100644 index 0000000000..3397404af8 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/ascend950/copy_l0c_to_ub.hpp @@ -0,0 +1,166 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_GEMM_TILE_ASCEND950_COPY_L0C_TO_UB_950_HPP +#define CATLASS_GEMM_TILE_ASCEND950_COPY_L0C_TO_UB_950_HPP + +#include "catlass/arch/kda_arch.hpp" +#include "catlass/kda_catlass.hpp" +#include "catlass/gemm/tile/ascend950/kda_ascend950_copy_l0c_to_dst.hpp" +#include "tla/kda_tensor.hpp" + +#if defined(__NPU_ARCH__) && (__NPU_ARCH__ == 3510) +constexpr AscendC::FixpipeConfig CFG_ROW_MAJOR_UB = {AscendC::CO2Layout::ROW_MAJOR, true}; +#endif + +namespace Catlass::Gemm::Tile { + +template +struct CopyL0CToUBTla< + Catlass::Arch::Ascend950, + TensorSrc_, + tla::Tensor, LayoutDst_, CoordDst_, AscendC::TPosition::VECCALC>, + CopyL0CToUBMode::NO_SPLIT, + ScaleGranularity::NO_QUANT, + ReluEnable_, + std::enable_if_t::value>> { + using ArchTag = Catlass::Arch::Ascend950; + using ElementDst = ElementDst_; + using ElementSrc = typename TensorSrc_::Element; + static constexpr auto quantPre = + CopyL0CToDstQuantMode::VALUE; + static constexpr auto reluEn = ReluEnable_; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, uint8_t unitFlag = 0) + { + static_assert( + tla::detail::isRowMajor::value && TensorSrc::position == AscendC::TPosition::CO1 && TensorDst::position == AscendC::TPosition::VECCALC, + "The input parameters do not match. TensorSrc must be L0C, while TensorDst must be UB and RowMajor"); + + AscendC::FixpipeParamsC310 intriParams; + + // Fixpipe layout information + intriParams.nSize = tla::get<1>(dstTensor.originShape()); + intriParams.mSize = tla::get<0>(dstTensor.originShape()); + intriParams.srcStride = tla::get<1, 1>(srcTensor.stride()) / tla::get<0, 0>(srcTensor.stride()); + intriParams.dstStride = tla::get<0>(dstTensor.stride()); + + // Fixpipe auxiliary arguments + intriParams.quantPre = quantPre; + intriParams.reluEn = reluEn; + intriParams.unitFlag = unitFlag; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + + // Call AscendC Fixpipe + AscendC::Fixpipe( + dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], intriParams); + } +}; + +template +struct CopyL0CToUBTla< + Catlass::Arch::Ascend950, + TensorSrc_, + tla::Tensor, LayoutDst_, CoordDst_, AscendC::TPosition::VECCALC>, + CopyL0CToUBMode::SPLIT_M, + ScaleGranularity::NO_QUANT, + ReluEnable_, + std::enable_if_t::value>> { + using ArchTag = Catlass::Arch::Ascend950; + using ElementDst = ElementDst_; + using ElementSrc = typename TensorSrc_::Element; + static constexpr auto quantPre = + CopyL0CToDstQuantMode::VALUE; + static constexpr auto reluEn = ReluEnable_; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, uint8_t unitFlag = 0) + { + static_assert( + tla::detail::isRowMajor::value && TensorSrc::position == AscendC::TPosition::CO1 && TensorDst::position == AscendC::TPosition::VECCALC, + "The input parameters do not match. TensorSrc must be L0C, while TensorDst must be UB and RowMajor"); + + AscendC::FixpipeParamsC310 intriParams; + + // Fixpipe layout information + intriParams.nSize = tla::get<1>(dstTensor.originShape()); + intriParams.mSize = RoundUp(tla::get<0>(dstTensor.originShape()), 2); // m must be even when spilt m + intriParams.srcStride = tla::get<1, 1>(srcTensor.stride()) / tla::get<0, 0>(srcTensor.stride()); + intriParams.dstStride = tla::get<0>(dstTensor.stride()); + + // Fixpipe auxiliary arguments + intriParams.quantPre = quantPre; + intriParams.reluEn = reluEn; + intriParams.unitFlag = unitFlag; + intriParams.dualDstCtl = 1; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + + // Call AscendC Fixpipe + AscendC::Fixpipe( + dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], intriParams); + } +}; + +template +struct CopyL0CToUBTla< + Catlass::Arch::Ascend950, + TensorSrc_, + tla::Tensor, LayoutDst_, CoordDst_, AscendC::TPosition::VECCALC>, + CopyL0CToUBMode::SPLIT_N, + ScaleGranularity::NO_QUANT, + ReluEnable_, + std::enable_if_t::value>> { + using ArchTag = Catlass::Arch::Ascend950; + using ElementDst = ElementDst_; + using ElementSrc = typename TensorSrc_::Element; + static constexpr auto quantPre = + CopyL0CToDstQuantMode::VALUE; + static constexpr auto reluEn = ReluEnable_; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, uint8_t unitFlag = 0) + { + static_assert( + tla::detail::isRowMajor::value && TensorSrc::position == AscendC::TPosition::CO1 && TensorDst::position == AscendC::TPosition::VECCALC, + "The input parameters do not match. TensorSrc must be L0C, while TensorDst must be UB and RowMajor"); + + AscendC::FixpipeParamsC310 intriParams; + + // Fixpipe layout information + intriParams.nSize = RoundUp(tla::get<1>(dstTensor.originShape()), 32); + intriParams.mSize = tla::get<0>(dstTensor.originShape()); // m must be even when spilt m + intriParams.srcStride = tla::get<1, 1>(srcTensor.stride()) / tla::get<0, 0>(srcTensor.stride()); + intriParams.dstStride = tla::get<0>(dstTensor.stride()); + + // Fixpipe auxiliary arguments + intriParams.quantPre = quantPre; + intriParams.reluEn = reluEn; + intriParams.unitFlag = unitFlag; + intriParams.dualDstCtl = 2; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + + // Call AscendC Fixpipe + AscendC::Fixpipe( + dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], intriParams); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace Catlass::Gemm::Tile + +#endif // CATLASS_GEMM_TILE_ASCEND950_COPY_L0C_TO_UB_950_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/ascend950/kda_ascend950_copy_gm_to_l1.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/ascend950/kda_ascend950_copy_gm_to_l1.hpp new file mode 100644 index 0000000000..862a427e49 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/ascend950/kda_ascend950_copy_gm_to_l1.hpp @@ -0,0 +1,727 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_ASCEND950_COPY_GM_TO_L1_HPP +#define CATLASS_KDA_ASCEND950_COPY_GM_TO_L1_HPP + +#include "catlass/arch/kda_arch.hpp" +#include "catlass/kda_catlass.hpp" +#include "catlass/gemm/tile/kda_tile_copy_tla.hpp" +#include "catlass/kda_numeric_size.hpp" +#include "catlass/layout/kda_layout.hpp" +#include "catlass/gemm/kda_gemm_type.hpp" +#include "tla/kda_tensor.hpp" + +namespace Catlass::Gemm::Tile { + +/// Partial specialization for CopyGmToL1, Ascend950, RowMajor in and zN out. +template +struct TileCopyTla< + Arch::Ascend950, + tla::Tensor, LayoutSrc, CoordSrc, AscendC::TPosition::GM>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::A1>, + std::enable_if_t::value && tla::detail::iszN::value>> { + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + TileCopyTla() {}; + + template + CATLASS_DEVICE void operator()( + TensorDst const &dstTensor, + TensorSrc const &srcTensor, + uint32_t ndNum = 1, + uint32_t srcNdMatrixStride = 0, + uint32_t dstNzMatrixStride = 0) + { + static_assert( + tla::detail::isRowMajor::value && tla::detail::iszN::value && TensorSrc::position == AscendC::TPosition::GM && TensorDst::position == AscendC::TPosition::A1, + "The input parameters do not match. TensorSrc must be GM and RowMajor, while TensorDst must be L1 and zN"); + + const uint32_t nValue = tla::get<0>(srcTensor.originShape()); + const uint32_t dValue = tla::get<1>(srcTensor.originShape()); + const uint32_t srcDValue = tla::get<0>(srcTensor.stride()); + const uint32_t dstInnerStrideRow = tla::get<0, 0>(dstTensor.stride()); + const uint32_t dstOuterStrideCol = tla::get<1, 1>(dstTensor.stride()); + + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = ndNum; + intriParams.nValue = nValue; + intriParams.dValue = dValue; + intriParams.srcNdMatrixStride = srcNdMatrixStride; + intriParams.srcDValue = srcDValue; + intriParams.dstNzC0Stride = dstOuterStrideCol / ELE_NUM_PER_C0; + intriParams.dstNzNStride = dstInnerStrideRow / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = dstNzMatrixStride; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + + AscendC::DataCopy(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], intriParams); + } +}; + +/// Partial specialization for CopyGmToL1, Ascend950, zN in and zN out. +template +struct TileCopyTla< + Arch::Ascend950, + tla::Tensor, LayoutSrc, CoordSrc, AscendC::TPosition::GM>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::A1>, + std::enable_if_t::value && tla::detail::iszN::value>> { + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + TileCopyTla() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert( + tla::detail::iszN::value && tla::detail::iszN::value && TensorSrc::position == AscendC::TPosition::GM && TensorDst::position == AscendC::TPosition::A1, + "The input parameters do not match. TensorSrc must be GM and zN, while TensorDst must be L1 and zN"); + + uint32_t blockCount = CeilDiv(tla::get<1>(srcTensor.originShape())); + uint32_t blockLen = tla::get<0>(srcTensor.originShape()); + + AscendC::DataCopyParams repeatParams; + + repeatParams.blockCount = blockCount; + repeatParams.blockLen = blockLen; + repeatParams.srcStride = tla::get<1, 1>(srcTensor.stride()) / ELE_NUM_PER_C0 - blockLen; + repeatParams.dstStride = tla::get<1, 1>(dstTensor.stride()) / ELE_NUM_PER_C0 - blockLen; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + + AscendC::DataCopy(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], repeatParams); + } +}; + +/// Partial specialization for CopyGmToL1, Ascend950, ColumnMajor in and nZ out. +template +struct TileCopyTla< + Arch::Ascend950, + tla::Tensor, LayoutSrc, CoordSrc, AscendC::TPosition::GM>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::A1>, + std::enable_if_t::value && tla::detail::isnZ::value>> { + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + TileCopyTla() {}; + + template + CATLASS_DEVICE void operator()( + TensorDst const &dstTensor, + TensorSrc const &srcTensor, + uint32_t ndNum = 1, + uint32_t srcNdMatrixStride = 0, + uint32_t dstNzMatrixStride = 0) + { + static_assert( + tla::detail::isColumnMajor::value && tla::detail::isnZ::value && TensorSrc::position == AscendC::TPosition::GM && TensorDst::position == AscendC::TPosition::A1, + "The input parameters do not match. TensorSrc must be GM and ColumnMajor, " + "while TensorDst must be L1 and nZ"); + + const uint32_t nValue = tla::get<1>(srcTensor.originShape()); + const uint32_t dValue = tla::get<0>(srcTensor.originShape()); + const uint32_t srcDValue = tla::get<1>(srcTensor.stride()); + const uint32_t dstInnerStrideCol = tla::get<1, 0>(dstTensor.stride()); + const uint32_t dstOuterStrideRow = tla::get<0, 1>(dstTensor.stride()); + + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = ndNum; + intriParams.nValue = nValue; + intriParams.dValue = dValue; + intriParams.srcNdMatrixStride = srcNdMatrixStride; + intriParams.srcDValue = srcDValue; + intriParams.dstNzC0Stride = dstOuterStrideRow / ELE_NUM_PER_C0; + intriParams.dstNzNStride = dstInnerStrideCol / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = dstNzMatrixStride; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + + AscendC::DataCopy(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], intriParams); + } +}; + +/// Partial specialization for CopyGmToL1, Ascend950, nZ in and nZ out. +template +struct TileCopyTla< + Arch::Ascend950, + tla::Tensor, LayoutSrc, CoordSrc, AscendC::TPosition::GM>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::A1>, + std::enable_if_t::value && tla::detail::isnZ::value>> { + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + TileCopyTla() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert( + tla::detail::isnZ::value && tla::detail::isnZ::value && TensorSrc::position == AscendC::TPosition::GM && TensorDst::position == AscendC::TPosition::A1, + "The input parameters do not match. TensorSrc must be GM and nZ, " + "while TensorDst must be L1 and nZ"); + + uint32_t blockCount = CeilDiv(tla::get<0>(srcTensor.originShape())); + uint32_t blockLen = tla::get<1>(srcTensor.originShape()); + + AscendC::DataCopyParams repeatParams; + + repeatParams.blockCount = blockCount; + repeatParams.blockLen = blockLen; + repeatParams.srcStride = tla::get<0, 1>(srcTensor.stride()) / ELE_NUM_PER_C0 - blockLen; + repeatParams.dstStride = tla::get<0, 1>(dstTensor.stride()) / ELE_NUM_PER_C0 - blockLen; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + + AscendC::DataCopy(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], repeatParams); + } +}; + +/// Partial specialization for CopyGmToL1, Ascend950, VectorLayout in and VectorLayout out. +template +struct TileCopyTla< + Arch::Ascend950, + tla::Tensor, LayoutSrc, CoordSrc, AscendC::TPosition::GM>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::A1>, + std::enable_if_t::value && tla::detail::isVector::value>> { + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + TileCopyTla() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert( + tla::detail::isVector::value && tla::detail::isVector::value && TensorSrc::position == AscendC::TPosition::GM && TensorDst::position == AscendC::TPosition::A1, + "The input parameters do not match. TensorSrc must be GM and Vector, " + "while TensorDst must be L1 and Vector"); + + AscendC::DataCopyParams intriParams; + intriParams.blockCount = 1; + intriParams.blockLen = CeilDiv(tla::get<0>(srcTensor.originShape()), ELE_NUM_PER_C0); + intriParams.srcStride = 0; + intriParams.dstStride = 0; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + + AscendC::DataCopy(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], intriParams); + } +}; + +////////////////////////////////////CopyGmToL1(No-TLA, Ascend950)//////////////////////////////////////////////// +template < + class ArchTag, + /// GemmType for matrix operand + class GmType, + class L1Type = void, + class Enable = void> +struct CopyGmToL1 { + static_assert(DEPENDENT_FALSE, "Unsupported copy gm to l1, can not find the specialization."); +}; + +template < + class ArchTag, + /// GemmType for matrix operand + class GmType, + class L1Type = void> +struct CopyGmToL1GMMPTD { + static_assert(DEPENDENT_FALSE, "Unsupported copy gm to l1, can not find the specialization."); +}; + +template < + class ArchTag, + /// GemmType for matrix operand + class GmType, + class L1Type = void> +struct CopyGmToL1DynamicOptimized { + static_assert(DEPENDENT_FALSE, "Unsupported copy gm to l1, can not find the specialization."); +}; + +template +struct CopyGmToL1> { + using LayoutDst = layout::zN; + using LayoutSrc = layout::RowMajor; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + CATLASS_DEVICE + CopyGmToL1() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + const uint32_t dstInnerStrideRow = layoutDst.stride(0); + const uint32_t dstOuterStrideCol = layoutDst.stride(3); + + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = 1; + intriParams.nValue = layoutSrc.shape(0); + intriParams.dValue = layoutSrc.shape(1); + intriParams.srcDValue = layoutSrc.stride(0); + + // strideColsByFractal --> + intriParams.dstNzC0Stride = layoutDst.stride(3) / ELE_NUM_PER_C0; + intriParams.dstNzNStride = layoutDst.stride(0) / ELE_NUM_PER_C0; + + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzMatrixStride = 0; + + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } + + // layoutSrc must be the layout of one of the src matrices + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc, + uint32_t ndNum, uint32_t srcNdMatrixStride, + uint32_t dstNzNStride, uint32_t dstNzMatrixStride, + uint32_t dstNzC0Stride) + { + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = ndNum; + intriParams.nValue = layoutSrc.shape(0); + intriParams.dValue = layoutSrc.shape(1); + intriParams.srcDValue = layoutSrc.stride(0); + + // Manually set the copy stride + intriParams.dstNzNStride = dstNzNStride; + intriParams.dstNzC0Stride = dstNzC0Stride; + intriParams.srcNdMatrixStride = srcNdMatrixStride; + intriParams.dstNzMatrixStride = dstNzMatrixStride; + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } +}; + +/// Partial specialization for CopyGmToL1(AtlasA5 no-tla), ColumnMajor in and nZ out. +template +struct CopyGmToL1> { + using LayoutDst = layout::nZ; + using LayoutSrc = layout::ColumnMajor; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + // Note: AscendC::int4b_t no longer supported on A5-platform + + // Methods + CATLASS_DEVICE + CopyGmToL1() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + // dstOuterStrideRow = tla::get<0, 1>(dstTensor.stride()) --> + // tag.stride(1) --> strideRowsByFractal[语义] + // dstInnerStrideCol = tla::get<1, 0>(dstTensor.stride()) --> + // tag.stride(2) --> strideColsInFractal[语义] + + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = 1; + intriParams.nValue = layoutSrc.shape(1); + intriParams.dValue = layoutSrc.shape(0); + intriParams.srcDValue = layoutSrc.stride(1); + + intriParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0; // Outer stride -- col + intriParams.dstNzNStride = layoutDst.stride(2) / ELE_NUM_PER_C0; + + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzMatrixStride = 0; + + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } +}; + +/// Partial specialization for CopyGmToL1(AtlasA5 no-tla), zN in and zN out. +template +struct CopyGmToL1> { + using LayoutDst = layout::zN; + using LayoutSrc = layout::zN; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + // Note: AscendC::int4b_t no longer supported on A5-platform + + // Methods + CATLASS_DEVICE + CopyGmToL1() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + const uint32_t blockLen = layoutSrc.shape(0) * layoutSrc.shape(1); + AscendC::DataCopyParams repeatParams; + + repeatParams.blockCount = layoutSrc.shape(3); + repeatParams.blockLen = blockLen; + + repeatParams.srcStride = layoutSrc.stride(3) / ELE_NUM_PER_C0 - blockLen; + repeatParams.dstStride = layoutDst.stride(3) / ELE_NUM_PER_C0 - blockLen; + + AscendC::DataCopy(dstTensor, srcTensor, repeatParams); + } +}; + +/// Partial specialization for CopyGmToL1(AtlasA5 no-tla), nZ in and nZ out. +template +struct CopyGmToL1> { + using LayoutDst = layout::nZ; + using LayoutSrc = layout::nZ; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + // Note: AscendC::int4b_t no longer supported on A5-platform + + CATLASS_DEVICE + CopyGmToL1() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + const uint32_t blockLen = layoutSrc.shape(2) * layoutSrc.shape(3); + AscendC::DataCopyParams repeatParams; + + repeatParams.blockCount = layoutSrc.shape(1); + repeatParams.blockLen = blockLen; + + repeatParams.srcStride = layoutSrc.stride(1) / ELE_NUM_PER_C0 - blockLen; + repeatParams.dstStride = layoutDst.stride(1) / ELE_NUM_PER_C0 - blockLen; + + AscendC::DataCopy(dstTensor, srcTensor, repeatParams); + } +}; + +/// Partial specialization for CopyGmToL1(no-tla), AtlasA5, fp8_e8m0_t, MxScaleA RowMajor in and zZ out. +template +struct CopyGmToL1, + Gemm::GemmType, + std::enable_if_t>> { + using LayoutDst = layout::zZ; + using LayoutSrc = layout::zN; + + static constexpr uint32_t ELE_NUM_PER_C0 = 2; + + CopyGmToL1() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + if (layoutSrc.shape(2) != ELE_NUM_PER_C0) { + // std::cerr << "layoutSrc.shape(2) != 2" << std::endl; + return; + } + + AscendC::Dn2NzParams intriParams; + + intriParams.dnNum = 1; + intriParams.nValue = CeilDiv(layoutSrc.shape(2) * layoutSrc.shape(3)); + intriParams.dValue = layoutSrc.shape(0); + intriParams.srcDnMatrixStride = 0; + + intriParams.srcDValue = CeilDiv(layoutSrc.stride(0)); + intriParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0; + intriParams.dstNzNStride = 1; + intriParams.dstNzMatrixStride = 0; + + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } +}; + +/// Partial specialization for AtlasA2, PaddingRowMajor in and zN out. +template +struct CopyGmToL1> { + using LayoutDst = layout::zN; + using LayoutSrc = layout::PaddingRowMajor; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Mehtods + + CATLASS_DEVICE + CopyGmToL1() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = 1; + intriParams.dValue = layoutSrc.orgShape(1); + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzC0Stride = layoutDst.stride(3) / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + + intriParams.nValue = layoutSrc.orgShape(0); + intriParams.srcDValue = layoutSrc.stride(0); + intriParams.dstNzNStride = layoutDst.stride(0) / ELE_NUM_PER_C0; + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } +}; + +/// Partial specialization for AtlasA2, ColumnMajor in and nZ out. +template < + class Element> +struct CopyGmToL1> { + using LayoutDst = layout::nZ; + using LayoutSrc = layout::PaddingColumnMajor; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Mehtods + + CATLASS_DEVICE + CopyGmToL1() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = 1; + intriParams.dValue = layoutSrc.orgShape(0); + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + + intriParams.nValue = layoutSrc.orgShape(1); + intriParams.srcDValue = layoutSrc.stride(2); + intriParams.dstNzNStride = layoutDst.stride(2) / ELE_NUM_PER_C0; + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } +}; + +//////////////////////////// CopyGmToL1DynamicOptimized(Ascend950, No TLA) //////////////////////////// + +/// Partial specialization for Ascend950, zN in and zN out. +template +struct CopyGmToL1DynamicOptimized> : public CopyGmToL1> {}; + +/// Partial specialization for Ascend950, nZ in and nZ out. +template +struct CopyGmToL1DynamicOptimized> : public CopyGmToL1> {}; + +/// Partial specialization for Ascend950, RowMajor in and zN out. +template +struct CopyGmToL1DynamicOptimized> { + using LayoutDst = layout::zN; + using LayoutSrc = layout::RowMajor; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + CATLASS_DEVICE + CopyGmToL1DynamicOptimized() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + if (layoutSrc.shape(0) <= 16) { + // If the number of matrix row is very small, call the regular interval-based data-copy + for (int i = 0; i < layoutSrc.shape(0); ++i) { + AscendC::DataCopyParams dataCopyParams( + CeilDiv(layoutSrc.shape(1), layoutDst.shape(2)), + layoutDst.shape(2) / ELE_NUM_PER_C0, + 0, + (layoutDst.stride(3) - layoutDst.shape(2)) / ELE_NUM_PER_C0); + AscendC::DataCopy(dstTensor[i * layoutDst.shape(2)], srcTensor[i * layoutSrc.stride(0)], dataCopyParams); + } + } else { + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = 1; + intriParams.nValue = layoutSrc.shape(0); + intriParams.dValue = layoutSrc.shape(1); + intriParams.srcDValue = layoutSrc.stride(0); + intriParams.dstNzC0Stride = layoutDst.stride(3) / ELE_NUM_PER_C0; + intriParams.dstNzNStride = layoutDst.stride(0) / ELE_NUM_PER_C0; + + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzMatrixStride = 0; + + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } + } +}; + +/// Partial specialization for Ascend950, ColumnMajor in and nZ out. +template +struct CopyGmToL1DynamicOptimized> { + using LayoutDst = layout::nZ; + using LayoutSrc = layout::ColumnMajor; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + CATLASS_DEVICE + CopyGmToL1DynamicOptimized() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + if (layoutSrc.shape(1) <= 16) { + // If the number of matrix cols is 1, the regular interval-based DataCopy interface can be used instead of + // the ND2NZ DataCopy interface, resulting in higher transfer efficiency. + for (int i = 0; i < layoutSrc.shape(1); ++i) { + AscendC::DataCopyParams dataCopyParams( + CeilDiv(layoutSrc.shape(0), layoutDst.shape(0)), + layoutDst.shape(0) / ELE_NUM_PER_C0, + 0, + (layoutDst.stride(1) - layoutDst.shape(0)) / ELE_NUM_PER_C0); + AscendC::DataCopy( + dstTensor[i * layoutDst.shape(0)], srcTensor[i * layoutSrc.stride(1)], dataCopyParams); + } + } else { + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = 1; + intriParams.nValue = layoutSrc.shape(1); + intriParams.dValue = layoutSrc.shape(0); + intriParams.srcDValue = layoutSrc.stride(1); + + intriParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0; // Outer stride -- col + intriParams.dstNzNStride = layoutDst.stride(2) / ELE_NUM_PER_C0; + + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzMatrixStride = 0; + + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } + } +}; + +//////////////////////////// CopyGmToL1GMMPTD(Ascend950, No TLA) //////////////////////////// +/// Partial specialization for Ascend950, RowMajor in and zN out. +template +struct CopyGmToL1GMMPTD> { + using LayoutDst = layout::zN; + using LayoutSrc = layout::RowMajor; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Mehtods + + CATLASS_DEVICE + CopyGmToL1GMMPTD() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = 1; + intriParams.dValue = layoutSrc.shape(1); + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzC0Stride = layoutDst.stride(3) / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + + if (layoutSrc.shape(0) == 1) { + // If the number of matrix rows is 1, the regular interval-based DataCopy interface can be used instead of + // the ND2NZ DataCopy interface, resulting in higher transfer efficiency. + AscendC::DataCopyParams dataCopyParams( + CeilDiv(layoutSrc.shape(1), layoutDst.shape(2)), + layoutDst.shape(2) / ELE_NUM_PER_C0, + 0, + (layoutDst.stride(3) - layoutDst.shape(2)) / ELE_NUM_PER_C0); + AscendC::DataCopy(dstTensor, srcTensor, dataCopyParams); + } else { + if (layoutSrc.shape(1) != ELE_NUM_PER_C0 || layoutSrc.stride(0) != ELE_NUM_PER_C0) { + intriParams.nValue = layoutSrc.shape(0); + intriParams.srcDValue = layoutSrc.stride(0); + intriParams.dstNzNStride = layoutDst.stride(0) / ELE_NUM_PER_C0; + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } else { + // If the matrix has ELE_NUM_PER_C0 columns and a stride of ELE_NUM_PER_C0, it follows a row-major + // layout in L1, allowing the use of the standard contiguous DataCopy interface for more efficient + // transfers. + AscendC::DataCopy(dstTensor, srcTensor, layoutSrc.shape(0) * layoutSrc.shape(1)); + } + } + } + + // layoutSrc must be the layout of one of the src matrices + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc, + uint32_t ndNum, uint32_t srcNdMatrixStride, + uint32_t dstNzNStride, uint32_t dstNzMatrixStride, + uint32_t dstNzC0Stride) + { + AscendC::Nd2NzParams intriParams; + + intriParams.nValue = layoutSrc.shape(0); + intriParams.dValue = layoutSrc.shape(1); + intriParams.srcDValue = layoutSrc.stride(0); + intriParams.dstNzNStride = dstNzNStride; + intriParams.dstNzC0Stride = dstNzC0Stride; + intriParams.ndNum = ndNum; + intriParams.srcNdMatrixStride = srcNdMatrixStride; + intriParams.dstNzMatrixStride = dstNzMatrixStride; + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace Catlass::Gemm::Tile + +#endif diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/ascend950/kda_ascend950_copy_l0c_to_dst.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/ascend950/kda_ascend950_copy_l0c_to_dst.hpp new file mode 100644 index 0000000000..53d5adcc3d --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/ascend950/kda_ascend950_copy_l0c_to_dst.hpp @@ -0,0 +1,274 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_ASCEND950_COPY_L0C_TO_DST_HPP +#define CATLASS_KDA_ASCEND950_COPY_L0C_TO_DST_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_arch.hpp" +#include "catlass/gemm/kda_gemm_type.hpp" +#include "tla/kda_tensor.hpp" + +namespace Catlass::Gemm::Tile { + +enum class ScaleGranularity { + UNDEFINED = -1, + NO_QUANT = 0, + PER_TENSOR, + PER_CHANNEL, + PER_GROUP +}; + +template < + class ArchTag, + class ElementSrc, + class ElementDst, + ScaleGranularity DEQUANT_GRANULARITY = ScaleGranularity::NO_QUANT> +struct CopyL0CToDstQuantMode { + static_assert(DEPENDENT_FALSE, "Unsupported copy l0c to dst quant mode, can not find the specialization."); +}; + +// CopyL0CToDst fp32 to fp32 +template <> +struct CopyL0CToDstQuantMode< + Arch::Ascend950, + float, float, + ScaleGranularity::NO_QUANT> { + static constexpr auto VALUE = QuantMode_t::NoQuant; +}; + +// CopyL0CToDst cast fp32 to fp16 +template <> +struct CopyL0CToDstQuantMode< + Arch::Ascend950, + float, half, + ScaleGranularity::NO_QUANT> { + static constexpr auto VALUE = QuantMode_t::F322F16; +}; + +// CopyL0CToDst cast fp32 to bf16 +template <> +struct CopyL0CToDstQuantMode< + Arch::Ascend950, + float, bfloat16_t, + ScaleGranularity::NO_QUANT> { + static constexpr auto VALUE = QuantMode_t::F322BF16; +}; + +// CopyL0CToDst cast float to uint8/int8 +template <> +struct CopyL0CToDstQuantMode< + Arch::Ascend950, + float, uint8_t, + ScaleGranularity::PER_TENSOR> { + static constexpr auto VALUE = QuantMode_t::QF322B8_PRE; +}; + +template <> +struct CopyL0CToDstQuantMode< + Arch::Ascend950, + float, uint8_t, + ScaleGranularity::PER_CHANNEL> { + static constexpr auto VALUE = QuantMode_t::VQF322B8_PRE; +}; + +template <> +struct CopyL0CToDstQuantMode< + Arch::Ascend950, + float, int8_t, + ScaleGranularity::PER_TENSOR> { + static constexpr auto VALUE = QuantMode_t::QF322B8_PRE; +}; + +template <> +struct CopyL0CToDstQuantMode< + Arch::Ascend950, + float, int8_t, + ScaleGranularity::PER_CHANNEL> { + static constexpr auto VALUE = QuantMode_t::VQF322B8_PRE; +}; + +// CopyL0CToDst output int32 +template <> +struct CopyL0CToDstQuantMode< + Arch::Ascend950, + int32_t, int32_t, + ScaleGranularity::NO_QUANT> { + static constexpr auto VALUE = QuantMode_t::NoQuant; +}; + +// CopyL0CToDst cast int32_t to fp16 +template <> +struct CopyL0CToDstQuantMode< + Arch::Ascend950, + int32_t, half, + ScaleGranularity::PER_TENSOR> { + static constexpr auto VALUE = QuantMode_t::DEQF16; +}; + +template <> +struct CopyL0CToDstQuantMode< + Arch::Ascend950, + int32_t, half, + ScaleGranularity::PER_CHANNEL> { + static constexpr auto VALUE = QuantMode_t::VDEQF16; +}; + +// CopyL0CToDst cast int32 to uint8/int8 +template <> +struct CopyL0CToDstQuantMode< + Arch::Ascend950, + int32_t, uint8_t, + ScaleGranularity::PER_TENSOR> { + static constexpr auto VALUE = QuantMode_t::REQ8; +}; + +template <> +struct CopyL0CToDstQuantMode< + Arch::Ascend950, + int32_t, uint8_t, + ScaleGranularity::PER_CHANNEL> { + static constexpr auto VALUE = QuantMode_t::VREQ8; +}; + +template <> +struct CopyL0CToDstQuantMode< + Arch::Ascend950, + int32_t, int8_t, + ScaleGranularity::PER_TENSOR> { + static constexpr auto VALUE = QuantMode_t::REQ8; +}; + +template <> +struct CopyL0CToDstQuantMode< + Arch::Ascend950, + int32_t, int8_t, + ScaleGranularity::PER_CHANNEL> { + static constexpr auto VALUE = QuantMode_t::VREQ8; +}; + +// CopyL0CToDst cast int32 to bf16 +template <> +struct CopyL0CToDstQuantMode< + Arch::Ascend950, + int32_t, bfloat16_t, + ScaleGranularity::PER_TENSOR> { + static constexpr auto VALUE = QuantMode_t::QS322BF16_PRE; +}; + +template <> +struct CopyL0CToDstQuantMode< + Arch::Ascend950, + int32_t, bfloat16_t, + ScaleGranularity::PER_CHANNEL> { + static constexpr auto VALUE = QuantMode_t::VQS322BF16_PRE; +}; + +// CopyL0CToDst cast fp32 to fp16 +template <> +struct CopyL0CToDstQuantMode< + Arch::Ascend950, + float, half, + ScaleGranularity::PER_TENSOR> { + static constexpr auto VALUE = QuantMode_t::QF322F16_PRE; +}; + +template <> +struct CopyL0CToDstQuantMode< + Arch::Ascend950, + float, half, + ScaleGranularity::PER_CHANNEL> { + static constexpr auto VALUE = QuantMode_t::VQF322F16_PRE; +}; + +// CopyL0CToDst cast fp32 to bf16 +template <> +struct CopyL0CToDstQuantMode< + Arch::Ascend950, + float, bfloat16_t, + ScaleGranularity::PER_TENSOR> { + static constexpr auto VALUE = QuantMode_t::QF322BF16_PRE; +}; + +template <> +struct CopyL0CToDstQuantMode< + Arch::Ascend950, + float, bfloat16_t, + ScaleGranularity::PER_CHANNEL> { + static constexpr auto VALUE = QuantMode_t::VQF322BF16_PRE; +}; + +// CopyL0CToDst cast fp32 to fp32 +template <> +struct CopyL0CToDstQuantMode< + Arch::Ascend950, + float, float, + ScaleGranularity::PER_TENSOR> { + static constexpr auto VALUE = QuantMode_t::QF322F32_PRE; +}; + +template <> +struct CopyL0CToDstQuantMode< + Arch::Ascend950, + float, float, + ScaleGranularity::PER_CHANNEL> { + static constexpr auto VALUE = QuantMode_t::VQF322F32_PRE; +}; + +template < + class ArchTag, + class ElementAccumulator, + class GmType, + ScaleGranularity DEQUANT_GRANULARITY = ScaleGranularity::NO_QUANT, + bool ReluEnable = false> +struct CopyL0CToGm { + static_assert(DEPENDENT_FALSE, "Unsupported copy l0c to gm, can not find the specialization."); +}; + +///////////////////////////////////////////CopyL0CToGmTla///////////////////////////////////////////////// +// L0C copy mode +struct CopyToGM {}; +struct CopyToL1 {}; + +template < + class ArchTag, + class TensorSrc, + class TensorDst, + ScaleGranularity DEQUANT_GRANULARITY = ScaleGranularity::NO_QUANT, + bool ReluEnable = false, + class Enable = void> +struct CopyL0CToGmTla { + static_assert(DEPENDENT_FALSE, "Unsupported copy l0c to gm, can not find the specialization."); +}; + +///////////////////////////////////////////////////////////////////////////////////////////////////////////// + +enum class CopyL0CToUBMode { + NO_SPLIT = 0, + SPLIT_M, + SPLIT_N, + RESERVED +}; + +template < + class ArchTag, + class TensorSrc, + class TensorDst, + CopyL0CToUBMode CopyMode = CopyL0CToUBMode::NO_SPLIT, + ScaleGranularity DEQUANT_GRANULARITY = ScaleGranularity::NO_QUANT, + bool ReluEnable = false, + class Enable = void> +struct CopyL0CToUBTla { + static_assert(DEPENDENT_FALSE, "Unsupported copy l0c to ub, can not find the specialization."); +}; + +} // namespace Catlass::Gemm::Tile + +#endif // CATLASS_KDA_ASCEND950_COPY_L0C_TO_DST_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/ascend950/kda_ascend950_copy_l0c_to_gm.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/ascend950/kda_ascend950_copy_l0c_to_gm.hpp new file mode 100644 index 0000000000..86266e6dd3 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/ascend950/kda_ascend950_copy_l0c_to_gm.hpp @@ -0,0 +1,533 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_ASCEND950_COPY_L0C_TO_GM_HPP +#define CATLASS_KDA_ASCEND950_COPY_L0C_TO_GM_HPP + +#include "catlass/arch/kda_arch.hpp" +#include "catlass/kda_catlass.hpp" +#include "catlass/gemm/tile/ascend950/kda_ascend950_copy_l0c_to_dst.hpp" +#include "tla/kda_tensor.hpp" + +namespace Catlass::Gemm::Tile { + +template +struct CopyL0CToGmTla< + Catlass::Arch::Ascend950, + TensorSrc_, + tla::Tensor, LayoutDst_, CoordDst_, AscendC::TPosition::GM>, + ScaleGranularity::NO_QUANT, + ReluEnable_, + std::enable_if_t::value>> { + using ArchTag = Catlass::Arch::Ascend950; + using ElementDst = ElementDst_; + using ElementSrc = typename TensorSrc_::Element; + static constexpr auto quantPre = + CopyL0CToDstQuantMode::VALUE; + static constexpr auto reluEn = ReluEnable_; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, uint8_t unitFlag = 0) + { + static_assert( + tla::detail::isRowMajor::value && TensorSrc::position == AscendC::TPosition::CO1 && TensorDst::position == AscendC::TPosition::GM, + "The input parameters do not match. TensorSrc must be L0C, while TensorDst must be GM and RowMajor"); + + AscendC::DataCopyCO12DstParams intriParams; + + intriParams.nSize = tla::get<1>(dstTensor.originShape()); + intriParams.mSize = tla::get<0>(dstTensor.originShape()); + intriParams.dstStride = tla::get<0>(dstTensor.stride()); + intriParams.srcStride = tla::get<1, 1>(srcTensor.stride()) / tla::get<0, 0>(srcTensor.stride()); + intriParams.quantPre = quantPre; + intriParams.nz2ndEn = true; + intriParams.reluPre = reluEn; + intriParams.unitFlag = unitFlag; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + + AscendC::SetFixpipeNz2ndFlag(1, 1, 1); + AscendC::DataCopy(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], intriParams); + } + + template + CATLASS_DEVICE void + operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, uint32_t l0Batch, uint32_t dstNdStride) + { + static_assert( + tla::detail::isRowMajor::value && TensorSrc::position == AscendC::TPosition::CO1 && TensorDst::position == AscendC::TPosition::GM, + "The input parameters do not match. TensorSrc must be L0C, while TensorDst must be GM and RowMajor"); + + const uint32_t L0CM = tla::get<0, 0>(srcTensor.shape()) * tla::get<0, 1>(srcTensor.shape()); + const uint32_t L0CN = tla::get<1, 0>(srcTensor.shape()) * tla::get<1, 1>(srcTensor.shape()); + + // AscendC::FixpipeParamsC310 intriParams; + + // // Fixpipe layout information + // intriParams.nSize = tla::get<1>(dstTensor.shape()); + // intriParams.mSize = tla::get<0>(dstTensor.shape()); + // intriParams.srcStride = tla::get<1, 1>(srcTensor.stride()) / tla::get<0, 0>(srcTensor.stride()); + // intriParams.dstStride = tla::get<0>(dstTensor.stride()); + + // // Fixpipe auxiliary arguments + // intriParams.quantPre = quantPre; + // intriParams.reluEn = reluEn; + // intriParams.unitFlag = 0; + + // intriParams.params.ndNum = l0Batch; + // intriParams.params.srcNdStride = L0CM * L0CN / tla::get<1, 0>(srcTensor.shape()); + // intriParams.params.dstNdStride = dstNdStride; + + // // Call AscendC Fixpipe + // AscendC::Fixpipe( + // dstTensor.data(), srcTensor.data(), intriParams); + + AscendC::DataCopyCO12DstParams intriParams; + + intriParams.nSize = tla::get<1>(dstTensor.originShape()); + intriParams.mSize = tla::get<0>(dstTensor.originShape()); + intriParams.dstStride = tla::get<0>(dstTensor.stride()); + intriParams.srcStride = tla::get<1, 1>(srcTensor.stride()) / tla::get<0, 0>(srcTensor.stride()); + intriParams.quantPre = quantPre; + intriParams.nz2ndEn = true; + intriParams.reluPre = reluEn; + intriParams.unitFlag = 0; + + AscendC::SetFixpipeNz2ndFlag(l0Batch, L0CM * L0CN / tla::get<1, 0>(srcTensor.shape()), dstNdStride); + AscendC::DataCopy(dstTensor.data(), srcTensor.data(), intriParams); + } +}; + +template +struct CopyL0CToGmTla< + Catlass::Arch::Ascend950, + TensorSrc_, + tla::Tensor, LayoutDst_, CoordDst_, AscendC::TPosition::GM>, + ScaleGranularity::NO_QUANT, + ReluEnable_, + std::enable_if_t::value>> { + using ArchTag = Catlass::Arch::Ascend950; + using ElementDst = ElementDst_; + using ElementSrc = typename TensorSrc_::Element; + static constexpr auto quantPre = + CopyL0CToDstQuantMode::VALUE; + static constexpr auto reluEn = ReluEnable_; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, uint8_t unitFlag = 0) + { + static_assert( + tla::detail::iszN::value && TensorSrc::position == AscendC::TPosition::CO1 && TensorDst::position == AscendC::TPosition::GM, + "The input parameters do not match. TensorSrc must be L0C, while TensorDst must be GM and zN"); + + AscendC::DataCopyCO12DstParams intriParams; + + intriParams.nSize = tla::get<1>(dstTensor.originShape()); + intriParams.mSize = tla::get<0>(dstTensor.originShape()); + intriParams.dstStride = tla::get<1, 1>(dstTensor.stride()) / (BYTE_PER_C0 / sizeof(ElementDst)); + intriParams.srcStride = tla::get<1, 1>(srcTensor.stride()) / tla::get<0, 0>(srcTensor.stride()); + intriParams.quantPre = quantPre; + intriParams.nz2ndEn = false; + intriParams.reluPre = reluEn; + intriParams.unitFlag = unitFlag; + + if constexpr (std::is_same_v && std::is_same_v) { + intriParams.channelSplit = true; + } + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + + AscendC::DataCopy(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], intriParams); + } +}; + +template +struct CopyL0CToGmTla< + Catlass::Arch::Ascend950, + TensorSrc_, + tla::Tensor, LayoutDst_, CoordDst_, AscendC::TPosition::GM>, + ScaleGranularity::PER_TENSOR, + ReluEnable_, + std::enable_if_t::value>> { + using ArchTag = Catlass::Arch::Ascend950; + using ElementDst = ElementDst_; + using ElementSrc = typename TensorSrc_::Element; + static constexpr auto quantPre = + CopyL0CToDstQuantMode::VALUE; + static constexpr auto reluEn = ReluEnable_; + + struct Params { + float scale = 1.0; + + CATLASS_HOST_DEVICE + Params() = default; + + CATLASS_HOST_DEVICE + Params(float scalar) + { + scale = scalar; + } + }; + Params params; + + CATLASS_DEVICE + CopyL0CToGmTla() = default; + + CATLASS_DEVICE + CopyL0CToGmTla(Params const ¶ms_) + : params(params_) {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, uint8_t unitFlag = 0) + { + static_assert( + tla::detail::isRowMajor::value && TensorSrc::position == AscendC::TPosition::CO1 && TensorDst::position == AscendC::TPosition::GM, + "The input parameters do not match. TensorSrc must be L0C, while TensorDst must be GM and RowMajor"); + + AscendC::FixpipeParamsC310 intriParams; + + // Fixpipe layout information + intriParams.nSize = tla::get<1>(dstTensor.shape()); + intriParams.mSize = tla::get<0>(dstTensor.shape()); + intriParams.srcStride = tla::get<1, 1>(srcTensor.stride()) / tla::get<0, 0>(srcTensor.stride()); + intriParams.dstStride = tla::get<0>(dstTensor.stride()); + + // Fixpipe auxiliary arguments + intriParams.quantPre = quantPre; + intriParams.deqScalar = static_cast(*reinterpret_cast(¶ms.scale)); + intriParams.reluEn = reluEn; + intriParams.unitFlag = unitFlag; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + + // Call AscendC Fixpipe + AscendC::Fixpipe( + dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], intriParams); + } +}; + +template +struct CopyL0CToGmTla< + Catlass::Arch::Ascend950, + TensorSrc_, + tla::Tensor, LayoutDst_, CoordDst_, AscendC::TPosition::GM>, + ScaleGranularity::PER_CHANNEL, + ReluEnable_, + std::enable_if_t::value>> { + using ArchTag = Catlass::Arch::Ascend950; + using ElementDst = ElementDst_; + using ElementSrc = typename TensorSrc_::Element; + static constexpr auto quantPre = + CopyL0CToDstQuantMode::VALUE; + static constexpr auto reluEn = ReluEnable_; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, TensorQuant const &quantTensor, uint8_t unitFlag = 0) + { + static_assert( + tla::detail::isRowMajor::value && TensorSrc::position == AscendC::TPosition::CO1 && TensorDst::position == AscendC::TPosition::GM, + "The input parameters do not match. TensorSrc must be L0C, while TensorDst must be GM and RowMajor"); + + AscendC::FixpipeParamsC310 intriParams; + + // Fixpipe layout information + intriParams.nSize = tla::get<1>(dstTensor.shape()); + intriParams.mSize = tla::get<0>(dstTensor.shape()); + intriParams.srcStride = tla::get<1, 1>(srcTensor.stride()) / tla::get<0, 0>(srcTensor.stride()); + intriParams.dstStride = tla::get<0>(dstTensor.stride()); + + // Fixpipe auxiliary arguments + intriParams.quantPre = quantPre; + intriParams.reluEn = reluEn; + intriParams.unitFlag = unitFlag; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + auto quantOffset = quantTensor.layout()(quantTensor.coord()); + + // Call AscendC Fixpipe + AscendC::Fixpipe( + dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], quantTensor.data()[quantOffset], intriParams); + } +}; + +////////////////////////////////////CopyL0CToGm(No-TLA, Ascend950)//////////////////////////////////////////////// +template < + class ElementAccumulator_, + class ElementDst_, + bool ReluEnable_> +struct CopyL0CToGm, + ScaleGranularity::NO_QUANT, + ReluEnable_> { + using ArchTag = Catlass::Arch::Ascend950; + using ElementDst = ElementDst_; + using ElementSrc = ElementAccumulator_; + using LayoutSrc = Catlass::layout::zN; + using LayoutDst = Catlass::layout::RowMajor; + static constexpr auto quantPre = CopyL0CToDstQuantMode::VALUE; + static constexpr auto reluEn = ReluEnable_; + + struct Params {}; + Params params; + + CATLASS_DEVICE + CopyL0CToGm() = default; + + CATLASS_DEVICE + CopyL0CToGm(Params const ¶ms_) + : params(params_) {}; + + CATLASS_DEVICE + void operator()(AscendC::GlobalTensor const &dst, AscendC::LocalTensor const &src, + LayoutDst const &dstLayout, LayoutSrc const &srcLayout, uint8_t unitFlag = 0) + { + AscendC::DataCopyCO12DstParams intriParams; + + // Fixpipe layout information + intriParams.nSize = dstLayout.shape(1); + intriParams.mSize = dstLayout.shape(0); + intriParams.srcStride = srcLayout.stride(3) / srcLayout.stride(0); + intriParams.dstStride = dstLayout.stride(0); + intriParams.quantPre = quantPre; + intriParams.nz2ndEn = true; + intriParams.reluPre = reluEn; + intriParams.unitFlag = unitFlag; + + AscendC::SetFixpipeNz2ndFlag(1, 1, 1); + AscendC::DataCopy(dst, src, intriParams); + } +}; + +template < + class ElementAccumulator_, + class ElementDst_, + bool ReluEnable_> +struct CopyL0CToGm, + ScaleGranularity::PER_TENSOR, + ReluEnable_> { + using ArchTag = Catlass::Arch::Ascend950; + using ElementDst = ElementDst_; + using ElementSrc = ElementAccumulator_; + using LayoutSrc = Catlass::layout::zN; + using LayoutDst = Catlass::layout::RowMajor; + static constexpr auto quantPre = CopyL0CToDstQuantMode::VALUE; + static constexpr auto reluEn = ReluEnable_; + + struct Params { + float scale = 1.0; + + CATLASS_HOST_DEVICE + Params() = default; + + CATLASS_HOST_DEVICE + Params(float scalar) + { + scale = scalar; + } + }; + Params params; + + CATLASS_DEVICE + CopyL0CToGm() = default; + + CATLASS_DEVICE + CopyL0CToGm(Params const ¶ms_) + : params(params_) {}; + + CATLASS_DEVICE + void operator()(AscendC::GlobalTensor const &dst, AscendC::LocalTensor const &src, + LayoutDst const &dstLayout, LayoutSrc const &srcLayout, uint8_t unitFlag = 0) + { + AscendC::FixpipeParamsC310 intriParams; + + // Fixpipe layout information + intriParams.nSize = dstLayout.shape(1); + intriParams.mSize = dstLayout.shape(0); + intriParams.srcStride = srcLayout.stride(3) / srcLayout.stride(0); + intriParams.dstStride = dstLayout.stride(0); + + // Fixpipe auxiliary arguments + intriParams.quantPre = quantPre; + intriParams.deqScalar = static_cast(*reinterpret_cast(¶ms.scale)); + intriParams.reluEn = reluEn; + intriParams.unitFlag = unitFlag; + + // Call AscendC Fixpipe + AscendC::Fixpipe(dst, src, intriParams); + } +}; + +template < + class ElementAccumulator_, + class ElementDst_, + bool ReluEnable_> +struct CopyL0CToGm, + ScaleGranularity::PER_CHANNEL, + ReluEnable_> { + using ArchTag = Catlass::Arch::Ascend950; + using ElementDst = ElementDst_; + using ElementSrc = ElementAccumulator_; + using LayoutSrc = Catlass::layout::zN; + using LayoutDst = Catlass::layout::RowMajor; + static constexpr auto quantPre = CopyL0CToDstQuantMode::VALUE; + static constexpr auto reluEn = ReluEnable_; + + struct Params {}; + Params params; + + CATLASS_DEVICE + CopyL0CToGm() = default; + + CATLASS_DEVICE + CopyL0CToGm(Params const ¶ms_) + : params(params_) {}; + + CATLASS_DEVICE + void operator()(AscendC::GlobalTensor const &dst, AscendC::LocalTensor const &src, AscendC::LocalTensor const &scale, + LayoutDst const &dstLayout, LayoutSrc const &srcLayout, uint8_t unitFlag = 0) + { + AscendC::FixpipeParamsC310 intriParams; + + // Fixpipe layout information + intriParams.nSize = dstLayout.shape(1); + intriParams.mSize = dstLayout.shape(0); + intriParams.srcStride = srcLayout.stride(3) / srcLayout.stride(0); + intriParams.dstStride = dstLayout.stride(0); + + // Fixpipe auxiliary arguments + intriParams.quantPre = quantPre; + intriParams.reluEn = reluEn; + intriParams.unitFlag = unitFlag; + + // Call AscendC Fixpipe + AscendC::Fixpipe(dst, src, scale, intriParams); + } +}; + +template < + class ElementAccumulator_, + class ElementDst_, + bool ReluEnable_> +struct CopyL0CToGm, + ScaleGranularity::NO_QUANT, + ReluEnable_> { + using ArchTag = Catlass::Arch::Ascend950; + using ElementDst = ElementDst_; + using ElementSrc = ElementAccumulator_; + using LayoutSrc = Catlass::layout::zN; + using LayoutDst = Catlass::layout::zN; + static constexpr auto quantPre = CopyL0CToDstQuantMode::VALUE; + static constexpr auto reluEn = ReluEnable_; + + struct Params {}; + Params params; + + CATLASS_DEVICE + CopyL0CToGm() = default; + + CATLASS_DEVICE + CopyL0CToGm(Params const ¶ms_) + : params(params_) {}; + + CATLASS_DEVICE + void operator()(AscendC::GlobalTensor const &dst, AscendC::LocalTensor const &src, + LayoutDst const &dstLayout, LayoutSrc const &srcLayout, uint8_t unitFlag = 0) + { + AscendC::DataCopyCO12DstParams intriParams; + + // Fixpipe layout information + intriParams.nSize = dstLayout.shape(2) * dstLayout.shape(3); + intriParams.mSize = dstLayout.shape(0) * dstLayout.shape(1); + intriParams.srcStride = srcLayout.stride(3) / srcLayout.shape(2); + intriParams.dstStride = dstLayout.stride(3) / (BYTE_PER_C0 / sizeof(ElementDst)); + intriParams.quantPre = quantPre; + intriParams.nz2ndEn = false; + intriParams.reluPre = reluEn; + intriParams.unitFlag = unitFlag; + + if constexpr (std::is_same_v && std::is_same_v) { + intriParams.channelSplit = true; + } + + AscendC::DataCopy(dst, src, intriParams); + } +}; + +template < + class ElementAccumulator_, + class ElementDst_, + bool ReluEnable_> +struct CopyL0CToGm, + ScaleGranularity::NO_QUANT, + ReluEnable_> { + using ArchTag = Catlass::Arch::Ascend950; + using ElementDst = ElementDst_; + using ElementSrc = ElementAccumulator_; + using LayoutSrc = Catlass::layout::zN; + using LayoutDst = Catlass::layout::NDC1HWC0; + static constexpr auto quantPre = CopyL0CToDstQuantMode::VALUE; + static constexpr auto reluEn = ReluEnable_; + + struct Params {}; + Params params; + + CATLASS_DEVICE + CopyL0CToGm() = default; + + CATLASS_DEVICE + CopyL0CToGm(Params const ¶ms_) + : params(params_) {}; + + CATLASS_DEVICE + void operator()(AscendC::GlobalTensor const &dst, AscendC::LocalTensor const &src, + LayoutDst const &dstLayout, LayoutSrc const &srcLayout, uint8_t unitFlag = 0) + { + AscendC::FixpipeParamsV220 intriParams; + + intriParams.nSize = srcLayout.orgShape(1); + intriParams.mSize = srcLayout.orgShape(0); + intriParams.srcStride = srcLayout.stride(3) / srcLayout.shape(2); + intriParams.dstStride = dstLayout.shape(1) * dstLayout.shape(2); + + if constexpr (AscendC::IsSameType::value && + AscendC::IsSameType::value) { + intriParams.isChannelSplit = true; + } + + intriParams.quantPre = quantPre; + intriParams.reluEn = false; + intriParams.unitFlag = unitFlag; + AscendC::Fixpipe(dst, src, intriParams); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace Catlass::Gemm::Tile + +#endif // CATLASS_KDA_ASCEND950_COPY_L0C_TO_GM_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/ascend950/kda_ascend950_copy_l1_to_bt.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/ascend950/kda_ascend950_copy_l1_to_bt.hpp new file mode 100644 index 0000000000..97d8004b48 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/ascend950/kda_ascend950_copy_l1_to_bt.hpp @@ -0,0 +1,107 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_ASCEND950_COPY_L1_TO_BT_HPP +#define CATLASS_KDA_ASCEND950_COPY_L1_TO_BT_HPP + +#include "catlass/arch/kda_arch.hpp" +#include "catlass/kda_catlass.hpp" +#include "catlass/gemm/tile/kda_tile_copy_tla.hpp" +#include "catlass/kda_numeric_size.hpp" +#include "catlass/layout/kda_layout.hpp" +#include "catlass/gemm/kda_gemm_type.hpp" +#include "tla/kda_tensor.hpp" + +namespace Catlass::Gemm::Tile { + +/// Partial specialization for CopyL1ToBT, Ascend950, VectorLayout in and VectorLayout out. +template +struct TileCopyTla< + Arch::Ascend950, + tla::Tensor, LayoutSrc, CoordSrc, AscendC::TPosition::A1>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::C2>, + std::enable_if_t::value && tla::detail::isVector::value>> { + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + TileCopyTla() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert( + tla::detail::isVector::value && tla::detail::isVector::value && TensorSrc::position == AscendC::TPosition::A1 && TensorDst::position == AscendC::TPosition::C2, + "The input parameters do not match. TensorSrc must be L1 and Vector, " + "while TensorDst must be BT and Vector"); + + AscendC::DataCopyParams intriParams; + intriParams.blockCount = 1; + intriParams.blockLen = CeilDiv(tla::get<0>(srcTensor.originShape()), ELE_NUM_PER_C0); + if (sizeof(ElementSrc) == 4) { + // the burst length should be even when B32 + intriParams.blockLen = RoundUp(intriParams.blockLen, 2); + } + intriParams.srcStride = 0; + intriParams.dstStride = 0; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + + AscendC::DataCopy(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], intriParams); + } +}; + +////////////////////////////////////TileCopy(No-TLA, Ascend950)//////////////////////////////////////////////// + +template < + class ArchTag, + class L1Type, + class L0Type = void> +struct CopyL1ToBT { + static_assert(DEPENDENT_FALSE, "Unsupported copy l1 to biasTable buffer, can not find the specialization."); +}; + +template +struct CopyL1ToBT, + Catlass::Gemm::GemmType> { + using LayoutDst = layout::VectorLayout; + using LayoutSrc = layout::VectorLayout; + + static constexpr uint32_t ELE_NUM_PER_C2 = BYTE_PER_C2 / sizeof(ElementSrc); + + CATLASS_DEVICE + CopyL1ToBT() {} + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor dstTensor, + AscendC::LocalTensor srcTensor, + LayoutDst layoutDst, LayoutSrc layoutSrc) + { + AscendC::DataCopyParams intriParams; + intriParams.blockCount = 1; + intriParams.blockLen = (layoutDst.shape(0) + ELE_NUM_PER_C2 - 1) / ELE_NUM_PER_C2; + if (sizeof(ElementSrc) == 4) { + // the burst length should be even when B32 + intriParams.blockLen = RoundUp(intriParams.blockLen, 2); + } + intriParams.srcStride = 0; + intriParams.dstStride = 0; + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace Catlass::Gemm::Tile + +#endif // CATLASS_KDA_ASCEND950_COPY_L1_TO_BT_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/ascend950/kda_ascend950_copy_l1_to_l0a.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/ascend950/kda_ascend950_copy_l1_to_l0a.hpp new file mode 100644 index 0000000000..bb2424a8cd --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/ascend950/kda_ascend950_copy_l1_to_l0a.hpp @@ -0,0 +1,467 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_ASCEND950_COPY_L1_TO_L0A_HPP +#define CATLASS_KDA_ASCEND950_COPY_L1_TO_L0A_HPP + +#include "catlass/arch/kda_arch.hpp" +#include "catlass/kda_catlass.hpp" +#include "catlass/kda_numeric_size.hpp" +#include "catlass/layout/kda_layout.hpp" +#include "catlass/gemm/kda_gemm_type.hpp" +#include "catlass/gemm/tile/kda_tile_copy_tla.hpp" +#include "tla/kda_tensor.hpp" + +namespace Catlass::Gemm::Tile { + +/// Partial specialization for CopyL1ToL0A, Ascend950, zN in and zN out. +template +struct TileCopyTla< + Arch::Ascend950, + tla::Tensor, LayoutSrc, CoordSrc, AscendC::TPosition::A1>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::A2>, + std::enable_if_t::value && tla::detail::iszN::value>> { + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + TileCopyTla() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert( + tla::detail::iszN::value && tla::detail::iszN::value && TensorSrc::position == AscendC::TPosition::A1 && TensorDst::position == AscendC::TPosition::A2, + "The input parameters do not match. TensorSrc must be L1 and zN, while TensorDst must be L0A and zN"); + + const uint32_t dstOuterShapeRow = CeilDiv(tla::get<0>(dstTensor.originShape()), tla::get<0, 0>(dstTensor.shape())); + const uint32_t dstOuterShapeCol = CeilDiv(tla::get<1>(dstTensor.originShape()), tla::get<1, 0>(dstTensor.shape())); + const uint32_t srcOuterStrideCol = tla::get<1, 1>(srcTensor.stride()); + const uint32_t dstOuterStrideCol = tla::get<1, 1>(dstTensor.stride()); + auto srcCoord = srcTensor.coord(); + + AscendC::LoadData2DParamsV2 loadDataParams; + loadDataParams.mStartPosition = CeilDiv(tla::get<0>(srcCoord)); + loadDataParams.kStartPosition = CeilDiv(tla::get<1>(srcCoord)); + loadDataParams.mStep = dstOuterShapeRow; + loadDataParams.kStep = dstOuterShapeCol; + loadDataParams.srcStride = CeilDiv(srcOuterStrideCol); + loadDataParams.dstStride = CeilDiv(dstOuterStrideCol); + loadDataParams.ifTranspose = false; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + AscendC::LoadData(dstTensor.data()[dstOffset], srcTensor.data(), loadDataParams); + } + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, uint32_t l0Batch) + { + static_assert( + tla::detail::iszN::value && tla::detail::iszN::value && TensorSrc::position == AscendC::TPosition::A1 && TensorDst::position == AscendC::TPosition::A2, + "The input parameters do not match. TensorSrc must be L1 and zN, while TensorDst must be L0A and zN"); + + const uint32_t dstOuterShapeRow = CeilDiv(tla::get<0>(dstTensor.originShape()), tla::get<0, 0>(dstTensor.shape())); + const uint32_t dstOuterShapeCol = CeilDiv(tla::get<1>(dstTensor.originShape()), tla::get<1, 0>(dstTensor.shape())); + const uint32_t srcOuterStrideCol = tla::get<1, 1>(srcTensor.stride()); + const uint32_t dstOuterStrideCol = tla::get<1, 1>(dstTensor.stride()); + + AscendC::LoadData2DParamsV2 loadDataParams; + loadDataParams.mStartPosition = 0; + loadDataParams.kStartPosition = 0; + loadDataParams.mStep = dstOuterShapeRow; + loadDataParams.kStep = dstOuterShapeCol * l0Batch; + loadDataParams.srcStride = CeilDiv(srcOuterStrideCol); + loadDataParams.dstStride = CeilDiv(dstOuterStrideCol); + loadDataParams.ifTranspose = false; + + AscendC::LoadData(dstTensor.data(), srcTensor.data(), loadDataParams); + } +}; + +/// Partial specialization for CopyL1ToL0A, Ascend950, not B8, nZ in and zN out. (Transpose A) +template +struct TileCopyTla< + Arch::Ascend950, + tla::Tensor, LayoutSrc, CoordSrc, AscendC::TPosition::A1>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::A2>, + std::enable_if_t< + !AscendC::Std::is_one_of_v && + !AscendC::Std::is_one_of_v && + tla::detail::isnZ::value && tla::detail::iszN::value>> { + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + TileCopyTla() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert( + !AscendC::Std::is_one_of_v && + !AscendC::Std::is_one_of_v && + tla::detail::isnZ::value && tla::detail::iszN::value && TensorSrc::position == AscendC::TPosition::A1 && TensorDst::position == AscendC::TPosition::A2, + "The input parameters do not match. TensorSrc must be L1 and nZ, while TensorDst must be L0A and zN"); + + const uint32_t L0M = tla::get<0>(dstTensor.originShape()); + const uint32_t L0K = tla::get<1>(dstTensor.originShape()); + const uint32_t srcOuterStrideRow = tla::get<0, 1>(srcTensor.stride()); + const uint32_t dstOuterStrideCol = tla::get<1, 1>(dstTensor.stride()); + auto srcCoord = srcTensor.coord(); + + AscendC::LoadData2DParamsV2 loadDataParams; + loadDataParams.mStartPosition = CeilDiv(tla::get<1>(srcCoord)); + loadDataParams.kStartPosition = CeilDiv(tla::get<0>(srcCoord)); + loadDataParams.mStep = CeilDiv(L0K); + loadDataParams.kStep = CeilDiv(L0M); + if constexpr (AscendC::Std::is_one_of_v) { + loadDataParams.kStep = RoundUp<2>(loadDataParams.kStep); // for b32 data types, ensure kStep is even + } + loadDataParams.srcStride = CeilDiv(srcOuterStrideRow); + loadDataParams.dstStride = CeilDiv(dstOuterStrideCol); + loadDataParams.ifTranspose = true; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + AscendC::LoadData(dstTensor.data()[dstOffset], srcTensor.data(), loadDataParams); + } + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, uint32_t l0Batch) + { + static_assert( + !AscendC::Std::is_one_of_v && + !AscendC::Std::is_one_of_v && + tla::detail::isnZ::value && tla::detail::iszN::value && TensorSrc::position == AscendC::TPosition::A1 && TensorDst::position == AscendC::TPosition::A2, + "The input parameters do not match. TensorSrc must be L1 and nZ, while TensorDst must be L0A and zN"); + + const uint32_t L1M = tla::get<0, 0>(srcTensor.shape()) * tla::get<0, 1>(srcTensor.shape()); + const uint32_t L1K = tla::get<1, 0>(srcTensor.shape()) * tla::get<1, 1>(srcTensor.shape()); + const uint32_t L0M = tla::get<0, 0>(dstTensor.shape()) * tla::get<0, 1>(dstTensor.shape()); + const uint32_t L0K = tla::get<1, 0>(dstTensor.shape()) * tla::get<1, 1>(dstTensor.shape()); + const uint32_t L0MOrigin = tla::get<0>(dstTensor.originShape()); + const uint32_t L0KOrigin = tla::get<1>(dstTensor.originShape()); + const uint32_t srcOuterStrideRow = tla::get<0, 1>(srcTensor.stride()); + const uint32_t dstOuterStrideCol = tla::get<1, 1>(dstTensor.stride()); + + AscendC::LoadData2DParamsV2 loadDataParams; + loadDataParams.mStartPosition = 0; + loadDataParams.kStartPosition = 0; + loadDataParams.mStep = CeilDiv(L0KOrigin); + loadDataParams.kStep = CeilDiv(L0MOrigin); + if constexpr (AscendC::Std::is_one_of_v) { + loadDataParams.kStep = RoundUp<2>(loadDataParams.kStep); // for b32 data types, ensure kStep is even + } + loadDataParams.srcStride = CeilDiv(srcOuterStrideRow); + loadDataParams.dstStride = CeilDiv(dstOuterStrideCol); + loadDataParams.ifTranspose = true; + + for (uint32_t l0BatchIdx = 0; l0BatchIdx < l0Batch; l0BatchIdx++) { + AscendC::LoadData( + dstTensor.data()[l0BatchIdx * L0M * L0K], srcTensor.data()[l0BatchIdx * L1M * L1K], loadDataParams); + } + } +}; + +/// Partial specialization for CopyL1ToL0A, Ascend950, B8, nZ in and zN out. (Transpose A) +template +struct TileCopyTla< + Arch::Ascend950, + tla::Tensor, LayoutSrc, CoordSrc, AscendC::TPosition::A1>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::A2>, + std::enable_if_t< + AscendC::Std::is_one_of_v && + AscendC::Std::is_one_of_v && + tla::detail::isnZ::value && tla::detail::iszN::value>> { + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + TileCopyTla() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert( + AscendC::Std::is_one_of_v && + AscendC::Std::is_one_of_v && + tla::detail::isnZ::value && tla::detail::iszN::value && TensorSrc::position == AscendC::TPosition::A1 && TensorDst::position == AscendC::TPosition::A2, + "The input parameters do not match. TensorSrc must be L1 and nZ, while TensorDst must be L0A and zN"); + + const uint32_t L0MPadded = tla::get<0, 0>(dstTensor.shape()) * tla::get<0, 1>(dstTensor.shape()); + const uint32_t L0MOrigin = tla::get<0>(dstTensor.originShape()); + const uint32_t L0KOrigin = tla::get<1>(dstTensor.originShape()); + const uint32_t srcOuterStrideRow = tla::get<0, 1>(srcTensor.stride()); + const uint32_t dstOuterStrideCol = tla::get<1, 1>(dstTensor.stride()); + auto srcCoord = srcTensor.coord(); + + AscendC::LoadData2DParamsV2 loadDataParams; + if (RoundUp(L0MOrigin) % ELE_NUM_PER_C0 == 0) { + loadDataParams.mStartPosition = CeilDiv(tla::get<1>(srcCoord)); + loadDataParams.kStartPosition = CeilDiv(tla::get<0>(srcCoord)); + loadDataParams.mStep = RoundUp<2>(CeilDiv(L0KOrigin)); + loadDataParams.kStep = CeilDiv(L0MOrigin); + loadDataParams.srcStride = CeilDiv(srcOuterStrideRow); + loadDataParams.dstStride = CeilDiv(dstOuterStrideCol); + loadDataParams.ifTranspose = true; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + AscendC::LoadData(dstTensor.data()[dstOffset], srcTensor.data(), loadDataParams); + } else { + for (uint32_t kIdx = 0; kIdx < CeilDiv(L0KOrigin); kIdx++) { + loadDataParams.mStartPosition = CeilDiv(tla::get<1>(srcCoord)) + kIdx * 2; + loadDataParams.kStartPosition = CeilDiv(tla::get<0>(srcCoord)); + loadDataParams.mStep = 2; + loadDataParams.kStep = CeilDiv(L0MOrigin); + loadDataParams.srcStride = CeilDiv(srcOuterStrideRow); + loadDataParams.dstStride = CeilDiv(dstOuterStrideCol); + loadDataParams.ifTranspose = true; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + AscendC::LoadData( + dstTensor.data()[dstOffset + kIdx * L0MPadded * ELE_NUM_PER_C0], srcTensor.data(), loadDataParams); + } + } + } + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, uint32_t l0Batch) + { + static_assert( + AscendC::Std::is_one_of_v && + AscendC::Std::is_one_of_v && + tla::detail::isnZ::value && tla::detail::iszN::value && TensorSrc::position == AscendC::TPosition::A1 && TensorDst::position == AscendC::TPosition::A2, + "The input parameters do not match. TensorSrc must be L1 and nZ, while TensorDst must be L0A and zN"); + + const uint32_t L1M = tla::get<0, 0>(srcTensor.shape()) * tla::get<0, 1>(srcTensor.shape()); + const uint32_t L1K = tla::get<1, 0>(srcTensor.shape()) * tla::get<1, 1>(srcTensor.shape()); + const uint32_t L0M = tla::get<0, 0>(dstTensor.shape()) * tla::get<0, 1>(dstTensor.shape()); + const uint32_t L0K = tla::get<1, 0>(dstTensor.shape()) * tla::get<1, 1>(dstTensor.shape()); + const uint32_t L0MPadded = L0M; + const uint32_t L0KPadded = L0K; + const uint32_t L0MOrigin = tla::get<0>(dstTensor.originShape()); + const uint32_t L0KOrigin = tla::get<1>(dstTensor.originShape()); + const uint32_t srcOuterStrideRow = tla::get<0, 1>(srcTensor.stride()); + const uint32_t dstOuterStrideCol = tla::get<1, 1>(dstTensor.stride()); + + AscendC::LoadData2DParamsV2 loadDataParams; + if (RoundUp(L0MOrigin) % ELE_NUM_PER_C0 == 0) { + loadDataParams.mStartPosition = 0; + loadDataParams.kStartPosition = 0; + loadDataParams.mStep = RoundUp<2>(CeilDiv(L0KOrigin)); + loadDataParams.kStep = CeilDiv(L0MOrigin); + loadDataParams.srcStride = CeilDiv(srcOuterStrideRow); + loadDataParams.dstStride = CeilDiv(dstOuterStrideCol); + loadDataParams.ifTranspose = true; + + for (uint32_t l0BatchIdx = 0; l0BatchIdx < l0Batch; l0BatchIdx++) { + AscendC::LoadData( + dstTensor.data()[l0BatchIdx * L0MPadded * L0KPadded], srcTensor.data()[l0BatchIdx * L1M * L1K], loadDataParams); + } + } else { + loadDataParams.mStartPosition = 0; + loadDataParams.kStartPosition = 0; + loadDataParams.mStep = 2; + loadDataParams.kStep = CeilDiv(L0MOrigin); + loadDataParams.srcStride = CeilDiv(srcOuterStrideRow); + loadDataParams.dstStride = CeilDiv(dstOuterStrideCol); + loadDataParams.ifTranspose = true; + for (uint32_t l0BatchIdx = 0; l0BatchIdx < l0Batch; l0BatchIdx++) { + for (uint32_t kIdx = 0; kIdx < CeilDiv(L0KOrigin); kIdx++) { + AscendC::LoadData( + dstTensor.data()[l0BatchIdx * L0MPadded * L0KPadded + kIdx * L0MPadded * ELE_NUM_PER_C0], + srcTensor.data()[l0BatchIdx * L1M * L1K + kIdx * ELE_NUM_PER_FRACTAL * 2], loadDataParams); + } + } + } + } +}; + +template +struct TileCopyTla< + Arch::Ascend950, + tla::Tensor, LayoutSrc, CoordSrc, AscendC::TPosition::A1>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::A2>, + std::enable_if_t::value>> { + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + CATLASS_DEVICE + TileCopyTla() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + uint16_t aL1M = tla::get<0, 0>(srcTensor.stride()); + uint16_t madM = tla::get<1, 1>(dstTensor.stride()); + uint16_t madK = CeilDiv(tla::get<1>(dstTensor.originShape()), tla::get<1, 0>(dstTensor.shape())); + + AscendC::LoadData2DParamsV2 loadDataParams; + loadDataParams.mStartPosition = 0; + loadDataParams.kStartPosition = 0; + loadDataParams.mStep = CeilDiv(madM); + loadDataParams.kStep = CeilDiv(madK); + loadDataParams.srcStride = CeilDiv(aL1M); + loadDataParams.dstStride = CeilDiv(madM); + + loadDataParams.ifTranspose = false; + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + AscendC::LoadData(dstTensor.data(), srcTensor.data()[srcOffset], loadDataParams); + } +}; + +////////////////////////////////////CopyL1ToL0A(No-TLA, Ascend950)//////////////////////////////////////////////// +template < + class ArchTag, + class L1Type, + class L0Type = void, + class Enable = void> +struct CopyL1ToL0A { + static_assert(DEPENDENT_FALSE, "Unsupported copy l1 to l0, can not find the specialization."); +}; + +/// Partial specialization for CopyL1ToL0A, AtlasA5, zN in and zN out. +template +struct CopyL1ToL0A> { + using LayoutDst = layout::zN; + using LayoutSrc = layout::zN; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyL1ToL0A() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::LoadData2DParamsV2 loadDataParams; + loadDataParams.mStartPosition = 0; + loadDataParams.kStartPosition = 0; + loadDataParams.mStep = layoutDst.shape(1); + loadDataParams.kStep = layoutDst.shape(3); + loadDataParams.srcStride = CeilDiv(layoutSrc.stride(3)); + loadDataParams.dstStride = CeilDiv(layoutDst.stride(3)); + loadDataParams.ifTranspose = false; + + AscendC::LoadData(dstTensor, srcTensor, loadDataParams); + } +}; + +/// Partial specialization for CopyL1ToL0A, AtlasA5, not B8 or B4, nZ in and zN out. (Transpose A) +template +struct CopyL1ToL0A, void, + std::enable_if_t< + !AscendC::Std::is_one_of_v>> { + using LayoutDst = layout::zN; + using LayoutSrc = layout::nZ; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyL1ToL0A() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + const uint32_t L0M = layoutDst.shape(0) * layoutDst.shape(1); + const uint32_t L0K = layoutDst.shape(2) * layoutDst.shape(3); + const uint32_t srcOuterStrideRow = layoutSrc.stride(1); + const uint32_t dstOuterStrideCol = layoutDst.stride(3); + + AscendC::LoadData2DParamsV2 loadDataParams; + loadDataParams.mStartPosition = 0; + loadDataParams.kStartPosition = 0; + loadDataParams.mStep = CeilDiv(L0K); + loadDataParams.kStep = CeilDiv(L0M); + loadDataParams.srcStride = CeilDiv(srcOuterStrideRow); + loadDataParams.dstStride = CeilDiv(dstOuterStrideCol); + loadDataParams.ifTranspose = true; + + AscendC::LoadData(dstTensor, srcTensor, loadDataParams); + } +}; + +/// Partial specialization for CopyL1ToL0A, AtlasA5, B8 or B4, nZ in and zN out. (Transpose A) +template +struct CopyL1ToL0A, void, + std::enable_if_t< + AscendC::Std::is_one_of_v>> { + using LayoutDst = layout::zN; + using LayoutSrc = layout::nZ; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyL1ToL0A() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + const uint32_t L0M = layoutDst.shape(0) * layoutDst.shape(1); + const uint32_t L0K = layoutDst.shape(2) * layoutDst.shape(3); + const uint32_t srcOuterStrideRow = layoutSrc.stride(1); + const uint32_t dstOuterStrideCol = layoutDst.stride(3); + + AscendC::LoadData2DParamsV2 loadDataParams; + if (L0M % ELE_NUM_PER_C0 == 0) { + loadDataParams.mStartPosition = 0; + loadDataParams.kStartPosition = 0; + loadDataParams.mStep = CeilDiv(L0K); + loadDataParams.kStep = CeilDiv(L0M); + loadDataParams.srcStride = CeilDiv(srcOuterStrideRow); + loadDataParams.dstStride = CeilDiv(dstOuterStrideCol); + loadDataParams.ifTranspose = true; + + AscendC::LoadData(dstTensor, srcTensor, loadDataParams); + } else { + for (uint32_t kIdx = 0; kIdx < L0K / ELE_NUM_PER_C0; kIdx++) { + loadDataParams.mStartPosition = kIdx * 2; + loadDataParams.kStartPosition = 0; + loadDataParams.mStep = 2; + loadDataParams.kStep = CeilDiv(L0M); + loadDataParams.srcStride = CeilDiv(srcOuterStrideRow); + loadDataParams.dstStride = CeilDiv(dstOuterStrideCol); + loadDataParams.ifTranspose = true; + + AscendC::LoadData( + dstTensor[kIdx * L0M * ELE_NUM_PER_C0], srcTensor, loadDataParams); + } + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace Catlass::Gemm::Tile + +#endif // CATLASS_KDA_ASCEND950_COPY_L1_TO_L0A_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/ascend950/kda_ascend950_copy_l1_to_l0b.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/ascend950/kda_ascend950_copy_l1_to_l0b.hpp new file mode 100644 index 0000000000..d4a59a1439 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/ascend950/kda_ascend950_copy_l1_to_l0b.hpp @@ -0,0 +1,437 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_ASCEND950_COPY_L1_TO_L0B_HPP +#define CATLASS_KDA_ASCEND950_COPY_L1_TO_L0B_HPP + +#include "catlass/arch/kda_arch.hpp" +#include "catlass/kda_catlass.hpp" +#include "catlass/kda_numeric_size.hpp" +#include "catlass/layout/kda_layout.hpp" +#include "catlass/gemm/kda_gemm_type.hpp" +#include "catlass/gemm/tile/kda_tile_copy_tla.hpp" +#include "tla/kda_tensor.hpp" + +namespace Catlass::Gemm::Tile { + +/// Partial specialization for CopyL1ToL0B, Ascend950, not B8, zN in and nZ out. +template +struct TileCopyTla< + Arch::Ascend950, + tla::Tensor, LayoutSrc, CoordSrc, AscendC::TPosition::A1>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::B2>, + std::enable_if_t< + !AscendC::Std::is_one_of_v && + !AscendC::Std::is_one_of_v && + tla::detail::iszN::value && tla::detail::isnZ::value>> { + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + TileCopyTla() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert( + !AscendC::Std::is_one_of_v && + !AscendC::Std::is_one_of_v && + tla::detail::iszN::value && tla::detail::isnZ::value && TensorSrc::position == AscendC::TPosition::A1 && TensorDst::position == AscendC::TPosition::B2, + "The input parameters do not match. TensorSrc must be L1 and zN, while TensorDst must be L0B and nZ"); + + const uint32_t L0KOrigin = tla::get<0>(dstTensor.originShape()); + const uint32_t L0NOrigin = tla::get<1>(dstTensor.originShape()); + const uint32_t srcOuterStrideCol = tla::get<1, 1>(srcTensor.stride()); + const uint32_t dstOuterStrideRow = tla::get<0, 1>(dstTensor.stride()); + auto srcCoord = srcTensor.coord(); + + AscendC::LoadData2DParamsV2 loadDataParams; + loadDataParams.mStartPosition = CeilDiv(tla::get<0>(srcCoord)); + loadDataParams.kStartPosition = CeilDiv(tla::get<1>(srcCoord)); + loadDataParams.mStep = CeilDiv(L0KOrigin); + loadDataParams.kStep = CeilDiv(L0NOrigin); + if constexpr (AscendC::Std::is_one_of_v) { + loadDataParams.kStep = RoundUp<2>(loadDataParams.kStep); // for b32 data types, ensure kStep is even + } + loadDataParams.srcStride = CeilDiv(srcOuterStrideCol); + loadDataParams.dstStride = CeilDiv(dstOuterStrideRow); + loadDataParams.ifTranspose = true; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + AscendC::LoadData(dstTensor.data()[dstOffset], srcTensor.data(), loadDataParams); + } + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, uint32_t l0Batch) + { + static_assert( + !AscendC::Std::is_one_of_v && + !AscendC::Std::is_one_of_v && + tla::detail::iszN::value && tla::detail::isnZ::value && TensorSrc::position == AscendC::TPosition::A1 && TensorDst::position == AscendC::TPosition::B2, + "The input parameters do not match. TensorSrc must be L1 and zN, while TensorDst must be L0B and nZ"); + + const uint32_t L1K = tla::get<0, 0>(srcTensor.shape()) * tla::get<0, 1>(srcTensor.shape()); + const uint32_t L1N = tla::get<1, 0>(srcTensor.shape()) * tla::get<1, 1>(srcTensor.shape()); + const uint32_t L0K = tla::get<0, 0>(dstTensor.shape()) * tla::get<0, 1>(dstTensor.shape()); + const uint32_t L0N = tla::get<1, 0>(dstTensor.shape()) * tla::get<1, 1>(dstTensor.shape()); + const uint32_t L0KOrigin = tla::get<0>(dstTensor.originShape()); + const uint32_t L0NOrigin = tla::get<1>(dstTensor.originShape()); + const uint32_t srcOuterStrideCol = tla::get<1, 1>(srcTensor.stride()); + const uint32_t dstOuterStrideRow = tla::get<0, 1>(dstTensor.stride()); + + AscendC::LoadData2DParamsV2 loadDataParams; + loadDataParams.mStartPosition = 0; + loadDataParams.kStartPosition = 0; + loadDataParams.mStep = CeilDiv(L0KOrigin); + loadDataParams.kStep = CeilDiv(L0NOrigin); + if constexpr (AscendC::Std::is_one_of_v) { + loadDataParams.kStep = RoundUp<2>(loadDataParams.kStep); // for b32 data types, ensure kStep is even + } + loadDataParams.srcStride = CeilDiv(srcOuterStrideCol); + loadDataParams.dstStride = CeilDiv(dstOuterStrideRow); + loadDataParams.ifTranspose = true; + + for (uint32_t l0BatchIdx = 0; l0BatchIdx < l0Batch; l0BatchIdx++) { + AscendC::LoadData( + dstTensor.data()[l0BatchIdx * L0N * L0K], srcTensor.data()[l0BatchIdx * L1N * L1K], loadDataParams); + } + } +}; + +/// Partial specialization for CopyL1ToL0B, Ascend950, B8, zN in and nZ out. +template +struct TileCopyTla< + Arch::Ascend950, + tla::Tensor, LayoutSrc, CoordSrc, AscendC::TPosition::A1>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::B2>, + std::enable_if_t< + AscendC::Std::is_one_of_v && + AscendC::Std::is_one_of_v && + tla::detail::iszN::value && tla::detail::isnZ::value>> { + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + TileCopyTla() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert( + AscendC::Std::is_one_of_v && + AscendC::Std::is_one_of_v && + tla::detail::iszN::value && tla::detail::isnZ::value && TensorSrc::position == AscendC::TPosition::A1 && TensorDst::position == AscendC::TPosition::B2, + "The input parameters do not match. TensorSrc must be L1 and zN, while TensorDst must be L0B and nZ"); + + const uint32_t L0NPadded = tla::get<1, 0>(dstTensor.shape()) * tla::get<1, 1>(dstTensor.shape()); + const uint32_t L0KOrigin = tla::get<0>(dstTensor.originShape()); + const uint32_t L0NOrigin = tla::get<1>(dstTensor.originShape()); + const uint32_t srcOuterStrideCol = tla::get<1, 1>(srcTensor.stride()); + const uint32_t dstOuterStrideRow = tla::get<0, 1>(dstTensor.stride()); + auto srcCoord = srcTensor.coord(); + + AscendC::LoadData2DParamsV2 loadDataParams; + if (RoundUp(L0NOrigin) % ELE_NUM_PER_C0 == 0) { + loadDataParams.mStartPosition = CeilDiv(tla::get<0>(srcCoord)); + loadDataParams.kStartPosition = CeilDiv(tla::get<1>(srcCoord)); + loadDataParams.mStep = RoundUp<2>(CeilDiv(L0KOrigin)); + loadDataParams.kStep = CeilDiv(L0NOrigin); + loadDataParams.srcStride = CeilDiv(srcOuterStrideCol); + loadDataParams.dstStride = CeilDiv(dstOuterStrideRow); + loadDataParams.ifTranspose = true; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + AscendC::LoadData(dstTensor.data()[dstOffset], srcTensor.data(), loadDataParams); + } else { + for (uint32_t kIdx = 0; kIdx < CeilDiv(L0KOrigin); kIdx++) { + loadDataParams.mStartPosition = CeilDiv(tla::get<0>(srcCoord)) + kIdx * 2; + loadDataParams.kStartPosition = CeilDiv(tla::get<1>(srcCoord)); + loadDataParams.mStep = 2; + loadDataParams.kStep = CeilDiv(L0NOrigin); + loadDataParams.srcStride = CeilDiv(srcOuterStrideCol); + loadDataParams.dstStride = CeilDiv(dstOuterStrideRow); + loadDataParams.ifTranspose = true; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + AscendC::LoadData( + dstTensor.data()[dstOffset + kIdx * L0NPadded * ELE_NUM_PER_C0], srcTensor.data(), loadDataParams); + } + } + } + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, uint32_t l0Batch) + { + static_assert( + AscendC::Std::is_one_of_v && + AscendC::Std::is_one_of_v && + tla::detail::iszN::value && tla::detail::isnZ::value && TensorSrc::position == AscendC::TPosition::A1 && TensorDst::position == AscendC::TPosition::B2, + "The input parameters do not match. TensorSrc must be L1 and zN, while TensorDst must be L0B and nZ"); + + const uint32_t L1K = tla::get<0, 0>(srcTensor.shape()) * tla::get<0, 1>(srcTensor.shape()); + const uint32_t L1N = tla::get<1, 0>(srcTensor.shape()) * tla::get<1, 1>(srcTensor.shape()); + const uint32_t L0K = tla::get<0, 0>(dstTensor.shape()) * tla::get<0, 1>(dstTensor.shape()); + const uint32_t L0N = tla::get<1, 0>(dstTensor.shape()) * tla::get<1, 1>(dstTensor.shape()); + const uint32_t L0KOrigin = tla::get<0>(dstTensor.originShape()); + const uint32_t L0NOrigin = tla::get<1>(dstTensor.originShape()); + const uint32_t srcOuterStrideCol = tla::get<1, 1>(srcTensor.stride()); + const uint32_t dstOuterStrideRow = tla::get<0, 1>(dstTensor.stride()); + + AscendC::LoadData2DParamsV2 loadDataParams; + if (RoundUp(L0NOrigin) % ELE_NUM_PER_C0 == 0) { + loadDataParams.mStartPosition = 0; + loadDataParams.kStartPosition = 0; + loadDataParams.mStep = RoundUp<2>(CeilDiv(L0KOrigin)); + loadDataParams.kStep = CeilDiv(L0NOrigin); + loadDataParams.srcStride = CeilDiv(srcOuterStrideCol); + loadDataParams.dstStride = CeilDiv(dstOuterStrideRow); + loadDataParams.ifTranspose = true; + + for (uint32_t l0BatchIdx = 0; l0BatchIdx < l0Batch; l0BatchIdx++) { + AscendC::LoadData( + dstTensor.data()[l0BatchIdx * L0N * L0K], srcTensor.data()[l0BatchIdx * L1N * L1K], loadDataParams); + } + } else { + loadDataParams.mStartPosition = 0; + loadDataParams.kStartPosition = 0; + loadDataParams.mStep = 2; + loadDataParams.kStep = CeilDiv(L0NOrigin); + loadDataParams.srcStride = CeilDiv(srcOuterStrideCol); + loadDataParams.dstStride = CeilDiv(dstOuterStrideRow); + loadDataParams.ifTranspose = true; + for (uint32_t l0BatchIdx = 0; l0BatchIdx < l0Batch; l0BatchIdx++) { + for (uint32_t kIdx = 0; kIdx < CeilDiv(L0KOrigin); kIdx++) { + AscendC::LoadData( + dstTensor.data()[l0BatchIdx * L0N * L0K + kIdx * L0N * ELE_NUM_PER_C0], + srcTensor.data()[l0BatchIdx * L1N * L1K + kIdx * ELE_NUM_PER_FRACTAL * 2], loadDataParams); + } + } + } + } +}; + +/// Partial specialization for CopyL1ToL0B, Ascend950, nZ in and nZ out. (Transpose B) +template +struct TileCopyTla< + Arch::Ascend950, + tla::Tensor, LayoutSrc, CoordSrc, AscendC::TPosition::A1>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::B2>, + std::enable_if_t::value && tla::detail::isnZ::value>> { + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + TileCopyTla() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert( + tla::detail::isnZ::value && tla::detail::isnZ::value && TensorSrc::position == AscendC::TPosition::A1 && TensorDst::position == AscendC::TPosition::B2, + "The input parameters do not match. TensorSrc must be L1 and nZ, while TensorDst must be L0B and nZ"); + + const uint32_t dstOuterShapeRow = CeilDiv(tla::get<0>(dstTensor.originShape()), tla::get<0, 0>(dstTensor.shape())); + const uint32_t dstOuterShapeCol = CeilDiv(tla::get<1>(dstTensor.originShape()), tla::get<1, 0>(dstTensor.shape())); + const uint32_t srcOuterStrideRow = tla::get<0, 1>(srcTensor.stride()); + const uint32_t dstOuterStrideRow = tla::get<0, 1>(dstTensor.stride()); + auto srcCoord = srcTensor.coord(); + + AscendC::LoadData2DParamsV2 loadDataParams; + loadDataParams.mStartPosition = CeilDiv(tla::get<1>(srcCoord)); + loadDataParams.kStartPosition = CeilDiv(tla::get<0>(srcCoord)); + loadDataParams.mStep = dstOuterShapeCol; + loadDataParams.kStep = dstOuterShapeRow; + loadDataParams.srcStride = CeilDiv(srcOuterStrideRow); + loadDataParams.dstStride = CeilDiv(dstOuterStrideRow); + loadDataParams.ifTranspose = false; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + AscendC::LoadData(dstTensor.data()[dstOffset], srcTensor.data(), loadDataParams); + } + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, uint32_t l0Batch) + { + static_assert( + tla::detail::isnZ::value && tla::detail::isnZ::value && TensorSrc::position == AscendC::TPosition::A1 && TensorDst::position == AscendC::TPosition::B2, + "The input parameters do not match. TensorSrc must be L1 and nZ, while TensorDst must be L0B and nZ"); + + const uint32_t dstOuterShapeRow = CeilDiv(tla::get<0>(dstTensor.originShape()), tla::get<0, 0>(dstTensor.shape())); + const uint32_t dstOuterShapeCol = CeilDiv(tla::get<1>(dstTensor.originShape()), tla::get<1, 0>(dstTensor.shape())); + const uint32_t srcOuterStrideRow = tla::get<0, 1>(srcTensor.stride()); + const uint32_t dstOuterStrideRow = tla::get<0, 1>(dstTensor.stride()); + + AscendC::LoadData2DParamsV2 loadDataParams; + loadDataParams.mStartPosition = 0; + loadDataParams.kStartPosition = 0; + loadDataParams.mStep = dstOuterShapeCol; + loadDataParams.kStep = dstOuterShapeRow * l0Batch; + loadDataParams.srcStride = CeilDiv(srcOuterStrideRow); + loadDataParams.dstStride = CeilDiv(dstOuterStrideRow); + loadDataParams.ifTranspose = false; + + AscendC::LoadData(dstTensor.data(), srcTensor.data(), loadDataParams); + } +}; + +////////////////////////////////////CopyL1ToL0B(No-TLA, Ascend950)//////////////////////////////////////////////// +template < + class ArchTag, + class L1Type, + class L0Type = void, + class Enable = void> +struct CopyL1ToL0B { + static_assert(DEPENDENT_FALSE, "Unsupported copy l1 to l0, can not find the specialization."); +}; + +/// Partial specialization for CopyL1ToL0B, Ascend950, nZ in and nZ out. (Transpose B) +template +struct CopyL1ToL0B> { + using LayoutDst = layout::nZ; + using LayoutSrc = layout::nZ; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyL1ToL0B() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + const uint32_t dstOuterShapeRow = layoutDst.shape(1); + const uint32_t dstOuterShapeCol = layoutDst.shape(3); + const uint32_t srcOuterStrideRow = layoutSrc.stride(1); + const uint32_t dstOuterStrideRow = layoutDst.stride(1); + + AscendC::LoadData2DParamsV2 loadDataParams; + loadDataParams.mStartPosition = 0; + loadDataParams.kStartPosition = 0; + loadDataParams.mStep = dstOuterShapeCol; + loadDataParams.kStep = dstOuterShapeRow; + loadDataParams.srcStride = CeilDiv(srcOuterStrideRow); + loadDataParams.dstStride = CeilDiv(dstOuterStrideRow); + loadDataParams.ifTranspose = false; + + AscendC::LoadData(dstTensor, srcTensor, loadDataParams); + } +}; + +/// Partial specialization for CopyL1ToL0B, AtlasA5, not B8 or B4, zN in and nZ out. +template +struct CopyL1ToL0B, void, + std::enable_if_t< + !AscendC::Std::is_one_of_v>> { + using LayoutDst = layout::nZ; + using LayoutSrc = layout::zN; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyL1ToL0B() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + const uint32_t L0K = layoutDst.shape(0) * layoutDst.shape(1); + const uint32_t L0N = layoutDst.shape(2) * layoutDst.shape(3); + const uint32_t srcOuterStrideCol = layoutSrc.stride(3); + const uint32_t dstOuterStrideRow = layoutDst.stride(1); + + AscendC::LoadData2DParamsV2 loadDataParams; + loadDataParams.mStartPosition = 0; + loadDataParams.kStartPosition = 0; + loadDataParams.mStep = CeilDiv(L0K); + loadDataParams.kStep = CeilDiv(L0N); + loadDataParams.srcStride = CeilDiv(srcOuterStrideCol); + loadDataParams.dstStride = CeilDiv(dstOuterStrideRow); + loadDataParams.ifTranspose = true; + + AscendC::LoadData(dstTensor, srcTensor, loadDataParams); + } +}; + +/// Partial specialization for CopyL1ToL0B, AtlasA5, B8 or B4, zN in and nZ out. +template +struct CopyL1ToL0B, void, + std::enable_if_t< + AscendC::Std::is_one_of_v>> { + using LayoutDst = layout::nZ; + using LayoutSrc = layout::zN; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyL1ToL0B() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + const uint32_t L0K = layoutDst.shape(0) * layoutDst.shape(1); + const uint32_t L0N = layoutDst.shape(2) * layoutDst.shape(3); + const uint32_t srcOuterStrideCol = layoutSrc.stride(3); + const uint32_t dstOuterStrideRow = layoutDst.stride(1); + + AscendC::LoadData2DParamsV2 loadDataParams; + if (L0N % ELE_NUM_PER_C0 == 0) { + loadDataParams.mStartPosition = 0; + loadDataParams.kStartPosition = 0; + loadDataParams.mStep = CeilDiv(L0K); + loadDataParams.kStep = CeilDiv(L0N); + loadDataParams.srcStride = CeilDiv(srcOuterStrideCol); + loadDataParams.dstStride = CeilDiv(dstOuterStrideRow); + loadDataParams.ifTranspose = true; + + AscendC::LoadData(dstTensor, srcTensor, loadDataParams); + } else { + for (uint32_t kIdx = 0; kIdx < L0K / ELE_NUM_PER_C0; kIdx++) { + loadDataParams.mStartPosition = kIdx * 2; + loadDataParams.kStartPosition = 0; + loadDataParams.mStep = 2; + loadDataParams.kStep = CeilDiv(L0N); + loadDataParams.srcStride = CeilDiv(srcOuterStrideCol); + loadDataParams.dstStride = CeilDiv(dstOuterStrideRow); + loadDataParams.ifTranspose = true; + + AscendC::LoadData( + dstTensor[kIdx * L0N * ELE_NUM_PER_C0], srcTensor, loadDataParams); + } + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace Catlass::Gemm::Tile + +#endif // CATLASS_KDA_ASCEND950_COPY_L1_TO_L0B_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_atlasa2_cast_int4_to_int8.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_atlasa2_cast_int4_to_int8.hpp new file mode 100644 index 0000000000..b11b3644fb --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_atlasa2_cast_int4_to_int8.hpp @@ -0,0 +1,180 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_ATLASA2_CAST_INT4_TO_INT8_HPP +#define CATLASS_KDA_ATLASA2_CAST_INT4_TO_INT8_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/kda_coord.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" +#include "catlass/gemm/kda_gemm_helper.hpp" + +namespace Catlass::Gemm::Tile { + +template < + class ArchTag, + class SrcType_, + class DstType_, + // Length of the compute elements + uint32_t COMPUTE_LEN_, + uint32_t STAGES = 2> +struct TileCastInt4ToInt8 { + using ElementSrc = typename SrcType_::Element; + using ElementDst = typename DstType_::Element; + using LayoutSrc = typename SrcType_::Layout; + using LayoutDst = typename DstType_::Layout; + + static constexpr uint32_t ELE_NUM_PER_BLK_INT8 = BYTE_PER_BLK / sizeof(ElementDst); + static constexpr uint32_t ELE_NUM_PER_BLK_INT4 = BYTE_PER_BLK / sizeof(ElementDst) * 2; + + static_assert( + sizeof(ElementSrc) == sizeof(ElementDst), + "Error: Src and Dst element sizes are equal (forbidden)"); + + static_assert( + std::is_same_v || + std::is_same_v, + "Unsupported layout, only can be RowMajor, ColumnMajor or RowMajorInt4 or ColumnMajorInt4"); + + static constexpr uint32_t COMPUTE_LEN = COMPUTE_LEN_; + static_assert(COMPUTE_LEN <= 24 * 1024, "COMPUTE_LEN cannot exceed 24 * 1024"); + + struct Params {}; + + /// Construct + CATLASS_DEVICE + TileCastInt4ToInt8(Arch::Resource const &resource, Params const ¶ms) + { + if constexpr (g_coreType == AscendC::AIV) { + uint32_t ubOffset = 0; + uint32_t ubInSize = COMPUTE_LEN * sizeof(int8_t) / 2; + uint32_t ubOutSize = COMPUTE_LEN * sizeof(ElementDst); + uint32_t ubWorkspaceSize = COMPUTE_LEN * sizeof(half); + // Init buffers + for (uint32_t i = 0; i < STAGES; i++) { + ubInTensorList[i] = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += ubInSize; + ubOutTensorList[i] = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += ubOutSize; + ubWorkspaceList[i] = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += ubWorkspaceSize; + + ubEventList[i] = i; + AscendC::SetFlag(ubEventList[i]); + AscendC::SetFlag(ubEventList[i]); + } + } + } + + /// Destructor + CATLASS_DEVICE + ~TileCastInt4ToInt8() + { + if constexpr (g_coreType == AscendC::AIV) { + for (uint32_t i = 0; i < STAGES; i++) { + AscendC::WaitFlag(ubEventList[i]); + AscendC::WaitFlag(ubEventList[i]); + } + } + } + + CATLASS_DEVICE + void operator()( + AscendC::GlobalTensor const &gmDst, LayoutDst const &layoutDst, + AscendC::GlobalTensor const &gmSrc, LayoutSrc const &layoutSrc) + { + uint32_t tilesNum = layoutSrc.shape(0); + uint32_t tileLen = layoutSrc.shape(1); + uint32_t tileLenRoundInt8 = RoundUp(layoutSrc.shape(1), ELE_NUM_PER_BLK_INT8); + uint32_t tileLenRoundInt4 = RoundUp(layoutSrc.shape(1), ELE_NUM_PER_BLK_INT4); + uint64_t tileStrideSrc = layoutSrc.stride(0); + uint64_t tileStrideDst = layoutDst.stride(0); + if constexpr (std::is_same_v) { + tilesNum = layoutSrc.shape(1); + tileLen = layoutSrc.shape(0); + tileLenRoundInt8 = RoundUp(layoutSrc.shape(0), ELE_NUM_PER_BLK_INT8); + tileLenRoundInt4 = RoundUp(layoutSrc.shape(0), ELE_NUM_PER_BLK_INT4); + tileStrideSrc = layoutSrc.stride(1); + tileStrideDst = layoutDst.stride(1); + } + uint32_t tilesPerAiv = tilesNum / AscendC::GetSubBlockNum(); + if (AscendC::GetSubBlockIdx() < (tilesNum % AscendC::GetSubBlockNum())) { + tilesPerAiv++; + } + uint64_t taskOffsetSrc = AscendC::GetSubBlockIdx() * tilesPerAiv * tileStrideSrc; + uint64_t taskOffsetDst = AscendC::GetSubBlockIdx() * tilesPerAiv * tileStrideDst; + if (AscendC::GetSubBlockIdx() >= (tilesNum % AscendC::GetSubBlockNum())) { + taskOffsetSrc += (tilesNum % AscendC::GetSubBlockNum()) * tileStrideSrc; + taskOffsetDst += (tilesNum % AscendC::GetSubBlockNum()) * tileStrideDst; + } + uint32_t tilesPerLoop = 32; + uint32_t loops = CeilDiv(tilesPerAiv, tilesPerLoop); + uint32_t pingpong = 0; + for (uint32_t loopIdx = 0; loopIdx < loops; loopIdx++) { + uint32_t actualTiles = tilesPerLoop; + if (loopIdx == loops - 1) { + actualTiles = tilesPerAiv - loopIdx * tilesPerLoop; + } + uint64_t tileOffsetSrc = loopIdx * tilesPerLoop * tileStrideSrc; + AscendC::DataCopyExtParams dataCopyParamsIn( + actualTiles, (tileLen + 1) / 2 * sizeof(int8_t), + (tileStrideSrc - tileLen) * sizeof(int8_t) / 2, + (tileLenRoundInt4 - tileLen) / ELE_NUM_PER_BLK_INT4, 0); + AscendC::DataCopyPadExtParams padParams(false, 0, 0, 0); + + AscendC::WaitFlag(ubEventList[pingpong]); + AscendC::DataCopyPad(ubInTensorList[pingpong], + gmSrc[taskOffsetSrc + tileOffsetSrc], dataCopyParamsIn, padParams); + + AscendC::SetFlag(ubEventList[pingpong]); + AscendC::WaitFlag(ubEventList[pingpong]); + + AscendC::WaitFlag(ubEventList[pingpong]); + + AscendC::Cast(ubWorkspaceList[pingpong], ubInTensorList[pingpong], + AscendC::RoundMode::CAST_NONE, actualTiles * tileLenRoundInt4); + + AscendC::PipeBarrier(); + + AscendC::SetFlag(ubEventList[pingpong]); + + AscendC::Cast(ubOutTensorList[pingpong], ubWorkspaceList[pingpong], + AscendC::RoundMode::CAST_NONE, actualTiles * tileLenRoundInt4); + + AscendC::SetFlag(ubEventList[pingpong]); + AscendC::WaitFlag(ubEventList[pingpong]); + + uint64_t tileOffsetDst = loopIdx * tilesPerLoop * tileStrideDst; + AscendC::DataCopyExtParams dataCopyParamsOut( + actualTiles, tileLen * sizeof(ElementDst), + (tileLenRoundInt4 - tileLen) / ELE_NUM_PER_BLK_INT8, + (tileStrideDst - tileLen) * sizeof(ElementDst), 0); + AscendC::DataCopyPad(gmDst[taskOffsetDst + tileOffsetDst], + ubOutTensorList[pingpong], dataCopyParamsOut); + AscendC::SetFlag(ubEventList[pingpong]); + + pingpong = (pingpong + 1) % STAGES; + } + } + +protected: + /// Data members + AscendC::LocalTensor ubInTensorList[STAGES]; + AscendC::LocalTensor ubOutTensorList[STAGES]; + AscendC::LocalTensor ubWorkspaceList[STAGES]; + + int32_t ubEventList[STAGES]; +}; + +} // namespace Catlass::Gemm::Tile + +#endif // CATLASS_KDA_ATLASA2_CAST_INT4_TO_INT8_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_atlasa2_copy_gm_to_l1.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_atlasa2_copy_gm_to_l1.hpp new file mode 100644 index 0000000000..cb35cafa6b --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_atlasa2_copy_gm_to_l1.hpp @@ -0,0 +1,2194 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_ATLASA2_COPY_GM_TO_L1_HPP +#define CATLASS_KDA_ATLASA2_COPY_GM_TO_L1_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/kda_numeric_size.hpp" +#include "catlass/arch/kda_arch.hpp" +#include "catlass/layout/kda_layout.hpp" +#include "catlass/gemm/kda_gemm_type.hpp" +#include "catlass/gemm/tile/kda_tile_copy_tla.hpp" +#include "tla/kda_tensor.hpp" + +namespace Catlass::Gemm::Tile { + +template < + class ArchTag, + /// GemmType for matrix operand + class GmType, + class L1Type = void> +struct CopyGmToL1 { + static_assert(DEPENDENT_FALSE, "Unsupported copy gm to l1, can not find the specialization."); +}; + +template < + class ArchTag, + /// GemmType for matrix operand + class GmType, + class L1Type = void> +struct CopyGmToL1IntervalDataCopy { + static_assert(DEPENDENT_FALSE, "Unsupported copy gm to l1, can not find the specialization."); +}; + +template < + class ArchTag, + /// GemmType for matrix operand + class GmType, + class L1Type = void> +struct CopyGmToL1GMMPTD { + static_assert(DEPENDENT_FALSE, "Unsupported copy gm to l1, can not find the specialization."); +}; + +template < + class ArchTag, + /// GemmType for matrix operand + class GmType, + class L1Type = void> +struct CopyGmToL1DynamicOptimized { + static_assert(DEPENDENT_FALSE, "Unsupported copy gm to l1, can not find the specialization."); +}; + +/// Partial specialization for AtlasA2, RowMajor in and zN out. +template +struct CopyGmToL1DynamicOptimized> { + using LayoutDst = layout::zN; + using LayoutSrc = layout::RowMajor; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyGmToL1DynamicOptimized() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = 1; + intriParams.dValue = layoutSrc.shape(1); + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzC0Stride = layoutDst.stride(3) / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + + if (layoutSrc.shape(0) <= 16) { + // If the number of matrix rows is 1, the regular interval-based DataCopy interface can be used instead of + // the ND2NZ DataCopy interface, resulting in higher transfer efficiency. + for (int i = 0; i < layoutSrc.shape(0); ++i) { + AscendC::DataCopyParams dataCopyParams( + CeilDiv(layoutSrc.shape(1), layoutDst.shape(2)), + layoutDst.shape(2) / ELE_NUM_PER_C0, + 0, + (layoutDst.stride(3) - layoutDst.shape(2)) / ELE_NUM_PER_C0); + AscendC::DataCopy( + dstTensor[i * layoutDst.shape(2)], srcTensor[i * layoutSrc.stride(0)], dataCopyParams); + } + } else { + if (layoutSrc.stride(0) < STRIDE_LIMIT) { + if (layoutSrc.shape(1) != ELE_NUM_PER_C0 || layoutSrc.stride(0) != ELE_NUM_PER_C0) { + intriParams.nValue = layoutSrc.shape(0); + intriParams.srcDValue = layoutSrc.stride(0); + intriParams.dstNzNStride = layoutDst.stride(0) / ELE_NUM_PER_C0; + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } else { + // If the matrix has ELE_NUM_PER_C0, columns and a stride of ELE_NUM_PER_C0, it follows a row-major + // layout in L1, allowing the use of standard contiguous DataCopy interface for more efficient + // transfers. + AscendC::DataCopy(dstTensor, srcTensor, layoutSrc.shape(0) * layoutSrc.shape(1)); + } + } else { + intriParams.nValue = 1; + intriParams.srcDValue = 0; + intriParams.dstNzNStride = 0; + for (uint32_t i = 0; i < layoutSrc.shape(0); i++) { + AscendC::DataCopy(dstTensor[i * ELE_NUM_PER_C0], srcTensor[i * layoutSrc.stride(0)], intriParams); + } + } + } + } +}; + +/// Partial specialization for AtlasA2, ColumnMajor in and nZ out. +template +struct CopyGmToL1DynamicOptimized> { + using LayoutDst = layout::nZ; + using LayoutSrc = layout::ColumnMajor; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyGmToL1DynamicOptimized() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = 1; + intriParams.dValue = layoutSrc.shape(0); + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + + if (layoutSrc.shape(1) <= 16) { + // If the number of matrix cols is 1, the regular interval-based DataCopy interface can be used instead of + // the ND2NZ DataCopy interface, resulting in higher transfer efficiency. + for (int i = 0; i < layoutSrc.shape(1); ++i) { + AscendC::DataCopyParams dataCopyParams( + CeilDiv(layoutSrc.shape(0), layoutDst.shape(0)), + layoutDst.shape(0) / ELE_NUM_PER_C0, + 0, + (layoutDst.stride(1) - layoutDst.shape(0)) / ELE_NUM_PER_C0); + AscendC::DataCopy( + dstTensor[i * layoutDst.shape(0)], srcTensor[i * layoutSrc.stride(1)], dataCopyParams); + } + } else { + if (layoutSrc.stride(1) < STRIDE_LIMIT) { + if (layoutSrc.shape(0) != ELE_NUM_PER_C0 || layoutSrc.stride(1) != ELE_NUM_PER_C0) { + intriParams.nValue = layoutSrc.shape(1); + intriParams.srcDValue = layoutSrc.stride(1); + intriParams.dstNzNStride = layoutDst.stride(2) / ELE_NUM_PER_C0; + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } else { + // If the matrix has ELE_NUM_PER_C0, rows and a stride of ELE_NUM_PER_C0, it follows a col-major + // layout in L1, allowing the use of standard contiguous DataCopy interface for more efficient + // transfers. + AscendC::DataCopy(dstTensor, srcTensor, layoutSrc.shape(0) * layoutSrc.shape(1)); + } + } else { + intriParams.nValue = 1; + intriParams.srcDValue = 0; + intriParams.dstNzNStride = 0; + for (uint32_t i = 0; i < layoutSrc.shape(1); i++) { + AscendC::DataCopy(dstTensor[i * ELE_NUM_PER_C0], srcTensor[i * layoutSrc.stride(1)], intriParams); + } + } + } + } +}; + +/// Partial specialization for AtlasA2, zN in and zN out. +template +struct CopyGmToL1DynamicOptimized> : public CopyGmToL1> {}; + +/// Partial specialization for AtlasA2, nZ in and nZ out. +template +struct CopyGmToL1DynamicOptimized> : public CopyGmToL1> {}; + +/// Partial specialization for AtlasA2, PaddingRowMajor in and zN out. +template +struct CopyGmToL1DynamicOptimized> : public CopyGmToL1> {}; + +/// Partial specialization for AtlasA2, PaddingColumnMajor in and nZ out. +template +struct CopyGmToL1DynamicOptimized> : public CopyGmToL1> {}; + +/// Partial specialization for AtlasA2, RowMajor in and zN out. +template +struct CopyGmToL1GMMPTD> { + using LayoutDst = layout::zN; + using LayoutSrc = layout::RowMajor; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyGmToL1GMMPTD() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = 1; + intriParams.dValue = layoutSrc.shape(1); + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzC0Stride = layoutDst.stride(3) / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + + if (layoutSrc.shape(0) == 1) { + // If the number of matrix rows is 1, the regular interval-based DataCopy interface can be used instead of + // the ND2NZ DataCopy interface, resulting in higher transfer efficiency. + AscendC::DataCopyParams dataCopyParams( + CeilDiv(layoutSrc.shape(1), layoutDst.shape(2)), + layoutDst.shape(2) / ELE_NUM_PER_C0, + 0, + (layoutDst.stride(3) - layoutDst.shape(2)) / ELE_NUM_PER_C0); + AscendC::DataCopy(dstTensor, srcTensor, dataCopyParams); + } else { + if (layoutSrc.stride(0) < STRIDE_LIMIT) { + if (layoutSrc.shape(1) != ELE_NUM_PER_C0 || layoutSrc.stride(0) != ELE_NUM_PER_C0) { + intriParams.nValue = layoutSrc.shape(0); + intriParams.srcDValue = layoutSrc.stride(0); + intriParams.dstNzNStride = layoutDst.stride(0) / ELE_NUM_PER_C0; + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } else { + // If the matrix has ELE_NUM_PER_C0 columns and a stride of ELE_NUM_PER_C0, it follows a row-major + // layout in L1, allowing the use of the standard contiguous DataCopy interface for more efficient + // transfers. + AscendC::DataCopy(dstTensor, srcTensor, layoutSrc.shape(0) * layoutSrc.shape(1)); + } + } else { + intriParams.nValue = 1; + intriParams.srcDValue = 0; + intriParams.dstNzNStride = 0; + for (uint32_t i = 0; i < layoutSrc.shape(0); i++) { + AscendC::DataCopy(dstTensor[i * ELE_NUM_PER_C0], srcTensor[i * layoutSrc.stride(0)], intriParams); + } + } + } + } + + // layoutSrc must be the layout of one of the src matrices + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc, + uint32_t ndNum, uint32_t srcNdMatrixStride, + uint32_t dstNzNStride, uint32_t dstNzMatrixStride, + uint32_t dstNzC0Stride) + { + AscendC::Nd2NzParams intriParams; + + intriParams.nValue = layoutSrc.shape(0); + intriParams.dValue = layoutSrc.shape(1); + intriParams.srcDValue = layoutSrc.stride(0); + intriParams.dstNzNStride = dstNzNStride; + intriParams.dstNzC0Stride = dstNzC0Stride; + if (srcNdMatrixStride < STRIDE_LIMIT) { + intriParams.ndNum = ndNum; + intriParams.srcNdMatrixStride = srcNdMatrixStride; + intriParams.dstNzMatrixStride = dstNzMatrixStride; + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } else { + intriParams.ndNum = 1; + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzMatrixStride = 0; + for (uint32_t i = 0; i < ndNum; i++) { + AscendC::DataCopy(dstTensor[i * ELE_NUM_PER_C0], srcTensor[i * srcNdMatrixStride], intriParams); + } + } + } +}; + +//////////////////////////////////////// +/// Using the standard strided DataCopy interface to implement nd2nz +/// transfer may achieve higher data transfer efficiency when the data block shape is short and wide +/// Partial specialization for AtlasA2, half, RowMajor in and zN out. +template <> +struct CopyGmToL1IntervalDataCopy> { + using LayoutDst = layout::zN; + using LayoutSrc = layout::RowMajor; + using Element = half; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyGmToL1IntervalDataCopy() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + for (int i = 0; i < layoutSrc.shape(0); ++i) { + AscendC::DataCopyParams dataCopyParams( + CeilDiv(layoutSrc.shape(1), layoutDst.shape(2)), + layoutDst.shape(2) / ELE_NUM_PER_C0, + 0, + (layoutDst.stride(3) - layoutDst.shape(2)) / ELE_NUM_PER_C0); + AscendC::DataCopy(dstTensor[i * layoutDst.shape(2)], srcTensor[i * layoutSrc.stride(0)], dataCopyParams); + } + } +}; + +/// Partial specialization for AtlasA2, half, PaddingRowMajor in and zN out. +/// Using the standard strided DataCopy interface to implement nd2nz +/// transfer may achieve higher data transfer efficiency when the data block shape is short and wide +template <> +struct CopyGmToL1IntervalDataCopy> { + using LayoutDst = layout::zN; + using LayoutSrc = layout::PaddingRowMajor; + using Element = half; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyGmToL1IntervalDataCopy() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + for (int i = 0; i < layoutSrc.orgShape(0); ++i) { + AscendC::DataCopyParams dataCopyParams( + CeilDiv(layoutSrc.orgShape(1), layoutDst.shape(2)), + layoutDst.shape(2) / ELE_NUM_PER_C0, + 0, + (layoutDst.stride(3) - layoutDst.shape(2)) / ELE_NUM_PER_C0); + AscendC::DataCopy(dstTensor[i * layoutDst.shape(2)], srcTensor[i * layoutSrc.stride(0)], dataCopyParams); + } + } +}; + +/// Partial specialization for AtlasA2, half, ColumnMajor in and zN out. +/// Using the standard strided DataCopy interface to implement nd2nz +/// transfer may achieve higher data transfer efficiency when the data block shape is tall and narrow +template <> +struct CopyGmToL1IntervalDataCopy> { + using LayoutDst = layout::nZ; + using LayoutSrc = layout::ColumnMajor; + using Element = half; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyGmToL1IntervalDataCopy() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + for (int i = 0; i < layoutSrc.shape(1); ++i) { + AscendC::DataCopyParams dataCopyParams( + CeilDiv(layoutSrc.shape(0), layoutDst.shape(0)), + layoutDst.shape(0) / ELE_NUM_PER_C0, + 0, + (layoutDst.stride(1) - layoutDst.shape(0)) / ELE_NUM_PER_C0); + AscendC::DataCopy(dstTensor[i * layoutDst.shape(0)], srcTensor[i * layoutSrc.stride(1)], dataCopyParams); + } + } +}; + +/// Partial specialization for AtlasA2, half, PaddingColumnMajor in and zN out. +/// Using the standard strided DataCopy interface to implement nd2nz +/// transfer may achieve higher data transfer efficiency when the data block shape is tall and narrow +template <> +struct CopyGmToL1IntervalDataCopy> { + using LayoutDst = layout::nZ; + using LayoutSrc = layout::PaddingColumnMajor; + using Element = half; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyGmToL1IntervalDataCopy() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + for (int i = 0; i < layoutSrc.orgShape(1); ++i) { + AscendC::DataCopyParams dataCopyParams( + CeilDiv(layoutSrc.orgShape(0), layoutDst.shape(0)), + layoutDst.shape(0) / ELE_NUM_PER_C0, + 0, + (layoutDst.stride(1) - layoutDst.shape(0)) / ELE_NUM_PER_C0); + AscendC::DataCopy(dstTensor[i * layoutDst.shape(0)], srcTensor[i * layoutSrc.stride(2)], dataCopyParams); + } + } +}; + +/// new add gemm +template +struct CopyGmToL1, Gemm::GemmType> { + using LayoutDst = layout::zN; + using LayoutSrc = layout::RowMajor; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyGmToL1() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = 1; + intriParams.dValue = layoutSrc.shape(1); + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzC0Stride = layoutDst.stride(3) / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + + if (layoutSrc.stride(0) < STRIDE_LIMIT) { + intriParams.nValue = layoutSrc.shape(0); + intriParams.srcDValue = layoutSrc.stride(0); + intriParams.dstNzNStride = layoutDst.stride(0) / ELE_NUM_PER_C0; + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } else { + intriParams.nValue = 1; + intriParams.srcDValue = 0; + intriParams.dstNzNStride = 0; + for (uint32_t i = 0; i < layoutSrc.shape(0); i++) { + AscendC::DataCopy(dstTensor[i * ELE_NUM_PER_C0], srcTensor[i * layoutSrc.stride(0)], intriParams); + } + } + } +}; + +template +struct CopyGmToL1, Gemm::GemmType> { + using LayoutDst = layout::zZ; + using LayoutSrc = layout::RowMajor; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyGmToL1() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::Nd2NzParams intriParams; + uint32_t srcNdStride = C0_NUM_PER_FRACTAL * layoutSrc.stride(0); + uint32_t ndNum = layoutSrc.shape(0) / C0_NUM_PER_FRACTAL; + uint32_t remains = layoutSrc.shape(0) % C0_NUM_PER_FRACTAL; + if (srcNdStride < STRIDE_LIMIT) { + if (ndNum) { + intriParams.ndNum = ndNum; + intriParams.nValue = C0_NUM_PER_FRACTAL; + intriParams.dValue = layoutSrc.shape(1); + intriParams.srcNdMatrixStride = srcNdStride; + intriParams.srcDValue = layoutSrc.stride(0); + + intriParams.dstNzC0Stride = layoutDst.stride(3) / ELE_NUM_PER_C0; + intriParams.dstNzNStride = layoutDst.stride(0) / ELE_NUM_PER_C0; + + intriParams.dstNzMatrixStride = layoutDst.stride(1); + + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } + + if (remains) { + AscendC::Nd2NzParams tailParams; + tailParams.ndNum = 1; + tailParams.nValue = remains; + tailParams.dValue = layoutSrc.shape(1); + tailParams.srcNdMatrixStride = srcNdStride; + tailParams.srcDValue = layoutSrc.stride(0); + + tailParams.dstNzC0Stride = layoutDst.stride(3) / ELE_NUM_PER_C0; + tailParams.dstNzNStride = layoutDst.stride(0) / ELE_NUM_PER_C0; + tailParams.dstNzMatrixStride = 0; //` + + AscendC::DataCopy(dstTensor[ndNum * layoutDst.stride(1)], srcTensor[ndNum * srcNdStride], tailParams); + } + } else if (layoutSrc.stride(0) < STRIDE_LIMIT) { + for (uint32_t i = 0; i < ndNum; i++) { + AscendC::Nd2NzParams intriParams; + intriParams.ndNum = 1; + intriParams.nValue = C0_NUM_PER_FRACTAL; + intriParams.dValue = layoutSrc.shape(1); + intriParams.srcNdMatrixStride = 0; + intriParams.srcDValue = layoutSrc.stride(0); + + intriParams.dstNzC0Stride = layoutDst.stride(3) / ELE_NUM_PER_C0; + intriParams.dstNzNStride = layoutDst.stride(0) / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + + AscendC::DataCopy(dstTensor[i * layoutDst.stride(1)], srcTensor[i * srcNdStride], intriParams); + } + if (remains) { + AscendC::Nd2NzParams tailParams; + tailParams.ndNum = 1; + tailParams.nValue = remains; + tailParams.dValue = layoutSrc.shape(1); + tailParams.srcNdMatrixStride = 0; + tailParams.srcDValue = layoutSrc.stride(0); + + tailParams.dstNzC0Stride = layoutDst.stride(3) / ELE_NUM_PER_C0; + tailParams.dstNzNStride = layoutDst.stride(0) / ELE_NUM_PER_C0; + tailParams.dstNzMatrixStride = 0; + + AscendC::DataCopy(dstTensor[ndNum * layoutDst.stride(1)], srcTensor[ndNum * srcNdStride], tailParams); + } + } else { + for (uint32_t i = 0; i < layoutSrc.shape(0); i++) { + uint32_t idxR0 = i / C0_NUM_PER_FRACTAL; + uint32_t idxInR0 = i % C0_NUM_PER_FRACTAL; + + AscendC::Nd2NzParams intriParams; + intriParams.ndNum = 1; + intriParams.nValue = 1; + intriParams.dValue = layoutSrc.shape(1); + intriParams.srcNdMatrixStride = 0; + intriParams.srcDValue = 0; + + intriParams.dstNzC0Stride = layoutDst.stride(3) / ELE_NUM_PER_C0; + intriParams.dstNzNStride = 0; + intriParams.dstNzMatrixStride = 0; + + uint32_t offsetDst = i * idxR0 * layoutDst.stride(1) + idxInR0 * ELE_NUM_PER_C0; + uint32_t offsetSrc = i * layoutSrc.stride(0); + AscendC::DataCopy(dstTensor[offsetDst], srcTensor[offsetSrc], intriParams); + } + } + } +}; + +template +struct CopyGmToL1, Gemm::GemmType> { + using LayoutDst = layout::nN; + using LayoutSrc = layout::ColumnMajor; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyGmToL1() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::Nd2NzParams intriParams; + uint32_t srcNdStride = C0_NUM_PER_FRACTAL * layoutSrc.stride(1); + uint32_t ndNum = layoutSrc.shape(1) / C0_NUM_PER_FRACTAL; + uint32_t remains = layoutSrc.shape(1) % C0_NUM_PER_FRACTAL; + if (srcNdStride < STRIDE_LIMIT) { + if (ndNum) { + intriParams.ndNum = ndNum; + intriParams.nValue = C0_NUM_PER_FRACTAL; + intriParams.dValue = layoutSrc.shape(0); + intriParams.srcNdMatrixStride = srcNdStride; + intriParams.srcDValue = layoutSrc.stride(1); + + intriParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0; + intriParams.dstNzNStride = layoutDst.stride(2) / ELE_NUM_PER_C0; + + intriParams.dstNzMatrixStride = layoutDst.stride(3); + + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } + + if (remains) { + AscendC::Nd2NzParams tailParams; + tailParams.ndNum = 1; + tailParams.nValue = remains; + tailParams.dValue = layoutSrc.shape(0); + tailParams.srcNdMatrixStride = srcNdStride; + tailParams.srcDValue = layoutSrc.stride(1); + + tailParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0; + tailParams.dstNzNStride = layoutDst.stride(2) / ELE_NUM_PER_C0; + tailParams.dstNzMatrixStride = 0; + + AscendC::DataCopy(dstTensor[ndNum * layoutDst.stride(3)], srcTensor[ndNum * srcNdStride], tailParams); + } + } else if (layoutSrc.stride(1) < STRIDE_LIMIT) { + for (uint32_t i = 0; i < ndNum; i++) { + AscendC::Nd2NzParams intriParams; + intriParams.ndNum = 1; + intriParams.nValue = C0_NUM_PER_FRACTAL; + intriParams.dValue = layoutSrc.shape(0); + intriParams.srcNdMatrixStride = 0; + intriParams.srcDValue = layoutSrc.stride(1); + + intriParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0; + intriParams.dstNzNStride = layoutDst.stride(2) / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + + AscendC::DataCopy(dstTensor[i * layoutDst.stride(3)], srcTensor[i * srcNdStride], intriParams); + } + if (remains) { + AscendC::Nd2NzParams tailParams; + tailParams.ndNum = 1; + tailParams.nValue = remains; + tailParams.dValue = layoutSrc.shape(0); + tailParams.srcNdMatrixStride = 0; + tailParams.srcDValue = layoutSrc.stride(1); + + tailParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0; + tailParams.dstNzNStride = layoutDst.stride(2) / ELE_NUM_PER_C0; + tailParams.dstNzMatrixStride = 0; + + AscendC::DataCopy(dstTensor[ndNum * layoutDst.stride(3)], srcTensor[ndNum * srcNdStride], tailParams); + } + } else { + for (uint32_t i = 0; i < layoutSrc.shape(1); i++) { + uint32_t idxR0 = i / C0_NUM_PER_FRACTAL; + uint32_t idxInR0 = i % C0_NUM_PER_FRACTAL; + + AscendC::Nd2NzParams intriParams; + intriParams.ndNum = 1; + intriParams.nValue = 1; + intriParams.dValue = layoutSrc.shape(0); + intriParams.srcNdMatrixStride = 0; + intriParams.srcDValue = 0; + + intriParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0; + intriParams.dstNzNStride = 0; + intriParams.dstNzMatrixStride = 0; + + uint32_t offsetDst = i * idxR0 * layoutDst.stride(3) + idxInR0 * ELE_NUM_PER_C0; + uint32_t offsetSrc = i * layoutSrc.stride(1); + AscendC::DataCopy(dstTensor[offsetDst], srcTensor[offsetSrc], intriParams); + } + } + } +}; + +template +struct CopyGmToL1, Gemm::GemmType> { + using LayoutDst = layout::nZ; + using LayoutSrc = layout::ColumnMajor; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyGmToL1() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = 1; + intriParams.dValue = layoutSrc.shape(0); + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + + if (layoutSrc.stride(1) < STRIDE_LIMIT) { + intriParams.nValue = layoutSrc.shape(1); + intriParams.srcDValue = layoutSrc.stride(1); + intriParams.dstNzNStride = layoutDst.stride(2) / ELE_NUM_PER_C0; + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } else { + intriParams.nValue = 1; + intriParams.srcDValue = 0; + intriParams.dstNzNStride = 0; + for (uint32_t i = 0; i < layoutSrc.shape(1); i++) { + AscendC::DataCopy(dstTensor[i * ELE_NUM_PER_C0], srcTensor[i * layoutSrc.stride(1)], intriParams); + } + } + } +}; + +template +struct CopyGmToL1, Gemm::GemmType> { + using LayoutDst = layout::nZ; + using LayoutSrc = layout::ColumnMajor; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyGmToL1() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = 1; + intriParams.dValue = layoutSrc.shape(0); + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + + if (layoutSrc.stride(1) < STRIDE_LIMIT) { + intriParams.nValue = layoutSrc.shape(1); + intriParams.srcDValue = layoutSrc.stride(1); + intriParams.dstNzNStride = layoutDst.stride(2) / ELE_NUM_PER_C0; + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } else { + intriParams.nValue = 1; + intriParams.srcDValue = 0; + intriParams.dstNzNStride = 0; + for (uint32_t i = 0; i < layoutSrc.shape(1); i++) { + AscendC::DataCopy(dstTensor[i * ELE_NUM_PER_C0], srcTensor[i * layoutSrc.stride(1)], intriParams); + } + } + } +}; +//////////////////////////////////////// + +/////////////////////////////////////// +/// new add gemv, VectorLayout -> zN +template +struct CopyGmToL1, Gemm::GemmType> { + using LayoutDst = layout::zN; + using LayoutSrc = layout::VectorLayout; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyGmToL1() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = 1; + intriParams.dValue = layoutSrc.shape(0); + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzC0Stride = layoutDst.stride(3) / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + intriParams.nValue = 1; + intriParams.srcDValue = layoutSrc.shape(0); + intriParams.dstNzNStride = layoutDst.stride(0) / ELE_NUM_PER_C0; + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } +}; + +template +struct CopyGmToL1> { + using LayoutDst = layout::NDC1HWC0; + using LayoutSrc = layout::NDC1HWC0; + + // Methods + + CATLASS_DEVICE + CopyGmToL1() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + const static uint64_t MAX_UINT16 = 65535; + + uint32_t cin1LoadL1 = layoutDst.orgShape(2); + uint32_t hiLoadL1 = layoutDst.orgShape(3); + + uint32_t dilationD = layoutSrc.orgShape(1); + uint32_t OriC1 = layoutSrc.orgShape(2); + uint32_t OriH = layoutSrc.orgShape(3); + uint32_t OriW = layoutSrc.orgShape(4); + uint32_t OriK0 = layoutSrc.orgShape(5); + + uint64_t dataCopyLoop = CeilDiv(cin1LoadL1, OriC1); + uint64_t dataCopySubLoop = 0; + uint64_t blockCount = OriC1; + bool srcStrideBeyondMaxU16 = false; + if (OriH * OriW - hiLoadL1 * OriW > MAX_UINT16) { + dataCopySubLoop = dataCopyLoop > 0 ? cin1LoadL1 / dataCopyLoop : 0; + blockCount = 1; + srcStrideBeyondMaxU16 = true; + } + + uint64_t aL1GmOffset = 0; + uint64_t aL1Offset = 0; + if (cin1LoadL1 > OriC1 || srcStrideBeyondMaxU16) { + repeatParams.blockCount = blockCount; + repeatParams.blockLen = hiLoadL1 * OriW; + repeatParams.srcStride = OriW * OriH - repeatParams.blockLen; + repeatParams.dstStride = 0; + for (uint64_t i = 0; i < dataCopyLoop; i++) { + if (srcStrideBeyondMaxU16) { + uint64_t aL1GmSubOffset = aL1GmOffset; + uint64_t aL1SubOffset = aL1Offset; + for (uint64_t j = 0; j < dataCopySubLoop; j++) { + AscendC::DataCopy(dstTensor[aL1SubOffset], srcTensor[aL1GmSubOffset], repeatParams); + aL1GmSubOffset += OriH * OriW * OriK0; + aL1SubOffset += hiLoadL1 * OriW * OriK0; + } + } else { + AscendC::DataCopy(dstTensor[aL1Offset], srcTensor[aL1GmOffset], repeatParams); + } + aL1GmOffset += dilationD * OriC1 * OriH * OriW * OriK0; + aL1Offset += OriC1 * hiLoadL1 * OriW * OriK0; + } + } else { + repeatParams.blockCount = cin1LoadL1; + repeatParams.blockLen = hiLoadL1 * OriW; + repeatParams.srcStride = OriW * OriH - repeatParams.blockLen; + repeatParams.dstStride = 0; + AscendC::DataCopy(dstTensor[aL1Offset], srcTensor[aL1GmOffset], repeatParams); + aL1Offset += cin1LoadL1 * hiLoadL1 * OriW * OriK0; + } + } + +private: + uint64_t aL1GmOffset = 0; + AscendC::DataCopyParams repeatParams; +}; + +template +struct CopyGmToL1> { + using LayoutDst = layout::nZ; + using LayoutSrc = layout::KDC1KHKWN1N0C0; + + // Methods + + CATLASS_DEVICE + CopyGmToL1() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + uint32_t currentNBL1 = layoutDst.orgShape(1); + uint32_t currentKBL1 = layoutDst.orgShape(0); + + uint32_t N1 = layoutSrc.shape(3); + uint32_t N0 = layoutSrc.shape(2); + uint32_t C0 = layoutSrc.shape(0); + uint32_t OriCoAlign = N1 * N0; + + const static uint32_t LOAD2D_MAX_REPEAT_TIMES = 255; + + if (currentNBL1 >= OriCoAlign) { + uint32_t repeatTimes = (currentKBL1 * currentNBL1) / (N0 * C0); + if (repeatTimes > LOAD2D_MAX_REPEAT_TIMES) { + repeatParams.blockCount = 1; + repeatParams.srcStride = 0; + repeatParams.blockLen = CeilDiv(currentKBL1 * currentNBL1, C0); + AscendC::DataCopy(dstTensor, srcTensor, repeatParams); + } else { + AscendC::LoadData2DParams loadData2dParams; + loadData2dParams.srcStride = 1; + loadData2dParams.repeatTimes = repeatTimes; + AscendC::LoadData(dstTensor, srcTensor, loadData2dParams); + } + } else { + repeatParams.blockCount = currentKBL1 / C0; + repeatParams.blockLen = currentNBL1; + repeatParams.srcStride = N1 * N0 - currentNBL1; + AscendC::DataCopy(dstTensor, srcTensor, repeatParams); + } + } + +private: + AscendC::DataCopyParams repeatParams; +}; + +/////////////////////////////////////// +/// new add gemv, ColumnMajor -> nN +template +struct CopyGmToL1, Gemm::GemmType> { + using LayoutDst = layout::nN; + using LayoutSrc = layout::ColumnMajor; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyGmToL1() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::Nd2NzParams intriParams; + uint32_t srcNdStride = C0_NUM_PER_FRACTAL * layoutSrc.stride(1); + uint32_t ndNum = layoutSrc.shape(1) / C0_NUM_PER_FRACTAL; + uint32_t remains = layoutSrc.shape(1) % C0_NUM_PER_FRACTAL; + if (srcNdStride < STRIDE_LIMIT) { + if (ndNum) { + intriParams.ndNum = ndNum; + intriParams.nValue = C0_NUM_PER_FRACTAL; + intriParams.dValue = layoutSrc.shape(0); + intriParams.srcNdMatrixStride = srcNdStride; + intriParams.srcDValue = layoutSrc.stride(1); + + intriParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0; + intriParams.dstNzNStride = layoutDst.stride(2) / ELE_NUM_PER_C0; + + intriParams.dstNzMatrixStride = layoutDst.stride(3); + + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } + + if (remains) { + AscendC::Nd2NzParams tailParams; + tailParams.ndNum = 1; + tailParams.nValue = remains; + tailParams.dValue = layoutSrc.shape(0); + tailParams.srcNdMatrixStride = srcNdStride; + tailParams.srcDValue = layoutSrc.stride(1); + + tailParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0; + tailParams.dstNzNStride = layoutDst.stride(2) / ELE_NUM_PER_C0; + tailParams.dstNzMatrixStride = 0; + + AscendC::DataCopy(dstTensor[ndNum * layoutDst.stride(3)], srcTensor[ndNum * srcNdStride], tailParams); + } + } else if (layoutSrc.stride(1) < STRIDE_LIMIT) { + for (uint32_t i = 0; i < ndNum; i++) { + AscendC::Nd2NzParams intriParams; + intriParams.ndNum = 1; + intriParams.nValue = C0_NUM_PER_FRACTAL; + intriParams.dValue = layoutSrc.shape(0); + intriParams.srcNdMatrixStride = 0; + intriParams.srcDValue = layoutSrc.stride(1); + + intriParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0; + intriParams.dstNzNStride = layoutDst.stride(2) / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + + AscendC::DataCopy(dstTensor[i * layoutDst.stride(3)], srcTensor[i * srcNdStride], intriParams); + } + if (remains) { + AscendC::Nd2NzParams tailParams; + tailParams.ndNum = 1; + tailParams.nValue = remains; + tailParams.dValue = layoutSrc.shape(0); + tailParams.srcNdMatrixStride = 0; + tailParams.srcDValue = layoutSrc.stride(1); + + tailParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0; + tailParams.dstNzNStride = layoutDst.stride(2) / ELE_NUM_PER_C0; + tailParams.dstNzMatrixStride = 0; + + AscendC::DataCopy(dstTensor[ndNum * layoutDst.stride(3)], srcTensor[ndNum * srcNdStride], tailParams); + } + } else { + for (uint32_t i = 0; i < layoutSrc.shape(1); i++) { + uint32_t idxR0 = i / C0_NUM_PER_FRACTAL; + uint32_t idxInR0 = i % C0_NUM_PER_FRACTAL; + + AscendC::Nd2NzParams intriParams; + intriParams.ndNum = 1; + intriParams.nValue = 1; + intriParams.dValue = layoutSrc.shape(0); + intriParams.srcNdMatrixStride = 0; + intriParams.srcDValue = 0; + + intriParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0; + intriParams.dstNzNStride = 0; + intriParams.dstNzMatrixStride = 0; + + uint32_t offsetDst = i * idxR0 * layoutDst.stride(3) + idxInR0 * ELE_NUM_PER_C0; + uint32_t offsetSrc = i * layoutSrc.stride(1); + AscendC::DataCopy(dstTensor[offsetDst], srcTensor[offsetSrc], intriParams); + } + } + } +}; + +template +struct CopyGmToL1, Gemm::GemmType> { + using LayoutDst = layout::zN; + using LayoutSrc = layout::RowMajor; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyGmToL1() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = 1; + intriParams.dValue = layoutSrc.shape(1); + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzC0Stride = layoutDst.stride(3) / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + + if (layoutSrc.stride(0) < STRIDE_LIMIT) { + intriParams.nValue = layoutSrc.shape(0); + intriParams.srcDValue = layoutSrc.stride(0); + intriParams.dstNzNStride = layoutDst.stride(0) / ELE_NUM_PER_C0; + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } else { + intriParams.nValue = 1; + intriParams.srcDValue = 0; + intriParams.dstNzNStride = 0; + for (uint32_t i = 0; i < layoutSrc.shape(0); i++) { + AscendC::DataCopy(dstTensor[i * ELE_NUM_PER_C0], srcTensor[i * layoutSrc.stride(0)], intriParams); + } + } + } +}; +///////////////////////////////// + +/// Partial specialization for AtlasA2, RowMajor in and zN out. +template +struct CopyGmToL1> { + using LayoutDst = layout::zN; + using LayoutSrc = layout::RowMajor; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyGmToL1() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = 1; + intriParams.dValue = layoutSrc.shape(1); + // There is an issue with AscenC in the data migration of nd2nz of type int4b_t: dValue needs to be passed + // with the length of int8, and it is uncertain whether this will be fixed in subsequent versions. + if constexpr (std::is_same_v) { + intriParams.dValue = CeilDiv(intriParams.dValue, 2); + } + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzC0Stride = layoutDst.stride(3) / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + + if (layoutSrc.stride(0) < STRIDE_LIMIT) { + intriParams.nValue = layoutSrc.shape(0); + intriParams.srcDValue = layoutSrc.stride(0); + if constexpr (std::is_same_v) { + intriParams.srcDValue = CeilDiv(intriParams.srcDValue, 2); + } + intriParams.dstNzNStride = layoutDst.stride(0) / ELE_NUM_PER_C0; + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } else { + intriParams.nValue = 1; + intriParams.srcDValue = 0; + intriParams.dstNzNStride = 0; + for (uint32_t i = 0; i < layoutSrc.shape(0); i++) { + AscendC::DataCopy(dstTensor[i * ELE_NUM_PER_C0], srcTensor[i * layoutSrc.stride(0)], intriParams); + } + } + } + + // layoutSrc must be the layout of one of the src matrices + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc, + uint32_t ndNum, uint32_t srcNdMatrixStride, + uint32_t dstNzNStride, uint32_t dstNzMatrixStride, + uint32_t dstNzC0Stride) + { + AscendC::Nd2NzParams intriParams; + + intriParams.nValue = layoutSrc.shape(0); + intriParams.dValue = layoutSrc.shape(1); + intriParams.srcDValue = layoutSrc.stride(0); + intriParams.dstNzNStride = dstNzNStride; + intriParams.dstNzC0Stride = dstNzC0Stride; + if (srcNdMatrixStride < STRIDE_LIMIT) { + intriParams.ndNum = ndNum; + intriParams.srcNdMatrixStride = srcNdMatrixStride; + intriParams.dstNzMatrixStride = dstNzMatrixStride; + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } else { + intriParams.ndNum = 1; + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzMatrixStride = 0; + for (uint32_t i = 0; i < ndNum; i++) { + AscendC::DataCopy(dstTensor[i * ELE_NUM_PER_C0], srcTensor[i * srcNdMatrixStride], intriParams); + } + } + } +}; +///////////////////////////////// + +/// Partial specialization for AtlasA2, ColumnMajor in and nZ out. +template < + class Element> +struct CopyGmToL1> { + using LayoutDst = layout::nZ; + using LayoutSrc = layout::ColumnMajor; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyGmToL1() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = 1; + intriParams.dValue = layoutSrc.shape(0); + // There is an issue with AscenC in the data migration of nd2nz of type int4b_t: dValue needs to be passed + // with the length of int8, and it is uncertain whether this will be fixed in subsequent versions. + if constexpr (std::is_same_v) { + intriParams.dValue = CeilDiv(intriParams.dValue, 2); + } + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + + if (layoutSrc.stride(1) < STRIDE_LIMIT) { + intriParams.nValue = layoutSrc.shape(1); + intriParams.srcDValue = layoutSrc.stride(1); + if constexpr (std::is_same_v) { + intriParams.srcDValue = CeilDiv(intriParams.srcDValue, 2); + } + intriParams.dstNzNStride = layoutDst.stride(2) / ELE_NUM_PER_C0; + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } else { + intriParams.nValue = 1; + intriParams.srcDValue = 0; + intriParams.dstNzNStride = 0; + for (uint32_t i = 0; i < layoutSrc.shape(1); i++) { + AscendC::DataCopy(dstTensor[i * ELE_NUM_PER_C0], srcTensor[i * layoutSrc.stride(1)], intriParams); + } + } + } +}; + +/// Partial specialization for zN in and zN out. +template < + class ArchTag, + class Element> +struct CopyGmToL1> { + using LayoutDst = layout::zN; + using LayoutSrc = layout::zN; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyGmToL1() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + uint32_t blockCount = CeilDiv(layoutSrc.orgShape(1)); + uint32_t blockLen = layoutSrc.orgShape(0); + + AscendC::DataCopyParams repeatParams; + + if (layoutSrc.stride(3) / ELE_NUM_PER_C0 < STRIDE_LIMIT) { + repeatParams.blockCount = blockCount; + repeatParams.blockLen = blockLen; + repeatParams.srcStride = layoutSrc.stride(3) / ELE_NUM_PER_C0 - blockLen; + repeatParams.dstStride = layoutDst.stride(3) / ELE_NUM_PER_C0 - blockLen; + AscendC::DataCopy(dstTensor, srcTensor, repeatParams); + } else { + repeatParams.blockCount = 1; + repeatParams.blockLen = blockLen; + repeatParams.srcStride = 0; + repeatParams.dstStride = 0; + for (uint32_t i = 0; i < blockCount; i++) { + uint64_t dstOffset = i * layoutDst.stride(3); + uint64_t srcOffset = i * layoutSrc.stride(3); + AscendC::DataCopy(dstTensor[dstOffset], srcTensor[srcOffset], repeatParams); + } + } + } +}; + +/// Partial specialization for nZ in and nZ out. +template < + class ArchTag, + class Element> +struct CopyGmToL1> { + using LayoutDst = layout::nZ; + using LayoutSrc = layout::nZ; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyGmToL1() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + uint32_t blockCount = CeilDiv(layoutSrc.orgShape(0)); + uint32_t blockLen = layoutSrc.orgShape(1); + + AscendC::DataCopyParams repeatParams; + + if (layoutSrc.stride(1) / ELE_NUM_PER_C0 < STRIDE_LIMIT) { + repeatParams.blockCount = blockCount; + repeatParams.blockLen = blockLen; + repeatParams.srcStride = layoutSrc.stride(1) / ELE_NUM_PER_C0 - blockLen; + repeatParams.dstStride = layoutDst.stride(1) / ELE_NUM_PER_C0 - blockLen; + AscendC::DataCopy(dstTensor, srcTensor, repeatParams); + } else { + repeatParams.blockCount = 1; + repeatParams.blockLen = blockLen; + repeatParams.srcStride = 0; + repeatParams.dstStride = 0; + for (uint32_t i = 0; i < blockCount; i++) { + uint64_t dstOffset = i * layoutDst.stride(1); + uint64_t srcOffset = i * layoutSrc.stride(1); + AscendC::DataCopy(dstTensor[dstOffset], srcTensor[srcOffset], repeatParams); + } + } + } +}; + +/// Partial specialization for AtlasA2, PaddingRowMajor in and zN out. +template +struct CopyGmToL1> { + using LayoutDst = layout::zN; + using LayoutSrc = layout::PaddingRowMajor; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyGmToL1() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = 1; + intriParams.dValue = layoutSrc.orgShape(1); + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzC0Stride = layoutDst.stride(3) / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + + intriParams.nValue = layoutSrc.orgShape(0); + intriParams.srcDValue = layoutSrc.stride(0); + intriParams.dstNzNStride = layoutDst.stride(0) / ELE_NUM_PER_C0; + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } +}; + +/// Partial specialization for AtlasA2, ColumnMajor in and nZ out. +template < + class Element> +struct CopyGmToL1> { + using LayoutDst = layout::nZ; + using LayoutSrc = layout::PaddingColumnMajor; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyGmToL1() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = 1; + intriParams.dValue = layoutSrc.orgShape(0); + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + + intriParams.nValue = layoutSrc.orgShape(1); + intriParams.srcDValue = layoutSrc.stride(2); + intriParams.dstNzNStride = layoutDst.stride(2) / ELE_NUM_PER_C0; + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } +}; + +/// Partial specialization for AtlasA2, RowMajor in and RowMajor out. +template +struct CopyGmToL1, + Gemm::GemmType> { + using LayoutDst = layout::RowMajor; + using LayoutSrc = layout::RowMajor; + + static constexpr uint32_t ELE_NUM_PER_BLK = BytesToBits(BYTE_PER_BLK) / SizeOfBits::value; + static constexpr uint32_t BLOCK_LEN_LIMIT = 65536; + static constexpr uint32_t MAX_REPEAT = 4095; + + // Methods + + CATLASS_DEVICE + CopyGmToL1() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + uint32_t rows = layoutSrc.shape(0); + uint32_t cols = layoutSrc.shape(1); + uint32_t srcStride = (layoutSrc.stride(0) - layoutSrc.shape(1)) / ELE_NUM_PER_BLK; + uint32_t dstStride = (layoutDst.stride(0) - layoutDst.shape(1)) / ELE_NUM_PER_BLK; + + if ((layoutSrc.shape(1) == layoutSrc.stride(0)) && (layoutDst.shape(1) == layoutDst.stride(0))) { + DataCopy(dstTensor, srcTensor, rows * cols); + } else if (srcStride < STRIDE_LIMIT && dstStride < STRIDE_LIMIT && (cols / ELE_NUM_PER_BLK) < BLOCK_LEN_LIMIT) { + uint32_t rLoops = CeilDiv(rows, MAX_REPEAT); + for (uint32_t i = 0; i < rLoops; ++i) { + uint32_t rActual = (i < rLoops - 1) ? MAX_REPEAT : rows - i * MAX_REPEAT; + AscendC::DataCopyParams dataCopyParams( + rActual, cols / ELE_NUM_PER_BLK, srcStride, dstStride); + DataCopy(dstTensor[i * MAX_REPEAT * layoutDst.stride(0)], + srcTensor[i * MAX_REPEAT * layoutSrc.stride(0)], dataCopyParams); + } + } else { + for (uint32_t i = 0; i < rows; ++i) { + DataCopy(dstTensor[i * layoutDst.stride(0)], srcTensor[i * layoutSrc.stride(0)], cols); + } + } + } +}; + +///////////////////////////////////////////TileCopyTla////////////////////////////////////////////////////// +/// Partial specialization for CopyGmToL1, AtlasA2, RowMajor in and zN out. +template +struct TileCopyTla, LayoutSrc, CoordSrc, AscendC::TPosition::GM>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::A1>, + std::enable_if_t::value && + tla::detail::iszN::value>> { + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + TileCopyTla() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert(tla::detail::isRowMajor::value && + tla::detail::iszN::value && + TensorSrc::position == AscendC::TPosition::GM && + TensorDst::position == AscendC::TPosition::A1, + "The input parameters do not match. TensorSrc must be GM and RowMajor, while TensorDst must be L1 and zN"); + + const uint32_t nValue = tla::get<0>(srcTensor.originShape()); + const uint32_t dValue = tla::get<1>(srcTensor.originShape()); + const uint32_t srcDValue = tla::get<0>(srcTensor.stride()); + const uint32_t dstInnerStrideRow = tla::get<0, 0>(dstTensor.stride()); + const uint32_t dstOuterStrideCol = tla::get<1, 1>(dstTensor.stride()); + + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = 1; + intriParams.dValue = dValue; + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzC0Stride = dstOuterStrideCol / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + + if (srcDValue < STRIDE_LIMIT) { + intriParams.nValue = nValue; + intriParams.srcDValue = srcDValue; + intriParams.dstNzNStride = dstInnerStrideRow / ELE_NUM_PER_C0; + AscendC::DataCopy(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], intriParams); + } else { + intriParams.nValue = 1; + intriParams.srcDValue = 0; + intriParams.dstNzNStride = 0; + for (uint32_t i = 0; i < nValue; i++) { + AscendC::DataCopy(dstTensor.data()[dstOffset + i * ELE_NUM_PER_C0], + srcTensor.data()[srcOffset + i * srcDValue], + intriParams); + } + } + } +}; + +/// Partial specialization for CopyGmToL1, AtlasA2, ColumnMajor in and nZ out. +template +struct TileCopyTla, LayoutSrc, CoordSrc, AscendC::TPosition::GM>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::A1>, + std::enable_if_t::value && + tla::detail::isnZ::value>> { + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + TileCopyTla() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert(tla::detail::isColumnMajor::value && + tla::detail::isnZ::value && + TensorSrc::position == AscendC::TPosition::GM && + TensorDst::position == AscendC::TPosition::A1, + "The input parameters do not match. TensorSrc must be GM and ColumnMajor, " + "while TensorDst must be L1 and nZ"); + + const uint32_t nValue = tla::get<1>(srcTensor.originShape()); + const uint32_t dValue = tla::get<0>(srcTensor.originShape()); + const uint32_t srcDValue = tla::get<1>(srcTensor.stride()); + const uint32_t dstInnerStrideRow = tla::get<1, 0>(dstTensor.stride()); + const uint32_t dstOuterStrideCol = tla::get<0, 1>(dstTensor.stride()); + + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = 1; + intriParams.dValue = dValue; + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzC0Stride = dstOuterStrideCol / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + + if (srcDValue < STRIDE_LIMIT) { + intriParams.nValue = nValue; + intriParams.srcDValue = srcDValue; + intriParams.dstNzNStride = dstInnerStrideRow / ELE_NUM_PER_C0; + AscendC::DataCopy(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], intriParams); + } else { + intriParams.nValue = 1; + intriParams.srcDValue = 0; + intriParams.dstNzNStride = 0; + for (uint32_t i = 0; i < nValue; i++) { + AscendC::DataCopy(dstTensor.data()[dstOffset + i * ELE_NUM_PER_C0], + srcTensor.data()[srcOffset + i * srcDValue], + intriParams); + } + } + } +}; + +/// Partial specialization for CopyGmToL1, AtlasA2, zN in and zN out. +template +struct TileCopyTla, LayoutSrc, CoordSrc, AscendC::TPosition::GM>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::A1>, + std::enable_if_t::value && + tla::detail::iszN::value>> { + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + TileCopyTla() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert(tla::detail::iszN::value && + tla::detail::iszN::value && + TensorSrc::position == AscendC::TPosition::GM && + TensorDst::position == AscendC::TPosition::A1, + "The input parameters do not match. TensorSrc must be GM and zN, while TensorDst must be L1 and zN"); + + const uint32_t srcOuterStrideCol = tla::get<1, 1>(srcTensor.stride()); + const uint32_t dstOuterStrideCol = tla::get<1, 1>(dstTensor.stride()); + + uint32_t blockCount = CeilDiv(tla::get<1>(srcTensor.originShape())); + uint32_t blockLen = tla::get<0>(srcTensor.originShape()); + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + + AscendC::DataCopyParams repeatParams; + if (srcOuterStrideCol / ELE_NUM_PER_C0 <= STRIDE_LIMIT) { + repeatParams.blockCount = blockCount; + repeatParams.blockLen = blockLen; + repeatParams.srcStride = srcOuterStrideCol / ELE_NUM_PER_C0 - blockLen; + repeatParams.dstStride = dstOuterStrideCol / ELE_NUM_PER_C0 - blockLen; + AscendC::DataCopy(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], repeatParams); + } else { + repeatParams.blockCount = 1; + repeatParams.blockLen = blockLen; + repeatParams.srcStride = 0; + repeatParams.dstStride = 0; + for (uint32_t i = 0; i < blockCount; i++) { + AscendC::DataCopy(dstTensor.data()[dstOffset + i * dstOuterStrideCol], + srcTensor.data()[srcOffset + i * srcOuterStrideCol], + repeatParams); + } + } + } +}; + +/// Partial specialization for CopyGmToL1, AtlasA2, nZ in and nZ out. +template +struct TileCopyTla, LayoutSrc, CoordSrc, AscendC::TPosition::GM>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::A1>, + std::enable_if_t::value && + tla::detail::isnZ::value>> { + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + TileCopyTla() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert(tla::detail::isnZ::value && + tla::detail::isnZ::value && + TensorSrc::position == AscendC::TPosition::GM && + TensorDst::position == AscendC::TPosition::A1, + "The input parameters do not match. TensorSrc must be GM and nZ, while TensorDst must be L1 and nZ"); + + const uint32_t srcOuterStrideRow = tla::get<0, 1>(srcTensor.stride()); + const uint32_t dstOuterStrideRow = tla::get<0, 1>(dstTensor.stride()); + + uint32_t blockCount = CeilDiv(tla::get<0>(srcTensor.originShape())); + uint32_t blockLen = tla::get<1>(srcTensor.originShape()); + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + + AscendC::DataCopyParams repeatParams; + if (srcOuterStrideRow / ELE_NUM_PER_C0 <= STRIDE_LIMIT) { + repeatParams.blockCount = blockCount; + repeatParams.blockLen = blockLen; + repeatParams.srcStride = srcOuterStrideRow / ELE_NUM_PER_C0 - blockLen; + repeatParams.dstStride = dstOuterStrideRow / ELE_NUM_PER_C0 - blockLen; + AscendC::DataCopy(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], repeatParams); + } else { + repeatParams.blockCount = 1; + repeatParams.blockLen = blockLen; + repeatParams.srcStride = 0; + repeatParams.dstStride = 0; + for (uint32_t i = 0; i < blockCount; i++) { + AscendC::DataCopy(dstTensor.data()[dstOffset + i * dstOuterStrideRow], + srcTensor.data()[srcOffset + i * srcOuterStrideRow], + repeatParams); + } + } + } +}; + +///////////////////////////////////////////TileCopyTlaExt////////////////////////////////////////////////////// +/// Partial specialization for TileCopyTlaExt, CopyGmToL1, AtlasA2, PaddingRowMajor in and zN out. +template +struct TileCopyTlaExt, LayoutSrc, CoordSrc, AscendC::TPosition::GM>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::A1>, + layout::RowMajor, layout::zN> { + using ActualShape = tla::Shape; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + TileCopyTlaExt() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, ActualShape actualShape) + { + static_assert(tla::detail::isRowMajor::value && + tla::detail::iszN::value && + TensorSrc::position == AscendC::TPosition::GM && + TensorDst::position == AscendC::TPosition::A1, + "The input parameters do not match. TensorSrc must be GM and RowMajor, while TensorDst must be L1 and zN"); + + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = 1; + intriParams.dValue = tla::get<1>(actualShape); + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzC0Stride = tla::get<1, 1>(dstTensor.stride()) / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + + intriParams.nValue = tla::get<0>(actualShape); + intriParams.srcDValue = tla::get<0>(srcTensor.stride()); + intriParams.dstNzNStride = tla::get<0, 0>(dstTensor.stride()) / ELE_NUM_PER_C0; + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + AscendC::DataCopy(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], intriParams); + } +}; + +/// Partial specialization for TileCopyTlaExt, CopyGmToL1, AtlasA2, PaddingRowMajor in and zN out. +template +struct TileCopyTlaExt, LayoutSrc, CoordSrc, AscendC::TPosition::GM>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::A1>, + layout::PaddingRowMajor, layout::zN> { + using ActualShape = tla::Shape; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + TileCopyTlaExt() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, ActualShape actualShape) + { + static_assert(tla::detail::iszN::value && + TensorSrc::position == AscendC::TPosition::GM && + TensorDst::position == AscendC::TPosition::A1, + "The input parameters do not match. TensorSrc must be GM and PaddingRowMajor, " + "while TensorDst must be L1 and zN"); + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = 1; + intriParams.dValue = tla::get<1>(actualShape); + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzC0Stride = tla::get<1, 1>(dstTensor.stride()) / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + + intriParams.nValue = tla::get<0>(actualShape); + intriParams.srcDValue = tla::get<0, 0>(srcTensor.stride()); + intriParams.dstNzNStride = tla::get<0, 0>(dstTensor.stride()) / ELE_NUM_PER_C0; + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + AscendC::DataCopy(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], intriParams); + } +}; + +/// Partial specialization for TileCopyTlaExt, CopyGmToL1, AtlasA2, ColumnMajor in and nZ out. +template +struct TileCopyTlaExt, LayoutSrc, CoordSrc, AscendC::TPosition::GM>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::A1>, + layout::ColumnMajor, layout::nZ> { + using ActualShape = tla::Shape; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + TileCopyTlaExt() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, ActualShape actualShape) + { + static_assert(tla::detail::isColumnMajor::value && + tla::detail::isnZ::value && + TensorSrc::position == AscendC::TPosition::GM && + TensorDst::position == AscendC::TPosition::A1, + "The input parameters do not match. TensorSrc must be GM and ColumnMajor, " + "while TensorDst must be L1 and nZ"); + + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = 1; + intriParams.dValue = tla::get<0>(actualShape); + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzC0Stride = tla::get<0, 1>(dstTensor.stride()) / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + + intriParams.nValue = tla::get<1>(actualShape); + intriParams.srcDValue = tla::get<1>(srcTensor.stride()); + intriParams.dstNzNStride = tla::get<1, 0>(dstTensor.stride()) / ELE_NUM_PER_C0; + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + AscendC::DataCopy(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], intriParams); + } +}; + +/// Partial specialization for TileCopyTlaExt, CopyGmToL1, AtlasA2, PaddingColumnMajor in and nZ out. +template +struct TileCopyTlaExt, LayoutSrc, CoordSrc, AscendC::TPosition::GM>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::A1>, + layout::PaddingColumnMajor, layout::nZ> { + using ActualShape = tla::Shape; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + TileCopyTlaExt() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, ActualShape actualShape) + { + static_assert(tla::detail::isnZ::value && + TensorSrc::position == AscendC::TPosition::GM && + TensorDst::position == AscendC::TPosition::A1, + "The input parameters do not match. TensorSrc must be GM and PaddingColumnMajor, " + "while TensorDst must be L1 and nZ"); + + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = 1; + intriParams.dValue = tla::get<0>(actualShape); + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzC0Stride = tla::get<0, 1>(dstTensor.stride()) / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + + intriParams.nValue = tla::get<1>(actualShape); + intriParams.srcDValue = tla::get<1, 0>(srcTensor.stride()); + intriParams.dstNzNStride = tla::get<1, 0>(dstTensor.stride()) / ELE_NUM_PER_C0; + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + AscendC::DataCopy(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], intriParams); + } +}; + +/// Partial specialization for TileCopyTlaExt, CopyGmToL1, AtlasA2, zN in and zN out. +template +struct TileCopyTlaExt, LayoutSrc, CoordSrc, AscendC::TPosition::GM>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::A1>, + layout::zN, layout::zN> { + using ActualShape = tla::Shape; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + TileCopyTlaExt() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, ActualShape actualShape) + { + static_assert(tla::detail::iszN::value && + tla::detail::iszN::value && + TensorSrc::position == AscendC::TPosition::GM && + TensorDst::position == AscendC::TPosition::A1, + "The input parameters do not match. TensorSrc must be GM and zN, while TensorDst must be L1 and zN"); + + const uint32_t srcOuterStrideCol = tla::get<1, 1>(srcTensor.stride()); + const uint32_t dstOuterStrideCol = tla::get<1, 1>(dstTensor.stride()); + + uint32_t blockCount = CeilDiv(tla::get<1>(actualShape)); + uint32_t blockLen = tla::get<0>(actualShape); + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + + AscendC::DataCopyParams repeatParams; + if (srcOuterStrideCol / ELE_NUM_PER_C0 <= STRIDE_LIMIT) { + repeatParams.blockCount = blockCount; + repeatParams.blockLen = blockLen; + repeatParams.srcStride = srcOuterStrideCol / ELE_NUM_PER_C0 - blockLen; + repeatParams.dstStride = dstOuterStrideCol / ELE_NUM_PER_C0 - blockLen; + AscendC::DataCopy(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], repeatParams); + } else { + repeatParams.blockCount = 1; + repeatParams.blockLen = blockLen; + repeatParams.srcStride = 0; + repeatParams.dstStride = 0; + for (uint32_t i = 0; i < blockCount; i++) { + AscendC::DataCopy(dstTensor.data()[dstOffset + i * dstOuterStrideCol], + srcTensor.data()[srcOffset + i * srcOuterStrideCol], + repeatParams); + } + } + } +}; + +/// Partial specialization for TileCopyTlaExt, CopyGmToL1, AtlasA2, nZ in and nZ out. +template +struct TileCopyTlaExt, LayoutSrc, CoordSrc, AscendC::TPosition::GM>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::A1>, + layout::nZ, layout::nZ> { + using ActualShape = tla::Shape; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + TileCopyTlaExt() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, ActualShape actualShape) + { + static_assert(tla::detail::isnZ::value && + tla::detail::isnZ::value && + TensorSrc::position == AscendC::TPosition::GM && + TensorDst::position == AscendC::TPosition::A1, + "The input parameters do not match. TensorSrc must be GM and nZ, while TensorDst must be L1 and nZ"); + + const uint32_t srcOuterStrideRow = tla::get<0, 1>(srcTensor.stride()); + const uint32_t dstOuterStrideRow = tla::get<0, 1>(dstTensor.stride()); + + uint32_t blockCount = CeilDiv(tla::get<0>(actualShape)); + uint32_t blockLen = tla::get<1>(actualShape); + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + + AscendC::DataCopyParams repeatParams; + if (srcOuterStrideRow / ELE_NUM_PER_C0 <= STRIDE_LIMIT) { + repeatParams.blockCount = blockCount; + repeatParams.blockLen = blockLen; + repeatParams.srcStride = srcOuterStrideRow / ELE_NUM_PER_C0 - blockLen; + repeatParams.dstStride = dstOuterStrideRow / ELE_NUM_PER_C0 - blockLen; + AscendC::DataCopy(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], repeatParams); + } else { + repeatParams.blockCount = 1; + repeatParams.blockLen = blockLen; + repeatParams.srcStride = 0; + repeatParams.dstStride = 0; + for (uint32_t i = 0; i < blockCount; i++) { + AscendC::DataCopy(dstTensor.data()[dstOffset + i * dstOuterStrideRow], + srcTensor.data()[srcOffset + i * srcOuterStrideRow], + repeatParams); + } + } + } +}; + +template +struct CopyGmToL1, + Gemm::GemmType> { + using LayoutDst = layout::VectorLayout; + using LayoutSrc = layout::VectorLayout; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyGmToL1() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::DataCopyParams intriParams; + intriParams.blockCount = 1; + intriParams.blockLen = layoutDst.shape(0) / ELE_NUM_PER_C0; + intriParams.srcStride = 0; + intriParams.dstStride = 0; + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } +}; + +/// Partial specialization for CopyGmToL1, AtlasA2, RowMajor or ColumnMajor in and zN out. +template +struct TileCopySparseTla, LayoutSrc, CoordSrc, AscendC::TPosition::GM>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::A1>, + std::enable_if_t<(tla::detail::isRowMajor::value || + tla::detail::isColumnMajor::value) && + tla::detail::iszN::value>> { + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(ElementSrc); + + // Methods + + CATLASS_DEVICE + TileCopySparseTla() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert((tla::detail::isRowMajor::value || + tla::detail::isColumnMajor::value) && + tla::detail::iszN::value && + TensorSrc::position == AscendC::TPosition::GM && + TensorDst::position == AscendC::TPosition::A1, + "The input parameters do not match. TensorSrc must be GM and RowMajor or ColumnMajor, while TensorDst must be L1 and zN"); + + auto srcShape = srcTensor.shape(); + auto srcStride = srcTensor.stride(); + auto dstShape = dstTensor.shape(); + auto coord = srcTensor.coord(); + + AscendC::Nd2NzParams nd2nzParams; + nd2nzParams.ndNum = 1; + nd2nzParams.nValue = + tla::min(tla::get<0, 0>(dstShape) * tla::get<0, 1>(dstShape), tla::get<0>(srcShape) - tla::get<0>(coord)); + nd2nzParams.dValue = + tla::min(tla::get<1, 0>(dstShape) * tla::get<1, 1>(dstShape), tla::get<1>(srcShape) - tla::get<1>(coord)); + nd2nzParams.srcNdMatrixStride = 0; + nd2nzParams.srcDValue = tla::get<0>(srcStride); + nd2nzParams.dstNzC0Stride = tla::get<0, 0>(dstShape) * tla::get<0, 1>(dstShape); + nd2nzParams.dstNzNStride = 1; + nd2nzParams.dstNzMatrixStride = 0; + + auto offset = srcTensor.layout()(coord); + + AscendC::DataCopy(dstTensor.data(), srcTensor.data()[offset], nd2nzParams); + } +}; + +template +struct TileCopySparseTla, LayoutSrc, CoordSrc, AscendC::TPosition::GM>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::A1>, + std::enable_if_t::value && + tla::detail::isnZ::value>> { + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(ElementSrc); + + // Methods + + CATLASS_DEVICE + TileCopySparseTla() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert((tla::detail::isColumnMajor::value) && + tla::detail::isnZ::value && + TensorSrc::position == AscendC::TPosition::GM && + TensorDst::position == AscendC::TPosition::A1, + "The input parameters do not match. TensorSrc must be GM and ColumnMajor, while TensorDst must be L1 and nZ"); + const uint32_t nValue = tla::get<1>(srcTensor.shape()); + const uint32_t dValue = tla::get<0>(srcTensor.shape()); + const uint32_t srcDValue = tla::get<1>(srcTensor.stride()); + const uint32_t dstInnerStrideRow = tla::get<1, 0>(dstTensor.stride()); + const uint32_t dstOuterStrideCol = tla::get<0, 1>(dstTensor.stride()); + + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = 1; + intriParams.dValue = dValue; + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzC0Stride = dstOuterStrideCol / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + + if (srcDValue < STRIDE_LIMIT) { + intriParams.nValue = nValue; + intriParams.srcDValue = srcDValue; + intriParams.dstNzNStride = dstInnerStrideRow / ELE_NUM_PER_C0; + AscendC::DataCopy(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], intriParams); + } else { + intriParams.nValue = 1; + intriParams.srcDValue = 0; + intriParams.dstNzNStride = 0; + for (uint32_t i = 0; i < nValue; i++) { + AscendC::DataCopy(dstTensor.data()[dstOffset + i * ELE_NUM_PER_C0], + srcTensor.data()[srcOffset + i * srcDValue], + intriParams); + } + } + } +}; + +/// Partial specialization for CopyGmToL1, AtlasA2, sparse zn in and zN out. +template +struct TileCopySparseTla, LayoutSrc, CoordSrc, AscendC::TPosition::GM>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::A1>, + std::enable_if_t::value && + tla::detail::iszN::value>> { + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(ElementSrc); + + // Methods + + CATLASS_DEVICE + TileCopySparseTla() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert(tla::detail::iszN::value && + tla::detail::iszN::value && + TensorSrc::position == AscendC::TPosition::GM && + TensorDst::position == AscendC::TPosition::A1, + "The input parameters do not match. TensorSrc must be GM and zN, while TensorDst must be L1 and zN"); + + using TransT = typename TensorSrc::Element; + + auto srcShape = srcTensor.shape(); + auto dstShape = dstTensor.shape(); + auto coord = srcTensor.coord(); + int32_t alignedGRow = tla::get<0, 0>(srcShape) * tla::get<0, 1>(srcShape); + // height direction need to be 16 aligned + auto alignHeight = tla::get<0, 0>(dstShape) * tla::get<0, 1>(dstShape); + int32_t blockLen = alignHeight * tla::get<1, 0>(dstShape) * sizeof(TransT) / BYTE_PER_C0; + int32_t srcStride = (alignedGRow - alignHeight) * tla::get<1, 0>(srcShape) * sizeof(TransT) / BYTE_PER_C0; + uint16_t nburst = tla::get<1, 1>(dstShape); + auto offset = srcTensor.layout()(coord); + AscendC::DataCopy( + dstTensor.data(), srcTensor.data()[offset], + {nburst, static_cast(blockLen), static_cast(srcStride), static_cast(0)}); + } +}; + +/// Partial specialization for CopyGmToL1, AtlasA2, sparse nZ in and nZ out. +template +struct TileCopySparseTla, LayoutSrc, CoordSrc, AscendC::TPosition::GM>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::A1>, + std::enable_if_t::value && + tla::detail::isnZ::value>> { + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(ElementSrc); + + // Methods + + CATLASS_DEVICE + TileCopySparseTla() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert(tla::detail::isnZ::value && + tla::detail::isnZ::value && + TensorSrc::position == AscendC::TPosition::GM && + TensorDst::position == AscendC::TPosition::A1, + "The input parameters do not match. TensorSrc must be GM and nZ, while TensorDst must be L1 and nZ"); + + using TransT = typename TensorSrc::Element; + auto srcShape = srcTensor.shape(); + auto dstShape = dstTensor.shape(); + auto coord = srcTensor.coord(); + int32_t alignedGCol = tla::get<1, 0>(srcShape) * tla::get<1, 1>(srcShape); + auto alignCol = tla::get<1, 0>(dstShape) * tla::get<1, 1>(dstShape); + int32_t blockLen = alignCol * tla::get<0, 0>(dstShape) * sizeof(TransT) / BYTE_PER_C0; + int32_t srcStride = (alignedGCol - alignCol) * tla::get<0, 0>(srcShape) * sizeof(TransT) / BYTE_PER_C0; + uint16_t nburst = tla::get<0, 1>(dstShape); + auto offset = srcTensor.layout()(coord); + AscendC::DataCopy( + dstTensor.data(), srcTensor.data()[offset], + {nburst, static_cast(blockLen), static_cast(srcStride), static_cast(0)}); + } +}; + +///////////////////////////////////////////TileCopyFAQTla////////////////////////////////////////////////////// +/// used for FA LoadQ, multi matrix datacopy ND->NZ +template +struct TileCopyFAQTla, LayoutSrc, CoordSrc, AscendC::TPosition::GM>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::A1>> { + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + TileCopyFAQTla() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert(tla::detail::iszN::value && + TensorSrc::position == AscendC::TPosition::GM && + TensorDst::position == AscendC::TPosition::A1, + "The input parameters do not match. TensorSrc must be GM, while TensorDst must be L1 and zN"); + + const uint32_t nValue = tla::get<1>(srcTensor.shape()); + const uint32_t dValue = tla::get<2>(srcTensor.shape()); + const uint32_t srcDValue = tla::get<1>(srcTensor.stride()); + const uint32_t ndNum = tla::get<0>(srcTensor.shape()); + const uint32_t srcNdMatrixStride = tla::get<0>(srcTensor.stride()); + const uint32_t dstNzNStride = tla::get<0>(srcTensor.shape()); + const uint32_t dstNzMatrixStride = ELE_NUM_PER_C0; + const uint32_t dstNzC0Stride = tla::get<0, 0>(dstTensor.shape()) * tla::get<0, 1>(dstTensor.shape()); + const uint32_t dstInnerStrideRow = tla::get<0, 0>(dstTensor.stride()); + const uint32_t dstOuterStrideCol = tla::get<1, 1>(dstTensor.stride()); + + AscendC::Nd2NzParams intriParams; + + intriParams.nValue = nValue; + intriParams.dValue = dValue; + intriParams.srcDValue = srcDValue; + intriParams.dstNzNStride = dstNzNStride; + intriParams.dstNzC0Stride = dstNzC0Stride; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + + if (srcNdMatrixStride < STRIDE_LIMIT) { + intriParams.ndNum = ndNum; + intriParams.srcNdMatrixStride = srcNdMatrixStride; + intriParams.dstNzMatrixStride = dstNzMatrixStride; + AscendC::DataCopy(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], intriParams); + } else { + intriParams.ndNum = 1; + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzMatrixStride = 0; + for (uint32_t i = 0; i < nValue; i++) { + AscendC::DataCopy(dstTensor.data()[dstOffset + i * ELE_NUM_PER_C0], + srcTensor.data()[srcOffset + i * srcDValue], + intriParams); + } + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace Catlass::Gemm::Tile + +#endif // CATLASS_KDA_ATLASA2_COPY_GM_TO_L1_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_atlasa2_copy_gm_to_ub.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_atlasa2_copy_gm_to_ub.hpp new file mode 100644 index 0000000000..35b1fb4002 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_atlasa2_copy_gm_to_ub.hpp @@ -0,0 +1,99 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_ATLASA2_COPY_GM_TO_UB_HPP +#define CATLASS_KDA_ATLASA2_COPY_GM_TO_UB_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_arch.hpp" +#include "catlass/gemm/tile/kda_tile_copy_tla.hpp" +#include "catlass/gemm/kda_gemm_type.hpp" +#include "tla/kda_tensor.hpp" + +namespace Catlass::Gemm::Tile { + +/// Partial specialization for AtlasA2, RowMajor in and RowMajor out. +template +struct TileCopyTla, LayoutSrc, CoordSrc, AscendC::TPosition::GM>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::VECCALC>, + std::enable_if_t::value && + tla::detail::isRowMajor::value>> { + static constexpr uint32_t ELE_NUM_PER_BLK = BYTE_PER_BLK / sizeof(ElementSrc); + + // Methods + + CATLASS_DEVICE + TileCopyTla() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert(tla::detail::isRowMajor::value && + tla::detail::isRowMajor::value && + TensorSrc::position == AscendC::TPosition::GM && + TensorDst::position == AscendC::TPosition::VECCALC, + "The input parameters do not match. TensorSrc must be GM and RowMajor, " + "while TensorDst must be UB and RowMajor"); + + const uint16_t row = tla::get<0>(srcTensor.originShape()); + const uint16_t col = tla::get<1>(srcTensor.originShape()); + + AscendC::DataCopyExtParams dataCopyParams( + row, + col * sizeof(ElementSrc), + (tla::get<0>(srcTensor.stride()) - col) * sizeof(ElementSrc), + (tla::get<0>(dstTensor.stride()) - col) / ELE_NUM_PER_BLK, + 0); + AscendC::DataCopyPadExtParams padParams(false, 0, 0, 0); + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + AscendC::DataCopyPad(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], dataCopyParams, padParams); + }; +}; + +template < + class ArchTag, + class GmType> +struct CopyGm2Ub { + static_assert(DEPENDENT_FALSE, "Unsupported copy gm to ub, can not find the specialization."); +}; + +template +struct CopyGm2Ub> { + using LayoutSrc = layout::VectorLayout; + using LayoutDst = layout::VectorLayout; + + static constexpr uint32_t ELE_NUM_PER_BLK = BYTE_PER_BLK / sizeof(Element); + + CATLASS_DEVICE + CopyGm2Ub() = default; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + layout::VectorLayout const &layoutDst, + layout::VectorLayout const &layoutSrc) + { + AscendC::DataCopyExtParams dataCopyParams( + 1, + layoutSrc.shape(0) * sizeof(Element), + 0, + 0, + 0); + AscendC::DataCopyPadExtParams padParams(false, 0, 0, 0); + AscendC::DataCopyPad(dstTensor, srcTensor, dataCopyParams, padParams); + }; +}; + +} // namespace Catlass::Gemm::Tile + +#endif // CATLASS_KDA_ATLASA2_COPY_GM_TO_UB_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_atlasa2_copy_l0c_to_gm.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_atlasa2_copy_l0c_to_gm.hpp new file mode 100644 index 0000000000..baa55940e6 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_atlasa2_copy_l0c_to_gm.hpp @@ -0,0 +1,632 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_ATLASA2_COPY_L0C_TO_GM_HPP +#define CATLASS_KDA_ATLASA2_COPY_L0C_TO_GM_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_arch.hpp" +#include "catlass/gemm/kda_gemm_type.hpp" +#include "tla/kda_tensor.hpp" + +namespace Catlass::Gemm::Tile { + +enum class ScaleGranularity { + UNDEFINED = -1, + NO_QUANT = 0, + PER_TENSOR, + PER_CHANNEL, + PER_GROUP +}; + +template < + class ArchTag, + class ElementSrc, + class ElementDst, + ScaleGranularity DEQUANT_GRANULARITY = ScaleGranularity::NO_QUANT> +struct CopyL0CToGmQuantMode { + static_assert(DEPENDENT_FALSE, "Unsupported copy l0c to gm, can not find the specialization."); +}; + +// CopyL0CToGm fp32 to fp32 +template <> +struct CopyL0CToGmQuantMode< + Catlass::Arch::AtlasA2, + float, float, + ScaleGranularity::NO_QUANT> { + static constexpr auto VALUE = QuantMode_t::NoQuant; +}; + +// CopyL0CToGm cast fp32 to fp16 +template <> +struct CopyL0CToGmQuantMode< + Catlass::Arch::AtlasA2, + float, half, + ScaleGranularity::NO_QUANT> { + static constexpr auto VALUE = QuantMode_t::F322F16; +}; + +template <> +struct CopyL0CToGmQuantMode< + Catlass::Arch::AtlasA2, + float, half, + ScaleGranularity::PER_TENSOR> { + static constexpr auto VALUE = QuantMode_t::F322F16; +}; + +// CopyL0CToGm cast fp32 to bf16 +template <> +struct CopyL0CToGmQuantMode< + Catlass::Arch::AtlasA2, + float, bfloat16_t, + ScaleGranularity::NO_QUANT> { + static constexpr auto VALUE = QuantMode_t::F322BF16; +}; + +template <> +struct CopyL0CToGmQuantMode< + Catlass::Arch::AtlasA2, + float, bfloat16_t, + ScaleGranularity::PER_TENSOR> { + static constexpr auto VALUE = QuantMode_t::F322BF16; +}; + +// CopyL0CToGm cast float to uint8/int8 +template <> +struct CopyL0CToGmQuantMode< + Catlass::Arch::AtlasA2, + float, uint8_t, + ScaleGranularity::PER_TENSOR> { + static constexpr auto VALUE = QuantMode_t::QF322B8_PRE; +}; + +template <> +struct CopyL0CToGmQuantMode< + Catlass::Arch::AtlasA2, + float, uint8_t, + ScaleGranularity::PER_CHANNEL> { + static constexpr auto VALUE = QuantMode_t::VQF322B8_PRE; +}; + +template <> +struct CopyL0CToGmQuantMode< + Catlass::Arch::AtlasA2, + float, int8_t, + ScaleGranularity::PER_TENSOR> { + static constexpr auto VALUE = QuantMode_t::QF322B8_PRE; +}; + +template <> +struct CopyL0CToGmQuantMode< + Catlass::Arch::AtlasA2, + float, int8_t, + ScaleGranularity::PER_CHANNEL> { + static constexpr auto VALUE = QuantMode_t::VQF322B8_PRE; +}; + +// CopyL0CToGm output int32 +template <> +struct CopyL0CToGmQuantMode< + Catlass::Arch::AtlasA2, + int32_t, int32_t, + ScaleGranularity::NO_QUANT> { + static constexpr auto VALUE = QuantMode_t::NoQuant; +}; + +// CopyL0CToGm cast int32_t to fp16 +template <> +struct CopyL0CToGmQuantMode< + Catlass::Arch::AtlasA2, + int32_t, half, + ScaleGranularity::PER_TENSOR> { + static constexpr auto VALUE = QuantMode_t::DEQF16; +}; + +template <> +struct CopyL0CToGmQuantMode< + Catlass::Arch::AtlasA2, + int32_t, half, + ScaleGranularity::PER_CHANNEL> { + static constexpr auto VALUE = QuantMode_t::VDEQF16; +}; + +// CopyL0CToGm cast int32 to uint8/int8 +template <> +struct CopyL0CToGmQuantMode< + Catlass::Arch::AtlasA2, + int32_t, uint8_t, + ScaleGranularity::PER_TENSOR> { + static constexpr auto VALUE = QuantMode_t::REQ8; +}; + +template <> +struct CopyL0CToGmQuantMode< + Catlass::Arch::AtlasA2, + int32_t, uint8_t, + ScaleGranularity::PER_CHANNEL> { + static constexpr auto VALUE = QuantMode_t::VREQ8; +}; + +template <> +struct CopyL0CToGmQuantMode< + Catlass::Arch::AtlasA2, + int32_t, int8_t, + ScaleGranularity::PER_TENSOR> { + static constexpr auto VALUE = QuantMode_t::REQ8; +}; + +template <> +struct CopyL0CToGmQuantMode< + Catlass::Arch::AtlasA2, + int32_t, int8_t, + ScaleGranularity::PER_CHANNEL> { + static constexpr auto VALUE = QuantMode_t::VREQ8; +}; + +template < + class ArchTag, + class ElementAccumulator, + class GmType, + ScaleGranularity DEQUANT_GRANULARITY = ScaleGranularity::NO_QUANT, + bool ReluEnable = false> +struct CopyL0CToGm { + static_assert(DEPENDENT_FALSE, "Unsupported copy l0c to gm, can not find the specialization."); +}; + +template < + class ElementAccumulator_, + class ElementDst_, + bool ReluEnable_> +struct CopyL0CToGm, + ScaleGranularity::NO_QUANT, + ReluEnable_> { + using ArchTag = Catlass::Arch::AtlasA2; + using ElementDst = ElementDst_; + using ElementSrc = ElementAccumulator_; + using LayoutSrc = Catlass::layout::zN; + using LayoutDst = Catlass::layout::RowMajor; + static constexpr auto quantPre = CopyL0CToGmQuantMode::VALUE; + static constexpr auto reluEn = ReluEnable_; + + struct Params {}; + Params params; + + CATLASS_DEVICE + CopyL0CToGm() = default; + + CATLASS_DEVICE + CopyL0CToGm(Params const ¶ms_) + : params(params_) {}; + + CATLASS_DEVICE + void operator()(AscendC::GlobalTensor const &dst, AscendC::LocalTensor const &src, + LayoutDst const &dstLayout, LayoutSrc const &srcLayout, uint8_t unitFlag = 0) + { + AscendC::FixpipeParamsV220 intriParams; + + // Fixpipe layout information + intriParams.nSize = dstLayout.shape(1); + intriParams.mSize = dstLayout.shape(0); + intriParams.srcStride = srcLayout.stride(3) / srcLayout.stride(0); + intriParams.dstStride = dstLayout.stride(0); + + // Fixpipe auxiliary arguments + intriParams.quantPre = quantPre; + intriParams.reluEn = reluEn; + intriParams.unitFlag = unitFlag; + + // Call AscendC Fixpipe + AscendC::Fixpipe(dst, src, intriParams); + } +}; + +template < + class ElementAccumulator_, + class ElementDst_, + bool ReluEnable_> +struct CopyL0CToGm, + ScaleGranularity::PER_TENSOR, + ReluEnable_> { + using ArchTag = Catlass::Arch::AtlasA2; + using ElementDst = ElementDst_; + using ElementSrc = ElementAccumulator_; + using LayoutSrc = Catlass::layout::zN; + using LayoutDst = Catlass::layout::RowMajor; + static constexpr auto quantPre = CopyL0CToGmQuantMode::VALUE; + static constexpr auto reluEn = ReluEnable_; + + struct Params { + float scale = 1.0; + + CATLASS_HOST_DEVICE + Params() = default; + + CATLASS_HOST_DEVICE + Params(float scalar) + { + scale = scalar; + } + }; + Params params; + + CATLASS_DEVICE + CopyL0CToGm() = default; + + CATLASS_DEVICE + CopyL0CToGm(Params const ¶ms_) + : params(params_) {}; + + CATLASS_DEVICE + void operator()(AscendC::GlobalTensor const &dst, AscendC::LocalTensor const &src, + LayoutDst const &dstLayout, LayoutSrc const &srcLayout, uint8_t unitFlag = 0) + { + AscendC::FixpipeParamsV220 intriParams; + + // Fixpipe layout information + intriParams.nSize = dstLayout.shape(1); + intriParams.mSize = dstLayout.shape(0); + intriParams.srcStride = srcLayout.stride(3) / srcLayout.stride(0); + intriParams.dstStride = dstLayout.stride(0); + + // Fixpipe auxiliary arguments + intriParams.quantPre = quantPre; + intriParams.deqScalar = static_cast(*reinterpret_cast(¶ms.scale)); + intriParams.reluEn = reluEn; + intriParams.unitFlag = unitFlag; + + // Call AscendC Fixpipe + AscendC::Fixpipe(dst, src, intriParams); + } +}; + +template < + class ElementAccumulator_, + class ElementDst_, + bool ReluEnable_> +struct CopyL0CToGm, + ScaleGranularity::PER_CHANNEL, + ReluEnable_> { + using ArchTag = Catlass::Arch::AtlasA2; + using ElementDst = ElementDst_; + using ElementSrc = ElementAccumulator_; + using LayoutSrc = Catlass::layout::zN; + using LayoutDst = Catlass::layout::RowMajor; + static constexpr auto quantPre = CopyL0CToGmQuantMode::VALUE; + static constexpr auto reluEn = ReluEnable_; + + struct Params {}; + Params params; + + CATLASS_DEVICE + CopyL0CToGm() = default; + + CATLASS_DEVICE + CopyL0CToGm(Params const ¶ms_) + : params(params_) {}; + + CATLASS_DEVICE + void operator()(AscendC::GlobalTensor const &dst, AscendC::LocalTensor const &src, AscendC::LocalTensor const &scale, + LayoutDst const &dstLayout, LayoutSrc const &srcLayout, uint8_t unitFlag = 0) + { + AscendC::FixpipeParamsV220 intriParams; + + // Fixpipe layout information + intriParams.nSize = dstLayout.shape(1); + intriParams.mSize = dstLayout.shape(0); + intriParams.srcStride = srcLayout.stride(3) / srcLayout.stride(0); + intriParams.dstStride = dstLayout.stride(0); + + // Fixpipe auxiliary arguments + intriParams.quantPre = quantPre; + intriParams.reluEn = reluEn; + intriParams.unitFlag = unitFlag; + + // Call AscendC Fixpipe + AscendC::SetFixPipeConfig(scale, false); + AscendC::PipeBarrier(); + AscendC::Fixpipe(dst, src, intriParams); + } +}; + +template < + class ElementAccumulator_, + class ElementDst_, + bool ReluEnable_> +struct CopyL0CToGm, + ScaleGranularity::NO_QUANT, + ReluEnable_> { + using ArchTag = Catlass::Arch::AtlasA2; + using ElementDst = ElementDst_; + using ElementSrc = ElementAccumulator_; + using LayoutSrc = Catlass::layout::zN; + using LayoutDst = Catlass::layout::zN; + static constexpr auto quantPre = CopyL0CToGmQuantMode::VALUE; + static constexpr auto reluEn = ReluEnable_; + + struct Params {}; + Params params; + + CATLASS_DEVICE + CopyL0CToGm() = default; + + CATLASS_DEVICE + CopyL0CToGm(Params const ¶ms_) + : params(params_) {}; + + CATLASS_DEVICE + void operator()(AscendC::GlobalTensor const &dst, AscendC::LocalTensor const &src, + LayoutDst const &dstLayout, LayoutSrc const &srcLayout, uint8_t unitFlag = 0) + { + AscendC::FixpipeParamsV220 intriParams; + + // Fixpipe layout information + intriParams.nSize = dstLayout.shape(2) * dstLayout.shape(3); + intriParams.mSize = dstLayout.shape(0) * dstLayout.shape(1); + intriParams.srcStride = srcLayout.stride(3) / srcLayout.shape(2); + intriParams.dstStride = dstLayout.stride(3) / (BYTE_PER_C0 / sizeof(ElementDst)); + + // Fixpipe auxiliary arguments + intriParams.quantPre = quantPre; + intriParams.reluEn = reluEn; + intriParams.unitFlag = unitFlag; + + if constexpr (std::is_same_v && std::is_same_v) { + intriParams.isChannelSplit = true; + } + + // Call AscendC Fixpipe + AscendC::Fixpipe(dst, src, intriParams); + } +}; + +template < + class ElementAccumulator_, + class ElementDst_, + bool ReluEnable_> +struct CopyL0CToGm, + ScaleGranularity::NO_QUANT, + ReluEnable_> { + using ArchTag = Catlass::Arch::AtlasA2; + using ElementDst = ElementDst_; + using ElementSrc = ElementAccumulator_; + using LayoutSrc = Catlass::layout::zN; + using LayoutDst = Catlass::layout::NDC1HWC0; + static constexpr auto quantPre = CopyL0CToGmQuantMode::VALUE; + static constexpr auto reluEn = ReluEnable_; + + struct Params {}; + Params params; + + CATLASS_DEVICE + CopyL0CToGm() = default; + + CATLASS_DEVICE + CopyL0CToGm(Params const ¶ms_) + : params(params_) {}; + + CATLASS_DEVICE + void operator()(AscendC::GlobalTensor const &dst, AscendC::LocalTensor const &src, + LayoutDst const &dstLayout, LayoutSrc const &srcLayout, uint8_t unitFlag = 0) + { + AscendC::FixpipeParamsV220 intriParams; + + intriParams.nSize = srcLayout.orgShape(1); + intriParams.mSize = srcLayout.orgShape(0); + intriParams.srcStride = srcLayout.stride(3) / srcLayout.shape(2); + intriParams.dstStride = dstLayout.shape(1) * dstLayout.shape(2); + + if constexpr (AscendC::IsSameType::value && + AscendC::IsSameType::value) { + intriParams.isChannelSplit = true; + } + + intriParams.quantPre = quantPre; + intriParams.reluEn = false; + intriParams.unitFlag = unitFlag; + AscendC::Fixpipe(dst, src, intriParams); + } +}; + +///////////////////////////////////////////CopyL0CToGmTla///////////////////////////////////////////////// +// L0C copy mode +struct CopyToGM {}; +struct CopyToL1 {}; + +template < + class ArchTag, + class TensorSrc, + class TensorDst, + ScaleGranularity DEQUANT_GRANULARITY = ScaleGranularity::NO_QUANT, + bool ReluEnable = false, + class Enable = void> +struct CopyL0CToGmTla { + static_assert(DEPENDENT_FALSE, "Unsupported copy l0c to gm, can not find the specialization."); +}; + +template < + class TensorSrc_, + class ElementDst_, + class LayoutDst_, + class CoordDst_, + bool ReluEnable_> +struct CopyL0CToGmTla, LayoutDst_, CoordDst_, AscendC::TPosition::GM>, + ScaleGranularity::NO_QUANT, + ReluEnable_, + std::enable_if_t::value>> { + using ArchTag = Catlass::Arch::AtlasA2; + using ElementDst = ElementDst_; + using ElementSrc = typename TensorSrc_::Element; + static constexpr auto quantPre = CopyL0CToGmQuantMode::VALUE; + static constexpr auto reluEn = ReluEnable_; + + struct Params {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, uint8_t unitFlag = 0) + { + static_assert(tla::detail::isRowMajor::value && + TensorSrc::position == AscendC::TPosition::CO1 && + TensorDst::position == AscendC::TPosition::GM, + "The input parameters do not match. TensorSrc must be L0C, while TensorDst must be GM and RowMajor"); + + AscendC::FixpipeParamsV220 intriParams; + + // Fixpipe layout information + intriParams.nSize = tla::get<1>(dstTensor.originShape()); + intriParams.mSize = tla::get<0>(dstTensor.originShape()); + intriParams.srcStride = tla::get<1, 1>(srcTensor.stride()) / tla::get<0, 0>(srcTensor.stride()); + intriParams.dstStride = tla::get<0>(dstTensor.stride()); + + // Fixpipe auxiliary arguments + intriParams.quantPre = quantPre; + intriParams.reluEn = reluEn; + intriParams.unitFlag = unitFlag; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + + // Call AscendC Fixpipe + AscendC::Fixpipe( + dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], intriParams); + } +}; + +template < + class TensorSrc_, + class ElementDst_, + class LayoutDst_, + class CoordDst_, + bool ReluEnable_> +struct CopyL0CToGmTla, LayoutDst_, CoordDst_, AscendC::TPosition::GM>, + ScaleGranularity::NO_QUANT, + ReluEnable_, + std::enable_if_t::value>> { + using ArchTag = Catlass::Arch::AtlasA2; + using ElementDst = ElementDst_; + using ElementSrc = typename TensorSrc_::Element; + static constexpr auto quantPre = CopyL0CToGmQuantMode::VALUE; + static constexpr auto reluEn = ReluEnable_; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, uint8_t unitFlag = 0) + { + static_assert(tla::detail::iszN::value && + TensorSrc::position == AscendC::TPosition::CO1 && + TensorDst::position == AscendC::TPosition::GM, + "The input parameters do not match. TensorSrc must be L0C, while TensorDst must be GM and zN"); + + AscendC::FixpipeParamsV220 intriParams; + + // Fixpipe layout information + intriParams.nSize = tla::get<1>(dstTensor.originShape()); + intriParams.mSize = tla::get<0>(dstTensor.originShape()); + intriParams.srcStride = tla::get<1, 1>(srcTensor.stride()) / tla::get<1, 0>(srcTensor.shape()); + intriParams.dstStride = tla::get<1, 1>(dstTensor.stride()) / (BYTE_PER_C0 / sizeof(ElementDst)); + + // Fixpipe auxiliary arguments + intriParams.quantPre = quantPre; + intriParams.reluEn = reluEn; + intriParams.unitFlag = unitFlag; + + if constexpr (std::is_same_v && std::is_same_v) { + intriParams.isChannelSplit = true; + } + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + + // Call AscendC Fixpipe + AscendC::Fixpipe( + dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], intriParams); + } +}; + +template < + class ArchTag, + class TensorSrc, + class TensorDst, + ScaleGranularity DEQUANT_GRANULARITY = ScaleGranularity::NO_QUANT, + class Enable = void> +struct CopyL0CToGmSparseTla { + static_assert(DEPENDENT_FALSE, "Unsupported copy l0c to gm, can not find the specialization."); +}; + +template < + class TensorSrc_, + class ElementDst_, + class LayoutDst_, + class CoordDst_> +struct CopyL0CToGmSparseTla, LayoutDst_, CoordDst_, AscendC::TPosition::GM>, + ScaleGranularity::NO_QUANT, + std::enable_if_t::value>> { + using ArchTag = Catlass::Arch::AtlasA2; + using ElementDst = ElementDst_; + using ElementSrc = typename TensorSrc_::Element; + static constexpr auto quantPre = CopyL0CToGmQuantMode::VALUE; + + struct Params {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert(tla::detail::isRowMajor::value && + TensorSrc::position == AscendC::TPosition::CO1 && + TensorDst::position == AscendC::TPosition::GM, + "The input parameters do not match. TensorSrc must be L0C, while TensorDst must be GM and RowMajor"); + + auto dstShape = dstTensor.shape(); + auto srcShape = srcTensor.shape(); + auto dstStride = dstTensor.stride(); + auto coord = dstTensor.coord(); + + AscendC::FixpipeParamsV220 params; + params.nSize = + tla::min(tla::get<1, 0>(srcShape) * tla::get<1, 1>(srcShape), tla::get<1>(dstShape) - tla::get<1>(coord)); + params.mSize = + tla::min(tla::get<0, 0>(srcShape) * tla::get<0, 1>(srcShape), tla::get<0>(dstShape) - tla::get<0>(coord)); + params.srcStride = tla::get<0, 0>(srcShape) * tla::get<0, 1>(srcShape); + params.dstStride = tla::get<0>(dstStride); + params.quantPre = quantPre; + + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + auto dstOffset = dstTensor.layout()(coord); + AscendC::Fixpipe(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], params); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace Catlass::Gemm::Tile + +#endif // CATLASS_KDA_ATLASA2_COPY_L0C_TO_GM_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_atlasa2_copy_l1_to_bt.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_atlasa2_copy_l1_to_bt.hpp new file mode 100644 index 0000000000..ebe5bc4fbd --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_atlasa2_copy_l1_to_bt.hpp @@ -0,0 +1,62 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_ATLASA2_COPY_L1_TO_BT_HPP +#define CATLASS_KDA_ATLASA2_COPY_L1_TO_BT_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_arch.hpp" +#include "catlass/layout/kda_layout.hpp" +#include "catlass/gemm/kda_gemm_type.hpp" +#include "tla/kda_tensor.hpp" + +using namespace tla; + +namespace Catlass::Gemm::Tile { + +template < + class ArchTag, + class L1Type, + class L0Type = void> +struct CopyL1ToBT { + static_assert(DEPENDENT_FALSE, "Unsupported copy l1 to biasTable buffer, can not find the specialization."); +}; + +template +struct CopyL1ToBT, + Catlass::Gemm::GemmType> { + using LayoutDst = layout::VectorLayout; + using LayoutSrc = layout::VectorLayout; + + static constexpr uint32_t ELE_NUM_PER_C2 = BYTE_PER_C2 / sizeof(ElementSrc); + + CATLASS_DEVICE + CopyL1ToBT() {} + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor dstTensor, + AscendC::LocalTensor srcTensor, + LayoutDst layoutDst, LayoutSrc layoutSrc) + { + AscendC::DataCopyParams intriParams; + intriParams.blockCount = 1; + intriParams.blockLen = (layoutDst.shape(0) + ELE_NUM_PER_C2 - 1) / ELE_NUM_PER_C2; + intriParams.srcStride = 0; + intriParams.dstStride = 0; + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace Catlass::Gemm::Tile + +#endif // CATLASS_KDA_ATLASA2_COPY_L1_TO_BT_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_atlasa2_copy_l1_to_fp.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_atlasa2_copy_l1_to_fp.hpp new file mode 100644 index 0000000000..52bc44969a --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_atlasa2_copy_l1_to_fp.hpp @@ -0,0 +1,62 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_ATLASA2_COPY_L1_TO_FP_HPP +#define CATLASS_KDA_ATLASA2_COPY_L1_TO_FP_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_arch.hpp" +#include "catlass/layout/kda_layout.hpp" +#include "catlass/gemm/kda_gemm_type.hpp" +#include "tla/kda_tensor.hpp" + +using namespace tla; + +namespace Catlass::Gemm::Tile { + +template < + class ArchTag, + class L1Type, + class L0Type = void> +struct CopyL1ToFP { + static_assert(DEPENDENT_FALSE, "Unsupported copy l1 to fixpipe buffer, can not find the specialization."); +}; + +template +struct CopyL1ToFP, + Catlass::Gemm::GemmType> { + using LayoutDst = layout::VectorLayout; + using LayoutSrc = layout::VectorLayout; + + static constexpr uint32_t ELE_NUM_PER_FP = BYTE_PER_BLK_FP / sizeof(ElementSrc); + + CATLASS_DEVICE + CopyL1ToFP() {} + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor dstTensor, + AscendC::LocalTensor srcTensor, + LayoutDst layoutDst, LayoutSrc layoutSrc) + { + AscendC::DataCopyParams intriParams; + intriParams.blockCount = 1; + intriParams.blockLen = (layoutDst.shape(0) + ELE_NUM_PER_FP - 1) / ELE_NUM_PER_FP; + intriParams.srcStride = 0; + intriParams.dstStride = 0; + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace Catlass::Gemm::Tile + +#endif // CATLASS_KDA_ATLASA2_COPY_L1_TO_FP_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_atlasa2_copy_l1_to_l0a.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_atlasa2_copy_l1_to_l0a.hpp new file mode 100644 index 0000000000..89b0637c43 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_atlasa2_copy_l1_to_l0a.hpp @@ -0,0 +1,719 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_ATLASA2_COPY_L1_TO_L0A_HPP +#define CATLASS_KDA_ATLASA2_COPY_L1_TO_L0A_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/kda_numeric_size.hpp" +#include "catlass/arch/kda_arch.hpp" +#include "catlass/layout/kda_layout.hpp" +#include "catlass/gemm/kda_gemm_type.hpp" +#include "catlass/gemm/tile/kda_tile_copy_tla.hpp" +#include "tla/kda_tensor.hpp" + +namespace Catlass::Gemm::Tile { + +template < + class ArchTag, + class L1Type, + class L0Type = void> +struct CopyL1ToL0A { + static_assert(DEPENDENT_FALSE, "Unsupported copy l1 to l0, can not find the specialization."); +}; + +//////////////////////////////// +/// new add gemm +template +struct CopyL1ToL0A, Catlass::Gemm::GemmType> { + using LayoutDst = layout::zZ; + using LayoutSrc = layout::zN; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + CATLASS_DEVICE + CopyL1ToL0A() {} + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor dstTensor, + AscendC::LocalTensor srcTensor, + LayoutDst layoutDst, LayoutSrc layoutSrc) + { + AscendC::LoadData2DParams loadDataParams; + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = static_cast(layoutDst.shape(3)); + loadDataParams.srcStride = layoutSrc.stride(3) / ELE_NUM_PER_FRACTAL; + loadDataParams.sid = 0; + loadDataParams.dstGap = layoutDst.stride(3) / ELE_NUM_PER_FRACTAL - 1; + loadDataParams.ifTranspose = false; + loadDataParams.addrMode = 0; + + for (uint32_t i = 0; i < layoutDst.shape(1); i++) { + AscendC::LoadData(dstTensor[i * layoutDst.stride(1)], srcTensor[i * layoutSrc.stride(1)], loadDataParams); + } + } +}; + +template +struct CopyL1ToL0A, Catlass::Gemm::GemmType> { + using LayoutDst = layout::zZ; + using LayoutSrc = layout::nN; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + CATLASS_DEVICE + CopyL1ToL0A() {} + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor dstTensor, + AscendC::LocalTensor srcTensor, + LayoutDst layoutDst, LayoutSrc layoutSrc) + { + AscendC::LoadData2DParams loadDataParams; + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = static_cast(CeilDiv(layoutDst.orgShape(1))); + loadDataParams.srcStride = static_cast(CeilDiv(layoutSrc.orgShape(0))); + ; + loadDataParams.sid = 0; + loadDataParams.dstGap = 0; + loadDataParams.ifTranspose = true; + loadDataParams.addrMode = 0; + for (uint32_t i = 0; i < CeilDiv(layoutSrc.orgShape(0)); i++) { + AscendC::LoadData(dstTensor[i * layoutDst.stride(1)], srcTensor[i * layoutSrc.stride(1)], loadDataParams); + } + } +}; + +template +struct CopyL1ToL0A, Catlass::Gemm::GemmType> { + using Element = float; + using LayoutDst = layout::zZ; + using LayoutSrc = layout::nN; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + CATLASS_DEVICE + CopyL1ToL0A() {} + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor dstTensor, + AscendC::LocalTensor srcTensor, + LayoutDst layoutDst, LayoutSrc layoutSrc) + { + AscendC::LoadData2dTransposeParams loadDataParams; + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = static_cast(CeilDiv(layoutDst.orgShape(1))); + loadDataParams.srcStride = static_cast(CeilDiv(layoutSrc.orgShape(0))); + loadDataParams.dstGap = 1; + loadDataParams.dstFracGap = 0; + for (uint32_t i = 0; i < CeilDiv(layoutSrc.orgShape(0)); i++) { + AscendC::LoadDataWithTranspose(dstTensor[i * layoutDst.stride(1)], srcTensor[i * layoutSrc.stride(1) * 2], loadDataParams); + } + } +}; + +template +struct CopyL1ToL0A, Catlass::Gemm::GemmType> { + using Element = int8_t; + using LayoutDst = layout::zZ; + using LayoutSrc = layout::nZ; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + CATLASS_DEVICE + CopyL1ToL0A() {} + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor dstTensor, + AscendC::LocalTensor srcTensor, + LayoutDst layoutDst, LayoutSrc layoutSrc) + { + AscendC::LoadData2dTransposeParams loadDataParams; + + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = static_cast(CeilDiv(layoutDst.orgShape(1))); + loadDataParams.srcStride = 1; + loadDataParams.dstGap = 0; + loadDataParams.dstFracGap = CeilDiv(layoutDst.orgShape(1)) - 1; + + for (uint32_t i = 0; i < CeilDiv(layoutDst.orgShape(0)); i++) { + AscendC::LoadDataWithTranspose(dstTensor[i * layoutDst.stride(1) * 2], + srcTensor[i * layoutSrc.stride(1)], + loadDataParams); + } + } +}; + +template +struct CopyL1ToL0A> { + using LayoutDst = layout::zZ; + using LayoutSrc = layout::NDC1HWC0; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint8_t RIGHT_MOVE_8 = 8; + + // Methods + + CATLASS_DEVICE + CopyL1ToL0A(uint32_t strideW = 0, + uint32_t strideH = 0, + uint32_t filterW = 0, + uint32_t filterH = 0, + uint32_t dilationFilterW = 0, + uint32_t dilationFilterH = 0) + { + loadData3Dv2Params.strideW = strideW; + loadData3Dv2Params.strideH = strideH; + loadData3Dv2Params.filterW = filterW; + loadData3Dv2Params.filterSizeW = filterW >> RIGHT_MOVE_8; + loadData3Dv2Params.filterH = filterH; + loadData3Dv2Params.filterSizeH = filterH >> RIGHT_MOVE_8; + loadData3Dv2Params.dilationFilterW = dilationFilterW; + loadData3Dv2Params.dilationFilterH = dilationFilterH; + } + + CATLASS_DEVICE + static CopyL1ToL0A MakeCopyL1ToL0A( + uint32_t strideW = 0, + uint32_t strideH = 0, + uint32_t filterW = 0, + uint32_t filterH = 0, + uint32_t dilationFilterW = 0, + uint32_t dilationFilterH = 0) + { + return CopyL1ToL0A(strideW, strideH, filterW, filterH, dilationFilterW, dilationFilterH); + } + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor dstTensor, + AscendC::LocalTensor srcTensor, + LayoutDst layoutDst, LayoutSrc layoutSrc, + uint32_t kStartPt, uint32_t mStartPt, + uint32_t l1H, uint32_t l1W, uint8_t *padList) + { + loadData3Dv2Params.l1H = l1H; + loadData3Dv2Params.l1W = l1W; + loadData3Dv2Params.padList[0] = padList[0]; + loadData3Dv2Params.padList[1] = padList[1]; + loadData3Dv2Params.padList[2] = padList[2]; + loadData3Dv2Params.padList[3] = padList[3]; + loadData3Dv2Params.kStartPt = kStartPt; + loadData3Dv2Params.mStartPt = mStartPt; + loadData3Dv2Params.kExtension = layoutDst.orgShape(1); + loadData3Dv2Params.mExtension = layoutDst.orgShape(0); + loadData3Dv2Params.channelSize = layoutSrc.orgShape(1) * layoutSrc.orgShape(2) * layoutSrc.orgShape(5); + static constexpr AscendC::IsResetLoad3dConfig CONV3D_LOAD3DV2_DEFAULT_CONFIG = {true, true}; + AscendC::LoadData(dstTensor, srcTensor, loadData3Dv2Params); + } + +private: + AscendC::LoadData3DParamsV2 loadData3Dv2Params; +}; + +////////////////////////////////////////// + +/// Partial specialization for zN in and zZ out. +template +struct CopyL1ToL0A> { + using LayoutDst = layout::zZ; + using LayoutSrc = layout::zN; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyL1ToL0A() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::LoadData2DParams loadDataParams; + + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = static_cast(layoutDst.shape(3)); + loadDataParams.srcStride = layoutSrc.stride(3) / ELE_NUM_PER_FRACTAL; + loadDataParams.sid = 0; + loadDataParams.dstGap = layoutDst.stride(3) / ELE_NUM_PER_FRACTAL - 1; + loadDataParams.ifTranspose = false; + loadDataParams.addrMode = 0; + + for (uint32_t i = 0; i < layoutDst.shape(1); i++) { + AscendC::LoadData(dstTensor[i * layoutDst.stride(1)], srcTensor[i * layoutSrc.stride(1)], loadDataParams); + } + } +}; + +/// Partial specialization for float, zN in and zZ out. +template +struct CopyL1ToL0A> { + using Element = float; + using LayoutDst = layout::zZ; + using LayoutSrc = layout::zN; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyL1ToL0A() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + constexpr uint8_t PAD_LIST[4] = {0, 0, 0, 0}; + uint16_t l1M = layoutSrc.shape(0) * layoutSrc.shape(1); + uint16_t l1K = layoutSrc.shape(2) * layoutSrc.shape(3); + uint16_t l0M = layoutDst.shape(0) * layoutDst.shape(1); + uint16_t l0K = layoutDst.shape(2) * layoutDst.shape(3); + AscendC::SetFmatrix(1, l1M, PAD_LIST, AscendC::FmatrixMode::FMATRIX_LEFT); + static constexpr AscendC::IsResetLoad3dConfig config = {false, false}; + AscendC::LoadData3DParamsV2 loadDataParams; + loadDataParams.kExtension = l0K; + loadDataParams.mExtension = l0M; + loadDataParams.channelSize = l1K; + + AscendC::LoadData(dstTensor, srcTensor, loadDataParams); + } +}; + +template +struct CopyL1ToL0A> { + using LayoutDst = layout::zZ; + using LayoutSrc = layout::nZ; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + CATLASS_DEVICE + CopyL1ToL0A() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::LoadData2DParams loadDataParams; + + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = static_cast(CeilDiv(layoutDst.orgShape(1))); + loadDataParams.srcStride = layoutSrc.stride(3) / ELE_NUM_PER_FRACTAL; + loadDataParams.sid = 0; + loadDataParams.dstGap = layoutDst.stride(3) / ELE_NUM_PER_FRACTAL - 1; + loadDataParams.ifTranspose = true; + loadDataParams.addrMode = 0; + + for (uint32_t i = 0; i < CeilDiv(layoutDst.orgShape(0)); i++) { + AscendC::LoadData(dstTensor[i * layoutDst.stride(1)], srcTensor[i * layoutSrc.stride(1)], loadDataParams); + } + } +}; + +/// Partial specialization for int8_t, nZ in and zZ out. (Transpose A) +template +struct CopyL1ToL0A> { + using Element = int8_t; + using LayoutDst = layout::zZ; + using LayoutSrc = layout::nZ; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyL1ToL0A() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::LoadData2dTransposeParams loadDataParams; + + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = static_cast(CeilDiv(layoutDst.orgShape(1))); + loadDataParams.srcStride = 1; + loadDataParams.dstGap = 0; + loadDataParams.dstFracGap = CeilDiv(layoutDst.orgShape(1)) - 1; + + for (uint32_t i = 0; i < CeilDiv(layoutDst.orgShape(0)); i++) { + AscendC::LoadDataWithTranspose(dstTensor[i * layoutDst.stride(1) * 2], + srcTensor[i * layoutSrc.stride(1)], + loadDataParams); + } + } +}; + +/// Partial specialization for float, nZ in and zZ out. (Transpose A) +template +struct CopyL1ToL0A> { + using Element = float; + using LayoutDst = layout::zZ; + using LayoutSrc = layout::nZ; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyL1ToL0A() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + constexpr uint8_t PAD_LIST[] = {0, 0, 0, 0}; + uint16_t l1M = layoutSrc.shape(0) * layoutSrc.shape(1); + uint16_t l1K = layoutSrc.shape(2) * layoutSrc.shape(3); + uint16_t l0M = layoutDst.shape(0) * layoutDst.shape(1); + uint16_t l0K = layoutDst.shape(2) * layoutDst.shape(3); + // K, M need to be 16 aligned for f32 + uint16_t l1MAlign = RoundUp(l1M); + uint16_t l1KAlign = RoundUp(l1K); + uint16_t l0MAlign = RoundUp(l0M); + uint16_t l0KAlign = RoundUp(l0K); + AscendC::SetFmatrix(1, l1KAlign, PAD_LIST, AscendC::FmatrixMode::FMATRIX_LEFT); + static constexpr AscendC::IsResetLoad3dConfig config = {false, false}; + AscendC::LoadData3DParamsV2 loadDataParams; + loadDataParams.kExtension = l0MAlign; + loadDataParams.mExtension = l0KAlign; + loadDataParams.enTranspose = true; + loadDataParams.channelSize = l1MAlign; + + AscendC::LoadData(dstTensor, srcTensor, loadDataParams); + } +}; + +///////////////////////////////////////////TileCopyTla////////////////////////////////////////////////////// + +/// Partial specialization for CopyL1ToL0A, AtlasA2, zN in and zZ out. +template +struct TileCopyTla, LayoutSrc, CoordSrc, AscendC::TPosition::A1>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::A2>, + std::enable_if_t::value && + tla::detail::iszZ::value>> { + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + TileCopyTla() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert(tla::detail::iszN::value && + tla::detail::iszZ::value && + TensorSrc::position == AscendC::TPosition::A1 && + TensorDst::position == AscendC::TPosition::A2, + "The input parameters do not match. TensorSrc must be L1 and zN, while TensorDst must be L0A and zZ"); + + const uint32_t srcOuterStrideRow = tla::get<0, 1>(srcTensor.stride()); + const uint32_t srcOuterStrideCol = tla::get<1, 1>(srcTensor.stride()); + const uint32_t dstOuterShapeRow = CeilDiv(tla::get<0>(dstTensor.originShape()), tla::get<0, 0>(dstTensor.shape())); + const uint32_t dstOuterShapeCol = CeilDiv(tla::get<1>(dstTensor.originShape()), tla::get<1, 0>(dstTensor.shape())); + const uint32_t dstOuterStrideRow = tla::get<0, 1>(dstTensor.stride()); + + AscendC::LoadData2DParams loadDataParams; + + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = dstOuterShapeCol; + loadDataParams.srcStride = srcOuterStrideCol / ELE_NUM_PER_FRACTAL; + loadDataParams.sid = 0; + loadDataParams.dstGap = 0; + loadDataParams.ifTranspose = false; + loadDataParams.addrMode = 0; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + + for (uint32_t i = 0; i < dstOuterShapeRow; i++) { + AscendC::LoadData(dstTensor.data()[dstOffset + i * dstOuterStrideRow], + srcTensor.data()[srcOffset + i * srcOuterStrideRow], + loadDataParams); + } + } +}; + +/// Partial specialization for CopyL1ToL0A, AtlasA2, float, zN in and zZ out. +template +struct TileCopyTla, LayoutSrc, CoordSrc, AscendC::TPosition::A1>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::A2>, + std::enable_if_t::value && tla::detail::iszZ::value>> { + using Element = float; + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + TileCopyTla() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert(std::is_same_v && + std::is_same_v && + tla::detail::iszN::value && + tla::detail::iszZ::value && + TensorSrc::position == AscendC::TPosition::A1 && + TensorDst::position == AscendC::TPosition::A2, + "The input parameters do not match. TensorSrc must be float, L1 and zN, " + "while TensorDst must be float, L0A and zZ"); + + constexpr uint8_t PAD_LIST[4] = {0, 0, 0, 0}; + uint16_t l1M = tla::get<1, 1>(srcTensor.stride()) / tla::get<1, 0>(srcTensor.shape()); + uint16_t l1K = RoundUp(tla::get<1>(srcTensor.originShape())); + uint16_t l0M = RoundUp(tla::get<0>(dstTensor.originShape())); + uint16_t l0K = RoundUp(tla::get<1>(dstTensor.originShape())); + AscendC::SetFmatrix(1, l1M, PAD_LIST, AscendC::FmatrixMode::FMATRIX_LEFT); + static constexpr AscendC::IsResetLoad3dConfig config = {false, false}; + AscendC::LoadData3DParamsV2 loadDataParams; + loadDataParams.kExtension = l0K; + loadDataParams.mExtension = l0M; + loadDataParams.channelSize = l1K; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + + AscendC::LoadData(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], loadDataParams); + } +}; + +/// Partial specialization for CopyL1ToL0A, AtlasA2, nZ in and zZ out. (Transpose A) +template +struct TileCopyTla, LayoutSrc, CoordSrc, AscendC::TPosition::A1>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::A2>, + std::enable_if_t::value && + tla::detail::iszZ::value>> { + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + TileCopyTla() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert(tla::detail::isnZ::value && + tla::detail::iszZ::value && + TensorSrc::position == AscendC::TPosition::A1 && + TensorDst::position == AscendC::TPosition::A2, + "The input parameters do not match. TensorSrc must be L1 and nZ, while TensorDst must be L0A and zZ"); + + const uint32_t srcOuterStrideRow = tla::get<0, 1>(srcTensor.stride()); + const uint32_t dstOuterShapeRow = CeilDiv(tla::get<0>(dstTensor.originShape()), tla::get<0, 0>(dstTensor.shape())); + const uint32_t dstOuterShapeCol = CeilDiv(tla::get<1>(dstTensor.originShape()), tla::get<1, 0>(dstTensor.shape())); + const uint32_t dstOuterStrideRow = tla::get<0, 1>(dstTensor.stride()); + + AscendC::LoadData2DParams loadDataParams; + + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = dstOuterShapeCol; + loadDataParams.srcStride = 1; + loadDataParams.sid = 0; + loadDataParams.dstGap = 0; + loadDataParams.ifTranspose = true; + loadDataParams.addrMode = 0; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + + for (uint32_t i = 0; i < dstOuterShapeRow; i++) { + AscendC::LoadData(dstTensor.data()[dstOffset + i * dstOuterStrideRow], + srcTensor.data()[srcOffset + i * srcOuterStrideRow], + loadDataParams); + } + } +}; + +/// Partial specialization for CopyL1ToL0A, AtlasA2, int8_t, nZ in and zZ out. (Transpose A) +template +struct TileCopyTla, LayoutSrc, CoordSrc, AscendC::TPosition::A1>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::A2>, + std::enable_if_t::value && tla::detail::iszZ::value>> { + using Element = int8_t; + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + TileCopyTla() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert(std::is_same_v && + std::is_same_v && + tla::detail::isnZ::value && + tla::detail::iszZ::value && + TensorSrc::position == AscendC::TPosition::A1 && + TensorDst::position == AscendC::TPosition::A2, + "The input parameters do not match. TensorSrc must be int8_t, L1 and nZ, " + "while TensorDst must be int8_t, L0A and zZ"); + + const uint32_t srcOuterShapeRow = CeilDiv(tla::get<0>(srcTensor.originShape()), tla::get<0, 0>(srcTensor.shape())); + const uint32_t srcOuterStrideRow = tla::get<0, 1>(srcTensor.stride()); + const uint32_t dstOuterShapeCol = CeilDiv(tla::get<1>(dstTensor.originShape()), tla::get<1, 0>(dstTensor.shape())); + const uint32_t dstOuterStrideRow = tla::get<0, 1>(dstTensor.stride()); + + AscendC::LoadData2dTransposeParams loadDataParams; + + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = dstOuterShapeCol; + loadDataParams.srcStride = 1; + loadDataParams.dstGap = 0; + loadDataParams.dstFracGap = dstOuterShapeCol - 1; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + + for (uint32_t i = 0; i < srcOuterShapeRow; i++) { + AscendC::LoadDataWithTranspose(dstTensor.data()[dstOffset + i * dstOuterStrideRow * 2], + srcTensor.data()[srcOffset + i * srcOuterStrideRow], + loadDataParams); + } + } +}; + +/// Partial specialization for CopyL1ToL0A, AtlasA2, float, nZ in and zZ out. (Transpose A) +template +struct TileCopyTla, LayoutSrc, CoordSrc, AscendC::TPosition::A1>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::A2>, + std::enable_if_t::value && tla::detail::iszZ::value>> { + using Element = float; + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + TileCopyTla() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert(std::is_same_v && + std::is_same_v && + tla::detail::isnZ::value && + tla::detail::iszZ::value && + TensorSrc::position == AscendC::TPosition::A1 && + TensorDst::position == AscendC::TPosition::A2, + "The input parameters do not match. TensorSrc must be float, L1 and nZ, " + "while TensorDst must be float, L0A and zZ"); + + constexpr uint8_t PAD_LIST[4] = {0, 0, 0, 0}; + uint16_t l1M = tla::get<0>(srcTensor.originShape()); + uint16_t l1K = tla::get<0, 1>(srcTensor.stride()) / tla::get<0, 0>(srcTensor.shape()); + uint16_t l0M = tla::get<0>(dstTensor.originShape()); + uint16_t l0K = tla::get<1>(dstTensor.originShape()); + // K, M need to be 16 aligned for f32 + uint16_t l1MAlign = RoundUp(l1M); + uint16_t l1KAlign = RoundUp(l1K); + uint16_t l0MAlign = RoundUp(l0M); + uint16_t l0KAlign = RoundUp(l0K); + AscendC::SetFmatrix(1, l1KAlign, PAD_LIST, AscendC::FmatrixMode::FMATRIX_LEFT); + static constexpr AscendC::IsResetLoad3dConfig config = {false, false}; + AscendC::LoadData3DParamsV2 loadDataParams; + loadDataParams.kExtension = l0MAlign; + loadDataParams.mExtension = l0KAlign; + loadDataParams.enTranspose = true; + loadDataParams.channelSize = l1MAlign; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + + AscendC::LoadData(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], loadDataParams); + } +}; + +/// Partial specialization for CopyL1ToL0A, AtlasA2, zN in and zZ out. +template +struct TileCopySparseTla, LayoutSrc, CoordSrc, AscendC::TPosition::A1>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::A2>, + std::enable_if_t::value && + tla::detail::iszZ::value>> { + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(ElementSrc); + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(ElementSrc); + static constexpr uint16_t HW_N0 = 16; ///< Hardware configuration parameter for N0, reserved for future hardware configuration + static constexpr uint16_t HW_M0 = 16; ///< Hardware configuration parameter for M0 + static constexpr uint64_t M_POS_BIT = 48; ///< Bit position for M index + static constexpr uint64_t K_POS_BIT = 32; ///< Bit position for K index + static constexpr uint64_t M_STEP_BIT = 16; ///< Bit step for M index + static constexpr uint8_t PAD_LIST[4] = {0, 0, 0, 0}; ///< Padding list for data alignment + + // Methods + + CATLASS_DEVICE + TileCopySparseTla() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert(tla::detail::iszN::value && + tla::detail::iszZ::value && + TensorSrc::position == AscendC::TPosition::A1 && + TensorDst::position == AscendC::TPosition::A2, + "The input parameters do not match. TensorSrc must be L1 and zN, while TensorDst must be L0A and zZ"); + + auto srcShape = srcTensor.shape(); + auto dstShape = dstTensor.shape(); + auto coord = srcTensor.coord(); + uint16_t aL1K = tla::get<1, 0>(srcShape) * tla::get<1, 1>(srcShape); + uint16_t aL1M = tla::get<0, 0>(srcShape) * tla::get<0, 1>(srcShape); + uint16_t wAlign = RoundUp(aL1M, HW_M0); + AscendC::Load3DSetFMatrixCal(1, wAlign, PAD_LIST); + // format(M, K), K_axis is k direction, and M_axis is m direction in load3d intrin + uint16_t madMAlign = tla::get<0, 0>(dstShape) * tla::get<0, 1>(dstShape); + uint16_t kStep = tla::get<1, 0>(dstShape) * tla::get<1, 1>(dstShape); + AscendC::LoadData3DParamsV2Pro loadData3DV2; + loadData3DV2.channelSize = aL1K; + loadData3DV2.extConfig = ((uint64_t)tla::get<0>(coord) << M_POS_BIT) | ((uint64_t)tla::get<1>(coord) << K_POS_BIT) | + ((uint64_t)madMAlign << M_STEP_BIT) | (uint64_t)kStep; + + AscendC::LoadData(dstTensor.data(), srcTensor.data(), loadData3DV2); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace Catlass::Gemm::Tile + +#endif // CATLASS_KDA_ATLASA2_COPY_L1_TO_L0A_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_atlasa2_copy_l1_to_l0b.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_atlasa2_copy_l1_to_l0b.hpp new file mode 100644 index 0000000000..b94c228834 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_atlasa2_copy_l1_to_l0b.hpp @@ -0,0 +1,881 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_ATLASA2_COPY_L1_TO_L0B_HPP +#define CATLASS_KDA_ATLASA2_COPY_L1_TO_L0B_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/kda_numeric_size.hpp" +#include "catlass/arch/kda_arch.hpp" +#include "catlass/layout/kda_layout.hpp" +#include "catlass/gemm/kda_gemm_type.hpp" +#include "catlass/gemm/tile/kda_tile_copy_tla.hpp" +#include "tla/kda_tensor.hpp" + +namespace Catlass::Gemm::Tile { + +template < + class ArchTag, + class L1Type, + class L0Type = void> +struct CopyL1ToL0B { + static_assert(DEPENDENT_FALSE, "Unsupported copy l1 to l0, can not find the specialization."); +}; + +//////////////////////////////////////// +/// new add gemm +template +struct CopyL1ToL0B, Catlass::Gemm::GemmType> { + using LayoutDst = layout::nZ; + using LayoutSrc = layout::zZ; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + CATLASS_DEVICE + CopyL1ToL0B() {} + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor dstTensor, + AscendC::LocalTensor srcTensor, + LayoutDst layoutDst, LayoutSrc layoutSrc) + { + AscendC::LoadData2DParams loadDataParams; + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = static_cast(CeilDiv(layoutSrc.orgShape(1))); + loadDataParams.srcStride = 1; + loadDataParams.sid = 0; + loadDataParams.dstGap = 0; + loadDataParams.ifTranspose = true; + loadDataParams.addrMode = 0; + for (uint32_t i = 0; i < CeilDiv(layoutDst.orgShape(0)); i++) { // K N + AscendC::LoadData(dstTensor[i * layoutDst.stride(1)], srcTensor[i * layoutSrc.stride(1)], loadDataParams); + } + } +}; + +template +struct CopyL1ToL0B, Catlass::Gemm::GemmType> { + using Element = float; + using LayoutDst = layout::nZ; + using LayoutSrc = layout::zZ; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + CATLASS_DEVICE + CopyL1ToL0B() {} + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor dstTensor, + AscendC::LocalTensor srcTensor, + LayoutDst layoutDst, LayoutSrc layoutSrc) + { + AscendC::LoadData2dTransposeParams loadDataParams; + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = static_cast(CeilDiv(layoutSrc.orgShape(1))); + loadDataParams.srcStride = 1; + loadDataParams.dstGap = 0; + loadDataParams.dstFracGap = static_cast(CeilDiv(layoutDst.orgShape(1))) - 1; + for (uint32_t i = 0; i < CeilDiv(layoutDst.orgShape(0)); i++) { // K N + AscendC::LoadDataWithTranspose(dstTensor[i * layoutDst.stride(1) * 2], srcTensor[i * layoutSrc.stride(1)], loadDataParams); + } + } +}; + +template +struct CopyL1ToL0B, Catlass::Gemm::GemmType> { + using Element = int8_t; + using LayoutDst = layout::nZ; + using LayoutSrc = layout::zN; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + CATLASS_DEVICE + CopyL1ToL0B() {} + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor dstTensor, + AscendC::LocalTensor srcTensor, + LayoutDst layoutDst, LayoutSrc layoutSrc) + { + AscendC::LoadData2dTransposeParams loadDataParams; + + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = static_cast(CeilDiv(layoutDst.orgShape(1))); + loadDataParams.srcStride = layoutSrc.stride(3) / ELE_NUM_PER_FRACTAL / 2; + loadDataParams.dstGap = 1; + loadDataParams.dstFracGap = 0; + + for (uint32_t i = 0; i < CeilDiv(layoutDst.orgShape(0)); i++) { + AscendC::LoadDataWithTranspose(dstTensor[i * layoutDst.stride(1)], + srcTensor[i * layoutSrc.stride(1) * 2], + loadDataParams); + } + } +}; + +template +struct CopyL1ToL0B, Catlass::Gemm::GemmType> { + using LayoutDst = layout::nZ; + using LayoutSrc = layout::nZ; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyL1ToL0B() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::LoadData2DParams loadDataParams; + + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = static_cast(layoutDst.shape(3)); + loadDataParams.srcStride = layoutSrc.stride(3) / ELE_NUM_PER_FRACTAL; + loadDataParams.sid = 0; + loadDataParams.dstGap = layoutDst.stride(3) / ELE_NUM_PER_FRACTAL - 1; + loadDataParams.ifTranspose = false; + loadDataParams.addrMode = 0; + + for (uint32_t i = 0; i < layoutDst.shape(1); i++) { + AscendC::LoadData(dstTensor[i * layoutDst.stride(1)], srcTensor[i * layoutSrc.stride(1)], loadDataParams); + } + } +}; +///////////////////////////////////////////// + +//////////////////////////////////////////// +/// new add gemv +template +struct CopyL1ToL0B, Catlass::Gemm::GemmType> { + using LayoutDst = layout::zN; + using LayoutSrc = layout::zN; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyL1ToL0B() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::LoadData2DParams loadDataParams; + + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = static_cast(layoutDst.shape(1)); + loadDataParams.srcStride = layoutSrc.stride(1) / ELE_NUM_PER_FRACTAL; + loadDataParams.sid = 0; + loadDataParams.dstGap = layoutDst.stride(1) / ELE_NUM_PER_FRACTAL - 1; + loadDataParams.ifTranspose = false; + loadDataParams.addrMode = 0; + + for (uint32_t i = 0; i < layoutDst.shape(3); i++) { + AscendC::LoadData(dstTensor[i * layoutDst.stride(3)], srcTensor[i * layoutSrc.stride(3)], loadDataParams); + } + } +}; + +template +struct CopyL1ToL0B, Catlass::Gemm::GemmType> { + using LayoutDst = layout::zN; + using LayoutSrc = layout::nN; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyL1ToL0B() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::LoadData2DParams loadDataParams; + + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = layoutDst.shape(1) * layoutDst.shape(3); + loadDataParams.srcStride = layoutSrc.stride(1) / ELE_NUM_PER_FRACTAL; + loadDataParams.sid = 0; + loadDataParams.dstGap = layoutDst.stride(1) / ELE_NUM_PER_FRACTAL - 1; + loadDataParams.ifTranspose = true; + loadDataParams.addrMode = 0; + AscendC::LoadData(dstTensor, srcTensor, loadDataParams); + }; +}; + +template +struct CopyL1ToL0B, Catlass::Gemm::GemmType> { + using LayoutDst = layout::zN; + using LayoutSrc = layout::nN; + using Element = float; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyL1ToL0B() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::LoadData2dTransposeParams loadDataParams; + + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = static_cast(CeilDiv(layoutDst.orgShape(0))); + loadDataParams.srcStride = 1; + loadDataParams.dstGap = 0; + loadDataParams.dstFracGap = CeilDiv(layoutDst.orgShape(0)) - 1; + + for (uint32_t i = 0; i < CeilDiv<2 * ELE_NUM_PER_C0>(layoutDst.orgShape(1)); i++) { + AscendC::LoadDataWithTranspose( + dstTensor[i * layoutDst.stride(3) * 2], + srcTensor[i * layoutSrc.stride(3)], + loadDataParams); + } + }; +}; + +template +struct CopyL1ToL0B, Catlass::Gemm::GemmType> { + using LayoutDst = layout::zN; + using LayoutSrc = layout::nZ; + using Element = int8_t; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyL1ToL0B() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::LoadData2dTransposeParams loadDataParams; + + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = static_cast(CeilDiv(layoutDst.orgShape(0))); + loadDataParams.srcStride = layoutSrc.stride(1) / ELE_NUM_PER_FRACTAL / 2; + loadDataParams.dstGap = 1; + loadDataParams.dstFracGap = 0; + + for (uint32_t i = 0; i < CeilDiv(layoutDst.orgShape(1)); i++) { + AscendC::LoadDataWithTranspose( + dstTensor[i * layoutDst.stride(3)], + srcTensor[i * layoutSrc.stride(3) * 2], + loadDataParams); + } + } +}; +//////////////////////////////////////////// + +/// Partial specialization for int8_t, zN in and nZ out. +template +struct CopyL1ToL0B> { + using Element = int8_t; + using LayoutDst = layout::nZ; + using LayoutSrc = layout::zN; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyL1ToL0B() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::LoadData2dTransposeParams loadDataParams; + + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = static_cast(CeilDiv(layoutDst.orgShape(1))); + loadDataParams.srcStride = layoutSrc.stride(3) / ELE_NUM_PER_FRACTAL / 2; + loadDataParams.dstGap = 1; + loadDataParams.dstFracGap = 0; + + for (uint32_t i = 0; i < CeilDiv(layoutDst.orgShape(0)); i++) { + AscendC::LoadDataWithTranspose(dstTensor[i * layoutDst.stride(1)], + srcTensor[i * layoutSrc.stride(1) * 2], + loadDataParams); + } + } +}; + +/// Partial specialization for float, zN in and nZ out. +template +struct CopyL1ToL0B> { + using Element = float; + using LayoutDst = layout::nZ; + using LayoutSrc = layout::zN; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyL1ToL0B() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + constexpr uint8_t PAD_LIST[4] = {0, 0, 0, 0}; + uint16_t l1K = layoutSrc.shape(0) * layoutSrc.shape(1); + uint16_t l1N = layoutSrc.shape(2) * layoutSrc.shape(3); + uint16_t l0K = layoutDst.shape(0) * layoutDst.shape(1); + uint16_t l0N = layoutDst.shape(2) * layoutDst.shape(3); + // K, N need to be 16 aligned for f32 + uint16_t l1KAlign = RoundUp(l1K); + uint16_t l1NAlign = RoundUp(l1N); + uint16_t l0KAlign = RoundUp(l0K); + uint16_t l0NAlign = RoundUp(l0N); + AscendC::SetFmatrix(1, l1KAlign, PAD_LIST, AscendC::FmatrixMode::FMATRIX_RIGHT); + static constexpr AscendC::IsResetLoad3dConfig config = {false, false}; + AscendC::LoadData3DParamsV2 loadDataParams; + loadDataParams.kExtension = l0NAlign; + loadDataParams.mExtension = l0KAlign; + loadDataParams.channelSize = l1NAlign; + loadDataParams.fMatrixCtrl = true; + + AscendC::LoadData(dstTensor, srcTensor, loadDataParams); + } +}; + +/// Partial specialization for zN in and nZ out. +template +struct CopyL1ToL0B> { + using LayoutDst = layout::nZ; + using LayoutSrc = layout::zN; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyL1ToL0B() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::LoadData2DParams loadDataParams; + + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = static_cast(CeilDiv(layoutDst.orgShape(1))); + loadDataParams.srcStride = layoutSrc.stride(3) / ELE_NUM_PER_FRACTAL; + loadDataParams.sid = 0; + loadDataParams.dstGap = layoutDst.stride(3) / ELE_NUM_PER_FRACTAL - 1; + loadDataParams.ifTranspose = true; + loadDataParams.addrMode = 0; + + for (uint32_t i = 0; i < CeilDiv(layoutDst.orgShape(0)); i++) { + AscendC::LoadData(dstTensor[i * layoutDst.stride(1)], srcTensor[i * layoutSrc.stride(1)], loadDataParams); + } + } +}; + +template +struct CopyL1ToL0B> { + using Element = AscendC::int4b_t; + using LayoutDst = layout::nZ; + using LayoutSrc = layout::zN; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyL1ToL0B() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::LoadData2dTransposeParams loadDataParams; + + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = static_cast(CeilDiv(layoutDst.orgShape(1))); + loadDataParams.srcStride = layoutSrc.stride(3) / ELE_NUM_PER_FRACTAL / 4; + loadDataParams.dstGap = 3; + loadDataParams.dstFracGap = 0; + + for (uint32_t i = 0; i < CeilDiv(layoutDst.orgShape(0)); i++) { + AscendC::LoadDataWithTranspose(dstTensor[i * layoutDst.stride(1)], + srcTensor[i * layoutSrc.stride(1) * 4], + loadDataParams); + } + } +}; + +/// Partial specialization for nZ in and nZ out. (Transpose B) +template +struct CopyL1ToL0B> { + using LayoutDst = layout::nZ; + using LayoutSrc = layout::nZ; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyL1ToL0B() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::LoadData2DParams loadDataParams; + if (layoutSrc.shape(3) == layoutDst.shape(3)) { + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = static_cast(layoutDst.shape(1) * layoutDst.shape(3)); + loadDataParams.srcStride = 1; + loadDataParams.sid = 0; + loadDataParams.dstGap = 0; + loadDataParams.ifTranspose = false; + loadDataParams.addrMode = 0; + + AscendC::LoadData(dstTensor, srcTensor, loadDataParams); + } else { + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = static_cast(layoutDst.shape(3)); + loadDataParams.srcStride = layoutSrc.stride(3) / ELE_NUM_PER_FRACTAL; + loadDataParams.sid = 0; + loadDataParams.dstGap = layoutDst.stride(3) / ELE_NUM_PER_FRACTAL - 1; + loadDataParams.ifTranspose = false; + loadDataParams.addrMode = 0; + + for (uint32_t i = 0; i < layoutDst.shape(1); i++) { + AscendC::LoadData(dstTensor[i * layoutDst.stride(1)], srcTensor[i * layoutSrc.stride(1)], loadDataParams); + } + } + } +}; + +///////////////////////////////////////////TileCopyTla////////////////////////////////////////////////////// +/// Partial specialization for CopyL1ToL0B, AtlasA2, zN in and nZ out. +template +struct TileCopyTla, LayoutSrc, CoordSrc, AscendC::TPosition::A1>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::B2>, + std::enable_if_t::value && + tla::detail::isnZ::value>> { + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + TileCopyTla() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert(tla::detail::iszN::value && + tla::detail::isnZ::value && + TensorSrc::position == AscendC::TPosition::A1 && + TensorDst::position == AscendC::TPosition::B2, + "The input parameters do not match. TensorSrc must be L1 and zN, while TensorDst must be L0B and nZ"); + + const uint32_t srcOuterStrideRow = tla::get<0, 1>(srcTensor.stride()); + const uint32_t srcOuterStrideCol = tla::get<1, 1>(srcTensor.stride()); + const uint32_t dstOuterShapeRow = CeilDiv(tla::get<0>(dstTensor.originShape()), tla::get<0, 0>(dstTensor.shape())); + const uint32_t dstOuterShapeCol = CeilDiv(tla::get<1>(dstTensor.originShape()), tla::get<1, 0>(dstTensor.shape())); + const uint32_t dstOuterStrideRow = tla::get<0, 1>(dstTensor.stride()); + + AscendC::LoadData2DParams loadDataParams; + + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = dstOuterShapeCol; + loadDataParams.srcStride = srcOuterStrideCol / ELE_NUM_PER_FRACTAL; + loadDataParams.sid = 0; + loadDataParams.dstGap = 0; + loadDataParams.ifTranspose = true; + loadDataParams.addrMode = 0; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + + for (uint32_t i = 0; i < dstOuterShapeRow; i++) { + AscendC::LoadData(dstTensor.data()[dstOffset + i * dstOuterStrideRow], + srcTensor.data()[srcOffset + i * srcOuterStrideRow], + loadDataParams); + } + } +}; + +/// Partial specialization for CopyL1ToL0B, AtlasA2, nZ in and nZ out. (Transpose B) +template +struct TileCopyTla, LayoutSrc, CoordSrc, AscendC::TPosition::A1>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::B2>, + std::enable_if_t::value && + tla::detail::isnZ::value>> { + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + TileCopyTla() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert(tla::detail::isnZ::value && + tla::detail::isnZ::value && + TensorSrc::position == AscendC::TPosition::A1 && + TensorDst::position == AscendC::TPosition::B2, + "The input parameters do not match. TensorSrc must be L1 and nZ, while TensorDst must be L0B and nZ"); + + const uint32_t srcOuterStrideRow = tla::get<0, 1>(srcTensor.stride()); + const uint32_t srcOuterStrideCol = tla::get<1, 1>(srcTensor.stride()); + const uint32_t dstOuterShapeRow = CeilDiv(tla::get<0>(dstTensor.originShape()), tla::get<0, 0>(dstTensor.shape())); + const uint32_t dstOuterShapeCol = CeilDiv(tla::get<1>(dstTensor.originShape()), tla::get<1, 0>(dstTensor.shape())); + const uint32_t dstOuterStrideRow = tla::get<0, 1>(dstTensor.stride()); + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + + AscendC::LoadData2DParams loadDataParams; + if (srcOuterStrideRow == dstOuterStrideRow) { + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = dstOuterShapeRow * dstOuterShapeCol; + loadDataParams.srcStride = 1; + loadDataParams.sid = 0; + loadDataParams.dstGap = 0; + loadDataParams.ifTranspose = false; + loadDataParams.addrMode = 0; + + AscendC::LoadData(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], loadDataParams); + } else { + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = dstOuterShapeCol; + loadDataParams.srcStride = 1; + loadDataParams.sid = 0; + loadDataParams.dstGap = 0; + loadDataParams.ifTranspose = false; + loadDataParams.addrMode = 0; + + for (uint32_t i = 0; i < dstOuterShapeRow; i++) { + AscendC::LoadData(dstTensor.data()[dstOffset + i * dstOuterStrideRow], + srcTensor.data()[srcOffset + i * srcOuterStrideRow], + loadDataParams); + } + } + } +}; + +/// Partial specialization for CopyL1ToL0B, AtlasA2, int8_t, zN in and nZ out. +template +struct TileCopyTla, LayoutSrc, CoordSrc, AscendC::TPosition::A1>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::B2>, + std::enable_if_t::value && tla::detail::isnZ::value>> { + using Element = int8_t; + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + TileCopyTla() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert(std::is_same_v && + std::is_same_v && + tla::detail::iszN::value && + tla::detail::isnZ::value && + TensorSrc::position == AscendC::TPosition::A1 && + TensorDst::position == AscendC::TPosition::B2, + "The input parameters do not match. TensorSrc must be int8_t, L1 and zN, " + "while TensorDst must be int8_t, L0B and nZ"); + + const uint32_t srcOuterShapeCol = CeilDiv(tla::get<1>(srcTensor.originShape()), tla::get<1, 0>(srcTensor.shape())); + const uint32_t srcOuterStrideRow = tla::get<0, 1>(srcTensor.stride()); + const uint32_t srcOuterStrideCol = tla::get<1, 1>(srcTensor.stride()); + const uint32_t dstOuterShapeRow = CeilDiv(tla::get<0>(dstTensor.originShape()), tla::get<0, 0>(dstTensor.shape())); + const uint32_t dstOuterStrideRow = tla::get<0, 1>(dstTensor.stride()); + + AscendC::LoadData2dTransposeParams loadDataParams; + + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = srcOuterShapeCol; + loadDataParams.srcStride = srcOuterStrideCol / ELE_NUM_PER_FRACTAL / 2; + loadDataParams.dstGap = 1; + loadDataParams.dstFracGap = 0; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + + for (uint32_t i = 0; i < dstOuterShapeRow; i++) { + AscendC::LoadDataWithTranspose(dstTensor.data()[dstOffset + i * dstOuterStrideRow], + srcTensor.data()[srcOffset + i * srcOuterStrideRow * 2], + loadDataParams); + } + } +}; + +/// Partial specialization for CopyL1ToL0B, AtlasA2, float, zN in and nZ out. +template +struct TileCopyTla, LayoutSrc, CoordSrc, AscendC::TPosition::A1>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::B2>, + std::enable_if_t::value && tla::detail::isnZ::value>> { + using Element = float; + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + TileCopyTla() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert(std::is_same_v && + std::is_same_v && + tla::detail::iszN::value && + tla::detail::isnZ::value && + TensorSrc::position == AscendC::TPosition::A1 && + TensorDst::position == AscendC::TPosition::B2, + "The input parameters do not match. TensorSrc must be float, L1 and zN, " + "while TensorDst must be float, L0B and nZ"); + + constexpr uint8_t PAD_LIST[4] = {0, 0, 0, 0}; + uint16_t l1K = tla::get<1, 1>(srcTensor.stride()) / tla::get<1, 0>(srcTensor.shape()); + uint16_t l1N = tla::get<1>(srcTensor.originShape()); + uint16_t l0K = tla::get<0>(dstTensor.originShape()); + uint16_t l0N = tla::get<1>(dstTensor.originShape()); + // K, N need to be 16 aligned for f32 + uint16_t l1KAlign = RoundUp(l1K); + uint16_t l1NAlign = RoundUp(l1N); + uint16_t l0KAlign = RoundUp(l0K); + uint16_t l0NAlign = RoundUp(l0N); + AscendC::SetFmatrix(1, l1KAlign, PAD_LIST, AscendC::FmatrixMode::FMATRIX_RIGHT); + static constexpr AscendC::IsResetLoad3dConfig config = {false, false}; + AscendC::LoadData3DParamsV2 loadDataParams; + loadDataParams.kExtension = l0NAlign; + loadDataParams.mExtension = l0KAlign; + loadDataParams.channelSize = l1NAlign; + loadDataParams.fMatrixCtrl = true; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + + AscendC::LoadData(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], loadDataParams); + } +}; + +/// Partial specialization for CopyL1ToL0B, AtlasA2, zN in and nZ out. +template +struct CopyL1ToL0BSparseTla, LayoutSrc, CoordSrc, AscendC::TPosition::A1>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::B2>, + tla::Tensor, LayoutIdx, CoordIdx, AscendC::TPosition::A1>, + std::enable_if_t::value && + tla::detail::isnZ::value && + tla::detail::iszN::value>> { + static constexpr uint8_t INDEX_SHIFT = 2; ///< Shift value for index manipulation + + // Methods + + CATLASS_DEVICE + CopyL1ToL0BSparseTla() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, TensorIdx const &idxTensor) + { + static_assert(tla::detail::iszN::value && + tla::detail::isnZ::value && + TensorSrc::position == AscendC::TPosition::A1 && + TensorDst::position == AscendC::TPosition::B2, + "The input parameters do not match. TensorSrc must be L1 and zN, while TensorDst must be L0B and nZ"); + + auto dstShape = dstTensor.shape(); + auto dstStride = dstTensor.stride(); + auto srcShape = srcTensor.shape(); + auto srcStride = srcTensor.stride(); + auto coord = srcTensor.coord(); + + // SET LOAD2D parameters , loop axis: K or M, or 1 + uint16_t madK = tla::get<1, 0>(dstShape) * tla::get<1, 1>(dstShape); + + // k is c0Size_ aligned for f32 + uint16_t kC0 = CeilDiv(madK, tla::get<1, 0>(dstShape)); + uint16_t nFraC0 = tla::get<0, 1>(dstShape); + uint16_t l0bLoop = 1; + uint64_t l0bDstAddrStride = 0; + uint64_t l0bSrcAddrStride = 0; + uint8_t l0bRepeat = kC0 * nFraC0; + uint16_t l0bDststride = 0; + uint16_t l0bSrcstride = 1; + + uint16_t bL1N = tla::get<0, 0>(srcShape) * tla::get<0, 1>(srcShape); + if (nFraC0 * tla::get<0, 0>(dstShape) == bL1N) { // loop=1 + l0bLoop = 1; + } else if (nFraC0 >= kC0) { // LOOP is K and repeat is n axis + l0bLoop = kC0; + l0bDstAddrStride = nFraC0 * tla::get<1, 1>(dstStride); + l0bSrcAddrStride = bL1N * tla::get<1, 0>(srcShape); + l0bRepeat = nFraC0; + + l0bDststride = 0; + l0bSrcstride = 1; + } else { // LOOP is N and repeat is K axis + l0bLoop = nFraC0; + l0bDstAddrStride = tla::get<1, 1>(dstStride); + l0bSrcAddrStride = tla::get<0, 1>(srcStride); + l0bRepeat = kC0; + + l0bDststride = nFraC0 - 1; + l0bSrcstride = bL1N; + } + // use load2d for L1_2_L0B + // startIndex, repeatTimes, srcStride, sid, dstGap, ifTranspose, addrmode + AscendC::LoadData2dParams loadDataParams{0, l0bRepeat, l0bSrcstride, 0, l0bDststride, 0, 0}; + uint16_t bL1KOffset = tla::get<1>(coord); + uint64_t l1bOffset = tla::get<0>(coord) * tla::get<1, 0>(srcShape) + bL1KOffset * bL1N; + uint64_t l0bOffset = 0; + AscendC::LocalTensor dstLocal; + AscendC::LocalTensor srcLocal; + AscendC::LocalTensor l1BIndex; + for (uint64_t i = 0; i < l0bLoop; i++) { + dstLocal.SetAddr(dstTensor.data()[l0bOffset].address_); + srcLocal.SetAddr(srcTensor.data()[l1bOffset].address_); + l1BIndex.SetAddr(idxTensor.data()[l1bOffset >> INDEX_SHIFT].address_); + AscendC::LoadDataWithSparse(dstLocal, srcLocal, l1BIndex, loadDataParams); + l1bOffset += l0bSrcAddrStride; + l0bOffset += l0bDstAddrStride; + } + } +}; + +// Partial specialization for CopyL1ToL0B, AtlasA2, nZ in and nZ out. (Transpose B) +template +struct CopyL1ToL0BSparseTla, LayoutSrc, CoordSrc, AscendC::TPosition::A1>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::B2>, + tla::Tensor, LayoutIdx, CoordIdx, AscendC::TPosition::A1>, + std::enable_if_t::value && + tla::detail::isnZ::value && + tla::detail::isnZ::value>> { + static constexpr uint8_t INDEX_SHIFT = 2; ///< Shift value for index manipulation + + // Methods + + CATLASS_DEVICE + CopyL1ToL0BSparseTla() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, TensorIdx const &idxTensor) + { + static_assert(tla::detail::isnZ::value && + tla::detail::isnZ::value && + TensorSrc::position == AscendC::TPosition::A1 && + TensorDst::position == AscendC::TPosition::B2, + "The input parameters do not match. TensorSrc must be L1 and nZ, while TensorDst must be L0B and nZ"); + + auto dstShape = dstTensor.shape(); + auto dstStride = dstTensor.stride(); + auto srcShape = srcTensor.shape(); + auto srcStride = srcTensor.stride(); + auto coord = srcTensor.coord(); + + // SET LOAD2D parameters , loop axis: K or M, or 1 + uint16_t madK = tla::get<0, 0>(dstShape) * tla::get<0, 1>(dstShape); + + // k is c0Size_ aligned for f32 + uint16_t kC0 = CeilDiv(madK, tla::get<0, 0>(dstShape)); + uint16_t nFraC0 = tla::get<1, 1>(dstShape); + uint16_t l0bLoop = 1; + uint64_t l0bDstAddrStride = 0; + uint64_t l0bSrcAddrStride = 0; + uint8_t l0bRepeat = kC0 * nFraC0; + + AscendC::LoadData2DParams loadDataParams; + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = static_cast(l0bRepeat); + loadDataParams.srcStride = 1; + loadDataParams.sid = 0; + loadDataParams.dstGap = 0; + loadDataParams.ifTranspose = false; + loadDataParams.addrMode = 0; + uint16_t bL1N = tla::get<1, 0>(srcShape) * tla::get<1, 1>(srcShape); + uint16_t bL1KOffset = tla::get<0>(coord); + uint64_t l1bOffset = tla::get<1>(coord) * tla::get<0, 0>(srcShape) + bL1KOffset * bL1N; + + AscendC::LocalTensor dstLocal; + AscendC::LocalTensor srcLocal; + AscendC::LocalTensor l1BIndex; + dstLocal.SetAddr(dstTensor.data().address_); + srcLocal.SetAddr(srcTensor.data()[l1bOffset].address_); + l1BIndex.SetAddr(idxTensor.data()[l1bOffset >> INDEX_SHIFT].address_); + + AscendC::LoadDataWithSparse(dstLocal, srcLocal, l1BIndex, loadDataParams); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace Catlass::Gemm::Tile + +#endif // CATLASS_KDA_ATLASA2_COPY_L1_TO_L0B_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_atlasa2_copy_ub_to_gm.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_atlasa2_copy_ub_to_gm.hpp new file mode 100644 index 0000000000..4d558811d3 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_atlasa2_copy_ub_to_gm.hpp @@ -0,0 +1,131 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_ATLASA2_COPY_UB_TO_GM_HPP +#define CATLASS_KDA_ATLASA2_COPY_UB_TO_GM_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_arch.hpp" +#include "catlass/gemm/tile/kda_tile_copy_tla.hpp" +#include "catlass/gemm/kda_gemm_type.hpp" +#include "tla/kda_tensor.hpp" + +namespace Catlass::Gemm::Tile { + +/// Partial specialization for AtlasA2, RowMajor in and RowMajor out. +template +struct TileCopyTla, LayoutSrc, CoordSrc, AscendC::TPosition::VECCALC>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::GM>, + std::enable_if_t::value && + tla::detail::isRowMajor::value>> { + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(ElementSrc); + + // Methods + + CATLASS_DEVICE + TileCopyTla() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert(tla::detail::isRowMajor::value && + tla::detail::isRowMajor::value && + TensorSrc::position == AscendC::TPosition::VECCALC && + TensorDst::position == AscendC::TPosition::GM, + "The input parameters do not match. TensorSrc must be GM and RowMajor, " + "while TensorDst must be UB and RowMajor"); + + const uint16_t row = tla::get<0>(dstTensor.originShape()); + const uint16_t col = tla::get<1>(dstTensor.originShape()); + + AscendC::DataCopyExtParams dataCopyParams( + row, + col * sizeof(ElementSrc), + (tla::get<0>(srcTensor.stride()) - col) / ELE_NUM_PER_C0, + (tla::get<0>(dstTensor.stride()) - col) * sizeof(ElementSrc), + 0); + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + AscendC::DataCopyPad(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], dataCopyParams); + }; +}; + +/// Partial specialization for AtlasA2, RowMajor in and PaddingRowMajor out. +template +struct TileCopyTlaExt, LayoutSrc, CoordSrc, AscendC::TPosition::VECCALC>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::GM>, + layout::RowMajor, layout::PaddingRowMajor> { + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(ElementSrc); + + // Methods + + CATLASS_DEVICE + TileCopyTlaExt() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert(tla::detail::isRowMajor::value && + TensorSrc::position == AscendC::TPosition::VECCALC && + TensorDst::position == AscendC::TPosition::GM, + "The input parameters do not match. TensorSrc must be GM and PaddingRowMajor, " + "while TensorDst must be UB and RowMajor"); + + AscendC::DataCopyExtParams dataCopyParams( + tla::get<1, 1>(dstTensor.shape()), + tla::get<1, 0>(dstTensor.shape()) * sizeof(ElementSrc), + (tla::get<0>(srcTensor.stride()) - tla::get<1>(srcTensor.shape())) / ELE_NUM_PER_C0, + (tla::get<1, 1>(dstTensor.stride()) - tla::get<1, 0>(dstTensor.shape())) * sizeof(ElementSrc), + 0); + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + AscendC::DataCopyPad(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], dataCopyParams); + }; +}; + +template < + class ArchTag, + class GmType> +struct CopyUb2Gm { + static_assert(DEPENDENT_FALSE, "Unsupported copy ub to gm, can not find the specialization."); +}; + +template +struct CopyUb2Gm> { + using LayoutDst = layout::RowMajor; + using LayoutSrc = layout::RowMajor; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + + CATLASS_DEVICE + CopyUb2Gm() = default; + + CATLASS_DEVICE + void operator()( + AscendC::GlobalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + layout::RowMajor const &layoutDst, + layout::RowMajor const &layoutSrc) + { + AscendC::DataCopyExtParams dataCopyParams( + layoutDst.shape(0), + layoutDst.shape(1) * sizeof(Element), + (layoutSrc.stride(0) - layoutSrc.shape(1)) / ELE_NUM_PER_C0, + (layoutDst.stride(0) - layoutDst.shape(1)) * sizeof(Element), + 0); + AscendC::DataCopyPad(dstTensor, srcTensor, dataCopyParams); + } +}; + +} // namespace Catlass::Gemm::Tile + +#endif // CATLASS_KDA_ATLASA2_COPY_UB_TO_GM_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_cast_fp8_to_fp16.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_cast_fp8_to_fp16.hpp new file mode 100644 index 0000000000..9ffb0a69ff --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_cast_fp8_to_fp16.hpp @@ -0,0 +1,415 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_GEMM_TILE_ATLASA2_KDA_CAST_FP8_TO_FP16_HPP +#define CATLASS_GEMM_TILE_ATLASA2_KDA_CAST_FP8_TO_FP16_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/kda_coord.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" +#include "catlass/gemm/kda_gemm_helper.hpp" + +namespace Catlass::Gemm::Tile { + +template < + class ArchTag, + class SrcType_, + class DstType_, + uint32_t COMPUTE_LENGTH> +struct TileCastFp8ToFp16Dequant { + using ElementSrc = typename SrcType_::Element; + using ElementDst = typename DstType_::Element; + using LayoutSrc = typename SrcType_::Layout; + using LayoutDst = typename DstType_::Layout; + using LayoutRowMajor = Catlass::layout::RowMajor; + + static_assert(std::is_same_v || std::is_same_v, + "Unsupported layout, only can be Row/Column Major."); + static_assert(std::is_same_v, "layout src and dst must be the same."); + + static const uint32_t Alignment = 256; + // static constexpr uint32_t ELE_NUM_PER_BLK = BYTE_PER_BLK / sizeof(int8_t); + + struct Params { + half scalar; + half zeroPoint; + + CATLASS_HOST_DEVICE + Params() = default; + + CATLASS_HOST_DEVICE + Params(half scalar_, half zeroPoint_) + { + scalar = scalar_; + zeroPoint = zeroPoint_; + } + }; + + CATLASS_DEVICE + TileCastFp8ToFp16Dequant() + {} + + CATLASS_DEVICE + TileCastFp8ToFp16Dequant(Arch::Resource &resource, Params const ¶ms_) + : params(params_) + { + int64_t bufferOffset = 0; + for (uint32_t i = 0; i < BUFFER_NUM; i++) { + inputBuffer[i] = resource.ubBuf.template GetBufferByByte(bufferOffset * sizeof(ElementSrc)); + bufferOffset += COMPUTE_LENGTH; + } + for (uint32_t i = 0; i < BUFFER_NUM; i++) { + outputBuffer[i] = (resource.ubBuf.template GetBufferByByte(bufferOffset * sizeof(ElementSrc))) + .template ReinterpretCast(); + bufferOffset += COMPUTE_LENGTH * 2; + } + for (uint32_t i = 0; i < BUFFER_NUM; i++) { + workspace[i] = (resource.ubBuf.template GetBufferByByte(bufferOffset * sizeof(ElementSrc))) + .template ReinterpretCast(); + bufferOffset += COMPUTE_LENGTH * 2; + } + InitLocalMaskVec(resource, bufferOffset); + } + + CATLASS_DEVICE + void operator()( + AscendC::GlobalTensor gmDst, LayoutDst const &layoutDst, + AscendC::GlobalTensor gmSrc, LayoutSrc const &layoutSrc, + uint32_t &bufferIndex) + { + uint32_t tilesNum, tileLen, srcStride, dstStride; + if constexpr (std::is_same_v) { + tilesNum = layoutSrc.shape(0); + tileLen = layoutSrc.shape(1); + srcStride = layoutSrc.stride(0); + dstStride = layoutDst.stride(0); + } else if constexpr (std::is_same_v) { + tilesNum = layoutSrc.shape(1); + tileLen = layoutSrc.shape(0); + srcStride = layoutSrc.stride(1); + dstStride = layoutDst.stride(1); + } + + uint32_t tilesPerAiv = tilesNum / AscendC::GetSubBlockNum(); + uint32_t tilesRemain = tilesNum % AscendC::GetSubBlockNum(); + if (AscendC::GetSubBlockIdx() < tilesRemain) { + tilesPerAiv++; + } + + uint64_t taskSrcOffset = AscendC::GetSubBlockIdx() * tilesPerAiv * srcStride; + uint64_t taskDstOffset = AscendC::GetSubBlockIdx() * tilesPerAiv * dstStride; + if (AscendC::GetSubBlockIdx() >= tilesRemain) { + taskSrcOffset += tilesRemain * srcStride; + taskDstOffset += tilesRemain * dstStride; + } + + uint32_t totalLoops = 0; + uint32_t loopsPerTile, tilesInALoop; + uint32_t tileLenRoundFp8 = RoundUp(tileLen); + if (tileLenRoundFp8 > COMPUTE_LENGTH / 2) { + // One signle tile length is bigger than COMPUTE_LENGTH, which should be clipped. + loopsPerTile = CeilDiv(tileLen, COMPUTE_LENGTH); + totalLoops = tilesPerAiv * loopsPerTile; + } else if (tileLenRoundFp8 != 0) { + // COMPUTE_LENGTH is bigger than tile length, such that more than one tiles can be arranged together. + tilesInALoop = COMPUTE_LENGTH / tileLenRoundFp8; + totalLoops = CeilDiv(tilesPerAiv, tilesInALoop); + } // tileLenRoundFp8 == 0 --> totalLoops = 0 + + uint32_t tileTailLen = tileLen % COMPUTE_LENGTH; + uint64_t srcProcessOffset, dstProcessOffset; + uint32_t loadLen = COMPUTE_LENGTH, storeLen, loadRepeat = 1, storeRepeat = 1; + // uint32_t srcLoadStride = srcStride, dstLoadStride = tileLenRoundFp8, + // srcStoreStride = tileLenRoundFp8, dstStoreStride = dstStride; + for (int ldx = 0; ldx < totalLoops; ldx++) { + // Dynamic compute length + if (tileLenRoundFp8 > COMPUTE_LENGTH / 2) { + uint32_t fullTileRounds = ldx / loopsPerTile; + uint32_t residueTileRounds = ldx % loopsPerTile; + srcProcessOffset = taskSrcOffset + fullTileRounds * srcStride + residueTileRounds * COMPUTE_LENGTH; + dstProcessOffset = taskDstOffset + fullTileRounds * dstStride + residueTileRounds * COMPUTE_LENGTH; + + if ((residueTileRounds == loopsPerTile - 1) && (tileTailLen != 0)) { + loadLen = tileTailLen; + } + } else { + uint32_t fullTileRounds = ldx * tilesInALoop; + srcProcessOffset = taskSrcOffset + fullTileRounds * srcStride; + dstProcessOffset = taskDstOffset + fullTileRounds * dstStride; + + loadLen = tileLen; + storeLen = tileLen; + loadRepeat = tilesInALoop; + storeRepeat = tilesInALoop; + + if ((ldx == totalLoops - 1) && (tilesPerAiv % tilesInALoop != 0)) { + loadRepeat = tilesPerAiv % tilesInALoop; + storeRepeat = loadRepeat; + } + } + // uint32_t srcLoadStride = srcStride; + // uint32_t dstLoadStride = tileLenRoundFp8; + // uint32_t srcStoreStride = tileLenRoundFp8; + + // GM -> UB + AscendC::DataCopyExtParams dataCopyParamsIn( + loadRepeat, + loadLen * sizeof(ElementSrc), + (srcStride - loadLen) * sizeof(ElementSrc), // + (tileLenRoundFp8 - loadLen) * sizeof(ElementSrc) / BYTE_PER_BLK, + 0); + AscendC::DataCopyPadExtParams padParams(false, 0, 0, 0); + + AscendC::WaitFlag(EventIdBuffer[bufferIndex]); + AscendC::DataCopyPad(inputBuffer[bufferIndex], gmSrc[srcProcessOffset], + dataCopyParamsIn, padParams); + + AscendC::SetFlag(EventIdBuffer[bufferIndex]); + AscendC::WaitFlag(EventIdBuffer[bufferIndex]); + AscendC::WaitFlag(EventIdBuffer[bufferIndex]); + + Dequant(inputBuffer[bufferIndex], + outputBuffer[bufferIndex], + value_vector1, + value_vector2, + workspace[bufferIndex]); + AscendC::SetFlag(EventIdBuffer[bufferIndex]); + AscendC::WaitFlag(EventIdBuffer[bufferIndex]); + AscendC::SetFlag(EventIdBuffer[bufferIndex]); + + AscendC::DataCopyExtParams dataCopyParams( + storeRepeat, + storeLen * sizeof(ElementDst), + (tileLenRoundFp8 - storeLen) * sizeof(ElementDst) / BYTE_PER_C0, + (dstStride - storeLen) * sizeof(ElementDst), + 0); + AscendC::DataCopyPad(gmDst[dstProcessOffset], outputBuffer[bufferIndex], dataCopyParams); + AscendC::SetFlag(EventIdBuffer[bufferIndex]); + bufferIndex = (bufferIndex + 1) % BUFFER_NUM; + } + } + + /* + * DATAFLOW: gmSrc -> ubInTensor -> ubOutTensor -> gmDst + */ + CATLASS_DEVICE + void EpCastFp32ToFp16( + AscendC::GlobalTensor gmDst, LayoutRowMajor layoutDst, + AscendC::GlobalTensor gmSrc, LayoutRowMajor layoutSrc) + { + AscendC::LocalTensor ubInTensor[BUFFER_NUM]; + AscendC::LocalTensor ubOutTensor[BUFFER_NUM]; + + Arch::Resource resource; + int64_t bufferOffset = 0; + const int64_t CAST_LENGTH = 32 * 1024 / sizeof(half); // 一次处理16K个数据 + for (int i = 0; i < BUFFER_NUM; i++) { + ubInTensor[i] = resource.ubBuf.template GetBufferByByte(bufferOffset); + bufferOffset += CAST_LENGTH * 4; // float 4字节 + } + for (int i = 0; i < BUFFER_NUM; i++) { + ubOutTensor[i] = resource.ubBuf.template GetBufferByByte(bufferOffset); + bufferOffset += CAST_LENGTH * 2; // half 2字节 + } + + uint32_t tilesNum, tileLen, srcStride, dstStride; + tilesNum = layoutSrc.shape(0); + tileLen = layoutSrc.shape(1); + srcStride = layoutSrc.stride(0); + dstStride = layoutDst.stride(0); // Always RowMajor + + uint32_t tilesPerAiv = tilesNum / AscendC::GetSubBlockNum(); + uint32_t tilesRemain = tilesNum % AscendC::GetSubBlockNum(); + if (AscendC::GetSubBlockIdx() < tilesRemain) { + tilesPerAiv++; + } + + uint64_t taskSrcOffset = AscendC::GetSubBlockIdx() * tilesPerAiv * srcStride; + uint64_t taskDstOffset = AscendC::GetSubBlockIdx() * tilesPerAiv * dstStride; + if (AscendC::GetSubBlockIdx() >= tilesRemain) { + taskSrcOffset += tilesRemain * srcStride; + taskDstOffset += tilesRemain * dstStride; + } + + uint32_t totalLoops = 0; + uint32_t loopsPerTile, tilesInALoop; + uint32_t tileLenRoundFp8 = RoundUp(tileLen); + if (tileLenRoundFp8 > COMPUTE_LENGTH / 2) { + // One signle tile length is bigger than COMPUTE_LENGTH, which should be clipped. + loopsPerTile = CeilDiv(tileLen, COMPUTE_LENGTH); + totalLoops = tilesPerAiv * loopsPerTile; + } else if (tileLenRoundFp8 != 0) { + // COMPUTE_LENGTH is bigger than tile length, such that more than one tiles can be arranged together. + tilesInALoop = COMPUTE_LENGTH / tileLenRoundFp8; + totalLoops = CeilDiv(tilesPerAiv, tilesInALoop); + } // tileLenRoundFp8 == 0 --> totalLoops = 0 + + uint32_t tileTailLen = tileLen % COMPUTE_LENGTH; + uint64_t srcProcessOffset, dstProcessOffset; + uint32_t loadLen = COMPUTE_LENGTH, storeLen, loadRepeat = 1, storeRepeat = 1; + // uint32_t srcLoadStride = srcStride, dstLoadStride = tileLenRoundFp8, + // srcStoreStride = tileLenRoundFp8, dstStoreStride = dstStride; + for (int ldx = 0; ldx < totalLoops; ldx++) { + // Dynamic compute length + if (tileLenRoundFp8 > COMPUTE_LENGTH / 2) { + uint32_t fullTileRounds = ldx / loopsPerTile; + uint32_t residueTileRounds = ldx % loopsPerTile; + srcProcessOffset = taskSrcOffset + fullTileRounds * srcStride + residueTileRounds * COMPUTE_LENGTH; + dstProcessOffset = taskDstOffset + fullTileRounds * dstStride + residueTileRounds * COMPUTE_LENGTH; + + if ((residueTileRounds == loopsPerTile - 1) && (tileTailLen != 0)) { + loadLen = tileTailLen; + } + } else { + uint32_t fullTileRounds = ldx * tilesInALoop; + srcProcessOffset = taskSrcOffset + fullTileRounds * srcStride; + dstProcessOffset = taskDstOffset + fullTileRounds * dstStride; + + loadLen = tileLen; + storeLen = tileLen; + loadRepeat = tilesInALoop; + storeRepeat = tilesInALoop; + + if ((ldx == totalLoops - 1) && (tilesPerAiv % tilesInALoop != 0)) { + loadRepeat = tilesPerAiv % tilesInALoop; + storeRepeat = loadRepeat; + } + } + + // copy GM -> UB + AscendC::WaitFlag(EventIdBufferForCast[bufferIndexForCast]); + AscendC::DataCopyExtParams dataCopyParamsIn( + loadRepeat, + loadLen * sizeof(float), // float loaded + (srcStride - loadLen) * sizeof(float), + (tileLenRoundFp8 - loadLen) * sizeof(float) / BYTE_PER_BLK, + 0); + AscendC::DataCopyPadExtParams padParams(false, 0, 0, 0); + AscendC::DataCopyPad(ubInTensor[bufferIndexForCast], gmSrc[srcProcessOffset], + dataCopyParamsIn, padParams); + // Begin casting ... + AscendC::SetFlag(EventIdBufferForCast[bufferIndexForCast]); + AscendC::WaitFlag(EventIdBufferForCast[bufferIndexForCast]); + AscendC::WaitFlag(EventIdBufferForCast[bufferIndexForCast]); + AscendC::Cast(ubOutTensor[bufferIndexForCast], ubInTensor[bufferIndexForCast], + AscendC::RoundMode::CAST_RINT, CAST_LENGTH); + AscendC::SetFlag(EventIdBufferForCast[bufferIndexForCast]); + AscendC::WaitFlag(EventIdBufferForCast[bufferIndexForCast]); + AscendC::SetFlag(EventIdBufferForCast[bufferIndexForCast]); + // End casting ... + + AscendC::DataCopyExtParams dataCopyParams( + storeRepeat, + storeLen * sizeof(half), + (tileLenRoundFp8 - storeLen) * sizeof(half) / BYTE_PER_C0, + (dstStride - storeLen) * sizeof(half), + 0); + AscendC::DataCopyPad(gmDst[dstProcessOffset], ubOutTensor[bufferIndexForCast], dataCopyParams); + AscendC::SetFlag(EventIdBufferForCast[bufferIndexForCast]); + bufferIndexForCast = (bufferIndexForCast + 1) % BUFFER_NUM; + } + } + +private: + CATLASS_DEVICE + void InitLocalMaskVec(Arch::Resource &resource, int64_t &bufferOffset) + { + int16_t value_uint = 0x4000; + value_vector1 = (resource.ubBuf.template GetBufferByByte(bufferOffset * sizeof(ElementSrc))) + .template ReinterpretCast(); + bufferOffset += 256; + AscendC::Duplicate(value_vector1, value_uint, 128); + AscendC::PipeBarrier(); + // AscendC::PipeBarrier(); + value_uint = 0x3FFF; + value_vector2 = (resource.ubBuf.template GetBufferByByte(bufferOffset * sizeof(ElementSrc))) + .template ReinterpretCast(); + bufferOffset += 256; + AscendC::Duplicate(value_vector2, value_uint, 128); + AscendC::PipeBarrier(); + } + + CATLASS_DEVICE + void Dequant(AscendC::LocalTensor &src, AscendC::LocalTensor &dst, + AscendC::LocalTensor &value_vector1, AscendC::LocalTensor &value_vector2, + AscendC::LocalTensor &workspace) + { + pipe_barrier(PIPE_V); + uint32_t num = COMPUTE_LENGTH; + num = (num + 128 - 1) / 128 * 128; + AscendC::Cast(dst.template ReinterpretCast(), + src.template ReinterpretCast(), + AscendC::RoundMode::CAST_NONE, + num); + pipe_barrier(PIPE_V); + + AscendC::Adds(dst, dst, 1024, num); + pipe_barrier(PIPE_V); + + AscendC::ShiftLeft( + dst.template ReinterpretCast(), dst.template ReinterpretCast(), 7, num); + pipe_barrier(PIPE_V); + + uint64_t mask = 128; + AscendC::And(workspace.template ReinterpretCast(), + dst.template ReinterpretCast(), + value_vector1, + mask, + num / 128, + {1, 1, 1, 8, 8, 0}); + pipe_barrier(PIPE_V); + + AscendC::ShiftLeft( + workspace.template ReinterpretCast(), workspace.template ReinterpretCast(), 1, num); + pipe_barrier(PIPE_V); + + AscendC::And(dst.template ReinterpretCast(), + dst.template ReinterpretCast(), + value_vector2, + mask, + num / 128, + {1, 1, 1, 8, 8, 0}); + pipe_barrier(PIPE_V); + + AscendC::Or(dst.template ReinterpretCast(), + dst.template ReinterpretCast(), + workspace.template ReinterpretCast(), + num); + pipe_barrier(PIPE_V); + + AscendC::Muls(dst.template ReinterpretCast(), dst.template ReinterpretCast(), 1 << 8, num); + pipe_barrier(PIPE_V); + + AscendC::Adds(dst, dst, params.zeroPoint, num); + pipe_barrier(PIPE_V); + + AscendC::Muls(dst, dst, params.scalar, num); + pipe_barrier(PIPE_V); + } + +private: + static const uint32_t BUFFER_NUM = 2; + AscendC::LocalTensor inputBuffer[BUFFER_NUM]; + AscendC::LocalTensor value_vector1; + AscendC::LocalTensor value_vector2; + AscendC::LocalTensor outputBuffer[BUFFER_NUM]; + AscendC::LocalTensor workspace[BUFFER_NUM]; + AscendC::TEventID EventIdBuffer[BUFFER_NUM] = {EVENT_ID0, EVENT_ID1}; + AscendC::TEventID EventIdBufferForCast[BUFFER_NUM] = {EVENT_ID2, EVENT_ID3}; + uint32_t bufferIndexForCast{0}; + + Params params; +}; + +} // namespace Catlass::Gemm::Tile + +#endif // CATLASS_GEMM_TILE_ATLASA2_KDA_CAST_FP8_TO_FP16_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_cast_int8_to_fp16.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_cast_int8_to_fp16.hpp new file mode 100644 index 0000000000..e68d4bf327 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_cast_int8_to_fp16.hpp @@ -0,0 +1,188 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_GEMM_TILE_ATLASA2_KDA_CAST_INT8_TO_FP16_HPP +#define CATLASS_GEMM_TILE_ATLASA2_KDA_CAST_INT8_TO_FP16_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/kda_coord.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" +#include "catlass/gemm/kda_gemm_helper.hpp" + +namespace Catlass::Gemm::Tile { + +template < + class ArchTag, + class SrcType_, + class DstType_, + // Length of the compute elements + uint32_t COMPUTE_LEN_, + uint32_t STAGES = 2> +struct TileCastInt8ToFp16Dequant { + using ElementSrc = typename SrcType_::Element; + using ElementDst = typename DstType_::Element; + using LayoutSrc = typename SrcType_::Layout; + using LayoutDst = typename DstType_::Layout; + + static_assert(std::is_same_v || std::is_same_v, + "Unsupported layout, only can be RowMajor and ColumnMajor"); + static_assert(std::is_same_v, "layoutDst and layoutSrc must be the same"); + + static constexpr uint32_t ELE_NUM_PER_BLK_INT8 = BYTE_PER_BLK / sizeof(ElementSrc); + static constexpr uint32_t ELE_NUM_PER_BLK_HALF = BYTE_PER_BLK / sizeof(ElementDst); + + static constexpr uint32_t COMPUTE_LEN = COMPUTE_LEN_; + + static_assert(COMPUTE_LEN <= 32 * 1024, "COMPUTE_LEN cannot exceed 32 * 1024"); + + struct Params { + half deqScalar; + half deqZeroPoint; + + CATLASS_HOST_DEVICE + Params() = default; + + CATLASS_HOST_DEVICE + Params(half deqScalar_, half deqZeroPoint_) + { + deqScalar = deqScalar_; + deqZeroPoint = deqZeroPoint_; + } + }; + + /// Construct + CATLASS_DEVICE + TileCastInt8ToFp16Dequant(Arch::Resource const &resource, Params const ¶ms_) + : params(params_) + { + if constexpr (g_coreType == AscendC::AIV) { + uint32_t ubOffset = 0; + uint32_t ubInSize = COMPUTE_LEN * sizeof(ElementSrc); + uint32_t ubOutSize = COMPUTE_LEN * sizeof(ElementDst); + // Init buffers + for (uint32_t i = 0; i < STAGES; i++) { + ubInTensorList[i] = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += ubInSize; + ubOutTensorList[i] = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += ubOutSize; + + ubEventList[i] = i; + AscendC::SetFlag(ubEventList[i]); + AscendC::SetFlag(ubEventList[i]); + } + } + } + + /// Destructor + CATLASS_DEVICE + ~TileCastInt8ToFp16Dequant() + { + if constexpr (g_coreType == AscendC::AIV) { + for (uint32_t i = 0; i < STAGES; i++) { + AscendC::WaitFlag(ubEventList[i]); + AscendC::WaitFlag(ubEventList[i]); + } + } + } + + CATLASS_DEVICE + void operator()( + AscendC::GlobalTensor const &gmDst, LayoutDst const &layoutDst, + AscendC::GlobalTensor const &gmSrc, LayoutSrc const &layoutSrc) + { + uint32_t tilesNum = layoutSrc.shape(0); + uint32_t tileLen = layoutSrc.shape(1); + uint32_t tileLenRoundInt8 = RoundUp(layoutSrc.shape(1), ELE_NUM_PER_BLK_INT8); + uint64_t tileStrideSrc = layoutSrc.stride(0); + uint64_t tileStrideDst = layoutDst.stride(0); + if constexpr (std::is_same_v) { + tilesNum = layoutSrc.shape(1); + tileLen = layoutSrc.shape(0); + tileLenRoundInt8 = RoundUp(layoutSrc.shape(0), ELE_NUM_PER_BLK_INT8); + tileStrideSrc = layoutSrc.stride(1); + tileStrideDst = layoutDst.stride(1); + } + uint32_t tilesPerAiv = tilesNum / AscendC::GetSubBlockNum(); + if (AscendC::GetSubBlockIdx() < (tilesNum % AscendC::GetSubBlockNum())) { + tilesPerAiv++; + } + uint64_t taskOffsetSrc = AscendC::GetSubBlockIdx() * tilesPerAiv * tileStrideSrc; + uint64_t taskOffsetDst = AscendC::GetSubBlockIdx() * tilesPerAiv * tileStrideDst; + if (AscendC::GetSubBlockIdx() >= (tilesNum % AscendC::GetSubBlockNum())) { + taskOffsetSrc += (tilesNum % AscendC::GetSubBlockNum()) * tileStrideSrc; + taskOffsetDst += (tilesNum % AscendC::GetSubBlockNum()) * tileStrideDst; + } + uint32_t tilesPerLoop = COMPUTE_LEN / tileLenRoundInt8; + uint32_t loops = CeilDiv(tilesPerAiv, tilesPerLoop); + uint32_t pingpong = 0; + for (uint32_t loopIdx = 0; loopIdx < loops; loopIdx++) { + uint32_t actualTiles = tilesPerLoop; + if (loopIdx == loops - 1) { + actualTiles = tilesPerAiv - loopIdx * tilesPerLoop; + } + uint64_t tileOffsetSrc = loopIdx * tilesPerLoop * tileStrideSrc; + AscendC::DataCopyExtParams dataCopyParamsIn( + actualTiles, tileLen * sizeof(ElementSrc), + (tileStrideSrc - tileLen) * sizeof(ElementSrc), + (tileLenRoundInt8 - tileLen) / ELE_NUM_PER_BLK_INT8, 0); + AscendC::DataCopyPadExtParams padParams(false, 0, 0, 0); + + AscendC::WaitFlag(ubEventList[pingpong]); + AscendC::DataCopyPad(ubInTensorList[pingpong], + gmSrc[taskOffsetSrc + tileOffsetSrc], dataCopyParamsIn, padParams); + + AscendC::SetFlag(ubEventList[pingpong]); + AscendC::WaitFlag(ubEventList[pingpong]); + + AscendC::WaitFlag(ubEventList[pingpong]); + + AscendC::Cast(ubOutTensorList[pingpong], ubInTensorList[pingpong], + AscendC::RoundMode::CAST_NONE, actualTiles * tileLenRoundInt8); + AscendC::PipeBarrier(); + AscendC::SetFlag(ubEventList[pingpong]); + + AscendC::Adds(ubOutTensorList[pingpong], ubOutTensorList[pingpong], + params.deqZeroPoint, actualTiles * tileLenRoundInt8); + AscendC::PipeBarrier(); + + AscendC::Muls(ubOutTensorList[pingpong], ubOutTensorList[pingpong], + params.deqScalar, actualTiles * tileLenRoundInt8); + AscendC::PipeBarrier(); + + AscendC::SetFlag(ubEventList[pingpong]); + AscendC::WaitFlag(ubEventList[pingpong]); + + uint64_t tileOffsetDst = loopIdx * tilesPerLoop * tileStrideDst; + AscendC::DataCopyExtParams dataCopyParamsOut( + actualTiles, tileLen * sizeof(ElementDst), + (tileLenRoundInt8 - tileLen) / ELE_NUM_PER_BLK_HALF, + (tileStrideDst - tileLen) * sizeof(ElementDst), 0); + AscendC::DataCopyPad(gmDst[taskOffsetDst + tileOffsetDst], + ubOutTensorList[pingpong], dataCopyParamsOut); + AscendC::SetFlag(ubEventList[pingpong]); + + pingpong = (pingpong + 1) % STAGES; + } + } + +protected: + /// Data members + AscendC::LocalTensor ubInTensorList[STAGES]; + AscendC::LocalTensor ubOutTensorList[STAGES]; + int32_t ubEventList[STAGES]; + + Params params; +}; + +} // namespace Catlass::Gemm::Tile + +#endif // CATLASS_GEMM_TILE_ATLASA2_KDA_CAST_INT8_TO_FP16_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/copy_l0c_to_ub.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/copy_l0c_to_ub.hpp new file mode 100644 index 0000000000..785a5c0408 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/copy_l0c_to_ub.hpp @@ -0,0 +1,18 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_GEMM_TILE_COPY_L0C_TO_UB_HPP +#define CATLASS_GEMM_TILE_COPY_L0C_TO_UB_HPP + +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 3510) +#include "catlass/gemm/tile/ascend950/copy_l0c_to_ub.hpp" +#endif + +#endif diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_cast_fp8_to_fp16.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_cast_fp8_to_fp16.hpp new file mode 100644 index 0000000000..95e6fa7197 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_cast_fp8_to_fp16.hpp @@ -0,0 +1,18 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_CAST_FP8_TO_FP16_HPP +#define CATLASS_KDA_CAST_FP8_TO_FP16_HPP + +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 2201) +#include "catlass/gemm/tile/atlasa2/kda_cast_fp8_to_fp16.hpp" +#endif + +#endif diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_cast_int4_to_int8.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_cast_int4_to_int8.hpp new file mode 100644 index 0000000000..75c265b5af --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_cast_int4_to_int8.hpp @@ -0,0 +1,18 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_CAST_INT4_TO_INT8_HPP +#define CATLASS_KDA_CAST_INT4_TO_INT8_HPP + +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 2201) +#include "catlass/gemm/tile/atlasa2/kda_atlasa2_cast_int4_to_int8.hpp" +#endif + +#endif diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_cast_int8_to_fp16.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_cast_int8_to_fp16.hpp new file mode 100644 index 0000000000..e99e12754d --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_cast_int8_to_fp16.hpp @@ -0,0 +1,18 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_CAST_INT8_TO_FP16_HPP +#define CATLASS_KDA_CAST_INT8_TO_FP16_HPP + +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 2201) +#include "catlass/gemm/tile/atlasa2/kda_cast_int8_to_fp16.hpp" +#endif + +#endif diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_copy_gm_to_l1.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_copy_gm_to_l1.hpp new file mode 100644 index 0000000000..6097bc8cb6 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_copy_gm_to_l1.hpp @@ -0,0 +1,21 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_COPY_GM_TO_L1_HPP +#define CATLASS_KDA_COPY_GM_TO_L1_HPP + +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 2201) +#include "catlass/gemm/tile/atlasa2/kda_atlasa2_copy_gm_to_l1.hpp" +#endif +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 3510) +#include "catlass/gemm/tile/ascend950/kda_ascend950_copy_gm_to_l1.hpp" +#endif + +#endif diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_copy_gm_to_ub.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_copy_gm_to_ub.hpp new file mode 100644 index 0000000000..013d24d116 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_copy_gm_to_ub.hpp @@ -0,0 +1,18 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_GEMM_COPY_GM_TO_UB_HPP +#define CATLASS_KDA_GEMM_COPY_GM_TO_UB_HPP + +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 2201) +#include "catlass/gemm/tile/atlasa2/kda_atlasa2_copy_gm_to_ub.hpp" +#endif + +#endif diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_copy_l0c_to_gm.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_copy_l0c_to_gm.hpp new file mode 100644 index 0000000000..6bfa7b1945 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_copy_l0c_to_gm.hpp @@ -0,0 +1,21 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_COPY_L0C_TO_GM_HPP +#define CATLASS_KDA_COPY_L0C_TO_GM_HPP + +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 2201) +#include "catlass/gemm/tile/atlasa2/kda_atlasa2_copy_l0c_to_gm.hpp" +#endif +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 3510) +#include "catlass/gemm/tile/ascend950/kda_ascend950_copy_l0c_to_gm.hpp" +#endif + +#endif diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_copy_l1_to_bt.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_copy_l1_to_bt.hpp new file mode 100644 index 0000000000..ba44dbd032 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_copy_l1_to_bt.hpp @@ -0,0 +1,21 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_COPY_L1_TO_BT_HPP +#define CATLASS_KDA_COPY_L1_TO_BT_HPP + +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 2201) +#include "catlass/gemm/tile/atlasa2/kda_atlasa2_copy_l1_to_bt.hpp" +#endif +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 3510) +#include "catlass/gemm/tile/ascend950/kda_ascend950_copy_l1_to_bt.hpp" +#endif + +#endif diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_copy_l1_to_fp.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_copy_l1_to_fp.hpp new file mode 100644 index 0000000000..547871d161 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_copy_l1_to_fp.hpp @@ -0,0 +1,18 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_COPY_L1_TO_FP_HPP +#define CATLASS_KDA_COPY_L1_TO_FP_HPP + +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 2201) +#include "catlass/gemm/tile/atlasa2/kda_atlasa2_copy_l1_to_fp.hpp" +#endif + +#endif diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_copy_l1_to_l0a.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_copy_l1_to_l0a.hpp new file mode 100644 index 0000000000..327897533a --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_copy_l1_to_l0a.hpp @@ -0,0 +1,21 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_COPY_L1_TO_L0A_HPP +#define CATLASS_KDA_COPY_L1_TO_L0A_HPP + +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 2201) +#include "catlass/gemm/tile/atlasa2/kda_atlasa2_copy_l1_to_l0a.hpp" +#endif +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 3510) +#include "catlass/gemm/tile/ascend950/kda_ascend950_copy_l1_to_l0a.hpp" +#endif + +#endif diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_copy_l1_to_l0b.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_copy_l1_to_l0b.hpp new file mode 100644 index 0000000000..97ca958a9a --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_copy_l1_to_l0b.hpp @@ -0,0 +1,21 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_COPY_L1_TO_L0B_HPP +#define CATLASS_KDA_COPY_L1_TO_L0B_HPP + +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 2201) +#include "catlass/gemm/tile/atlasa2/kda_atlasa2_copy_l1_to_l0b.hpp" +#endif +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 3510) +#include "catlass/gemm/tile/ascend950/kda_ascend950_copy_l1_to_l0b.hpp" +#endif + +#endif diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_copy_ub_to_gm.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_copy_ub_to_gm.hpp new file mode 100644 index 0000000000..02d9303a06 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_copy_ub_to_gm.hpp @@ -0,0 +1,18 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_GEMM_COPY_UB_TO_GM_HPP +#define CATLASS_KDA_GEMM_COPY_UB_TO_GM_HPP + +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 2201) +#include "catlass/gemm/tile/atlasa2/kda_atlasa2_copy_ub_to_gm.hpp" +#endif + +#endif diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_gemm_tile_copy.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_gemm_tile_copy.hpp new file mode 100644 index 0000000000..9d463a5a6c --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_gemm_tile_copy.hpp @@ -0,0 +1,520 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_GEMM_TILE_COPY_HPP +#define CATLASS_KDA_GEMM_TILE_COPY_HPP + +#include + +#include "catlass/kda_catlass.hpp" +#include "catlass/detail/kda_tag_to_layout.hpp" +#include "catlass/gemm/kda_gemm_helper.hpp" +#include "catlass/gemm/tile/kda_cast_fp8_to_fp16.hpp" +#include "catlass/gemm/tile/kda_cast_int4_to_int8.hpp" +#include "catlass/gemm/tile/kda_cast_int8_to_fp16.hpp" +#include "catlass/gemm/tile/kda_copy_gm_to_l1.hpp" +#include "catlass/gemm/tile/kda_copy_gm_to_ub.hpp" +#include "catlass/gemm/tile/kda_copy_l0c_to_gm.hpp" +#include "catlass/gemm/tile/copy_l0c_to_ub.hpp" +#include "catlass/gemm/tile/kda_copy_l1_to_bt.hpp" +#include "catlass/gemm/tile/kda_copy_l1_to_fp.hpp" +#include "catlass/gemm/tile/kda_copy_l1_to_l0a.hpp" +#include "catlass/gemm/tile/kda_copy_l1_to_l0b.hpp" +#include "catlass/gemm/tile/kda_copy_ub_to_gm.hpp" +#include "catlass/gemm/tile/kda_tile_copy_tla.hpp" +#include "tla/kda_tensor.hpp" + +namespace Catlass::Gemm::Tile { + +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 2201) +template < + /// Tag indicating architecture + class ArchTag, + /// GemmType for A matrix operand + class AType, + /// GemmType type for B matrix operand + class BType, + /// GemmType type for C matrix operand + class CType, + /// GemmType type for Bias operand + class BiasType = void> +struct TileCopy { + using ElementA = typename AType::Element; + using ElementB = typename BType::Element; + using ElementAccumulator = typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; + + using CopyGmToL1A = Gemm::Tile::CopyGmToL1; + using CopyGmToL1B = Gemm::Tile::CopyGmToL1; + using CopyL1ToL0A = Gemm::Tile::CopyL1ToL0A::L1AType>; + using CopyL1ToL0B = Gemm::Tile::CopyL1ToL0B::L1BType>; + using CopyL0CToGm = Gemm::Tile::CopyL0CToGm; + using BiasTypeSelector = helper::L1BiasTypeSelector; + using CopyGmToL1Bias = std::conditional_t< + std::is_same_v, + void, + Gemm::Tile::CopyGmToL1>; + using CopyL1ToBT = std::conditional_t< + std::is_same_v, + void, + Gemm::Tile::CopyL1ToBT>; +}; + +template +struct TileCopyWithPrologueDeqPerTensor { + using ElementA = typename AType::Element; + using ElementB = typename BType::Element; + + using ElementAccumulator = typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; + + using CopyGmToL1A = Gemm::Tile::CopyGmToL1; + using CopyGmToL1B = Gemm::Tile::CopyGmToL1; + + using PrologueA = PrologueA_; + using PrologueB = PrologueB_; + + using CopyL1ToL0A = Gemm::Tile::CopyL1ToL0A::L1AType>; + using CopyL1ToL0B = Gemm::Tile::CopyL1ToL0B::L1BType>; + using CopyL0CToGm = Gemm::Tile::CopyL0CToGm; + using BiasTypeSelector = helper::L1BiasTypeSelector; + using CopyGmToL1Bias = std::conditional_t< + std::is_same_v, + void, + Gemm::Tile::CopyGmToL1>; + using CopyL1ToBT = std::conditional_t< + std::is_same_v, + void, + Gemm::Tile::CopyL1ToBT>; +}; + +template +struct TileCopyWithPrologue { + using ElementA = typename AType::Element; + using ElementB = typename BType::Element; + + using ElementAccumulator = typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; + + using CopyGmToL1A = Gemm::Tile::CopyGmToL1; + using CopyGmToL1B = Gemm::Tile::CopyGmToL1; + + using PrologueA = PrologueA_; + using PrologueB = PrologueB_; + + using CopyL1ToL0A = Gemm::Tile::CopyL1ToL0A::L1AType>; + using CopyL1ToL0B = Gemm::Tile::CopyL1ToL0B::L1BType>; + using CopyL0CToGm = Gemm::Tile::CopyL0CToGm; + using BiasTypeSelector = helper::L1BiasTypeSelector; + using CopyGmToL1Bias = std::conditional_t< + std::is_same_v, + void, + Gemm::Tile::CopyGmToL1>; + using CopyL1ToBT = std::conditional_t< + std::is_same_v, + void, + Gemm::Tile::CopyL1ToBT>; +}; + +/////////////////////////////////// +/// new add +template < + /// Tag indicating architecture + class ArchTag, + /// GemmType for A matrix operand + class AType, + /// GemmType type for B matrix operand + class BType, + /// GemmType type for C matrix operand + class CType, + /// GemmTpe type for Bias operand + class BiasType = void> +struct TileCopyGemm { + using ElementA = typename AType::Element; + using ElementB = typename BType::Element; + using ElementAccumulator = typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; + // change structual + using L1AType = typename Gemm::helper::L1AndL0TypeSelectorGemm::L1AType; + using L1BType = typename Gemm::helper::L1AndL0TypeSelectorGemm::L1BType; + using L0AType = typename Gemm::helper::L1AndL0TypeSelectorGemm::L0AType; + using L0BType = typename Gemm::helper::L1AndL0TypeSelectorGemm::L0BType; + + using CopyGmToL1A = Gemm::Tile::CopyGmToL1; + using CopyGmToL1B = Gemm::Tile::CopyGmToL1; + using CopyL1ToL0A = Gemm::Tile::CopyL1ToL0A; + using CopyL1ToL0B = Gemm::Tile::CopyL1ToL0B; + using CopyL0CToGm = Gemm::Tile::CopyL0CToGm; +}; +////////////////////////////// +template < + /// Tag indicating architecture + class ArchTag, + /// GemmType for A matrix operand + class AType, + /// GemmType type for B matrix operand + class BType, + /// GemmType type for C matrix operand + class CType, + /// GemmType type for Bias operand + class BiasType> +struct ConvTileCopy { + using ElementA = typename AType::Element; + using ElementB = typename BType::Element; + using ElementAccumulator = typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; + + using CopyGmToL1A = Gemm::Tile::CopyGmToL1; + using CopyGmToL1B = Gemm::Tile::CopyGmToL1; + using CopyL1ToL0A = Gemm::Tile::CopyL1ToL0A::L1AType>; + using CopyL1ToL0B = Gemm::Tile::CopyL1ToL0B::L1BType>; + using CopyL0CToGm = Gemm::Tile::CopyL0CToGm; + using BiasTypeSelector = helper::L1BiasTypeSelector; + using CopyGmToL1Bias = std::conditional_t< + std::is_same_v, + void, + Gemm::Tile::CopyGmToL1>; + using CopyL1ToBT = std::conditional_t< + std::is_same_v, + void, + Gemm::Tile::CopyL1ToBT>; +}; + +// fixpipe开启relu开关 +template < + /// Tag indicating architecture + class ArchTag, + /// GemmType for A matrix operand + class AType, + /// GemmType type for B matrix operand + class BType, + /// GemmType type for C matrix operand + class CType, + /// GemmType type for Bias operand + class BiasType = void> +struct ReluTileCopy : public TileCopy { + // 重写 CopyL0CToGm + using ElementAccumulator = typename TileCopy::ElementAccumulator; + using CopyL0CToGm = Gemm::Tile:: + CopyL0CToGm; +}; + +// fixpipe开启随路量化 +template < + /// Tag indicating architecture + class ArchTag, + /// GemmType for A matrix operand + class AType, + /// GemmType type for B matrix operand + class BType, + /// GemmType type for C matrix operand + class CType, + /// GemmType type for Bias operand + class BiasType = void, + /// GemmType type for Bias operand + ScaleGranularity SCALE_GRANU = ScaleGranularity::PER_TENSOR> +struct QuantTileCopy : public TileCopy { + // 重写 CopyL0CToGm + using ElementAccumulator = typename TileCopy::ElementAccumulator; + using CopyL0CToGm = Gemm::Tile::CopyL0CToGm; + + using CopyGmToL1Scale = Gemm::Tile::CopyGmToL1< + ArchTag, + Gemm::GemmType, + Gemm::GemmType>; + + using CopyL1ToFP = Gemm::Tile::CopyL1ToFP< + ArchTag, + Gemm::GemmType, + Gemm::GemmType>; +}; + +// sparse Tile +template < + /// Tag indicating architecture + class ArchTag, + class ElementA_, + class LayoutTagA, + class ElementB_, + class LayoutTagB, + class ElementC_, + class LayoutTagC> +struct SparseTileCopyTla { + using ElementA = ElementA_; + using ElementB = ElementB_; + using ElementAccumulator = typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; + + using LayoutTagL1A = typename helper::L1ATypeSelector>::L1AType::Layout; + using LayoutTagL1B = typename helper::L1BTypeSelector>::L1BType::Layout; + using LayoutTagL0A = layout::zZ; + using LayoutTagL0B = layout::nZ; + using LayoutTagL1BIdx = typename helper::L1BTypeSelector>::L1BType::Layout; + + using LayoutA = detail::TagToLayout_t; + using LayoutB = detail::TagToLayout_t; + using LayoutC = detail::TagToLayout_t; + + using LayoutL1A = detail::TagToLayout_t; + using LayoutL1B = detail::TagToLayout_t; + using LayoutL1BIdx = detail::TagToLayout_t; + using LayoutL0A = detail::TagToLayout_t; + using LayoutL0B = detail::TagToLayout_t; + using LayoutL0C = typename detail::LayoutL0C; + + using TensorL1A = + tla::Tensor, LayoutL1A, tla::Coord, AscendC::TPosition::A1>; + using TensorL1B = + tla::Tensor, LayoutL1B, tla::Coord, AscendC::TPosition::A1>; + using TensorL1BIdx = + tla::Tensor, LayoutL1BIdx, tla::Coord, AscendC::TPosition::A1>; + using TensorL0A = + tla::Tensor, LayoutL0A, tla::Coord, AscendC::TPosition::A2>; + using TensorL0B = + tla::Tensor, LayoutL0B, tla::Coord, AscendC::TPosition::B2>; + using TensorL0C = tla:: + Tensor, LayoutL0C, tla::Coord, AscendC::TPosition::CO1>; + + using L1AAlignHelper = Gemm::helper::L1AlignHelper; + using L1BAlignHelper = Gemm::helper::L1AlignHelper; + + template + using CopyGmToL1A = Gemm::Tile::TileCopySparseTla; + + template + using CopyGmToL1B = Gemm::Tile::TileCopySparseTla; + + template + using CopyGmToL1BIdx = Gemm::Tile::TileCopySparseTla; + + using CopyL1ToL0A = Gemm::Tile::TileCopySparseTla; + + using CopyL1ToL0B = Gemm::Tile::CopyL1ToL0BSparseTla; + + template + using CopyL0CToGm = Gemm::Tile::CopyL0CToGmSparseTla; +}; + +#endif + +template < + /// Tag indicating architecture + class ArchTag, + class ElementA_, + class LayoutTagA_, + class ElementB_, + class LayoutTagB_, + class ElementC_, + class LayoutTagC_, + class ElementBias = void, + bool ReluEnable_ = false, + ScaleGranularity DEQUANT_GRANULARITY_ = ScaleGranularity::NO_QUANT, + class L0CCopyMode = CopyToGM> +struct PackedTileCopyTla { + using ElementA = ElementA_; + using ElementB = ElementB_; + using LayoutTagA = LayoutTagA_; + using LayoutTagB = LayoutTagB_; + using LayoutTagC = LayoutTagC_; + using ElementAccumulator = typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; + static constexpr bool ReluEnable = ReluEnable_; + static constexpr ScaleGranularity DEQUANT_GRANULARITY = DEQUANT_GRANULARITY_; + + static constexpr bool HAS_BIAS = !std::is_void_v; + static constexpr bool HAS_QUANT_TENSOR = (DEQUANT_GRANULARITY == ScaleGranularity::PER_CHANNEL); + + using LayoutTagL1A = typename helper::L1ATypeSelector>::L1AType::Layout; + using LayoutTagL1B = typename helper::L1BTypeSelector>::L1BType::Layout; + using LayoutTagL0A = typename helper::L0ALayoutSelector::Layout; + using LayoutTagL0B = layout::nZ; + using LayoutTagL0C = layout::L0C; + + using LayoutA = detail::TagToLayout_t; + using LayoutB = detail::TagToLayout_t; + using LayoutC = detail::TagToLayout_t; + + using LayoutL1A = detail::TagToLayout_t; + using LayoutL1B = detail::TagToLayout_t; + using LayoutL0A = detail::TagToLayout_t; + using LayoutL0B = detail::TagToLayout_t; + using LayoutL0C = typename detail::LayoutL0C; + + using TensorL1AVectorLayout = + tla::Tensor, LayoutL1A, tla::Coord, AscendC::TPosition::A1>; + using TensorL1ALayout = + tla::Tensor, LayoutL1A, tla::Coord, AscendC::TPosition::A1>; + + using TensorL1A = + std::conditional_t::value, TensorL1AVectorLayout, TensorL1ALayout>; + using TensorL1B = + tla::Tensor, LayoutL1B, tla::Coord, AscendC::TPosition::A1>; + using TensorL0A = + tla::Tensor, LayoutL0A, tla::Coord, AscendC::TPosition::A2>; + using TensorL0B = + tla::Tensor, LayoutL0B, tla::Coord, AscendC::TPosition::B2>; + using TensorL0C = tla:: + Tensor, LayoutL0C, tla::Coord, AscendC::TPosition::CO1>; + using TensorL1Bias = std::conditional_t< + HAS_BIAS, + tla::Tensor< + AscendC::LocalTensor, + detail::TagToLayout_t, + tla::Coord, + AscendC::TPosition::A1>, + EmptyClass>; + using TensorL0Bias = tla::Tensor< + AscendC::LocalTensor, + detail::TagToLayout_t, + tla::Coord, + AscendC::TPosition::C2>; + using TensorL1Quant = std::conditional_t< + HAS_QUANT_TENSOR, + tla::Tensor, + detail::TagToLayout_t, + tla::Coord, + AscendC::TPosition::A1>, + EmptyClass>; + + using L1AAlignHelper = Gemm::helper::L1AlignHelper; + using L1BAlignHelper = Gemm::helper::L1AlignHelper; + + template + using CopyGmToL1A = Gemm::Tile::TileCopyTla; + + template + using CopyGmToL1B = Gemm::Tile::TileCopyTla; + + template + using CopyGmToL1Bias = + std::conditional_t, EmptyClass>; + + template + using CopyGmToL1Scale = + std::conditional_t, EmptyClass>; + + using CopyL1ToL0A = Gemm::Tile::TileCopyTla; + using CopyL1ToL0B = Gemm::Tile::TileCopyTla; + using CopyL1ToBT = + std::conditional_t, EmptyClass>; + +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 3510) + template + using CopyL0CToDst = Gemm::Tile::CopyL0CToGmTla; +#endif +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 2201) + template + using CopyL0CToGm = Gemm::Tile::CopyL0CToGmTla; +#endif +}; + +template < + /// Tag indicating architecture + class ArchTag, + class TensorA, + class LayoutTagA, + class TensorB, + class LayoutTagB, + class TensorC, + class LayoutTagC, + class TensorBias = void, + class LayoutTagBias = void, + bool IS_PADDING_A = false, + bool IS_PADDING_B = false> +struct PaddingPackedTileCopyTla { + static_assert( + std::is_same_v || std::is_same_v, + "Unsupported layout, only can be RowMajor and ColumnMajor"); + static_assert( + std::is_same_v || std::is_same_v, + "Unsupported layout, only can be RowMajor and ColumnMajor"); + using ElementA = typename TensorA::Element; + using ElementB = typename TensorB::Element; + using ElementAccumulator = typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; + + using LayoutTagL1A = typename helper::L1ATypeSelector>::L1AType::Layout; + using LayoutTagL1B = typename helper::L1BTypeSelector>::L1BType::Layout; + using LayoutTagL0A = layout::zZ; + using LayoutTagL0B = layout::nZ; + + using LayoutL1A = detail::TagToLayout_t; + using LayoutL1B = detail::TagToLayout_t; + using LayoutL0A = detail::TagToLayout_t; + using LayoutL0B = detail::TagToLayout_t; + using LayoutL0C = typename detail::LayoutL0C; + + using TensorL1A = + tla::Tensor, LayoutL1A, tla::Coord, AscendC::TPosition::A1>; + using TensorL1B = + tla::Tensor, LayoutL1B, tla::Coord, AscendC::TPosition::A1>; + using TensorL0A = + tla::Tensor, LayoutL0A, tla::Coord, AscendC::TPosition::A2>; + using TensorL0B = + tla::Tensor, LayoutL0B, tla::Coord, AscendC::TPosition::B2>; + using TensorL0C = tla:: + Tensor, LayoutL0C, tla::Coord, AscendC::TPosition::CO1>; + + using L1AAlignHelper = Gemm::helper::L1AlignHelper; + using L1BAlignHelper = Gemm::helper::L1AlignHelper; + + using LayoutPaddingTagA = std:: + conditional_t, layout::PaddingRowMajor, layout::PaddingColumnMajor>; + using LayoutPaddingTagB = std:: + conditional_t, layout::PaddingRowMajor, layout::PaddingColumnMajor>; + + using CopyGmToL1A = std::conditional_t< + IS_PADDING_A, + Gemm::Tile::TileCopyTlaExt, + Gemm::Tile::TileCopyTlaExt>; + using CopyGmToL1B = std::conditional_t< + IS_PADDING_B, + Gemm::Tile::TileCopyTlaExt, + Gemm::Tile::TileCopyTlaExt>; + + using CopyL1ToL0A = Gemm::Tile::TileCopyTla; + using CopyL1ToL0B = Gemm::Tile::TileCopyTla; +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 3510) + using CopyL0CToDst = Gemm::Tile::CopyL0CToGmTla; +#endif +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 2201) + using CopyL0CToGm = Gemm::Tile::CopyL0CToGmTla; +#endif +}; + +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 3510) +template < + /// Tag indicating architecture + class ArchTag, + class ElementA_, + class LayoutTagA, + class ElementB_, + class LayoutTagB, + class ElementC_, + class LayoutTagC, + class ElementBias = void, + CopyL0CToUBMode CopyMode_ = CopyL0CToUBMode::NO_SPLIT, + bool ReluEnable = false, + ScaleGranularity DEQUANT_GRANULARITY = ScaleGranularity::NO_QUANT> +struct PackedTileCopyTlaToUB + : public PackedTileCopyTla { + static constexpr CopyL0CToUBMode CopyMode = CopyMode_; + // 重写 CopyL0CToDst + using TensorL0C = typename PackedTileCopyTla< + ArchTag, + ElementA_, + LayoutTagA, + ElementB_, + LayoutTagB, + ElementC_, + LayoutTagC, + ElementBias>::TensorL0C; + + template + using CopyL0CToDst = + Gemm::Tile::CopyL0CToUBTla; +}; +#endif + +} // namespace Catlass::Gemm::Tile + +#endif // CATLASS_KDA_GEMM_TILE_COPY_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_tile_copy_tla.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_tile_copy_tla.hpp new file mode 100644 index 0000000000..46e2406361 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_tile_copy_tla.hpp @@ -0,0 +1,70 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_TILE_COPY_TLA_HPP +#define CATLASS_KDA_TILE_COPY_TLA_HPP + +#include "catlass/kda_catlass.hpp" + +namespace Catlass::Gemm::Tile { + +template < + class ArchTag, + class TensorSrc, + class TensorDst, + class Enable = void> +struct TileCopyTla { + static_assert(DEPENDENT_FALSE, "Unsupported TileCopyTla, can not find the specialization."); +}; + +// Extended template for TileCopyTla that supports manually specifying LayoutTagSrc and LayoutTagDst. +// Users can specialize the copy class by LayoutTagSrc and LayoutTagDst. +template < + class ArchTag, + class TensorSrc, + class TensorDst, + class LayoutTagSrc, + class LayoutTagDst> +struct TileCopyTlaExt { + static_assert(DEPENDENT_FALSE, "Unsupported TileCopyTlaExt, can not find the specialization."); +}; + +template < + class ArchTag, + class TensorSrc, + class TensorDst, + class Enable = void> +struct TileCopySparseTla { + static_assert(DEPENDENT_FALSE, "Unsupported TileCopySparseTla, can not find the specialization."); +}; + +template < + class ArchTag, + class Element, + class TensorSrc, + class TensorDst, + class TensorIdx, + class Enable = void> +struct CopyL1ToL0BSparseTla { + static_assert(DEPENDENT_FALSE, "Unsupported CopyL1ToL0BSparseTla, can not find the specialization."); +}; + +// used for FA LoadQ, multi matrix datacopy ND->NZ +template < + class ArchTag, + class TensorSrc, + class TensorDst> +struct TileCopyFAQTla { + static_assert(DEPENDENT_FALSE, "Unsupported TileCopyFAQTla, can not find the specialization."); +}; + +} // namespace Catlass::Gemm::Tile + +#endif // CATLASS_KDA_TILE_COPY_TLA_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_tile_mmad.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_tile_mmad.hpp new file mode 100644 index 0000000000..cca5eca621 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_tile_mmad.hpp @@ -0,0 +1,285 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_TILE_MMAD_HPP +#define CATLASS_KDA_TILE_MMAD_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/gemm/kda_gemm_helper.hpp" +#include "tla/kda_tensor.hpp" +namespace Catlass::Gemm::Tile { + +/////////////////////////////////////////////////////////// + +template < + /// Tag indicating architecture + class ArchTag_, + /// GemmType for A matrix operand + class AType_, + /// GemmType type for B matrix operand + class BType_, + /// GemmType type for Bias operand + class BiasType_> +struct TileMmad { + using ElementA = typename AType_::Element; + using ElementB = typename BType_::Element; + using ElementAccumulator = + typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; + + // Methods + + CATLASS_DEVICE + TileMmad() {} + + CATLASS_DEVICE + void operator()(AscendC::LocalTensor const &l0CTensor, + AscendC::LocalTensor const &l0ATensor, + AscendC::LocalTensor const &l0BTensor, + uint32_t m, uint32_t n, uint32_t k, + bool initC = true, uint8_t unitFlag = 0) + { + AscendC::MmadParams mmadParams; + mmadParams.m = m; + mmadParams.n = n; + mmadParams.k = k; + mmadParams.unitFlag = unitFlag; + mmadParams.cmatrixInitVal = initC; +#if (defined(__NPU_ARCH__) && __NPU_ARCH__ == 2201) + if constexpr (std::is_same_v && + (std::is_same_v || + std::is_same_v)) { + mmadParams.kDirectionAlign = true; + } +#endif +#if (defined(__NPU_ARCH__) && __NPU_ARCH__ == 3510) + if constexpr (std::is_same_v) { + mmadParams.disableGemv = false; + } else { + mmadParams.disableGemv = true; + } +#endif + + AscendC::Mmad(l0CTensor, + l0ATensor, + l0BTensor, + mmadParams); + + const uint32_t PIPE_M_BARRIER_THRESHOLD = 10; + if ((m / C0_NUM_PER_FRACTAL) * (n / C0_NUM_PER_FRACTAL) < PIPE_M_BARRIER_THRESHOLD) { + AscendC::PipeBarrier(); + } + } + + CATLASS_DEVICE + void operator()(AscendC::LocalTensor const &l0CTensor, + AscendC::LocalTensor const &l0ATensor, + AscendC::LocalTensor const &l0BTensor, + AscendC::LocalTensor const &l0BiasTensor, + uint32_t m, uint32_t n, uint32_t k, + bool initC = true, uint8_t unitFlag = 0) + { + AscendC::MmadParams mmadParams; + mmadParams.m = m; + mmadParams.n = n; + mmadParams.k = k; + mmadParams.unitFlag = unitFlag; + mmadParams.cmatrixInitVal = false; +#if (defined(__NPU_ARCH__) && __NPU_ARCH__ == 2201) + using LayoutL1A = typename helper::L1ATypeSelector::L1AType::Layout; + if constexpr (std::is_same_v && std::is_same_v) { + mmadParams.kDirectionAlign = true; + } +#endif +#if (defined(__NPU_ARCH__) && __NPU_ARCH__ == 3510) + mmadParams.disableGemv = true; +#endif + + AscendC::Mmad(l0CTensor, + l0ATensor, + l0BTensor, + l0BiasTensor, + mmadParams); + + const uint32_t PIPE_M_BARRIER_THRESHOLD = 10; + if ((m / C0_NUM_PER_FRACTAL) * (n / C0_NUM_PER_FRACTAL) < PIPE_M_BARRIER_THRESHOLD) { + AscendC::PipeBarrier(); + } + } +}; + +///////////////////////////////////////////TileMmadTla///////////////////////////////////////////////// + +template < + /// Tag indicating architecture + class ArchTag, + /// Element for A matrix operand + class ElementA, + /// LayoutTag for A matrix operand in L1 + class LayoutTagL1A> +struct TileMmadTla { + // Methods + + CATLASS_DEVICE + TileMmadTla() {} + + template + CATLASS_DEVICE void operator()(TensorC const &l0CTensor, + TensorA const &l0ATensor, + TensorB const &l0BTensor, + uint32_t m, uint32_t n, uint32_t k, + bool initC = true, uint8_t unitFlag = 0) + { + AscendC::MmadParams mmadParams; + mmadParams.m = m; + mmadParams.n = n; + mmadParams.k = k; + mmadParams.unitFlag = unitFlag; + mmadParams.cmatrixInitVal = initC; +#if (defined(__NPU_ARCH__) && __NPU_ARCH__ == 2201) + if constexpr (std::is_same_v && std::is_same_v) { + mmadParams.kDirectionAlign = true; + } +#endif +#if (defined(__NPU_ARCH__) && __NPU_ARCH__ == 3510) + if constexpr (std::is_same_v) { + mmadParams.disableGemv = false; + } else { + mmadParams.disableGemv = true; + } +#endif + + AscendC::Mmad(l0CTensor.data(), + l0ATensor.data(), + l0BTensor.data(), + mmadParams); + + const uint32_t PIPE_M_BARRIER_THRESHOLD = 10; + if ((m / C0_NUM_PER_FRACTAL) * (n / C0_NUM_PER_FRACTAL) < PIPE_M_BARRIER_THRESHOLD) { + AscendC::PipeBarrier(); + } + } + + template + CATLASS_DEVICE void operator()(TensorC const &l0CTensor, + TensorA const &l0ATensor, + TensorB const &l0BTensor, + TensorBias const &l0BiasTensor, + uint32_t m, uint32_t n, uint32_t k, + bool initC = true, uint8_t unitFlag = 0) + { + AscendC::MmadParams mmadParams; + mmadParams.m = m; + mmadParams.n = n; + mmadParams.k = k; + mmadParams.unitFlag = unitFlag; + mmadParams.cmatrixInitVal = false; +#if (defined(__NPU_ARCH__) && __NPU_ARCH__ == 2201) + if constexpr (std::is_same_v && std::is_same_v) { + mmadParams.kDirectionAlign = true; + } +#endif +#if (defined(__NPU_ARCH__) && __NPU_ARCH__ == 3510) + mmadParams.disableGemv = true; +#endif + + AscendC::Mmad(l0CTensor.data(), + l0ATensor.data(), + l0BTensor.data(), + l0BiasTensor.data(), + mmadParams); + + const uint32_t PIPE_M_BARRIER_THRESHOLD = 10; + if ((m / C0_NUM_PER_FRACTAL) * (n / C0_NUM_PER_FRACTAL) < PIPE_M_BARRIER_THRESHOLD) { + AscendC::PipeBarrier(); + } + } + + template + CATLASS_DEVICE void operator()(TensorC const &l0CTensor, + TensorA const &l0ATensor, + TensorB const &l0BTensor, + uint32_t m, uint32_t n, uint32_t k, + uint32_t l0Batch) + { + const uint32_t L0AM = tla::get<0, 0>(l0ATensor.shape()) * tla::get<0, 1>(l0ATensor.shape()); + const uint32_t L0AK = tla::get<1, 0>(l0ATensor.shape()) * tla::get<1, 1>(l0ATensor.shape()); + const uint32_t L0BK = tla::get<0, 0>(l0BTensor.shape()) * tla::get<0, 1>(l0BTensor.shape()); + const uint32_t L0BN = tla::get<1, 0>(l0BTensor.shape()) * tla::get<1, 1>(l0BTensor.shape()); + const uint32_t L0CM = tla::get<0, 0>(l0CTensor.shape()) * tla::get<0, 1>(l0CTensor.shape()); + const uint32_t L0CN = tla::get<1, 0>(l0CTensor.shape()) * tla::get<1, 1>(l0CTensor.shape()); + + AscendC::MmadParams mmadParams; + mmadParams.m = m; + mmadParams.n = n; + mmadParams.k = k; + mmadParams.unitFlag = 0; + mmadParams.cmatrixInitVal = true; +#if (defined(__NPU_ARCH__) && __NPU_ARCH__ == 3510) + mmadParams.disableGemv = true; +#endif + for (uint32_t l0BatchIdx = 0; l0BatchIdx < l0Batch; l0BatchIdx++) { + AscendC::Mmad(l0CTensor.data()[l0BatchIdx * L0CM * L0CN], + l0ATensor.data()[l0BatchIdx * L0AM * L0AK], + l0BTensor.data()[l0BatchIdx * L0BK * L0BN], + mmadParams); + } + } + + // automatically extracts actual sizes from tensor originShape + template + CATLASS_DEVICE void operator()(TensorC const &l0CTensor, + TensorA const &l0ATensor, + TensorB const &l0BTensor, + bool initC = true, uint8_t unitFlag = 0) + { + // Get actual sizes from tensor originShape + uint32_t m = tla::get<0>(l0CTensor.layout().originShape()); + uint32_t n = tla::get<1>(l0CTensor.layout().originShape()); + uint32_t k = tla::get<1>(l0ATensor.layout().originShape()); + AscendC::MmadParams mmadParams; +#if (defined(__NPU_ARCH__) && __NPU_ARCH__ == 2201) + if constexpr (std::is_same_v && std::is_same_v) { + mmadParams.kDirectionAlign = true; + } + if constexpr (!std::is_same_v) { + if (m == 1) + m = 16; // avoid gemv mode + } +#endif +#if (defined(__NPU_ARCH__) && __NPU_ARCH__ == 3510) + if constexpr (std::is_same_v) { + mmadParams.disableGemv = false; + } else { + mmadParams.disableGemv = true; + } +#endif + mmadParams.m = m; + mmadParams.n = n; + mmadParams.k = k; + mmadParams.unitFlag = unitFlag; + mmadParams.cmatrixInitVal = initC; + + AscendC::Mmad(l0CTensor.data(), + l0ATensor.data(), + l0BTensor.data(), + mmadParams); + + const uint32_t PIPE_M_BARRIER_THRESHOLD = 10; + if ((m / C0_NUM_PER_FRACTAL) * (n / C0_NUM_PER_FRACTAL) < PIPE_M_BARRIER_THRESHOLD) { + AscendC::PipeBarrier(); + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace Catlass::Gemm::Tile + +#endif // CATLASS_KDA_TILE_MMAD_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_tile_muls.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_tile_muls.hpp new file mode 100644 index 0000000000..bf4c313aa7 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_tile_muls.hpp @@ -0,0 +1,49 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ +#ifndef CATLASS_KDA_TILE_MULS_HPP +#define CATLASS_KDA_TILE_MULS_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/gemm/kda_gemm_helper.hpp" +#include "tla/kda_tensor.hpp" +namespace Catlass::Gemm::Tile { + +/////////////////////////////////////////////////////// +template +struct TileMuls { + using ArchTag = ArchTag_; + using Element = typename ComputeType_::Element; + + static constexpr uint32_t COMPUTE_LENGTH = COMPUTE_LENGTH_; + + CATLASS_DEVICE + TileMuls() + { + } + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor dstTensor, + AscendC::LocalTensor srcTensor, + Element scalar, + uint32_t len) + { + AscendC::SetMaskCount(); + AscendC::SetVectorMask(len); + AscendC::Muls( + dstTensor, srcTensor, scalar, AscendC::MASK_PLACEHOLDER, 1, AscendC::UnaryRepeatParams{}); + AscendC::SetMaskNorm(); + AscendC::ResetMask(); + } +}; +/////////////////////////////////////////////////////// +} // namespace Catlass::Gemm::Tile + +#endif // CATLASS_KDA_TILE_MULS_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_tile_traits.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_tile_traits.hpp new file mode 100644 index 0000000000..8a1c3017ad --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_tile_traits.hpp @@ -0,0 +1,48 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_TILE_TRAITS_HPP +#define CATLASS_KDA_TILE_TRAITS_HPP + +#include "catlass/kda_catlass.hpp" + +namespace Catlass::Gemm::Tile { + +struct EmptyType {}; + +template +struct PrologueTraits : public Prologue { + using Prologue::Prologue; + + using TensorSrc = AscendC::GlobalTensor; + using TensorDst = AscendC::GlobalTensor; +}; + +template <> +struct PrologueTraits { + using ElementSrc = EmptyType; + using LayoutSrc = EmptyType; + using ElementDst = EmptyType; + using LayoutDst = EmptyType; + + using TensorSrc = EmptyType; + using TensorDst = EmptyType; + + using Params = EmptyType; + + template + CATLASS_DEVICE + PrologueTraits(Args...) + {} +}; + +} // namespace Catlass::Gemm::Tile + +#endif diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemv/kda_gemv_helper.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemv/kda_gemv_helper.hpp new file mode 100644 index 0000000000..aef1fb9e18 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemv/kda_gemv_helper.hpp @@ -0,0 +1,122 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_GEMV_HELPER_HPP +#define CATLASS_KDA_GEMV_HELPER_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/layout/kda_layout.hpp" +#include "catlass/gemm/kda_gemm_type.hpp" + +namespace Catlass::Gemv::helper { + +template +struct UBAlignHelper { + static constexpr uint32_t ALIGN = BYTE_PER_BLK / sizeof(Element); +}; + +template +struct AtomicAddSelector { + static_assert(DEPENDENT_FALSE, + "Unsupported layout selector, can not find the specialization."); +}; + +template +struct AtomicAddSelector> { + static constexpr bool value = false; +}; + +template +struct AtomicAddSelector> { + static constexpr bool value = true; +}; + +template +struct L1AlignHelper { + static_assert(DEPENDENT_FALSE, "Unsupported align helper, can not find the specialization."); +}; + +template +struct L1AlignHelper { + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + static constexpr uint32_t M_ALIGNED = C0_NUM_PER_FRACTAL; + static constexpr uint32_t N_ALIGNED = ELE_NUM_PER_C0; +}; + +template +struct L1AlignHelper { + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + + static constexpr uint32_t getNAligned() + { + if constexpr (std::is_same::value) { + return ELE_NUM_PER_C0 / sizeof(Element); + } else { + return C0_NUM_PER_FRACTAL; + } + } + + static constexpr uint32_t getMAligned() + { + if constexpr (std::is_same::value) { + return ELE_NUM_PER_C0 / sizeof(Element); + } else { + return C0_NUM_PER_FRACTAL; + } + } + + static constexpr uint32_t N_ALIGNED = getNAligned(); + static constexpr uint32_t M_ALIGNED = getMAligned(); +}; + +template +struct L1AlignHelper { + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + static constexpr uint32_t M_ALIGNED = C0_NUM_PER_FRACTAL; + static constexpr uint32_t N_ALIGNED = ELE_NUM_PER_C0; +}; + +//////////////////////////////// +// new add gemvaic selector +template +struct L1AndL0TypeSelectorGemv { + static_assert(DEPENDENT_FALSE, + "Unsupported layout selector, can not find the specialization."); + static_assert(DEPENDENT_FALSE, + "Unsupported layout selector, can not find the specialization."); +}; + +template +struct L1AndL0TypeSelectorGemv, Gemm::GemmType> { + using L1AType = Gemm::GemmType; + using L1BType = Gemm::GemmType; + using L0AType = Gemm::GemmType; + using L0BType = Gemm::GemmType; +}; + +template +struct L1AndL0TypeSelectorGemv, Gemm::GemmType> { + using L1AType = Gemm::GemmType; + using L1BType = Gemm::GemmType; + using L0AType = Gemm::GemmType; + using L0BType = Gemm::GemmType; +}; + +template <> +struct L1AndL0TypeSelectorGemv, Gemm::GemmType> { + using L1AType = Gemm::GemmType; + using L1BType = Gemm::GemmType; + using L0AType = Gemm::GemmType; + using L0BType = Gemm::GemmType; +}; + +} // namespace Catlass::Gemv::helper + +#endif // CATLASS_KDA_GEMV_HELPER_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/kda_catlass.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/kda_catlass.hpp new file mode 100644 index 0000000000..fd8c3b6845 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/kda_catlass.hpp @@ -0,0 +1,55 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_CATLASS_HPP +#define CATLASS_KDA_CATLASS_HPP + +#include + +#if defined(__CCE__) +#include +#endif + +#include "catlass/detail/kda_alignment.hpp" +#include "catlass/detail/kda_dependent_false.hpp" +#include "catlass/detail/kda_macros.hpp" + +namespace Catlass { + +constexpr uint32_t BYTE_PER_C0 = 32; +constexpr uint32_t BYTE_PER_C2 = 64; +constexpr uint32_t C0_NUM_PER_FRACTAL = 16; +constexpr uint32_t BYTE_PER_FRACTAL = BYTE_PER_C0 * C0_NUM_PER_FRACTAL; + +constexpr uint32_t BYTE_PER_BLK = 32; +constexpr uint32_t BLK_NUM_PER_VECTOR_FRACTAL = 8; +constexpr uint32_t BYTE_PER_VECTOR_FRACTAL = BYTE_PER_BLK * BLK_NUM_PER_VECTOR_FRACTAL; + +constexpr uint64_t L2_OFFSET = 0; +constexpr uint32_t STRIDE_LIMIT = 65536; + +constexpr uint32_t BYTE_PER_BLK_FP = 128; /// datablock size of A1->C2PiPE2GM + +class EmptyClass {}; + +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 3510) +constexpr uint32_t MX_SCALE_COPY_GROUP_NUM = 2; // Mx-scale matrix 2-byte aligned +constexpr uint32_t MX_SCALE_GROUP_NUM = 32; // Data count for one MX-scale factor per group +constexpr uint32_t MX_BASEK_FACTOR = 64; // Data matrix alignment at K-dimension +#endif +} // namespace Catlass + +#if defined(__CCE__) && defined(L2_CACHE_HINT) && defined(CATLASS_BUILD_LEGACY) +inline __gm__ struct OpSystemRunCfg g_opSystemRunCfg { + Catlass::L2_OFFSET +}; +#endif + +#endif // CATLASS_KDA_CATLASS_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/kda_conv_coord.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/kda_conv_coord.hpp new file mode 100644 index 0000000000..d179ef0510 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/kda_conv_coord.hpp @@ -0,0 +1,1320 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_CONV_COORD_HPP +#define CATLASS_KDA_CONV_COORD_HPP + +#include "catlass/kda_coord.hpp" + +namespace Catlass { + +/// Shape of conv3d operation +struct Conv3dParams { +public: + typedef uint32_t Index; + static constexpr uint32_t N0 = 16; + using Fmap6HDShape = Coord<6, Index>; // {batch, di, cin1, hi, wi, cin0} + using FilterFracZ3DShape = Coord<7, Index>; // {kd, cin1, kh, kw, n1, n0, cin0} + using Out6HDShape = Coord<6, Index>; // {batch, do, cout1, ho, wo, cout0} + using Strides = Coord<3, Index>; + using Pads = Coord<3, Index>; + using Dilations = Coord<3, Index>; + +private: + Fmap6HDShape fmap6HDShape_; + FilterFracZ3DShape filterFracZ3DShape_; + Out6HDShape out6HDShape_; + Strides strides_; + Pads pads_; + Dilations dilations_; + Index cout_; + +public: + CATLASS_HOST_DEVICE + Conv3dParams( + Index BATCH = 1, + Index Di = 1, + Index Cin1 = 1, + Index Hi = 1, + Index Wi = 1, + Index C0 = 16, + Index Kd = 1, + Index Kh = 1, + Index Kw = 1, + Index N1 = 1, + Index Do = 1, + Index Ho = 1, + Index Wo = 1, + Index Cout1 = 1, + Index Cout = 1, + Index padHead = 0, + Index padTop = 0, + Index padLeft = 0, + Index strideD = 1, + Index strideH = 1, + Index strideW = 1, + Index dilationD = 1, + Index dilationH = 1, + Index dilationW = 1) + : fmap6HDShape_(MakeCoord(BATCH, Di, Cin1, Hi, Wi, C0)), + filterFracZ3DShape_(MakeCoord(Kd, Cin1, Kh, Kw, N1, N0, C0)), + out6HDShape_(MakeCoord(BATCH, Do, Cout1, Ho, Wo, C0)), + cout_(Cout), + pads_(MakeCoord(padHead, padTop, padLeft)), + strides_(MakeCoord(strideD, strideH, strideW)), + dilations_(MakeCoord(dilationD, dilationH, dilationW)) + {} + + CATLASS_HOST_DEVICE + static Conv3dParams MakeConvCoord( + const uint32_t *fmapShape, + const uint32_t *filterShape, + const uint32_t *paddings, + const uint32_t *strides, + const uint32_t *dilations) + { + return Conv3dParams( + fmapShape[0], + fmapShape[1], + fmapShape[2], + fmapShape[3], + fmapShape[4], + fmapShape[5], + filterShape[0], + filterShape[1], + filterShape[2], + CeilDiv(filterShape[3], N0), + (fmapShape[1] + paddings[0] * 2 - dilations[0] * (filterShape[0] - 1) - 1) / strides[0] + 1, // Do + (fmapShape[3] + paddings[1] * 2 - dilations[1] * (filterShape[1] - 1) - 1) / strides[1] + 1, // Ho + (fmapShape[4] + paddings[2] * 2 - dilations[2] * (filterShape[2] - 1) - 1) / strides[2] + 1, // Wo + CeilDiv(filterShape[3], fmapShape[5]), + filterShape[3], + paddings[0], + paddings[1], + paddings[2], + strides[0], + strides[1], + strides[2], + dilations[0], + dilations[1], + dilations[2]); + } + + // fmapShape + CATLASS_HOST_DEVICE + Index const &batch() const { return fmap6HDShape_[0]; } + CATLASS_HOST_DEVICE + Index const &cin1() const { return fmap6HDShape_[2]; } + CATLASS_HOST_DEVICE + Index const &di() const { return fmap6HDShape_[1]; } + CATLASS_HOST_DEVICE + Index const &hi() const { return fmap6HDShape_[3]; } + CATLASS_HOST_DEVICE + Index const &wi() const { return fmap6HDShape_[4]; } + CATLASS_HOST_DEVICE + Index const &cin0() const { return fmap6HDShape_[5]; } + CATLASS_HOST_DEVICE + Index const hiwi() const { return fmap6HDShape_[3] * fmap6HDShape_[4]; } + + // filterShape + CATLASS_HOST_DEVICE + Index const &kd() const { return filterFracZ3DShape_[0]; } + CATLASS_HOST_DEVICE + Index const &kh() const { return filterFracZ3DShape_[2]; } + CATLASS_HOST_DEVICE + Index const &kw() const { return filterFracZ3DShape_[3]; } + CATLASS_HOST_DEVICE + Index const khkw() const + { + return filterFracZ3DShape_[2] * filterFracZ3DShape_[3]; + } + CATLASS_HOST_DEVICE + Index const kdc1khkw() const + { + return filterFracZ3DShape_[0] * filterFracZ3DShape_[1] * filterFracZ3DShape_[2] * filterFracZ3DShape_[3]; + } + CATLASS_HOST_DEVICE + Index const &n1() const { return filterFracZ3DShape_[4]; } + CATLASS_HOST_DEVICE + Index const &n0() const { return filterFracZ3DShape_[5]; } + + // outShape + CATLASS_HOST_DEVICE + Index const &dout() const { return out6HDShape_[1]; } + CATLASS_HOST_DEVICE + Index const &ho() const { return out6HDShape_[3]; } + CATLASS_HOST_DEVICE + Index const &wo() const { return out6HDShape_[4]; } + CATLASS_HOST_DEVICE + Index const &cout1() const { return out6HDShape_[2]; } + CATLASS_HOST_DEVICE + Index const &cout0() const { return out6HDShape_[5]; } + CATLASS_HOST_DEVICE + Index const &cout() const { return cout_; } + + /// paddings + CATLASS_HOST_DEVICE + Index const &padhead() const { return pads_[0]; } + CATLASS_HOST_DEVICE + Index const &padtail() const { return pads_[0]; } + CATLASS_HOST_DEVICE + Index const &padtop() const { return pads_[1]; } + CATLASS_HOST_DEVICE + Index const &padbottom() const { return pads_[1]; } + CATLASS_HOST_DEVICE + Index const &padleft() const { return pads_[2]; } + CATLASS_HOST_DEVICE + Index const &padright() const { return pads_[2]; } + + /// strideSize + CATLASS_HOST_DEVICE + Index const &sD() const { return strides_[0]; } + CATLASS_HOST_DEVICE + Index const &sH() const { return strides_[1]; } + CATLASS_HOST_DEVICE + Index const &sW() const { return strides_[2]; } + + /// dilationSize + CATLASS_HOST_DEVICE + Index const &dD() const { return dilations_[0]; } + CATLASS_HOST_DEVICE + Index const dilatedKernelD() const { return 1 + (filterFracZ3DShape_[0] - 1) * dilations_[0]; } + CATLASS_HOST_DEVICE + Index const &dH() const { return dilations_[1]; } + CATLASS_HOST_DEVICE + Index const dilatedKernelH() const { return 1 + (filterFracZ3DShape_[2] - 1) * dilations_[1]; } + CATLASS_HOST_DEVICE + Index const &dW() const { return dilations_[2]; } + CATLASS_HOST_DEVICE + Index const dilatedKernelW() const { return 1 + (filterFracZ3DShape_[3] - 1) * dilations_[2]; } + + ///// used in block + CATLASS_HOST_DEVICE + Index const howo() const { return out6HDShape_[3] * out6HDShape_[4]; } + CATLASS_HOST_DEVICE + Index const alignCout() const { return out6HDShape_[2] * out6HDShape_[5]; } + CATLASS_HOST_DEVICE + Index const wicin0() const { return fmap6HDShape_[4] * fmap6HDShape_[5]; } + CATLASS_HOST_DEVICE + Index const khkwcin0() const { return filterFracZ3DShape_[2] * filterFracZ3DShape_[3] * filterFracZ3DShape_[6]; } + CATLASS_HOST_DEVICE + Index const alignCinKhKwKd() const + { + return filterFracZ3DShape_[0] * filterFracZ3DShape_[1] * filterFracZ3DShape_[2] * filterFracZ3DShape_[3] * filterFracZ3DShape_[6]; + } + CATLASS_HOST_DEVICE + Index const kdcin1() const { return filterFracZ3DShape_[0] * filterFracZ3DShape_[1]; } + CATLASS_HOST_DEVICE + Index const fmapOneBatchSize() const + { + return fmap6HDShape_[1] * fmap6HDShape_[2] * fmap6HDShape_[3] * fmap6HDShape_[4] * fmap6HDShape_[5]; + } + CATLASS_HOST_DEVICE + Index const outputOneBatchSize() const + { + return out6HDShape_[1] * out6HDShape_[2] * out6HDShape_[3] * out6HDShape_[4] * out6HDShape_[5]; + } +}; + +template < + uint32_t noCnt_ = 1, + uint32_t doCnt_ = 1, + uint32_t co1Cnt_ = 1, + uint32_t howoCnt_ = 1> +struct ConvCoreShape { + static uint32_t const noCnt = noCnt_; + static uint32_t const doCnt = doCnt_; + static uint32_t const co1Cnt = co1Cnt_; + static uint32_t const howoCnt = howoCnt_; + + /// Returns a Coord object + CATLASS_HOST_DEVICE + static Coord<4> ToCoord() + { + return MakeCoord(noCnt, doCnt, co1Cnt, howoCnt); + } +}; + +template < + uint32_t mAL1_ = 1, + uint32_t Kd_ = 1, + uint32_t Ci1_ = 1> +struct ConvFmapL1Shape { + static uint32_t constexpr mAL1 = mAL1_; + static uint32_t constexpr Kd = Kd_; + static uint32_t constexpr Ci1 = Ci1_; + + /// Returns a Coord object + CATLASS_HOST_DEVICE + static Coord<3> ToCoord() + { + return MakeCoord(mAL1, Kd, Ci1); + } +}; + +template < + uint32_t Kd_ = 1, + uint32_t Ci1_ = 1, + uint32_t nBL1_ = 1> +struct ConvFilterL1Shape { + static uint32_t constexpr Kd = Kd_; + static uint32_t constexpr Ci1 = Ci1_; + static uint32_t constexpr nBL1 = nBL1_; + + /// Returns a Coord object + CATLASS_HOST_DEVICE + static Coord<3> ToCoord() + { + return MakeCoord(Kd, Ci1, nBL1); + } +}; + +template < + uint32_t mL0_ = 1, + uint32_t kL0_ = 1, + uint32_t nL0_ = 1> +struct ConvL0Shape { + static uint32_t constexpr mL0 = mL0_; + static uint32_t constexpr kL0 = kL0_; + static uint32_t constexpr nL0 = nL0_; + + /// Returns a Coord object + CATLASS_HOST_DEVICE + static Coord<3> ToCoord() + { + return MakeCoord(mL0, kL0, nL0); + } +}; + +struct Conv3d6HdCoord : public Coord<4, uint32_t> { + using Index = uint32_t; + + using Base = Coord<4, Index>; + + static constexpr int N_INDEX = 0; + static constexpr int D_INDEX = 1; + static constexpr int C1_INDEX = 2; + static constexpr int HW_INDEX = 3; + + /// Default ctor + CATLASS_HOST_DEVICE + Conv3d6HdCoord() {} + + CATLASS_HOST_DEVICE + Conv3d6HdCoord(Coord<4, Index> const &coord) + : Base(coord) + {} + + CATLASS_HOST_DEVICE + Conv3d6HdCoord(Index n, Index d, Index c1, Index hw) + : Base(MakeCoord(n, d, c1, hw)) + {} + + CATLASS_HOST_DEVICE + Index const &n() const + { + return this->At(N_INDEX); + } + CATLASS_HOST_DEVICE + Index &n() + { + return this->At(N_INDEX); + } + + CATLASS_HOST_DEVICE + Index const &d() const + { + return this->At(D_INDEX); + } + CATLASS_HOST_DEVICE + Index &d() + { + return this->At(D_INDEX); + } + + CATLASS_HOST_DEVICE + Index const &c1() const + { + return this->At(C1_INDEX); + } + CATLASS_HOST_DEVICE + Index &c1() + { + return this->At(C1_INDEX); + } + + CATLASS_HOST_DEVICE + Index const &hw() const + { + return this->At(HW_INDEX); + } + CATLASS_HOST_DEVICE + Index &hw() + { + return this->At(HW_INDEX); + } +}; + +struct Conv3dFracZ3dCoord : public Coord<2, uint32_t> { + using Index = uint32_t; + + using Base = Coord<2, Index>; + + static constexpr int KDC1KHKW_INDEX = 0; + static constexpr int N1_INDEX = 1; + + /// Default ctor + CATLASS_HOST_DEVICE + Conv3dFracZ3dCoord() {} + + CATLASS_HOST_DEVICE + Conv3dFracZ3dCoord(Index kdc1khkw, Index n1) + : Base(MakeCoord(kdc1khkw, n1)) + {} + + CATLASS_HOST_DEVICE + Index const &kdc1khkw() const + { + return this->At(KDC1KHKW_INDEX); + } + CATLASS_HOST_DEVICE + Index &kdc1khkw() + { + return this->At(KDC1KHKW_INDEX); + } + + CATLASS_HOST_DEVICE + Index const &n1() const + { + return this->At(N1_INDEX); + } + CATLASS_HOST_DEVICE + Index &n1() + { + return this->At(N1_INDEX); + } +}; + +/////////////////// Shapes and Coords for Conv2d /////////////////// + +template < + uint32_t Ho_ = 1, + uint32_t Wo_ = 1, + uint32_t Cin1_ = 1> +struct Conv2dFmapL1Shape { // (Ho, Wo, Cin1) + static constexpr uint32_t Ho = Ho_; + static constexpr uint32_t Wo = Wo_; + static constexpr uint32_t Cin1 = Cin1_; + + /// Returns a Coord object + CATLASS_HOST_DEVICE + static Coord<3> ToCoord() + { + return MakeCoord(Ho, Wo, Cin1); + } +}; + +template < + uint32_t Cout_ = 16, + uint32_t Cin1_ = 1> +struct Conv2dFilterL1Shape { // (Cout, Cin1) + static constexpr uint32_t Cout = Cout_; + static constexpr uint32_t Cin1 = Cin1_; + + /// Returns a Coord object + CATLASS_HOST_DEVICE + static Coord<2> ToCoord() + { + return MakeCoord(Cout, Cin1); + } +}; + +template < + uint32_t M_ = 16, + uint32_t N_ = 16, + uint32_t K_ = 16> +struct Conv2dL0Shape { + static constexpr uint32_t M = M_; + static constexpr uint32_t N = N_; + static constexpr uint32_t K = K_; + + /// Returns a Coord object + CATLASS_HOST_DEVICE + static Coord<3> ToCoord() + { + return MakeCoord(M, N, K); + } +}; + +struct Conv2dFmapCoord : public Coord<5, uint32_t> { // (Batch, C1, H, W, C0) +public: + /// Integer-valued index + using Index = uint32_t; + + /// Base type is a Coord of rank=5 + using Base = Coord<5, Index>; + + /// LongIndex type + using LongIndex = typename Base::LongIndex; + + /// dimensions + static constexpr uint32_t BATCH_INDEX = 0; + static constexpr uint32_t C1_INDEX = 1; + static constexpr uint32_t H_INDEX = 2; + static constexpr uint32_t W_INDEX = 3; + static constexpr uint32_t C0_INDEX = 4; + + /// Default ctor + CATLASS_HOST_DEVICE + Conv2dFmapCoord() {} + + /// Constructs from Coord<5> + CATLASS_HOST_DEVICE + Conv2dFmapCoord(Coord<5, Index> const &coord) + : Base(coord) + {} + + /// Helper to construct from C1, H, W, C0 + CATLASS_HOST_DEVICE + Conv2dFmapCoord(Index batch, Index c1, Index h, Index w, Index c0) + : Base(MakeCoord(batch, c1, h, w, c0)) + {} + + CATLASS_HOST_DEVICE + Conv2dFmapCoord(LongIndex batch, LongIndex c1, LongIndex h, LongIndex w, LongIndex c0) + : Base(MakeCoord(Index(batch), Index(c1), Index(h), Index(w), Index(c0))) + {} + + CATLASS_HOST_DEVICE + Index const &c1() const + { + return this->At(C1_INDEX); + } + CATLASS_HOST_DEVICE + Index &c1() + { + return this->At(C1_INDEX); + } + + CATLASS_HOST_DEVICE + Index const &batch() const + { + return this->At(BATCH_INDEX); + } + CATLASS_HOST_DEVICE + Index &batch() + { + return this->At(BATCH_INDEX); + } + + CATLASS_HOST_DEVICE + Index const &h() const + { + return this->At(H_INDEX); + } + CATLASS_HOST_DEVICE + Index &h() + { + return this->At(H_INDEX); + } + + CATLASS_HOST_DEVICE + Index const &w() const + { + return this->At(W_INDEX); + } + CATLASS_HOST_DEVICE + Index &w() + { + return this->At(W_INDEX); + } + + CATLASS_HOST_DEVICE + Index const &c0() const + { + return this->At(C0_INDEX); + } + CATLASS_HOST_DEVICE + Index &c0() + { + return this->At(C0_INDEX); + } + + /// Element-wise addition + CATLASS_HOST_DEVICE + Conv2dFmapCoord operator+(Base const &b) const + { + return Conv2dFmapCoord(Base::operator+(b)); + } + + /// In-place addition + CATLASS_HOST_DEVICE + Conv2dFmapCoord &operator+=(Base const &b) + { + Base::operator+=(b); + return *this; + } +}; + +struct Conv2dFilterCoord : public Coord<5, uint32_t> { // (Cin1, Kh, Kw, Cout, C0) +public: + /// Integer-valued index + using Index = uint32_t; + + /// Base type is a Coord of rank=5 + using Base = Coord<5, Index>; + + /// LongIndex type + using LongIndex = typename Base::LongIndex; + + /// dimensions + static constexpr uint32_t CIN1_INDEX = 0; + static constexpr uint32_t KH_INDEX = 1; + static constexpr uint32_t KW_INDEX = 2; + static constexpr uint32_t COUT_INDEX = 3; + static constexpr uint32_t C0_INDEX = 4; + + /// Default ctor + CATLASS_HOST_DEVICE + Conv2dFilterCoord() {} + + /// Constructs from Coord<5> + CATLASS_HOST_DEVICE + Conv2dFilterCoord(Coord<5, Index> const &coord) + : Base(coord) + {} + + /// Helper to construct from Cin1, Kh, Kw, Cout, C0 + CATLASS_HOST_DEVICE + Conv2dFilterCoord(Index cin1, Index kh, Index kw, Index cout, Index c0) + : Base(MakeCoord(cin1, kh, kw, cout, c0)) + {} + + CATLASS_HOST_DEVICE + Conv2dFilterCoord(LongIndex cin1, LongIndex kh, LongIndex kw, LongIndex cout, LongIndex c0) + : Base(MakeCoord(Index(cin1), Index(kh), Index(kw), Index(cout), Index(c0))) + {} + + CATLASS_HOST_DEVICE + Index const &cin1() const + { + return this->At(CIN1_INDEX); + } + CATLASS_HOST_DEVICE + Index &cin1() + { + return this->At(CIN1_INDEX); + } + + CATLASS_HOST_DEVICE + Index const &kh() const + { + return this->At(KH_INDEX); + } + CATLASS_HOST_DEVICE + Index &kh() + { + return this->At(KH_INDEX); + } + + CATLASS_HOST_DEVICE + Index const &kw() const + { + return this->At(KW_INDEX); + } + CATLASS_HOST_DEVICE + Index &kw() + { + return this->At(KW_INDEX); + } + + CATLASS_HOST_DEVICE + Index const &cout() const + { + return this->At(COUT_INDEX); + } + CATLASS_HOST_DEVICE + Index &cout() + { + return this->At(COUT_INDEX); + } + + CATLASS_HOST_DEVICE + Index const &c0() const + { + return this->At(C0_INDEX); + } + CATLASS_HOST_DEVICE + Index &c0() + { + return this->At(C0_INDEX); + } + + /// Element-wise addition + CATLASS_HOST_DEVICE + Conv2dFilterCoord operator+(Base const &b) const + { + return Conv2dFilterCoord(Base::operator+(b)); + } + + /// In-place addition + CATLASS_HOST_DEVICE + Conv2dFilterCoord &operator+=(Base const &b) + { + Base::operator+=(b); + return *this; + } +}; + +struct Conv2dHoWoCoCoord : public Coord<3, uint32_t> { // (Ho, Wo, Cout) +public: + /// Integer-valued index + using Index = uint32_t; + + /// Base type is a Coord of rank=3 + using Base = Coord<3, Index>; + + /// LongIndex type + using LongIndex = typename Base::LongIndex; + + /// dimensions + static constexpr uint32_t HO_INDEX = 0; + static constexpr uint32_t WO_INDEX = 1; + static constexpr uint32_t COUT_INDEX = 2; + + /// Default ctor + CATLASS_HOST_DEVICE + Conv2dHoWoCoCoord() {} + + /// Constructs from Coord<3> + CATLASS_HOST_DEVICE + Conv2dHoWoCoCoord(Coord<3, Index> const &coord) + : Base(coord) + {} + + /// Helper to construct from Ho, Wo, Cout + CATLASS_HOST_DEVICE + Conv2dHoWoCoCoord(Index ho, Index wo, Index cout) + : Base(MakeCoord(ho, wo, cout)) + {} + + CATLASS_HOST_DEVICE + Conv2dHoWoCoCoord(LongIndex ho, LongIndex wo, LongIndex cout) + : Base(MakeCoord(Index(ho), Index(wo), Index(cout))) + {} + + CATLASS_HOST_DEVICE + Index const &ho() const + { + return this->At(HO_INDEX); + } + + CATLASS_HOST_DEVICE + Index &ho() + { + return this->At(HO_INDEX); + } + + CATLASS_HOST_DEVICE + Index const &wo() const + { + return this->At(WO_INDEX); + } + + CATLASS_HOST_DEVICE + Index &wo() + { + return this->At(WO_INDEX); + } + + CATLASS_HOST_DEVICE + Index howo() const + { + return this->At(HO_INDEX) * this->At(WO_INDEX); + } + + CATLASS_HOST_DEVICE + Index const &cout() const + { + return this->At(COUT_INDEX); + } + + CATLASS_HOST_DEVICE + Index &cout() + { + return this->At(COUT_INDEX); + } + + /// Element-wise addition + CATLASS_HOST_DEVICE + Conv2dHoWoCoCoord operator+(Base const &b) const + { + return Conv2dHoWoCoCoord(Base::operator+(b)); + } + + /// In-place addition + CATLASS_HOST_DEVICE + Conv2dHoWoCoCoord &operator+=(Base const &b) + { + Base::operator+=(b); + return *this; + } +}; + +struct Conv2dCoord : public Coord<5, uint32_t> { // (Batch, H, W, Cout, Cin1) +public: + /// Integer-valued index + using Index = uint32_t; + + /// Base type is a Coord of rank=5 + using Base = Coord<5, Index>; + + /// LongIndex type + using LongIndex = typename Base::LongIndex; + + /// dimensions + static constexpr uint32_t BATCH_INDEX = 0; + static constexpr uint32_t H_INDEX = 1; + static constexpr uint32_t W_INDEX = 2; + static constexpr uint32_t COUT_INDEX = 3; + static constexpr uint32_t CIN1_INDEX = 4; + + /// Default ctor + CATLASS_HOST_DEVICE + Conv2dCoord() {} + + /// Constructs from Coord<5> + CATLASS_HOST_DEVICE + Conv2dCoord(Coord<5, Index> const &coord) + : Base(coord) + {} + + /// Helper to construct from Batch, H, W, Cout, Cin1 + CATLASS_HOST_DEVICE + Conv2dCoord(Index batch, Index h, Index w, Index cout, Index cin1) + : Base(MakeCoord(batch, h, w, cout, cin1)) + {} + + CATLASS_HOST_DEVICE + Conv2dCoord(LongIndex batch, LongIndex h, LongIndex w, LongIndex cout, LongIndex cin1) + : Base(MakeCoord(Index(batch), Index(h), Index(w), Index(cout), Index(cin1))) + {} + + CATLASS_HOST_DEVICE + Index const &batch() const + { + return this->At(BATCH_INDEX); + } + CATLASS_HOST_DEVICE + Index &batch() + { + return this->At(BATCH_INDEX); + } + + CATLASS_HOST_DEVICE + Index const &h() const + { + return this->At(H_INDEX); + } + CATLASS_HOST_DEVICE + Index &h() + { + return this->At(H_INDEX); + } + + CATLASS_HOST_DEVICE + Index const &w() const + { + return this->At(W_INDEX); + } + CATLASS_HOST_DEVICE + Index &w() + { + return this->At(W_INDEX); + } + + CATLASS_HOST_DEVICE + Index const &cout() const + { + return this->At(COUT_INDEX); + } + CATLASS_HOST_DEVICE + Index &cout() + { + return this->At(COUT_INDEX); + } + + CATLASS_HOST_DEVICE + Index const &cin1() const + { + return this->At(CIN1_INDEX); + } + CATLASS_HOST_DEVICE + Index &cin1() + { + return this->At(CIN1_INDEX); + } + + /// Element-wise addition + CATLASS_HOST_DEVICE + Conv2dCoord operator+(Base const &b) const + { + return Conv2dCoord(Base::operator+(b)); + } + + /// In-place addition + CATLASS_HOST_DEVICE + Conv2dCoord &operator+=(Base const &b) + { + Base::operator+=(b); + return *this; + } + + CATLASS_HOST_DEVICE + auto GetHoWoCoCoord() const + { + return this->GetCoordByAxis(); + } +}; + +class Conv2dFilterParams { +public: + typedef uint8_t ShortIndex; + typedef uint32_t Index; + static constexpr uint32_t C0 = 16; + using Ks = Coord<2, ShortIndex>; + using Pads = Coord<4, ShortIndex>; + using Strides = Coord<2, ShortIndex>; + using Dilations = Coord<2, ShortIndex>; + +private: + Ks ks; + Pads pads; + Strides strides; + Dilations dilations; + +public: + Conv2dFilterParams(ShortIndex kh = 0, ShortIndex kw = 0, + ShortIndex padLeft = 0, ShortIndex padRight = 0, ShortIndex padTop = 0, ShortIndex padBottom = 0, + ShortIndex strideH = 0, ShortIndex strideW = 0, ShortIndex dilationH = 0, ShortIndex dilationW = 0) + : ks(MakeCoord(kh, kw)), + pads(MakeCoord(padLeft, padRight, padTop, padBottom)), + strides(MakeCoord(strideH, strideW)), + dilations(MakeCoord(dilationH, dilationW)) + {} + + CATLASS_HOST_DEVICE + ShortIndex const &kh() const + { + return this->ks[0]; + } + CATLASS_HOST_DEVICE + ShortIndex &kh() + { + return this->ks[0]; + } + + CATLASS_HOST_DEVICE + ShortIndex const &kw() const + { + return this->ks[1]; + } + CATLASS_HOST_DEVICE + ShortIndex &kw() + { + return this->ks[1]; + } + + CATLASS_HOST_DEVICE + ShortIndex const &padLeft() const + { + return this->pads[0]; + } + CATLASS_HOST_DEVICE + ShortIndex const &padLeft() + { + return this->pads[0]; + } + + CATLASS_HOST_DEVICE + ShortIndex const &padRight() const + { + return this->pads[1]; + } + CATLASS_HOST_DEVICE + ShortIndex const &padRight() + { + return this->pads[1]; + } + + CATLASS_HOST_DEVICE + ShortIndex const &padTop() const + { + return this->pads[2]; + } + CATLASS_HOST_DEVICE + ShortIndex const &padTop() + { + return this->pads[2]; + } + + CATLASS_HOST_DEVICE + ShortIndex const &padBottom() const + { + return this->pads[3]; + } + CATLASS_HOST_DEVICE + ShortIndex const &padBottom() + { + return this->pads[3]; + } + + CATLASS_HOST_DEVICE + ShortIndex const &strideH() const + { + return this->strides[0]; + } + CATLASS_HOST_DEVICE + ShortIndex const &strideH() + { + return this->strides[0]; + } + + CATLASS_HOST_DEVICE + ShortIndex const &strideW() const + { + return this->strides[1]; + } + CATLASS_HOST_DEVICE + ShortIndex const &strideW() + { + return this->strides[1]; + } + + CATLASS_HOST_DEVICE + ShortIndex const &dilationH() const + { + return this->dilations[0]; + } + CATLASS_HOST_DEVICE + ShortIndex const &dilationH() + { + return this->dilations[0]; + } + + CATLASS_HOST_DEVICE + ShortIndex const &dilationW() const + { + return this->dilations[1]; + } + CATLASS_HOST_DEVICE + ShortIndex const &dilationW() + { + return this->dilations[1]; + } +}; + +struct Conv2dParams { +public: + typedef uint8_t ShortIndex; + typedef uint32_t Index; + static constexpr uint32_t C0 = 16; + using Pads = Coord<4, ShortIndex>; + using Strides = Coord<2, ShortIndex>; + using Dilations = Coord<2, ShortIndex>; + +private: + // Batch, Hi, Wi, Cin, Cout, Kh, Kw + Conv2dFmapCoord fmapShape; // {Batch, Cin1, Hi, Wi, C0} + Conv2dFilterCoord filterShape; // {Cin1, Kh, Kw, Cout, C0} + Conv2dFmapCoord outputShape; // {Batch, Cout1, Ho, Wo, C0} + Conv2dFilterParams configs; // {Ks, Pads, Strides, Dilations} + Conv2dCoord postIm2colShape; // {Batch, Ho, Wo, Cout, Cin1} +public: + /// Default ctor + CATLASS_HOST_DEVICE + Conv2dParams() {} + + CATLASS_HOST_DEVICE + Conv2dParams(Index batch, Index hi, Index wi, Index cin, Index cout, + ShortIndex kh, ShortIndex kw, + ShortIndex padLeft, ShortIndex padRight, ShortIndex padTop, ShortIndex padBottom, + ShortIndex strideH, ShortIndex strideW, + ShortIndex dilationH, ShortIndex dilationW) + : fmapShape(MakeCoord(batch, CeilDiv(cin, C0), hi, wi, C0)), + filterShape(MakeCoord(CeilDiv(cin, C0), (Index)kh, (Index)kw, cout, C0)), + configs(kh, kw, padLeft, padRight, padTop, padBottom, strideH, strideW, dilationH, dilationW) + { + Index cout1 = CeilDiv(cout, C0); + Index ho = (hi + padTop + padBottom - dilationH * (kh - 1) - 1) / strideH + 1; + Index wo = (wi + padLeft + padRight - dilationW * (kw - 1) - 1) / strideW + 1; + outputShape = MakeCoord(batch, cout1, ho, wo, C0); + postIm2colShape = MakeCoord(batch, ho, wo, cout, filterShape.cin1()); + } + + CATLASS_HOST_DEVICE + static Conv2dParams MakeConv2dParams( + const Index *dataSizes, // [Batch, Hi, Wi, Cin, Cout] + const ShortIndex *filterSizes, // [Kh, Kw] + const ShortIndex *pads, // [padLeft, padRight, padTop, padBottom] + const ShortIndex *strides, // [strideH, strideW] + const ShortIndex *dilations) // [dilationH, dilationW] + { + return Conv2dParams( + dataSizes[0], dataSizes[1], dataSizes[2], dataSizes[3], dataSizes[4], + filterSizes[0], filterSizes[1], + pads[0], pads[1], pads[2], pads[3], + strides[0], strides[1], strides[2], strides[3]); + } + + CATLASS_HOST_DEVICE + Conv2dFilterParams const &getFilterParams() const + { + return this->configs; + } + + CATLASS_HOST_DEVICE + Conv2dFmapCoord const &getOutputShape() const + { + return this->outputShape; + } + + CATLASS_HOST_DEVICE + Conv2dCoord const &getPostIm2colShape() const + { + return this->postIm2colShape; + } + + CATLASS_HOST_DEVICE + Index const &batch() const + { + return this->fmapShape.batch(); + } + CATLASS_HOST_DEVICE + Index &batch() + { + return this->fmapShape.batch(); + } + + CATLASS_HOST_DEVICE + Index const &hi() const + { + return this->fmapShape.h(); + } + CATLASS_HOST_DEVICE + Index &hi() + { + return this->fmapShape.h(); + } + + CATLASS_HOST_DEVICE + Index const &wi() const + { + return this->fmapShape.w(); + } + CATLASS_HOST_DEVICE + Index &wi() + { + return this->fmapShape.w(); + } + + CATLASS_HOST_DEVICE + Index cin() const + { + return this->fmapShape.c1() * C0; + } + + CATLASS_HOST_DEVICE + Index const &cin1() const + { + return this->fmapShape.c1(); + } + CATLASS_HOST_DEVICE + Index &cin1() + { + return this->fmapShape.c1(); + } + + CATLASS_HOST_DEVICE + Index const &cout() const + { + return this->filterShape.cout(); + } + CATLASS_HOST_DEVICE + Index &cout() + { + return this->filterShape.cout(); + } + + CATLASS_HOST_DEVICE + Index const &cout1() const + { + return this->outputShape.c1(); + } + CATLASS_HOST_DEVICE + Index &cout1() + { + return this->outputShape.c1(); + } + + CATLASS_HOST_DEVICE + Index coutRound() const + { + return (this->filterShape.cout() + C0 - 1) / C0 * C0; + } + + CATLASS_HOST_DEVICE + Index const &ho() const + { + return this->outputShape.h(); + } + CATLASS_HOST_DEVICE + Index &ho() + { + return this->outputShape.h(); + } + + CATLASS_HOST_DEVICE + Index const &wo() const + { + return this->outputShape.w(); + } + CATLASS_HOST_DEVICE + Index &wo() + { + return this->outputShape.w(); + } + + CATLASS_HOST_DEVICE + Index howo() const + { + return this->outputShape.h() * this->outputShape.w(); + } + + CATLASS_HOST_DEVICE + Index howoRound() const + { + return (this->howo() + C0 - 1) / C0 * C0; + } + + CATLASS_HOST_DEVICE + ShortIndex const &kh() const + { + return this->configs.kh(); + } + CATLASS_HOST_DEVICE + ShortIndex &kh() + { + return this->configs.kh(); + } + + CATLASS_HOST_DEVICE + ShortIndex const &kw() const + { + return this->configs.kw(); + } + CATLASS_HOST_DEVICE + ShortIndex &kw() + { + return this->configs.kw(); + } + + CATLASS_HOST_DEVICE + ShortIndex const &padTop() const + { + return this->configs.padTop(); + } + CATLASS_HOST_DEVICE + ShortIndex const &padTop() + { + return this->configs.padTop(); + } + + CATLASS_HOST_DEVICE + ShortIndex const &padBottom() const + { + return this->configs.padBottom(); + } + CATLASS_HOST_DEVICE + ShortIndex const &padBottom() + { + return this->configs.padBottom(); + } + + CATLASS_HOST_DEVICE + ShortIndex const &padLeft() const + { + return this->configs.padLeft(); + } + CATLASS_HOST_DEVICE + ShortIndex const &padLeft() + { + return this->configs.padLeft(); + } + + CATLASS_HOST_DEVICE + ShortIndex const &padRight() const + { + return this->configs.padRight(); + } + CATLASS_HOST_DEVICE + ShortIndex const &padRight() + { + return this->configs.padRight(); + } + + CATLASS_HOST_DEVICE + ShortIndex const &strideH() const + { + return this->configs.strideH(); + } + CATLASS_HOST_DEVICE + ShortIndex const &strideH() + { + return this->configs.strideH(); + } + + CATLASS_HOST_DEVICE + ShortIndex const &strideW() const + { + return this->configs.strideW(); + } + CATLASS_HOST_DEVICE + ShortIndex const &strideW() + { + return this->configs.strideW(); + } + + CATLASS_HOST_DEVICE + ShortIndex const &dilationH() const + { + return this->configs.dilationH(); + } + CATLASS_HOST_DEVICE + ShortIndex const &dilationH() + { + return this->configs.dilationH(); + } + + CATLASS_HOST_DEVICE + ShortIndex const &dilationW() const + { + return this->configs.dilationW(); + } + CATLASS_HOST_DEVICE + ShortIndex const &dilationW() + { + return this->configs.dilationW(); + } +}; + +} // namespace Catlass + +#endif // CATLASS_KDA_CONV_COORD_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/kda_coord.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/kda_coord.hpp new file mode 100644 index 0000000000..d68fdb1b28 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/kda_coord.hpp @@ -0,0 +1,336 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_COORD_HPP +#define CATLASS_KDA_COORD_HPP + +#include "catlass/kda_catlass.hpp" + +namespace Catlass { + +/// Statically-sized array specifying Coords within a tensor +template < + int RANK_, ///< Logical rank of coordinate + class Index_ = uint32_t, ///< Index type used for each dimension + class LongIndex_ = int64_t ///< Long index type used for linear offsets + > +struct Coord { +public: + // Number of elements in Coord + static const int RANK = RANK_; + + // Index typen used to store elements + using Index = Index_; + + // Type used to represent linear offsets + using LongIndex = LongIndex_; + + // Default ctor initializes uniformly + CATLASS_HOST_DEVICE constexpr explicit Coord(Index value = Index(0)) + { + for (int i = 0; i < RANK; ++i) { + idx[i] = value; + } + } + + // Constructs from an array of integers + CATLASS_HOST_DEVICE constexpr Coord(Index const (&idx_)[RANK]) + { + for (int i = 0; i < RANK; ++i) { + idx[i] = idx_[i]; + } + } + + // Constructs frrom an array of integers + CATLASS_HOST_DEVICE + int Argmin() const + { + int i = 0; + for (int j = 1; j < RANK; ++j) { + if (idx[j] < idx[i]) { + i = j; + } + } + return i; + } + + // Returns the index of the dimension with greatest value + CATLASS_HOST_DEVICE + int Argmax() const + { + int i = 0; + for (int j = 1; j < RANK; ++j) { + if (idx[j] > idx[i]) { + i = j; + } + } + return i; + } + + // Returns true if Coord is non-zero + CATLASS_HOST_DEVICE + explicit operator bool() const + { + for (int i = 0; i < RANK; ++i) { + if (idx[i]) { + return true; + } + } + return false; + } + + // Return true if Coord is uniformly zero. + CATLASS_HOST_DEVICE + bool operator!() const + { + for (int i = 0; i < RANK; ++i) { + if (idx[i]) { + return false; + } + } + return true; + } + + // Element-wise addition + CATLASS_HOST_DEVICE + Coord operator+(Coord const &b) const + { + Coord c; + for (int i = 0; i < RANK; ++i) { + c.idx[i] = idx[i] + b.idx[i]; + } + return c; + } + + // Add a scalar to each element + CATLASS_HOST_DEVICE + Coord operator+(const Index val) const + { + Coord c; + for (int i = 0; i < RANK; ++i) { + c.idx[i] = idx[i] + val; + } + return c; + } + + // Element-wise subtraction + CATLASS_HOST_DEVICE + Coord operator-(Coord const &b) const + { + Coord c; + for (int i = 0; i < RANK; i++) { + c.idx[i] = idx[i] - b.idx[i]; + } + return c; + } + + // Subtract a scalar from each element + CATLASS_HOST_DEVICE + Coord operator-(Index const val) const + { + Coord c; + for (int i = 0; i < RANK; ++i) { + c.idx[i] = idx[i] - val; + } + return c; + } + + // Element-wise multiply + CATLASS_HOST_DEVICE + Coord operator*(Coord const &b) const + { + Coord c; + for (int i = 0; i < RANK; i++) { + c.idx[i] = idx[i] * b.idx[i]; + } + return c; + } + + // Element-wise division + CATLASS_HOST_DEVICE + Coord operator/(Coord const &b) const + { + Coord c; + for (int i = 0; i < RANK; i++) { + c.idx[i] = idx[i] / b.idx[i]; + } + return c; + } + + // Element-wise mod + CATLASS_HOST_DEVICE + Coord operator%(Coord const &b) const + { + Coord c; + for (int i = 0; i < RANK; i++) { + c.idx[i] = idx[i] % b.idx[i]; + } + return c; + } + + // In-place addition + CATLASS_HOST_DEVICE + Coord &operator+=(Coord const &b) + { + for (int i = 0; i < RANK; ++i) { + idx[i] += b.idx[i]; + } + return *this; + } + + // In-place equal + CATLASS_HOST_DEVICE + bool operator==(Coord const &b) const + { + for (int i = 0; i < RANK; ++i) { + if (idx[i] != b.idx[i]) { + return false; + } + } + return true; + } + + // In-place equal + CATLASS_HOST_DEVICE + bool operator==(Index const val) const + { + for (int i = 0; i < RANK; ++i) { + if (idx[i] != val) { + return false; + } + } + return true; + } + + // Member acces operator + CATLASS_HOST_DEVICE + Index &operator[](int dim) + { + return idx[dim]; + } + + // Member access operator + CATLASS_HOST_DEVICE + Index const &operator[](int dim) const + { + return idx[dim]; + } + + // Gets the index of a given Coord element + template + CATLASS_HOST_DEVICE Index &At() + { + return idx[DIM]; + } + + // Access via index; may limit unrolling potential + CATLASS_HOST_DEVICE + Index &At(int dim) + { + return idx[dim]; + } + + // Gets the index of a given Coord element + template + CATLASS_HOST_DEVICE + Index const & + At() const + { + return idx[DIM]; + } + + // Access via index; may limit unrolling potential + CATLASS_HOST_DEVICE + Index const &At(int dim) const + { + return idx[dim]; + } + + template + CATLASS_HOST_DEVICE auto GetCoordByAxis() const + { + Index idx_[sizeof...(Is)]{idx[Is]...}; + return Coord{idx_}; + } + + CATLASS_HOST_DEVICE + static Coord Min(Coord const &a, Coord const &b) + { + Coord res; + for (int i = 0; i < RANK; ++i) { + res[i] = a[i] < b[i] ? a[i] : b[i]; + } + return res; + } + +private: + // Indices + Index idx[RANK]; +}; + +// Helper to make a 1-element coordinate +template +CATLASS_HOST_DEVICE constexpr Coord<1, T> MakeCoord(T dim0) +{ + T values[1] = {dim0}; + return Coord<1, T>(values); +} + +/// Helper to make a 2-element coordinate +template +CATLASS_HOST_DEVICE constexpr Coord<2, T> MakeCoord(T dim0, T dim1) +{ + T values[2] = {dim0, dim1}; + return Coord<2, T>(values); +} + +/// Helper to make a 3-element coordinate +template +CATLASS_HOST_DEVICE constexpr Coord<3, T> MakeCoord(T dim0, T dim1, T dim2) +{ + T values[3] = {dim0, dim1, dim2}; + return Coord<3, T>(values); +} + +/// Helper to make a 4-element coordinate +template +CATLASS_HOST_DEVICE constexpr Coord<4, T> MakeCoord(T dim0, T dim1, T dim2, T dim3) +{ + T values[4] = {dim0, dim1, dim2, dim3}; + return Coord<4, T>(values); +} + +/// Helper to make a 5-element coordinate +template +CATLASS_HOST_DEVICE constexpr Coord<5, T> MakeCoord(T dim0, T dim1, T dim2, T dim3, T dim4) +{ + T values[5] = {dim0, dim1, dim2, dim3, dim4}; + return Coord<5, T>(values); +} + +/// Helper to make a 6-element coordinate +template +CATLASS_HOST_DEVICE constexpr Coord<6, T> MakeCoord(T dim0, T dim1, T dim2, T dim3, T dim4, T dim5) +{ + T values[6] = {dim0, dim1, dim2, dim3, dim4, dim5}; + return Coord<6, T>(values); +} + +/// Helper to make a 7-element coordinate +template +CATLASS_HOST_DEVICE constexpr Coord<7, T> MakeCoord(T dim0, T dim1, T dim2, T dim3, T dim4, T dim5, T dim6) +{ + T values[7] = {dim0, dim1, dim2, dim3, dim4, dim5, dim6}; + return Coord<7, T>(values); +} + +} // namespace Catlass + +#endif // CATLASS_KDA_COORD_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/kda_debug.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/kda_debug.hpp new file mode 100644 index 0000000000..6e0ecc9ecd --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/kda_debug.hpp @@ -0,0 +1,73 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_DEBUG_HPP +#define CATLASS_KDA_DEBUG_HPP + +#include +#include +#include + +#include + +#define SINGLE_CORE_DUMPSIZE (1024 * 1024) +#ifdef CATLASS_ARCH +#if CATLASS_ARCH == 2201 +#define DUMP_CORE_NUM 75 +#endif +#if CATLASS_ARCH == 3510 +#define DUMP_CORE_NUM 108 +#endif +#endif +#ifndef DUMP_CORE_NUM +#define DUMP_CORE_NUM 0 +#endif +// DUMP_CORE_NUM is from AscendC host stub +#define ALL_DUMPSIZE (DUMP_CORE_NUM * SINGLE_CORE_DUMPSIZE) + +using LogFuncType = std::function; +/** + * @brief Check acl api status code. + * @param status The return code of acl api. + * @param logFunc Log function, which receives a C-Style string. + * @return + */ +inline void aclCheck( + aclError status, + LogFuncType logFunc = [](const char *logStrPtr) { std::cerr << logStrPtr; }) +{ + if (status != ACL_SUCCESS) { + std::stringstream ss; + ss << "AclError: " << status; + logFunc(ss.str().c_str()); + } +} +/** + * @brief Check rt api status code. + * @param status The return code of rt api. + * @param logFunc Log function, which receives a C-Style string. + * @return + */ +inline void rtCheck( + int status, + LogFuncType logFunc = [](const char *logStrPtr) { std::cerr << logStrPtr; }) +{ + if (status != 0) { + std::stringstream ss; + ss << "RtError: " << status; + logFunc(ss.str().c_str()); + } +} + +namespace Adx { +void AdumpPrintWorkSpace(const void *dumpBufferAddr, const size_t dumpBufferSize, aclrtStream stream, const char *opType); +} // namespace Adx + +#endif // CATLASS_KDA_DEBUG_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/kda_gemm_coord.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/kda_gemm_coord.hpp new file mode 100644 index 0000000000..0c23e4ac8e --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/kda_gemm_coord.hpp @@ -0,0 +1,161 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_GEMM_COORD_HPP +#define CATLASS_KDA_GEMM_COORD_HPP + +#include "catlass/kda_coord.hpp" + +namespace Catlass { + +/// Shape of a matrix multiply-add operation +template < + /// Rows of matrix product + uint32_t M_ = 1, + /// Columns of matrix product + uint32_t N_ = 1, + /// Inner dimension of matrix product + uint32_t K_ = 1> +struct GemmShape { + static constexpr uint32_t M = M_; + static constexpr uint32_t N = N_; + static constexpr uint32_t K = K_; + + static constexpr int64_t MN = M * N; + static constexpr int64_t MK = M * K; + static constexpr int64_t KN = N * K; + static constexpr int64_t MNK = M * N * K; + + static constexpr int64_t COUNT = MNK; + + /// Returns a Coord object + CATLASS_HOST_DEVICE + static Coord<3> ToCoord() + { + return MakeCoord(M, N, K); + } + + CATLASS_HOST_DEVICE + static Coord<2> ToCoordMN() + { + return MakeCoord(M, N); + } + + CATLASS_HOST_DEVICE + static Coord<2> ToCoordMK() + { + return MakeCoord(M, K); + } + + CATLASS_HOST_DEVICE + static Coord<2> ToCoordKN() + { + return MakeCoord(K, N); + } +}; + +/// GemmCoord is a structure derived from Coord<3> that specifies a location within the +/// coordinate space of a Gemm problem. +struct GemmCoord : public Coord<3, uint32_t> { + /// Integer-valued index + using Index = uint32_t; + + /// Base type is a Coord of rank=3 + using Base = Coord<3, Index>; + + /// Gemm M dimension - rows of the output C matrix + static constexpr int M_INDEX = 0; + + /// Gemm N dimension - columns of the output C matrix + static constexpr int N_INDEX = 1; + + /// Gemm K dimension - inner dimension of the Gemm problem + static constexpr int K_INDEX = 2; + + /// Default ctor + CATLASS_HOST_DEVICE + GemmCoord() {} + + /// Constructs from Coord<3> and a batch + CATLASS_HOST_DEVICE + GemmCoord(Coord<3, Index> const &coord) + : Base(coord) + {} + + /// Helper to construct from a K, N, M, batch variables + CATLASS_HOST_DEVICE + GemmCoord(Index m, Index n, Index k) + : Base(MakeCoord(m, n, k)) + {} + + /// Returns the Gemm M coordinate + CATLASS_HOST_DEVICE + Index const &m() const + { + return this->At(M_INDEX); + } + + /// Returns reference to the Gemm M coordinate + CATLASS_HOST_DEVICE + Index &m() + { + return this->At(M_INDEX); + } + + /// Returns the Gemm N coordinate + CATLASS_HOST_DEVICE + Index const &n() const + { + return this->At(N_INDEX); + } + + /// Returns reference to the Gemm N coordinate + CATLASS_HOST_DEVICE + Index &n() + { + return this->At(N_INDEX); + } + + /// Returns the Gemm K coordinate + CATLASS_HOST_DEVICE + Index const &k() const + { + return this->At(K_INDEX); + } + + /// Returns reference to the Gemm K coordinate + CATLASS_HOST_DEVICE + Index &k() + { + return this->At(K_INDEX); + } + + CATLASS_HOST_DEVICE + auto GetCoordMN() const + { + return this->GetCoordByAxis(); + } + + CATLASS_HOST_DEVICE + auto GetCoordMK() const + { + return this->GetCoordByAxis(); + } + + CATLASS_HOST_DEVICE + auto GetCoordKN() const + { + return this->GetCoordByAxis(); + } +}; + +} // namespace Catlass + +#endif // CATLASS_KDA_GEMM_COORD_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/kda_gemv_coord.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/kda_gemv_coord.hpp new file mode 100644 index 0000000000..f209676924 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/kda_gemv_coord.hpp @@ -0,0 +1,108 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_GEMV_COORD_HPP +#define CATLASS_KDA_GEMV_COORD_HPP + +#include "catlass/kda_coord.hpp" + +namespace Catlass { + +/// Shape of a matrix multiply-add operation +template < + /// Rows of matrix product + uint32_t M_ = 1, + /// Columns of the matrix (number of elements in the input vector) + uint32_t N_ = 1> +struct GemvShape { + static constexpr uint32_t M = M_; + static constexpr uint32_t N = N_; + + static constexpr int64_t MN = M * N; + + static constexpr int64_t COUNT = MN; + + /// Returns a Coord object + CATLASS_HOST_DEVICE + static Coord<2> ToCoord() + { + return MakeCoord(M, N); + } +}; + +/// GemvCoord is a structure derived from Coord<2> that specifies a location within the +/// coordinate space of a GEMV problem. +struct GemvCoord : public Coord<2, uint32_t> { + /// Integer-valued index + using Index = uint32_t; + + /// Base type is a Coord of rank=2 + using Base = Coord<2, Index>; + + /// GEMV M dimension - rows of the output vector (y) + static constexpr int M_INDEX = 0; + + /// GEMV N dimension - columns of the matrix (length of the input vector x) + static constexpr int N_INDEX = 1; + + /// Default ctor + CATLASS_HOST_DEVICE + GemvCoord() {} + + /// Constructs from Coord<2> and a batch + CATLASS_HOST_DEVICE + GemvCoord(Coord<2, Index> const &coord) + : Base(coord) + {} + + /// Helper to construct from M, N coordinates + CATLASS_HOST_DEVICE + GemvCoord(Index m, Index n) + : Base(MakeCoord(m, n)) + {} + + /// Returns the GEMV M coordinate (row of the result y) + CATLASS_HOST_DEVICE + Index const &m() const + { + return this->At(M_INDEX); + } + + /// Returns reference to the GEMV M coordinate + CATLASS_HOST_DEVICE + Index &m() + { + return this->At(M_INDEX); + } + + /// Returns the GEMV N coordinate (column of the matrix A or the input vector x) + CATLASS_HOST_DEVICE + Index const &n() const + { + return this->At(N_INDEX); + } + + /// Returns reference to the GEMV N coordinate + CATLASS_HOST_DEVICE + Index &n() + { + return this->At(N_INDEX); + } + + CATLASS_HOST_DEVICE + auto GetCoordMN() const + { + return this->GetCoordByAxis(); + } +}; + +} // namespace Catlass + +#endif // CATLASS_KDA_GEMV_COORD_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/kda_matrix_coord.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/kda_matrix_coord.hpp new file mode 100644 index 0000000000..fccd956347 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/kda_matrix_coord.hpp @@ -0,0 +1,108 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_MATRIX_COORD_HPP +#define CATLASS_KDA_MATRIX_COORD_HPP + +#include "catlass/kda_coord.hpp" + +namespace Catlass { + +template < + uint32_t ROW_ = 1, + uint32_t COLUMN_ = 1> +struct MatrixShape { + static constexpr uint32_t ROW = ROW_; + static constexpr uint32_t COLUMN = COLUMN_; + + static constexpr int64_t COUNT = ROW * COLUMN; + + CATLASS_HOST_DEVICE + static Coord<2> ToCoord() + { + return MakeCoord(ROW, COLUMN); + } +}; + +/// MatrixCoord wraps Coord<2, uint32_t> to provide a helper for accessing named dimensions. Classes +/// expecting a coordinate in the rank=2 index space of a matrix should use MatrixCoord. +struct MatrixCoord : public Coord<2, uint32_t> { + /// Integer-valued index + using Index = uint32_t; + + /// Base type is a Coord of rank=2 + using Base = Coord<2, Index>; + + /// LongIndex type + using LongIndex = typename Base::LongIndex; + + /// Rows dimension + static constexpr uint32_t ROW_INDEX = 0; + + /// Columns dimension + static constexpr uint32_t COLUMN_INDEX = 1; + + /// Default ctor + CATLASS_HOST_DEVICE + MatrixCoord() {} + + /// Constructs from Coord<2> + CATLASS_HOST_DEVICE + MatrixCoord(Coord<2, Index> const &coord) + : Base(coord) + {} + + /// Helper to construct from a row and column + CATLASS_HOST_DEVICE + MatrixCoord(Index row, Index column) + : Base(MakeCoord(row, column)) + {} + + /// Helper to construct from a row and column, which are LongIndex based + CATLASS_HOST_DEVICE + MatrixCoord(LongIndex row, LongIndex column) + : Base(MakeCoord(Index(row), Index(column))) + {} + + /// Returns the row of the coordinate + CATLASS_HOST_DEVICE + Index const &row() const { return this->At(ROW_INDEX); } + + /// Returns the row of the coordinate + CATLASS_HOST_DEVICE + Index &row() { return this->At(ROW_INDEX); } + + /// Returns the column of the coordinate + CATLASS_HOST_DEVICE + Index const &column() const { return this->At(COLUMN_INDEX); } + + /// Returns the column of the coordinate + CATLASS_HOST_DEVICE + Index &column() { return this->At(COLUMN_INDEX); } + + /// Element-wise addition + CATLASS_HOST_DEVICE + MatrixCoord operator+(Base const &b) const + { + return MatrixCoord(Base::operator+(b)); + } + + /// In-place addition + CATLASS_HOST_DEVICE + MatrixCoord &operator+=(Base const &b) + { + Base::operator+=(b); + return *this; + } +}; + +} // namespace Catlass + +#endif diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/kda_numeric_size.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/kda_numeric_size.hpp new file mode 100644 index 0000000000..11b5a05cf2 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/kda_numeric_size.hpp @@ -0,0 +1,42 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ +#ifndef CATLASS_KDA_NUMERIC_SIZE_HPP +#define CATLASS_KDA_NUMERIC_SIZE_HPP + +#include "catlass/kda_catlass.hpp" + +namespace Catlass { + +#if defined(__CCE__) +using AscendC::SizeOfBits; +#else +template +struct SizeOfBits { + static constexpr size_t value = sizeof(T) * 8; +}; +#endif + +/// Returns the number of bytes required to hold a specified number of bits +template +CATLASS_HOST_DEVICE constexpr ReturnType BitsToBytes(T bits) +{ + return (static_cast(bits) + static_cast(7)) / static_cast(8); +} + +/// Returns the number of bits required to hold a specified number of bytes +template +CATLASS_HOST_DEVICE constexpr ReturnType BytesToBits(T bytes) +{ + return static_cast(bytes) * static_cast(8); +} + +} // namespace Catlass + +#endif // CATLASS_KDA_NUMERIC_SIZE_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/layout/kda_layout.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/layout/kda_layout.hpp new file mode 100644 index 0000000000..1f0c85e2c1 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/layout/kda_layout.hpp @@ -0,0 +1,19 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_LAYOUT_HPP +#define CATLASS_KDA_LAYOUT_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/layout/kda_matrix.hpp" +#include "catlass/layout/kda_vector.hpp" +#include "catlass/layout/kda_tensor.hpp" + +#endif // CATLASS_KDA_LAYOUT_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/layout/kda_matrix.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/layout/kda_matrix.hpp new file mode 100644 index 0000000000..6e2efbeeeb --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/layout/kda_matrix.hpp @@ -0,0 +1,1818 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_MATRIX_HPP +#define CATLASS_KDA_MATRIX_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/kda_coord.hpp" +#include "catlass/kda_numeric_size.hpp" +#include "catlass/detail/kda_alignment.hpp" +#include "catlass/kda_matrix_coord.hpp" +#include "catlass/kda_conv_coord.hpp" + +namespace Catlass::layout { + +/// Mapping function for row-major matrices +struct RowMajor { +public: + /// Logical rank of tensor + static constexpr int RANK = 2; + + /// Index type used for coordinates + using Index = uint32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate + using Shape = Coord; + + /// Stride vector + using Stride = Coord; + +public: + /// Constructor + CATLASS_HOST_DEVICE + RowMajor(Index rows = 0, Index cols = 0) + : shape_(MakeCoord(rows, cols)), + stride_(MakeCoord(LongIndex(cols), LongIndex(1))) + {} + + /// Constructor + CATLASS_HOST_DEVICE + RowMajor(Index rows, Index cols, LongIndex ldm) + : shape_(MakeCoord(rows, cols)), + stride_(MakeCoord(ldm, LongIndex(1))) + {} + + /// Ctor + CATLASS_HOST_DEVICE + RowMajor(Shape shape, Stride stride) + : shape_(shape), + stride_(stride) + {} + + template + CATLASS_HOST_DEVICE static RowMajor MakeLayout(Index rows, Index cols) + { + return RowMajor(rows, cols); + } + + template + CATLASS_HOST_DEVICE static RowMajor MakeLayoutInUb(MatrixCoord const &shape) + { + constexpr uint32_t ELE_NUM_PER_BLK = BytesToBits(BYTE_PER_BLK) / SizeOfBits::value; + return RowMajor(shape.row(), shape.column(), RoundUp(shape.column())); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (row, column) + CATLASS_HOST_DEVICE + LongIndex GetOffset(MatrixCoord const &coord) const + { + return LongIndex(coord.row()) * stride_[0] + LongIndex(coord.column()); + } + + /// Returns the layout of a tile. + CATLASS_HOST_DEVICE + RowMajor GetTileLayout(MatrixCoord const &tileShape) const + { + return RowMajor(tileShape, stride()); + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape shape() const + { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape &shape() + { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index shape(int idx) const + { + return shape_[idx]; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index &shape(int idx) + { + return shape_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride stride() const + { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride &stride() + { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index stride(int idx) const + { + return stride_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index &stride(int idx) + { + return stride_[idx]; + } + + /// Returns the length of the layout + CATLASS_HOST_DEVICE + LongIndex Capacity() const + { + return static_cast(shape_[0]) * stride_[0]; + } + +protected: + // + // Data members + // + + /// Shape data member + Shape shape_; + + /// Stride data member + Stride stride_; +}; + +/// Mapping function for col-major matrices +struct ColumnMajor { +public: + /// Logical rank of tensor + static constexpr int RANK = 2; + + /// Index type used for coordinates + using Index = uint32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate + using Shape = Coord; + + /// Stride vector + using Stride = Coord; + +public: + // Methods + + /// Constructor + CATLASS_HOST_DEVICE + ColumnMajor(Index rows = 0, Index cols = 0) + : shape_(MakeCoord(rows, cols)), + stride_(MakeCoord(LongIndex(1), LongIndex(rows))) + {} + + /// Constructor + CATLASS_HOST_DEVICE + ColumnMajor(Index rows, Index cols, LongIndex ldm) + : shape_(MakeCoord(rows, cols)), + stride_(MakeCoord(LongIndex(1), ldm)) + {} + + /// Ctor + CATLASS_HOST_DEVICE + ColumnMajor(Shape shape, Stride stride) + : shape_(shape), + stride_(stride) + {} + + template + CATLASS_HOST_DEVICE static ColumnMajor MakeLayout(Index rows, Index cols) + { + return ColumnMajor(rows, cols); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (row, column) + CATLASS_HOST_DEVICE + LongIndex GetOffset(MatrixCoord const &coord) const + { + return LongIndex(coord.row()) + LongIndex(coord.column()) * stride_[1]; + } + + /// Returns the layout of a tile. + CATLASS_HOST_DEVICE + ColumnMajor GetTileLayout(MatrixCoord const &tileShape) const + { + return ColumnMajor(tileShape, stride()); + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape shape() const + { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape &shape() + { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index shape(int idx) const + { + return shape_[idx]; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index &shape(int idx) + { + return shape_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride stride() const + { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride &stride() + { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index stride(int idx) const + { + return stride_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index &stride(int idx) + { + return stride_[idx]; + } + + /// Returns the length of the layout + CATLASS_HOST_DEVICE + LongIndex Capacity() const + { + return static_cast(shape_[1]) * stride_[1]; + } + +protected: + // + // Data members + // + + /// Shape data member + Shape shape_; + + /// Stride data member + Stride stride_; +}; + +/// Mapping function for nZ matrices which is col-major inside fractal and row-major between fractal +struct nZ { +public: + /// Logical rank of tensor + static constexpr int RANK = 4; + + /// Index type used for coordinates + using Index = uint32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical rank of orgshape + static constexpr int ORG_SHAPE_RANK = 2; + + /// Logical coordinate + using OrgShape = Coord; + + /// Logical coordinate + using Shape = Coord; + + /// Stride vector + using Stride = Coord; + +public: + // Methods + + /// Constructor + CATLASS_HOST_DEVICE constexpr nZ(Index orgRows = 0, /// Number of rows of origin matrices + Index orgCols = 0, /// Number of cols of origin matrices + Index rowsInFractal = 0, /// Number of rows inside the fractal + Index rowsByFractal = 0, /// number of rows by the fractal + Index colsInFractal = 0, /// number of cols inside the fractal + Index colsByFractal = 0, /// number of cols by the fractal + LongIndex strideRowsInFractal = 0, /// number of elements between adjacent rows inside the fractal + LongIndex strideRowsByFractal = 0, /// number of elements between adjacent fractal rows + LongIndex strideColsInFractal = 0, /// number of elements between adjacent cols inside the fractal + LongIndex strideColsByFractal = 0) /// number of elements between adjacent fractal cols + : orgShape_(MakeCoord(orgRows, orgCols)), + shape_(MakeCoord(rowsInFractal, rowsByFractal, colsInFractal, colsByFractal)), + stride_(MakeCoord(strideRowsInFractal, strideRowsByFractal, strideColsInFractal, strideColsByFractal)) + {} + + /// Ctor + CATLASS_HOST_DEVICE constexpr nZ(OrgShape orgShape, Shape shape, Stride stride) + : orgShape_(orgShape), + shape_(shape), + stride_(stride) + {} + + /// Make the layout of a coordinate (row, column) + template + CATLASS_HOST_DEVICE constexpr static nZ MakeLayout(Index orgRows, Index orgCols) + { + constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + Index rowsRound = RoundUp(orgRows); + Index colsRound = RoundUp(orgCols); + return nZ(orgRows, + orgCols, + ELE_NUM_PER_C0, + rowsRound / ELE_NUM_PER_C0, + C0_NUM_PER_FRACTAL, + colsRound / C0_NUM_PER_FRACTAL, + 1, + colsRound * ELE_NUM_PER_C0, + ELE_NUM_PER_C0, + ELE_NUM_PER_FRACTAL); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (row, column) + CATLASS_HOST_DEVICE + LongIndex GetOffset(MatrixCoord const &coord) const + { + return LongIndex(coord.row()) / shape_[0] * stride_[1] + LongIndex(coord.column()) / shape_[2] * stride_[3] + + (LongIndex(coord.row()) % shape_[0]) * stride_[0] + (LongIndex(coord.column()) % shape_[2]) * stride_[2]; + } + + /// Returns the layout of a tile. + CATLASS_HOST_DEVICE + nZ GetTileLayout(MatrixCoord const &tileOriShape) const + { + auto tileShape = MakeCoord( + shape(0), CeilDiv(tileOriShape.row(), shape(0)), + shape(2), CeilDiv(tileOriShape.column(), shape(2))); + return nZ(tileOriShape, tileShape, stride()); + } + + /// Returns the origin shape of the layout + CATLASS_HOST_DEVICE + typename OrgShape::Index orgShape(int idx) const + { + return orgShape_[idx]; + } + + /// Returns the origin shape of the layout + CATLASS_HOST_DEVICE + typename OrgShape::Index &orgShape(int idx) + { + return orgShape_[idx]; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape shape() const + { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape &shape() + { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index shape(int idx) const + { + return shape_[idx]; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index &shape(int idx) + { + return shape_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride stride() const + { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride &stride() + { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index stride(int idx) const + { + return stride_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index &stride(int idx) + { + return stride_[idx]; + } + + /// Returns the length of the layout + CATLASS_HOST_DEVICE + LongIndex Capacity() const + { + return static_cast(stride_[1]) * shape_[1]; + } + +private: + /// Origin Shape data member + OrgShape orgShape_; + + /// Shape data member + Shape shape_; + + /// Stride data member + Stride stride_; +}; + +/// Mapping function for zN matrices which is row-major inside fractal and col-major between fractal +struct zN { +public: + /// Logical rank of tensor + static constexpr int RANK = 4; + + /// Index type used for coordinates + using Index = uint32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical rank of orgshape + static constexpr int ORG_SHAPE_RANK = 2; + + /// Logical coordinate + using OrgShape = Coord; + + /// Logical coordinate + using Shape = Coord; + + /// Stride vector + using Stride = Coord; + +public: + // Methods + + /// Constructor + CATLASS_HOST_DEVICE constexpr zN(Index orgRows = 0, /// Number of rows of origin matrices + Index orgCols = 0, /// Number of cols of origin matrices + Index rowsInFractal = 0, /// Number of rows inside the fractal + Index rowsByFractal = 0, /// number of rows by the fractal + Index colsInFractal = 0, /// number of cols inside the fractal + Index colsByFractal = 0, /// number of cols by the fractal + LongIndex strideRowsInFractal = 0, /// number of elements between adjacent rows inside the fractal + LongIndex strideRowsByFractal = 0, /// number of elements between adjacent fractal rows + LongIndex strideColsInFractal = 0, /// number of elements between adjacent cols inside the fractal + LongIndex strideColsByFractal = 0) /// number of elements between adjacent fractal cols + : orgShape_(MakeCoord(orgRows, orgCols)), + shape_(MakeCoord(rowsInFractal, rowsByFractal, colsInFractal, colsByFractal)), + stride_(MakeCoord(strideRowsInFractal, strideRowsByFractal, strideColsInFractal, strideColsByFractal)) + {} + + /// Ctor + CATLASS_HOST_DEVICE constexpr zN(OrgShape orgShape, Shape shape, Stride stride) + : orgShape_(orgShape), + shape_(shape), + stride_(stride) + {} + + /// Make the layout of a coordinate (row, column) + template + CATLASS_HOST_DEVICE constexpr static zN MakeLayout(Index orgRows, Index orgCols) + { + constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + Index rowsRound = RoundUp(orgRows); + Index colsRound = RoundUp(orgCols); + return zN(orgRows, + orgCols, + C0_NUM_PER_FRACTAL, + rowsRound / C0_NUM_PER_FRACTAL, + ELE_NUM_PER_C0, + colsRound / ELE_NUM_PER_C0, + ELE_NUM_PER_C0, + ELE_NUM_PER_FRACTAL, + 1, + rowsRound * ELE_NUM_PER_C0); + } + + CATLASS_HOST_DEVICE + static zN MakeLayoutInL0C(MatrixCoord const &shape) + { + return zN(shape.row(), + shape.column(), + C0_NUM_PER_FRACTAL, + CeilDiv(shape.row()), + C0_NUM_PER_FRACTAL, + CeilDiv(shape.column()), + C0_NUM_PER_FRACTAL, + C0_NUM_PER_FRACTAL * C0_NUM_PER_FRACTAL, + 1, + RoundUp(shape.row()) * C0_NUM_PER_FRACTAL); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (row, column) + CATLASS_HOST_DEVICE + LongIndex GetOffset(MatrixCoord const &coord) const + { + return LongIndex(coord.row()) / shape_[0] * stride_[1] + LongIndex(coord.column()) / shape_[2] * stride_[3] + + (LongIndex(coord.row()) % shape_[0]) * stride_[0] + (LongIndex(coord.column()) % shape_[2]) * stride_[2]; + } + + /// Returns the layout of a tile. + CATLASS_HOST_DEVICE + zN GetTileLayout(MatrixCoord const &tileOriShape) const + { + auto tileShape = MakeCoord( + shape(0), CeilDiv(tileOriShape.row(), shape(0)), + shape(2), CeilDiv(tileOriShape.column(), shape(2))); + return zN(tileOriShape, tileShape, stride()); + } + + /// Returns the origin shape of the layout + CATLASS_HOST_DEVICE + typename OrgShape::Index orgShape(int idx) const + { + return orgShape_[idx]; + } + + /// Returns the origin shape of the layout + CATLASS_HOST_DEVICE + typename OrgShape::Index &orgShape(int idx) + { + return orgShape_[idx]; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape shape() const + { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape &shape() + { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index shape(int idx) const + { + return shape_[idx]; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index &shape(int idx) + { + return shape_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride stride() const + { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride &stride() + { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index stride(int idx) const + { + return stride_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index &stride(int idx) + { + return stride_[idx]; + } + + /// Returns the length of the layout + CATLASS_HOST_DEVICE + LongIndex Capacity() const + { + return static_cast(stride_[3]) * shape_[3]; + } + +private: + /// Origin Shape data member + OrgShape orgShape_; + + /// Shape data member + Shape shape_; + + /// Stride data member + Stride stride_; +}; + +/// Mapping function for zN matrices which is row-major inside fractal and row-major between fractal +struct zZ { +public: + /// Logical rank of tensor + static constexpr int RANK = 4; + + /// Index type used for coordinates + using Index = uint32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical rank of orgshape + static constexpr int ORG_SHAPE_RANK = 2; + + /// Logical coordinate + using OrgShape = Coord; + + /// Logical coordinate + using Shape = Coord; + + /// Stride vector + using Stride = Coord; + +public: + // Methods + + /// Constructor + CATLASS_HOST_DEVICE constexpr zZ(Index orgRows = 0, /// Number of rows of origin matrices + Index orgCols = 0, /// Number of cols of origin matrices + Index rowsInFractal = 0, /// Number of rows inside the fractal + Index rowsByFractal = 0, /// number of rows by the fractal + Index colsInFractal = 0, /// number of cols inside the fractal + Index colsByFractal = 0, /// number of cols by the fractal + LongIndex strideRowsInFractal = 0, /// number of elements between adjacent rows inside the fractal + LongIndex strideRowsByFractal = 0, /// number of elements between adjacent fractal rows + LongIndex strideColsInFractal = 0, /// number of elements between adjacent cols inside the fractal + LongIndex strideColsByFractal = 0) /// number of elements between adjacent fractal cols + : orgShape_(MakeCoord(orgRows, orgCols)), + shape_(MakeCoord(rowsInFractal, rowsByFractal, colsInFractal, colsByFractal)), + stride_(MakeCoord(strideRowsInFractal, strideRowsByFractal, strideColsInFractal, strideColsByFractal)) + {} + + /// Ctor + CATLASS_HOST_DEVICE constexpr zZ(OrgShape orgShape, Shape shape, Stride stride) + : orgShape_(orgShape), + shape_(shape), + stride_(stride) + {} + + /// Make the layout of a coordinate (row, column) + template + CATLASS_HOST_DEVICE constexpr static zZ MakeLayout(Index orgRows, Index orgCols) + { + constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + Index rowsRound = RoundUp(orgRows); + Index colsRound = RoundUp(orgCols); + return zZ(orgRows, + orgCols, + C0_NUM_PER_FRACTAL, + rowsRound / C0_NUM_PER_FRACTAL, + ELE_NUM_PER_C0, + colsRound / ELE_NUM_PER_C0, + ELE_NUM_PER_C0, + colsRound * C0_NUM_PER_FRACTAL, + 1, + ELE_NUM_PER_FRACTAL); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (row, column) + CATLASS_HOST_DEVICE + LongIndex GetOffset(MatrixCoord const &coord) const + { + return LongIndex(coord.row()) / shape_[0] * stride_[1] + LongIndex(coord.column()) / shape_[2] * stride_[3]; + } + + /// Returns the origin shape of the layout + CATLASS_HOST_DEVICE + typename OrgShape::Index orgShape(int idx) const + { + return orgShape_[idx]; + } + + /// Returns the origin shape of the layout + CATLASS_HOST_DEVICE + typename OrgShape::Index &orgShape(int idx) + { + return orgShape_[idx]; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape shape() const + { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape &shape() + { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index shape(int idx) const + { + return shape_[idx]; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index &shape(int idx) + { + return shape_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride stride() const + { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride &stride() + { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index stride(int idx) const + { + return stride_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index &stride(int idx) + { + return stride_[idx]; + } + +private: + /// Origin Shape data member + OrgShape orgShape_; + + /// Shape data member + Shape shape_; + + /// Stride data member + Stride stride_; +}; + +/// Mapping function for L0C matrices +/// A special data layout for L0C memory, used for accumulator storage. +/// This layout similar to zN layout, but uses C0_NUM_PER_FRACTAL for both row and column inner dimensions. +struct L0C { +public: + /// Logical rank of tensor + static constexpr int RANK = 4; + + /// Index type used for coordinates + using Index = uint32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical rank of orgshape + static constexpr int ORG_SHAPE_RANK = 2; + + /// Logical coordinate + using OrgShape = Coord; + + /// Logical coordinate + using Shape = Coord; + + /// Stride vector + using Stride = Coord; + +public: + // Methods + + /// Constructor + CATLASS_HOST_DEVICE constexpr L0C(Index orgRows = 0, /// Number of rows of origin matrices + Index orgCols = 0, /// Number of cols of origin matrices + Index rowsInFractal = 0, /// Number of rows inside the fractal + Index rowsByFractal = 0, /// number of rows by the fractal + Index colsInFractal = 0, /// number of cols inside the fractal + Index colsByFractal = 0, /// number of cols by the fractal + LongIndex strideRowsInFractal = 0, /// number of elements between adjacent rows inside the fractal + LongIndex strideRowsByFractal = 0, /// number of elements between adjacent fractal rows + LongIndex strideColsInFractal = 0, /// number of elements between adjacent cols inside the fractal + LongIndex strideColsByFractal = 0) /// number of elements between adjacent fractal cols + : orgShape_(MakeCoord(orgRows, orgCols)), + shape_(MakeCoord(rowsInFractal, rowsByFractal, colsInFractal, colsByFractal)), + stride_(MakeCoord(strideRowsInFractal, strideRowsByFractal, strideColsInFractal, strideColsByFractal)) + {} + + /// Ctor + CATLASS_HOST_DEVICE constexpr L0C(OrgShape orgShape, Shape shape, Stride stride) + : orgShape_(orgShape), + shape_(shape), + stride_(stride) + {} + + /// Make the layout of a coordinate (row, column) + template + CATLASS_HOST_DEVICE constexpr static L0C MakeLayout(Index orgRows, Index orgCols) + { + constexpr uint32_t ELE_NUM_PER_FRACTAL = 256; + Index rowsRound = RoundUp(orgRows); + Index colsRound = RoundUp(orgCols); + return L0C(orgRows, + orgCols, + C0_NUM_PER_FRACTAL, + rowsRound / C0_NUM_PER_FRACTAL, + C0_NUM_PER_FRACTAL, + colsRound / C0_NUM_PER_FRACTAL, + C0_NUM_PER_FRACTAL, + ELE_NUM_PER_FRACTAL, + 1, + rowsRound * C0_NUM_PER_FRACTAL); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (row, column) + CATLASS_HOST_DEVICE + LongIndex GetOffset(MatrixCoord const &coord) const + { + return LongIndex(coord.row()) / shape_[0] * stride_[1] + LongIndex(coord.column()) / shape_[2] * stride_[3] + + (LongIndex(coord.row()) % shape_[0]) * stride_[0] + (LongIndex(coord.column()) % shape_[2]) * stride_[2]; + } + + /// Returns the layout of a tile. + CATLASS_HOST_DEVICE + L0C GetTileLayout(MatrixCoord const &tileOriShape) const + { + auto tileShape = MakeCoord( + shape(0), CeilDiv(tileOriShape.row(), shape(0)), + shape(2), CeilDiv(tileOriShape.column(), shape(2))); + return L0C(tileOriShape, tileShape, stride()); + } + + /// Returns the origin shape of the layout + CATLASS_HOST_DEVICE + typename OrgShape::Index orgShape(int idx) const + { + return orgShape_[idx]; + } + + /// Returns the origin shape of the layout + CATLASS_HOST_DEVICE + typename OrgShape::Index &orgShape(int idx) + { + return orgShape_[idx]; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape shape() const + { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape &shape() + { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index shape(int idx) const + { + return shape_[idx]; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index &shape(int idx) + { + return shape_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride stride() const + { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride &stride() + { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index stride(int idx) const + { + return stride_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index &stride(int idx) + { + return stride_[idx]; + } + + /// Returns the length of the layout + CATLASS_HOST_DEVICE + LongIndex Capacity() const + { + return static_cast(stride_[3]) * shape_[3]; + } + +private: + /// Origin Shape data member + OrgShape orgShape_; + + /// Shape data member + Shape shape_; + + /// Stride data member + Stride stride_; +}; + +/// Mapping function for padding rowmajor matrices +/// A special data layout designed to improve the efficiency of matrix operations in non-512B aligned scenarios. +/// This layout is row-major within blocks and also row-major between blocks. +struct PaddingRowMajor { +public: + /// Logical rank of tensor + static constexpr int RANK = 4; + + /// Logical rank of orgshape + static constexpr int ORG_SHAPE_RANK = 2; + + /// Index type used for coordinates + using Index = uint32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate + using OrgShape = Coord; + + /// Logical coordinate + using Shape = Coord; + + /// Stride vector + using Stride = Coord; + +public: + /// Constructor + CATLASS_HOST_DEVICE + PaddingRowMajor(Index orgRows = 0, Index orgCols = 0, Index blockRows = 0, Index blockCols = 0) + : orgShape_(MakeCoord(orgRows, orgCols)), + shape_(MakeCoord(blockRows, CeilDiv(orgRows, blockRows), blockCols, CeilDiv(orgCols, blockCols))), + stride_(MakeCoord((LongIndex)blockCols, (LongIndex)blockRows * (LongIndex)RoundUp(orgCols, blockCols), + (LongIndex)1, (LongIndex)blockRows * (LongIndex)blockCols)) + {} + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (row, column) + CATLASS_HOST_DEVICE + LongIndex GetOffset(MatrixCoord const &coord) const + { + LongIndex blockRows = (LongIndex)shape_[0]; + LongIndex blockCols = (LongIndex)shape_[2]; + return (LongIndex)coord.row() / blockRows * stride_[1] + (LongIndex)coord.column() / blockCols * stride_[3] + (LongIndex)coord.row() % blockRows * stride_[0] + (LongIndex)coord.column() % blockCols; + } + + CATLASS_HOST_DEVICE + PaddingRowMajor GetTileLayout(MatrixCoord const &tileShape) const + { + return PaddingRowMajor(tileShape.row(), tileShape.column(), shape_[0], shape_[2]); + } + + /// Returns the origin shape of the layout + CATLASS_HOST_DEVICE + typename OrgShape::Index orgShape(int idx) const + { + return orgShape_[idx]; + } + + /// Returns the origin shape of the layout + CATLASS_HOST_DEVICE + typename OrgShape::Index &orgShape(int idx) + { + return orgShape_[idx]; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape shape() const + { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape &shape() + { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index shape(int idx) const + { + return shape_[idx]; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index &shape(int idx) + { + return shape_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride stride() const + { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride &stride() + { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index stride(int idx) const + { + return stride_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index &stride(int idx) + { + return stride_[idx]; + } + +private: + // + // Data members + // + + /// Origin Shape data member + OrgShape orgShape_; + + /// Shape data member + Shape shape_; + + /// Stride data member + Stride stride_; +}; + +/// Mapping function for padding columnmajor matrices +/// A special data layout designed to improve the efficiency of matrix operations in non-512B aligned scenarios. +/// This layout is column-major within blocks and also column-major between blocks. +struct PaddingColumnMajor { +public: + /// Logical rank of tensor + static constexpr int RANK = 4; + + /// Logical rank of orgshape + static constexpr int ORG_SHAPE_RANK = 2; + + /// Index type used for coordinates + using Index = uint32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate + using OrgShape = Coord; + + /// Logical coordinate + using Shape = Coord; + + /// Stride vector + using Stride = Coord; + +public: + /// Constructor + CATLASS_HOST_DEVICE + PaddingColumnMajor(Index orgRows = 0, Index orgCols = 0, Index blockRows = 0, Index blockCols = 0) + : orgShape_(MakeCoord(orgRows, orgCols)), + shape_(MakeCoord(blockRows, CeilDiv(orgRows, blockRows), blockCols, CeilDiv(orgCols, blockCols))), + stride_(MakeCoord((LongIndex)1, (LongIndex)blockRows * (LongIndex)blockCols, (LongIndex)blockRows, + (LongIndex)RoundUp(orgRows, blockRows) * (LongIndex)blockCols)) + {} + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (row, column) + CATLASS_HOST_DEVICE + LongIndex GetOffset(MatrixCoord const &coord) const + { + LongIndex blockRows = (LongIndex)shape_[0]; + LongIndex blockCols = (LongIndex)shape_[2]; + return (LongIndex)coord.row() / blockRows * stride_[1] + (LongIndex)coord.column() / blockCols * stride_[3] + (LongIndex)coord.row() % blockRows + (LongIndex)coord.column() % blockCols * stride_[2]; + } + + CATLASS_HOST_DEVICE + PaddingColumnMajor GetTileLayout(MatrixCoord const &tileShape) const + { + return PaddingColumnMajor(tileShape.row(), tileShape.column(), shape_[0], shape_[2]); + } + + /// Returns the origin shape of the layout + CATLASS_HOST_DEVICE + typename OrgShape::Index orgShape(int idx) const + { + return orgShape_[idx]; + } + + /// Returns the origin shape of the layout + CATLASS_HOST_DEVICE + typename OrgShape::Index &orgShape(int idx) + { + return orgShape_[idx]; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape shape() const + { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape &shape() + { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index shape(int idx) const + { + return shape_[idx]; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index &shape(int idx) + { + return shape_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride stride() const + { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride &stride() + { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index stride(int idx) const + { + return stride_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index &stride(int idx) + { + return stride_[idx]; + } + +private: + // + // Data members + // + + /// Origin Shape data member + OrgShape orgShape_; + + /// Shape data member + Shape shape_; + + /// Stride data member + Stride stride_; +}; + +/////////////////////// +// new add layout nN +// nN layout +struct nN { +public: + /// Logical rank of tensor + static constexpr int RANK = 4; + + /// Index type used for coordinates + using Index = uint32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical rank of orgshape + static constexpr int ORG_SHAPE_RANK = 2; + + /// Logical coordinate + using OrgShape = Coord; + + /// Logical coordinate + using Shape = Coord; + + /// Stride vector + using Stride = Coord; + +public: + // Methods + + /// Constructor + CATLASS_HOST_DEVICE + nN(Index orgRows = 0, /// Number of rows of origin matrices + Index orgCols = 0, /// Number of cols of origin matrices + + Index rowsInFractal = 0, /// Number of rows inside the fractal + Index rowsByFractal = 0, /// number of rows by the fractal + Index colsInFractal = 0, /// number of cols inside the fractal + Index colsByFractal = 0, /// number of cols by the fractal + + LongIndex strideRowsInFractal = 0, /// number of elements between adjacent rows inside the fractal + LongIndex strideRowsByFractal = 0, /// number of elements between adjacent fractal rows + LongIndex strideColsInFractal = 0, /// number of elements between adjacent cols inside the fractal + LongIndex strideColsByFractal = 0) /// number of elements between adjacent fractal cols + : orgShape_(MakeCoord(orgRows, orgCols)), + shape_(MakeCoord(rowsInFractal, rowsByFractal, colsInFractal, colsByFractal)), + stride_(MakeCoord(strideRowsInFractal, strideRowsByFractal, strideColsInFractal, strideColsByFractal)) + { + } + + /// Ctor + CATLASS_HOST_DEVICE + nN(OrgShape orgShape, Shape shape, Stride stride) + : orgShape_(orgShape), + shape_(shape), + stride_(stride) + {} + + /// Make the layout of a coordinate (row, column) + template + CATLASS_HOST_DEVICE static nN MakeLayout(Index orgRows, Index orgCols) + { + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + Index rowsRound = RoundUp(orgRows); + Index colsRound = RoundUp(orgCols); + return nN(orgRows, + orgCols, + + ELE_NUM_PER_C0, + rowsRound / ELE_NUM_PER_C0, + C0_NUM_PER_FRACTAL, + colsRound / C0_NUM_PER_FRACTAL, + + 1, + ELE_NUM_PER_FRACTAL, + ELE_NUM_PER_C0, + rowsRound * C0_NUM_PER_FRACTAL); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (row, column) + CATLASS_HOST_DEVICE + LongIndex GetOffset(MatrixCoord const &coord) const + { + return LongIndex(coord.row()) / shape_[0] * stride_[1] + LongIndex(coord.column()) / shape_[2] * stride_[3]; + } + + /// Returns the origin shape of the layout + CATLASS_HOST_DEVICE + typename OrgShape::Index orgShape(int idx) const + { + return orgShape_[idx]; + } + + /// Returns the origin shape of the layout + CATLASS_HOST_DEVICE + typename OrgShape::Index &orgShape(int idx) + { + return orgShape_[idx]; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape shape() const + { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape &shape() + { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index shape(int idx) const + { + return shape_[idx]; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index &shape(int idx) + { + return shape_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride stride() const + { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride &stride() + { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index stride(int idx) const + { + return stride_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index &stride(int idx) + { + return stride_[idx]; + } + +private: + /// Origin Shape data member + OrgShape orgShape_; + + /// Shape data member + Shape shape_; + + /// Stride data member + Stride stride_; +}; + +struct NDC1HWC0 { +public: + /// Index type used for coordinates + using Index = uint32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical rank of orgshape + /// (N,D,C1,H,W,C0) + static constexpr int ORG_SHAPE_RANK = 6; + + static constexpr int RANK = 5; + using OrgShape = Coord; + /// Logical coordinate + using Shape = Coord; + + /// Stride vector + using Stride = Coord; + +public: + // Methods + + /// Constructor + CATLASS_HOST_DEVICE constexpr NDC1HWC0( + Index batch = 0, + Index D = 0, + Index C1 = 0, + Index H = 0, + Index W = 0, + Index C0 = 0, + + Index rowsInFractal = 0, /// Number of rows inside the fractal + Index rowsByFractal = 0, /// number of rows by the fractal + Index colsInFractal = 0, /// number of cols inside the fractal + Index colsByFractal = 0, /// number of cols by the fractal + + LongIndex strideC0 = 0, /// number of elements between adjacent C0 cols + LongIndex strideHW = 0, /// number of elements between adjacent W rows + LongIndex StrideC1 = 0, /// number of elements between adjacent C1 cols + LongIndex StrideD = 0, /// number of elements between adjacent D batchCols + LongIndex StrideN = 0 /// number of elements between adjacent batch + ) + : orgShape_(MakeCoord(batch, D, C1, H, W, C0)), + shape_(MakeCoord(batch, rowsInFractal, rowsByFractal, colsInFractal, colsByFractal)), + stride_(MakeCoord(strideC0, strideHW, StrideC1, StrideD, StrideN)) + {} + + /// Ctor + CATLASS_HOST_DEVICE constexpr NDC1HWC0(OrgShape orgshape, Shape shape, Stride stride) + : orgShape_(orgshape), + shape_(shape), + stride_(stride) + {} + + CATLASS_HOST_DEVICE constexpr static NDC1HWC0 MakeLayout(Index Batch, Index D, Index C1, Index H, Index W, Index C0) + { + return NDC1HWC0(Batch, + D, + C1, + H, + W, + C0, + + W, + H, + C0, + D * C1, + + 1, /// StrideC0 + C0, /// StrideHW + H * W * C0, /// StrideC1 + H * W * C0 * C1, /// StrideD + H * W * C0 * C1 * D /// StrideN + ); + } + + // CATLASS_HOST_DEVICE + /// Returns the offset of a coordinate in linear memory. + CATLASS_HOST_DEVICE + LongIndex GetOffset(Conv3d6HdCoord const &coord) const + { + return LongIndex(coord.n()) * stride_[4] + LongIndex(coord.d()) * stride_[3] + + LongIndex(coord.c1()) * stride_[2] + LongIndex(coord.hw()) * stride_[1]; + } + + /// Returns the layout of a tile. + CATLASS_HOST_DEVICE + NDC1HWC0 GetTileLayout(OrgShape const &tileOriShape) const + { + Shape tileShape = MakeCoord( + tileOriShape[0], + tileOriShape[4], + tileOriShape[3], + shape(3), + tileOriShape[1] * tileOriShape[2]); + + Stride tileStride = MakeCoord(stride(0), + stride(1), + (LongIndex)(tileOriShape[3] * tileOriShape[4] * shape(3)), + (LongIndex)(tileOriShape[2] * tileOriShape[3] * tileOriShape[4] * shape(3)), + (LongIndex)(tileOriShape[1] * tileOriShape[2] * tileOriShape[3] * tileOriShape[4] * shape(3))); + return NDC1HWC0(tileOriShape, tileShape, tileStride); + } + + /// Returns the origin shape of the layout + CATLASS_HOST_DEVICE + typename OrgShape::Index orgShape(int idx) const + { + return orgShape_[idx]; + } + + /// Returns the origin shape of the layout + CATLASS_HOST_DEVICE + typename OrgShape::Index &orgShape(int idx) + { + return orgShape_[idx]; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape shape() const + { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape &shape() + { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index shape(int idx) const + { + return shape_[idx]; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index &shape(int idx) + { + return shape_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride stride() const + { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride &stride() + { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index stride(int idx) const + { + return stride_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index &stride(int idx) + { + return stride_[idx]; + } + +private: + /// Origin Shape data member + OrgShape orgShape_; + + /// Shape data member + Shape shape_; + + /// Stride data member + Stride stride_; +}; + +struct KDC1KHKWN1N0C0 { +public: + static constexpr int RANK = 4; + + /// Index type used for coordinates + using Index = uint32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical rank of orgshape + static constexpr int ORG_SHAPE_RANK = 4; + + /// Logical coordinate + using OrgShape = Coord; + + /// Logical coordinate + using Shape = Coord; + + /// Stride vector + using Stride = Coord; + +public: + // Methods + + /// Constructor + CATLASS_HOST_DEVICE constexpr KDC1KHKWN1N0C0(Index KdC1KhKw = 0, /// Merging Kd,Kh,Kw,C1 axes of KDC1KHKWN1N0C0 + Index N1 = 0, /// Cout = N1*N0 + Index N0 = 0, + Index C0 = 0, + + Index rowsInFractal = 0, /// Number of rows inside the fractal + Index rowsByFractal = 0, /// number of rows by the fractal + Index colsInFractal = 0, /// number of cols inside the fractal + Index colsByFractal = 0, /// number of cols by the fractal + + LongIndex strideC0 = 0, /// number of elements between adjacent rows inside the fractal + LongIndex StrideDC1HW = 0, /// number of elements between adjacent fractal rows + LongIndex strideN0 = 0, /// number of elements between adjacent cols inside the fractal + LongIndex strideN1 = 0 /// number of elements between adjacent fractal cols + ) + : orgShape_(MakeCoord(KdC1KhKw, N1, N0, C0)), + shape_(MakeCoord(rowsInFractal, rowsByFractal, colsInFractal, colsByFractal)), + stride_(MakeCoord(strideC0, strideN0, strideN1, StrideDC1HW)) + {} + + /// Ctor + CATLASS_HOST_DEVICE constexpr KDC1KHKWN1N0C0(OrgShape orgShape, Shape shape, Stride stride) + : orgShape_(orgShape), + shape_(shape), + stride_(stride) + {} + + /// Make the layout of a coordinate (Kd*C1*Kh*Kw,N1,N0,C0) + CATLASS_HOST_DEVICE constexpr static KDC1KHKWN1N0C0 MakeLayout(Index KdC1KhKw, Index N1, Index N0, Index C0) + { + return KDC1KHKWN1N0C0(KdC1KhKw, + N1, + N0, + C0, + + C0, + KdC1KhKw, + N0, + N1, + + 1, /// StrideC0 + C0 * N0 * N1, /// StrideDC1HW + C0, /// StrideN0 + C0 * N0 /// StrideN1 + ); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (KdC1KhKw_idx, N1_idx) + CATLASS_HOST_DEVICE + LongIndex GetOffset(Conv3dFracZ3dCoord const &coord) const + { + return LongIndex(coord.kdc1khkw()) * stride_[3] + LongIndex(coord.n1()) * stride_[1]; + } + + /// Returns the layout of a tile. + CATLASS_HOST_DEVICE + KDC1KHKWN1N0C0 GetTileLayout(OrgShape const &tileOriShape) const + { + Shape tileShape = MakeCoord(shape(0), /// C0 + tileOriShape[0], /// Kd*C1*Kh*Kw + shape(2), /// N0 + tileOriShape[1] /// N1 + ); + Stride tileStride = MakeCoord(stride(0), /// TileStrideC0 + stride(2) * tileOriShape[1] * tileOriShape[2], /// TileStrideDC1HW + (LongIndex)shape(0), /// TileStrideN0 + stride(2) * tileOriShape[2] /// TileStrideN1 + ); + return KDC1KHKWN1N0C0(tileOriShape, tileShape, tileStride); + } + + /// Returns the origin shape of the layout + CATLASS_HOST_DEVICE + typename OrgShape::Index orgShape(int idx) const + { + return orgShape_[idx]; + } + + /// Returns the origin shape of the layout + CATLASS_HOST_DEVICE + typename OrgShape::Index &orgShape(int idx) + { + return orgShape_[idx]; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape shape() const + { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape &shape() + { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index shape(int idx) const + { + return shape_[idx]; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index &shape(int idx) + { + return shape_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride stride() const + { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride &stride() + { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index stride(int idx) const + { + return stride_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index &stride(int idx) + { + return stride_[idx]; + } + +private: + /// Origin Shape data member + OrgShape orgShape_; + + /// Shape data member + Shape shape_; + + /// Stride data member + Stride stride_; +}; +} // namespace Catlass::layout + +#endif // CATLASS_KDA_MATRIX_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/layout/kda_tensor.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/layout/kda_tensor.hpp new file mode 100644 index 0000000000..3e8d607104 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/layout/kda_tensor.hpp @@ -0,0 +1,315 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ +#ifndef CATLASS_KDA_TENSOR_HPP +#define CATLASS_KDA_TENSOR_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/kda_conv_coord.hpp" +#include "catlass/detail/kda_alignment.hpp" + +namespace Catlass::layout { + +struct NC1HWC0 { // (Batch, C1, H, W, C0) + /// Logical rank of tensor + static constexpr int RANK = 5; + + /// Index type used for coordinates + using Index = uint32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate + using Shape = Coord; + + /// Stride vector + using Stride = Coord; + +public: + /// Ctor + CATLASS_HOST_DEVICE + NC1HWC0(Index batch = 0, Index c1 = 0, Index h = 0, Index w = 0, Index c0 = 0) + : shape_(MakeCoord(batch, c1, h, w, c0)) + { + LongIndex strideBatch = c1 * h * w * c0; + LongIndex strideC1 = h * w * c0; + LongIndex strideH = w * c0; + LongIndex strideW = c0; + LongIndex strideC0 = 1; + stride_ = MakeCoord(strideBatch, strideC1, strideH, strideW, strideC0); + } + + CATLASS_HOST_DEVICE + NC1HWC0( + Index batch, + Index c1, + Index h, + Index w, + Index c0, + LongIndex strideBatch, + LongIndex strideC1, + LongIndex strideH, + LongIndex strideW, + LongIndex strideC0) + : shape_(MakeCoord(batch, c1, h, w, c0)), + stride_(MakeCoord(strideBatch, strideC1, strideH, strideW, strideC0)) + { + } + + CATLASS_HOST_DEVICE + NC1HWC0(Shape shape, Stride stride) + : shape_(shape), + stride_(stride) + { + } + + /// Make the layout of a coordinate (batch, c1, h, w, c0) + template + CATLASS_HOST_DEVICE static NC1HWC0 MakeLayout(Index batch, Index c1, Index h, Index w, Index c0) + { + return NC1HWC0(batch, c1, h, w, c0); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (batch, c1, h, w, c0) + CATLASS_HOST_DEVICE + LongIndex GetOffset(Conv2dFmapCoord const &coord) const + { + return LongIndex(coord.batch()) * stride_[0] + LongIndex(coord.c1()) * stride_[1] + LongIndex(coord.h()) * stride_[2] + LongIndex(coord.w()) * stride_[3] + LongIndex(coord.c0()) * stride_[4]; + } + + /// Returns the layout of a tile. + CATLASS_HOST_DEVICE + NC1HWC0 GetTileLayout(Conv2dFmapCoord const &tileShape) const + { + return NC1HWC0(tileShape, stride()); + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape shape() const + { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape &shape() + { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index shape(int idx) const + { + return shape_[idx]; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index &shape(int idx) + { + return shape_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride stride() const + { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride &stride() + { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index stride(int idx) const + { + return stride_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index &stride(int idx) + { + return stride_[idx]; + } + + /// Returns the length of the layout + CATLASS_HOST_DEVICE + size_t Capacity() + { + return static_cast(shape_[0]) * stride_[0]; + } + +private: + /// Shape data member + Shape shape_; + + /// Stride data member + Stride stride_; +}; + +struct CI1KHKWCOCI0 { // (Cin1, Kh, Kw, Cout, C0) + /// Logical rank of tensor + static constexpr int RANK = 5; + + /// Index type used for coordinates + using Index = uint32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate + using Shape = Coord; + + /// Stride vector + using Stride = Coord; + +public: + /// Ctor + CATLASS_HOST_DEVICE + CI1KHKWCOCI0(Index cin1 = 0, Index kh = 0, Index kw = 0, Index cout = 0, Index c0 = 0) + : shape_(MakeCoord(cin1, kh, kw, cout, c0)) + { + LongIndex strideCin1 = kh * kw * cout * c0; + LongIndex strideKh = kw * cout * c0; + LongIndex strideKw = cout * c0; + LongIndex strideCout = c0; + LongIndex strideC0 = 1; + stride_ = MakeCoord(strideCin1, strideKh, strideKw, strideCout, strideC0); + } + + CATLASS_HOST_DEVICE + CI1KHKWCOCI0( + Index cin1, + Index kh, + Index kw, + Index cout, + Index c0, + LongIndex strideCin1, + LongIndex strideKh, + LongIndex strideKw, + LongIndex strideCout, + LongIndex strideC0) + : shape_(MakeCoord(cin1, kh, kw, cout, c0)), + stride_(MakeCoord(strideCin1, strideKh, strideKw, strideCout, strideC0)) + { + } + + CATLASS_HOST_DEVICE + CI1KHKWCOCI0(Shape shape, Stride stride) + : shape_(shape), + stride_(stride) + { + } + + /// Make the layout of a coordinate (cin1, h, w, c0) + template + CATLASS_HOST_DEVICE static CI1KHKWCOCI0 MakeLayout(Index cin1, Index kh, Index kw, Index cout, Index c0) + { + return CI1KHKWCOCI0(cin1, kh, kw, cout, c0); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (cin1, kh, kw, cout, c0) + CATLASS_HOST_DEVICE + LongIndex GetOffset(Conv2dFilterCoord const &coord) const + { + return LongIndex(coord.cin1()) * stride_[0] + LongIndex(coord.kh()) * stride_[1] + LongIndex(coord.kw()) * stride_[2] + LongIndex(coord.cout()) * stride_[3] + LongIndex(coord.c0()) * stride_[4]; + } + + /// Returns the layout of a tile. + CATLASS_HOST_DEVICE + CI1KHKWCOCI0 GetTileLayout(Conv2dFilterCoord const &tileShape) const + { + return CI1KHKWCOCI0(tileShape, stride()); + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape shape() const + { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape &shape() + { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index shape(int idx) const + { + return shape_[idx]; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index &shape(int idx) + { + return shape_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride stride() const + { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride &stride() + { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index stride(int idx) const + { + return stride_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index &stride(int idx) + { + return stride_[idx]; + } + + /// Returns the length of the layout + CATLASS_HOST_DEVICE + size_t Capacity() + { + return static_cast(shape_[0]) * stride_[0]; + } + +private: + /// Shape data member + Shape shape_; + + /// Stride data member + Stride stride_; +}; + +} // namespace Catlass::layout + +#endif // CATLASS_KDA_TENSOR_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/layout/kda_vector.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/layout/kda_vector.hpp new file mode 100644 index 0000000000..2d7750b0c5 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/layout/kda_vector.hpp @@ -0,0 +1,139 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_VECTOR_HPP +#define CATLASS_KDA_VECTOR_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/kda_coord.hpp" +#include "catlass/kda_numeric_size.hpp" + +namespace Catlass::layout { + +struct VectorLayout { +public: + /// Logical rank of tensor + static constexpr int RANK = 1; + + /// Index type used for coordinates + using Index = uint32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Shape vector + using Shape = Coord; + + /// Stride vector + using Stride = Coord; + + /// Logical coordinate + using TensorCoord = Coord; + +public: + // Methods + + CATLASS_HOST_DEVICE + VectorLayout(Index size = 0) + : shape_(MakeCoord(size)), + stride_(MakeCoord(LongIndex(1))) + {} + + CATLASS_HOST_DEVICE + VectorLayout(Shape shape, Stride stride) + : shape_(shape), + stride_(stride) + {} + + template + CATLASS_HOST_DEVICE static VectorLayout MakeLayoutInUb(TensorCoord const &tileShape) + { + constexpr uint32_t ELE_NUM_PER_BLK = BytesToBits(BYTE_PER_BLK) / SizeOfBits::value; + return VectorLayout{ELE_NUM_PER_BLK > (tileShape[0])}; + } + + CATLASS_HOST_DEVICE + LongIndex GetOffset(TensorCoord const &coord) const + { + return stride_[0] * coord[0]; + } + + /// Returns the layout of a tile. + CATLASS_HOST_DEVICE + VectorLayout GetTileLayout(TensorCoord const &tileShape) const + { + return VectorLayout(tileShape, stride()); + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape shape() const + { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape &shape() + { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index shape(int idx) const + { + return shape_[idx]; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index &shape(int idx) + { + return shape_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride stride() const + { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride &stride() + { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index stride(int idx) const + { + return stride_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index &stride(int idx) + { + return stride_[idx]; + } + +private: + /// Stride data member + Shape shape_; + Stride stride_; +}; + +} // namespace Catlass::layout + +#endif // CATLASS_KDA_VECTOR_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd.cpp b/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd.cpp new file mode 100644 index 0000000000..fc2321360a --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd.cpp @@ -0,0 +1,188 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#include "kernel_operator.h" +#include "lib/matmul_intf.h" + +#include "chunk_kda_fwd_common.h" +#include "chunk_kda_fwd_tiling_data.h" +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 +#define KDA_COMPILE_ARCH35_FAST_PATH 1 +#include "arch35/chunk_kda_fwd_impl.h" +#else +#define KDA_COMPILE_ARCH35_FAST_PATH 0 +#endif + +namespace KdaForward { + +template +__aicore__ inline void DispatchGeneric( + GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR g, GM_ADDR beta, + GM_ADDR aLog, GM_ADDR dtBias, GM_ADDR initialState, + GM_ADDR cuSeqlens, GM_ADDR chunkIndices, GM_ADDR attnOut, + GM_ADDR finalState, GM_ADDR gk, GM_ADDR aqk, GM_ADDR akk, + GM_ADDR w, GM_ADDR u, GM_ADDR qg, GM_ADDR kg, GM_ADDR vNew, GM_ADDR h, + GM_ADDR userWorkspace, const TilingData &tiling) +{ + RunGeneric( + q, k, v, g, beta, aLog, dtBias, initialState, cuSeqlens, + chunkIndices, attnOut, finalState, gk, aqk, akk, w, u, qg, kg, + vNew, h, userWorkspace, tiling); +} + +#if KDA_COMPILE_ARCH35_FAST_PATH +template +__aicore__ inline void DispatchArch35SafeGate( + GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR g, GM_ADDR beta, + GM_ADDR aLog, GM_ADDR dtBias, GM_ADDR initialState, + GM_ADDR cuSeqlens, GM_ADDR chunkIndices, GM_ADDR attnOut, + GM_ADDR finalState, GM_ADDR gk, GM_ADDR aqk, GM_ADDR akk, + GM_ADDR w, GM_ADDR u, GM_ADDR qg, GM_ADDR kg, GM_ADDR vNew, GM_ADDR h, + GM_ADDR userWorkspace, const TilingData &tiling) +{ + AscendC::TPipe pipe; + if (tiling.safeGate) { + arch35::Run( + q, k, v, g, beta, aLog, dtBias, initialState, cuSeqlens, + chunkIndices, attnOut, finalState, gk, aqk, akk, w, u, qg, + kg, vNew, h, userWorkspace, tiling, pipe); + } else { + arch35::Run( + q, k, v, g, beta, aLog, dtBias, initialState, cuSeqlens, + chunkIndices, attnOut, finalState, gk, aqk, akk, w, u, qg, + kg, vNew, h, userWorkspace, tiling, pipe); + } +} +#elif defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 +template +__aicore__ inline void DispatchArch35SafeGate( + GM_ADDR, GM_ADDR, GM_ADDR, GM_ADDR, GM_ADDR, + GM_ADDR, GM_ADDR, GM_ADDR, + GM_ADDR, GM_ADDR, + GM_ADDR, GM_ADDR, GM_ADDR, GM_ADDR, GM_ADDR, + GM_ADDR, GM_ADDR, GM_ADDR, GM_ADDR, GM_ADDR, GM_ADDR, + GM_ADDR, + const TilingData &) +{ +} +#endif + +template +__aicore__ inline void DispatchGenericSafeGate( + GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR g, GM_ADDR beta, + GM_ADDR aLog, GM_ADDR dtBias, GM_ADDR initialState, + GM_ADDR cuSeqlens, GM_ADDR chunkIndices, GM_ADDR attnOut, + GM_ADDR finalState, GM_ADDR gk, GM_ADDR aqk, GM_ADDR akk, + GM_ADDR w, GM_ADDR u, GM_ADDR qg, GM_ADDR kg, GM_ADDR vNew, GM_ADDR h, + GM_ADDR userWorkspace, const TilingData &tiling) +{ + if (tiling.safeGate) { + DispatchGeneric( + q, k, v, g, beta, aLog, dtBias, initialState, cuSeqlens, + chunkIndices, attnOut, finalState, gk, aqk, akk, w, u, qg, + kg, vNew, h, userWorkspace, tiling); + } else { + DispatchGeneric( + q, k, v, g, beta, aLog, dtBias, initialState, cuSeqlens, + chunkIndices, attnOut, finalState, gk, aqk, akk, w, u, qg, + kg, vNew, h, userWorkspace, tiling); + } +} + +} // namespace KdaForward + +using ChunkKdaFwd::ChunkKdaFwdTilingData; + +extern "C" __global__ __aicore__ void chunk_kda_fwd( + GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR g, GM_ADDR beta, + GM_ADDR a_log, GM_ADDR dt_bias, GM_ADDR initial_state, + GM_ADDR cu_seqlens, GM_ADDR chunk_indices, GM_ADDR attn_out, + GM_ADDR final_state, GM_ADDR gk, GM_ADDR aqk, GM_ADDR akk, + GM_ADDR w, GM_ADDR u, GM_ADDR qg, GM_ADDR kg, GM_ADDR v_new, GM_ADDR h, + GM_ADDR workspace, GM_ADDR tiling) +{ + KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIC_1_2); + REGISTER_TILING_DEFAULT(ChunkKdaFwdTilingData); + GM_ADDR userWorkspace = AscendC::GetUserWorkspace(workspace); + + // Direct-launch: copy the tiling blob from GM to the stack. The aclnn + // GET_TILING_DATA_WITH_STRUCT macro is unavailable under the plain AscendC + // direct-launch compilation, so the host-serialized struct is rehydrated + // here exactly like chunk_gated_delta_rule does. + constexpr size_t TILING_WORDS = (sizeof(ChunkKdaFwdTilingData) + 7) / 8; + alignas(8) uint64_t tilingBuf[TILING_WORDS]; + const __gm__ uint64_t *gmTilingWords = reinterpret_cast(tiling); + for (size_t i = 0; i < TILING_WORDS; ++i) { + tilingBuf[i] = gmTilingWords[i]; + } + ChunkKdaFwdTilingData *tilingData = reinterpret_cast(tilingBuf); + + // Tiling key 2 is selected when chunk_size == 64 and K == V == 128; the + // host encodes it implicitly through these shape fields, so dispatch on + // them instead of the aclnn TILING_KEY_IS compile-time key. + const bool useFastPath = tilingData->chunkSize == 64 && tilingData->kHeadDim == 128 && + tilingData->vHeadDim == 128; + if (useFastPath) { +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if (tilingData->betaDataType == 0) { + KdaForward::DispatchArch35SafeGate( + q, k, v, g, beta, a_log, dt_bias, initial_state, cu_seqlens, + chunk_indices, attn_out, final_state, gk, aqk, akk, w, u, qg, + kg, v_new, h, userWorkspace, *tilingData); + } else { + KdaForward::DispatchArch35SafeGate( + q, k, v, g, beta, a_log, dt_bias, initial_state, cu_seqlens, + chunk_indices, attn_out, final_state, gk, aqk, akk, w, u, qg, + kg, v_new, h, userWorkspace, *tilingData); + } +#else + if (tilingData->betaDataType == 0) { + KdaForward::DispatchGenericSafeGate( + q, k, v, g, beta, a_log, dt_bias, initial_state, cu_seqlens, + chunk_indices, attn_out, final_state, gk, aqk, akk, w, u, qg, + kg, v_new, h, userWorkspace, *tilingData); + } else { + KdaForward::DispatchGenericSafeGate( + q, k, v, g, beta, a_log, dt_bias, initial_state, cu_seqlens, + chunk_indices, attn_out, final_state, gk, aqk, akk, w, u, qg, + kg, v_new, h, userWorkspace, *tilingData); + } +#endif + } else { + if (tilingData->betaDataType == 0) { + KdaForward::DispatchGenericSafeGate( + q, k, v, g, beta, a_log, dt_bias, initial_state, cu_seqlens, + chunk_indices, attn_out, final_state, gk, aqk, akk, w, u, qg, + kg, v_new, h, userWorkspace, *tilingData); + } else { + KdaForward::DispatchGenericSafeGate( + q, k, v, g, beta, a_log, dt_bias, initial_state, cu_seqlens, + chunk_indices, attn_out, final_state, gk, aqk, akk, w, u, qg, + kg, v_new, h, userWorkspace, *tilingData); + } + } +} + +#undef KDA_COMPILE_ARCH35_FAST_PATH diff --git a/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_common.h b/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_common.h new file mode 100644 index 0000000000..29f5f234a8 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_common.h @@ -0,0 +1,331 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#pragma once + +#include "kernel_operator.h" +#include "kda_gate_cumsum_kernel.h" +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 +#include "arch35/chunk_kda_fwd_prepare_arch35.h" +#include "arch35/chunk_kda_fwd_post_wu_arch35.h" +#include "arch35/chunk_kda_fwd_finalize_arch35.h" +#else +#include "chunk_kda_fwd_prepare.h" +#include "chunk_kda_fwd_post_wu.h" +#include "chunk_kda_fwd_finalize.h" +#endif + +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 +#include "arch35/fwd_h/gemm/kernel/kda_fwd_h_kernel_arch35.hpp" +#else +#include "arch22/fwd_h/gemm/kernel/kda_fwd_h_kernel_arch22.hpp" +#endif + +namespace KdaForward { + +using namespace AscendC; + +struct GateRuntimeTiling { + int64_t batch; + int64_t t; + int64_t hv; + int64_t k; + int64_t rank; + int64_t chunkSize; + int64_t seqNum; + int64_t hasCuSeqlens; + int64_t hasALog; + int64_t hasDtBias; + int64_t dataType; + int64_t useGateInKernel; + int64_t safeGate; + int64_t inputSequenceMajor; + float lowerBound; + int64_t usedCoreNum; +}; + +struct ChunkKdaFwdAddresses { + GM_ADDR gk; + GM_ADDR finalState; + GM_ADDR w; + GM_ADDR u; + GM_ADDR qg; + GM_ADDR kg; + GM_ADDR vNew; + GM_ADDR h; + GM_ADDR qgScaled; + GM_ADDR uSeed; +}; + +struct FwdHTilingView { + int64_t batch; + int64_t seqlen; + int64_t kNumHead; + int64_t vNumHead; + int64_t kHeadDim; + int64_t vHeadDim; + int64_t chunkSize; + bool useInitialState; + bool storeFinalState; + int64_t isVariedLen; + int64_t shapeBatch; + int64_t tokenBatch; + int64_t vWorkspaceOffset; + int64_t vUpdateWorkspaceOffset; + int64_t kDecayWorkspaceOffset; + int64_t hWorkspaceOffset; + int64_t numSeqWorkspaceOffset; + int64_t numChunksWorkspaceOffset; +}; + +template +__aicore__ inline FwdHTilingView MakeFwdHTiling(const TilingData &tiling) +{ + constexpr bool persistFwdHStateInFp32 = true; + return { + tiling.isVarLen ? tiling.seqNum : tiling.batch, + tiling.seqlen, + tiling.vHeadNum, + tiling.vHeadNum, + tiling.kHeadDim, + tiling.vHeadDim, + tiling.chunkSize, + tiling.hasInitialState, + persistFwdHStateInFp32, + tiling.isVarLen ? 1 : 0, + tiling.isVarLen ? 1 : tiling.batch, + tiling.isVarLen ? tiling.seqNum : 1, + tiling.vWorkspaceOffset, + tiling.vUpdateWorkspaceOffset, + tiling.kDecayWorkspaceOffset, + tiling.hWorkspaceOffset, + tiling.numSeqWorkspaceOffset, + tiling.numChunksWorkspaceOffset, + }; +} + +__aicore__ inline GM_ADDR ResolveStorage( + GM_ADDR output, GM_ADDR userWorkspace, int64_t offset, bool storeOutput) +{ + return storeOutput ? output : userWorkspace + offset; +} + +template +__aicore__ inline ChunkKdaFwdAddresses ResolveAddresses( + GM_ADDR finalState, GM_ADDR gk, GM_ADDR w, GM_ADDR u, GM_ADDR qg, + GM_ADDR kg, GM_ADDR vNew, GM_ADDR h, GM_ADDR userWorkspace, + const TilingData &tiling) +{ + return { + ResolveStorage(gk, userWorkspace, tiling.gkStorageOffset, tiling.storeGk), + ResolveStorage(finalState, userWorkspace, tiling.finalStateStorageOffset, + tiling.storeFinalState), + ResolveStorage(w, userWorkspace, tiling.wStorageOffset, tiling.storeW), + ResolveStorage(u, userWorkspace, tiling.uStorageOffset, tiling.storeU), + ResolveStorage(qg, userWorkspace, tiling.qgStorageOffset, tiling.storeQG), + ResolveStorage(kg, userWorkspace, tiling.kgStorageOffset, tiling.storeKg), + ResolveStorage(vNew, userWorkspace, tiling.vNewStorageOffset, tiling.storeVNew), + ResolveStorage(h, userWorkspace, tiling.hStorageOffset, tiling.storeH), + userWorkspace + tiling.qgScaledOffset, + userWorkspace + tiling.outputScratchOffset, + }; +} + +template +__aicore__ inline GateRuntimeTiling MakeGateTiling(const TilingData &tiling) +{ + return { + tiling.batch, + tiling.seqlen, + tiling.vHeadNum, + tiling.kHeadDim, + tiling.inputRank, + tiling.chunkSize, + tiling.seqNum, + tiling.isVarLen ? 1 : 0, + tiling.hasALog ? 1 : 0, + tiling.hasDtBias ? 1 : 0, + tiling.gateDataType, + tiling.useGateInKernel ? 1 : 0, + tiling.safeGate ? 1 : 0, + tiling.inputSequenceMajor ? 1 : 0, + tiling.lowerBound, + tiling.gateUsedCoreNum, + }; +} + +template +__aicore__ inline void RunGateCumsum( + GM_ADDR g, GM_ADDR aLog, GM_ADDR dtBias, GM_ADDR cuSeqlens, + GM_ADDR gk, const TilingData &tiling) +{ + if (tiling.computeGateInPrepare) { + return; + } + if ASCEND_IS_AIV { + GateRuntimeTiling gateTiling = MakeGateTiling(tiling); + TPipe gatePipe; + if (gateTiling.dataType == 2) { + KdaGateCumsum::DispatchKdaGateCumsum( + g, aLog, dtBias, cuSeqlens, gk, gateTiling, &gatePipe); + } else if (gateTiling.dataType == 1) { + KdaGateCumsum::DispatchKdaGateCumsum( + g, aLog, dtBias, cuSeqlens, gk, gateTiling, &gatePipe); + } else { + KdaGateCumsum::DispatchKdaGateCumsum( + g, aLog, dtBias, cuSeqlens, gk, gateTiling, &gatePipe); + } + } +} + +template +__aicore__ inline void RunFrontEnd( + GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR g, GM_ADDR beta, + GM_ADDR aLog, GM_ADDR dtBias, GM_ADDR initialState, + GM_ADDR cuSeqlens, GM_ADDR chunkIndices, GM_ADDR aqk, GM_ADDR akk, + const ChunkKdaFwdAddresses &addresses, GM_ADDR userWorkspace, + const TilingData &tiling, TPipe &pipe) +{ + RunGateCumsum(g, aLog, dtBias, cuSeqlens, addresses.gk, tiling); + if (!tiling.computeGateInPrepare) { + SyncAll(); + } + GM_ADDR uSeed = (tiling.fusePostWu || tiling.fusePostWuIntoFwdH) ? addresses.u : addresses.uSeed; + + KdaPrepare::RunChunkKdaPrepare( + q, k, v, addresses.gk, g, aLog, dtBias, beta, initialState, + cuSeqlens, chunkIndices, aqk, akk, addresses.qg, + addresses.qgScaled, addresses.w, uSeed, addresses.kg, + userWorkspace, tiling, pipe, tiling.storeQG); + SyncAll(); + pipe.Reset(); + + if (!tiling.fusePostWu && !tiling.fusePostWuIntoFwdH) { + KdaPostWu::RunChunkKdaPostWu( + q, k, v, addresses.gk, beta, initialState, cuSeqlens, + chunkIndices, addresses.w, akk, uSeed, + addresses.w, addresses.u, addresses.kg, addresses.vNew, + userWorkspace, tiling, pipe); + SyncAll(); + pipe.Reset(); + } +} + +template +__aicore__ inline void RunFwdH( + GM_ADDR initialState, GM_ADDR cuSeqlens, GM_ADDR chunkIndices, + const ChunkKdaFwdAddresses &addresses, GM_ADDR userWorkspace, + const TilingData &tiling) +{ + using FwdHKernel = Catlass::Gemm::Kernel::KdaFwdHKernel< + T, float, float, float, TileShapes, true, false, true>; + const auto fwdHTiling = MakeFwdHTiling(tiling); + FwdHKernel stateOp; + stateOp.InitFromData( + addresses.kg, addresses.w, addresses.u, addresses.gk, addresses.gk, + initialState, cuSeqlens, chunkIndices, addresses.h, addresses.vNew, + addresses.finalState, fwdHTiling, + userWorkspace + tiling.fwdHWorkspaceBaseOffset); + stateOp.Process(); +} + +template +__aicore__ inline void RunGenericBackEnd( + GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR beta, GM_ADDR initialState, + GM_ADDR cuSeqlens, GM_ADDR chunkIndices, GM_ADDR aqk, GM_ADDR attnOut, + const ChunkKdaFwdAddresses &addresses, GM_ADDR userWorkspace, + const TilingData &tiling) +{ + if (tiling.vHeadDim > 128) { + RunFwdH( + initialState, cuSeqlens, chunkIndices, addresses, + userWorkspace, tiling); + } else { + RunFwdH( + initialState, cuSeqlens, chunkIndices, addresses, + userWorkspace, tiling); + } + SyncAll(); + TPipe pipe; + KdaFinalize::RunChunkKdaOutput( + q, k, v, addresses.gk, beta, initialState, cuSeqlens, + chunkIndices, addresses.qgScaled, aqk, + addresses.vNew, addresses.h, attnOut, userWorkspace, tiling, pipe); +} + +template +__aicore__ inline void RunGenericBackEnd( + GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR beta, GM_ADDR initialState, + GM_ADDR cuSeqlens, GM_ADDR chunkIndices, GM_ADDR aqk, GM_ADDR attnOut, + const ChunkKdaFwdAddresses &addresses, GM_ADDR userWorkspace, + const TilingData &tiling, TPipe &pipe) +{ + if (tiling.vHeadDim > 128) { + RunFwdH( + initialState, cuSeqlens, chunkIndices, addresses, + userWorkspace, tiling); + } else { + RunFwdH( + initialState, cuSeqlens, chunkIndices, addresses, + userWorkspace, tiling); + } + SyncAll(); + KdaFinalize::RunChunkKdaOutput( + q, k, v, addresses.gk, beta, initialState, cuSeqlens, + chunkIndices, addresses.qgScaled, aqk, + addresses.vNew, addresses.h, attnOut, userWorkspace, tiling, pipe); +} + +template +__aicore__ inline void RunGeneric( + GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR g, GM_ADDR beta, + GM_ADDR aLog, GM_ADDR dtBias, GM_ADDR initialState, + GM_ADDR cuSeqlens, GM_ADDR chunkIndices, GM_ADDR attnOut, + GM_ADDR finalState, GM_ADDR gk, GM_ADDR aqk, GM_ADDR akk, + GM_ADDR w, GM_ADDR u, GM_ADDR qg, GM_ADDR kg, GM_ADDR vNew, GM_ADDR h, + GM_ADDR userWorkspace, const TilingData &tiling) +{ + const auto addresses = ResolveAddresses( + finalState, gk, w, u, qg, kg, vNew, h, userWorkspace, tiling); +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + TPipe pipe; + RunFrontEnd( + q, k, v, g, beta, aLog, dtBias, initialState, cuSeqlens, + chunkIndices, aqk, akk, addresses, userWorkspace, tiling, pipe); + if (!tiling.isVarLen && tiling.seqlen % tiling.chunkSize == 0) { + pipe.Destroy(); + RunGenericBackEnd( + q, k, v, beta, initialState, cuSeqlens, chunkIndices, aqk, + attnOut, addresses, userWorkspace, tiling); + } else { + RunGenericBackEnd( + q, k, v, beta, initialState, cuSeqlens, chunkIndices, aqk, + attnOut, addresses, userWorkspace, tiling, pipe); + } +#else + { + TPipe pipe; + RunFrontEnd( + q, k, v, g, beta, aLog, dtBias, initialState, cuSeqlens, + chunkIndices, aqk, akk, addresses, userWorkspace, tiling, pipe); + } + RunGenericBackEnd( + q, k, v, beta, initialState, cuSeqlens, chunkIndices, aqk, + attnOut, addresses, userWorkspace, tiling); +#endif +} + +} // namespace KdaForward diff --git a/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_finalize.h b/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_finalize.h new file mode 100644 index 0000000000..bc99f96057 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_finalize.h @@ -0,0 +1,810 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#pragma once + +#ifndef CATLASS_ARCH +#define CATLASS_ARCH 2201 +#endif + +#include "catlass/arch/kda_arch.hpp" +#include "catlass/arch/kda_cross_core_sync.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/kda_catlass.hpp" +#include "catlass/gemm/block/kda_block_mmad.hpp" +#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" +#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "kernel_utils/block/block_mmad_pingpong_tla_multi.hpp" +#include "catlass/layout/kda_layout.hpp" +#include "kernel_operator.h" +#include "chunk_kda_fwd_varlen.h" +#include "tla/kda_tla_layout.hpp" +#include "tla/kda_tensor.hpp" + +using namespace AscendC; + +namespace KdaFinalize { +namespace { +using KdaInt64 = tla::Int<64>; +using KdaInt128 = tla::Int<128>; +constexpr float LN2 = 0.69314718055994530942f; +constexpr float KDA_EXP2_CLAMP = 80.0f; +constexpr float KDA_EXP_INPUT_MAX = KDA_EXP2_CLAMP * LN2; +constexpr float KDA_EXP_INPUT_MIN = -KDA_EXP2_CLAMP * LN2; +constexpr float KDA_FP16_MAX = 65504.0f; +constexpr uint32_t EXP2_UB_ELEMENTS = 256; +constexpr uint32_t EXP2_UB_BYTES = EXP2_UB_ELEMENTS * (sizeof(float) + sizeof(uint16_t)); +constexpr uint32_t EXP2_EVENT_ID = 0; +constexpr uint32_t KDA_SOLVE_BT = 64; +constexpr uint32_t KDA_SOLVE_MATRIX_ELEMENTS = KDA_SOLVE_BT * KDA_SOLVE_BT; +constexpr uint32_t KDA_SOLVE_SCRATCH_X = 0; +constexpr uint32_t KDA_SOLVE_SCRATCH_Y0 = 1; +constexpr uint32_t KDA_SOLVE_SCRATCH_TMP = 2; +constexpr uint32_t KDA_SOLVE_SCRATCH_Y1 = 3; +constexpr uint32_t KDA_SOLVE_SCRATCH_IDENTITY = 4; +constexpr uint32_t KDA_SOLVE_SCRATCH_SLOTS = 5; +constexpr uint32_t KDA_SOLVE_DIAG_BT = 16; +constexpr uint32_t KDA_SOLVE_DIAG_BLOCKS = KDA_SOLVE_BT / KDA_SOLVE_DIAG_BT; +constexpr uint32_t KDA_SOLVE_DIAG_MCH_ITERS = 3; +constexpr uint32_t KDA_SCORE_REF_BC = 16; +constexpr uint32_t KDA_VEC_ARENA_ELEMENTS = 32768; +constexpr uint32_t KDA_BITS_PER_MASK_BYTE = 8; +constexpr uint32_t KDA_SELECT_COL_BLOCKS = 2; +constexpr uint32_t KDA_SELECT_COL_MASK_BYTES = KDA_SOLVE_MATRIX_ELEMENTS / KDA_BITS_PER_MASK_BYTE; +constexpr uint32_t KDA_SELECT_MASK_BYTES = KDA_SELECT_COL_BLOCKS * KDA_SELECT_COL_MASK_BYTES; +constexpr uint32_t KDA_SELECT_AQK_MASK_BYTE_OFFSET = 120 * 1024; +constexpr uint32_t KDA_SELECT_AKK_MASK_BYTE_OFFSET = KDA_SELECT_AQK_MASK_BYTE_OFFSET + KDA_SELECT_MASK_BYTES; +constexpr uint32_t KDA_SELECT_ZERO_BYTE_OFFSET = KDA_SELECT_AKK_MASK_BYTE_OFFSET + KDA_SELECT_MASK_BYTES; +constexpr uint32_t KDA_SELECT_ZERO_FLOAT_OFFSET = KDA_SELECT_ZERO_BYTE_OFFSET / sizeof(float); +constexpr uint8_t KDA_SCORE_DONE_FLAG0 = 2; +constexpr uint8_t KDA_SCORE_DONE_FLAG1 = 3; +constexpr uint8_t KDA_SCORE_READY_FLAG0 = 4; +constexpr uint8_t KDA_SCORE_READY_FLAG1 = 5; +constexpr uint32_t KDA_SCORE_QUEUE_DEPTH = 2; +constexpr uint32_t KDA_SYNC_REVERSE_DEPTH = 1; +constexpr uint32_t KDA_SCORE_SCRATCH_PLANES = 3; +constexpr uint32_t KDA_SCORE_SCRATCH_QG = 0; +constexpr uint32_t KDA_SCORE_SCRATCH_W = 1; +constexpr uint32_t KDA_SCORE_SCRATCH_KG = 2; +constexpr uint64_t KDA_WORKSPACE_ALIGN = 512; +constexpr uint32_t KDA_GATE_TILE_ROWS = 32; + +using KdaArchTag = Catlass::Arch::AtlasA2; +using KdaDispatchPolicy = Catlass::Gemm::MmadPingpong; +using KdaScoreDispatchPolicy = Catlass::Gemm::MmadPingpongTlaMulti; +static_assert(KdaScoreDispatchPolicy::ENABLE_L1_RESIDENT, + "KDA Aqk/Akk score MMAD must keep the shared right matrix resident in L1"); +static_assert(KdaScoreDispatchPolicy::L1B_STAGES == 1, + "KDA Aqk/Akk score MMAD needs one L1 B slot so the second MMAD reuses it"); +using KdaSolveDispatchPolicy = Catlass::Gemm::MmadPingpong; +static_assert(!KdaSolveDispatchPolicy::USE_HF32_MODE, "KDA triangular solve must use IEEE FP32 Cube mode"); +using KdaL1TileShape = tla::Shape; +using KdaL0TileShape = KdaL1TileShape; +using KdaSolveL1TileShape = tla::Shape; +using KdaSolveL0TileShape = KdaSolveL1TileShape; + +__aicore__ inline uint32_t FloatToBits(float value) +{ + union Bits { + __aicore__ Bits() {} + float f; + uint32_t u; + } bits; + bits.f = value; + return bits.u; +} + +__aicore__ inline float BitsToFloat(uint32_t value) +{ + union Bits { + __aicore__ Bits() {} + uint32_t u; + float f; + } bits; + bits.u = value; + return bits.f; +} + +__aicore__ inline uint16_t Bf16ToBits(bfloat16_t value) +{ + union Bits { + __aicore__ Bits() {} + bfloat16_t f; + uint16_t u; + } bits; + bits.f = value; + return bits.u; +} + +__aicore__ inline bfloat16_t BitsToBf16(uint16_t value) +{ + union Bits { + __aicore__ Bits() {} + uint16_t u; + bfloat16_t f; + } bits; + bits.u = value; + return bits.f; +} + +template +__aicore__ inline T FloatToType(float value) +{ + if constexpr (IsSameType::value) { + uint32_t bits = FloatToBits(value); + uint32_t bias = 0x7FFFu + ((bits >> 16) & 1u); + return BitsToBf16(static_cast((bits + bias) >> 16)); + } + return static_cast(value); +} + +template +class ChunkKdaFwdFinalizeKernel { +public: + using OUT_T = float; + using AKK_T = float; + template + __aicore__ inline void Init(GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR gk, GM_ADDR beta, GM_ADDR initialState, + GM_ADDR cuSeqlens, GM_ADDR chunkIndices, GM_ADDR preparedQG, GM_ADDR preparedAqk, + GM_ADDR propagatedVNew, GM_ADDR propagatedH, GM_ADDR o, GM_ADDR finalState, GM_ADDR aqk, + GM_ADDR akk, GM_ADDR w, GM_ADDR u, GM_ADDR qg, GM_ADDR kg, GM_ADDR vNew, GM_ADDR h, + GM_ADDR workspace, const TilingData &tiling, TPipe *pipe, bool initVecBuffers = true) + { + pipe_ = pipe; + q_.SetGlobalBuffer((__gm__ T *)q); + k_.SetGlobalBuffer((__gm__ T *)k); + v_.SetGlobalBuffer((__gm__ T *)v); + gk_.SetGlobalBuffer((__gm__ GK_T *)gk); + beta_.SetGlobalBuffer((__gm__ BETA_T *)beta); + if (initialState != nullptr) { + initialState_.SetGlobalBuffer((__gm__ float *)initialState); + } + cuSeqlensAddr_ = reinterpret_cast<__gm__ int64_t *>(cuSeqlens); + if (preparedQG != nullptr) { + preparedQG_.SetGlobalBuffer((__gm__ T *)preparedQG); + } + if (preparedAqk != nullptr) { + preparedAqk_.SetGlobalBuffer((__gm__ T *)preparedAqk); + } + if (propagatedVNew != nullptr) { + propagatedVNew_.SetGlobalBuffer((__gm__ T *)propagatedVNew); + } + if (propagatedH != nullptr) { + propagatedH_.SetGlobalBuffer((__gm__ T *)propagatedH); + } + chunkIndicesAddr_ = reinterpret_cast<__gm__ int64_t *>(chunkIndices); + o_.SetGlobalBuffer((__gm__ OUT_T *)o); + finalState_.SetGlobalBuffer((__gm__ float *)finalState); + aqk_.SetGlobalBuffer((__gm__ float *)aqk); + akk_.SetGlobalBuffer((__gm__ AKK_T *)akk); + w_.SetGlobalBuffer((__gm__ T *)w); + u_.SetGlobalBuffer((__gm__ OUT_T *)u); + qg_.SetGlobalBuffer((__gm__ T *)qg); + kg_.SetGlobalBuffer((__gm__ T *)kg); + vNew_.SetGlobalBuffer((__gm__ T *)vNew); + h_.SetGlobalBuffer((__gm__ float *)h); + solveWorkspace_.SetGlobalBuffer((__gm__ float *)workspace); + + B_ = tiling.batch; + N_ = tiling.seqNum; + H_ = tiling.qHeadNum; + HV_ = tiling.vHeadNum; + T_ = tiling.seqlen; + K_ = tiling.kHeadDim; + V_ = tiling.vHeadDim; + BT_ = tiling.chunkSize; + NT_ = tiling.totalChunks; + scale_ = tiling.scale; + hasInitial_ = tiling.hasInitialState; + isVarLen_ = tiling.isVarLen; + usedCoreNum_ = tiling.outputUsedCoreNum; + const uint64_t outputElements = B_ * HV_ * T_ * V_; + o_.SetGlobalBuffer((__gm__ OUT_T *)workspace); + u_.SetGlobalBuffer((__gm__ OUT_T *)workspace + outputElements); + if ASCEND_IS_AIV { + uint64_t subBlockNum = static_cast(GetSubBlockNum()); + solveCoreIdx_ = subBlockNum == 0 ? 0 : static_cast(GetBlockIdx()) / subBlockNum; + } else { + solveCoreIdx_ = static_cast(GetBlockIdx()); + } + if (pipe_ != nullptr && initVecBuffers) { + pipe_->InitBuffer(exp2Buf_, EXP2_UB_BYTES); + pipe_->InitBuffer(vecBuf_, KDA_VEC_ARENA_ELEMENTS * sizeof(float)); + const uint64_t gateWritebackRows = ScoreVectorMaxRows(5 * sizeof(float) + 2 * sizeof(T) + sizeof(GK_T)); + pipe_->InitBuffer(gateWritebackBuf_, + static_cast(gateWritebackRows * K_ * (3 * sizeof(T) + sizeof(GK_T)))); + AllocVectorEvents(); + } + } + __aicore__ inline void ProcessAiv() + { + ProcessOutAiv(); + ReleaseVectorEvents(); + } + + __aicore__ inline void ProcessAic() { ProcessOutAic(); } + +private: + __aicore__ inline void AllocVectorEvents() + { + mte2ToVEvent_ = pipe_->AllocEventID(); + vToMte2Event_ = pipe_->AllocEventID(); + vToMte3Event_ = pipe_->AllocEventID(); + mte3ToVEvent_ = pipe_->AllocEventID(); + mte2ToMte3Event_ = pipe_->AllocEventID(); + mte3ToMte2Event_ = pipe_->AllocEventID(); + vectorEventsAllocated_ = true; + } + + __aicore__ inline void ReleaseVectorEvents() + { + if (!vectorEventsAllocated_) { + return; + } + pipe_->ReleaseEventID(mte2ToVEvent_); + pipe_->ReleaseEventID(vToMte2Event_); + pipe_->ReleaseEventID(vToMte3Event_); + pipe_->ReleaseEventID(mte3ToVEvent_); + pipe_->ReleaseEventID(mte2ToMte3Event_); + pipe_->ReleaseEventID(mte3ToMte2Event_); + vectorEventsAllocated_ = false; + } + + __aicore__ inline uint64_t QOffset(uint64_t b, uint64_t h, uint64_t t, uint64_t d) const + { + return ((b * H_ + h) * T_ + t) * K_ + d; + } + + __aicore__ inline uint64_t KVOffset(uint64_t b, uint64_t hv, uint64_t t, uint64_t d, uint64_t dim) const + { + return ((b * HV_ + hv) * T_ + t) * dim + d; + } + + __aicore__ inline uint64_t OutputOffset(uint64_t b, uint64_t hv, uint64_t t, uint64_t d) const + { + return ((b * T_ + t) * HV_ + hv) * V_ + d; + } + + __aicore__ inline uint64_t BetaOffset(uint64_t b, uint64_t hv, uint64_t t) const { return (b * HV_ + hv) * T_ + t; } + + __aicore__ inline uint64_t AOffset(uint64_t b, uint64_t hv, uint64_t t, uint64_t j) const + { + return ((b * HV_ + hv) * T_ + t) * BT_ + j; + } + + __aicore__ inline uint64_t HOffset(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t d, uint64_t r) const + { + return (((b * HV_ + hv) * NT_ + chunkIdx) * K_ + d) * V_ + r; + } + + __aicore__ inline uint64_t WScratchOffset(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t t, uint64_t d) const + { + return (((b * HV_ + hv) * NT_ + chunkIdx) * BT_ + t) * K_ + d; + } + + __aicore__ inline uint64_t SolveScratchOffset(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t slot) const + { + (void)b; + (void)hv; + (void)chunkIdx; + uint64_t matrixElements = BT_ * BT_; + return solveCoreIdx_ * KDA_SOLVE_SCRATCH_SLOTS * matrixElements + slot * matrixElements; + } + + __aicore__ inline uint64_t ScoreScratchOffset(uint64_t slot, uint64_t plane, uint64_t t = 0, uint64_t d = 0) const + { + return (((solveCoreIdx_ * KDA_SCORE_QUEUE_DEPTH + slot) * KDA_SCORE_SCRATCH_PLANES + plane) * BT_ + t) * K_ + d; + } + + __aicore__ inline uint64_t ScoreRefBlockSize() const { return KDA_SCORE_REF_BC; } + + __aicore__ inline uint64_t ScoreRowBlockCount(uint64_t curT, uint64_t rowBegin) const + { + uint64_t blockSize = ScoreRefBlockSize(); + uint64_t rowCount = curT - rowBegin; + if (rowCount > blockSize) { + rowCount = blockSize; + } + return rowCount; + } + + __aicore__ inline uint64_t ScoreRefToken(uint64_t start, uint64_t curT, uint64_t rowBegin, uint64_t rowCount) const + { + uint64_t ref = rowBegin + rowCount / 2; + if (ref >= curT) { + ref = curT - 1; + } + return start + ref; + } + + __aicore__ inline void RunExp2(LocalTensor &tensor, uint32_t count) + { + SetFlag(EXP2_EVENT_ID); + WaitFlag(EXP2_EVENT_ID); + ClampExpInput(tensor, count); + Exp(tensor, tensor, count); + PipeBarrier(); + SetFlag(EXP2_EVENT_ID); + WaitFlag(EXP2_EVENT_ID); + } + + __aicore__ inline void ClampExpInput(LocalTensor &tensor, uint32_t count) + { + Mins(tensor, tensor, KDA_EXP_INPUT_MAX, count); + PipeBarrier(); + Maxs(tensor, tensor, KDA_EXP_INPUT_MIN, count); + PipeBarrier(); + } + + __aicore__ inline void ClampFp32ToOutputType(LocalTensor &tensor, uint32_t count) + { + if constexpr (IsSameType::value) { + Mins(tensor, tensor, KDA_FP16_MAX, count); + PipeBarrier(); + Maxs(tensor, tensor, -KDA_FP16_MAX, count); + PipeBarrier(); + } + } + + template + __aicore__ inline void CopyVectorIn(LocalTensor &dst, GlobalTensor &src, uint64_t offset, + uint64_t count) + { + uint64_t rowBytes = count * static_cast(sizeof(CopyT)); + if (rowBytes >= 32 && rowBytes % 32 == 0) { + DataCopy(dst, src[offset], static_cast(count)); + return; + } + DataCopyParams params{1, static_cast(rowBytes), 0, 0}; + DataCopyPadParams padParams{false, 0, 0, 0}; + DataCopyPad(dst, src[offset], params, padParams); + } + + template + __aicore__ inline void CopyVectorOut(GlobalTensor &dst, uint64_t offset, LocalTensor &src, + uint64_t count) + { + uint64_t rowBytes = count * static_cast(sizeof(CopyT)); + if (rowBytes >= 32 && rowBytes % 32 == 0) { + DataCopy(dst[offset], src, static_cast(count)); + return; + } + DataCopyParams params{1, static_cast(rowBytes), 0, 0}; + DataCopyPad(dst[offset], src, params); + } + + template + __aicore__ inline void CopyRowsOut(GlobalTensor &dst, uint64_t offset, LocalTensor &src, + uint64_t rows, uint64_t cols, uint64_t dstStride) + { + if (cols == dstStride) { + CopyVectorOut(dst, offset, src, rows * cols); + return; + } + constexpr uint64_t blockBytes = 32; + const uint64_t rowBytes = cols * sizeof(CopyT); + const uint64_t gapBytes = (dstStride - cols) * sizeof(CopyT); + DataCopyParams params{static_cast(rows), static_cast(rowBytes / blockBytes), 0, + static_cast(gapBytes / blockBytes)}; + DataCopy(dst[offset], src, params); + } + + template + __aicore__ inline void CopyRowIn(LocalTensor &dst, GlobalTensor &src, uint64_t offset) + { + CopyVectorIn(dst, src, offset, K_); + } + + template + __aicore__ inline void CopyRowOut(GlobalTensor &dst, uint64_t offset, LocalTensor &src) + { + CopyVectorOut(dst, offset, src, K_); + } + + __aicore__ inline LocalTensor VecScratch(uint64_t slot) + { + return vecBuf_.Get()[slot * EXP2_UB_ELEMENTS]; + } + + template + __aicore__ inline void LoadAsFloatRow(GlobalTensor &src, uint64_t srcOffset, LocalTensor &dst, + uint64_t count) + { + if constexpr (IsSameType::value) { + CopyVectorIn(dst, src, srcOffset, count); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + Adds(dst, dst, 0.0f, static_cast(count)); + PipeBarrier(); + SetFlag(vToMte2Event_); + WaitFlag(vToMte2Event_); + } else { + constexpr uint32_t typedOffset = EXP2_UB_ELEMENTS * sizeof(float) / sizeof(CopyT); + LocalTensor rowLocal = exp2Buf_.Get()[typedOffset]; + CopyVectorIn(rowLocal, src, srcOffset, count); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + Cast(dst, rowLocal, RoundMode::CAST_NONE, static_cast(count)); + PipeBarrier(); + SetFlag(vToMte2Event_); + WaitFlag(vToMte2Event_); + } + PipeBarrier(); + } + + template + __aicore__ inline void LoadAsFloatVector(GlobalTensor &src, uint64_t srcOffset, LocalTensor &dst, + LocalTensor &typedScratch, uint64_t count) + { + if constexpr (IsSameType::value) { + CopyVectorIn(dst, src, srcOffset, count); + } else { + CopyVectorIn(typedScratch, src, srcOffset, count); + } + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + if constexpr (!IsSameType::value) { + Cast(dst, typedScratch, RoundMode::CAST_NONE, static_cast(count)); + PipeBarrier(); + } + } + + template + __aicore__ inline void StoreFloatRow(GlobalTensor &dst, uint64_t dstOffset, LocalTensor &src, + uint64_t count) + { + if constexpr (IsSameType::value) { + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + CopyVectorOut(dst, dstOffset, src, count); + } else { + constexpr uint32_t typedOffset = EXP2_UB_ELEMENTS * sizeof(float) / sizeof(CopyT); + LocalTensor rowLocal = exp2Buf_.Get()[typedOffset]; + Cast(rowLocal, src, RoundMode::CAST_RINT, static_cast(count)); + PipeBarrier(); + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + CopyVectorOut(dst, dstOffset, rowLocal, count); + } + SetFlag(mte3ToMte2Event_); + WaitFlag(mte3ToMte2Event_); + SetFlag(mte3ToVEvent_); + WaitFlag(mte3ToVEvent_); + } + + __aicore__ inline LocalTensor Exp2NegG(uint64_t b, uint64_t hv, uint64_t t) + { + LocalTensor exp2Local = exp2Buf_.Get(); + LoadAsFloatRow(gk_, KVOffset(b, hv, t, 0, K_), exp2Local, K_); + Muls(exp2Local, exp2Local, -LN2, static_cast(K_)); + PipeBarrier(); + RunExp2(exp2Local, static_cast(K_)); + return exp2Local; + } + + __aicore__ inline uint64_t ScoreVectorMaxRows(uint64_t bytesPerElem) const + { + constexpr uint64_t arenaBytes = static_cast(KDA_VEC_ARENA_ELEMENTS) * sizeof(float); + uint64_t maxRows = (arenaBytes / bytesPerElem) / K_; + if (K_ >= 128 && maxRows > 32) { + maxRows = 32; + } + return maxRows; + } + + __aicore__ inline void ComputeOutputCube(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, uint64_t curT) + { + using ElementA = T; + using ElementB = T; + using ElementC = OUT_T; + using LayoutTagA = Catlass::layout::RowMajor; + using LayoutTagB = Catlass::layout::RowMajor; + using LayoutTagC = Catlass::layout::RowMajor; + using TileCopy = Catlass::Gemm::Tile::PackedTileCopyTla; + using BlockMmad = Catlass::Gemm::Block::BlockMmadTla; + + Catlass::Arch::Resource resource; + BlockMmad blockMmad(resource); + + auto layoutQ = tla::MakeLayout(BT_, K_); + auto layoutH = tla::MakeLayout(K_, V_); + auto layoutO = tla::MakeLayout(BT_, V_); + for (uint64_t nOffset = 0; nOffset < V_; nOffset += 128) { + uint32_t curN = static_cast((V_ - nOffset) > 128 ? 128 : (V_ - nOffset)); + auto tensorH = tla::MakeTensor(propagatedH_[HOffset(b, hv, chunkIdx, 0, nOffset)], layoutH, + Catlass::Arch::PositionGM{}); + for (uint64_t mOffset = 0; mOffset < curT; mOffset += 64) { + uint32_t curM = static_cast((curT - mOffset) > 64 ? 64 : (curT - mOffset)); + Catlass::GemmCoord shapeQH{curM, curN, static_cast(K_)}; + auto tensorQ = tla::MakeTensor(preparedQG_[KVOffset(b, hv, start + mOffset, 0, K_)], layoutQ, + Catlass::Arch::PositionGM{}); + auto tensorO = tla::MakeTensor(o_[KVOffset(b, hv, start + mOffset, nOffset, V_)], layoutO, + Catlass::Arch::PositionGM{}); + auto blockQ = GetTile(tensorQ, tla::MakeCoord(0, 0), tla::MakeShape(shapeQH.m(), shapeQH.k())); + auto blockH = GetTile(tensorH, tla::MakeCoord(0, 0), tla::MakeShape(shapeQH.k(), shapeQH.n())); + auto blockO = GetTile(tensorO, tla::MakeCoord(0, 0), tla::MakeShape(shapeQH.m(), shapeQH.n())); + blockMmad(blockQ, blockH, blockO, shapeQH); + PipeBarrier(); + } + } + + auto layoutAqk = tla::MakeLayout(BT_, BT_); + auto layoutV = tla::MakeLayout(BT_, V_); + for (uint64_t nOffset = 0; nOffset < V_; nOffset += 128) { + uint32_t curN = static_cast((V_ - nOffset) > 128 ? 128 : (V_ - nOffset)); + auto tensorVNew = tla::MakeTensor(propagatedVNew_[KVOffset(b, hv, start, nOffset, V_)], layoutV, + Catlass::Arch::PositionGM{}); + for (uint64_t mOffset = 0; mOffset < curT; mOffset += 64) { + uint32_t curM = static_cast((curT - mOffset) > 64 ? 64 : (curT - mOffset)); + Catlass::GemmCoord shapeAV{curM, curN, static_cast(curT)}; + auto tensorAqk = tla::MakeTensor(preparedAqk_[AOffset(b, hv, start + mOffset, 0)], layoutAqk, + Catlass::Arch::PositionGM{}); + auto tensorLocal = tla::MakeTensor(u_[KVOffset(b, hv, start + mOffset, nOffset, V_)], layoutO, + Catlass::Arch::PositionGM{}); + auto blockAqk = GetTile(tensorAqk, tla::MakeCoord(0, 0), tla::MakeShape(shapeAV.m(), shapeAV.k())); + auto blockVNew = GetTile(tensorVNew, tla::MakeCoord(0, 0), tla::MakeShape(shapeAV.k(), shapeAV.n())); + auto blockLocal = GetTile(tensorLocal, tla::MakeCoord(0, 0), tla::MakeShape(shapeAV.m(), shapeAV.n())); + blockMmad(blockAqk, blockVNew, blockLocal, shapeAV); + PipeBarrier(); + } + } + } + + __aicore__ inline void FinalizeOutputRows(uint64_t b, uint64_t hv, uint64_t start, uint64_t curT, + uint64_t subBlockIdx, uint64_t subBlockNum) + { + if (subBlockNum == 0 || subBlockIdx >= subBlockNum || V_ == 0) { + return; + } + const uint64_t rowBegin = (curT * subBlockIdx) / subBlockNum; + const uint64_t rowEnd = (curT * (subBlockIdx + 1)) / subBlockNum; + const uint64_t gateWritebackRows = ScoreVectorMaxRows(5 * sizeof(float) + 2 * sizeof(T) + sizeof(GK_T)); + const uint64_t gateWritebackBytes = gateWritebackRows * K_ * (3 * sizeof(T) + sizeof(GK_T)); + uint64_t maxRows = KDA_VEC_ARENA_ELEMENTS / (3 * V_); + const uint64_t typedMaxRows = gateWritebackBytes / (V_ * sizeof(T)); + if (maxRows > typedMaxRows) { + maxRows = typedMaxRows; + } + if (maxRows == 0) { + return; + } + + for (uint64_t tileRow = rowBegin; tileRow < rowEnd; tileRow += maxRows) { + uint64_t tileRows = rowEnd - tileRow; + if (tileRows > maxRows) { + tileRows = maxRows; + } + const uint64_t elems = tileRows * V_; + const uint64_t ti = start + tileRow; + LocalTensor arena = vecBuf_.Get(); + LocalTensor stateLocal = arena; + LocalTensor localLocal = arena[elems]; + LocalTensor outLocal = arena[2 * elems]; + LocalTensor outTyped = gateWritebackBuf_.Get(); + + CopyVectorIn(stateLocal, o_, KVOffset(b, hv, ti, 0, V_), elems); + CopyVectorIn(localLocal, u_, KVOffset(b, hv, ti, 0, V_), elems); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + Muls(stateLocal, stateLocal, scale_, static_cast(elems)); + PipeBarrier(); + Add(outLocal, stateLocal, localLocal, static_cast(elems)); + PipeBarrier(); + ClampFp32ToOutputType(outLocal, static_cast(elems)); + Cast(outTyped, outLocal, RoundMode::CAST_RINT, static_cast(elems)); + PipeBarrier(); + + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + CopyRowsOut(vNew_, OutputOffset(b, hv, ti, 0), outTyped, tileRows, V_, HV_ * V_); + SetFlag(mte3ToMte2Event_); + WaitFlag(mte3ToMte2Event_); + SetFlag(mte3ToVEvent_); + WaitFlag(mte3ToVEvent_); + } + } + __aicore__ inline bool ResolveFlatChunk(uint64_t task, uint64_t &seq, uint64_t &b, uint64_t &h, uint64_t &hv, + uint64_t &chunkIdx, uint64_t &start, uint64_t &end) + { + hv = task % HV_; + uint64_t flatChunk = task / HV_; + if (!isVarLen_) { + seq = flatChunk / NT_; + b = seq; + chunkIdx = flatChunk % NT_; + start = chunkIdx * BT_; + end = start + BT_; + if (end > T_) { + end = T_; + } + } else { + if (!KdaVarlen::ResolveChunkRange(cuSeqlensAddr_, chunkIndicesAddr_, N_, T_, BT_, flatChunk, seq, start, + end)) { + return false; + } + b = 0; + chunkIdx = flatChunk; + } + h = hv / (HV_ / H_); + return start < end; + } + + __aicore__ inline void ProcessChunkOutAiv(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, uint64_t end, + uint64_t subBlockIdx, uint64_t subBlockNum) + { + uint64_t curT = end - start; + if (curT == 0) { + return; + } + if constexpr (IsSameType::value) { + return; + } + Catlass::Arch::CrossCoreWaitFlagWithReverse<0x2, PIPE_MTE2>(syncDoneFlag_); + FinalizeOutputRows(b, hv, start, curT, subBlockIdx, subBlockNum); + } + + __aicore__ inline void ProcessChunkOutAic(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, uint64_t end) + { + uint64_t curT = end - start; + if (curT == 0) { + return; + } + ComputeOutputCube(b, hv, chunkIdx, start, curT); + Catlass::Arch::CrossCoreSetFlagWithReverse<0x2, PIPE_FIX>(syncDoneFlag_); + } + + __aicore__ inline void ProcessOutAiv() + { + if constexpr (IsSameType::value) { + return; + } + uint64_t subBlockNum = static_cast(GetSubBlockNum()); + if (subBlockNum == 0) { + return; + } + uint64_t subBlockIdx = static_cast(GetSubBlockIdx()); + uint64_t coreNum = usedCoreNum_ == 0 ? 1 : usedCoreNum_; + uint64_t coreIdx = static_cast(GetBlockIdx()) / subBlockNum; + uint64_t taskNum = static_cast((isVarLen_ ? NT_ : B_ * NT_) * HV_); + for (uint64_t task = coreIdx; task < taskNum; task += coreNum) { + uint64_t seq = 0; + uint64_t b = 0; + uint64_t h = 0; + uint64_t hv = 0; + uint64_t chunkIdx = 0; + uint64_t start = 0; + uint64_t end = 0; + if (ResolveFlatChunk(task, seq, b, h, hv, chunkIdx, start, end)) { + (void)seq; + (void)h; + (void)chunkIdx; + ProcessChunkOutAiv(b, hv, chunkIdx, start, end, subBlockIdx, subBlockNum); + } + } + } + + __aicore__ inline void ProcessOutAic() + { + if constexpr (IsSameType::value) { + return; + } + uint64_t taskNum = static_cast((isVarLen_ ? NT_ : B_ * NT_) * HV_); + uint64_t coreNum = usedCoreNum_ == 0 ? 1 : usedCoreNum_; + for (uint64_t task = GetBlockIdx(); task < taskNum; task += coreNum) { + uint64_t seq = 0; + uint64_t b = 0; + uint64_t h = 0; + uint64_t hv = 0; + uint64_t chunkIdx = 0; + uint64_t start = 0; + uint64_t end = 0; + if (ResolveFlatChunk(task, seq, b, h, hv, chunkIdx, start, end)) { + (void)seq; + (void)h; + ProcessChunkOutAic(b, hv, chunkIdx, start, end); + } + } + } + +private: + GlobalTensor q_; + GlobalTensor k_; + GlobalTensor v_; + GlobalTensor gk_; + GlobalTensor beta_; + GlobalTensor initialState_; + GlobalTensor o_; + GlobalTensor finalState_; + GlobalTensor aqk_; + GlobalTensor akk_; + GlobalTensor w_; + GlobalTensor u_; + GlobalTensor qg_; + GlobalTensor kg_; + GlobalTensor vNew_; + GlobalTensor h_; + GlobalTensor preparedQG_; + GlobalTensor preparedAqk_; + GlobalTensor propagatedVNew_; + GlobalTensor propagatedH_; + GlobalTensor solveWorkspace_; + GlobalTensor scoreWorkspace_; + TPipe *pipe_ = nullptr; + TBuf exp2Buf_; + TBuf vecBuf_; + TBuf gateWritebackBuf_; + TEventID mte2ToVEvent_ = 0; + TEventID vToMte2Event_ = 0; + TEventID vToMte3Event_ = 0; + TEventID mte3ToVEvent_ = 0; + TEventID mte2ToMte3Event_ = 0; + TEventID mte3ToMte2Event_ = 0; + bool vectorEventsAllocated_ = false; + Catlass::Arch::CrossCoreFlagWithReverse scoreReadyFlag_{KDA_SCORE_READY_FLAG0, + KDA_SCORE_READY_FLAG1}; + Catlass::Arch::CrossCoreFlagWithReverse scoreDoneFlag_{KDA_SCORE_DONE_FLAG0, + KDA_SCORE_DONE_FLAG1}; + // Score production is fully drained before solve starts, so the solve handshake can safely reuse + // the existing score flags without consuming additional hardware flag IDs. + Catlass::Arch::CrossCoreFlagWithReverse syncReadyFlag_{KDA_SCORE_READY_FLAG0, + KDA_SCORE_READY_FLAG1}; + Catlass::Arch::CrossCoreFlagWithReverse syncDoneFlag_{KDA_SCORE_DONE_FLAG0, + KDA_SCORE_DONE_FLAG1}; + uint64_t B_ = 0; + uint64_t N_ = 0; + uint64_t H_ = 0; + uint64_t HV_ = 0; + uint64_t T_ = 0; + uint64_t K_ = 0; + uint64_t V_ = 0; + uint64_t BT_ = 0; + uint64_t NT_ = 0; + float scale_ = 1.0f; + bool hasInitial_ = false; + bool isVarLen_ = false; + bool isAivOnly_ = false; + uint64_t usedCoreNum_ = 1; + uint64_t solveCoreIdx_ = 0; + __gm__ int64_t *chunkIndicesAddr_ = nullptr; + __gm__ int64_t *cuSeqlensAddr_ = nullptr; +}; +} // namespace + +template +__aicore__ inline void RunChunkKdaOutput(GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR gk, GM_ADDR beta, + GM_ADDR initialState, GM_ADDR cuSeqlens, GM_ADDR chunkIndices, + GM_ADDR qgScaled, GM_ADDR aqk, GM_ADDR propagatedVNew, GM_ADDR propagatedH, + GM_ADDR o, GM_ADDR userWorkspace, const TilingData &tiling, TPipe &pipe) +{ + GM_ADDR outputScratch = userWorkspace + tiling.outputScratchOffset; + uint64_t outputElements = static_cast(tiling.batch) * static_cast(tiling.vHeadNum) * + static_cast(tiling.seqlen) * static_cast(tiling.vHeadDim); + GM_ADDR stateScratch = outputScratch; + GM_ADDR localScratch = outputScratch + outputElements * sizeof(float); + if ASCEND_IS_AIC { + ChunkKdaFwdFinalizeKernel op; + op.Init(q, k, v, gk, beta, initialState, cuSeqlens, chunkIndices, qgScaled, aqk, propagatedVNew, propagatedH, + stateScratch, userWorkspace, aqk, userWorkspace, userWorkspace, localScratch, userWorkspace, + userWorkspace, o, propagatedH, outputScratch, tiling, &pipe, false); + op.ProcessAic(); + } + if ASCEND_IS_AIV { + ChunkKdaFwdFinalizeKernel op; + op.Init(q, k, v, gk, beta, initialState, cuSeqlens, chunkIndices, qgScaled, aqk, propagatedVNew, propagatedH, + stateScratch, userWorkspace, aqk, userWorkspace, userWorkspace, localScratch, userWorkspace, + userWorkspace, o, propagatedH, outputScratch, tiling, &pipe); + op.ProcessAiv(); + } +} + +} // namespace KdaFinalize diff --git a/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_h_runtime_tiling.h b/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_h_runtime_tiling.h new file mode 100644 index 0000000000..9ef1bd15fe --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_h_runtime_tiling.h @@ -0,0 +1,42 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CHUNK_KDA_FWD_H_RUNTIME_TILING_H +#define CHUNK_KDA_FWD_H_RUNTIME_TILING_H + +#include + +struct ChunkKdaFwdHRuntimeTiling { + int64_t batch; + int64_t seqlen; + int64_t kNumHead; + int64_t vNumHead; + int64_t kHeadDim; + int64_t vHeadDim; + int64_t chunkSize; + bool useInitialState; + bool storeFinalState; + int64_t dataType; + int64_t gDataType; + int64_t stateDataType; + int64_t isVariedLen; + int64_t shapeBatch; + int64_t tokenBatch; + bool useG; + bool useGk; + int64_t vWorkspaceOffset; + int64_t vUpdateWorkspaceOffset; + int64_t kDecayWorkspaceOffset; + int64_t hWorkspaceOffset; + int64_t numSeqWorkspaceOffset; + int64_t numChunksWorkspaceOffset; +}; + +#endif // CHUNK_KDA_FWD_H_RUNTIME_TILING_H diff --git a/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_post_wu.h b/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_post_wu.h new file mode 100644 index 0000000000..b3942a08d4 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_post_wu.h @@ -0,0 +1,926 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#pragma once + +#ifndef CATLASS_ARCH +#define CATLASS_ARCH 2201 +#endif + +#include "catlass/arch/kda_arch.hpp" +#include "catlass/arch/kda_cross_core_sync.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/kda_catlass.hpp" +#include "catlass/gemm/block/kda_block_mmad.hpp" +#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" +#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" +#include "catlass/gemm/tile/kda_tile_mmad.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "kernel_utils/block/block_mmad_pingpong_tla_multi.hpp" +#include "catlass/layout/kda_layout.hpp" +#include "kernel_operator.h" +#include "chunk_kda_fwd_varlen.h" +#include "tla/kda_tla_layout.hpp" +#include "tla/kda_tensor.hpp" + +using namespace AscendC; + +namespace KdaPostWu { +namespace { +using KdaInt64 = tla::Int<64>; +using KdaInt128 = tla::Int<128>; +constexpr float LN2 = 0.69314718055994530942f; +constexpr float KDA_EXP2_CLAMP = 80.0f; +constexpr float KDA_EXP_INPUT_MAX = KDA_EXP2_CLAMP * LN2; +constexpr float KDA_EXP_INPUT_MIN = -KDA_EXP2_CLAMP * LN2; +constexpr float KDA_FP16_MAX = 65504.0f; +constexpr uint32_t EXP2_UB_ELEMENTS = 256; +constexpr uint32_t EXP2_UB_BYTES = EXP2_UB_ELEMENTS * (sizeof(float) + sizeof(uint16_t)); +constexpr uint32_t EXP2_EVENT_ID = 0; +constexpr uint32_t KDA_SOLVE_BT = 64; +constexpr uint32_t KDA_SOLVE_MATRIX_ELEMENTS = KDA_SOLVE_BT * KDA_SOLVE_BT; +constexpr uint32_t KDA_SOLVE_SCRATCH_X = 0; +constexpr uint32_t KDA_SOLVE_SCRATCH_Y0 = 1; +constexpr uint32_t KDA_SOLVE_SCRATCH_TMP = 2; +constexpr uint32_t KDA_SOLVE_SCRATCH_Y1 = 3; +constexpr uint32_t KDA_SOLVE_SCRATCH_IDENTITY = 4; +constexpr uint32_t KDA_SOLVE_SCRATCH_SLOTS = 5; +constexpr uint32_t KDA_SOLVE_DIAG_BT = 16; +constexpr uint32_t KDA_SOLVE_DIAG_BLOCKS = KDA_SOLVE_BT / KDA_SOLVE_DIAG_BT; +constexpr uint32_t KDA_SOLVE_DIAG_MCH_ITERS = 3; +constexpr uint32_t KDA_SCORE_REF_BC = 16; +constexpr uint32_t KDA_VEC_ARENA_ELEMENTS = 32768; +constexpr uint32_t KDA_BITS_PER_MASK_BYTE = 8; +constexpr uint32_t KDA_SELECT_COL_BLOCKS = 2; +constexpr uint32_t KDA_SELECT_COL_MASK_BYTES = KDA_SOLVE_MATRIX_ELEMENTS / KDA_BITS_PER_MASK_BYTE; +constexpr uint32_t KDA_SELECT_MASK_BYTES = KDA_SELECT_COL_BLOCKS * KDA_SELECT_COL_MASK_BYTES; +constexpr uint32_t KDA_SELECT_AQK_MASK_BYTE_OFFSET = 120 * 1024; +constexpr uint32_t KDA_SELECT_AKK_MASK_BYTE_OFFSET = KDA_SELECT_AQK_MASK_BYTE_OFFSET + KDA_SELECT_MASK_BYTES; +constexpr uint32_t KDA_SELECT_ZERO_BYTE_OFFSET = KDA_SELECT_AKK_MASK_BYTE_OFFSET + KDA_SELECT_MASK_BYTES; +constexpr uint32_t KDA_SELECT_ZERO_FLOAT_OFFSET = KDA_SELECT_ZERO_BYTE_OFFSET / sizeof(float); +constexpr uint8_t KDA_SCORE_DONE_FLAG0 = 2; +constexpr uint8_t KDA_SCORE_DONE_FLAG1 = 3; +constexpr uint8_t KDA_SCORE_READY_FLAG0 = 4; +constexpr uint8_t KDA_SCORE_READY_FLAG1 = 5; +constexpr uint32_t KDA_SCORE_QUEUE_DEPTH = 2; +constexpr uint32_t KDA_SYNC_REVERSE_DEPTH = 1; +constexpr uint32_t KDA_SCORE_SCRATCH_PLANES = 3; +constexpr uint32_t KDA_SCORE_SCRATCH_QG = 0; +constexpr uint32_t KDA_SCORE_SCRATCH_W = 1; +constexpr uint32_t KDA_SCORE_SCRATCH_KG = 2; +constexpr uint64_t KDA_WORKSPACE_ALIGN = 512; +constexpr uint32_t KDA_GATE_TILE_ROWS = 32; +constexpr uint32_t KDA_POST_EVENT = 3; +constexpr uint32_t KDA_POST_EVENT_NEXT = 4; +constexpr uint32_t KDA_POST_EVENT_FIX = 5; + +using KdaArchTag = Catlass::Arch::AtlasA2; +using KdaDispatchPolicy = Catlass::Gemm::MmadPingpong; +using KdaScoreDispatchPolicy = Catlass::Gemm::MmadPingpongTlaMulti; +static_assert(KdaScoreDispatchPolicy::ENABLE_L1_RESIDENT, + "KDA Aqk/Akk score MMAD must keep the shared right matrix resident in L1"); +static_assert(KdaScoreDispatchPolicy::L1B_STAGES == 1, + "KDA Aqk/Akk score MMAD needs one L1 B slot so the second MMAD reuses it"); +using KdaSolveDispatchPolicy = Catlass::Gemm::MmadPingpong; +static_assert(!KdaSolveDispatchPolicy::USE_HF32_MODE, "KDA triangular solve must use IEEE FP32 Cube mode"); +using KdaL1TileShape = tla::Shape; +using KdaL0TileShape = KdaL1TileShape; +using KdaSolveL1TileShape = tla::Shape; +using KdaSolveL0TileShape = KdaSolveL1TileShape; + +__aicore__ inline uint32_t FloatToBits(float value) +{ + union Bits { + __aicore__ Bits() {} + float f; + uint32_t u; + } bits; + bits.f = value; + return bits.u; +} + +__aicore__ inline float BitsToFloat(uint32_t value) +{ + union Bits { + __aicore__ Bits() {} + uint32_t u; + float f; + } bits; + bits.u = value; + return bits.f; +} + +__aicore__ inline uint16_t Bf16ToBits(bfloat16_t value) +{ + union Bits { + __aicore__ Bits() {} + bfloat16_t f; + uint16_t u; + } bits; + bits.f = value; + return bits.u; +} + +__aicore__ inline bfloat16_t BitsToBf16(uint16_t value) +{ + union Bits { + __aicore__ Bits() {} + uint16_t u; + bfloat16_t f; + } bits; + bits.u = value; + return bits.f; +} + +template +__aicore__ inline T FloatToType(float value) +{ + if constexpr (IsSameType::value) { + uint32_t bits = FloatToBits(value); + uint32_t bias = 0x7FFFu + ((bits >> 16) & 1u); + return BitsToBf16(static_cast((bits + bias) >> 16)); + } + return static_cast(value); +} + +template +class ChunkKdaFwdPostWuKernel { +public: + using OUT_T = T; + using AKK_T = T; + template + __aicore__ inline void Init(GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR gk, GM_ADDR beta, GM_ADDR initialState, + GM_ADDR cuSeqlens, GM_ADDR chunkIndices, GM_ADDR preparedQG, GM_ADDR preparedAqk, + GM_ADDR propagatedVNew, GM_ADDR propagatedH, GM_ADDR o, GM_ADDR finalState, GM_ADDR aqk, + GM_ADDR akk, GM_ADDR w, GM_ADDR u, GM_ADDR qg, GM_ADDR kg, GM_ADDR vNew, GM_ADDR h, + GM_ADDR workspace, const TilingData &tiling, TPipe *pipe, bool initVecBuffers = true) + { + pipe_ = pipe; + q_.SetGlobalBuffer((__gm__ T *)q); + k_.SetGlobalBuffer((__gm__ T *)k); + v_.SetGlobalBuffer((__gm__ T *)v); + gk_.SetGlobalBuffer((__gm__ GK_T *)gk); + beta_.SetGlobalBuffer((__gm__ BETA_T *)beta); + if (initialState != nullptr) { + initialState_.SetGlobalBuffer((__gm__ float *)initialState); + } + cuSeqlensAddr_ = reinterpret_cast<__gm__ int64_t *>(cuSeqlens); + if (preparedQG != nullptr) { + preparedQG_.SetGlobalBuffer((__gm__ T *)preparedQG); + } + if (preparedAqk != nullptr) { + preparedAqk_.SetGlobalBuffer((__gm__ T *)preparedAqk); + } + if (propagatedVNew != nullptr) { + propagatedVNew_.SetGlobalBuffer((__gm__ T *)propagatedVNew); + } + if (propagatedH != nullptr) { + propagatedH_.SetGlobalBuffer((__gm__ T *)propagatedH); + } + chunkIndicesAddr_ = reinterpret_cast<__gm__ int64_t *>(chunkIndices); + o_.SetGlobalBuffer((__gm__ OUT_T *)o); + finalState_.SetGlobalBuffer((__gm__ float *)finalState); + aqk_.SetGlobalBuffer((__gm__ float *)aqk); + akk_.SetGlobalBuffer((__gm__ AKK_T *)akk); + w_.SetGlobalBuffer((__gm__ T *)w); + u_.SetGlobalBuffer((__gm__ OUT_T *)u); + qg_.SetGlobalBuffer((__gm__ T *)qg); + kg_.SetGlobalBuffer((__gm__ T *)kg); + vNew_.SetGlobalBuffer((__gm__ T *)vNew); + h_.SetGlobalBuffer((__gm__ float *)h); + solveWorkspace_.SetGlobalBuffer((__gm__ float *)workspace); + + B_ = tiling.batch; + N_ = tiling.seqNum; + H_ = tiling.qHeadNum; + HV_ = tiling.vHeadNum; + T_ = tiling.seqlen; + K_ = tiling.kHeadDim; + V_ = tiling.vHeadDim; + BT_ = tiling.chunkSize; + NT_ = tiling.totalChunks; + scale_ = tiling.scale; + hasInitial_ = tiling.hasInitialState; + isVarLen_ = tiling.isVarLen; + inputSequenceMajor_ = tiling.inputSequenceMajor; + usedCoreNum_ = tiling.postWuUsedCoreNum; + if ASCEND_IS_AIV { + uint64_t subBlockNum = static_cast(GetSubBlockNum()); + solveCoreIdx_ = subBlockNum == 0 ? 0 : static_cast(GetBlockIdx()) / subBlockNum; + } else { + solveCoreIdx_ = static_cast(GetBlockIdx()); + } + if (pipe_ != nullptr && initVecBuffers) { + pipe_->InitBuffer(exp2Buf_, EXP2_UB_BYTES); + pipe_->InitBuffer(vecBuf_, KDA_VEC_ARENA_ELEMENTS * sizeof(float)); + const uint64_t gateWritebackRows = ScoreVectorMaxRows(5 * sizeof(float) + 2 * sizeof(T) + sizeof(GK_T)); + pipe_->InitBuffer(gateWritebackBuf_, + static_cast(gateWritebackRows * K_ * (3 * sizeof(T) + sizeof(GK_T)))); + AllocVectorEvents(); + } + } + __aicore__ inline void ProcessAiv() + { + ProcessPostAiv(); + ReleaseVectorEvents(); + } + + __aicore__ inline void ProcessAic() { ProcessPostAic(); } + +private: + __aicore__ inline void AllocVectorEvents() + { + mte2ToVEvent_ = pipe_->AllocEventID(); + vToMte2Event_ = pipe_->AllocEventID(); + vToMte3Event_ = pipe_->AllocEventID(); + mte3ToVEvent_ = pipe_->AllocEventID(); + mte2ToMte3Event_ = pipe_->AllocEventID(); + mte3ToMte2Event_ = pipe_->AllocEventID(); + vectorEventsAllocated_ = true; + } + + __aicore__ inline void ReleaseVectorEvents() + { + if (!vectorEventsAllocated_) { + return; + } + pipe_->ReleaseEventID(mte2ToVEvent_); + pipe_->ReleaseEventID(vToMte2Event_); + pipe_->ReleaseEventID(vToMte3Event_); + pipe_->ReleaseEventID(mte3ToVEvent_); + pipe_->ReleaseEventID(mte2ToMte3Event_); + pipe_->ReleaseEventID(mte3ToMte2Event_); + vectorEventsAllocated_ = false; + } + + __aicore__ inline uint64_t QOffset(uint64_t b, uint64_t h, uint64_t t, uint64_t d) const + { + if (inputSequenceMajor_) { + return ((b * T_ + t) * H_ + h) * K_ + d; + } + return ((b * H_ + h) * T_ + t) * K_ + d; + } + + __aicore__ inline uint64_t KVOffset(uint64_t b, uint64_t hv, uint64_t t, uint64_t d, uint64_t dim) const + { + return ((b * HV_ + hv) * T_ + t) * dim + d; + } + + __aicore__ inline uint64_t BetaOffset(uint64_t b, uint64_t hv, uint64_t t) const { return (b * HV_ + hv) * T_ + t; } + + __aicore__ inline uint64_t AOffset(uint64_t b, uint64_t hv, uint64_t t, uint64_t j) const + { + return ((b * HV_ + hv) * T_ + t) * BT_ + j; + } + + __aicore__ inline uint64_t HOffset(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t d, uint64_t r) const + { + return (((b * HV_ + hv) * NT_ + chunkIdx) * K_ + d) * V_ + r; + } + + __aicore__ inline uint64_t WScratchOffset(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t t, uint64_t d) const + { + return (((b * HV_ + hv) * NT_ + chunkIdx) * BT_ + t) * K_ + d; + } + + __aicore__ inline uint64_t SolveScratchOffset(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t slot) const + { + (void)b; + (void)hv; + (void)chunkIdx; + uint64_t matrixElements = BT_ * BT_; + return solveCoreIdx_ * KDA_SOLVE_SCRATCH_SLOTS * matrixElements + slot * matrixElements; + } + + __aicore__ inline uint64_t ScoreScratchOffset(uint64_t slot, uint64_t plane, uint64_t t = 0, uint64_t d = 0) const + { + return (((solveCoreIdx_ * KDA_SCORE_QUEUE_DEPTH + slot) * KDA_SCORE_SCRATCH_PLANES + plane) * BT_ + t) * K_ + d; + } + + __aicore__ inline uint64_t ScoreRefBlockSize() const { return KDA_SCORE_REF_BC; } + + __aicore__ inline uint64_t ScoreRowBlockCount(uint64_t curT, uint64_t rowBegin) const + { + uint64_t blockSize = ScoreRefBlockSize(); + uint64_t rowCount = curT - rowBegin; + if (rowCount > blockSize) { + rowCount = blockSize; + } + return rowCount; + } + + __aicore__ inline uint64_t ScoreRefToken(uint64_t start, uint64_t curT, uint64_t rowBegin, uint64_t rowCount) const + { + uint64_t ref = rowBegin + rowCount / 2; + if (ref >= curT) { + ref = curT - 1; + } + return start + ref; + } + + __aicore__ inline void RunExp2(LocalTensor &tensor, uint32_t count) + { + SetFlag(EXP2_EVENT_ID); + WaitFlag(EXP2_EVENT_ID); + ClampExpInput(tensor, count); + Exp(tensor, tensor, count); + PipeBarrier(); + SetFlag(EXP2_EVENT_ID); + WaitFlag(EXP2_EVENT_ID); + } + + __aicore__ inline void ClampExpInput(LocalTensor &tensor, uint32_t count) + { + Mins(tensor, tensor, KDA_EXP_INPUT_MAX, count); + PipeBarrier(); + Maxs(tensor, tensor, KDA_EXP_INPUT_MIN, count); + PipeBarrier(); + } + + __aicore__ inline void ClampFp32ToOutputType(LocalTensor &tensor, uint32_t count) + { + if constexpr (IsSameType::value) { + Mins(tensor, tensor, KDA_FP16_MAX, count); + PipeBarrier(); + Maxs(tensor, tensor, -KDA_FP16_MAX, count); + PipeBarrier(); + } + } + + template + __aicore__ inline void CopyVectorIn(LocalTensor &dst, GlobalTensor &src, uint64_t offset, + uint64_t count) + { + uint64_t rowBytes = count * static_cast(sizeof(CopyT)); + if (rowBytes >= 32 && rowBytes % 32 == 0) { + DataCopy(dst, src[offset], static_cast(count)); + return; + } + DataCopyParams params{1, static_cast(rowBytes), 0, 0}; + DataCopyPadParams padParams{false, 0, 0, 0}; + DataCopyPad(dst, src[offset], params, padParams); + } + + template + __aicore__ inline void CopyVectorOut(GlobalTensor &dst, uint64_t offset, LocalTensor &src, + uint64_t count) + { + uint64_t rowBytes = count * static_cast(sizeof(CopyT)); + if (rowBytes >= 32 && rowBytes % 32 == 0) { + DataCopy(dst[offset], src, static_cast(count)); + return; + } + DataCopyParams params{1, static_cast(rowBytes), 0, 0}; + DataCopyPad(dst[offset], src, params); + } + + template + __aicore__ inline void CopyRowsIn(LocalTensor &dst, GlobalTensor &src, uint64_t offset, uint64_t rows, + uint64_t cols, uint64_t rowStride) + { + if (rows == 0 || cols == 0) { + return; + } + if (rowStride == cols) { + CopyVectorIn(dst, src, offset, rows * cols); + return; + } + DataCopyExtParams params{static_cast(rows), static_cast(cols * sizeof(CopyT)), + static_cast((rowStride - cols) * sizeof(CopyT)), 0, 0}; + DataCopyPadExtParams padParams{false, 0, 0, 0}; + DataCopyPad(dst, src[offset], params, padParams); + } + + template + __aicore__ inline void CopyRowIn(LocalTensor &dst, GlobalTensor &src, uint64_t offset) + { + CopyVectorIn(dst, src, offset, K_); + } + + template + __aicore__ inline void CopyRowOut(GlobalTensor &dst, uint64_t offset, LocalTensor &src) + { + CopyVectorOut(dst, offset, src, K_); + } + + __aicore__ inline LocalTensor VecScratch(uint64_t slot) + { + return vecBuf_.Get()[slot * EXP2_UB_ELEMENTS]; + } + + template + __aicore__ inline void LoadAsFloatRow(GlobalTensor &src, uint64_t srcOffset, LocalTensor &dst, + uint64_t count) + { + if constexpr (IsSameType::value) { + CopyVectorIn(dst, src, srcOffset, count); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + Adds(dst, dst, 0.0f, static_cast(count)); + PipeBarrier(); + SetFlag(vToMte2Event_); + WaitFlag(vToMte2Event_); + } else { + constexpr uint32_t typedOffset = EXP2_UB_ELEMENTS * sizeof(float) / sizeof(CopyT); + LocalTensor rowLocal = exp2Buf_.Get()[typedOffset]; + CopyVectorIn(rowLocal, src, srcOffset, count); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + Cast(dst, rowLocal, RoundMode::CAST_NONE, static_cast(count)); + PipeBarrier(); + SetFlag(vToMte2Event_); + WaitFlag(vToMte2Event_); + } + PipeBarrier(); + } + + template + __aicore__ inline void LoadAsFloatVector(GlobalTensor &src, uint64_t srcOffset, LocalTensor &dst, + LocalTensor &typedScratch, uint64_t count) + { + if constexpr (IsSameType::value) { + CopyVectorIn(dst, src, srcOffset, count); + } else { + CopyVectorIn(typedScratch, src, srcOffset, count); + } + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + if constexpr (!IsSameType::value) { + Cast(dst, typedScratch, RoundMode::CAST_NONE, static_cast(count)); + PipeBarrier(); + } + } + + template + __aicore__ inline void StoreFloatRow(GlobalTensor &dst, uint64_t dstOffset, LocalTensor &src, + uint64_t count) + { + if constexpr (IsSameType::value) { + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + CopyVectorOut(dst, dstOffset, src, count); + } else { + constexpr uint32_t typedOffset = EXP2_UB_ELEMENTS * sizeof(float) / sizeof(CopyT); + LocalTensor rowLocal = exp2Buf_.Get()[typedOffset]; + Cast(rowLocal, src, RoundMode::CAST_RINT, static_cast(count)); + PipeBarrier(); + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + CopyVectorOut(dst, dstOffset, rowLocal, count); + } + SetFlag(mte3ToMte2Event_); + WaitFlag(mte3ToMte2Event_); + SetFlag(mte3ToVEvent_); + WaitFlag(mte3ToVEvent_); + } + + __aicore__ inline LocalTensor Exp2NegG(uint64_t b, uint64_t hv, uint64_t t) + { + LocalTensor exp2Local = exp2Buf_.Get(); + LoadAsFloatRow(gk_, KVOffset(b, hv, t, 0, K_), exp2Local, K_); + Muls(exp2Local, exp2Local, -LN2, static_cast(K_)); + PipeBarrier(); + RunExp2(exp2Local, static_cast(K_)); + return exp2Local; + } + + __aicore__ inline uint64_t ScoreVectorMaxRows(uint64_t bytesPerElem) const + { + constexpr uint64_t arenaBytes = static_cast(KDA_VEC_ARENA_ELEMENTS) * sizeof(float); + uint64_t maxRows = (arenaBytes / bytesPerElem) / K_; + if (K_ >= 128 && maxRows > 32) { + maxRows = 32; + } + return maxRows; + } + __aicore__ inline bool UsePostWuCube(uint64_t curT) const + { + return curT > 0 && curT <= BT_ && (BT_ == 64 || BT_ == 128) && K_ >= 16 && V_ >= 16 && V_ <= 256 && + K_ % 16 == 0 && V_ % 16 == 0; + } + + __aicore__ inline void ComputePostWuCube(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, uint64_t curT) + { + using ElementA = AKK_T; + using ElementB = T; + using LayoutTagA = Catlass::layout::RowMajor; + using LayoutTagB = Catlass::layout::RowMajor; + using LayoutTagC = Catlass::layout::RowMajor; + using WTileCopy = Catlass::Gemm::Tile::PackedTileCopyTla; + using UTileCopy = Catlass::Gemm::Tile::PackedTileCopyTla; + using PostL1TileShape128 = tla::Shape; + using PostL0TileShape128 = tla::Shape; + using PostL1TileShape256 = tla::Shape; + using PostL0TileShape256 = tla::Shape; + using WBlockMmad = Catlass::Gemm::Block::BlockMmadTla; + using UBlockMmad128 = + Catlass::Gemm::Block::BlockMmadTla; + using UBlockMmad256 = + Catlass::Gemm::Block::BlockMmadTla; + LayoutTagA tagA = LayoutTagA::template MakeLayout(BT_, BT_); + auto layoutA = tla::MakeLayoutFromTag(tagA); + auto tensorA = tla::MakeTensor(preparedAqk_[AOffset(b, hv, start, 0)], layoutA, Catlass::Arch::PositionGM{}); + + { + LayoutTagB tagB = LayoutTagB::template MakeLayout(BT_, K_); + LayoutTagC tagC = LayoutTagC::template MakeLayout(BT_, K_); + auto layoutB = tla::MakeLayoutFromTag(tagB); + auto layoutC = tla::MakeLayoutFromTag(tagC); + Catlass::GemmCoord shape{static_cast(curT), static_cast(K_), + static_cast(curT)}; + auto tensorB = + tla::MakeTensor(preparedQG_[KVOffset(b, hv, start, 0, K_)], layoutB, Catlass::Arch::PositionGM{}); + auto tensorC = + tla::MakeTensor(h_[WScratchOffset(b, hv, chunkIdx, 0, 0)], layoutC, Catlass::Arch::PositionGM{}); + auto blockA = GetTile(tensorA, tla::MakeCoord(0, 0), tla::MakeShape(shape.m(), shape.k())); + auto blockB = GetTile(tensorB, tla::MakeCoord(0, 0), tla::MakeShape(shape.k(), shape.n())); + auto blockC = GetTile(tensorC, tla::MakeCoord(0, 0), tla::MakeShape(shape.m(), shape.n())); + Catlass::Arch::Resource wResource; + WBlockMmad wBlockMmad(wResource); + wBlockMmad(blockA, blockB, blockC, shape); + PipeBarrier(); + } + + { + LayoutTagB tagB = LayoutTagB::template MakeLayout(BT_, V_); + LayoutTagC tagC = LayoutTagC::template MakeLayout(BT_, V_); + auto layoutB = tla::MakeLayoutFromTag(tagB); + auto layoutC = tla::MakeLayoutFromTag(tagC); + Catlass::GemmCoord shape{static_cast(curT), static_cast(V_), + static_cast(curT)}; + auto tensorB = + tla::MakeTensor(propagatedVNew_[KVOffset(b, hv, start, 0, V_)], layoutB, Catlass::Arch::PositionGM{}); + auto tensorC = tla::MakeTensor(u_[KVOffset(b, hv, start, 0, V_)], layoutC, Catlass::Arch::PositionGM{}); + auto blockA = GetTile(tensorA, tla::MakeCoord(0, 0), tla::MakeShape(shape.m(), shape.k())); + auto blockB = GetTile(tensorB, tla::MakeCoord(0, 0), tla::MakeShape(shape.k(), shape.n())); + auto blockC = GetTile(tensorC, tla::MakeCoord(0, 0), tla::MakeShape(shape.m(), shape.n())); + Catlass::Arch::Resource uResource; + if (V_ <= 128) { + UBlockMmad128 uBlockMmad(uResource); + uBlockMmad(blockA, blockB, blockC, shape); + } else { + UBlockMmad256 uBlockMmad(uResource); + uBlockMmad(blockA, blockB, blockC, shape); + } + PipeBarrier(); + } + } + + __aicore__ inline void CopyScratchWAndFinalizeKg(uint64_t b, uint64_t h, uint64_t hv, uint64_t chunkIdx, + uint64_t start, uint64_t curT, uint64_t subBlockIdx, + uint64_t subBlockNum, bool copyScratchW) + { + constexpr uint64_t typedOffsetFloats = 20480; + constexpr uint64_t typedOffset = typedOffsetFloats * sizeof(float) / sizeof(T); + constexpr uint64_t kgFp32Planes = 4; + uint64_t rowBegin = (curT * subBlockIdx) / subBlockNum; + uint64_t rowEnd = (curT * (subBlockIdx + 1)) / subBlockNum; + if (rowBegin >= rowEnd) { + return; + } + uint64_t maxRows = (typedOffsetFloats / kgFp32Planes) / K_; + if (maxRows > 32) { + maxRows = 32; + } + if (maxRows == 0) { + return; + } + + uint64_t last = start + curT - 1; + LocalTensor arena = vecBuf_.Get(); + LocalTensor gateLast = exp2Buf_.Get(); + LocalTensor typedLocal = vecBuf_.Get()[typedOffset]; + LoadAsFloatRow(gk_, KVOffset(b, hv, last, 0, K_), gateLast, K_); + + for (uint64_t tileRow = rowBegin; tileRow < rowEnd; tileRow += maxRows) { + uint64_t tileRows = rowEnd - tileRow; + if (tileRows > maxRows) { + tileRows = maxRows; + } + uint64_t elemCount = tileRows * K_; + uint64_t token = start + tileRow; + + if (copyScratchW) { + uint64_t scratchBase = WScratchOffset(b, hv, chunkIdx, tileRow, 0); + DataCopy(arena, h_[scratchBase], static_cast(elemCount)); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + Cast(typedLocal, arena, RoundMode::CAST_RINT, static_cast(elemCount)); + PipeBarrier(); + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + DataCopy(w_[KVOffset(b, hv, token, 0, K_)], typedLocal, static_cast(elemCount)); + SetFlag(mte3ToMte2Event_); + WaitFlag(mte3ToMte2Event_); + } + + LocalTensor kLocal = arena; + LocalTensor gLocal = arena[elemCount]; + LocalTensor expLocal = arena[2 * elemCount]; + LocalTensor outLocal = arena[3 * elemCount]; + const uint64_t gateOffsetBytes = (typedOffset + elemCount) * sizeof(T); + LocalTensor gateTyped = vecBuf_.Get()[(gateOffsetBytes + sizeof(GK_T) - 1) / sizeof(GK_T)]; + CopyRowsIn(typedLocal, k_, QOffset(b, h, token, 0), tileRows, K_, inputSequenceMajor_ ? H_ * K_ : K_); + LoadAsFloatVector(gk_, KVOffset(b, hv, token, 0, K_), gLocal, gateTyped, elemCount); + Cast(kLocal, typedLocal, RoundMode::CAST_NONE, static_cast(elemCount)); + PipeBarrier(); + + for (uint64_t row = 0; row < tileRows; ++row) { + Sub(expLocal[row * K_], gateLast, gLocal[row * K_], static_cast(K_)); + } + PipeBarrier(); + Muls(expLocal, expLocal, LN2, static_cast(elemCount)); + PipeBarrier(); + ClampExpInput(expLocal, static_cast(elemCount)); + Exp(expLocal, expLocal, static_cast(elemCount)); + PipeBarrier(); + Mul(outLocal, kLocal, expLocal, static_cast(elemCount)); + PipeBarrier(); + ClampFp32ToOutputType(outLocal, static_cast(elemCount)); + Cast(typedLocal, outLocal, RoundMode::CAST_RINT, static_cast(elemCount)); + PipeBarrier(); + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + CopyVectorOut(kg_, KVOffset(b, hv, token, 0, K_), typedLocal, elemCount); + SetFlag(mte3ToMte2Event_); + WaitFlag(mte3ToMte2Event_); + } + if (rowEnd == curT) { + CopyVectorIn(typedLocal, k_, QOffset(b, h, last, 0), K_); + SetFlag(mte2ToMte3Event_); + WaitFlag(mte2ToMte3Event_); + CopyVectorOut(kg_, KVOffset(b, hv, last, 0, K_), typedLocal, K_); + SetFlag(mte3ToMte2Event_); + WaitFlag(mte3ToMte2Event_); + } + SetFlag(mte3ToVEvent_); + WaitFlag(mte3ToVEvent_); + } + + template + __aicore__ inline void ComputeTailWuRow(GlobalTensor &src, GlobalTensor &dst, + uint64_t akkBase, uint64_t srcBase, uint64_t dstBase, uint64_t curT, + uint64_t dim) + { + LocalTensor acc = vecBuf_.Get(); + LocalTensor value = vecBuf_.Get()[512]; + LocalTensor typed = vecBuf_.Get()[4096]; + LocalTensor coefficientTyped = exp2Buf_.Get(); + LocalTensor coefficients = exp2Buf_.Get()[128]; + CopyVectorIn(coefficientTyped, preparedAqk_, akkBase, curT); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + Cast(coefficients, coefficientTyped, RoundMode::CAST_NONE, static_cast(curT)); + PipeBarrier(); + SetFlag(EXP2_EVENT_ID); + WaitFlag(EXP2_EVENT_ID); + for (uint64_t j = 0; j < curT; ++j) { + LoadAsFloatVector(src, srcBase + j * dim, value, typed, dim); + float coefficient = coefficients.GetValue(j); + SetFlag(EXP2_EVENT_ID); + WaitFlag(EXP2_EVENT_ID); + Muls(value, value, coefficient, static_cast(dim)); + PipeBarrier(); + if (j == 0) { + Adds(acc, value, 0.0f, static_cast(dim)); + } else { + Add(acc, acc, value, static_cast(dim)); + } + PipeBarrier(); + SetFlag(vToMte2Event_); + WaitFlag(vToMte2Event_); + } + SetFlag(EXP2_EVENT_ID); + WaitFlag(EXP2_EVENT_ID); + ClampFp32ToOutputType(acc, static_cast(dim)); + StoreFloatRow(dst, dstBase, acc, dim); + } + + __aicore__ inline void ComputeTailWuVector(uint64_t b, uint64_t hv, uint64_t start, uint64_t curT, + uint64_t subBlockIdx, uint64_t subBlockNum) + { + uint64_t rowBegin = (curT * subBlockIdx) / subBlockNum; + uint64_t rowEnd = (curT * (subBlockIdx + 1)) / subBlockNum; + for (uint64_t row = rowBegin; row < rowEnd; ++row) { + ComputeTailWuRow(preparedQG_, w_, AOffset(b, hv, start + row, 0), KVOffset(b, hv, start, 0, K_), + KVOffset(b, hv, start + row, 0, K_), curT, K_); + ComputeTailWuRow(propagatedVNew_, u_, AOffset(b, hv, start + row, 0), KVOffset(b, hv, start, 0, V_), + KVOffset(b, hv, start + row, 0, V_), curT, V_); + } + } + + __aicore__ inline bool ResolveFlatChunk(uint64_t task, uint64_t &seq, uint64_t &b, uint64_t &h, uint64_t &hv, + uint64_t &chunkIdx, uint64_t &start, uint64_t &end) + { + hv = task % HV_; + uint64_t flatChunk = task / HV_; + if (!isVarLen_) { + seq = flatChunk / NT_; + b = seq; + chunkIdx = flatChunk % NT_; + start = chunkIdx * BT_; + end = start + BT_; + if (end > T_) { + end = T_; + } + } else { + if (!KdaVarlen::ResolveChunkRange(cuSeqlensAddr_, chunkIndicesAddr_, N_, T_, BT_, flatChunk, seq, start, + end)) { + return false; + } + b = 0; + chunkIdx = flatChunk; + } + h = hv / (HV_ / H_); + return start < end; + } + + __aicore__ inline void ProcessChunkPostAiv(uint64_t b, uint64_t h, uint64_t hv, uint64_t chunkIdx, uint64_t start, + uint64_t end, uint64_t subBlockIdx, uint64_t subBlockNum) + { + uint64_t curT = end - start; + if (curT == 0 || !UsePostWuCube(curT)) { + return; + } + if (curT < BT_) { + ComputeTailWuVector(b, hv, start, curT, subBlockIdx, subBlockNum); + CopyScratchWAndFinalizeKg(b, h, hv, chunkIdx, start, curT, subBlockIdx, subBlockNum, false); + return; + } + Catlass::Arch::CrossCoreWaitFlagWithReverse<0x2, PIPE_MTE2>(syncDoneFlag_); + CopyScratchWAndFinalizeKg(b, h, hv, chunkIdx, start, curT, subBlockIdx, subBlockNum, true); + } + + __aicore__ inline void ProcessChunkPostAic(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, uint64_t end) + { + if constexpr (IsSameType::value) { + ProcessChunkPostAicTyped(b, hv, chunkIdx, start, end); + } + } + + __aicore__ inline void ProcessChunkPostAicTyped(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, + uint64_t end) + { + uint64_t curT = end - start; + if (curT == 0 || !UsePostWuCube(curT)) { + return; + } + if (curT < BT_) { + return; + } + ComputePostWuCube(b, hv, chunkIdx, start, curT); + Catlass::Arch::CrossCoreSetFlagWithReverse<0x2, PIPE_FIX>(syncDoneFlag_); + } + + __aicore__ inline void ProcessPostAiv() + { + if constexpr (IsSameType::value) { + return; + } + uint64_t subBlockNum = static_cast(GetSubBlockNum()); + if (subBlockNum == 0) { + return; + } + uint64_t subBlockIdx = static_cast(GetSubBlockIdx()); + uint64_t coreNum = usedCoreNum_ == 0 ? 1 : usedCoreNum_; + uint64_t coreIdx = static_cast(GetBlockIdx()) / subBlockNum; + uint64_t taskNum = static_cast((isVarLen_ ? NT_ : B_ * NT_) * HV_); + for (uint64_t task = coreIdx; task < taskNum; task += coreNum) { + uint64_t seq = 0; + uint64_t b = 0; + uint64_t h = 0; + uint64_t hv = 0; + uint64_t chunkIdx = 0; + uint64_t start = 0; + uint64_t end = 0; + if (ResolveFlatChunk(task, seq, b, h, hv, chunkIdx, start, end)) { + (void)seq; + ProcessChunkPostAiv(b, h, hv, chunkIdx, start, end, subBlockIdx, subBlockNum); + } + } + } + + __aicore__ inline void ProcessPostAic() + { + if constexpr (IsSameType::value) { + return; + } + uint64_t taskNum = static_cast((isVarLen_ ? NT_ : B_ * NT_) * HV_); + uint64_t coreNum = usedCoreNum_ == 0 ? 1 : usedCoreNum_; + for (uint64_t task = GetBlockIdx(); task < taskNum; task += coreNum) { + uint64_t seq = 0; + uint64_t b = 0; + uint64_t h = 0; + uint64_t hv = 0; + uint64_t chunkIdx = 0; + uint64_t start = 0; + uint64_t end = 0; + if (ResolveFlatChunk(task, seq, b, h, hv, chunkIdx, start, end)) { + (void)seq; + (void)h; + ProcessChunkPostAic(b, hv, chunkIdx, start, end); + } + } + } + +private: + GlobalTensor q_; + GlobalTensor k_; + GlobalTensor v_; + GlobalTensor gk_; + GlobalTensor beta_; + GlobalTensor initialState_; + GlobalTensor o_; + GlobalTensor finalState_; + GlobalTensor aqk_; + GlobalTensor akk_; + GlobalTensor w_; + GlobalTensor u_; + GlobalTensor qg_; + GlobalTensor kg_; + GlobalTensor vNew_; + GlobalTensor h_; + GlobalTensor preparedQG_; + GlobalTensor preparedAqk_; + GlobalTensor propagatedVNew_; + GlobalTensor propagatedH_; + GlobalTensor solveWorkspace_; + GlobalTensor scoreWorkspace_; + TPipe *pipe_ = nullptr; + TBuf exp2Buf_; + TBuf vecBuf_; + TBuf gateWritebackBuf_; + TEventID mte2ToVEvent_ = 0; + TEventID vToMte2Event_ = 0; + TEventID vToMte3Event_ = 0; + TEventID mte3ToVEvent_ = 0; + TEventID mte2ToMte3Event_ = 0; + TEventID mte3ToMte2Event_ = 0; + bool vectorEventsAllocated_ = false; + Catlass::Arch::CrossCoreFlagWithReverse scoreReadyFlag_{KDA_SCORE_READY_FLAG0, + KDA_SCORE_READY_FLAG1}; + Catlass::Arch::CrossCoreFlagWithReverse scoreDoneFlag_{KDA_SCORE_DONE_FLAG0, + KDA_SCORE_DONE_FLAG1}; + // Score production is fully drained before solve starts, so the solve handshake can safely reuse + // the existing score flags without consuming additional hardware flag IDs. + Catlass::Arch::CrossCoreFlagWithReverse syncReadyFlag_{KDA_SCORE_READY_FLAG0, + KDA_SCORE_READY_FLAG1}; + Catlass::Arch::CrossCoreFlagWithReverse syncDoneFlag_{KDA_SCORE_DONE_FLAG0, + KDA_SCORE_DONE_FLAG1}; + uint64_t B_ = 0; + uint64_t N_ = 0; + uint64_t H_ = 0; + uint64_t HV_ = 0; + uint64_t T_ = 0; + uint64_t K_ = 0; + uint64_t V_ = 0; + uint64_t BT_ = 0; + uint64_t NT_ = 0; + float scale_ = 1.0f; + bool hasInitial_ = false; + bool isVarLen_ = false; + bool isAivOnly_ = false; + bool inputSequenceMajor_ = false; + uint64_t usedCoreNum_ = 1; + uint64_t solveCoreIdx_ = 0; + __gm__ int64_t *chunkIndicesAddr_ = nullptr; + __gm__ int64_t *cuSeqlensAddr_ = nullptr; +}; +} // namespace + +template +__aicore__ inline void RunChunkKdaPostWu(GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR gk, GM_ADDR beta, + GM_ADDR initialState, GM_ADDR cuSeqlens, GM_ADDR chunkIndices, GM_ADDR wSeed, + GM_ADDR akk, GM_ADDR uSeed, GM_ADDR w, GM_ADDR u, GM_ADDR kg, GM_ADDR vNew, + GM_ADDR userWorkspace, const TilingData &tiling, TPipe &pipe) +{ + GM_ADDR postScratch = userWorkspace + tiling.postWuScratchOffset; + if ASCEND_IS_AIC { + ChunkKdaFwdPostWuKernel op; + op.Init(q, k, v, gk, beta, initialState, cuSeqlens, chunkIndices, wSeed, akk, uSeed, nullptr, userWorkspace, + userWorkspace, userWorkspace, akk, w, u, userWorkspace, kg, vNew, postScratch, postScratch, tiling, + &pipe, false); + op.ProcessAic(); + } + if ASCEND_IS_AIV { + ChunkKdaFwdPostWuKernel op; + op.Init(q, k, v, gk, beta, initialState, cuSeqlens, chunkIndices, wSeed, akk, uSeed, nullptr, userWorkspace, + userWorkspace, userWorkspace, akk, w, u, userWorkspace, kg, vNew, postScratch, postScratch, tiling, + &pipe); + op.ProcessAiv(); + } +} + +} // namespace KdaPostWu diff --git a/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_prepare.h b/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_prepare.h new file mode 100644 index 0000000000..c5f6004802 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_prepare.h @@ -0,0 +1,2518 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#pragma once + +#ifndef CATLASS_ARCH +#define CATLASS_ARCH 2201 +#endif + +#include "catlass/arch/kda_arch.hpp" +#include "catlass/arch/kda_cross_core_sync.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/kda_catlass.hpp" +#include "catlass/gemm/block/kda_block_mmad.hpp" +#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" +#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "kernel_utils/block/block_mmad_pingpong_tla_multi.hpp" +#include "catlass/layout/kda_layout.hpp" +#include "kernel_operator.h" +#include "chunk_kda_fwd_varlen.h" +#include "tla/kda_tla_layout.hpp" +#include "tla/kda_tensor.hpp" + +using namespace AscendC; + +namespace KdaPrepare { +namespace { +using KdaInt64 = tla::Int<64>; +using KdaInt128 = tla::Int<128>; +constexpr float LN2 = 0.69314718055994530942f; +constexpr float KDA_EXP2_CLAMP = 80.0f; +constexpr float KDA_EXP_INPUT_MAX = KDA_EXP2_CLAMP * LN2; +constexpr float KDA_EXP_INPUT_MIN = -KDA_EXP2_CLAMP * LN2; +constexpr float KDA_SCORE_EXP2_CLAMP = 120.0f; +constexpr float KDA_SCORE_EXP2_MIN_CLAMP = 126.0f; +constexpr float KDA_SCORE_EXP_INPUT_MAX = KDA_SCORE_EXP2_CLAMP * LN2; +constexpr float KDA_SCORE_EXP_INPUT_MIN = -KDA_SCORE_EXP2_MIN_CLAMP * LN2; +constexpr float KDA_FP16_MAX = 65504.0f; +constexpr uint32_t EXP2_UB_ELEMENTS = 256; +constexpr uint32_t EXP2_UB_BYTES = EXP2_UB_ELEMENTS * (sizeof(float) + sizeof(uint16_t)); +constexpr uint32_t EXP2_EVENT_ID = 0; +constexpr uint32_t KDA_SOLVE_BT = 64; +constexpr uint32_t KDA_SOLVE_MATRIX_ELEMENTS = KDA_SOLVE_BT * KDA_SOLVE_BT; +constexpr uint32_t KDA_SOLVE_SCRATCH_X = 0; +constexpr uint32_t KDA_SOLVE_SCRATCH_Y0 = 1; +constexpr uint32_t KDA_SOLVE_SCRATCH_TMP = 2; +constexpr uint32_t KDA_SOLVE_SCRATCH_Y1 = 3; +constexpr uint32_t KDA_SOLVE_SCRATCH_IDENTITY = 4; +constexpr uint32_t KDA_SOLVE_SCRATCH_SLOTS = 5; +constexpr uint32_t KDA_SOLVE_PIPELINE_DEPTH = 4; +constexpr uint32_t KDA_SOLVE_DIAG_BT = 16; +constexpr uint32_t KDA_SOLVE_DIAG_BLOCKS = KDA_SOLVE_BT / KDA_SOLVE_DIAG_BT; +constexpr uint32_t KDA_SOLVE_DIAG_MCH_ITERS = 3; +// Keep the local safe-gate exponent span within the BF16 score range while +// reducing repeated gate-factor work and AIV/AIC handshakes. +constexpr uint32_t KDA_SCORE_REF_BC = 32; +constexpr uint32_t KDA_SAFE_SCORE_REF_BC = 32; +constexpr uint32_t KDA_VEC_ARENA_ELEMENTS = 32768; +constexpr uint32_t KDA_BITS_PER_MASK_BYTE = 8; +constexpr uint32_t KDA_SELECT_COL_BLOCKS = 2; +constexpr uint32_t KDA_SELECT_COL_MASK_BYTES = KDA_SOLVE_MATRIX_ELEMENTS / KDA_BITS_PER_MASK_BYTE; +constexpr uint32_t KDA_SELECT_MASK_BYTES = KDA_SELECT_COL_BLOCKS * KDA_SELECT_COL_MASK_BYTES; +constexpr uint32_t KDA_SELECT_AQK_MASK_BYTE_OFFSET = 120 * 1024; +constexpr uint32_t KDA_SELECT_AKK_MASK_BYTE_OFFSET = KDA_SELECT_AQK_MASK_BYTE_OFFSET + KDA_SELECT_MASK_BYTES; +constexpr uint32_t KDA_SELECT_ZERO_BYTE_OFFSET = KDA_SELECT_AKK_MASK_BYTE_OFFSET + KDA_SELECT_MASK_BYTES; +constexpr uint32_t KDA_SELECT_ZERO_FLOAT_OFFSET = KDA_SELECT_ZERO_BYTE_OFFSET / sizeof(float); +constexpr uint8_t KDA_SCORE_DONE_FLAG0 = 2; +constexpr uint8_t KDA_SCORE_DONE_FLAG1 = 3; +constexpr uint8_t KDA_SCORE_READY_FLAG0 = 4; +constexpr uint8_t KDA_SCORE_READY_FLAG1 = 5; +constexpr uint8_t KDA_SOLVE_DONE_FLAG = 6; +constexpr uint8_t KDA_SOLVE_READY_FLAG = 7; +constexpr uint32_t KDA_SCORE_QUEUE_DEPTH = 2; +constexpr uint32_t KDA_SCORE_LANES = 2; +constexpr uint32_t KDA_SCORE_SCRATCH_SLOTS = KDA_SCORE_QUEUE_DEPTH * KDA_SCORE_LANES; +constexpr uint32_t KDA_SYNC_REVERSE_DEPTH = 1; +constexpr uint32_t KDA_SCORE_SCRATCH_PLANES = 3; +constexpr uint32_t KDA_SCORE_SCRATCH_QG = 0; +constexpr uint32_t KDA_SCORE_SCRATCH_W = 1; +constexpr uint32_t KDA_SCORE_SCRATCH_KG = 2; +constexpr uint64_t KDA_WORKSPACE_ALIGN = 512; +constexpr uint32_t KDA_GATE_TILE_ROWS = 16; +constexpr uint32_t KDA_GATE_PIPELINE_DEPTH = 3; +constexpr uint32_t KDA_AIV_UB_BUDGET_BYTES = 192 * 1024; +using KdaArchTag = Catlass::Arch::AtlasA2; +using KdaDispatchPolicy = Catlass::Gemm::MmadPingpong; +using KdaScoreDispatchPolicy = Catlass::Gemm::MmadPingpongTlaMulti; +static_assert(KdaScoreDispatchPolicy::ENABLE_L1_RESIDENT, + "KDA Aqk/Akk score MMAD must keep the shared right matrix resident in L1"); +static_assert(KdaScoreDispatchPolicy::L1B_STAGES == 1, + "KDA Aqk/Akk score MMAD needs one L1 B slot so the second MMAD reuses it"); +using KdaSolveDispatchPolicy = Catlass::Gemm::MmadPingpong; +static_assert(!KdaSolveDispatchPolicy::USE_HF32_MODE, "KDA triangular solve must use IEEE FP32 Cube mode"); +using KdaL1TileShape = tla::Shape; +using KdaL0TileShape = KdaL1TileShape; +using KdaSolveL1TileShape = tla::Shape; +using KdaSolveL0TileShape = KdaSolveL1TileShape; + +__aicore__ inline uint32_t FloatToBits(float value) +{ + union Bits { + __aicore__ Bits() {} + float f; + uint32_t u; + } bits; + bits.f = value; + return bits.u; +} + +__aicore__ inline float BitsToFloat(uint32_t value) +{ + union Bits { + __aicore__ Bits() {} + uint32_t u; + float f; + } bits; + bits.u = value; + return bits.f; +} + +__aicore__ inline uint16_t Bf16ToBits(bfloat16_t value) +{ + union Bits { + __aicore__ Bits() {} + bfloat16_t f; + uint16_t u; + } bits; + bits.f = value; + return bits.u; +} + +__aicore__ inline bfloat16_t BitsToBf16(uint16_t value) +{ + union Bits { + __aicore__ Bits() {} + uint16_t u; + bfloat16_t f; + } bits; + bits.u = value; + return bits.f; +} + +template +__aicore__ inline T FloatToType(float value) +{ + if constexpr (IsSameType::value) { + uint32_t bits = FloatToBits(value); + uint32_t bias = 0x7FFFu + ((bits >> 16) & 1u); + return BitsToBf16(static_cast((bits + bias) >> 16)); + } + return static_cast(value); +} + +template +class ChunkKdaFwdPrepareKernel { +public: + using OUT_T = T; + using AKK_T = float; + using SCORE_T = std::conditional_t::value, bfloat16_t, T>; + template + __aicore__ inline void Init(GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR gk, GM_ADDR beta, GM_ADDR initialState, + GM_ADDR cuSeqlens, GM_ADDR chunkIndices, GM_ADDR preparedQG, GM_ADDR preparedAqk, + GM_ADDR propagatedVNew, GM_ADDR propagatedH, GM_ADDR o, GM_ADDR finalState, GM_ADDR aqk, + GM_ADDR akk, GM_ADDR w, GM_ADDR u, GM_ADDR qg, GM_ADDR kg, GM_ADDR vNew, GM_ADDR h, + GM_ADDR workspace, const TilingData &tiling, TPipe *pipe, bool initVecBuffers = true) + { + pipe_ = pipe; + q_.SetGlobalBuffer((__gm__ T *)q); + k_.SetGlobalBuffer((__gm__ T *)k); + v_.SetGlobalBuffer((__gm__ T *)v); + gk_.SetGlobalBuffer((__gm__ GK_T *)gk); + beta_.SetGlobalBuffer((__gm__ BETA_T *)beta); + if (initialState != nullptr) { + initialState_.SetGlobalBuffer((__gm__ float *)initialState); + } + cuSeqlensAddr_ = reinterpret_cast<__gm__ int64_t *>(cuSeqlens); + if (preparedQG != nullptr) { + preparedQG_.SetGlobalBuffer((__gm__ T *)preparedQG); + } + if (preparedAqk != nullptr) { + preparedAqk_.SetGlobalBuffer((__gm__ T *)preparedAqk); + } + if (propagatedVNew != nullptr) { + propagatedVNew_.SetGlobalBuffer((__gm__ T *)propagatedVNew); + } + if (propagatedH != nullptr) { + propagatedH_.SetGlobalBuffer((__gm__ T *)propagatedH); + } + chunkIndicesAddr_ = reinterpret_cast<__gm__ int64_t *>(chunkIndices); + o_.SetGlobalBuffer((__gm__ OUT_T *)o); + finalState_.SetGlobalBuffer((__gm__ float *)finalState); + aqk_.SetGlobalBuffer((__gm__ float *)aqk); + akk_.SetGlobalBuffer((__gm__ AKK_T *)akk); + w_.SetGlobalBuffer((__gm__ T *)w); + u_.SetGlobalBuffer((__gm__ OUT_T *)u); + qg_.SetGlobalBuffer((__gm__ T *)qg); + kg_.SetGlobalBuffer((__gm__ T *)kg); + vNew_.SetGlobalBuffer((__gm__ T *)vNew); + h_.SetGlobalBuffer((__gm__ float *)h); + solveWorkspace_.SetGlobalBuffer((__gm__ float *)workspace); + + B_ = tiling.batch; + N_ = tiling.seqNum; + H_ = tiling.qHeadNum; + HV_ = tiling.vHeadNum; + T_ = tiling.seqlen; + K_ = tiling.kHeadDim; + V_ = tiling.vHeadDim; + BT_ = tiling.chunkSize; + NT_ = tiling.totalChunks; + scale_ = tiling.scale; + hasInitial_ = tiling.hasInitialState; + isVarLen_ = tiling.isVarLen; + inputSequenceMajor_ = tiling.inputSequenceMajor; + usedCoreNum_ = tiling.prepareUsedCoreNum; + constexpr uint64_t solvePipelineDepth = SAFE_GATE ? KDA_SOLVE_PIPELINE_DEPTH : 1; + const uint64_t solveBytes = + usedCoreNum_ * solvePipelineDepth * KDA_SOLVE_SCRATCH_SLOTS * BT_ * BT_ * sizeof(float); + const uint64_t alignedSolveBytes = + (solveBytes + KDA_WORKSPACE_ALIGN - 1) / KDA_WORKSPACE_ALIGN * KDA_WORKSPACE_ALIGN; + scoreWorkspace_.SetGlobalBuffer((__gm__ SCORE_T *)(workspace + alignedSolveBytes)); + if ASCEND_IS_AIV { + uint64_t subBlockNum = static_cast(GetSubBlockNum()); + solveCoreIdx_ = subBlockNum == 0 ? 0 : static_cast(GetBlockIdx()) / subBlockNum; + } else { + solveCoreIdx_ = static_cast(GetBlockIdx()); + } + if (pipe_ != nullptr && initVecBuffers) { + pipe_->InitBuffer(exp2Buf_, EXP2_UB_BYTES); + pipe_->InitBuffer(vecBuf_, KDA_VEC_ARENA_ELEMENTS * sizeof(float)); + const uint64_t gateStageElems = GatePipelineRows() * K_; + const uint64_t gateInputSlotBytes = gateStageElems * (2 * sizeof(T) + sizeof(GK_T)); + const uint64_t gatePipelineBytes = + KDA_GATE_PIPELINE_DEPTH * (gateInputSlotBytes + gateStageElems * sizeof(T)); + pipe_->InitBuffer(gateWritebackBuf_, static_cast(gatePipelineBytes)); + AllocVectorEvents(); + } + } + __aicore__ inline void ProcessAivOnly() + { + isAivOnly_ = true; + ProcessPreAiv(); + ReleaseVectorEvents(); + } + + __aicore__ inline void ProcessAiv() + { + ProcessPreAiv(); + ReleaseVectorEvents(); + } + + __aicore__ inline void ProcessAic() { ProcessPreAic(); } + +private: + __aicore__ inline void AllocVectorEvents() + { + mte2ToVEvent_ = pipe_->AllocEventID(); + vToMte2Event_ = pipe_->AllocEventID(); + vToMte3Event_ = pipe_->AllocEventID(); + mte3ToVEvent_ = pipe_->AllocEventID(); + mte2ToMte3Event_ = pipe_->AllocEventID(); + for (uint32_t slot = 0; slot < KDA_GATE_PIPELINE_DEPTH; ++slot) { + mte3ToMte2Events_[slot] = pipe_->AllocEventID(); + } + vectorEventsAllocated_ = true; + } + + __aicore__ inline void ReleaseVectorEvents() + { + if (!vectorEventsAllocated_) { + return; + } + pipe_->ReleaseEventID(mte2ToVEvent_); + pipe_->ReleaseEventID(vToMte2Event_); + pipe_->ReleaseEventID(vToMte3Event_); + pipe_->ReleaseEventID(mte3ToVEvent_); + pipe_->ReleaseEventID(mte2ToMte3Event_); + for (uint32_t slot = 0; slot < KDA_GATE_PIPELINE_DEPTH; ++slot) { + pipe_->ReleaseEventID(mte3ToMte2Events_[slot]); + } + vectorEventsAllocated_ = false; + } + + __aicore__ inline uint64_t QOffset(uint64_t b, uint64_t h, uint64_t t, uint64_t d) const + { + if (inputSequenceMajor_) { + return ((b * T_ + t) * H_ + h) * K_ + d; + } + return ((b * H_ + h) * T_ + t) * K_ + d; + } + + __aicore__ inline uint64_t VInputOffset(uint64_t b, uint64_t hv, uint64_t t, uint64_t d) const + { + if (inputSequenceMajor_) { + return ((b * T_ + t) * HV_ + hv) * V_ + d; + } + return ((b * HV_ + hv) * T_ + t) * V_ + d; + } + + __aicore__ inline uint64_t KVOffset(uint64_t b, uint64_t hv, uint64_t t, uint64_t d, uint64_t dim) const + { + return ((b * HV_ + hv) * T_ + t) * dim + d; + } + + __aicore__ inline uint64_t BetaOffset(uint64_t b, uint64_t hv, uint64_t t) const { return (b * HV_ + hv) * T_ + t; } + + __aicore__ inline uint64_t AOffset(uint64_t b, uint64_t hv, uint64_t t, uint64_t j) const + { + return ((b * HV_ + hv) * T_ + t) * BT_ + j; + } + + __aicore__ inline uint64_t HOffset(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t d, uint64_t r) const + { + return (((b * HV_ + hv) * NT_ + chunkIdx) * K_ + d) * V_ + r; + } + + __aicore__ inline uint64_t WScratchOffset(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t t, uint64_t d) const + { + return (((b * HV_ + hv) * NT_ + chunkIdx) * BT_ + t) * K_ + d; + } + + __aicore__ inline uint64_t SolveScratchOffset(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t slot) const + { + (void)b; + (void)hv; + (void)chunkIdx; + constexpr uint64_t solvePipelineDepth = SAFE_GATE ? KDA_SOLVE_PIPELINE_DEPTH : 1; + uint64_t matrixElements = BT_ * BT_; + return ((solveCoreIdx_ * solvePipelineDepth + activeSolveSlot_) * KDA_SOLVE_SCRATCH_SLOTS + slot) * + matrixElements; + } + + __aicore__ inline uint64_t ScoreScratchOffset(uint64_t slot, uint64_t plane, uint64_t t = 0, uint64_t d = 0) const + { + return (((solveCoreIdx_ * KDA_SCORE_SCRATCH_SLOTS + slot) * KDA_SCORE_SCRATCH_PLANES + plane) * BT_ + t) * K_ + + d; + } + + __aicore__ inline uint64_t ScoreScratchSlot(uint64_t queueSlot, uint64_t lane, bool pairHeads) const + { + return pairHeads ? queueSlot * KDA_SCORE_LANES + lane : queueSlot; + } + + __aicore__ inline uint64_t ScoreRefBlockSize() const + { + if constexpr (SAFE_GATE) { + return KDA_SAFE_SCORE_REF_BC; + } + return KDA_SCORE_REF_BC; + } + + __aicore__ inline uint64_t ScoreRowBlockCount(uint64_t curT, uint64_t rowBegin) const + { + uint64_t blockSize = ScoreRefBlockSize(); + uint64_t rowCount = curT - rowBegin; + if (rowCount > blockSize) { + rowCount = blockSize; + } + return rowCount; + } + + __aicore__ inline uint64_t ScoreRefToken(uint64_t start, uint64_t curT, uint64_t rowBegin, uint64_t rowCount) const + { + uint64_t ref = rowBegin + rowCount / 2; + if (ref >= curT) { + ref = curT - 1; + } + return start + ref; + } + + __aicore__ inline void RunExp2(LocalTensor &tensor, uint32_t count) + { + SetFlag(EXP2_EVENT_ID); + WaitFlag(EXP2_EVENT_ID); + ClampExpInput(tensor, count); + Exp(tensor, tensor, count); + PipeBarrier(); + SetFlag(EXP2_EVENT_ID); + WaitFlag(EXP2_EVENT_ID); + } + + __aicore__ inline void ClampExpInput(LocalTensor &tensor, uint32_t count) + { + Mins(tensor, tensor, KDA_EXP_INPUT_MAX, count); + PipeBarrier(); + Maxs(tensor, tensor, KDA_EXP_INPUT_MIN, count); + PipeBarrier(); + } + + __aicore__ inline void ClampScoreExpInput(LocalTensor &tensor, uint32_t count) + { + constexpr float expInputMax = + IsSameType::value ? KDA_SCORE_EXP_INPUT_MAX : KDA_EXP_INPUT_MAX; + constexpr float expInputMin = + IsSameType::value ? KDA_SCORE_EXP_INPUT_MIN : KDA_EXP_INPUT_MIN; + Mins(tensor, tensor, expInputMax, count); + PipeBarrier(); + Maxs(tensor, tensor, expInputMin, count); + PipeBarrier(); + } + + template + __aicore__ inline void ClampFp32ForCast(LocalTensor &tensor, uint32_t count) + { + if constexpr (IsSameType::value) { + Mins(tensor, tensor, KDA_FP16_MAX, count); + PipeBarrier(); + Maxs(tensor, tensor, -KDA_FP16_MAX, count); + PipeBarrier(); + } + } + + __aicore__ inline void ClampFp32ToOutputType(LocalTensor &tensor, uint32_t count) + { + ClampFp32ForCast(tensor, count); + } + + template + __aicore__ inline void CopyVectorIn(LocalTensor &dst, GlobalTensor &src, uint64_t offset, + uint64_t count) + { + uint64_t rowBytes = count * static_cast(sizeof(CopyT)); + if (rowBytes >= 32 && rowBytes % 32 == 0) { + DataCopy(dst, src[offset], static_cast(count)); + return; + } + DataCopyParams params{1, static_cast(rowBytes), 0, 0}; + DataCopyPadParams padParams{false, 0, 0, 0}; + DataCopyPad(dst, src[offset], params, padParams); + } + + template + __aicore__ inline void CopyRowsIn(LocalTensor &dst, GlobalTensor &src, uint64_t offset, uint64_t rows, + uint64_t cols, uint64_t rowStride) + { + if (rows == 0 || cols == 0) { + return; + } + if (rowStride == cols) { + CopyVectorIn(dst, src, offset, rows * cols); + return; + } + DataCopyExtParams params{static_cast(rows), static_cast(cols * sizeof(CopyT)), + static_cast((rowStride - cols) * sizeof(CopyT)), 0, 0}; + DataCopyPadExtParams padParams{false, 0, 0, 0}; + DataCopyPad(dst, src[offset], params, padParams); + } + + template + __aicore__ inline void CopyVectorOut(GlobalTensor &dst, uint64_t offset, LocalTensor &src, + uint64_t count) + { + uint64_t rowBytes = count * static_cast(sizeof(CopyT)); + if (rowBytes >= 32 && rowBytes % 32 == 0) { + DataCopy(dst[offset], src, static_cast(count)); + return; + } + DataCopyParams params{1, static_cast(rowBytes), 0, 0}; + DataCopyPad(dst[offset], src, params); + } + + template + __aicore__ inline void CopyRowIn(LocalTensor &dst, GlobalTensor &src, uint64_t offset) + { + CopyVectorIn(dst, src, offset, K_); + } + + template + __aicore__ inline void CopyRowOut(GlobalTensor &dst, uint64_t offset, LocalTensor &src) + { + CopyVectorOut(dst, offset, src, K_); + } + + __aicore__ inline LocalTensor VecScratch(uint64_t slot) + { + return vecBuf_.Get()[slot * EXP2_UB_ELEMENTS]; + } + + __aicore__ inline uint64_t GateStageElems() const { return GatePipelineRows() * K_; } + + __aicore__ inline uint64_t GatePipelineRows() const + { + constexpr uint64_t fixedBytes = static_cast(KDA_VEC_ARENA_ELEMENTS) * sizeof(float) + EXP2_UB_BYTES; + constexpr uint64_t availableBytes = KDA_AIV_UB_BUDGET_BYTES - fixedBytes; + uint64_t bytesPerRow = K_ * KDA_GATE_PIPELINE_DEPTH * (3 * sizeof(T) + sizeof(GK_T)); + uint64_t rows = bytesPerRow == 0 ? 0 : availableBytes / bytesPerRow; + return rows < KDA_GATE_TILE_ROWS ? rows : KDA_GATE_TILE_ROWS; + } + + __aicore__ inline uint64_t GateInputSlotBytes() const { return GateStageElems() * (2 * sizeof(T) + sizeof(GK_T)); } + + __aicore__ inline LocalTensor GateQTyped(uint64_t slot) + { + uint64_t byteOffset = slot * GateInputSlotBytes(); + return gateWritebackBuf_.Get()[byteOffset / sizeof(T)]; + } + + __aicore__ inline LocalTensor GateKTyped(uint64_t slot) + { + uint64_t byteOffset = slot * GateInputSlotBytes() + GateStageElems() * sizeof(T); + return gateWritebackBuf_.Get()[byteOffset / sizeof(T)]; + } + + __aicore__ inline LocalTensor GateGTyped(uint64_t slot) + { + uint64_t byteOffset = slot * GateInputSlotBytes() + 2 * GateStageElems() * sizeof(T); + return gateWritebackBuf_.Get()[byteOffset / sizeof(GK_T)]; + } + + __aicore__ inline LocalTensor GateKgTyped(uint64_t slot) + { + uint64_t byteOffset = KDA_GATE_PIPELINE_DEPTH * GateInputSlotBytes() + slot * GateStageElems() * sizeof(T); + return gateWritebackBuf_.Get()[byteOffset / sizeof(T)]; + } + + __aicore__ inline void PrefetchQKGate(uint64_t slot, uint64_t b, uint64_t h, uint64_t hv, uint64_t token, + uint64_t elems) + { + const uint64_t rows = elems / K_; + LocalTensor qTyped = GateQTyped(slot); + LocalTensor kTyped = GateKTyped(slot); + LocalTensor gateTyped = GateGTyped(slot); + CopyRowsIn(qTyped, q_, QOffset(b, h, token, 0), rows, K_, inputSequenceMajor_ ? H_ * K_ : K_); + CopyRowsIn(kTyped, k_, QOffset(b, h, token, 0), rows, K_, inputSequenceMajor_ ? H_ * K_ : K_); + CopyVectorIn(gateTyped, gk_, KVOffset(b, hv, token, 0, K_), elems); + SetFlag(mte2ToVEvent_); + } + + __aicore__ inline void PrefetchKGate(uint64_t slot, uint64_t b, uint64_t h, uint64_t hv, uint64_t token, + uint64_t elems) + { + const uint64_t rows = elems / K_; + LocalTensor kTyped = GateQTyped(slot); + LocalTensor gateTyped = GateGTyped(slot); + CopyRowsIn(kTyped, k_, QOffset(b, h, token, 0), rows, K_, inputSequenceMajor_ ? H_ * K_ : K_); + CopyVectorIn(gateTyped, gk_, KVOffset(b, hv, token, 0, K_), elems); + SetFlag(mte2ToVEvent_); + } + + __aicore__ inline void WaitGateInputReady() { WaitFlag(mte2ToVEvent_); } + + __aicore__ inline void WaitGateOutputForMte2(uint64_t slot = 0) + { + WaitFlag(mte3ToMte2Events_[slot]); + } + + __aicore__ inline void WaitGateOutputForVector() { WaitFlag(mte3ToVEvent_); } + + __aicore__ inline void SignalGateOutputDone() + { + SetFlag(mte3ToMte2Events_[0]); + SetFlag(mte3ToVEvent_); + } + + __aicore__ inline void SignalGateOutputDoneForMte2(uint64_t slot) + { + SetFlag(mte3ToMte2Events_[slot]); + } + + template + __aicore__ inline void LoadAsFloatRow(GlobalTensor &src, uint64_t srcOffset, LocalTensor &dst, + uint64_t count) + { + if constexpr (IsSameType::value) { + CopyVectorIn(dst, src, srcOffset, count); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + Adds(dst, dst, 0.0f, static_cast(count)); + PipeBarrier(); + SetFlag(vToMte2Event_); + WaitFlag(vToMte2Event_); + } else { + constexpr uint32_t typedOffset = EXP2_UB_ELEMENTS * sizeof(float) / sizeof(CopyT); + LocalTensor rowLocal = exp2Buf_.Get()[typedOffset]; + CopyVectorIn(rowLocal, src, srcOffset, count); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + Cast(dst, rowLocal, RoundMode::CAST_NONE, static_cast(count)); + PipeBarrier(); + SetFlag(vToMte2Event_); + WaitFlag(vToMte2Event_); + } + PipeBarrier(); + } + + template + __aicore__ inline void LoadAsFloatVector(GlobalTensor &src, uint64_t srcOffset, LocalTensor &dst, + LocalTensor &typedScratch, uint64_t count) + { + if constexpr (IsSameType::value) { + CopyVectorIn(dst, src, srcOffset, count); + } else { + CopyVectorIn(typedScratch, src, srcOffset, count); + } + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + if constexpr (!IsSameType::value) { + Cast(dst, typedScratch, RoundMode::CAST_NONE, static_cast(count)); + PipeBarrier(); + } + } + + template + __aicore__ inline void StoreFloatRow(GlobalTensor &dst, uint64_t dstOffset, LocalTensor &src, + uint64_t count) + { + if constexpr (IsSameType::value) { + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + CopyVectorOut(dst, dstOffset, src, count); + } else { + constexpr uint32_t typedOffset = EXP2_UB_ELEMENTS * sizeof(float) / sizeof(CopyT); + LocalTensor rowLocal = exp2Buf_.Get()[typedOffset]; + Cast(rowLocal, src, RoundMode::CAST_RINT, static_cast(count)); + PipeBarrier(); + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + CopyVectorOut(dst, dstOffset, rowLocal, count); + } + SetFlag(mte3ToMte2Events_[0]); + WaitFlag(mte3ToMte2Events_[0]); + SetFlag(mte3ToVEvent_); + WaitFlag(mte3ToVEvent_); + } + + __aicore__ inline LocalTensor Exp2NegG(uint64_t b, uint64_t hv, uint64_t t) + { + LocalTensor exp2Local = exp2Buf_.Get(); + LoadAsFloatRow(gk_, KVOffset(b, hv, t, 0, K_), exp2Local, K_); + Muls(exp2Local, exp2Local, -LN2, static_cast(K_)); + PipeBarrier(); + RunExp2(exp2Local, static_cast(K_)); + return exp2Local; + } + + __aicore__ inline void PrepareScoreFactorsBulk(uint64_t b, uint64_t h, uint64_t hv, uint64_t start, + uint64_t subBlockIdx, uint64_t subBlockNum, uint64_t refToken, + uint64_t scoreRowBegin, uint64_t scoreRowCount, uint64_t validColEnd, + uint64_t scoreSlot) + { + LocalTensor refFp32 = exp2Buf_.Get(); + LoadAsFloatRow(gk_, KVOffset(b, hv, refToken, 0, K_), refFp32, K_); + + uint64_t qwBegin = scoreRowBegin + (scoreRowCount * subBlockIdx) / subBlockNum; + uint64_t qwEnd = scoreRowBegin + (scoreRowCount * (subBlockIdx + 1)) / subBlockNum; + uint64_t qwMaxRows = GatePipelineRows(); + bool qwOutputPending = false; + uint64_t qwSlot = 0; + if (qwBegin < qwEnd && qwMaxRows > 0) { + uint64_t firstRows = qwEnd - qwBegin; + if (firstRows > qwMaxRows) { + firstRows = qwMaxRows; + } + PrefetchQKGate(qwSlot, b, h, hv, start + qwBegin, firstRows * K_); + } + for (uint64_t tileRow = qwBegin; tileRow < qwEnd && qwMaxRows > 0; tileRow += qwMaxRows) { + uint64_t tileRows = qwEnd - tileRow; + if (tileRows > qwMaxRows) { + tileRows = qwMaxRows; + } + uint64_t elems = tileRows * K_; + LocalTensor qTyped = GateQTyped(qwSlot); + LocalTensor kTyped = GateKTyped(qwSlot); + LocalTensor qScore = qTyped.template ReinterpretCast(); + LocalTensor kScore = kTyped.template ReinterpretCast(); + LocalTensor gateTyped = GateGTyped(qwSlot); + LocalTensor arena = vecBuf_.Get(); + LocalTensor qFp32 = arena; + LocalTensor kFp32 = arena[elems]; + LocalTensor gFp32 = arena[2 * elems]; + LocalTensor expFp32 = arena[3 * elems]; + LocalTensor outFp32 = arena[4 * elems]; + + WaitGateInputReady(); + Cast(qFp32, qTyped, RoundMode::CAST_NONE, static_cast(elems)); + Cast(kFp32, kTyped, RoundMode::CAST_NONE, static_cast(elems)); + if constexpr (IsSameType::value) { + gFp32 = gateTyped; + } else { + Cast(gFp32, gateTyped, RoundMode::CAST_NONE, static_cast(elems)); + } + if (qwOutputPending) { + WaitGateOutputForMte2(); + } + uint64_t nextTileRow = tileRow + qwMaxRows; + if (nextTileRow < qwEnd) { + uint64_t nextRows = qwEnd - nextTileRow; + if (nextRows > qwMaxRows) { + nextRows = qwMaxRows; + } + PrefetchQKGate(qwSlot ^ 1, b, h, hv, start + nextTileRow, nextRows * K_); + } + PipeBarrier(); + for (uint64_t row = 0; row < tileRows; ++row) { + Sub(expFp32[row * K_], gFp32[row * K_], refFp32, static_cast(K_)); + } + PipeBarrier(); + Muls(expFp32, expFp32, LN2, static_cast(elems)); + PipeBarrier(); + ClampScoreExpInput(expFp32, static_cast(elems)); + Exp(expFp32, expFp32, static_cast(elems)); + PipeBarrier(); + + Mul(outFp32, qFp32, expFp32, static_cast(elems)); + PipeBarrier(); + ClampFp32ForCast(outFp32, static_cast(elems)); + Cast(qScore, outFp32, RoundMode::CAST_RINT, static_cast(elems)); + PipeBarrier(); + Mul(outFp32, kFp32, expFp32, static_cast(elems)); + PipeBarrier(); + ClampFp32ForCast(outFp32, static_cast(elems)); + Cast(kScore, outFp32, RoundMode::CAST_RINT, static_cast(elems)); + PipeBarrier(); + + if (qwOutputPending) { + WaitGateOutputForVector(); + } + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + CopyVectorOut(scoreWorkspace_, ScoreScratchOffset(scoreSlot, KDA_SCORE_SCRATCH_QG, tileRow), qScore, elems); + CopyVectorOut(scoreWorkspace_, ScoreScratchOffset(scoreSlot, KDA_SCORE_SCRATCH_W, tileRow), kScore, elems); + SignalGateOutputDone(); + qwOutputPending = true; + qwSlot ^= 1; + } + if (qwOutputPending) { + WaitGateOutputForMte2(); + WaitGateOutputForVector(); + } + + uint64_t kgBegin = (validColEnd * subBlockIdx) / subBlockNum; + uint64_t kgEnd = (validColEnd * (subBlockIdx + 1)) / subBlockNum; + uint64_t kgMaxRows = GatePipelineRows(); + bool kgOutputPending = false; + uint64_t kgSlot = 0; + if (kgBegin < kgEnd && kgMaxRows > 0) { + uint64_t firstRows = kgEnd - kgBegin; + if (firstRows > kgMaxRows) { + firstRows = kgMaxRows; + } + PrefetchKGate(kgSlot, b, h, hv, start + kgBegin, firstRows * K_); + } + for (uint64_t tileRow = kgBegin; tileRow < kgEnd && kgMaxRows > 0; tileRow += kgMaxRows) { + uint64_t tileRows = kgEnd - tileRow; + if (tileRows > kgMaxRows) { + tileRows = kgMaxRows; + } + uint64_t elems = tileRows * K_; + LocalTensor kTyped = GateQTyped(kgSlot); + LocalTensor kgScore = kTyped.template ReinterpretCast(); + LocalTensor gateTyped = GateGTyped(kgSlot); + LocalTensor arena = vecBuf_.Get(); + LocalTensor kFp32 = arena; + LocalTensor gFp32 = arena[elems]; + LocalTensor expFp32 = arena[2 * elems]; + LocalTensor outFp32 = arena[3 * elems]; + + WaitGateInputReady(); + Cast(kFp32, kTyped, RoundMode::CAST_NONE, static_cast(elems)); + if constexpr (IsSameType::value) { + gFp32 = gateTyped; + } else { + Cast(gFp32, gateTyped, RoundMode::CAST_NONE, static_cast(elems)); + } + if (kgOutputPending) { + WaitGateOutputForMte2(); + } + uint64_t nextTileRow = tileRow + kgMaxRows; + if (nextTileRow < kgEnd) { + uint64_t nextRows = kgEnd - nextTileRow; + if (nextRows > kgMaxRows) { + nextRows = kgMaxRows; + } + PrefetchKGate(kgSlot ^ 1, b, h, hv, start + nextTileRow, nextRows * K_); + } + PipeBarrier(); + for (uint64_t row = 0; row < tileRows; ++row) { + Sub(expFp32[row * K_], refFp32, gFp32[row * K_], static_cast(K_)); + } + PipeBarrier(); + Muls(expFp32, expFp32, LN2, static_cast(elems)); + PipeBarrier(); + ClampScoreExpInput(expFp32, static_cast(elems)); + Exp(expFp32, expFp32, static_cast(elems)); + PipeBarrier(); + Mul(outFp32, kFp32, expFp32, static_cast(elems)); + PipeBarrier(); + ClampFp32ForCast(outFp32, static_cast(elems)); + Cast(kgScore, outFp32, RoundMode::CAST_RINT, static_cast(elems)); + PipeBarrier(); + + if (kgOutputPending) { + WaitGateOutputForVector(); + } + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + CopyVectorOut(scoreWorkspace_, ScoreScratchOffset(scoreSlot, KDA_SCORE_SCRATCH_KG, tileRow), kgScore, + elems); + SignalGateOutputDone(); + kgOutputPending = true; + kgSlot ^= 1; + } + if (kgOutputPending) { + WaitGateOutputForMte2(); + WaitGateOutputForVector(); + } + } + + __aicore__ inline void PrepareGateProductsBulk(uint64_t b, uint64_t h, uint64_t hv, uint64_t start, uint64_t curT, + uint64_t subBlockIdx, uint64_t subBlockNum, bool useRef, + uint64_t refToken, uint64_t validColEnd, bool writeScoreScratch, + uint64_t scoreSlot) + { + if constexpr (IsSameType::value) { + return; + } + if (subBlockNum == 0 || subBlockIdx >= subBlockNum || K_ == 0) { + return; + } + uint64_t rowBegin = (curT * subBlockIdx) / subBlockNum; + uint64_t rowEnd = (curT * (subBlockIdx + 1)) / subBlockNum; + if (rowBegin >= rowEnd) { + return; + } + + uint64_t maxRows = GatePipelineRows(); + if (maxRows == 0) { + return; + } + LocalTensor refFp32 = exp2Buf_.Get(); + if (useRef) { + LoadAsFloatRow(gk_, KVOffset(b, hv, refToken, 0, K_), refFp32, K_); + } + + bool outputPending = false; + uint64_t gateSlot = 0; + uint64_t firstRows = rowEnd - rowBegin; + if (firstRows > maxRows) { + firstRows = maxRows; + } + PrefetchQKGate(gateSlot, b, h, hv, start + rowBegin, firstRows * K_); + for (uint64_t tileRow = rowBegin; tileRow < rowEnd; tileRow += maxRows) { + uint64_t tileRows = rowEnd - tileRow; + if (tileRows > maxRows) { + tileRows = maxRows; + } + uint64_t elems = tileRows * K_; + LocalTensor qTyped = GateQTyped(gateSlot); + LocalTensor kTyped = GateKTyped(gateSlot); + LocalTensor kgTyped = GateKgTyped(gateSlot); + LocalTensor qScore = qTyped.template ReinterpretCast(); + LocalTensor wScore = kTyped.template ReinterpretCast(); + LocalTensor kgScore = kgTyped.template ReinterpretCast(); + LocalTensor gateTyped = GateGTyped(gateSlot); + LocalTensor arena = vecBuf_.Get(); + LocalTensor qFp32 = arena; + LocalTensor kFp32 = arena[elems]; + LocalTensor gFp32 = arena[2 * elems]; + LocalTensor expFp32 = arena[3 * elems]; + LocalTensor outFp32 = arena[4 * elems]; + + uint64_t token = start + tileRow; + WaitGateInputReady(); + Cast(qFp32, qTyped, RoundMode::CAST_NONE, static_cast(elems)); + Cast(kFp32, kTyped, RoundMode::CAST_NONE, static_cast(elems)); + if constexpr (IsSameType::value) { + gFp32 = gateTyped; + } else { + Cast(gFp32, gateTyped, RoundMode::CAST_NONE, static_cast(elems)); + } + uint64_t nextTileRow = tileRow + maxRows; + if (outputPending) { + WaitGateOutputForMte2(); + } + if (nextTileRow < rowEnd) { + uint64_t nextRows = rowEnd - nextTileRow; + if (nextRows > maxRows) { + nextRows = maxRows; + } + PrefetchQKGate(gateSlot ^ 1, b, h, hv, start + nextTileRow, nextRows * K_); + } + PipeBarrier(); + + if (useRef) { + for (uint64_t row = 0; row < tileRows; ++row) { + Sub(expFp32[row * K_], gFp32[row * K_], refFp32, static_cast(K_)); + } + } else { + Adds(expFp32, gFp32, 0.0f, static_cast(elems)); + } + PipeBarrier(); + Muls(expFp32, expFp32, LN2, static_cast(elems)); + PipeBarrier(); + if (writeScoreScratch) { + ClampScoreExpInput(expFp32, static_cast(elems)); + } else { + ClampExpInput(expFp32, static_cast(elems)); + } + Exp(expFp32, expFp32, static_cast(elems)); + PipeBarrier(); + + Mul(outFp32, qFp32, expFp32, static_cast(elems)); + PipeBarrier(); + if (writeScoreScratch) { + ClampFp32ForCast(outFp32, static_cast(elems)); + Cast(qScore, outFp32, RoundMode::CAST_RINT, static_cast(elems)); + } else { + ClampFp32ToOutputType(outFp32, static_cast(elems)); + Cast(qTyped, outFp32, RoundMode::CAST_RINT, static_cast(elems)); + } + PipeBarrier(); + + Mul(outFp32, kFp32, expFp32, static_cast(elems)); + PipeBarrier(); + if (writeScoreScratch) { + ClampFp32ForCast(outFp32, static_cast(elems)); + Cast(wScore, outFp32, RoundMode::CAST_RINT, static_cast(elems)); + } else { + ClampFp32ToOutputType(outFp32, static_cast(elems)); + Cast(kTyped, outFp32, RoundMode::CAST_RINT, static_cast(elems)); + } + PipeBarrier(); + + if (useRef) { + for (uint64_t row = 0; row < tileRows; ++row) { + Sub(expFp32[row * K_], refFp32, gFp32[row * K_], static_cast(K_)); + } + } else { + Muls(expFp32, gFp32, -1.0f, static_cast(elems)); + } + PipeBarrier(); + Muls(expFp32, expFp32, LN2, static_cast(elems)); + PipeBarrier(); + if (writeScoreScratch) { + ClampScoreExpInput(expFp32, static_cast(elems)); + } else { + ClampExpInput(expFp32, static_cast(elems)); + } + Exp(expFp32, expFp32, static_cast(elems)); + PipeBarrier(); + Mul(outFp32, kFp32, expFp32, static_cast(elems)); + PipeBarrier(); + if (useRef && tileRow + tileRows > validColEnd) { + for (uint64_t row = 0; row < tileRows; ++row) { + if (tileRow + row >= validColEnd) { + Duplicate(outFp32[row * K_], 0.0f, static_cast(K_)); + } + } + PipeBarrier(); + } + if (writeScoreScratch) { + ClampFp32ForCast(outFp32, static_cast(elems)); + } else { + ClampFp32ToOutputType(outFp32, static_cast(elems)); + } + if (outputPending) { + WaitGateOutputForVector(); + } + if (writeScoreScratch) { + Cast(kgScore, outFp32, RoundMode::CAST_RINT, static_cast(elems)); + } else { + Cast(kgTyped, outFp32, RoundMode::CAST_RINT, static_cast(elems)); + } + PipeBarrier(); + + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + if (writeScoreScratch) { + CopyVectorOut(scoreWorkspace_, ScoreScratchOffset(scoreSlot, KDA_SCORE_SCRATCH_QG, tileRow), qScore, + elems); + CopyVectorOut(scoreWorkspace_, ScoreScratchOffset(scoreSlot, KDA_SCORE_SCRATCH_W, tileRow), wScore, + elems); + CopyVectorOut(scoreWorkspace_, ScoreScratchOffset(scoreSlot, KDA_SCORE_SCRATCH_KG, tileRow), kgScore, + elems); + } else { + CopyVectorOut(qg_, KVOffset(b, hv, token, 0, K_), qTyped, elems); + CopyVectorOut(w_, KVOffset(b, hv, token, 0, K_), kTyped, elems); + CopyVectorOut(kg_, KVOffset(b, hv, token, 0, K_), kgTyped, elems); + } + SignalGateOutputDone(); + outputPending = true; + gateSlot ^= 1; + } + if (outputPending) { + WaitGateOutputForMte2(); + WaitGateOutputForVector(); + } + return; + } + + __aicore__ inline void PrepareGateProducts(uint64_t b, uint64_t h, uint64_t hv, uint64_t start, uint64_t curT, + uint64_t subBlockIdx, uint64_t subBlockNum, bool useRef = false, + uint64_t refToken = 0, uint64_t validColEnd = 0, + bool writeScoreScratch = false, uint64_t scoreSlot = 0, + uint64_t scoreRowBegin = 0, uint64_t scoreRowCount = 0) + { + if (subBlockNum == 0 || subBlockIdx >= subBlockNum) { + return; + } + if (validColEnd == 0 || validColEnd > curT) { + validColEnd = curT; + } + if (writeScoreScratch) { + PrepareScoreFactorsBulk(b, h, hv, start, subBlockIdx, subBlockNum, refToken, scoreRowBegin, scoreRowCount, + validColEnd, scoreSlot); + return; + } + PrepareGateProductsBulk(b, h, hv, start, curT, subBlockIdx, subBlockNum, useRef, refToken, validColEnd, + writeScoreScratch, scoreSlot); + } + + __aicore__ inline void ComputeRawAqkAkkCube(uint64_t b, uint64_t hv, uint64_t start, uint64_t curT) + { + ComputeRawAqkAkkCubeBlock(b, hv, start, curT, 0, curT); + } + + __aicore__ inline void ComputeRawAqkAkkCubeBlock(uint64_t b, uint64_t hv, uint64_t start, uint64_t curT, + uint64_t rowBegin, uint64_t rowCount, + bool readScoreScratch = false, uint64_t scoreSlot = 0, + uint64_t colCount = 0) + { + using ElementA = SCORE_T; + using ElementB = SCORE_T; + using ElementC = float; + using LayoutTagA = Catlass::layout::RowMajor; + using LayoutTagB = Catlass::layout::ColumnMajor; + using LayoutTagC = Catlass::layout::RowMajor; + using TileCopy = Catlass::Gemm::Tile::PackedTileCopyTla; + using BlockMmad = Catlass::Gemm::Block::BlockMmadTla; + + Catlass::Arch::Resource resource; + BlockMmad blockMmad(resource); + auto layoutA = tla::MakeLayout(BT_, K_); + auto layoutB = tla::MakeLayout(K_, BT_); + auto layoutC = tla::MakeLayout(BT_, BT_); + if (colCount == 0 || colCount > curT) { + colCount = curT; + } + Catlass::GemmCoord shape{static_cast(rowCount), static_cast(colCount), + static_cast(K_)}; + + (void)readScoreScratch; + auto tensorQPos = tla::MakeTensor(scoreWorkspace_[ScoreScratchOffset(scoreSlot, KDA_SCORE_SCRATCH_QG)], layoutA, + Catlass::Arch::PositionGM{}); + auto tensorKPos = tla::MakeTensor(scoreWorkspace_[ScoreScratchOffset(scoreSlot, KDA_SCORE_SCRATCH_W)], layoutA, + Catlass::Arch::PositionGM{}); + auto tensorKNeg = tla::MakeTensor(scoreWorkspace_[ScoreScratchOffset(scoreSlot, KDA_SCORE_SCRATCH_KG)], layoutB, + Catlass::Arch::PositionGM{}); + auto tensorAqk = tla::MakeTensor(aqk_[AOffset(b, hv, start, 0)], layoutC, Catlass::Arch::PositionGM{}); + auto tensorAkk = tla::MakeTensor(akk_[AOffset(b, hv, start, 0)], layoutC, Catlass::Arch::PositionGM{}); + + auto blockQPos = GetTile(tensorQPos, tla::MakeCoord(rowBegin, 0), tla::MakeShape(shape.m(), shape.k())); + auto blockKPos = GetTile(tensorKPos, tla::MakeCoord(rowBegin, 0), tla::MakeShape(shape.m(), shape.k())); + auto blockKNeg = GetTile(tensorKNeg, tla::MakeCoord(0, 0), tla::MakeShape(shape.k(), shape.n())); + auto blockAqk = GetTile(tensorAqk, tla::MakeCoord(rowBegin, 0), tla::MakeShape(shape.m(), shape.n())); + auto blockAkk = GetTile(tensorAkk, tla::MakeCoord(rowBegin, 0), tla::MakeShape(shape.m(), shape.n())); + + blockMmad.preSetFlags(); + blockMmad(blockQPos, blockKNeg, blockAqk, shape); + blockMmad(blockKPos, blockKNeg, blockAkk, shape); + blockMmad.finalWaitFlags(); + } + + __aicore__ inline bool UseAkkCubeSolve(uint64_t curT) const + { + return curT > 0 && curT <= BT_ && (BT_ == 64 || BT_ == 128) && K_ >= 16 && V_ >= 16 && V_ <= 256 && + K_ % 16 == 0 && V_ % 16 == 0; + } + + __aicore__ inline bool UsePostWuCube(uint64_t curT) const + { + return curT > 0 && curT <= BT_ && (BT_ == 64 || BT_ == 128) && K_ >= 16 && V_ >= 16 && V_ <= 256 && + K_ % 16 == 0 && V_ % 16 == 0; + } + + __aicore__ inline void CopyLocalFloat(LocalTensor dst, LocalTensor src, uint64_t count) + { + if (count == 0) { + return; + } + Adds(dst, src, 0.0f, static_cast(count)); + PipeBarrier(); + } + + __aicore__ inline void FillLocalFloat(LocalTensor dst, float value, uint64_t count) + { + if (count == 0) { + return; + } + Duplicate(dst, value, static_cast(count)); + PipeBarrier(); + } + + __aicore__ inline void ForwardSubDiag16(LocalTensor diag, LocalTensor row, LocalTensor prod, + LocalTensor rowBrcb, LocalTensor reduced, uint64_t valid) + { + constexpr uint32_t brcbStride = 8; + constexpr uint32_t diagSize = KDA_SOLVE_DIAG_BT; + constexpr uint8_t rowBlk = diagSize * sizeof(float) / 32; + + for (uint64_t i = 2; i < valid; ++i) { + uint32_t rowOffset = static_cast(i * diagSize); + DataCopy(row, diag[rowOffset], diagSize); + PipeBarrier(); + + Brcb(rowBrcb, row, diagSize / brcbStride, {1, 8}); + PipeBarrier(); + for (uint32_t col = 0; col < diagSize; col += brcbStride) { + Mul(prod[col], diag[col], rowBrcb, brcbStride, static_cast(diagSize), + {1, 1, 0, rowBlk, rowBlk, 1}); + } + PipeBarrier(); + + uint32_t remain = diagSize; + while (remain > 1) { + uint32_t calcCount = (remain / 2) * diagSize; + remain = (remain + 1) / 2; + Add(prod, prod, prod[remain * diagSize], calcCount); + PipeBarrier(); + } + DataCopy(reduced, prod, diagSize); + PipeBarrier(); + Add(row, row, reduced, diagSize); + PipeBarrier(); + DataCopy(diag[rowOffset], row, diagSize); + PipeBarrier(); + } + + SetFlag(EXP2_EVENT_ID); + WaitFlag(EXP2_EVENT_ID); + for (uint32_t i = 0; i < diagSize; ++i) { + uint32_t diagOffset = i * diagSize + i; + if (i < valid) { + diag.SetValue(diagOffset, diag.GetValue(diagOffset) + 1.0f); + } else { + diag.SetValue(diagOffset, 1.0f); + } + } + SetFlag(EXP2_EVENT_ID); + WaitFlag(EXP2_EVENT_ID); + } + + __aicore__ inline void SolveDiagonalBlocksInRows(LocalTensor akkMat, LocalTensor xMat, + LocalTensor arena, uint64_t scratchBase, uint64_t curT, + uint64_t rowBegin, uint64_t rowCount) + { + constexpr uint32_t diagSize = KDA_SOLVE_DIAG_BT; + constexpr uint32_t diagElements = diagSize * diagSize; + constexpr uint32_t brcbElements = diagSize * 8; + + LocalTensor diag = arena[scratchBase]; + LocalTensor row = diag[diagElements]; + LocalTensor prod = row[diagSize]; + LocalTensor rowBrcb = prod[diagElements]; + LocalTensor reduced = rowBrcb[brcbElements]; + + uint64_t rowEnd = rowBegin + rowCount; + for (uint64_t blockBegin = 0; blockBegin < BT_; blockBegin += diagSize) { + if (blockBegin < rowBegin || blockBegin + diagSize > rowEnd) { + continue; + } + Duplicate(diag, 0.0f, diagElements); + PipeBarrier(); + + uint64_t localBlockRow = blockBegin - rowBegin; + uint64_t valid = blockBegin < curT ? curT - blockBegin : 0; + if (valid > diagSize) { + valid = diagSize; + } + for (uint32_t rowIdx = 0; rowIdx < diagSize; ++rowIdx) { + uint64_t srcOffset = (localBlockRow + rowIdx) * BT_ + blockBegin; + Muls(diag[rowIdx * diagSize], akkMat[srcOffset], -1.0f, diagSize); + } + PipeBarrier(); + + ForwardSubDiag16(diag, row, prod, rowBrcb, reduced, valid); + for (uint32_t rowIdx = 0; rowIdx < diagSize; ++rowIdx) { + uint64_t dstOffset = (localBlockRow + rowIdx) * BT_ + blockBegin; + Adds(xMat[dstOffset], diag[rowIdx * diagSize], 0.0f, diagSize); + } + PipeBarrier(); + } + } + + __aicore__ inline void BuildPrefixMask(LocalTensor dst, uint64_t prefix, uint64_t count) + { + if (prefix > count) { + prefix = count; + } + Duplicate(dst, 0.0f, static_cast(count)); + if (prefix > 0) { + Duplicate(dst, 1.0f, static_cast(prefix)); + } + PipeBarrier(); + } + + __aicore__ inline uint64_t BuildCausalMask(uint64_t threshold, uint64_t colBegin) const + { + if (threshold <= colBegin) { + return ~0ULL; + } + if (threshold >= colBegin + KDA_SOLVE_BT) { + return 0ULL; + } + return ~0ULL << (threshold - colBegin); + } + + __aicore__ inline void BuildCausalSelectMasks(LocalTensor aqkMask, LocalTensor akkMask, + uint64_t rowBegin, uint64_t rowCount, uint64_t colBegin) + { + __ubuf__ uint64_t *aqkMaskPtr = reinterpret_cast<__ubuf__ uint64_t *>(aqkMask.GetPhyAddr()); + __ubuf__ uint64_t *akkMaskPtr = reinterpret_cast<__ubuf__ uint64_t *>(akkMask.GetPhyAddr()); + for (uint32_t localRow = 0; localRow < rowCount; ++localRow) { + uint32_t row = static_cast(rowBegin + localRow); + aqkMaskPtr[localRow] = BuildCausalMask(static_cast(row) + 1, colBegin); + akkMaskPtr[localRow] = BuildCausalMask(static_cast(row), colBegin); + } + } + + __aicore__ inline void SelectCausalRows(LocalTensor aqkMat, LocalTensor akkMat, uint64_t rowBegin, + uint64_t rowCount) + { + LocalTensor aqkMask = vecBuf_.Get()[KDA_SELECT_AQK_MASK_BYTE_OFFSET]; + LocalTensor akkMask = vecBuf_.Get()[KDA_SELECT_AKK_MASK_BYTE_OFFSET]; + LocalTensor zeroLocal = vecBuf_.Get()[KDA_SELECT_ZERO_FLOAT_OFFSET]; + Duplicate(zeroLocal, 0.0f, 8); + PipeBarrier(); + + uint64_t colBlockCount = (BT_ + KDA_SOLVE_BT - 1) / KDA_SOLVE_BT; + for (uint64_t colBlock = 0; colBlock < colBlockCount; ++colBlock) { + uint64_t maskOffset = colBlock * KDA_SELECT_COL_MASK_BYTES; + uint64_t colBegin = colBlock * KDA_SOLVE_BT; + BuildCausalSelectMasks(aqkMask[maskOffset], akkMask[maskOffset], rowBegin, rowCount, colBegin); + } + SetFlag(EXP2_EVENT_ID); + WaitFlag(EXP2_EVENT_ID); + + uint8_t rowStride = static_cast(BT_ * sizeof(float) / 32); + BinaryRepeatParams repeatParams = {1, 0, 1, rowStride, 0, rowStride}; + for (uint64_t colBlock = 0; colBlock < colBlockCount; ++colBlock) { + uint64_t maskOffset = colBlock * KDA_SELECT_COL_MASK_BYTES; + uint64_t colBegin = colBlock * KDA_SOLVE_BT; + Select(aqkMat[colBegin], aqkMask[maskOffset], zeroLocal, aqkMat[colBegin], SELMODE::VSEL_TENSOR_TENSOR_MODE, + KDA_SOLVE_BT, static_cast(rowCount), repeatParams); + Select(akkMat[colBegin], akkMask[maskOffset], zeroLocal, akkMat[colBegin], SELMODE::VSEL_TENSOR_TENSOR_MODE, + KDA_SOLVE_BT, static_cast(rowCount), repeatParams); + } + PipeBarrier(); + SetFlag(EXP2_EVENT_ID); + WaitFlag(EXP2_EVENT_ID); + } + + __aicore__ inline void PrepareAqkAkkSolveInput64(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start) + { + LocalTensor arena = vecBuf_.Get(); + LocalTensor aqkMat = arena; + LocalTensor akkMat = arena[KDA_SOLVE_MATRIX_ELEMENTS]; + LocalTensor xMat = arena[2 * KDA_SOLVE_MATRIX_ELEMENTS]; + LocalTensor betaLocal = arena[3 * KDA_SOLVE_MATRIX_ELEMENTS]; + LocalTensor betaBrcb = arena[3 * KDA_SOLVE_MATRIX_ELEMENTS + KDA_SOLVE_BT]; + LocalTensor maskLocal = arena[3 * KDA_SOLVE_MATRIX_ELEMENTS + KDA_SOLVE_BT + 512]; + LocalTensor oneHotLocal = arena[3 * KDA_SOLVE_MATRIX_ELEMENTS + KDA_SOLVE_BT + 512 + KDA_SOLVE_BT]; + + LoadAsFloatRow(beta_, BetaOffset(b, hv, start), betaLocal, KDA_SOLVE_BT); + Brcb(betaBrcb, betaLocal, 8, {1, 8}); + PipeBarrier(); + + DataCopy(aqkMat, aqk_[AOffset(b, hv, start, 0)], KDA_SOLVE_MATRIX_ELEMENTS); + DataCopy(akkMat, akk_[AOffset(b, hv, start, 0)], KDA_SOLVE_MATRIX_ELEMENTS); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + + for (uint64_t col = 0; col < KDA_SOLVE_BT; col += 8) { + Mul(akkMat[col], akkMat[col], betaBrcb, 8, KDA_SOLVE_BT, {1, 1, 1, 8, 8, 1}); + PipeBarrier(); + } + SelectCausalRows(aqkMat, akkMat, 0, KDA_SOLVE_BT); + + Muls(xMat, akkMat, -1.0f, KDA_SOLVE_MATRIX_ELEMENTS); + PipeBarrier(); + for (uint64_t row = 0; row < KDA_SOLVE_BT; ++row) { + BuildPrefixMask(maskLocal, row + 1, KDA_SOLVE_BT); + BuildPrefixMask(oneHotLocal, row, KDA_SOLVE_BT); + Sub(maskLocal, maskLocal, oneHotLocal, KDA_SOLVE_BT); + PipeBarrier(); + Add(xMat[row * KDA_SOLVE_BT], xMat[row * KDA_SOLVE_BT], maskLocal, KDA_SOLVE_BT); + PipeBarrier(); + } + + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + DataCopy(aqk_[AOffset(b, hv, start, 0)], aqkMat, KDA_SOLVE_MATRIX_ELEMENTS); + DataCopy(akk_[AOffset(b, hv, start, 0)], akkMat, KDA_SOLVE_MATRIX_ELEMENTS); + DataCopy(h_[SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_X)], xMat, KDA_SOLVE_MATRIX_ELEMENTS); + SetFlag(mte3ToMte2Events_[0]); + WaitFlag(mte3ToMte2Events_[0]); + SetFlag(mte3ToVEvent_); + WaitFlag(mte3ToVEvent_); + } + + __aicore__ inline void PrepareAqkAkkSolveInputTail(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, + uint64_t curT) + { + uint64_t elemCount = curT * KDA_SOLVE_BT; + LocalTensor arena = vecBuf_.Get(); + LocalTensor aqkMat = arena; + LocalTensor akkMat = arena[KDA_SOLVE_MATRIX_ELEMENTS]; + LocalTensor xMat = arena[2 * KDA_SOLVE_MATRIX_ELEMENTS]; + LocalTensor betaLocal = arena[3 * KDA_SOLVE_MATRIX_ELEMENTS]; + LocalTensor betaBrcb = arena[3 * KDA_SOLVE_MATRIX_ELEMENTS + KDA_SOLVE_BT]; + LocalTensor maskLocal = arena[3 * KDA_SOLVE_MATRIX_ELEMENTS + KDA_SOLVE_BT + 512]; + LocalTensor oneHotLocal = arena[3 * KDA_SOLVE_MATRIX_ELEMENTS + KDA_SOLVE_BT + 512 + KDA_SOLVE_BT]; + + FillLocalFloat(betaLocal, 0.0f, KDA_SOLVE_BT); + SetFlag(vToMte2Event_); + WaitFlag(vToMte2Event_); + LoadAsFloatRow(beta_, BetaOffset(b, hv, start), betaLocal, curT); + Brcb(betaBrcb, betaLocal, 8, {1, 8}); + PipeBarrier(); + + DataCopy(aqkMat, aqk_[AOffset(b, hv, start, 0)], static_cast(elemCount)); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + if (elemCount < KDA_SOLVE_MATRIX_ELEMENTS) { + FillLocalFloat(aqkMat[elemCount], 0.0f, KDA_SOLVE_MATRIX_ELEMENTS - elemCount); + } + DataCopy(akkMat, akk_[AOffset(b, hv, start, 0)], static_cast(elemCount)); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + if (elemCount < KDA_SOLVE_MATRIX_ELEMENTS) { + FillLocalFloat(akkMat[elemCount], 0.0f, KDA_SOLVE_MATRIX_ELEMENTS - elemCount); + } + + for (uint64_t col = 0; col < KDA_SOLVE_BT; col += 8) { + Mul(akkMat[col], akkMat[col], betaBrcb, 8, KDA_SOLVE_BT, {1, 1, 1, 8, 8, 1}); + PipeBarrier(); + } + SelectCausalRows(aqkMat, akkMat, 0, KDA_SOLVE_BT); + + Muls(xMat, akkMat, -1.0f, KDA_SOLVE_MATRIX_ELEMENTS); + PipeBarrier(); + for (uint64_t row = 0; row < KDA_SOLVE_BT; ++row) { + BuildPrefixMask(maskLocal, row + 1, KDA_SOLVE_BT); + BuildPrefixMask(oneHotLocal, row, KDA_SOLVE_BT); + Sub(maskLocal, maskLocal, oneHotLocal, KDA_SOLVE_BT); + PipeBarrier(); + Add(xMat[row * KDA_SOLVE_BT], xMat[row * KDA_SOLVE_BT], maskLocal, KDA_SOLVE_BT); + PipeBarrier(); + } + + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + DataCopy(aqk_[AOffset(b, hv, start, 0)], aqkMat, static_cast(elemCount)); + DataCopy(h_[SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_X)], xMat, KDA_SOLVE_MATRIX_ELEMENTS); + DataCopy(h_[SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_Y0)], akkMat, KDA_SOLVE_MATRIX_ELEMENTS); + SetFlag(mte3ToMte2Events_[0]); + WaitFlag(mte3ToMte2Events_[0]); + SetFlag(mte3ToVEvent_); + WaitFlag(mte3ToVEvent_); + } + + __aicore__ inline void GetSolveRowRange(uint64_t curT, uint64_t subBlockIdx, uint64_t subBlockNum, + uint64_t &rowBegin, uint64_t &rowEnd) const + { + if (subBlockNum == 0 || subBlockIdx >= subBlockNum) { + rowBegin = 0; + rowEnd = 0; + return; + } + rowBegin = (curT * subBlockIdx) / subBlockNum; + rowEnd = (curT * (subBlockIdx + 1)) / subBlockNum; + } + + __aicore__ inline void PrepareAqkAkkSolveInputRows(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, + uint64_t curT, uint64_t rowBegin, uint64_t rowEnd, + bool storeLToAkk, bool storeLToScratch) + { + uint64_t rowCount = rowEnd - rowBegin; + if (rowCount == 0) { + return; + } + uint64_t validRowCount = rowBegin < curT ? curT - rowBegin : 0; + if (validRowCount > rowCount) { + validRowCount = rowCount; + } + uint64_t elemCount = rowCount * BT_; + uint64_t validElemCount = validRowCount * BT_; + LocalTensor arena = vecBuf_.Get(); + LocalTensor aqkMat = arena; + LocalTensor akkMat = arena[elemCount]; + LocalTensor xMat = arena[2 * elemCount]; + LocalTensor betaLocal = arena[3 * elemCount]; + LocalTensor betaBrcb = arena[3 * elemCount + BT_]; + LocalTensor maskLocal = arena[3 * elemCount + BT_ + 512]; + LocalTensor oneHotLocal = arena[3 * elemCount + BT_ + 512 + BT_]; + + uint64_t token = start + rowBegin; + + if (validRowCount < rowCount) { + FillLocalFloat(aqkMat, 0.0f, elemCount); + FillLocalFloat(akkMat, 0.0f, elemCount); + FillLocalFloat(betaLocal, 0.0f, rowCount); + } + SetFlag(vToMte2Event_); + WaitFlag(vToMte2Event_); + if (validRowCount > 0) { + LoadAsFloatRow(beta_, BetaOffset(b, hv, token), betaLocal, validRowCount); + DataCopy(aqkMat, aqk_[AOffset(b, hv, token, 0)], static_cast(validElemCount)); + DataCopy(akkMat, akk_[AOffset(b, hv, token, 0)], static_cast(validElemCount)); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + } + Brcb(betaBrcb, betaLocal, static_cast((rowCount + 7) / 8), {1, 8}); + PipeBarrier(); + uint8_t rowStride = static_cast(BT_ * sizeof(float) / 32); + for (uint64_t col = 0; col < BT_; col += 8) { + Mul(akkMat[col], akkMat[col], betaBrcb, 8, static_cast(rowCount), + {1, 1, 0, rowStride, rowStride, 1}); + } + PipeBarrier(); + if (validRowCount > 0) { + SelectCausalRows(aqkMat, akkMat, rowBegin, validRowCount); + } + + Muls(xMat, akkMat, -1.0f, static_cast(elemCount)); + PipeBarrier(); + if constexpr (SAFE_GATE) { + uint64_t scratchBase = 3 * elemCount + BT_ + 512 + 2 * BT_; + SolveDiagonalBlocksInRows(akkMat, xMat, arena, scratchBase, curT, rowBegin, rowCount); + } else { + for (uint64_t localRow = 0; localRow < rowCount; ++localRow) { + uint64_t row = rowBegin + localRow; + BuildPrefixMask(maskLocal, row + 1, BT_); + BuildPrefixMask(oneHotLocal, row, BT_); + Sub(maskLocal, maskLocal, oneHotLocal, static_cast(BT_)); + PipeBarrier(); + Add(xMat[localRow * BT_], xMat[localRow * BT_], maskLocal, static_cast(BT_)); + PipeBarrier(); + } + } + + uint64_t xBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_X) + rowBegin * BT_; + uint64_t lBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_Y0) + rowBegin * BT_; + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + if (validRowCount > 0) { + DataCopy(aqk_[AOffset(b, hv, token, 0)], aqkMat, static_cast(validElemCount)); + if (storeLToAkk) { + DataCopy(akk_[AOffset(b, hv, token, 0)], akkMat, static_cast(validElemCount)); + } + } + DataCopy(solveWorkspace_[xBase], xMat, static_cast(elemCount)); + if (storeLToScratch) { + DataCopy(solveWorkspace_[lBase], akkMat, static_cast(elemCount)); + } + SetFlag(mte3ToMte2Events_[0]); + WaitFlag(mte3ToMte2Events_[0]); + SetFlag(mte3ToVEvent_); + WaitFlag(mte3ToVEvent_); + } + + __aicore__ inline void CubeGemmSolveSub(GlobalTensor &tensorA, uint64_t baseA, uint64_t rowA, uint64_t colA, + GlobalTensor &tensorB, uint64_t baseB, uint64_t rowB, uint64_t colB, + GlobalTensor &tensorC, uint64_t baseC, uint64_t rowC, uint64_t colC, + uint32_t m, uint32_t n, uint32_t k) + { + using ElementA = float; + using ElementB = float; + using ElementC = float; + using LayoutTagA = Catlass::layout::RowMajor; + using LayoutTagB = Catlass::layout::RowMajor; + using LayoutTagC = Catlass::layout::RowMajor; + using TileCopy = Catlass::Gemm::Tile::PackedTileCopyTla; + using BlockMmad = + Catlass::Gemm::Block::BlockMmadTla; + Catlass::Arch::Resource resource; + auto layoutA = tla::MakeLayout(BT_, BT_); + auto layoutB = tla::MakeLayout(BT_, BT_); + auto layoutC = tla::MakeLayout(BT_, BT_); + auto tensorLayoutA = tla::MakeTensor(tensorA[baseA], layoutA, Catlass::Arch::PositionGM{}); + auto tensorLayoutB = tla::MakeTensor(tensorB[baseB], layoutB, Catlass::Arch::PositionGM{}); + auto tensorLayoutC = tla::MakeTensor(tensorC[baseC], layoutC, Catlass::Arch::PositionGM{}); + Catlass::GemmCoord shape{m, n, k}; + auto blockA = GetTile(tensorLayoutA, tla::MakeCoord(rowA, colA), tla::MakeShape(shape.m(), shape.k())); + auto blockB = GetTile(tensorLayoutB, tla::MakeCoord(rowB, colB), tla::MakeShape(shape.k(), shape.n())); + auto blockC = GetTile(tensorLayoutC, tla::MakeCoord(rowC, colC), tla::MakeShape(shape.m(), shape.n())); + BlockMmad blockMmad(resource); + blockMmad(blockA, blockB, blockC, shape); + PipeBarrier(); + } + + __aicore__ inline void AddSolveTmpToX(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, bool storeAkk) + { + LocalTensor arena = vecBuf_.Get(); + LocalTensor xLocal = arena; + LocalTensor tmpLocal = arena[KDA_SOLVE_MATRIX_ELEMENTS]; + uint64_t xBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_X); + uint64_t tmpBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_TMP); + + DataCopy(xLocal, h_[xBase], KDA_SOLVE_MATRIX_ELEMENTS); + DataCopy(tmpLocal, h_[tmpBase], KDA_SOLVE_MATRIX_ELEMENTS); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + + Add(xLocal, xLocal, tmpLocal, KDA_SOLVE_MATRIX_ELEMENTS); + PipeBarrier(); + + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + DataCopy(h_[xBase], xLocal, KDA_SOLVE_MATRIX_ELEMENTS); + if (storeAkk) { + DataCopy(akk_[AOffset(b, hv, start, 0)], xLocal, KDA_SOLVE_MATRIX_ELEMENTS); + } + SetFlag(mte3ToMte2Events_[0]); + WaitFlag(mte3ToMte2Events_[0]); + SetFlag(mte3ToVEvent_); + WaitFlag(mte3ToVEvent_); + } + + __aicore__ inline void AddSolveTmpToXTail(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, uint64_t curT, + bool storeAkk) + { + uint64_t elemCount = curT * KDA_SOLVE_BT; + LocalTensor arena = vecBuf_.Get(); + LocalTensor xLocal = arena; + LocalTensor tmpLocal = arena[KDA_SOLVE_MATRIX_ELEMENTS]; + uint64_t xBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_X); + uint64_t tmpBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_TMP); + + DataCopy(xLocal, h_[xBase], KDA_SOLVE_MATRIX_ELEMENTS); + DataCopy(tmpLocal, h_[tmpBase], KDA_SOLVE_MATRIX_ELEMENTS); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + + Add(xLocal, xLocal, tmpLocal, KDA_SOLVE_MATRIX_ELEMENTS); + PipeBarrier(); + + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + DataCopy(h_[xBase], xLocal, KDA_SOLVE_MATRIX_ELEMENTS); + if (storeAkk) { + DataCopy(akk_[AOffset(b, hv, start, 0)], xLocal, static_cast(elemCount)); + } + SetFlag(mte3ToMte2Events_[0]); + WaitFlag(mte3ToMte2Events_[0]); + SetFlag(mte3ToVEvent_); + WaitFlag(mte3ToVEvent_); + } + + __aicore__ inline void AddSolveTmpToXRows(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, uint64_t curT, + uint64_t rowBegin, uint64_t rowEnd, bool storeAkk) + { + uint64_t rowCount = rowEnd - rowBegin; + if (rowCount == 0) { + return; + } + uint64_t validRowCount = rowBegin < curT ? curT - rowBegin : 0; + if (validRowCount > rowCount) { + validRowCount = rowCount; + } + uint64_t elemCount = rowCount * BT_; + uint64_t validElemCount = validRowCount * BT_; + LocalTensor arena = vecBuf_.Get(); + LocalTensor xLocal = arena; + LocalTensor tmpLocal = arena[elemCount]; + uint64_t xBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_X) + rowBegin * BT_; + uint64_t tmpBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_TMP) + rowBegin * BT_; + uint64_t token = start + rowBegin; + + DataCopy(xLocal, solveWorkspace_[xBase], static_cast(elemCount)); + DataCopy(tmpLocal, solveWorkspace_[tmpBase], static_cast(elemCount)); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + + Add(xLocal, xLocal, tmpLocal, static_cast(elemCount)); + PipeBarrier(); + + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + DataCopy(solveWorkspace_[xBase], xLocal, static_cast(elemCount)); + if (storeAkk && validRowCount > 0) { + DataCopy(akk_[AOffset(b, hv, token, 0)], xLocal, static_cast(validElemCount)); + } + SetFlag(mte3ToMte2Events_[0]); + WaitFlag(mte3ToMte2Events_[0]); + SetFlag(mte3ToVEvent_); + WaitFlag(mte3ToVEvent_); + } + + __aicore__ inline void AddSolveTmpToXDiagRows(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, + uint64_t rowBegin, uint64_t rowEnd, bool storeAkk) + { + uint64_t rowCount = rowEnd - rowBegin; + if (rowCount == 0) { + return; + } + uint64_t elemCount = rowCount * BT_; + LocalTensor arena = vecBuf_.Get(); + LocalTensor xLocal = arena; + LocalTensor tmpLocal = arena[elemCount]; + uint64_t xBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_X) + rowBegin * BT_; + uint64_t tmpBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_TMP) + rowBegin * BT_; + uint64_t token = start + rowBegin; + + DataCopy(xLocal, solveWorkspace_[xBase], static_cast(elemCount)); + DataCopy(tmpLocal, solveWorkspace_[tmpBase], static_cast(elemCount)); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + + for (uint64_t localRow = 0; localRow < rowCount; ++localRow) { + uint64_t row = rowBegin + localRow; + uint64_t col = (row / KDA_SOLVE_DIAG_BT) * KDA_SOLVE_DIAG_BT; + uint64_t offset = localRow * BT_ + col; + Add(xLocal[offset], xLocal[offset], tmpLocal[offset], KDA_SOLVE_DIAG_BT); + PipeBarrier(); + } + + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + DataCopy(solveWorkspace_[xBase], xLocal, static_cast(elemCount)); + if (storeAkk) { + DataCopy(akk_[AOffset(b, hv, token, 0)], xLocal, static_cast(elemCount)); + } + SetFlag(mte3ToMte2Events_[0]); + WaitFlag(mte3ToMte2Events_[0]); + SetFlag(mte3ToVEvent_); + WaitFlag(mte3ToVEvent_); + } + + __aicore__ inline void StoreSolveXRowsToAkk(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, + uint64_t curT, uint64_t rowBegin, uint64_t rowEnd) + { + uint64_t validRowCount = rowBegin < curT ? curT - rowBegin : 0; + uint64_t rowCount = rowEnd - rowBegin; + if (validRowCount > rowCount) { + validRowCount = rowCount; + } + if (validRowCount == 0) { + return; + } + uint64_t elemCount = validRowCount * BT_; + LocalTensor xLocal = vecBuf_.Get(); + uint64_t xBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_X) + rowBegin * BT_; + + DataCopy(xLocal, solveWorkspace_[xBase], static_cast(elemCount)); + SetFlag(mte2ToMte3Event_); + WaitFlag(mte2ToMte3Event_); + DataCopy(akk_[AOffset(b, hv, start + rowBegin, 0)], xLocal, static_cast(elemCount)); + SetFlag(mte3ToMte2Events_[0]); + WaitFlag(mte3ToMte2Events_[0]); + } + + __aicore__ inline void ComputeAkkMergeCube(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start) + { + uint64_t aiBase = AOffset(b, hv, start, 0); + uint64_t negABase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_X); + uint64_t tmpBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_TMP); + + for (uint32_t mergeSize = 2 * KDA_SOLVE_DIAG_BT; mergeSize <= BT_; mergeSize *= 2) { + uint32_t half = mergeSize / 2; + for (uint32_t block = 0; block < BT_; block += mergeSize) { + uint32_t lower = block + half; + CubeGemmSolveSub(akk_, aiBase, lower, lower, solveWorkspace_, negABase, lower, block, solveWorkspace_, + tmpBase, 0, 0, half, half, half); + CubeGemmSolveSub(solveWorkspace_, tmpBase, 0, 0, akk_, aiBase, block, block, akk_, aiBase, lower, block, + half, half, half); + } + } + } + + __aicore__ inline void ComputeAkkMergeCubeWorkspace(uint64_t b, uint64_t hv, uint64_t chunkIdx) + { + uint64_t xBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_X); + uint64_t tmpBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_TMP); + + for (uint32_t mergeSize = 2 * KDA_SOLVE_DIAG_BT; mergeSize <= BT_; mergeSize *= 2) { + uint32_t half = mergeSize / 2; + for (uint32_t block = 0; block < BT_; block += mergeSize) { + uint32_t lower = block + half; + CubeGemmSolveSub(solveWorkspace_, xBase, lower, lower, solveWorkspace_, xBase, lower, block, + solveWorkspace_, tmpBase, 0, 0, half, half, half); + CubeGemmSolveSub(solveWorkspace_, tmpBase, 0, 0, solveWorkspace_, xBase, block, block, solveWorkspace_, + xBase, lower, block, half, half, half); + } + } + } + + __aicore__ inline void ComputeAkkInverseMchFull(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start) + { + uint64_t aBase = AOffset(b, hv, start, 0); + uint64_t xBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_X); + uint64_t yBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_Y0); + uint64_t yNextBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_Y1); + uint64_t tmpBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_TMP); + + uint32_t diagBlocks = static_cast(BT_ / KDA_SOLVE_DIAG_BT); + for (uint32_t block = 0; block < diagBlocks; ++block) { + uint32_t off = block * KDA_SOLVE_DIAG_BT; + CubeGemmSolveSub(akk_, aBase, off, off, akk_, aBase, off, off, solveWorkspace_, yBase, off, off, + KDA_SOLVE_DIAG_BT, KDA_SOLVE_DIAG_BT, KDA_SOLVE_DIAG_BT); + } + for (uint32_t iter = 0; iter < KDA_SOLVE_DIAG_MCH_ITERS; ++iter) { + for (uint32_t block = 0; block < diagBlocks; ++block) { + uint32_t off = block * KDA_SOLVE_DIAG_BT; + CubeGemmSolveSub(solveWorkspace_, xBase, off, off, solveWorkspace_, yBase, off, off, solveWorkspace_, + tmpBase, off, off, KDA_SOLVE_DIAG_BT, KDA_SOLVE_DIAG_BT, KDA_SOLVE_DIAG_BT); + } + Catlass::Arch::CrossCoreSetFlagWithReverse<0x2, PIPE_FIX>(mchSyncDoneFlag_); + if (iter + 1 < KDA_SOLVE_DIAG_MCH_ITERS) { + for (uint32_t block = 0; block < diagBlocks; ++block) { + uint32_t off = block * KDA_SOLVE_DIAG_BT; + CubeGemmSolveSub(solveWorkspace_, yBase, off, off, solveWorkspace_, yBase, off, off, + solveWorkspace_, yNextBase, off, off, KDA_SOLVE_DIAG_BT, KDA_SOLVE_DIAG_BT, + KDA_SOLVE_DIAG_BT); + } + } + Catlass::Arch::CrossCoreWaitFlagWithReverse<0x2, PIPE_FIX>(mchSyncReadyFlag_); + if (iter + 1 < KDA_SOLVE_DIAG_MCH_ITERS) { + uint64_t oldYBase = yBase; + yBase = yNextBase; + yNextBase = oldYBase; + } + } + ComputeAkkMergeCube(b, hv, chunkIdx, start); + Catlass::Arch::CrossCoreSetFlagWithReverse<0x2, PIPE_FIX>(mchSyncDoneFlag_); + } + + __aicore__ inline void ComputeAkkInverseMchTail(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, + uint64_t curT) + { + uint64_t xBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_X); + uint64_t lBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_Y0); + uint64_t yBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_Y1); + uint64_t yNextBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_Y0); + uint64_t tmpBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_TMP); + (void)start; + (void)curT; + + uint32_t diagBlocks = static_cast(BT_ / KDA_SOLVE_DIAG_BT); + for (uint32_t block = 0; block < diagBlocks; ++block) { + uint32_t off = block * KDA_SOLVE_DIAG_BT; + CubeGemmSolveSub(solveWorkspace_, lBase, off, off, solveWorkspace_, lBase, off, off, solveWorkspace_, yBase, + off, off, KDA_SOLVE_DIAG_BT, KDA_SOLVE_DIAG_BT, KDA_SOLVE_DIAG_BT); + } + for (uint32_t iter = 0; iter < KDA_SOLVE_DIAG_MCH_ITERS; ++iter) { + for (uint32_t block = 0; block < diagBlocks; ++block) { + uint32_t off = block * KDA_SOLVE_DIAG_BT; + CubeGemmSolveSub(solveWorkspace_, xBase, off, off, solveWorkspace_, yBase, off, off, solveWorkspace_, + tmpBase, off, off, KDA_SOLVE_DIAG_BT, KDA_SOLVE_DIAG_BT, KDA_SOLVE_DIAG_BT); + } + Catlass::Arch::CrossCoreSetFlagWithReverse<0x2, PIPE_FIX>(mchSyncDoneFlag_); + if (iter + 1 < KDA_SOLVE_DIAG_MCH_ITERS) { + for (uint32_t block = 0; block < diagBlocks; ++block) { + uint32_t off = block * KDA_SOLVE_DIAG_BT; + CubeGemmSolveSub(solveWorkspace_, yBase, off, off, solveWorkspace_, yBase, off, off, + solveWorkspace_, yNextBase, off, off, KDA_SOLVE_DIAG_BT, KDA_SOLVE_DIAG_BT, + KDA_SOLVE_DIAG_BT); + } + } + Catlass::Arch::CrossCoreWaitFlagWithReverse<0x2, PIPE_FIX>(mchSyncReadyFlag_); + if (iter + 1 < KDA_SOLVE_DIAG_MCH_ITERS) { + uint64_t oldYBase = yBase; + yBase = yNextBase; + yNextBase = oldYBase; + } + } + ComputeAkkMergeCubeWorkspace(b, hv, chunkIdx); + Catlass::Arch::CrossCoreSetFlagWithReverse<0x2, PIPE_FIX>(mchSyncDoneFlag_); + } + + __aicore__ inline void ScaleRowsByBeta(GlobalTensor &src, GlobalTensor &dst, uint64_t b, uint64_t hv, + uint64_t start, uint64_t rowBegin, uint64_t rowCount, uint64_t dim, + LocalTensor &betaLocal, LocalTensor &betaBrcb, + LocalTensor &matrixLocal, bool sourceSequenceMajor = false) + { + constexpr uint64_t vecElemsPerRepeat = 64; + constexpr uint64_t typedOffsetFloats = 20480; + constexpr uint64_t typedOffset = typedOffsetFloats * sizeof(float) / sizeof(T); + uint64_t elemCount = rowCount * dim; + uint64_t baseOffset = KVOffset(b, hv, start + rowBegin, 0, dim); + uint64_t sourceOffset = sourceSequenceMajor ? VInputOffset(b, hv, start + rowBegin, 0) : baseOffset; + uint64_t sourceStride = sourceSequenceMajor ? HV_ * dim : dim; + + if constexpr (IsSameType::value) { + CopyRowsIn(matrixLocal, src, sourceOffset, rowCount, dim, sourceStride); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + } else { + LocalTensor matrixTyped = vecBuf_.Get()[typedOffset]; + CopyRowsIn(matrixTyped, src, sourceOffset, rowCount, dim, sourceStride); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + Cast(matrixLocal, matrixTyped, RoundMode::CAST_NONE, static_cast(elemCount)); + PipeBarrier(); + } + + uint8_t repeatStride = static_cast(dim * sizeof(float) / 32); + for (uint64_t col = 0; col < dim; col += vecElemsPerRepeat) { + uint64_t mask = dim - col; + if (mask > vecElemsPerRepeat) { + mask = vecElemsPerRepeat; + } + Mul(matrixLocal[col], matrixLocal[col], betaBrcb, mask, static_cast(rowCount), + {1, 1, 0, repeatStride, repeatStride, 1}); + } + PipeBarrier(); + + if constexpr (IsSameType::value) { + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + DataCopy(dst[baseOffset], matrixLocal, static_cast(elemCount)); + } else { + LocalTensor matrixTyped = vecBuf_.Get()[typedOffset]; + Cast(matrixTyped, matrixLocal, RoundMode::CAST_RINT, static_cast(elemCount)); + PipeBarrier(); + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + DataCopy(dst[baseOffset], matrixTyped, static_cast(elemCount)); + } + SetFlag(mte3ToMte2Events_[0]); + WaitFlag(mte3ToMte2Events_[0]); + SetFlag(mte3ToVEvent_); + WaitFlag(mte3ToVEvent_); + } + + __aicore__ inline void PrepareWuCubeInputs(uint64_t b, uint64_t hv, uint64_t start, uint64_t curT, + uint64_t subBlockIdx, uint64_t subBlockNum) + { + uint64_t rowsPerSubBlock = (curT + subBlockNum - 1) / subBlockNum; + uint64_t rowBegin = subBlockIdx * rowsPerSubBlock; + if (rowBegin >= curT) { + return; + } + uint64_t rowCount = curT - rowBegin; + if (rowCount > rowsPerSubBlock) { + rowCount = rowsPerSubBlock; + } + LocalTensor arena = vecBuf_.Get(); + LocalTensor betaLocal = arena; + LocalTensor betaBrcb = arena[KDA_SOLVE_BT]; + LocalTensor matrixLocal = arena[KDA_SOLVE_BT + 512]; + LoadAsFloatRow(beta_, BetaOffset(b, hv, start + rowBegin), betaLocal, rowCount); + Brcb(betaBrcb, betaLocal, static_cast((rowCount + 7) / 8), {1, 8}); + PipeBarrier(); + ScaleRowsByBeta(w_, w_, b, hv, start, rowBegin, rowCount, K_, betaLocal, betaBrcb, matrixLocal); + ScaleRowsByBeta(v_, vNew_, b, hv, start, rowBegin, rowCount, V_, betaLocal, betaBrcb, matrixLocal, + inputSequenceMajor_); + } + + __aicore__ inline void FinalizePrepareIntermediates(uint64_t b, uint64_t hv, uint64_t start, uint64_t curT, + uint64_t subBlockIdx, uint64_t subBlockNum) + { + constexpr uint64_t tileRows = 32; + // Keep tail rows on the same AIV that owns their padded solve rows. Splitting by curT would + // move short-tail export to AIV1 while AIV0 is still writing the solved matrix. + const uint64_t rowBegin = (BT_ * subBlockIdx) / subBlockNum; + uint64_t rowEnd = (BT_ * (subBlockIdx + 1)) / subBlockNum; + if (rowEnd > curT) { + rowEnd = curT; + } + if (rowBegin >= rowEnd) { + return; + } + for (uint64_t tileRow = rowBegin; tileRow < rowEnd; tileRow += tileRows) { + const uint64_t rows = (rowEnd - tileRow) > tileRows ? tileRows : (rowEnd - tileRow); + const uint64_t matrixElems = rows * BT_; + const uint64_t qgElems = rows * K_; + LocalTensor arena = vecBuf_.Get(); + LocalTensor aqkLocal = arena; + LocalTensor akkLocal = arena[matrixElems]; + const uint64_t typedOffset = (2 * matrixElems + qgElems) * sizeof(float) / sizeof(T); + LocalTensor typedBase = vecBuf_.Get()[typedOffset]; + LocalTensor aqkTyped = typedBase; + LocalTensor akkTyped = typedBase[matrixElems]; + LocalTensor qgTyped = typedBase[2 * matrixElems]; + + CopyVectorIn(aqkLocal, aqk_, AOffset(b, hv, start + tileRow, 0), matrixElems); + CopyVectorIn(akkLocal, akk_, AOffset(b, hv, start + tileRow, 0), matrixElems); + CopyVectorIn(qgTyped, qg_, KVOffset(b, hv, start + tileRow, 0, K_), qgElems); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + + Muls(aqkLocal, aqkLocal, scale_, static_cast(matrixElems)); + ClampFp32ToOutputType(aqkLocal, static_cast(matrixElems)); + ClampFp32ToOutputType(akkLocal, static_cast(matrixElems)); + Cast(aqkTyped, aqkLocal, RoundMode::CAST_RINT, static_cast(matrixElems)); + Cast(akkTyped, akkLocal, RoundMode::CAST_RINT, static_cast(matrixElems)); + PipeBarrier(); + + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + CopyVectorOut(o_, AOffset(b, hv, start + tileRow, 0), aqkTyped, matrixElems); + CopyVectorOut(u_, AOffset(b, hv, start + tileRow, 0), akkTyped, matrixElems); + CopyVectorOut(kg_, KVOffset(b, hv, start + tileRow, 0, K_), qgTyped, qgElems); + SetFlag(mte3ToMte2Events_[0]); + WaitFlag(mte3ToMte2Events_[0]); + SetFlag(mte3ToVEvent_); + WaitFlag(mte3ToVEvent_); + } + } + + __aicore__ inline bool ResolveFlatChunk(uint64_t task, uint64_t &seq, uint64_t &b, uint64_t &h, uint64_t &hv, + uint64_t &chunkIdx, uint64_t &start, uint64_t &end) + { + hv = task % HV_; + uint64_t flatChunk = task / HV_; + if (!isVarLen_) { + seq = flatChunk / NT_; + b = seq; + chunkIdx = flatChunk % NT_; + start = chunkIdx * BT_; + end = start + BT_; + if (end > T_) { + end = T_; + } + } else { + if (!KdaVarlen::ResolveChunkRange(cuSeqlensAddr_, chunkIndicesAddr_, N_, T_, BT_, flatChunk, seq, start, + end)) { + return false; + } + b = 0; + chunkIdx = flatChunk; + } + h = hv / (HV_ / H_); + return start < end; + } + + __aicore__ inline void ProcessChunkPreAiv(uint64_t b, uint64_t h, uint64_t hv, uint64_t chunkIdx, uint64_t start, + uint64_t end, uint64_t subBlockIdx, uint64_t subBlockNum) + { + if constexpr (IsSameType::value) { + ProcessChunkPreAivFp32(b, h, hv, chunkIdx, start, end, subBlockIdx, subBlockNum); + } + } + + template + __aicore__ inline void JoinAivMte3() + { + if constexpr (CORE_TYPE == AscendC::AIV) { + } + } + + template + __aicore__ inline void RunAicAfterBothAivReady(uint64_t subBlockIdx, uint64_t subBlockNum) + { + if constexpr (CORE_TYPE == AscendC::AIV) { + (void)subBlockIdx; + (void)subBlockNum; + JoinAivMte3(); + if constexpr (SAFE_GATE) { + Catlass::Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(syncReadyFlag_); + Catlass::Arch::CrossCoreWaitFlag(syncDoneFlag_); + } else { + Catlass::Arch::CrossCoreSetFlagWithReverse<0x2, PIPE_MTE3>(mchSyncReadyFlag_); + Catlass::Arch::CrossCoreWaitFlagWithReverse<0x2, PIPE_MTE2>(mchSyncDoneFlag_); + } + } + } + + template + __aicore__ inline void SignalAicSolveReady() + { + if constexpr (CORE_TYPE == AscendC::AIV) { + JoinAivMte3(); + Catlass::Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(syncReadyFlag_); + } + } + + template + __aicore__ inline void WaitAicSolveDone() + { + if constexpr (CORE_TYPE == AscendC::AIV) { + Catlass::Arch::CrossCoreWaitFlag(syncDoneFlag_); + } + } + + __aicore__ inline void ProcessChunkPreAivFp32(uint64_t b, uint64_t h, uint64_t hv, uint64_t chunkIdx, + uint64_t start, uint64_t end, uint64_t subBlockIdx, + uint64_t subBlockNum, bool deferSafeSolve = false, + bool waitPendingSafeSolve = false, uint64_t scoreLane = 0, + bool pairHeads = false) + { + uint64_t curT = end - start; + if (curT == 0) { + return; + } + if constexpr (IsSameType::value) { + return; + } + + if (K_ < 16) { + return; + } + bool usePostWuCube = UsePostWuCube(curT); + bool useAkkCubeSolve = UseAkkCubeSolve(curT); + uint64_t solveRowBegin = 0; + uint64_t solveRowEnd = 0; + GetSolveRowRange(BT_, subBlockIdx, subBlockNum, solveRowBegin, solveRowEnd); + uint64_t scoreBlockSize = ScoreRefBlockSize(); + uint64_t scoreBlockCount = (curT + scoreBlockSize - 1) / scoreBlockSize; + uint64_t pipelineBlockCount = + (scoreBlockCount + KDA_SCORE_QUEUE_DEPTH - 1) / KDA_SCORE_QUEUE_DEPTH * KDA_SCORE_QUEUE_DEPTH; + for (uint64_t block = 0; block < pipelineBlockCount; ++block) { + if (block < scoreBlockCount) { + uint64_t rowBegin = block * scoreBlockSize; + uint64_t rowCount = ScoreRowBlockCount(curT, rowBegin); + uint64_t refToken = ScoreRefToken(start, curT, rowBegin, rowCount); + uint64_t scoreSlot = ScoreScratchSlot(block % KDA_SCORE_QUEUE_DEPTH, scoreLane, pairHeads); + PrepareGateProducts(b, h, hv, start, curT, subBlockIdx, subBlockNum, true, refToken, + rowBegin + rowCount, true, scoreSlot, rowBegin, rowCount); + } + JoinAivMte3(); + Catlass::Arch::CrossCoreSetFlagWithReverse<0x2, PIPE_MTE3>(scoreReadyFlag_); + if (block > 0) { + if constexpr (SAFE_GATE) { + if (waitPendingSafeSolve) { + WaitAicSolveDone(); + waitPendingSafeSolve = false; + } + } + Catlass::Arch::CrossCoreWaitFlagWithReverse<0x2, PIPE_MTE2>(scoreDoneFlag_); + } + } + bool fusedScoreWriteback = false; + if (!fusedScoreWriteback) { + // The final score MMAD only consumes scoreWorkspace_. Run the + // independent gate writeback while AIC drains its MMAD/Fixpipe path. + PrepareGateProducts(b, h, hv, start, curT, subBlockIdx, subBlockNum); + } + if (pipelineBlockCount > 0) { + Catlass::Arch::CrossCoreWaitFlagWithReverse<0x2, PIPE_MTE2>(scoreDoneFlag_); + } + if constexpr (SAFE_GATE) { + if (waitPendingSafeSolve) { + WaitAicSolveDone(); + waitPendingSafeSolve = false; + } + } + if (useAkkCubeSolve) { + bool fullChunk = curT == BT_; + if constexpr (SAFE_GATE) { + if (pairHeads) { + for (uint64_t rowPart = 0; rowPart < KDA_SCORE_LANES; ++rowPart) { + uint64_t pairRowBegin = 0; + uint64_t pairRowEnd = 0; + GetSolveRowRange(BT_, rowPart, KDA_SCORE_LANES, pairRowBegin, pairRowEnd); + PrepareAqkAkkSolveInputRows(b, hv, chunkIdx, start, curT, pairRowBegin, pairRowEnd, false, + false); + } + } else { + PrepareAqkAkkSolveInputRows(b, hv, chunkIdx, start, curT, solveRowBegin, solveRowEnd, false, false); + } + if (deferSafeSolve) { + SignalAicSolveReady(); + return; + } + RunAicAfterBothAivReady(subBlockIdx, subBlockNum); + StoreSolveXRowsToAkk(b, hv, chunkIdx, start, curT, solveRowBegin, solveRowEnd); + } else { + PrepareAqkAkkSolveInputRows(b, hv, chunkIdx, start, curT, solveRowBegin, solveRowEnd, fullChunk, + !fullChunk); + uint32_t solveIters = KDA_SOLVE_DIAG_MCH_ITERS; + RunAicAfterBothAivReady(subBlockIdx, subBlockNum); + for (uint32_t iter = 0; iter < solveIters; ++iter) { + AddSolveTmpToXDiagRows(b, hv, chunkIdx, start, solveRowBegin, solveRowEnd, + fullChunk && iter + 1 == solveIters); + RunAicAfterBothAivReady(subBlockIdx, subBlockNum); + } + if (!fullChunk) { + StoreSolveXRowsToAkk(b, hv, chunkIdx, start, curT, solveRowBegin, solveRowEnd); + } + } + } + // Host validation guarantees every accepted shape has enough workspace for this cube path. + PrepareWuCubeInputs(b, hv, start, curT, subBlockIdx, subBlockNum); + FinalizePrepareIntermediates(b, hv, start, curT, subBlockIdx, subBlockNum); + } + + __aicore__ inline void FinishDeferredSafeChunk(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, + uint64_t end, uint64_t subBlockIdx, uint64_t subBlockNum) + { + uint64_t curT = end - start; + uint64_t solveRowBegin = 0; + uint64_t solveRowEnd = 0; + GetSolveRowRange(BT_, subBlockIdx, subBlockNum, solveRowBegin, solveRowEnd); + StoreSolveXRowsToAkk(b, hv, chunkIdx, start, curT, solveRowBegin, solveRowEnd); + PrepareWuCubeInputs(b, hv, start, curT, subBlockIdx, subBlockNum); + FinalizePrepareIntermediates(b, hv, start, curT, subBlockIdx, subBlockNum); + } + + __aicore__ inline void FinishDeferredSafeChunkPair(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, + uint64_t end) + { + FinishDeferredSafeChunk(b, hv, chunkIdx, start, end, 0, 1); + } + + __aicore__ inline void ProcessChunkPreAic(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, uint64_t end) + { + if constexpr (IsSameType::value) { + ProcessChunkPreAicFp32(b, hv, chunkIdx, start, end); + } + } + + __aicore__ inline void ProcessChunkPreAicFp32(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, + uint64_t end) + { + uint64_t curT = end - start; + if (curT == 0 || K_ < 16) { + return; + } + uint64_t scoreBlockSize = ScoreRefBlockSize(); + uint64_t scoreBlockCount = (curT + scoreBlockSize - 1) / scoreBlockSize; + uint64_t pipelineBlockCount = + (scoreBlockCount + KDA_SCORE_QUEUE_DEPTH - 1) / KDA_SCORE_QUEUE_DEPTH * KDA_SCORE_QUEUE_DEPTH; + for (uint64_t block = 0; block < pipelineBlockCount; ++block) { + Catlass::Arch::CrossCoreWaitFlagWithReverse<0x2, PIPE_FIX>(scoreReadyFlag_); + if (block < scoreBlockCount) { + uint64_t rowBegin = block * scoreBlockSize; + uint64_t rowCount = ScoreRowBlockCount(curT, rowBegin); + ComputeRawAqkAkkCubeBlock(b, hv, start, curT, rowBegin, rowCount, true, block % KDA_SCORE_QUEUE_DEPTH, + rowBegin + rowCount); + } + Catlass::Arch::CrossCoreSetFlagWithReverse<0x2, PIPE_FIX>(scoreDoneFlag_); + } + bool usePostWuCube = UsePostWuCube(curT); + bool useAkkCubeSolve = UseAkkCubeSolve(curT); + if (useAkkCubeSolve) { + if constexpr (SAFE_GATE) { + Catlass::Arch::CrossCoreWaitFlag(syncReadyFlag_); + ComputeAkkMergeCubeWorkspace(b, hv, chunkIdx); + Catlass::Arch::CrossCoreSetFlag<0x2, PIPE_FIX>(syncDoneFlag_); + } else { + Catlass::Arch::CrossCoreWaitFlagWithReverse<0x2, PIPE_FIX>(mchSyncReadyFlag_); + if (curT == BT_) { + ComputeAkkInverseMchFull(b, hv, chunkIdx, start); + } else { + ComputeAkkInverseMchTail(b, hv, chunkIdx, start, curT); + } + } + } + (void)usePostWuCube; + (void)chunkIdx; + } + + __aicore__ inline void ProcessChunkPreAicHeadPairFp32(uint64_t b, uint64_t hvBase, uint64_t chunkIdx, + uint64_t start, uint64_t end, uint64_t localTaskIdx) + { + uint64_t curT = end - start; + if (curT == 0 || K_ < 16) { + return; + } + uint64_t scoreBlockSize = ScoreRefBlockSize(); + uint64_t scoreBlockCount = (curT + scoreBlockSize - 1) / scoreBlockSize; + uint64_t pipelineBlockCount = + (scoreBlockCount + KDA_SCORE_QUEUE_DEPTH - 1) / KDA_SCORE_QUEUE_DEPTH * KDA_SCORE_QUEUE_DEPTH; + for (uint64_t block = 0; block < pipelineBlockCount; ++block) { + Catlass::Arch::CrossCoreWaitFlagWithReverse<0x2, PIPE_FIX>(scoreReadyFlag_); + if (block < scoreBlockCount) { + uint64_t rowBegin = block * scoreBlockSize; + uint64_t rowCount = ScoreRowBlockCount(curT, rowBegin); + for (uint64_t lane = 0; lane < KDA_SCORE_LANES; ++lane) { + uint64_t hv = hvBase + lane; + uint64_t scoreSlot = ScoreScratchSlot(block % KDA_SCORE_QUEUE_DEPTH, lane, true); + ComputeRawAqkAkkCubeBlock(b, hv, start, curT, rowBegin, rowCount, true, scoreSlot, + rowBegin + rowCount); + } + } + Catlass::Arch::CrossCoreSetFlagWithReverse<0x2, PIPE_FIX>(scoreDoneFlag_); + } + + if (UseAkkCubeSolve(curT)) { + Catlass::Arch::CrossCoreWaitFlag(syncReadyFlag_); + for (uint64_t lane = 0; lane < KDA_SCORE_LANES; ++lane) { + activeSolveSlot_ = + (localTaskIdx % (KDA_SOLVE_PIPELINE_DEPTH / KDA_SCORE_LANES)) * KDA_SCORE_LANES + lane; + ComputeAkkMergeCubeWorkspace(b, hvBase + lane, chunkIdx); + } + Catlass::Arch::CrossCoreSetFlag<0x2, PIPE_FIX>(syncDoneFlag_); + } + } + + __aicore__ inline bool ResolveFlatChunkForHv(uint64_t flatChunk, uint64_t hv, uint64_t &seq, uint64_t &b, + uint64_t &h, uint64_t &chunkIdx, uint64_t &start, uint64_t &end) + { + if (!isVarLen_) { + seq = flatChunk / NT_; + b = seq; + chunkIdx = flatChunk % NT_; + start = chunkIdx * BT_; + end = start + BT_; + if (end > T_) { + end = T_; + } + } else { + if (!KdaVarlen::ResolveChunkRange(cuSeqlensAddr_, chunkIndicesAddr_, N_, T_, BT_, flatChunk, seq, start, + end)) { + return false; + } + b = 0; + chunkIdx = flatChunk; + } + h = hv / (HV_ / H_); + return start < end; + } + + __aicore__ inline void ProcessPreAivHeadPair() + { + const uint64_t subBlockIdx = static_cast(GetSubBlockIdx()); + const uint64_t subBlockNum = static_cast(GetSubBlockNum()); + const uint64_t coreNum = usedCoreNum_; + const uint64_t coreIdx = static_cast(GetBlockIdx()) / subBlockNum; + const uint64_t chunkCount = isVarLen_ ? NT_ : B_ * NT_; + const uint64_t headWindows = HV_ / KDA_SCORE_LANES; + const uint64_t taskNum = chunkCount * headWindows; + bool pendingValid = false; + uint64_t pendingB = 0; + uint64_t pendingHv = 0; + uint64_t pendingChunkIdx = 0; + uint64_t pendingStart = 0; + uint64_t pendingEnd = 0; + uint64_t pendingSlot = 0; + uint64_t localTaskIdx = 0; + + for (uint64_t task = coreIdx; task < taskNum; task += coreNum, ++localTaskIdx) { + uint64_t flatChunk = task / headWindows; + uint64_t hv = (task % headWindows) * KDA_SCORE_LANES + subBlockIdx; + uint64_t seq = 0; + uint64_t b = 0; + uint64_t h = 0; + uint64_t chunkIdx = 0; + uint64_t start = 0; + uint64_t end = 0; + if (!ResolveFlatChunkForHv(flatChunk, hv, seq, b, h, chunkIdx, start, end)) { + continue; + } + (void)seq; + uint64_t currentSlot = + (localTaskIdx % (KDA_SOLVE_PIPELINE_DEPTH / KDA_SCORE_LANES)) * KDA_SCORE_LANES + subBlockIdx; + activeSolveSlot_ = currentSlot; + bool deferSolve = UseAkkCubeSolve(end - start); + ProcessChunkPreAivFp32(b, h, hv, chunkIdx, start, end, 0, 1, deferSolve, pendingValid, subBlockIdx, true); + if (pendingValid) { + activeSolveSlot_ = pendingSlot; + FinishDeferredSafeChunkPair(pendingB, pendingHv, pendingChunkIdx, pendingStart, pendingEnd); + } + pendingValid = deferSolve; + if (pendingValid) { + pendingB = b; + pendingHv = hv; + pendingChunkIdx = chunkIdx; + pendingStart = start; + pendingEnd = end; + pendingSlot = currentSlot; + } + } + if (pendingValid) { + WaitAicSolveDone(); + activeSolveSlot_ = pendingSlot; + FinishDeferredSafeChunkPair(pendingB, pendingHv, pendingChunkIdx, pendingStart, pendingEnd); + } + } + + __aicore__ inline void ProcessPreAicHeadPair() + { + const uint64_t chunkCount = isVarLen_ ? NT_ : B_ * NT_; + const uint64_t headWindows = HV_ / KDA_SCORE_LANES; + const uint64_t taskNum = chunkCount * headWindows; + const uint64_t coreNum = usedCoreNum_ == 0 ? 1 : usedCoreNum_; + uint64_t localTaskIdx = 0; + for (uint64_t task = GetBlockIdx(); task < taskNum; task += coreNum, ++localTaskIdx) { + uint64_t flatChunk = task / headWindows; + uint64_t hvBase = (task % headWindows) * KDA_SCORE_LANES; + uint64_t seq = 0; + uint64_t b = 0; + uint64_t h = 0; + uint64_t chunkIdx = 0; + uint64_t start = 0; + uint64_t end = 0; + if (ResolveFlatChunkForHv(flatChunk, hvBase, seq, b, h, chunkIdx, start, end)) { + (void)seq; + (void)h; + ProcessChunkPreAicHeadPairFp32(b, hvBase, chunkIdx, start, end, localTaskIdx); + } + } + } + + __aicore__ inline void ProcessPreAiv() + { + if constexpr (IsSameType::value) { + isAivOnly_ = true; + } + uint64_t subBlockNum = isAivOnly_ ? 1 : static_cast(GetSubBlockNum()); + if (subBlockNum == 0) { + return; + } + uint64_t subBlockIdx = isAivOnly_ ? 0 : static_cast(GetSubBlockIdx()); + uint64_t coreNum = isAivOnly_ ? static_cast(GetBlockNum()) : usedCoreNum_; + uint64_t coreIdx = + isAivOnly_ ? static_cast(GetBlockIdx()) : static_cast(GetBlockIdx()) / subBlockNum; + uint64_t taskNum = static_cast((isVarLen_ ? NT_ : B_ * NT_) * HV_); + if constexpr (SAFE_GATE && !IsSameType::value) { + bool pendingValid = false; + uint64_t pendingB = 0; + uint64_t pendingHv = 0; + uint64_t pendingChunkIdx = 0; + uint64_t pendingStart = 0; + uint64_t pendingEnd = 0; + uint64_t pendingSlot = 0; + uint64_t localTaskIdx = 0; + for (uint64_t task = coreIdx; task < taskNum; task += coreNum, ++localTaskIdx) { + uint64_t seq = 0; + uint64_t b = 0; + uint64_t h = 0; + uint64_t hv = 0; + uint64_t chunkIdx = 0; + uint64_t start = 0; + uint64_t end = 0; + if (!ResolveFlatChunk(task, seq, b, h, hv, chunkIdx, start, end)) { + continue; + } + (void)seq; + uint64_t currentSlot = localTaskIdx % KDA_SOLVE_PIPELINE_DEPTH; + activeSolveSlot_ = currentSlot; + bool deferSolve = UseAkkCubeSolve(end - start); + if (!deferSolve && pendingValid) { + WaitAicSolveDone(); + activeSolveSlot_ = pendingSlot; + FinishDeferredSafeChunk(pendingB, pendingHv, pendingChunkIdx, pendingStart, pendingEnd, subBlockIdx, + subBlockNum); + pendingValid = false; + activeSolveSlot_ = currentSlot; + } + ProcessChunkPreAivFp32(b, h, hv, chunkIdx, start, end, subBlockIdx, subBlockNum, deferSolve, + pendingValid); + if (pendingValid) { + activeSolveSlot_ = pendingSlot; + FinishDeferredSafeChunk(pendingB, pendingHv, pendingChunkIdx, pendingStart, pendingEnd, subBlockIdx, + subBlockNum); + } + pendingValid = deferSolve; + if (pendingValid) { + pendingB = b; + pendingHv = hv; + pendingChunkIdx = chunkIdx; + pendingStart = start; + pendingEnd = end; + pendingSlot = currentSlot; + } + } + if (pendingValid) { + WaitAicSolveDone(); + activeSolveSlot_ = pendingSlot; + FinishDeferredSafeChunk(pendingB, pendingHv, pendingChunkIdx, pendingStart, pendingEnd, subBlockIdx, + subBlockNum); + } + return; + } + for (uint64_t task = coreIdx; task < taskNum; task += coreNum) { + uint64_t seq = 0; + uint64_t b = 0; + uint64_t h = 0; + uint64_t hv = 0; + uint64_t chunkIdx = 0; + uint64_t start = 0; + uint64_t end = 0; + if (ResolveFlatChunk(task, seq, b, h, hv, chunkIdx, start, end)) { + (void)seq; + ProcessChunkPreAiv(b, h, hv, chunkIdx, start, end, subBlockIdx, subBlockNum); + } + } + } + + __aicore__ inline void ProcessPreAic() + { + if constexpr (IsSameType::value) { + return; + } + uint64_t taskNum = static_cast((isVarLen_ ? NT_ : B_ * NT_) * HV_); + uint64_t coreNum = usedCoreNum_ == 0 ? 1 : usedCoreNum_; + uint64_t localTaskIdx = 0; + for (uint64_t task = GetBlockIdx(); task < taskNum; task += coreNum, ++localTaskIdx) { + uint64_t seq = 0; + uint64_t b = 0; + uint64_t h = 0; + uint64_t hv = 0; + uint64_t chunkIdx = 0; + uint64_t start = 0; + uint64_t end = 0; + if (ResolveFlatChunk(task, seq, b, h, hv, chunkIdx, start, end)) { + if constexpr (SAFE_GATE) { + activeSolveSlot_ = localTaskIdx % KDA_SOLVE_PIPELINE_DEPTH; + } + (void)seq; + (void)h; + ProcessChunkPreAic(b, hv, chunkIdx, start, end); + } + } + } + +private: + GlobalTensor q_; + GlobalTensor k_; + GlobalTensor v_; + GlobalTensor gk_; + GlobalTensor beta_; + GlobalTensor initialState_; + GlobalTensor o_; + GlobalTensor finalState_; + GlobalTensor aqk_; + GlobalTensor akk_; + GlobalTensor w_; + GlobalTensor u_; + GlobalTensor qg_; + GlobalTensor kg_; + GlobalTensor vNew_; + GlobalTensor h_; + GlobalTensor preparedQG_; + GlobalTensor preparedAqk_; + GlobalTensor propagatedVNew_; + GlobalTensor propagatedH_; + GlobalTensor solveWorkspace_; + GlobalTensor scoreWorkspace_; + TPipe *pipe_ = nullptr; + TBuf exp2Buf_; + TBuf vecBuf_; + TBuf gateWritebackBuf_; + TEventID mte2ToVEvent_ = 0; + TEventID vToMte2Event_ = 0; + TEventID vToMte3Event_ = 0; + TEventID mte3ToVEvent_ = 0; + TEventID mte2ToMte3Event_ = 0; + TEventID mte3ToMte2Events_[KDA_GATE_PIPELINE_DEPTH] = {0, 0, 0}; + bool vectorEventsAllocated_ = false; + Catlass::Arch::CrossCoreFlagWithReverse scoreReadyFlag_{KDA_SCORE_READY_FLAG0, + KDA_SCORE_READY_FLAG1}; + Catlass::Arch::CrossCoreFlagWithReverse scoreDoneFlag_{KDA_SCORE_DONE_FLAG0, + KDA_SCORE_DONE_FLAG1}; + // Solve has one outstanding task per core. Reuse the primary score IDs as an ordered token stream; + // score credits remain on the reverse IDs, so no additional hardware flag IDs are consumed. + Catlass::Arch::CrossCoreFlag syncReadyFlag_{KDA_SOLVE_READY_FLAG}; + Catlass::Arch::CrossCoreFlag syncDoneFlag_{KDA_SOLVE_DONE_FLAG}; + Catlass::Arch::CrossCoreFlagWithReverse mchSyncReadyFlag_{KDA_SCORE_READY_FLAG0, + KDA_SCORE_READY_FLAG1}; + Catlass::Arch::CrossCoreFlagWithReverse mchSyncDoneFlag_{KDA_SCORE_DONE_FLAG0, + KDA_SCORE_DONE_FLAG1}; + uint64_t B_ = 0; + uint64_t N_ = 0; + uint64_t H_ = 0; + uint64_t HV_ = 0; + uint64_t T_ = 0; + uint64_t K_ = 0; + uint64_t V_ = 0; + uint64_t BT_ = 0; + uint64_t NT_ = 0; + float scale_ = 1.0f; + bool hasInitial_ = false; + bool isVarLen_ = false; + bool inputSequenceMajor_ = false; + bool isAivOnly_ = false; + uint64_t usedCoreNum_ = 1; + uint64_t solveCoreIdx_ = 0; + uint64_t activeSolveSlot_ = 0; + __gm__ int64_t *chunkIndicesAddr_ = nullptr; + __gm__ int64_t *cuSeqlensAddr_ = nullptr; +}; +} // namespace + +template +__aicore__ inline void RunChunkKdaPrepare(GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR gk, GM_ADDR beta, + GM_ADDR initialState, GM_ADDR cuSeqlens, GM_ADDR chunkIndices, GM_ADDR aqk, + GM_ADDR akk, GM_ADDR qg, GM_ADDR qgScaled, GM_ADDR wSeed, GM_ADDR uSeed, + GM_ADDR userWorkspace, const TilingData &tiling, TPipe &pipe) +{ + GM_ADDR aqkFp32 = userWorkspace + tiling.prepareAqkFp32Offset; + GM_ADDR akkFp32 = userWorkspace + tiling.prepareAkkFp32Offset; + GM_ADDR prepareScratch = userWorkspace + tiling.prepareScratchOffset; + + if ASCEND_IS_AIC { + ChunkKdaFwdPrepareKernel op; + op.Init(q, k, v, gk, beta, initialState, cuSeqlens, chunkIndices, nullptr, nullptr, nullptr, nullptr, aqk, + userWorkspace, aqkFp32, akkFp32, wSeed, akk, qg, qgScaled, uSeed, userWorkspace, prepareScratch, tiling, + &pipe, false); + op.ProcessAic(); + } + if ASCEND_IS_AIV { + ChunkKdaFwdPrepareKernel op; + op.Init(q, k, v, gk, beta, initialState, cuSeqlens, chunkIndices, nullptr, nullptr, nullptr, nullptr, aqk, + userWorkspace, aqkFp32, akkFp32, wSeed, akk, qg, qgScaled, uSeed, userWorkspace, prepareScratch, tiling, + &pipe); + op.ProcessAiv(); + } +} + +template +__aicore__ inline void RunChunkKdaPrepare(GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR gk, GM_ADDR, GM_ADDR, GM_ADDR, + GM_ADDR beta, GM_ADDR initialState, GM_ADDR cuSeqlens, GM_ADDR chunkIndices, + GM_ADDR aqk, GM_ADDR akk, GM_ADDR qg, GM_ADDR qgScaled, GM_ADDR wSeed, + GM_ADDR uSeed, GM_ADDR, GM_ADDR userWorkspace, const TilingData &tiling, + TPipe &pipe, bool = true) +{ + RunChunkKdaPrepare(q, k, v, gk, beta, initialState, cuSeqlens, chunkIndices, aqk, akk, + qg, qgScaled, wSeed, uSeed, userWorkspace, tiling, pipe); +} + +} // namespace KdaPrepare diff --git a/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_tiling_data.h b/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_tiling_data.h new file mode 100644 index 0000000000..a3d284b49e --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_tiling_data.h @@ -0,0 +1,99 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file chunk_kda_fwd_tiling_data.h + * \brief direct-launch tiling data of chunk_kda_fwd + * + * Plain-struct mirror of the aclnn ChunkKdaFwdTilingData (register/tilingdata_base.h + * TILING_DATA_FIELD_DEF layout). Field order/type must be kept in sync with the host + * side filling routine so that the packed blob copied GM->stack in the kernel entry + * matches exactly what the host serializes. + */ + +#ifndef CHUNK_KDA_FWD_TILING_DATA_H +#define CHUNK_KDA_FWD_TILING_DATA_H + +#include + +namespace ChunkKdaFwd { + +#pragma pack(push, 8) +struct ChunkKdaFwdTilingData { + int64_t batch; + int64_t seqNum; + int64_t qHeadNum; + int64_t vHeadNum; + int64_t seqlen; + int64_t kHeadDim; + int64_t vHeadDim; + int64_t chunkSize; + int64_t totalChunks; + int64_t inputRank; + + float scale; + float lowerBound; + + bool hasInitialState; + bool isVarLen; + bool safeGate; + bool inputSequenceMajor; + bool useGateInKernel; + bool hasALog; + bool hasDtBias; + bool computeGateInPrepare; + bool fusePostWu; + bool fusePostWuIntoFwdH; + bool useDenseFwdH; + bool storeFinalState; + bool storeGk; + bool storeW; + bool storeU; + bool storeQG; + bool storeKg; + bool storeVNew; + bool storeH; + + int64_t gateDataType; + int64_t gateUsedCoreNum; + int64_t prepareUsedCoreNum; + int64_t postWuUsedCoreNum; + int64_t outputUsedCoreNum; + + int64_t gkStorageOffset; + int64_t finalStateStorageOffset; + int64_t wStorageOffset; + int64_t uStorageOffset; + int64_t qgStorageOffset; + int64_t kgStorageOffset; + int64_t vNewStorageOffset; + int64_t hStorageOffset; + int64_t qgScaledOffset; + int64_t prepareAqkFp32Offset; + int64_t prepareAkkFp32Offset; + int64_t prepareScratchOffset; + int64_t postWuScratchOffset; + int64_t outputScratchOffset; + int64_t fwdHWorkspaceBaseOffset; + int64_t vWorkspaceOffset; + int64_t vUpdateWorkspaceOffset; + int64_t kDecayWorkspaceOffset; + int64_t hWorkspaceOffset; + int64_t numSeqWorkspaceOffset; + int64_t numChunksWorkspaceOffset; + + // direct-launch only: beta storage dtype (0 = fp32, 1 = bf16) + int64_t betaDataType; +}; +#pragma pack(pop) + +} // namespace ChunkKdaFwd + +#endif // CHUNK_KDA_FWD_TILING_DATA_H diff --git a/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_varlen.h b/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_varlen.h new file mode 100644 index 0000000000..69fabfb14f --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_varlen.h @@ -0,0 +1,78 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#pragma once + +#include "kernel_operator.h" + +namespace KdaVarlen { + +__aicore__ inline bool ResolveChunkRange( + const __gm__ int64_t *cuSeqlens, const __gm__ int64_t *chunkIndices, + uint64_t seqNum, uint64_t totalTokens, uint64_t chunkSize, + uint64_t flatChunk, uint64_t &seq, uint64_t &start, uint64_t &end) +{ + if (cuSeqlens == nullptr || chunkSize == 0) { + return false; + } + + int64_t seqValue = -1; + int64_t localChunkValue = -1; + if (chunkIndices != nullptr) { + const uint64_t metadataOffset = flatChunk * 2; + seqValue = chunkIndices[metadataOffset]; + localChunkValue = chunkIndices[metadataOffset + 1]; + } else { + uint64_t chunkPrefix = 0; + for (uint64_t seqIdx = 0; seqIdx < seqNum; ++seqIdx) { + const int64_t seqStart = cuSeqlens[seqIdx]; + const int64_t seqEnd = cuSeqlens[seqIdx + 1]; + if (seqStart < 0 || seqEnd < seqStart) { + return false; + } + const uint64_t seqLength = static_cast(seqEnd - seqStart); + const uint64_t seqChunks = (seqLength + chunkSize - 1) / chunkSize; + if (flatChunk < chunkPrefix + seqChunks) { + seqValue = static_cast(seqIdx); + localChunkValue = static_cast(flatChunk - chunkPrefix); + break; + } + chunkPrefix += seqChunks; + } + } + + if (seqValue < 0 || localChunkValue < 0 || + static_cast(seqValue) >= seqNum) { + return false; + } + const int64_t seqStartValue = cuSeqlens[seqValue]; + const int64_t seqEndValue = cuSeqlens[seqValue + 1]; + if (seqStartValue < 0 || seqEndValue < seqStartValue) { + return false; + } + const uint64_t seqStart = static_cast(seqStartValue); + const uint64_t seqEnd = static_cast(seqEndValue); + const uint64_t localChunk = static_cast(localChunkValue); + start = seqStart + localChunk * chunkSize; + if (start >= seqEnd || start >= totalTokens) { + return false; + } + end = start + chunkSize; + if (end > seqEnd) { + end = seqEnd; + } + if (end > totalTokens) { + return false; + } + seq = static_cast(seqValue); + return start < end; +} + +} // namespace KdaVarlen diff --git a/csrc/chunk_kda_fwd/op_kernel/kda_gate_cumsum_kernel.h b/csrc/chunk_kda_fwd/op_kernel/kda_gate_cumsum_kernel.h new file mode 100644 index 0000000000..886322d9db --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/kda_gate_cumsum_kernel.h @@ -0,0 +1,749 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#pragma once + +#include "kernel_operator.h" +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 +#ifndef FLA_NPU_REGBASE_HPP_INCLUDED +#define FLA_NPU_REGBASE_HPP_INCLUDED +#include "kernel_utils/vector/regbase.hpp" +#endif +#endif + +namespace KdaGateCumsum { + +using namespace AscendC; + +constexpr float RCP_LN2 = 1.4426950408889634f; +constexpr uint32_t GATE_ROW_ELEMENTS = 256; +constexpr uint32_t GATE_PIPELINE_DEPTH = 2; +constexpr uint32_t GATE_BULK_ROWS = 64; +constexpr uint32_t GATE_BULK_COLS = 128; +constexpr uint32_t GATE_BULK_ELEMENTS = GATE_BULK_ROWS * GATE_BULK_COLS; + +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 +static __simd_vf__ inline void AccumulateGateRowRegbase(__ubuf__ float *input, __ubuf__ float *acc, + __ubuf__ float *output, uint16_t count) +{ + using namespace AscendC::MicroAPI; + constexpr uint16_t FLOAT_ELEMENTS_PER_REG = AscendC::VECTOR_REG_WIDTH / sizeof(float); + constexpr uint16_t ELEMENTS_PER_PAIR = 2 * FLOAT_ELEMENTS_PER_REG; + + for (uint16_t offset = 0; offset < count; offset += ELEMENTS_PER_PAIR) { + RegTensor inputZeroReg; + RegTensor inputOneReg; + RegTensor accZeroReg; + RegTensor accOneReg; + MaskReg floatMask = CreateMask(); + + LoadAlign(inputZeroReg, input + offset); + LoadAlign(inputOneReg, input + offset + FLOAT_ELEMENTS_PER_REG); + LoadAlign(accZeroReg, acc + offset); + LoadAlign(accOneReg, acc + offset + FLOAT_ELEMENTS_PER_REG); + + Muls(inputZeroReg, inputZeroReg, RCP_LN2, floatMask); + Muls(inputOneReg, inputOneReg, RCP_LN2, floatMask); + Add(accZeroReg, accZeroReg, inputZeroReg, floatMask); + Add(accOneReg, accOneReg, inputOneReg, floatMask); + + StoreAlign(acc + offset, accZeroReg, floatMask); + StoreAlign(acc + offset + FLOAT_ELEMENTS_PER_REG, accOneReg, floatMask); + StoreAlign(output + offset, accZeroReg, floatMask); + StoreAlign(output + offset + FLOAT_ELEMENTS_PER_REG, accOneReg, floatMask); + } +} + +static __simd_vf__ inline void AccumulateGateChunk128Regbase(__ubuf__ float *input, + __ubuf__ float *output, + uint16_t rows) +{ + using namespace AscendC::MicroAPI; + constexpr uint16_t FLOAT_ELEMENTS_PER_REG = AscendC::VECTOR_REG_WIDTH / sizeof(float); + constexpr uint16_t ROW_ELEMENTS = 2 * FLOAT_ELEMENTS_PER_REG; + + MaskReg floatMask = CreateMask(); + RegTensor accZeroReg; + RegTensor accOneReg; + Duplicate(accZeroReg, 0.0f, floatMask); + Duplicate(accOneReg, 0.0f, floatMask); + for (uint16_t row = 0; row < rows; ++row) { + uint32_t rowOffset = static_cast(row) * ROW_ELEMENTS; + RegTensor inputZeroReg; + RegTensor inputOneReg; + LoadAlign(inputZeroReg, input + rowOffset); + LoadAlign( + inputOneReg, input + rowOffset + FLOAT_ELEMENTS_PER_REG); + Muls(inputZeroReg, inputZeroReg, RCP_LN2, floatMask); + Muls(inputOneReg, inputOneReg, RCP_LN2, floatMask); + Add(accZeroReg, accZeroReg, inputZeroReg, floatMask); + Add(accOneReg, accOneReg, inputOneReg, floatMask); + StoreAlign(output + rowOffset, accZeroReg, floatMask); + StoreAlign(output + rowOffset + FLOAT_ELEMENTS_PER_REG, accOneReg, floatMask); + } +} + +template +static __simd_vf__ inline void AccumulateSafeGateChunk128Regbase( + __ubuf__ float *input, __ubuf__ float *bias, __ubuf__ float *output, + uint16_t rows, float expA, float lowerBound) +{ + using namespace AscendC::MicroAPI; + constexpr uint16_t FLOAT_ELEMENTS_PER_REG = AscendC::VECTOR_REG_WIDTH / sizeof(float); + constexpr uint16_t ROW_ELEMENTS = 2 * FLOAT_ELEMENTS_PER_REG; + + MaskReg floatMask = CreateMask(); + RegTensor accZeroReg; + RegTensor accOneReg; + RegTensor oneZeroReg; + RegTensor oneOneReg; + RegTensor biasZeroReg; + RegTensor biasOneReg; + Duplicate(accZeroReg, 0.0f, floatMask); + Duplicate(accOneReg, 0.0f, floatMask); + Duplicate(oneZeroReg, 1.0f, floatMask); + Duplicate(oneOneReg, 1.0f, floatMask); + if constexpr (HAS_BIAS) { + LoadAlign(biasZeroReg, bias); + LoadAlign(biasOneReg, bias + FLOAT_ELEMENTS_PER_REG); + } + + const float gateScale = lowerBound * RCP_LN2; + for (uint16_t row = 0; row < rows; ++row) { + uint32_t rowOffset = static_cast(row) * ROW_ELEMENTS; + RegTensor gateZeroReg; + RegTensor gateOneReg; + RegTensor sigmoidZeroReg; + RegTensor sigmoidOneReg; + LoadAlign(gateZeroReg, input + rowOffset); + LoadAlign( + gateOneReg, input + rowOffset + FLOAT_ELEMENTS_PER_REG); + if constexpr (HAS_BIAS) { + Add(gateZeroReg, gateZeroReg, biasZeroReg, floatMask); + Add(gateOneReg, gateOneReg, biasOneReg, floatMask); + } + Muls(gateZeroReg, gateZeroReg, -expA, floatMask); + Muls(gateOneReg, gateOneReg, -expA, floatMask); + Exp(gateZeroReg, gateZeroReg, floatMask); + Exp(gateOneReg, gateOneReg, floatMask); + Adds(gateZeroReg, gateZeroReg, 1.0f, floatMask); + Adds(gateOneReg, gateOneReg, 1.0f, floatMask); + Div(sigmoidZeroReg, oneZeroReg, gateZeroReg, floatMask); + Div(sigmoidOneReg, oneOneReg, gateOneReg, floatMask); + Muls(sigmoidZeroReg, sigmoidZeroReg, gateScale, floatMask); + Muls(sigmoidOneReg, sigmoidOneReg, gateScale, floatMask); + Add(accZeroReg, accZeroReg, sigmoidZeroReg, floatMask); + Add(accOneReg, accOneReg, sigmoidOneReg, floatMask); + StoreAlign(output + rowOffset, accZeroReg, floatMask); + StoreAlign(output + rowOffset + FLOAT_ELEMENTS_PER_REG, accOneReg, floatMask); + } +} + +#endif + +template +class KdaGateCumsumKernel { +public: + template + __aicore__ inline void Init(GM_ADDR g, GM_ADDR aLog, GM_ADDR dtBias, GM_ADDR cuSeqlens, GM_ADDR gk, + const TilingData &tiling, TPipe *pipe) + { + g_.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(g)); + aLog_.SetGlobalBuffer(reinterpret_cast<__gm__ float *>(aLog)); + dtBias_.SetGlobalBuffer(reinterpret_cast<__gm__ float *>(dtBias)); + cuSeqlens_.SetGlobalBuffer(reinterpret_cast<__gm__ int64_t *>(cuSeqlens)); + gk_.SetGlobalBuffer(reinterpret_cast<__gm__ float *>(gk)); + pipe_ = pipe; + batch_ = static_cast(tiling.batch); + t_ = static_cast(tiling.t); + hv_ = static_cast(tiling.hv); + k_ = static_cast(tiling.k); + rank_ = static_cast(tiling.rank); + chunkSize_ = static_cast(tiling.chunkSize); + seqNum_ = static_cast(tiling.seqNum); + hasCuSeqlens_ = tiling.hasCuSeqlens != 0; + hasALog_ = tiling.hasALog != 0; + hasDtBias_ = tiling.hasDtBias != 0; + inputSequenceMajor_ = tiling.inputSequenceMajor != 0; + lowerBound_ = tiling.lowerBound; + usedCoreNum_ = static_cast(tiling.usedCoreNum); + maxChunks_ = (t_ + chunkSize_ - 1) / chunkSize_; + pipe_->InitBuffer(rowBuf_, GATE_ROW_ELEMENTS * sizeof(float)); + pipe_->InitBuffer(accBuf_, GATE_ROW_ELEMENTS * sizeof(float)); + pipe_->InitBuffer(outBuf_, GATE_PIPELINE_DEPTH * GATE_ROW_ELEMENTS * sizeof(float)); + pipe_->InitBuffer(tmpBuf_, GATE_ROW_ELEMENTS * sizeof(float)); + pipe_->InitBuffer(oneBuf_, GATE_ROW_ELEMENTS * sizeof(float)); + pipe_->InitBuffer(biasBuf_, GATE_ROW_ELEMENTS * sizeof(float)); + pipe_->InitBuffer(inBuf_, GATE_PIPELINE_DEPTH * GATE_ROW_ELEMENTS * sizeof(T)); + pipe_->InitBuffer(chunkBuf_, 2 * GATE_BULK_ELEMENTS * sizeof(float)); + pipe_->InitBuffer(scalarBuf_, 32); + pipe_->InitBuffer(scalarI64Buf_, 32); + AllocEvents(); + } + + __aicore__ inline void Process() + { + uint64_t taskCount = hasCuSeqlens_ ? seqNum_ * hv_ : batch_ * hv_ * maxChunks_; + uint64_t coreIdx = static_cast(GetBlockIdx()); + for (uint64_t task = coreIdx; task < taskCount; task += usedCoreNum_) { + ProcessTask(task); + } + ReleaseEvents(); + } + +private: + __aicore__ inline void AllocEvents() + { + for (uint32_t slot = 0; slot < GATE_PIPELINE_DEPTH; ++slot) { + inputMte2ToVEvent_[slot] = pipe_->AllocEventID(); + inputVToMte2Event_[slot] = pipe_->AllocEventID(); + outputVToMte3Event_[slot] = pipe_->AllocEventID(); + outputMte3ToVEvent_[slot] = pipe_->AllocEventID(); + } + auxMte2ToVEvent_ = pipe_->AllocEventID(); + scalarVToSEvent_ = pipe_->AllocEventID(); + bulkMte3ToMte2Event_ = pipe_->AllocEventID(); + } + + __aicore__ inline void ReleaseEvents() + { + for (uint32_t slot = 0; slot < GATE_PIPELINE_DEPTH; ++slot) { + pipe_->ReleaseEventID(inputMte2ToVEvent_[slot]); + pipe_->ReleaseEventID(inputVToMte2Event_[slot]); + pipe_->ReleaseEventID(outputVToMte3Event_[slot]); + pipe_->ReleaseEventID(outputMte3ToVEvent_[slot]); + } + pipe_->ReleaseEventID(auxMte2ToVEvent_); + pipe_->ReleaseEventID(scalarVToSEvent_); + pipe_->ReleaseEventID(bulkMte3ToMte2Event_); + } + + __aicore__ inline uint64_t InputOffset(uint64_t b, uint64_t t, uint64_t hv, uint64_t k) const + { + if (inputSequenceMajor_) { + if (rank_ == 4) { + return ((b * t_ + t) * hv_ + hv) * k_ + k; + } + return (t * hv_ + hv) * k_ + k; + } + if (rank_ == 4) { + return ((b * hv_ + hv) * t_ + t) * k_ + k; + } + return (hv * t_ + t) * k_ + k; + } + + __aicore__ inline uint64_t OutputOffset(uint64_t b, uint64_t t, uint64_t hv, uint64_t k) const + { + if (rank_ == 4) { + return ((b * hv_ + hv) * t_ + t) * k_ + k; + } + return (hv * t_ + t) * k_ + k; + } + + __aicore__ inline void CopyVectorIn(LocalTensor &dst, GlobalTensor &src, uint64_t offset, uint64_t count) + { + uint64_t rowBytes = count * static_cast(sizeof(T)); + if (rowBytes >= 32 && rowBytes % 32 == 0) { + DataCopy(dst, src[offset], static_cast(count)); + } else { + DataCopyParams params{1, static_cast(rowBytes), 0, 0}; + DataCopyPadParams padParams{false, 0, 0, 0}; + DataCopyPad(dst, src[offset], params, padParams); + } + } + + __aicore__ inline void CopyFloatVectorIn(LocalTensor &dst, GlobalTensor &src, uint64_t offset, + uint64_t count) + { + uint64_t rowBytes = count * sizeof(float); + if (rowBytes >= 32 && rowBytes % 32 == 0) { + DataCopy(dst, src[offset], static_cast(count)); + } else { + DataCopyParams params{1, static_cast(rowBytes), 0, 0}; + DataCopyPadParams padParams{false, 0, 0, 0}; + DataCopyPad(dst, src[offset], params, padParams); + } + } + + __aicore__ inline void CopyFloatVectorOut(GlobalTensor &dst, uint64_t offset, LocalTensor &src, + uint64_t count) + { + uint64_t rowBytes = count * sizeof(float); + if (rowBytes >= 32 && rowBytes % 32 == 0) { + DataCopy(dst[offset], src, static_cast(count)); + } else { + DataCopyParams params{1, static_cast(rowBytes), 0, 0}; + DataCopyPad(dst[offset], src, params); + } + } + + __aicore__ inline void CopyGateRowsIn(LocalTensor &dst, uint64_t b, uint64_t start, + uint64_t hv, uint64_t rows) + { + if (!inputSequenceMajor_) { + CopyVectorIn(dst, g_, InputOffset(b, start, hv, 0), rows * k_); + return; + } + DataCopyExtParams params{ + static_cast(rows), + static_cast(k_ * sizeof(T)), + static_cast((hv_ - 1) * k_ * sizeof(T)), + 0, + 0}; + DataCopyPadExtParams padParams{false, 0, 0, 0}; + DataCopyPad(dst, g_[InputOffset(b, start, hv, 0)], params, padParams); + } + + __aicore__ inline void PrefetchGateRow(uint64_t offset, uint32_t slot) + { + LocalTensor input = inBuf_.Get()[slot * GATE_ROW_ELEMENTS]; + CopyVectorIn(input, g_, offset, k_); + SetFlag(inputMte2ToVEvent_[slot]); + } + + __aicore__ inline void MaterializeGateRow(uint32_t slot, LocalTensor &row) + { + LocalTensor input = inBuf_.Get()[slot * GATE_ROW_ELEMENTS]; + if constexpr (IsSameType::value) { + Adds(row, input, 0.0f, static_cast(k_)); + } else { + Cast(row, input, RoundMode::CAST_NONE, static_cast(k_)); + } + PipeBarrier(); + } + + __aicore__ inline float ReadFloat(GlobalTensor &tensor, uint64_t offset) + { + LocalTensor scalar = scalarBuf_.Get(); + DataCopyParams params{1, static_cast(sizeof(float)), 0, 0}; + DataCopyPadParams padParams{false, 0, 0, 0}; + DataCopyPad(scalar, tensor[offset], params, padParams); + SetFlag(auxMte2ToVEvent_); + WaitFlag(auxMte2ToVEvent_); + Adds(scalar, scalar, 0.0f, 1); + PipeBarrier(); + SetFlag(scalarVToSEvent_); + WaitFlag(scalarVToSEvent_); + __ubuf__ float *ptr = (__ubuf__ float *)scalar.GetPhyAddr(); + return ptr[0]; + } + + __aicore__ inline int64_t ReadInt64(GlobalTensor &tensor, uint64_t offset) + { + LocalTensor scalar = scalarI64Buf_.Get(); + DataCopyParams params{1, static_cast(sizeof(int64_t)), 0, 0}; + DataCopyPadParams padParams{false, 0, 0, 0}; + DataCopyPad(scalar, tensor[offset], params, padParams); + SetFlag(auxMte2ToVEvent_); + WaitFlag(auxMte2ToVEvent_); + SetFlag(scalarVToSEvent_); + WaitFlag(scalarVToSEvent_); + __ubuf__ int64_t *ptr = (__ubuf__ int64_t *)scalar.GetPhyAddr(); + return ptr[0]; + } + + __aicore__ inline float ExpScalar(float x) + { + LocalTensor scalar = scalarBuf_.Get(); + Duplicate(scalar, x, 1); + PipeBarrier(); + Exp(scalar, scalar, 1); + PipeBarrier(); + SetFlag(scalarVToSEvent_); + WaitFlag(scalarVToSEvent_); + __ubuf__ float *ptr = (__ubuf__ float *)scalar.GetPhyAddr(); + return ptr[0]; + } + + __aicore__ inline void PrepareGate(uint64_t hv) + { + if constexpr (USE_GATE_IN_KERNEL) { + expA_ = ExpScalar(ReadFloat(aLog_, hv)); + if (hasDtBias_) { + LocalTensor bias = biasBuf_.Get(); + CopyFloatVectorIn(bias, dtBias_, hv * k_, k_); + SetFlag(auxMte2ToVEvent_); + WaitFlag(auxMte2ToVEvent_); + } + } + } + + __aicore__ inline void ApplyGate(LocalTensor &row) + { + if constexpr (USE_GATE_IN_KERNEL) { + if (hasDtBias_) { + LocalTensor bias = biasBuf_.Get(); + Add(row, row, bias, static_cast(k_)); + PipeBarrier(); + } + if constexpr (SAFE_GATE) { + Muls(row, row, expA_, static_cast(k_)); + PipeBarrier(); + + LocalTensor tmp = tmpBuf_.Get(); + Muls(tmp, row, -1.0f, static_cast(k_)); + PipeBarrier(); + Exp(tmp, tmp, static_cast(k_)); + PipeBarrier(); + Adds(tmp, tmp, 1.0f, static_cast(k_)); + PipeBarrier(); + + LocalTensor one = oneBuf_.Get(); + Duplicate(one, 1.0f, static_cast(k_)); + PipeBarrier(); + Div(row, one, tmp, static_cast(k_)); + PipeBarrier(); + Muls(row, row, lowerBound_, static_cast(k_)); + PipeBarrier(); + } else { + LocalTensor positive = oneBuf_.Get(); + LocalTensor tmp = tmpBuf_.Get(); + Maxs(positive, row, 0.0f, static_cast(k_)); + Abs(tmp, row, static_cast(k_)); + PipeBarrier(); + Muls(tmp, tmp, -1.0f, static_cast(k_)); + PipeBarrier(); + Exp(tmp, tmp, static_cast(k_)); + PipeBarrier(); + Adds(tmp, tmp, 1.0f, static_cast(k_)); + PipeBarrier(); + Ln(tmp, tmp, static_cast(k_)); + PipeBarrier(); + Add(row, positive, tmp, static_cast(k_)); + PipeBarrier(); + Muls(row, row, -expA_, static_cast(k_)); + PipeBarrier(); + } + } + } + + __aicore__ inline void ProcessTask(uint64_t task) + { + uint64_t hv = hasCuSeqlens_ ? task % hv_ : (task / maxChunks_) % hv_; + PrepareGate(hv); + if (!hasCuSeqlens_) { + uint64_t chunk = task % maxChunks_; + uint64_t b = task / (maxChunks_ * hv_); + uint64_t start = chunk * chunkSize_; + uint64_t end = start + chunkSize_; + if (end > t_) { + end = t_; + } + ProcessChunk(b, hv, start, end); + return; + } + uint64_t seq = task / hv_; + uint64_t seqStart = static_cast(ReadInt64(cuSeqlens_, seq)); + uint64_t seqEnd = static_cast(ReadInt64(cuSeqlens_, seq + 1)); + for (uint64_t start = seqStart; start < seqEnd; start += chunkSize_) { + uint64_t end = start + chunkSize_; + if (end > seqEnd) { + end = seqEnd; + } + ProcessChunk(0, hv, start, end); + } + } + + __aicore__ inline void ProcessChunk(uint64_t b, uint64_t hv, uint64_t start, uint64_t end) + { +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 220 + if constexpr (USE_GATE_IN_KERNEL && SAFE_GATE && IsSameType::value) { + if (k_ == GATE_BULK_COLS && chunkSize_ == GATE_BULK_ROWS && end - start == GATE_BULK_ROWS) { + ProcessChunkBulkFp32SafeVec(b, hv, start); + return; + } + } +#endif +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if constexpr (USE_GATE_IN_KERNEL && SAFE_GATE && IsSameType::value) { + if (k_ == GATE_BULK_COLS && chunkSize_ == GATE_BULK_ROWS) { + ProcessChunkBulkFp32Safe(b, hv, start, end); + return; + } + } + if constexpr (!USE_GATE_IN_KERNEL && IsSameType::value) { + if (k_ == GATE_BULK_COLS && chunkSize_ == GATE_BULK_ROWS) { + ProcessChunkBulkFp32(b, hv, start, end); + return; + } + } +#endif + LocalTensor acc = accBuf_.Get(); + LocalTensor row = rowBuf_.Get(); + Duplicate(acc, 0.0f, static_cast(k_)); + PipeBarrier(); + uint64_t rows = end - start; + if (rows == 0) { + return; + } + + PrefetchGateRow(InputOffset(b, start, hv, 0), 0); + + for (uint64_t rowIdx = 0; rowIdx < rows; ++rowIdx) { + uint64_t token = start + rowIdx; + uint32_t slot = static_cast(rowIdx & 1); + WaitFlag(inputMte2ToVEvent_[slot]); + + if (rowIdx + 1 < rows) { + uint32_t nextSlot = slot ^ 1; + if (rowIdx >= 1) { + WaitFlag(inputVToMte2Event_[nextSlot]); + } + PrefetchGateRow(InputOffset(b, token + 1, hv, 0), nextSlot); + } + + uint32_t outputSlot = slot; + if (rowIdx >= GATE_PIPELINE_DEPTH) { + WaitFlag(outputMte3ToVEvent_[outputSlot]); + } + LocalTensor output = + outBuf_.Get()[outputSlot * GATE_ROW_ELEMENTS]; + +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if constexpr (!USE_GATE_IN_KERNEL && IsSameType::value) { + if ((k_ % 128) == 0) { + LocalTensor input = inBuf_.Get()[slot * GATE_ROW_ELEMENTS]; + AccumulateGateRowRegbase( + (__ubuf__ float *)input.GetPhyAddr(), (__ubuf__ float *)acc.GetPhyAddr(), + (__ubuf__ float *)output.GetPhyAddr(), static_cast(k_)); + PipeBarrier(); + } else { + MaterializeGateRow(slot, row); + Muls(row, row, RCP_LN2, static_cast(k_)); + PipeBarrier(); + Add(acc, acc, row, static_cast(k_)); + PipeBarrier(); + Adds(output, acc, 0.0f, static_cast(k_)); + } + } else { + MaterializeGateRow(slot, row); + ApplyGate(row); + Muls(row, row, RCP_LN2, static_cast(k_)); + PipeBarrier(); + Add(acc, acc, row, static_cast(k_)); + PipeBarrier(); + Adds(output, acc, 0.0f, static_cast(k_)); + } +#else + MaterializeGateRow(slot, row); + ApplyGate(row); + Muls(row, row, RCP_LN2, static_cast(k_)); + PipeBarrier(); + Add(acc, acc, row, static_cast(k_)); + PipeBarrier(); + Adds(output, acc, 0.0f, static_cast(k_)); +#endif + + PipeBarrier(); + SetFlag(inputVToMte2Event_[slot]); + SetFlag(outputVToMte3Event_[outputSlot]); + WaitFlag(outputVToMte3Event_[outputSlot]); + CopyFloatVectorOut(gk_, OutputOffset(b, token, hv, 0), output, k_); + SetFlag(outputMte3ToVEvent_[outputSlot]); + } + + uint64_t drainStart = rows > GATE_PIPELINE_DEPTH ? rows - GATE_PIPELINE_DEPTH : 0; + for (uint64_t rowIdx = drainStart; rowIdx < rows; ++rowIdx) { + uint32_t outputSlot = static_cast(rowIdx & 1); + WaitFlag(inputVToMte2Event_[outputSlot]); + WaitFlag(outputMte3ToVEvent_[outputSlot]); + } + } + +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 220 + __aicore__ inline void ScanGateChunk64x128(LocalTensor input, LocalTensor output) + { + constexpr uint32_t rowElements = GATE_BULK_COLS; + Adds(input, output, 0.0f, rowElements); + Add(input[rowElements], output[rowElements], output, (GATE_BULK_ROWS - 1) * rowElements); + PipeBarrier(); + + Adds(output, input, 0.0f, 2 * rowElements); + Add(output[2 * rowElements], input[2 * rowElements], input, (GATE_BULK_ROWS - 2) * rowElements); + PipeBarrier(); + + Adds(input, output, 0.0f, 4 * rowElements); + Add(input[4 * rowElements], output[4 * rowElements], output, (GATE_BULK_ROWS - 4) * rowElements); + PipeBarrier(); + + Adds(output, input, 0.0f, 8 * rowElements); + Add(output[8 * rowElements], input[8 * rowElements], input, (GATE_BULK_ROWS - 8) * rowElements); + PipeBarrier(); + + Adds(input, output, 0.0f, 16 * rowElements); + Add(input[16 * rowElements], output[16 * rowElements], output, (GATE_BULK_ROWS - 16) * rowElements); + PipeBarrier(); + + Adds(output, input, 0.0f, 32 * rowElements); + Add(output[32 * rowElements], input[32 * rowElements], input, (GATE_BULK_ROWS - 32) * rowElements); + PipeBarrier(); + } + + __aicore__ inline void ProcessChunkBulkFp32SafeVec(uint64_t b, uint64_t hv, uint64_t start) + { + constexpr uint32_t elems = GATE_BULK_ELEMENTS; + LocalTensor input = chunkBuf_.Get(); + LocalTensor output = chunkBuf_.Get()[GATE_BULK_ELEMENTS]; + CopyGateRowsIn(input, b, start, hv, GATE_BULK_ROWS); + SetFlag(inputMte2ToVEvent_[0]); + WaitFlag(inputMte2ToVEvent_[0]); + + Duplicate(output, 1.0f, elems); + if (hasDtBias_) { + LocalTensor bias = biasBuf_.Get(); + for (uint32_t row = 0; row < GATE_BULK_ROWS; ++row) { + Add(input[row * GATE_BULK_COLS], input[row * GATE_BULK_COLS], bias, GATE_BULK_COLS); + } + PipeBarrier(); + } + Muls(input, input, -expA_, elems); + PipeBarrier(); + Exp(input, input, elems); + PipeBarrier(); + Adds(input, input, 1.0f, elems); + PipeBarrier(); + Div(output, output, input, elems); + PipeBarrier(); + Muls(output, output, lowerBound_ * RCP_LN2, elems); + PipeBarrier(); + ScanGateChunk64x128(input, output); + + SetFlag(outputVToMte3Event_[0]); + WaitFlag(outputVToMte3Event_[0]); + CopyFloatVectorOut(gk_, OutputOffset(b, start, hv, 0), output, elems); + SetFlag(bulkMte3ToMte2Event_); + WaitFlag(bulkMte3ToMte2Event_); + SetFlag(outputMte3ToVEvent_[0]); + WaitFlag(outputMte3ToVEvent_[0]); + } +#endif + +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + __aicore__ inline void ProcessChunkBulkFp32Safe(uint64_t b, uint64_t hv, uint64_t start, uint64_t end) + { + uint64_t rows = end - start; + if (rows == 0) { + return; + } + uint32_t elems = static_cast(rows * k_); + LocalTensor input = chunkBuf_.Get(); + LocalTensor output = chunkBuf_.Get()[GATE_BULK_ELEMENTS]; + CopyGateRowsIn(input, b, start, hv, rows); + SetFlag(inputMte2ToVEvent_[0]); + WaitFlag(inputMte2ToVEvent_[0]); + LocalTensor bias = biasBuf_.Get(); + if (hasDtBias_) { + AccumulateSafeGateChunk128Regbase( + (__ubuf__ float *)input.GetPhyAddr(), (__ubuf__ float *)bias.GetPhyAddr(), + (__ubuf__ float *)output.GetPhyAddr(), static_cast(rows), expA_, lowerBound_); + } else { + AccumulateSafeGateChunk128Regbase( + (__ubuf__ float *)input.GetPhyAddr(), (__ubuf__ float *)bias.GetPhyAddr(), + (__ubuf__ float *)output.GetPhyAddr(), static_cast(rows), expA_, lowerBound_); + } + PipeBarrier(); + SetFlag(outputVToMte3Event_[0]); + WaitFlag(outputVToMte3Event_[0]); + CopyFloatVectorOut(gk_, OutputOffset(b, start, hv, 0), output, elems); + SetFlag(bulkMte3ToMte2Event_); + WaitFlag(bulkMte3ToMte2Event_); + SetFlag(outputMte3ToVEvent_[0]); + WaitFlag(outputMte3ToVEvent_[0]); + } + + __aicore__ inline void ProcessChunkBulkFp32(uint64_t b, uint64_t hv, uint64_t start, uint64_t end) + { + uint64_t rows = end - start; + if (rows == 0) { + return; + } + uint32_t elems = static_cast(rows * k_); + LocalTensor buffer0 = chunkBuf_.Get(); + CopyGateRowsIn(buffer0, b, start, hv, rows); + SetFlag(inputMte2ToVEvent_[0]); + WaitFlag(inputMte2ToVEvent_[0]); + AccumulateGateChunk128Regbase( + (__ubuf__ float *)buffer0.GetPhyAddr(), (__ubuf__ float *)buffer0.GetPhyAddr(), + static_cast(rows)); + PipeBarrier(); + SetFlag(outputVToMte3Event_[0]); + WaitFlag(outputVToMte3Event_[0]); + CopyFloatVectorOut(gk_, OutputOffset(b, start, hv, 0), buffer0, elems); + SetFlag(bulkMte3ToMte2Event_); + WaitFlag(bulkMte3ToMte2Event_); + SetFlag(outputMte3ToVEvent_[0]); + WaitFlag(outputMte3ToVEvent_[0]); + } +#endif + + GlobalTensor g_; + GlobalTensor aLog_; + GlobalTensor dtBias_; + GlobalTensor cuSeqlens_; + GlobalTensor gk_; + TPipe *pipe_ = nullptr; + TBuf rowBuf_; + TBuf accBuf_; + TBuf outBuf_; + TBuf tmpBuf_; + TBuf oneBuf_; + TBuf biasBuf_; + TBuf inBuf_; + TBuf chunkBuf_; + TBuf scalarBuf_; + TBuf scalarI64Buf_; + TEventID inputMte2ToVEvent_[GATE_PIPELINE_DEPTH]; + TEventID inputVToMte2Event_[GATE_PIPELINE_DEPTH]; + TEventID outputVToMte3Event_[GATE_PIPELINE_DEPTH]; + TEventID outputMte3ToVEvent_[GATE_PIPELINE_DEPTH]; + TEventID auxMte2ToVEvent_; + TEventID scalarVToSEvent_; + TEventID bulkMte3ToMte2Event_; + uint64_t batch_ = 0; + uint64_t t_ = 0; + uint64_t hv_ = 0; + uint64_t k_ = 0; + uint64_t rank_ = 0; + uint64_t chunkSize_ = 0; + uint64_t seqNum_ = 0; + uint64_t maxChunks_ = 0; + bool hasCuSeqlens_ = false; + bool hasALog_ = false; + bool hasDtBias_ = false; + bool inputSequenceMajor_ = false; + float expA_ = 1.0f; + float lowerBound_ = -5.0f; + uint64_t usedCoreNum_ = 1; +}; + +template +__aicore__ inline void RunKdaGateCumsum(GM_ADDR g, GM_ADDR aLog, GM_ADDR dtBias, GM_ADDR cuSeqlens, GM_ADDR gk, + const TilingData &tilingData, TPipe *pipe) +{ + KdaGateCumsumKernel op; + op.Init(g, aLog, dtBias, cuSeqlens, gk, tilingData, pipe); + op.Process(); +} + +template +__aicore__ inline void DispatchKdaGateCumsum(GM_ADDR g, GM_ADDR aLog, GM_ADDR dtBias, GM_ADDR cuSeqlens, + GM_ADDR gk, const TilingData &tilingData, TPipe *pipe) +{ + if (tilingData.useGateInKernel != 0) { + if (tilingData.safeGate != 0) { + RunKdaGateCumsum(g, aLog, dtBias, cuSeqlens, gk, tilingData, pipe); + } else { + RunKdaGateCumsum(g, aLog, dtBias, cuSeqlens, gk, tilingData, pipe); + } + } else { + RunKdaGateCumsum(g, aLog, dtBias, cuSeqlens, gk, tilingData, pipe); + } +} +} // namespace KdaGateCumsum diff --git a/csrc/chunk_kda_fwd/op_kernel/kernel_utils/block/block_mmad_pingpong_tla_kernel_utils.hpp b/csrc/chunk_kda_fwd/op_kernel/kernel_utils/block/block_mmad_pingpong_tla_kernel_utils.hpp new file mode 100644 index 0000000000..f099eb38d3 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/kernel_utils/block/block_mmad_pingpong_tla_kernel_utils.hpp @@ -0,0 +1,1014 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef COMMON_BLOCK_MMAD_PINGPONG_TLA_KERNEL_UTILS_HPP +#define COMMON_BLOCK_MMAD_PINGPONG_TLA_KERNEL_UTILS_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/kda_coord.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" +#include "catlass/gemm/kda_gemm_helper.hpp" +#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" +#include "catlass/gemm/tile/kda_tile_mmad.hpp" +#include "../tile/copy_l0c_to_ub.hpp" +#include "tla/kda_tla_layout.hpp" +#include "tla/kda_tensor.hpp" + +namespace Common { + +template , + class TileMmad = Catlass::Gemm::Tile::TileMmadTla > +struct BlockMmadTla { + static_assert(DEPENDENT_FALSE, "BlockMmadTla is not implemented for this DispatchPolicy"); +}; + +// Now ENABLE_UNIT_FLAG_ must be false when intput element is int8 +template +struct MmadPingpong : public Catlass::Gemm::MmadBase { + static constexpr uint32_t L1A_STAGES = L1A_STAGES_; + static constexpr uint32_t L1B_STAGES = L1B_STAGES_; + static constexpr uint32_t L0A_STAGES = L0A_STAGES_; + static constexpr uint32_t L0B_STAGES = L0B_STAGES_; + static constexpr uint32_t L0C_STAGES = L0C_STAGES_; + static constexpr uint32_t UB_STAGES = UB_STAGES_; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; + static constexpr bool USE_HF32_MODE = USE_HF32_MODE_; + static constexpr bool ENABLE_L1_RESIDENT = ENABLE_L1_RESIDENT_; +}; + +template +struct BlockMmadTla, + L1TileShape_, L0TileShape_, ElementA_, ElementB_, ElementC_, ElementBias_, TileCopy_, TileMmad_> { +public: + // Type Aliases + using DispatchPolicy = MmadPingpong; + using ArchTag = typename DispatchPolicy::ArchTag; + using TileCopy = TileCopy_; + using L1TileShape = L1TileShape_; + using L0TileShape = L0TileShape_; + using ElementA = ElementA_; + using LayoutA = typename TileCopy::LayoutA; + using ElementB = ElementB_; + using LayoutB = typename TileCopy::LayoutB; + using ElementC = ElementC_; + using LayoutC = typename TileCopy::LayoutC; + using ElementBias = ElementBias_; + + using TileMmad = TileMmad_; + + using CopyL1ToL0A = typename TileCopy::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopy::CopyL1ToL0B; + using CopyL1ToBT = typename TileCopy::CopyL1ToBT; + + using ElementAccumulator = typename TileCopy::ElementAccumulator; + + static constexpr bool HAS_BIAS = TileCopy::HAS_BIAS; + + using LayoutTagL1A = typename TileCopy::LayoutTagL1A; + using LayoutTagL1B = typename TileCopy::LayoutTagL1B; + using LayoutTagL0A = typename TileCopy::LayoutTagL0A; + using LayoutTagL0B = typename TileCopy::LayoutTagL0B; + + using L1AAlignHelper = typename TileCopy_::L1AAlignHelper; + using L1BAlignHelper = typename TileCopy_::L1BAlignHelper; + + static_assert(tla::is_tuple::value && tla::is_static::value, + "L1TileShape must be tla::tuple and static!"); + static_assert(tla::is_tuple::value && tla::is_static::value, + "L0TileShape must be tla::tuple and static!"); + + static constexpr uint64_t FLAG_ID_MAX = 16; + static constexpr bool ENABLE_UNIT_FLAG = DispatchPolicy::ENABLE_UNIT_FLAG; + static constexpr bool USE_HF32_MODE = DispatchPolicy::USE_HF32_MODE; + static constexpr bool ENABLE_L1_RESIDENT = DispatchPolicy::ENABLE_L1_RESIDENT; + static constexpr uint32_t L1A_STAGES = DispatchPolicy::L1A_STAGES; + static constexpr uint32_t L1B_STAGES = DispatchPolicy::L1B_STAGES; + static constexpr uint32_t L0A_STAGES = DispatchPolicy::L0A_STAGES; + static constexpr uint32_t L0B_STAGES = DispatchPolicy::L0B_STAGES; + static constexpr uint32_t L0C_STAGES = DispatchPolicy::L0C_STAGES; + static constexpr uint32_t L1_TILE_M = tla::get<0>(L1TileShape{}); + static constexpr uint32_t L1_TILE_N = tla::get<1>(L1TileShape{}); + static constexpr uint32_t L1_TILE_K = tla::get<2>(L1TileShape{}); + static constexpr uint32_t L0_TILE_M = tla::get<0>(L0TileShape{}); + static constexpr uint32_t L0_TILE_N = tla::get<1>(L0TileShape{}); + static constexpr uint32_t L0_TILE_K = tla::get<2>(L0TileShape{}); + static constexpr uint32_t UB_STAGES = UB_STAGES_; + static constexpr uint32_t MAX_CUBE_VEC_SYNC_NUM = 5; + + // L1 tile size + static constexpr uint32_t L1A_TILE_SIZE = L1_TILE_M * L1_TILE_K * sizeof(ElementA); + static constexpr uint32_t L1B_TILE_SIZE = L1_TILE_N * L1_TILE_K * sizeof(ElementB); + // L0 tile size + static constexpr uint32_t L0A_TILE_SIZE = L0_TILE_M * L0_TILE_K * sizeof(ElementA); + static constexpr uint32_t L0B_TILE_SIZE = L0_TILE_K * L0_TILE_N * sizeof(ElementB); + static constexpr uint32_t L0C_TILE_SIZE = L1_TILE_M * L1_TILE_N * sizeof(ElementAccumulator); + + // Check HF32_MODE + static_assert(!USE_HF32_MODE || + (USE_HF32_MODE && std::is_same_v && std::is_same_v), + "HF32 MODE only supports in float!"); + + // Check L0C_STAGES + static_assert(!(ENABLE_UNIT_FLAG && L0C_STAGES != 1), "L0C_STAGES must be 1 when UnitFlag is true!"); + + // Check LayoutC + static_assert(tla::detail::isRowMajor::value || + ((std::is_same_v || std::is_same_v || + std::is_same_v) && + tla::detail::iszN::value), + "LayoutC only supports zN in half or bfloat16 or float, RowMajor in all dtype yet!"); + + // Check L1TileShape + static_assert(L1A_TILE_SIZE * L1A_STAGES + L1B_TILE_SIZE * L1B_STAGES <= ArchTag::L1_SIZE, + "L1TileShape exceeding the L1 space!"); + + // Check L0TileShape + static_assert(L0A_TILE_SIZE * L0A_STAGES <= ArchTag::L0A_SIZE, "L0TileShape exceeding the L0A space!"); + static_assert(L0B_TILE_SIZE * L0B_STAGES <= ArchTag::L0B_SIZE, "L0TileShape exceeding the L0B space!"); + static_assert(L0C_TILE_SIZE * L0C_STAGES <= ArchTag::L0C_SIZE, "L0TileShape exceeding the L0C space!"); + + static constexpr uint32_t _32B = 32 * 8; // in bits + static_assert(L1_TILE_M == L0_TILE_M && L1_TILE_N == L0_TILE_N, + "The situation where the basic blocks of L1 and L0 differ on the m and n axes is not supported yet"); + static_assert(L0_TILE_K <= L1_TILE_K, "L0TileShape::K cannot exceed L1TileShape::K"); +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 2201) + static_assert(L1_TILE_M * SizeOfBits::value % _32B == 0, "L1TileShape::M must be 32B aligned."); + static_assert(L1_TILE_K * SizeOfBits::value % _32B == 0, "L1TileShape::K must be 32B aligned."); + static_assert(L1_TILE_K * SizeOfBits::value % _32B == 0, "L1TileShape::K must be 32B aligned."); + static_assert(L1_TILE_N * SizeOfBits::value % _32B == 0, "L1TileShape::N must be 32B aligned."); + static_assert(L0_TILE_K * SizeOfBits::value % _32B == 0, "L0TileShape::K must be 32B aligned."); +#endif + + static_assert((!HAS_BIAS && (L1A_STAGES + L1B_STAGES) <= 8) || (HAS_BIAS && (L1A_STAGES + L1B_STAGES) <= 7), + "L1 Buffer overflow: Exceeds the supported range of EVENT(0~7)"); + + static_assert((!HAS_BIAS && (L0A_STAGES + L0B_STAGES) <= 8) || (HAS_BIAS && (L0A_STAGES + L0B_STAGES) <= 7), + "L0 Buffer overflow: Exceeds the supported range of EVENT_ID(0~7)"); + + static constexpr auto L1A_LAYOUT = + tla::MakeLayout(tla::Int{}, tla::Int{}); + static constexpr auto L1B_LAYOUT = + tla::MakeLayout(tla::Int{}, tla::Int{}); + static constexpr auto L1BIAS_LAYOUT = tla::MakeLayout(tla::Int{}); + static constexpr auto L0BIAS_LAYOUT = tla::MakeLayout(tla::Int{}); + + // When enabling L1 resident mode, restore the pointer and coordinates that record the last state + // to the initial state. if two blockmmad instances need to be consecutively invoked at the kernel layer, + // RestoreStatus() must be inserted between them. + CATLASS_DEVICE + void RestoreStatus() + { + for (int i = 0; i < L1A_STAGES; ++i) { + lastAddrA[i] = nullptr; + lastCoordA[i] = Catlass::MatrixCoord{0U, 0U}; + } + for (int i = 0; i < L1B_STAGES; ++i) { + lastAddrB[i] = nullptr; + lastCoordB[i] = Catlass::MatrixCoord{0U, 0U}; + } + } + + /// Construct + CATLASS_DEVICE + BlockMmadTla(Catlass::Arch::Resource &resource, uint32_t l1BufAddrStart = 0) + { + if ASCEND_IS_AIC { + // use HF32 when USE_HF32_MODE is true + if constexpr (USE_HF32_MODE) { + AscendC::SetHF32Mode(true); + } else { + AscendC::SetHF32Mode(false); + } + if constexpr (ENABLE_UNIT_FLAG && tla::detail::isRowMajor::value) { + AscendC::SetMMLayoutTransform(true); + } + uint32_t l1AOffset = l1BufAddrStart; + uint32_t l1BOffset = l1BufAddrStart + L1A_TILE_SIZE * L1A_STAGES; + // Init buffers + for (uint32_t i = 0; i < L1A_STAGES; i++) { + // Assign L1/L0A/L0B space for each stages + l1ATensorList[i] = resource.l1Buf.template GetBufferByByte(l1AOffset + L1A_TILE_SIZE * i); + // Assign event ID for each stages + l1AEventList[i] = i; + // The event id that needs to be set before the loop + AscendC::SetFlag(l1AEventList[i]); + } + for (uint32_t i = 0; i < L1B_STAGES; i++) { + // Assign L1/L0A/L0B space for each stages + l1BTensorList[i] = resource.l1Buf.template GetBufferByByte(l1BOffset + L1B_TILE_SIZE * i); + // Assign event ID for each stages + l1BEventList[i] = i + L1A_STAGES; + // The event id that needs to be set before the loop + AscendC::SetFlag(l1BEventList[i]); + } + for (uint32_t i = 0; i < L0A_STAGES; i++) { + // Assign L1/L0A/L0B space for each stages + l0ATensorList[i] = resource.l0ABuf.template GetBufferByByte(L0A_TILE_SIZE * i); + // Assign event ID for each stages + l0AEventList[i] = i; + // The event id that needs to be set before the loop + AscendC::SetFlag(l0AEventList[i]); + } + for (uint32_t i = 0; i < L0B_STAGES; i++) { + // Assign L1/L0A/L0B space for each stages + l0BTensorList[i] = resource.l0BBuf.template GetBufferByByte(L0B_TILE_SIZE * i); + // Assign event ID for each stages + l0BEventList[i] = i + L0A_STAGES; + // The event id that needs to be set before the loop + AscendC::SetFlag(l0BEventList[i]); + } + if constexpr (!ENABLE_UNIT_FLAG) { + for (uint32_t i = 0; i < L0C_STAGES; i++) { + l0CTensorList[i] = resource.l0CBuf.template GetBufferByByte(L0C_TILE_SIZE * i); + l0CEventList[i] = i; + AscendC::SetFlag(l0CEventList[i]); + } + } else { + l0CTensorList[0] = resource.l0CBuf.template GetBufferByByte(0); + } + if constexpr (HAS_BIAS) { + uint32_t l1BiasOffset = l1BOffset + L1B_TILE_SIZE * L1B_STAGES; + l1BiasTensor = resource.l1Buf.template GetBufferByByte(l1BiasOffset); + l0BiasTensor = resource.btBuf.template GetBufferByByte(0); + AscendC::SetFlag(L1A_STAGES + L1B_STAGES); + AscendC::SetFlag(L0A_STAGES + L0B_STAGES); + } + + if constexpr (ENABLE_L1_RESIDENT) { + RestoreStatus(); + } + } + } + + /// Destructor + CATLASS_DEVICE + ~BlockMmadTla() + { + if ASCEND_IS_AIC { + if constexpr (USE_HF32_MODE) { + AscendC::SetHF32Mode(false); + } + if constexpr (ENABLE_UNIT_FLAG && tla::detail::isRowMajor::value) { + AscendC::SetMMLayoutTransform(false); + } + for (uint32_t i = 0; i < L1A_STAGES; i++) { + AscendC::WaitFlag(l1AEventList[i]); + } + for (uint32_t i = 0; i < L1B_STAGES; i++) { + AscendC::WaitFlag(l1BEventList[i]); + } + for (uint32_t i = 0; i < L0A_STAGES; i++) { + AscendC::WaitFlag(l0AEventList[i]); + } + for (uint32_t i = 0; i < L0B_STAGES; i++) { + AscendC::WaitFlag(l0BEventList[i]); + } + if constexpr (!ENABLE_UNIT_FLAG) { + for (uint32_t i = 0; i < L0C_STAGES; i++) { + AscendC::WaitFlag(l0CEventList[i]); + } + } + if constexpr (HAS_BIAS) { + AscendC::WaitFlag(L1A_STAGES + L1B_STAGES); + AscendC::WaitFlag(L0A_STAGES + L0B_STAGES); + } + } + } + + /// Perform a block-scoped matrix multiply-accumulate + template + CATLASS_DEVICE void operator()(TensorA &tensorA, TensorB &tensorB, TensorC &tensorC, + Catlass::GemmCoord const &actualShape, TensorBias const &tensorBias = {}) + { + // Check L1TileShape + if constexpr (HAS_BIAS) { + static constexpr uint32_t BIAS_BUF_SIZE = L0_TILE_N * sizeof(ElementAccumulator); + static constexpr uint32_t L1BIAS_SIZE = L1_TILE_N * sizeof(ElementBias); + static_assert(BIAS_BUF_SIZE <= ArchTag::BIAS_SIZE, + "BIAS_BUF_SIZE exceeding the BT space! Reduce L0_TILE_N"); + static_assert(L1A_TILE_SIZE * L1A_STAGES + L1B_TILE_SIZE * L1B_STAGES + L1BIAS_SIZE <= ArchTag::L1_SIZE, + "L1TileShape exceeding the L1 space!"); + } + + using CopyGmToL1A = typename TileCopy_::template CopyGmToL1A; + using CopyGmToL1B = typename TileCopy_::template CopyGmToL1B; + CopyGmToL1A copyGmToL1A; + CopyGmToL1B copyGmToL1B; +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 2201) + using CopyL0CToGm = typename TileCopy_::template CopyL0CToGm; + CopyL0CToGm copyL0CToDst; +#endif +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 3510) + using CopyL0CToDst = typename TileCopy_::template CopyL0CToDst; + CopyL0CToDst copyL0CToDst; +#endif + + uint32_t mBlockActual = actualShape.m(); + uint32_t kBlockActual = actualShape.k(); + uint32_t nBlockActual = actualShape.n(); + + uint32_t mL1Actual = mBlockActual; + if constexpr (std::is_same_v) { + // Avoid using the gemv mode in mmad + if (mL1Actual == 1) { + mL1Actual = 16; + } + } + uint32_t nL1Actual = nBlockActual; + + auto layoutInL0C = tla::MakeLayoutL0C(mL1Actual, nL1Actual); + auto tensorL0C = tla::MakeTensor(l0CTensorList[l0CListId], layoutInL0C, Catlass::Arch::PositionL0C{}); + auto tensorL0Bias = tla::MakeTensor(l0BiasTensor, L0BIAS_LAYOUT, Catlass::Arch::PositionBias{}); + + uint32_t kL1Actual = min(kBlockActual, L1_TILE_K); + // load first matrix A tile from GM to L1 + AscendC::WaitFlag(l1AEventList[l1AListId]); + auto tensorL1A = tla::MakeTensor(l1ATensorList[l1AListId], L1A_LAYOUT, Catlass::Arch::PositionL1{}); + auto tensorTileA = GetTileA(tensorA, 0, 0, mBlockActual, kL1Actual); + if constexpr (ENABLE_L1_RESIDENT) { + // If the currently loaded GM pointer and block coordinates are the same as the last loaded ones, + // skip this loading. + if (lastAddrA[l1AListId] != tensorTileA.data().GetPhyAddr() || + tla::get<0>(tensorTileA.coord()) != lastCoordA[l1AListId].row() || + tla::get<1>(tensorTileA.coord()) != lastCoordA[l1AListId].column()) { + copyGmToL1A(tensorL1A, tensorTileA); + lastCoordA[l1AListId] = + Catlass::MatrixCoord{tla::get<0>(tensorTileA.coord()), tla::get<1>(tensorTileA.coord())}; + lastAddrA[l1AListId] = const_cast<__gm__ typename AscendC::GlobalTensor::PrimType *>( + tensorTileA.data().GetPhyAddr()); + } + } else { + copyGmToL1A(tensorL1A, tensorTileA); + } + AscendC::SetFlag(l1AEventList[l1AListId]); + + // load first matrix B tile from GM to L1 + AscendC::WaitFlag(l1BEventList[l1BListId]); + auto tensorL1B = tla::MakeTensor(l1BTensorList[l1BListId], L1B_LAYOUT, Catlass::Arch::PositionL1{}); + auto tensorTileB = GetTile(tensorB, tla::MakeCoord(0, 0), tla::MakeShape(kL1Actual, nBlockActual)); + if constexpr (ENABLE_L1_RESIDENT) { + if (lastAddrB[l1BListId] != tensorTileB.data().GetPhyAddr() || + tla::get<0>(tensorTileB.coord()) != lastCoordB[l1BListId].row() || + tla::get<1>(tensorTileB.coord()) != lastCoordB[l1BListId].column()) { + copyGmToL1B(tensorL1B, tensorTileB); + lastCoordB[l1BListId] = + Catlass::MatrixCoord{tla::get<0>(tensorTileB.coord()), tla::get<1>(tensorTileB.coord())}; + lastAddrB[l1BListId] = const_cast<__gm__ typename AscendC::GlobalTensor::PrimType *>( + tensorTileB.data().GetPhyAddr()); + } + } else { + copyGmToL1B(tensorL1B, tensorTileB); + } + AscendC::SetFlag(l1BEventList[l1BListId]); + + if constexpr (HAS_BIAS && !std::is_same_v) { + using CopyGmToL1Bias = typename TileCopy::template CopyGmToL1Bias; + CopyGmToL1Bias copyGmToL1Bias; + AscendC::WaitFlag(L1A_STAGES + L1B_STAGES); + auto l1Bias = l1BiasTensor.template ReinterpretCast(); + auto tensorL1Bias = tla::MakeTensor(l1Bias, L1BIAS_LAYOUT, Catlass::Arch::PositionL1{}); + copyGmToL1Bias(tensorL1Bias, tensorBias); + AscendC::SetFlag(L1A_STAGES + L1B_STAGES); + } + + if constexpr (!ENABLE_UNIT_FLAG) { + AscendC::WaitFlag(l0CEventList[l0CListId]); + } + + uint32_t mL0Loop = CeilDiv(mL1Actual); + uint32_t nL0Loop = CeilDiv(nL1Actual); + + // main loop + uint32_t kL1Loop = CeilDiv(kBlockActual); + for (uint32_t kL1Idx = 0; kL1Idx < kL1Loop; kL1Idx++) { + uint32_t l1AListIdNext = (l1AListId + 1 < L1A_STAGES) ? (l1AListId + 1) : 0; + uint32_t l1BListIdNext = (l1BListId + 1 < L1B_STAGES) ? (l1BListId + 1) : 0; + uint32_t kL1ActualNext{0}; + // preload next tile from GM to L1 + if (kL1Idx < kL1Loop - 1) { + uint32_t kL1IdxNext = kL1Idx + 1; + kL1ActualNext = (kL1IdxNext < kL1Loop - 1) ? L1_TILE_K : (kBlockActual - kL1IdxNext * L1_TILE_K); + + // Get L1 tensor for next stage + auto l1ATensor = l1ATensorList[l1AListIdNext]; + auto l1BTensor = l1BTensorList[l1BListIdNext]; + auto tensorL1A = tla::MakeTensor(l1ATensor, L1A_LAYOUT, Catlass::Arch::PositionL1{}); + auto tensorL1B = tla::MakeTensor(l1BTensor, L1B_LAYOUT, Catlass::Arch::PositionL1{}); + // Get GM tile for next stage + auto tensorTileA = GetTileA(tensorA, 0, kL1IdxNext * L1_TILE_K, mBlockActual, kL1ActualNext); + auto tensorTileB = GetTile(tensorB, tla::MakeCoord(kL1IdxNext * L1_TILE_K, 0), + tla::MakeShape(kL1ActualNext, nBlockActual)); + + // load next matrix A tile from GM to L1 + AscendC::WaitFlag(l1AEventList[l1AListIdNext]); + if constexpr (ENABLE_L1_RESIDENT) { + if (lastAddrA[l1AListIdNext] != tensorTileA.data().GetPhyAddr() || + tla::get<0>(tensorTileA.coord()) != lastCoordA[l1AListIdNext].row() || + tla::get<1>(tensorTileA.coord()) != lastCoordA[l1AListIdNext].column()) { + copyGmToL1A(tensorL1A, tensorTileA); + lastCoordA[l1AListIdNext] = + Catlass::MatrixCoord{tla::get<0>(tensorTileA.coord()), tla::get<1>(tensorTileA.coord())}; + lastAddrA[l1AListIdNext] = + const_cast<__gm__ typename AscendC::GlobalTensor::PrimType *>( + tensorTileA.data().GetPhyAddr()); + } + } else { + copyGmToL1A(tensorL1A, tensorTileA); + } + AscendC::SetFlag(l1AEventList[l1AListIdNext]); + + // load next matrix B tile from GM to L1 + AscendC::WaitFlag(l1BEventList[l1BListIdNext]); + if constexpr (ENABLE_L1_RESIDENT) { + if (lastAddrB[l1BListIdNext] != tensorTileB.data().GetPhyAddr() || + tla::get<0>(tensorTileB.coord()) != lastCoordB[l1BListIdNext].row() || + tla::get<1>(tensorTileB.coord()) != lastCoordB[l1BListIdNext].column()) { + copyGmToL1B(tensorL1B, tensorTileB); + lastCoordB[l1BListIdNext] = + Catlass::MatrixCoord{tla::get<0>(tensorTileB.coord()), tla::get<1>(tensorTileB.coord())}; + lastAddrB[l1BListIdNext] = + const_cast<__gm__ typename AscendC::GlobalTensor::PrimType *>( + tensorTileB.data().GetPhyAddr()); + } + } else { + copyGmToL1B(tensorL1B, tensorTileB); + } + AscendC::SetFlag(l1BEventList[l1BListIdNext]); + } + + // Get L1 tensor for current stage + auto l1ATensor = l1ATensorList[l1AListId]; + auto l1BTensor = l1BTensorList[l1BListId]; + tensorL1A = tla::MakeTensor(l1ATensor, L1A_LAYOUT, Catlass::Arch::PositionL1{}); + tensorL1B = tla::MakeTensor(l1BTensor, L1B_LAYOUT, Catlass::Arch::PositionL1{}); + // Get the loop nums on L0 + uint32_t kL0Loop = CeilDiv(kL1Actual); + + for (int mL0Idx = 0; mL0Idx < mL0Loop; mL0Idx++) { + uint32_t mL0Actual = (mL0Idx < mL0Loop - 1) ? L0_TILE_M : (mL1Actual - mL0Idx * L0_TILE_M); + + for (int kL0Idx = 0; kL0Idx < kL0Loop; kL0Idx++) { + uint32_t kL0Actual = (kL0Idx < kL0Loop - 1) ? L0_TILE_K : (kL1Actual - kL0Idx * L0_TILE_K); + + // Locate the current tile on L0A + auto l0ATile = l0ATensorList[l0AListId]; + auto layoutAInL0 = tla::MakeLayout(mL0Actual, kL0Actual); + auto tensorL0A = tla::MakeTensor(l0ATile, layoutAInL0, Catlass::Arch::PositionL0A{}); + // Locate the current tile of matrix A on L1 + auto tensorTileL1A = + GetTileA(tensorL1A, mL0Idx * L0_TILE_M, kL0Idx * L0_TILE_K, mL0Actual, kL0Actual); + + AscendC::WaitFlag(l0AEventList[l0AListId]); + if ((mL0Idx == 0) && (kL0Idx == 0)) { + AscendC::WaitFlag(l1AEventList[l1AListId]); + } + + // Load current tile from L1 to L0A + copyL1ToL0A(tensorL0A, tensorTileL1A); + + if ((mL0Idx == mL0Loop - 1) && (kL0Idx == kL0Loop - 1)) { + AscendC::SetFlag(l1AEventList[l1AListId]); + } + + bool initC = ((kL1Idx == 0) && (kL0Idx == 0)); + for (int nL0Idx = 0; nL0Idx < nL0Loop; nL0Idx++) { + uint32_t nL0Actual = (nL0Idx < nL0Loop - 1) ? L0_TILE_N : (nL1Actual - nL0Idx * L0_TILE_N); + + // Locate the current tile on L0B + auto l0BTile = l0BTensorList[l0BListId]; + auto layoutBInL0 = tla::MakeLayout(kL0Actual, nL0Actual); + auto tensorL0B = tla::MakeTensor(l0BTile, layoutBInL0, Catlass::Arch::PositionL0B{}); + // Locate the current tile of matrix B on L1 + auto tensorTileL1B = GetTile(tensorL1B, tla::MakeCoord(kL0Idx * L0_TILE_K, nL0Idx * L0_TILE_N), + tla::MakeShape(kL0Actual, nL0Actual)); + + // Wait for mmad finished + AscendC::WaitFlag(l0BEventList[l0BListId]); + // If the current tile is the first one on the k&n axis, wait for loading matrix B from GM to L1 + if ((mL0Idx == 0) && (kL0Idx == 0) && (nL0Idx == 0)) { + AscendC::WaitFlag(l1BEventList[l1BListId]); + } + + // Load current tile from L1 to L0B + copyL1ToL0B(tensorL0B, tensorTileL1B); + + // If the current tile is the last one on the k&n axis, notify to load matrix B from GM to L1 + if ((mL0Idx == mL0Loop - 1) && (kL0Idx == kL0Loop - 1) && (nL0Idx == nL0Loop - 1)) { + AscendC::SetFlag(l1BEventList[l1BListId]); + } + + if constexpr (HAS_BIAS && !std::is_same_v) { + if (initC) { + if (nL0Idx == 0) { + AscendC::WaitFlag(L1A_STAGES + L1B_STAGES); + } + AscendC::WaitFlag(L0A_STAGES + L0B_STAGES); + auto l1Bias = l1BiasTensor.template ReinterpretCast(); + auto tensorL1Bias = tla::MakeTensor(l1Bias, L1BIAS_LAYOUT, Catlass::Arch::PositionL1{}); + auto tensorTileL1Bias = GetTile(tensorL1Bias, tla::MakeCoord(nL0Idx * L0_TILE_N), + tla::MakeShape(nL0Actual)); + // Load bias to l0 biasTable + copyL1ToBT(tensorL0Bias, tensorTileL1Bias); + if (nL0Idx == nL0Loop - 1) { + AscendC::SetFlag(L1A_STAGES + L1B_STAGES); + } + } + } + + // Notify to do mmad + AscendC::SetFlag(l0CEventList[l0CListId]); + + // Locate the current tile on L0C + auto tensorTileL0C = GetTile(tensorL0C, tla::MakeCoord(mL0Idx * L0_TILE_M, nL0Idx * L0_TILE_N), + tla::MakeShape(mL0Actual, nL0Actual)); + + // Compute the matrix multiplication on L0A and L0B and write the result to the accumulator + // Wait for loading L0B + AscendC::WaitFlag(l0CEventList[l0CListId]); + + // If the unit flag is enabled, the unit flag is set according to the calculation progress + uint8_t unitFlag = 0b00; + if constexpr (ENABLE_UNIT_FLAG) { + if ((kL1Idx == kL1Loop - 1) && (mL0Idx == mL0Loop - 1) && (kL0Idx == kL0Loop - 1) && + (nL0Idx == nL0Loop - 1)) { + unitFlag = 0b11; + } else { + unitFlag = 0b10; + } + } + + if constexpr (HAS_BIAS && !std::is_same_v) { + if (initC) { + tileMmad(tensorTileL0C, tensorL0A, tensorL0B, tensorL0Bias, mL0Actual, nL0Actual, + kL0Actual, initC, unitFlag); + AscendC::SetFlag(L0A_STAGES + L0B_STAGES); + } else { + tileMmad(tensorTileL0C, tensorL0A, tensorL0B, mL0Actual, nL0Actual, kL0Actual, initC, + unitFlag); + } + } else { + tileMmad(tensorTileL0C, tensorL0A, tensorL0B, mL0Actual, nL0Actual, kL0Actual, initC, + unitFlag); + } + + // Notify to move the next L0B tile + AscendC::SetFlag(l0BEventList[l0BListId]); + l0BListId = (l0BListId + 1 < L0B_STAGES) ? (l0BListId + 1) : 0; + } + AscendC::SetFlag(l0AEventList[l0AListId]); + l0AListId = (l0AListId + 1 < L0A_STAGES) ? (l0AListId + 1) : 0; + } + } + l1AListId = l1AListIdNext; + l1BListId = l1BListIdNext; + kL1Actual = kL1ActualNext; + } + + // copy block out + if constexpr (!ENABLE_UNIT_FLAG) { + AscendC::SetFlag(l0CEventList[l0CListId]); + AscendC::WaitFlag(l0CEventList[l0CListId]); + copyL0CToDst(tensorC, tensorL0C); + AscendC::SetFlag(l0CEventList[l0CListId]); + l0CListId = (l0CListId + 1 < L0C_STAGES) ? (l0CListId + 1) : 0; + } else { + copyL0CToDst(tensorC, tensorL0C, 0b11); + } + } + + /// Perform a block-scoped matrix multiply-accumulate + template + CATLASS_DEVICE void operator()(TensorA &tensorA, TensorB &tensorB, TensorC tensorCList[MAX_CUBE_VEC_SYNC_NUM], + Catlass::GemmCoord const &actualShape, uint32_t ubRowNum, uint8_t beginSubBlockIdx, + uint64_t l0C2UBAIVAICFlag, uint64_t l0C2UBAICAIVFlag, uint32_t &ubListId, + uint8_t sendVecNum, uint64_t ubNum = 2, TensorBias const &tensorBias = {}) + { + // Check L1TileShape + if constexpr (HAS_BIAS) { + static constexpr uint32_t BIAS_BUF_SIZE = L0_TILE_N * sizeof(ElementAccumulator); + static constexpr uint32_t L1BIAS_SIZE = L1_TILE_N * sizeof(ElementBias); + static_assert(BIAS_BUF_SIZE <= ArchTag::BIAS_SIZE, + "BIAS_BUF_SIZE exceeding the BT space! Reduce L0_TILE_N"); + static_assert(L1A_TILE_SIZE * L1A_STAGES + L1B_TILE_SIZE * L1B_STAGES + L1BIAS_SIZE <= ArchTag::L1_SIZE, + "L1TileShape exceeding the L1 space!"); + } + + using CopyGmToL1A = typename TileCopy_::template CopyGmToL1A; + using CopyGmToL1B = typename TileCopy_::template CopyGmToL1B; + using CopyL0CToDst = typename TileCopy_::template CopyL0CToDst; + CopyGmToL1A copyGmToL1A; + CopyGmToL1B copyGmToL1B; + CopyL0CToDst copyL0CToDst; + + uint32_t mBlockActual = actualShape.m(); + uint32_t kBlockActual = actualShape.k(); + uint32_t nBlockActual = actualShape.n(); + + uint32_t mL1Actual = mBlockActual; + if constexpr (std::is_same_v) { + // Avoid using the gemv mode in mmad + if (mL1Actual == 1) { + mL1Actual = 16; + } + } + uint32_t nL1Actual = nBlockActual; + + auto layoutInL0C = tla::MakeLayoutL0C(mL1Actual, nL1Actual); + auto tensorL0C = tla::MakeTensor(l0CTensorList[l0CListId], layoutInL0C, Catlass::Arch::PositionL0C{}); + auto tensorL0Bias = tla::MakeTensor(l0BiasTensor, L0BIAS_LAYOUT, Catlass::Arch::PositionBias{}); + + uint32_t kL1Actual = min(kBlockActual, L1_TILE_K); + // load first matrix A tile from GM to L1 + AscendC::WaitFlag(l1AEventList[l1AListId]); + auto tensorL1A = tla::MakeTensor(l1ATensorList[l1AListId], L1A_LAYOUT, Catlass::Arch::PositionL1{}); + auto tensorTileA = GetTileA(tensorA, 0, 0, mBlockActual, kL1Actual); + if constexpr (ENABLE_L1_RESIDENT) { + // If the currently loaded GM pointer and block coordinates are the same as the last loaded ones, + // skip this loadding. + if (lastAddrA[l1AListId] != tensorTileA.data().GetPhyAddr() || + tla::get<0>(tensorTileA.coord()) != lastCoordA[l1AListId].row() || + tla::get<1>(tensorTileA.coord()) != lastCoordA[l1AListId].column()) { + copyGmToL1A(tensorL1A, tensorTileA); + lastCoordA[l1AListId] = + Catlass::MatrixCoord{tla::get<0>(tensorTileA.coord()), tla::get<1>(tensorTileA.coord())}; + lastAddrA[l1AListId] = const_cast<__gm__ typename AscendC::GlobalTensor::PrimType *>( + tensorTileA.data().GetPhyAddr()); + } + } else { + copyGmToL1A(tensorL1A, tensorTileA); + } + AscendC::SetFlag(l1AEventList[l1AListId]); + + // load first matrix B tile from GM to L1 + AscendC::WaitFlag(l1BEventList[l1BListId]); + auto tensorL1B = tla::MakeTensor(l1BTensorList[l1BListId], L1B_LAYOUT, Catlass::Arch::PositionL1{}); + auto tensorTileB = GetTile(tensorB, tla::MakeCoord(0, 0), tla::MakeShape(kL1Actual, nBlockActual)); + if constexpr (ENABLE_L1_RESIDENT) { + if (lastAddrB[l1BListId] != tensorTileB.data().GetPhyAddr() || + tla::get<0>(tensorTileB.coord()) != lastCoordB[l1BListId].row() || + tla::get<1>(tensorTileB.coord()) != lastCoordB[l1BListId].column()) { + copyGmToL1B(tensorL1B, tensorTileB); + lastCoordB[l1BListId] = + Catlass::MatrixCoord{tla::get<0>(tensorTileB.coord()), tla::get<1>(tensorTileB.coord())}; + lastAddrB[l1BListId] = const_cast<__gm__ typename AscendC::GlobalTensor::PrimType *>( + tensorTileB.data().GetPhyAddr()); + } + } else { + copyGmToL1B(tensorL1B, tensorTileB); + } + AscendC::SetFlag(l1BEventList[l1BListId]); + + if constexpr (HAS_BIAS && !std::is_same_v) { + using CopyGmToL1Bias = typename TileCopy::template CopyGmToL1Bias; + CopyGmToL1Bias copyGmToL1Bias; + AscendC::WaitFlag(L1A_STAGES + L1B_STAGES); + auto l1Bias = l1BiasTensor.template ReinterpretCast(); + auto tensorL1Bias = tla::MakeTensor(l1Bias, L1BIAS_LAYOUT, Catlass::Arch::PositionL1{}); + copyGmToL1Bias(tensorL1Bias, tensorBias); + AscendC::SetFlag(L1A_STAGES + L1B_STAGES); + } + + if constexpr (!ENABLE_UNIT_FLAG) { + AscendC::WaitFlag(l0CEventList[l0CListId]); + } + + uint32_t mL0Loop = CeilDiv(mL1Actual); + uint32_t nL0Loop = CeilDiv(nL1Actual); + + // main loop + uint32_t kL1Loop = CeilDiv(kBlockActual); + for (uint32_t kL1Idx = 0; kL1Idx < kL1Loop; kL1Idx++) { + uint32_t l1AListIdNext = (l1AListId + 1 < L1A_STAGES) ? (l1AListId + 1) : 0; + uint32_t l1BListIdNext = (l1BListId + 1 < L1B_STAGES) ? (l1BListId + 1) : 0; + uint32_t kL1ActualNext{0}; + // preload next tile from GM to L1 + if (kL1Idx < kL1Loop - 1) { + uint32_t kL1IdxNext = kL1Idx + 1; + kL1ActualNext = (kL1IdxNext < kL1Loop - 1) ? L1_TILE_K : (kBlockActual - kL1IdxNext * L1_TILE_K); + + // Get L1 tensor for next stage + auto l1ATensor = l1ATensorList[l1AListIdNext]; + auto l1BTensor = l1BTensorList[l1BListIdNext]; + auto tensorL1A = tla::MakeTensor(l1ATensor, L1A_LAYOUT, Catlass::Arch::PositionL1{}); + auto tensorL1B = tla::MakeTensor(l1BTensor, L1B_LAYOUT, Catlass::Arch::PositionL1{}); + // Get GM tile for next stage + auto tensorTileA = GetTileA(tensorA, 0, kL1IdxNext * L1_TILE_K, mBlockActual, kL1ActualNext); + auto tensorTileB = GetTile(tensorB, tla::MakeCoord(kL1IdxNext * L1_TILE_K, 0), + tla::MakeShape(kL1ActualNext, nBlockActual)); + + // load next matrix A tile from GM to L1 + AscendC::WaitFlag(l1AEventList[l1AListIdNext]); + if constexpr (ENABLE_L1_RESIDENT) { + if (lastAddrA[l1AListIdNext] != tensorTileA.data().GetPhyAddr() || + tla::get<0>(tensorTileA.coord()) != lastCoordA[l1AListIdNext].row() || + tla::get<1>(tensorTileA.coord()) != lastCoordA[l1AListIdNext].column()) { + copyGmToL1A(tensorL1A, tensorTileA); + lastCoordA[l1AListIdNext] = + Catlass::MatrixCoord{tla::get<0>(tensorTileA.coord()), tla::get<1>(tensorTileA.coord())}; + lastAddrA[l1AListIdNext] = + const_cast<__gm__ typename AscendC::GlobalTensor::PrimType *>( + tensorTileA.data().GetPhyAddr()); + } + } else { + copyGmToL1A(tensorL1A, tensorTileA); + } + AscendC::SetFlag(l1AEventList[l1AListIdNext]); + + // load next matrix B tile from GM to L1 + AscendC::WaitFlag(l1BEventList[l1BListIdNext]); + if constexpr (ENABLE_L1_RESIDENT) { + if (lastAddrB[l1BListIdNext] != tensorTileB.data().GetPhyAddr() || + tla::get<0>(tensorTileB.coord()) != lastCoordB[l1BListIdNext].row() || + tla::get<1>(tensorTileB.coord()) != lastCoordB[l1BListIdNext].column()) { + copyGmToL1B(tensorL1B, tensorTileB); + lastCoordB[l1BListIdNext] = + Catlass::MatrixCoord{tla::get<0>(tensorTileB.coord()), tla::get<1>(tensorTileB.coord())}; + lastAddrB[l1BListIdNext] = + const_cast<__gm__ typename AscendC::GlobalTensor::PrimType *>( + tensorTileB.data().GetPhyAddr()); + } + } else { + copyGmToL1B(tensorL1B, tensorTileB); + } + AscendC::SetFlag(l1BEventList[l1BListIdNext]); + } + + // Get L1 tensor for current stage + auto l1ATensor = l1ATensorList[l1AListId]; + auto l1BTensor = l1BTensorList[l1BListId]; + tensorL1A = tla::MakeTensor(l1ATensor, L1A_LAYOUT, Catlass::Arch::PositionL1{}); + tensorL1B = tla::MakeTensor(l1BTensor, L1B_LAYOUT, Catlass::Arch::PositionL1{}); + // Get the loop nums on L0 + uint32_t kL0Loop = CeilDiv(kL1Actual); + + for (int mL0Idx = 0; mL0Idx < mL0Loop; mL0Idx++) { + uint32_t mL0Actual = (mL0Idx < mL0Loop - 1) ? L0_TILE_M : (mL1Actual - mL0Idx * L0_TILE_M); + + for (int kL0Idx = 0; kL0Idx < kL0Loop; kL0Idx++) { + uint32_t kL0Actual = (kL0Idx < kL0Loop - 1) ? L0_TILE_K : (kL1Actual - kL0Idx * L0_TILE_K); + + // Locate the current tile on L0A + auto l0ATile = l0ATensorList[l0AListId]; + auto layoutAInL0 = tla::MakeLayout(mL0Actual, kL0Actual); + auto tensorL0A = tla::MakeTensor(l0ATile, layoutAInL0, Catlass::Arch::PositionL0A{}); + // Locate the current tile of matrix A on L1 + auto tensorTileL1A = + GetTileA(tensorL1A, mL0Idx * L0_TILE_M, kL0Idx * L0_TILE_K, mL0Actual, kL0Actual); + + AscendC::WaitFlag(l0AEventList[l0AListId]); + if ((mL0Idx == 0) && (kL0Idx == 0)) { + AscendC::WaitFlag(l1AEventList[l1AListId]); + } + + // Load current tile from L1 to L0A + copyL1ToL0A(tensorL0A, tensorTileL1A); + + if ((mL0Idx == mL0Loop - 1) && (kL0Idx == kL0Loop - 1)) { + AscendC::SetFlag(l1AEventList[l1AListId]); + } + + bool initC = ((kL1Idx == 0) && (kL0Idx == 0)); + for (int nL0Idx = 0; nL0Idx < nL0Loop; nL0Idx++) { + uint32_t nL0Actual = (nL0Idx < nL0Loop - 1) ? L0_TILE_N : (nL1Actual - nL0Idx * L0_TILE_N); + + // Locate the current tile on L0B + auto l0BTile = l0BTensorList[l0BListId]; + auto layoutBInL0 = tla::MakeLayout(kL0Actual, nL0Actual); + auto tensorL0B = tla::MakeTensor(l0BTile, layoutBInL0, Catlass::Arch::PositionL0B{}); + // Locate the current tile of matrix B on L1 + auto tensorTileL1B = GetTile(tensorL1B, tla::MakeCoord(kL0Idx * L0_TILE_K, nL0Idx * L0_TILE_N), + tla::MakeShape(kL0Actual, nL0Actual)); + + // Wait for mmad finished + AscendC::WaitFlag(l0BEventList[l0BListId]); + // If the current tile is the first one on the k&n axis, wait for loading matrix B from GM to L1 + if ((mL0Idx == 0) && (kL0Idx == 0) && (nL0Idx == 0)) { + AscendC::WaitFlag(l1BEventList[l1BListId]); + } + + // Load current tile from L1 to L0B + copyL1ToL0B(tensorL0B, tensorTileL1B); + + // If the current tile is the last one on the k&n axis, notify to load matrix B from GM to L1 + if ((mL0Idx == mL0Loop - 1) && (kL0Idx == kL0Loop - 1) && (nL0Idx == nL0Loop - 1)) { + AscendC::SetFlag(l1BEventList[l1BListId]); + } + + if constexpr (HAS_BIAS && !std::is_same_v) { + if (initC) { + if (nL0Idx == 0) { + AscendC::WaitFlag(L1A_STAGES + L1B_STAGES); + } + AscendC::WaitFlag(L0A_STAGES + L0B_STAGES); + auto l1Bias = l1BiasTensor.template ReinterpretCast(); + auto tensorL1Bias = tla::MakeTensor(l1Bias, L1BIAS_LAYOUT, Catlass::Arch::PositionL1{}); + auto tensorTileL1Bias = GetTile(tensorL1Bias, tla::MakeCoord(nL0Idx * L0_TILE_N), + tla::MakeShape(nL0Actual)); + // Load bias to l0 biastable + copyL1ToBT(tensorL0Bias, tensorTileL1Bias); + if (nL0Idx == nL0Loop - 1) { + AscendC::SetFlag(L1A_STAGES + L1B_STAGES); + } + } + } + + // Notify to do mmad + AscendC::SetFlag(l0CEventList[l0CListId]); + + // Locate the current tile on L0C + auto tensorTileL0C = GetTile(tensorL0C, tla::MakeCoord(mL0Idx * L0_TILE_M, nL0Idx * L0_TILE_N), + tla::MakeShape(mL0Actual, nL0Actual)); + + // Compute the matrix multiplication on L0A and L0B and write the result to the accumulator + // Wait for loading L0B + AscendC::WaitFlag(l0CEventList[l0CListId]); + + // If the unit flag is enabled, the unit flag is set according to the calculation progress + uint8_t unitFlag = 0b00; + if constexpr (ENABLE_UNIT_FLAG) { + if ((kL1Idx == kL1Loop - 1) && (mL0Idx == mL0Loop - 1) && (kL0Idx == kL0Loop - 1) && + (nL0Idx == nL0Loop - 1)) { + unitFlag = 0b11; + } else { + unitFlag = 0b10; + } + } + + if constexpr (HAS_BIAS && !std::is_same_v) { + if (initC) { + tileMmad(tensorTileL0C, tensorL0A, tensorL0B, tensorL0Bias, mL0Actual, nL0Actual, + kL0Actual, initC, unitFlag); + AscendC::SetFlag(L0A_STAGES + L0B_STAGES); + } else { + tileMmad(tensorTileL0C, tensorL0A, tensorL0B, mL0Actual, nL0Actual, kL0Actual, initC, + unitFlag); + } + } else { + tileMmad(tensorTileL0C, tensorL0A, tensorL0B, mL0Actual, nL0Actual, kL0Actual, initC, + unitFlag); + } + + // Notify to move the next L0B tile + AscendC::SetFlag(l0BEventList[l0BListId]); + l0BListId = (l0BListId + 1 < L0B_STAGES) ? (l0BListId + 1) : 0; + } + AscendC::SetFlag(l0AEventList[l0AListId]); + l0AListId = (l0AListId + 1 < L0A_STAGES) ? (l0AListId + 1) : 0; + } + } + l1AListId = l1AListIdNext; + l1BListId = l1BListIdNext; + kL1Actual = kL1ActualNext; + } + + // copy block out + if constexpr (!ENABLE_UNIT_FLAG) { + AscendC::SetFlag(l0CEventList[l0CListId]); + AscendC::WaitFlag(l0CEventList[l0CListId]); + uint32_t leftRowNum = mL1Actual; + uint32_t twoVecRowNum = ubRowNum * sendVecNum; + uint32_t curRowNum = min(twoVecRowNum, leftRowNum); + uint32_t rowIdx = 0; + uint32_t tileNum = CeilDiv(mL1Actual, twoVecRowNum); + for (uint32_t i = 0; i < tileNum; i++) { + auto tensorTileL0C = + GetTile(tensorL0C, tla::MakeCoord(rowIdx, 0), tla::MakeShape(curRowNum, nL1Actual)); + if (sendVecNum == 2) { + AscendC::CrossCoreWaitFlag<0x4, PIPE_FIX>(l0C2UBAIVAICFlag + ubListId); + AscendC::CrossCoreWaitFlag<0x4, PIPE_FIX>(l0C2UBAIVAICFlag + FLAG_ID_MAX + ubListId); + copyL0CToDst(tensorCList[ubListId], tensorTileL0C, ubRowNum, beginSubBlockIdx, sendVecNum); + AscendC::CrossCoreSetFlag<0x4, PIPE_FIX>(l0C2UBAICAIVFlag + ubListId); + AscendC::CrossCoreSetFlag<0x4, PIPE_FIX>(l0C2UBAICAIVFlag + FLAG_ID_MAX + ubListId); + rowIdx += curRowNum; + leftRowNum -= curRowNum; + curRowNum = min(twoVecRowNum, leftRowNum); + ubListId = (ubListId + 1 < ubNum) ? (ubListId + 1) : 0; + } else if (sendVecNum == 1) { + if (beginSubBlockIdx == 0) { + AscendC::CrossCoreWaitFlag<0x4, PIPE_FIX>(l0C2UBAIVAICFlag + ubListId); + } else if (beginSubBlockIdx == 1) { + AscendC::CrossCoreWaitFlag<0x4, PIPE_FIX>(l0C2UBAIVAICFlag + FLAG_ID_MAX + ubListId); + } + copyL0CToDst(tensorCList[ubListId], tensorTileL0C, ubRowNum, beginSubBlockIdx, sendVecNum); + if (beginSubBlockIdx == 0) { + AscendC::CrossCoreSetFlag<0x4, PIPE_FIX>(l0C2UBAICAIVFlag + ubListId); + } else if (beginSubBlockIdx == 1) { + AscendC::CrossCoreSetFlag<0x4, PIPE_FIX>(l0C2UBAICAIVFlag + FLAG_ID_MAX + ubListId); + } + rowIdx += curRowNum; + leftRowNum -= curRowNum; + curRowNum = min(twoVecRowNum, leftRowNum); + ubListId = (ubListId + 1 < ubNum) ? (ubListId + 1) : 0; + } + } + AscendC::SetFlag(l0CEventList[l0CListId]); + l0CListId = (l0CListId + 1 < L0C_STAGES) ? (l0CListId + 1) : 0; + } else { + uint32_t leftRowNum = mL1Actual; + uint32_t twoVecRowNum = ubRowNum * sendVecNum; + uint32_t curRowNum = min(twoVecRowNum, leftRowNum); + uint32_t rowIdx = 0; + uint32_t tileNum = CeilDiv(mL1Actual, twoVecRowNum); + for (uint32_t i = 0; i < tileNum; i++) { + auto tensorTileL0C = + GetTile(tensorL0C, tla::MakeCoord(rowIdx, 0), tla::MakeShape(curRowNum, nL1Actual)); + if (sendVecNum == 2) { + AscendC::CrossCoreWaitFlag<0x4, PIPE_FIX>(l0C2UBAIVAICFlag + ubListId); + AscendC::CrossCoreWaitFlag<0x4, PIPE_FIX>(l0C2UBAIVAICFlag + FLAG_ID_MAX + ubListId); + copyL0CToDst(tensorCList[ubListId], tensorTileL0C, ubRowNum, beginSubBlockIdx, sendVecNum); + AscendC::CrossCoreSetFlag<0x4, PIPE_FIX>(l0C2UBAICAIVFlag + ubListId); + AscendC::CrossCoreSetFlag<0x4, PIPE_FIX>(l0C2UBAICAIVFlag + FLAG_ID_MAX + ubListId); + rowIdx += curRowNum; + leftRowNum -= curRowNum; + curRowNum = min(twoVecRowNum, leftRowNum); + ubListId = (ubListId + 1 < ubNum) ? (ubListId + 1) : 0; + } else if (sendVecNum == 1) { + if (beginSubBlockIdx == 0) { + AscendC::CrossCoreWaitFlag<0x4, PIPE_FIX>(l0C2UBAIVAICFlag + ubListId); + } else if (beginSubBlockIdx == 1) { + AscendC::CrossCoreWaitFlag<0x4, PIPE_FIX>(l0C2UBAIVAICFlag + FLAG_ID_MAX + ubListId); + } + copyL0CToDst(tensorCList[ubListId], tensorTileL0C, ubRowNum, beginSubBlockIdx, sendVecNum); + if (beginSubBlockIdx == 0) { + AscendC::CrossCoreSetFlag<0x4, PIPE_FIX>(l0C2UBAICAIVFlag + ubListId); + } else if (beginSubBlockIdx == 1) { + AscendC::CrossCoreSetFlag<0x4, PIPE_FIX>(l0C2UBAICAIVFlag + FLAG_ID_MAX + ubListId); + } + rowIdx += curRowNum; + leftRowNum -= curRowNum; + curRowNum = min(twoVecRowNum, leftRowNum); + ubListId = (ubListId + 1 < ubNum) ? (ubListId + 1) : 0; + } + } + } + } + +protected: + template + CATLASS_DEVICE auto GetTileA(TensorA &tensorA, uint32_t mIndex, uint32_t kIndex, uint32_t mSize, uint32_t kSize) + { + if constexpr (tla::detail::isVector::value) { + return GetTile(tensorA, tla::MakeCoord(kIndex), tla::MakeShape(kSize)); + } else { + return GetTile(tensorA, tla::MakeCoord(mIndex, kIndex), tla::MakeShape(mSize, kSize)); + } + } + + // Multi-stage tensors list + AscendC::LocalTensor l1ATensorList[L1A_STAGES]; + AscendC::LocalTensor l1BTensorList[L1B_STAGES]; + AscendC::LocalTensor l0ATensorList[L0A_STAGES]; + AscendC::LocalTensor l0BTensorList[L0B_STAGES]; + AscendC::LocalTensor l0CTensorList[L0C_STAGES]; + AscendC::LocalTensor l1BiasTensor; + AscendC::LocalTensor l0BiasTensor; + + // Multi-stage event id list + int32_t l1AEventList[L1A_STAGES]; + int32_t l1BEventList[L1B_STAGES]; + int32_t l0AEventList[L0A_STAGES]; + int32_t l0BEventList[L0B_STAGES]; + int32_t l0CEventList[L0C_STAGES]; + + __gm__ typename AscendC::GlobalTensor::PrimType *lastAddrA[L1A_STAGES]; + __gm__ typename AscendC::GlobalTensor::PrimType *lastAddrB[L1B_STAGES]; + Catlass::MatrixCoord lastCoordA[L1A_STAGES]; + Catlass::MatrixCoord lastCoordB[L1B_STAGES]; + + // The id of current stage + uint32_t l1AListId{0}; + uint32_t l1BListId{0}; + uint32_t l0AListId{0}; + uint32_t l0BListId{0}; + uint32_t l0CListId{0}; + + TileMmad tileMmad; + CopyL1ToL0A copyL1ToL0A; + CopyL1ToL0B copyL1ToL0B; + CopyL1ToBT copyL1ToBT; +}; + +} // namespace Common + +#endif // CATLASS_GEMM_BLOCK_BLOCK_MMAD_PINGPONG_TLA_KERNEL_UTILS_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/kernel_utils/block/block_mmad_pingpong_tla_multi.hpp b/csrc/chunk_kda_fwd/op_kernel/kernel_utils/block/block_mmad_pingpong_tla_multi.hpp new file mode 100644 index 0000000000..08dd67ff11 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/kernel_utils/block/block_mmad_pingpong_tla_multi.hpp @@ -0,0 +1,677 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_GEMM_BLOCK_BLOCK_MMAD_PINGPONG_TLA_MULTI_HPP +#define CATLASS_GEMM_BLOCK_BLOCK_MMAD_PINGPONG_TLA_MULTI_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/kda_coord.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" +#include "catlass/gemm/kda_gemm_helper.hpp" +#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" +#include "catlass/gemm/tile/kda_tile_mmad.hpp" +#include "tla/kda_tla_layout.hpp" +#include "tla/kda_tensor.hpp" + +namespace Catlass::Gemm { + +template +struct MmadPingpongTlaMulti : public MmadBase { + static constexpr uint32_t L1A_STAGES = L1A_STAGES_; + static constexpr uint32_t L1B_STAGES = L1B_STAGES_; + static constexpr uint32_t L0A_STAGES = L0A_STAGES_; + static constexpr uint32_t L0B_STAGES = L0B_STAGES_; + static constexpr uint32_t L0C_STAGES = L0C_STAGES_; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; + static constexpr bool USE_HF32_MODE = USE_HF32_MODE_; + static constexpr bool ENABLE_L1_RESIDENT = ENABLE_L1_RESIDENT_; +}; + +} // namespace Catlass::Gemm + +namespace Catlass::Gemm::Block { + +template +struct BlockMmadTla, + L1TileShape_, L0TileShape_, ElementA_, ElementB_, ElementC_, ElementBias_, TileCopy_, TileMmad_> { +public: + // Type Aliases + using DispatchPolicy = + MmadPingpongTlaMulti; + using ArchTag = typename DispatchPolicy::ArchTag; + using TileCopy = TileCopy_; + using L1TileShape = L1TileShape_; + using L0TileShape = L0TileShape_; + using ElementA = ElementA_; + using LayoutA = typename TileCopy::LayoutA; + using ElementB = ElementB_; + using LayoutB = typename TileCopy::LayoutB; + using ElementC = ElementC_; + using LayoutC = typename TileCopy::LayoutC; + using ElementBias = ElementBias_; + + using TileMmad = TileMmad_; + + using CopyL1ToL0A = typename TileCopy::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopy::CopyL1ToL0B; + using CopyL1ToBT = typename TileCopy::CopyL1ToBT; + + using ElementAccumulator = typename TileCopy::ElementAccumulator; + + static constexpr bool HAS_BIAS = TileCopy::HAS_BIAS; + + using LayoutTagL1A = typename TileCopy::LayoutTagL1A; + using LayoutTagL1B = typename TileCopy::LayoutTagL1B; + using LayoutTagL0A = typename TileCopy::LayoutTagL0A; + using LayoutTagL0B = typename TileCopy::LayoutTagL0B; + + using L1AAlignHelper = typename TileCopy_::L1AAlignHelper; + using L1BAlignHelper = typename TileCopy_::L1BAlignHelper; + + static_assert(tla::is_tuple::value && tla::is_static::value, + "L1TileShape must be tla::tuple and static!"); + static_assert(tla::is_tuple::value && tla::is_static::value, + "L0TileShape must be tla::tuple and static!"); + + static constexpr bool ENABLE_UNIT_FLAG = DispatchPolicy::ENABLE_UNIT_FLAG; + static constexpr bool USE_HF32_MODE = DispatchPolicy::USE_HF32_MODE; + static constexpr bool ENABLE_L1_RESIDENT = DispatchPolicy::ENABLE_L1_RESIDENT; + static constexpr uint32_t L1A_STAGES = DispatchPolicy::L1A_STAGES; + static constexpr uint32_t L1B_STAGES = DispatchPolicy::L1B_STAGES; + static constexpr uint32_t L0A_STAGES = DispatchPolicy::L0A_STAGES; + static constexpr uint32_t L0B_STAGES = DispatchPolicy::L0B_STAGES; + static constexpr uint32_t L0C_STAGES = DispatchPolicy::L0C_STAGES; + static constexpr uint32_t L1_TILE_M = tla::get<0>(L1TileShape{}); + static constexpr uint32_t L1_TILE_N = tla::get<1>(L1TileShape{}); + static constexpr uint32_t L1_TILE_K = tla::get<2>(L1TileShape{}); + static constexpr uint32_t L0_TILE_M = tla::get<0>(L0TileShape{}); + static constexpr uint32_t L0_TILE_N = tla::get<1>(L0TileShape{}); + static constexpr uint32_t L0_TILE_K = tla::get<2>(L0TileShape{}); + + // L1 tile size + static constexpr uint32_t L1A_TILE_SIZE = L1_TILE_M * L1_TILE_K * sizeof(ElementA); + static constexpr uint32_t L1B_TILE_SIZE = L1_TILE_N * L1_TILE_K * sizeof(ElementB); + // L0 tile size + static constexpr uint32_t L0A_TILE_SIZE = L0_TILE_M * L0_TILE_K * sizeof(ElementA); + static constexpr uint32_t L0B_TILE_SIZE = L0_TILE_K * L0_TILE_N * sizeof(ElementB); + static constexpr uint32_t L0C_TILE_SIZE = L1_TILE_M * L1_TILE_N * sizeof(ElementAccumulator); + + // Check HF32_MODE + static_assert(!USE_HF32_MODE || + (USE_HF32_MODE && std::is_same_v && std::is_same_v), + "HF32 MODE only supports in float!"); + + // Check L0C_STAGES + static_assert(!(ENABLE_UNIT_FLAG && L0C_STAGES != 1), "L0C_STAGES must be 1 when UnitFlag is true!"); + + // Check LayoutC + static_assert(tla::detail::isRowMajor::value || + ((std::is_same_v || std::is_same_v || + std::is_same_v) && + tla::detail::iszN::value), + "LayoutC only supports zN in half or bfloat16 or float, RowMajor in all dtype yet!"); + + // Check L1TileShape + static_assert(L1A_TILE_SIZE * L1A_STAGES + L1B_TILE_SIZE * L1B_STAGES <= ArchTag::L1_SIZE, + "L1TileShape exceeding the L1 space!"); + + // Check L0TileShape + static_assert(L0A_TILE_SIZE * L0A_STAGES <= ArchTag::L0A_SIZE, "L0TileShape exceeding the L0A space!"); + static_assert(L0B_TILE_SIZE * L0B_STAGES <= ArchTag::L0B_SIZE, "L0TileShape exceeding the L0B space!"); + static_assert(L0C_TILE_SIZE * L0C_STAGES <= ArchTag::L0C_SIZE, "L0TileShape exceeding the L0C space!"); + + static constexpr uint32_t _32B = 32 * 8; // in bits + static_assert(L1_TILE_M == L0_TILE_M && L1_TILE_N == L0_TILE_N, + "The situation where the basic blocks of L1 and L0 differ on the m and n axes is not supported yet"); + static_assert(L0_TILE_K <= L1_TILE_K, "L0TileShape::K cannot exceed L1TileShape::K"); +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 2201) + static_assert(L1_TILE_M * SizeOfBits::value % _32B == 0, "L1TileShape::M must be 32B aligned."); + static_assert(L1_TILE_K * SizeOfBits::value % _32B == 0, "L1TileShape::K must be 32B aligned."); + static_assert(L1_TILE_K * SizeOfBits::value % _32B == 0, "L1TileShape::K must be 32B aligned."); + static_assert(L1_TILE_N * SizeOfBits::value % _32B == 0, "L1TileShape::N must be 32B aligned."); + static_assert(L0_TILE_K * SizeOfBits::value % _32B == 0, "L0TileShape::K must be 32B aligned."); +#endif + + static_assert((!HAS_BIAS && (L1A_STAGES + L1B_STAGES) <= 8) || (HAS_BIAS && (L1A_STAGES + L1B_STAGES) <= 7), + "L1 Buffer overflow: Exceeds the supported range of EVENT(0~7)"); + + static_assert((!HAS_BIAS && (L0A_STAGES + L0B_STAGES) <= 8) || (HAS_BIAS && (L0A_STAGES + L0B_STAGES) <= 7), + "L0 Buffer overflow: Exceeds the supported range of EVENT_ID(0~7)"); + + static constexpr auto L1A_LAYOUT = + tla::MakeLayout(tla::Int{}, tla::Int{}); + static constexpr auto L1B_LAYOUT = + tla::MakeLayout(tla::Int{}, tla::Int{}); + static constexpr auto L1BIAS_LAYOUT = tla::MakeLayout(tla::Int{}); + static constexpr auto L0BIAS_LAYOUT = tla::MakeLayout(tla::Int{}); + + // When enabling L1 resident mode, restore the pointer and coordinates that record the last state + // to the initial state. if two blockmmad instances need to be consecutively invoked at the kernel layer, + // RestoreStatus() must be inserted between them. + CATLASS_DEVICE + void RestoreStatus() + { + for (int i = 0; i < L1A_STAGES; ++i) { + lastAddrA[i] = nullptr; + lastCoordA[i] = MatrixCoord{0U, 0U}; + } + for (int i = 0; i < L1B_STAGES; ++i) { + lastAddrB[i] = nullptr; + lastCoordB[i] = MatrixCoord{0U, 0U}; + } + } + + /// Construct + CATLASS_DEVICE + BlockMmadTla(Arch::Resource &resource, uint32_t l1BufAddrStart = 0) + { + if ASCEND_IS_AIC { + uint32_t l1AOffset = l1BufAddrStart; + uint32_t l1BOffset = l1BufAddrStart + L1A_TILE_SIZE * L1A_STAGES; + // Init buffers + for (uint32_t i = 0; i < L1A_STAGES; i++) { + // Assign L1/L0A/L0B space for each stages + l1ATensorList[i] = resource.l1Buf.template GetBufferByByte(l1AOffset + L1A_TILE_SIZE * i); + // Assign event ID for each stages + l1AEventList[i] = i; + } + for (uint32_t i = 0; i < L1B_STAGES; i++) { + // Assign L1/L0A/L0B space for each stages + l1BTensorList[i] = resource.l1Buf.template GetBufferByByte(l1BOffset + L1B_TILE_SIZE * i); + // Assign event ID for each stages + l1BEventList[i] = i + L1A_STAGES; + } + for (uint32_t i = 0; i < L0A_STAGES; i++) { + // Assign L1/L0A/L0B space for each stages + l0ATensorList[i] = resource.l0ABuf.template GetBufferByByte(L0A_TILE_SIZE * i); + // Assign event ID for each stages + l0AEventList[i] = i; + } + for (uint32_t i = 0; i < L0B_STAGES; i++) { + // Assign L1/L0A/L0B space for each stages + l0BTensorList[i] = resource.l0BBuf.template GetBufferByByte(L0B_TILE_SIZE * i); + // Assign event ID for each stages + l0BEventList[i] = i + L0A_STAGES; + } + if constexpr (!ENABLE_UNIT_FLAG) { + for (uint32_t i = 0; i < L0C_STAGES; i++) { + l0CTensorList[i] = resource.l0CBuf.template GetBufferByByte(L0C_TILE_SIZE * i); + l0CEventList[i] = i; + } + } else { + l0CTensorList[0] = resource.l0CBuf.template GetBufferByByte(0); + } + if constexpr (HAS_BIAS) { + uint32_t l1BiasOffset = l1BOffset + L1B_TILE_SIZE * L1B_STAGES; + l1BiasTensor = resource.l1Buf.template GetBufferByByte(l1BiasOffset); + l0BiasTensor = resource.btBuf.template GetBufferByByte(0); + } + } + } + + /// Destructor + CATLASS_DEVICE + ~BlockMmadTla() {} + + CATLASS_DEVICE + void preSetFlags() + { + if ASCEND_IS_AIC { + // use HF32 when USE_HF32_MODE is true + if constexpr (USE_HF32_MODE) { + AscendC::SetHF32Mode(true); + } else { + AscendC::SetHF32Mode(false); + } + if constexpr (ENABLE_UNIT_FLAG && tla::detail::isRowMajor::value) { + AscendC::SetMMLayoutTransform(true); + } + for (uint32_t i = 0; i < L1A_STAGES; i++) { + AscendC::SetFlag(l1AEventList[i]); + } + for (uint32_t i = 0; i < L1B_STAGES; i++) { + AscendC::SetFlag(l1BEventList[i]); + } + for (uint32_t i = 0; i < L0A_STAGES; i++) { + AscendC::SetFlag(l0AEventList[i]); + } + for (uint32_t i = 0; i < L0B_STAGES; i++) { + AscendC::SetFlag(l0BEventList[i]); + } + if constexpr (!ENABLE_UNIT_FLAG) { + for (uint32_t i = 0; i < L0C_STAGES; i++) { + AscendC::SetFlag(l0CEventList[i]); + } + } + if constexpr (HAS_BIAS) { + AscendC::SetFlag(L1A_STAGES + L1B_STAGES); + AscendC::SetFlag(L0A_STAGES + L0B_STAGES); + } + + if constexpr (ENABLE_L1_RESIDENT) { + RestoreStatus(); + } + } + } + + CATLASS_DEVICE + void finalWaitFlags() + { + if ASCEND_IS_AIC { + if constexpr (USE_HF32_MODE) { + AscendC::SetHF32Mode(false); + } + if constexpr (ENABLE_UNIT_FLAG && tla::detail::isRowMajor::value) { + AscendC::SetMMLayoutTransform(false); + } + for (uint32_t i = 0; i < L1A_STAGES; i++) { + AscendC::WaitFlag(l1AEventList[i]); + } + for (uint32_t i = 0; i < L1B_STAGES; i++) { + AscendC::WaitFlag(l1BEventList[i]); + } + for (uint32_t i = 0; i < L0A_STAGES; i++) { + AscendC::WaitFlag(l0AEventList[i]); + } + for (uint32_t i = 0; i < L0B_STAGES; i++) { + AscendC::WaitFlag(l0BEventList[i]); + } + if constexpr (!ENABLE_UNIT_FLAG) { + for (uint32_t i = 0; i < L0C_STAGES; i++) { + AscendC::WaitFlag(l0CEventList[i]); + } + } + if constexpr (HAS_BIAS) { + AscendC::WaitFlag(L1A_STAGES + L1B_STAGES); + AscendC::WaitFlag(L0A_STAGES + L0B_STAGES); + } + } + } + + /// Perform a block-scoped matrix multiply-accumulate + template + CATLASS_DEVICE void operator()(TensorA &tensorA, TensorB &tensorB, TensorC &tensorC, GemmCoord const &actualShape, + TensorBias const &tensorBias = {}, bool clearL1Padding = false) + { + // Check L1TileShape + if constexpr (HAS_BIAS) { + static constexpr uint32_t BIAS_BUF_SIZE = L0_TILE_N * sizeof(ElementAccumulator); + static constexpr uint32_t L1BIAS_SIZE = L1_TILE_N * sizeof(ElementBias); + static_assert(BIAS_BUF_SIZE <= ArchTag::BIAS_SIZE, + "BIAS_BUF_SIZE exceeding the BT space! Reduce L0_TILE_N"); + static_assert(L1A_TILE_SIZE * L1A_STAGES + L1B_TILE_SIZE * L1B_STAGES + L1BIAS_SIZE <= ArchTag::L1_SIZE, + "L1TileShape exceeding the L1 space!"); + } + + using CopyGmToL1A = typename TileCopy_::template CopyGmToL1A; + using CopyGmToL1B = typename TileCopy_::template CopyGmToL1B; + CopyGmToL1A copyGmToL1A; + CopyGmToL1B copyGmToL1B; +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 2201) + using CopyL0CToGm = typename TileCopy_::template CopyL0CToGm; + CopyL0CToGm copyL0CToDst; +#endif +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 3510) + using CopyL0CToDst = typename TileCopy_::template CopyL0CToDst; + CopyL0CToDst copyL0CToDst; +#endif + + uint32_t mBlockActual = actualShape.m(); + uint32_t kBlockActual = actualShape.k(); + uint32_t nBlockActual = actualShape.n(); + + uint32_t mL1Actual = mBlockActual; + if constexpr (std::is_same_v) { + // Avoid using the gemv mode in mmad + if (mL1Actual == 1) { + mL1Actual = 16; + } + } + uint32_t nL1Actual = nBlockActual; + + auto layoutInL0C = tla::MakeLayoutL0C(mL1Actual, nL1Actual); + auto tensorL0C = tla::MakeTensor(l0CTensorList[l0CListId], layoutInL0C, Arch::PositionL0C{}); + auto tensorL0Bias = tla::MakeTensor(l0BiasTensor, L0BIAS_LAYOUT, Arch::PositionBias{}); + + uint32_t kL1Actual = min(kBlockActual, L1_TILE_K); + // load first matrix A tile from GM to L1 + AscendC::WaitFlag(l1AEventList[l1AListId]); + if (clearL1Padding) { + AscendC::InitConstValueParams clearParams(1, static_cast(L1A_TILE_SIZE / 32), 0, + static_cast(0)); + AscendC::InitConstValue(l1ATensorList[l1AListId], clearParams); + } + auto tensorL1A = tla::MakeTensor(l1ATensorList[l1AListId], L1A_LAYOUT, Arch::PositionL1{}); + auto tensorTileA = GetTileA(tensorA, 0, 0, mBlockActual, kL1Actual); + if constexpr (ENABLE_L1_RESIDENT) { + // If the currently loaded GM pointer and block coordinates are the same as the last loaded ones, + // skip this loading. + if (lastAddrA[l1AListId] != tensorTileA.data().GetPhyAddr() || + tla::get<0>(tensorTileA.coord()) != lastCoordA[l1AListId].row() || + tla::get<1>(tensorTileA.coord()) != lastCoordA[l1AListId].column()) { + copyGmToL1A(tensorL1A, tensorTileA); + lastCoordA[l1AListId] = MatrixCoord{static_cast(tla::get<0>(tensorTileA.coord())), + static_cast(tla::get<1>(tensorTileA.coord()))}; + lastAddrA[l1AListId] = const_cast<__gm__ typename AscendC::GlobalTensor::PrimType *>( + tensorTileA.data().GetPhyAddr()); + } + } else { + copyGmToL1A(tensorL1A, tensorTileA); + } + AscendC::SetFlag(l1AEventList[l1AListId]); + + // load first matrix B tile from GM to L1 + AscendC::WaitFlag(l1BEventList[l1BListId]); + if (clearL1Padding) { + AscendC::InitConstValueParams clearParams(1, static_cast(L1B_TILE_SIZE / 32), 0, + static_cast(0)); + AscendC::InitConstValue(l1BTensorList[l1BListId], clearParams); + } + auto tensorL1B = tla::MakeTensor(l1BTensorList[l1BListId], L1B_LAYOUT, Arch::PositionL1{}); + auto tensorTileB = GetTile(tensorB, tla::MakeCoord(0, 0), tla::MakeShape(kL1Actual, nBlockActual)); + if constexpr (ENABLE_L1_RESIDENT) { + if (lastAddrB[l1BListId] != tensorTileB.data().GetPhyAddr() || + tla::get<0>(tensorTileB.coord()) != lastCoordB[l1BListId].row() || + tla::get<1>(tensorTileB.coord()) != lastCoordB[l1BListId].column()) { + copyGmToL1B(tensorL1B, tensorTileB); + lastCoordB[l1BListId] = MatrixCoord{static_cast(tla::get<0>(tensorTileB.coord())), + static_cast(tla::get<1>(tensorTileB.coord()))}; + lastAddrB[l1BListId] = const_cast<__gm__ typename AscendC::GlobalTensor::PrimType *>( + tensorTileB.data().GetPhyAddr()); + } + } else { + copyGmToL1B(tensorL1B, tensorTileB); + } + AscendC::SetFlag(l1BEventList[l1BListId]); + + if constexpr (HAS_BIAS && !std::is_same_v) { + using CopyGmToL1Bias = typename TileCopy::template CopyGmToL1Bias; + CopyGmToL1Bias copyGmToL1Bias; + AscendC::WaitFlag(L1A_STAGES + L1B_STAGES); + auto l1Bias = l1BiasTensor.template ReinterpretCast(); + auto tensorL1Bias = tla::MakeTensor(l1Bias, L1BIAS_LAYOUT, Arch::PositionL1{}); + copyGmToL1Bias(tensorL1Bias, tensorBias); + AscendC::SetFlag(L1A_STAGES + L1B_STAGES); + } + + if constexpr (!ENABLE_UNIT_FLAG) { + AscendC::WaitFlag(l0CEventList[l0CListId]); + } + + uint32_t mL0Loop = CeilDiv(mL1Actual); + uint32_t nL0Loop = CeilDiv(nL1Actual); + + // main loop + uint32_t kL1Loop = CeilDiv(kBlockActual); + for (uint32_t kL1Idx = 0; kL1Idx < kL1Loop; kL1Idx++) { + uint32_t l1AListIdNext = (l1AListId + 1 < L1A_STAGES) ? (l1AListId + 1) : 0; + uint32_t l1BListIdNext = (l1BListId + 1 < L1B_STAGES) ? (l1BListId + 1) : 0; + uint32_t kL1ActualNext{0}; + // preload next tile from GM to L1 + if (kL1Idx < kL1Loop - 1) { + uint32_t kL1IdxNext = kL1Idx + 1; + kL1ActualNext = (kL1IdxNext < kL1Loop - 1) ? L1_TILE_K : (kBlockActual - kL1IdxNext * L1_TILE_K); + + // Get L1 tensor for next stage + auto l1ATensor = l1ATensorList[l1AListIdNext]; + auto l1BTensor = l1BTensorList[l1BListIdNext]; + auto tensorL1A = tla::MakeTensor(l1ATensor, L1A_LAYOUT, Arch::PositionL1{}); + auto tensorL1B = tla::MakeTensor(l1BTensor, L1B_LAYOUT, Arch::PositionL1{}); + // Get GM tile for next stage + auto tensorTileA = GetTileA(tensorA, 0, kL1IdxNext * L1_TILE_K, mBlockActual, kL1ActualNext); + auto tensorTileB = GetTile(tensorB, tla::MakeCoord(kL1IdxNext * L1_TILE_K, 0), + tla::MakeShape(kL1ActualNext, nBlockActual)); + + // load next matrix A tile from GM to L1 + AscendC::WaitFlag(l1AEventList[l1AListIdNext]); + if (clearL1Padding) { + AscendC::InitConstValueParams clearParams(1, static_cast(L1A_TILE_SIZE / 32), 0, + static_cast(0)); + AscendC::InitConstValue(l1ATensorList[l1AListIdNext], clearParams); + } + if constexpr (ENABLE_L1_RESIDENT) { + if (lastAddrA[l1AListIdNext] != tensorTileA.data().GetPhyAddr() || + tla::get<0>(tensorTileA.coord()) != lastCoordA[l1AListIdNext].row() || + tla::get<1>(tensorTileA.coord()) != lastCoordA[l1AListIdNext].column()) { + copyGmToL1A(tensorL1A, tensorTileA); + lastCoordA[l1AListIdNext] = + MatrixCoord{static_cast(tla::get<0>(tensorTileA.coord())), + static_cast(tla::get<1>(tensorTileA.coord()))}; + lastAddrA[l1AListIdNext] = + const_cast<__gm__ typename AscendC::GlobalTensor::PrimType *>( + tensorTileA.data().GetPhyAddr()); + } + } else { + copyGmToL1A(tensorL1A, tensorTileA); + } + AscendC::SetFlag(l1AEventList[l1AListIdNext]); + + // load next matrix B tile from GM to L1 + AscendC::WaitFlag(l1BEventList[l1BListIdNext]); + if (clearL1Padding) { + AscendC::InitConstValueParams clearParams(1, static_cast(L1B_TILE_SIZE / 32), 0, + static_cast(0)); + AscendC::InitConstValue(l1BTensorList[l1BListIdNext], clearParams); + } + if constexpr (ENABLE_L1_RESIDENT) { + if (lastAddrB[l1BListIdNext] != tensorTileB.data().GetPhyAddr() || + tla::get<0>(tensorTileB.coord()) != lastCoordB[l1BListIdNext].row() || + tla::get<1>(tensorTileB.coord()) != lastCoordB[l1BListIdNext].column()) { + copyGmToL1B(tensorL1B, tensorTileB); + lastCoordB[l1BListIdNext] = + MatrixCoord{static_cast(tla::get<0>(tensorTileB.coord())), + static_cast(tla::get<1>(tensorTileB.coord()))}; + lastAddrB[l1BListIdNext] = + const_cast<__gm__ typename AscendC::GlobalTensor::PrimType *>( + tensorTileB.data().GetPhyAddr()); + } + } else { + copyGmToL1B(tensorL1B, tensorTileB); + } + AscendC::SetFlag(l1BEventList[l1BListIdNext]); + } + + // Get L1 tensor for current stage + auto l1ATensor = l1ATensorList[l1AListId]; + auto l1BTensor = l1BTensorList[l1BListId]; + tensorL1A = tla::MakeTensor(l1ATensor, L1A_LAYOUT, Arch::PositionL1{}); + tensorL1B = tla::MakeTensor(l1BTensor, L1B_LAYOUT, Arch::PositionL1{}); + // Get the loop nums on L0 + uint32_t kL0Loop = CeilDiv(kL1Actual); + + for (int mL0Idx = 0; mL0Idx < mL0Loop; mL0Idx++) { + uint32_t mL0Actual = (mL0Idx < mL0Loop - 1) ? L0_TILE_M : (mL1Actual - mL0Idx * L0_TILE_M); + + for (int kL0Idx = 0; kL0Idx < kL0Loop; kL0Idx++) { + uint32_t kL0Actual = (kL0Idx < kL0Loop - 1) ? L0_TILE_K : (kL1Actual - kL0Idx * L0_TILE_K); + + // Locate the current tile on L0A + auto l0ATile = l0ATensorList[l0AListId]; + auto layoutAInL0 = tla::MakeLayout(mL0Actual, kL0Actual); + auto tensorL0A = tla::MakeTensor(l0ATile, layoutAInL0, Arch::PositionL0A{}); + // Locate the current tile of matrix A on L1 + auto tensorTileL1A = + GetTileA(tensorL1A, mL0Idx * L0_TILE_M, kL0Idx * L0_TILE_K, mL0Actual, kL0Actual); + + AscendC::WaitFlag(l0AEventList[l0AListId]); + if ((mL0Idx == 0) && (kL0Idx == 0)) { + AscendC::WaitFlag(l1AEventList[l1AListId]); + } + + // Load current tile from L1 to L0A + copyL1ToL0A(tensorL0A, tensorTileL1A); + + if ((mL0Idx == mL0Loop - 1) && (kL0Idx == kL0Loop - 1)) { + AscendC::SetFlag(l1AEventList[l1AListId]); + } + + bool initC = ((kL1Idx == 0) && (kL0Idx == 0)); + for (int nL0Idx = 0; nL0Idx < nL0Loop; nL0Idx++) { + uint32_t nL0Actual = (nL0Idx < nL0Loop - 1) ? L0_TILE_N : (nL1Actual - nL0Idx * L0_TILE_N); + + // Locate the current tile on L0B + auto l0BTile = l0BTensorList[l0BListId]; + auto layoutBInL0 = tla::MakeLayout(kL0Actual, nL0Actual); + auto tensorL0B = tla::MakeTensor(l0BTile, layoutBInL0, Arch::PositionL0B{}); + // Locate the current tile of matrix B on L1 + auto tensorTileL1B = GetTile(tensorL1B, tla::MakeCoord(kL0Idx * L0_TILE_K, nL0Idx * L0_TILE_N), + tla::MakeShape(kL0Actual, nL0Actual)); + + // Wait for mmad finished + AscendC::WaitFlag(l0BEventList[l0BListId]); + // If the current tile is the first one on the k&n axis, wait for loading matrix B from GM to L1 + if ((mL0Idx == 0) && (kL0Idx == 0) && (nL0Idx == 0)) { + AscendC::WaitFlag(l1BEventList[l1BListId]); + } + + // Load current tile from L1 to L0B + copyL1ToL0B(tensorL0B, tensorTileL1B); + + // If the current tile is the last one on the k&n axis, notify to load matrix B from GM to L1 + if ((mL0Idx == mL0Loop - 1) && (kL0Idx == kL0Loop - 1) && (nL0Idx == nL0Loop - 1)) { + AscendC::SetFlag(l1BEventList[l1BListId]); + } + + if constexpr (HAS_BIAS && !std::is_same_v) { + if (initC) { + if (nL0Idx == 0) { + AscendC::WaitFlag(L1A_STAGES + L1B_STAGES); + } + AscendC::WaitFlag(L0A_STAGES + L0B_STAGES); + auto l1Bias = l1BiasTensor.template ReinterpretCast(); + auto tensorL1Bias = tla::MakeTensor(l1Bias, L1BIAS_LAYOUT, Arch::PositionL1{}); + auto tensorTileL1Bias = GetTile(tensorL1Bias, tla::MakeCoord(nL0Idx * L0_TILE_N), + tla::MakeShape(nL0Actual)); + // Load bias to l0 biasTable + copyL1ToBT(tensorL0Bias, tensorTileL1Bias); + if (nL0Idx == nL0Loop - 1) { + AscendC::SetFlag(L1A_STAGES + L1B_STAGES); + } + } + } + + // Notify to do mmad + AscendC::SetFlag(l0CEventList[l0CListId]); + + // Locate the current tile on L0C + auto tensorTileL0C = GetTile(tensorL0C, tla::MakeCoord(mL0Idx * L0_TILE_M, nL0Idx * L0_TILE_N), + tla::MakeShape(mL0Actual, nL0Actual)); + + // Compute the matrix multiplication on L0A and L0B and write the result to the accumulator + // Wait for loading L0B + AscendC::WaitFlag(l0CEventList[l0CListId]); + + // If the unit flag is enabled, the unit flag is set according to the calculation progress + uint8_t unitFlag = 0b00; + if constexpr (ENABLE_UNIT_FLAG) { + if ((kL1Idx == kL1Loop - 1) && (mL0Idx == mL0Loop - 1) && (kL0Idx == kL0Loop - 1) && + (nL0Idx == nL0Loop - 1)) { + unitFlag = 0b11; + } else { + unitFlag = 0b10; + } + } + + if constexpr (HAS_BIAS && !std::is_same_v) { + if (initC) { + tileMmad(tensorTileL0C, tensorL0A, tensorL0B, tensorL0Bias, mL0Actual, nL0Actual, + kL0Actual, initC, unitFlag); + AscendC::SetFlag(L0A_STAGES + L0B_STAGES); + } else { + tileMmad(tensorTileL0C, tensorL0A, tensorL0B, mL0Actual, nL0Actual, kL0Actual, initC, + unitFlag); + } + } else { + tileMmad(tensorTileL0C, tensorL0A, tensorL0B, mL0Actual, nL0Actual, kL0Actual, initC, + unitFlag); + } + + // Notify to move the next L0B tile + AscendC::SetFlag(l0BEventList[l0BListId]); + l0BListId = (l0BListId + 1 < L0B_STAGES) ? (l0BListId + 1) : 0; + } + AscendC::SetFlag(l0AEventList[l0AListId]); + l0AListId = (l0AListId + 1 < L0A_STAGES) ? (l0AListId + 1) : 0; + } + } + l1AListId = l1AListIdNext; + l1BListId = l1BListIdNext; + kL1Actual = kL1ActualNext; + } + + // copy block out + if constexpr (!ENABLE_UNIT_FLAG) { + AscendC::SetFlag(l0CEventList[l0CListId]); + AscendC::WaitFlag(l0CEventList[l0CListId]); + copyL0CToDst(tensorC, tensorL0C); + AscendC::SetFlag(l0CEventList[l0CListId]); + l0CListId = (l0CListId + 1 < L0C_STAGES) ? (l0CListId + 1) : 0; + } else { + copyL0CToDst(tensorC, tensorL0C, 0b11); + } + } + +protected: + template + CATLASS_DEVICE auto GetTileA(TensorA &tensorA, uint32_t mIndex, uint32_t kIndex, uint32_t mSize, uint32_t kSize) + { + if constexpr (tla::detail::isVector::value) { + return GetTile(tensorA, tla::MakeCoord(kIndex), tla::MakeShape(kSize)); + } else { + return GetTile(tensorA, tla::MakeCoord(mIndex, kIndex), tla::MakeShape(mSize, kSize)); + } + } + + // Multi-stage tensors list + AscendC::LocalTensor l1ATensorList[L1A_STAGES]; + AscendC::LocalTensor l1BTensorList[L1B_STAGES]; + AscendC::LocalTensor l0ATensorList[L0A_STAGES]; + AscendC::LocalTensor l0BTensorList[L0B_STAGES]; + AscendC::LocalTensor l0CTensorList[L0C_STAGES]; + AscendC::LocalTensor l1BiasTensor; + AscendC::LocalTensor l0BiasTensor; + + // Multi-stage event id list + int32_t l1AEventList[L1A_STAGES]; + int32_t l1BEventList[L1B_STAGES]; + int32_t l0AEventList[L0A_STAGES]; + int32_t l0BEventList[L0B_STAGES]; + int32_t l0CEventList[L0C_STAGES]; + + __gm__ typename AscendC::GlobalTensor::PrimType *lastAddrA[L1A_STAGES]; + __gm__ typename AscendC::GlobalTensor::PrimType *lastAddrB[L1B_STAGES]; + MatrixCoord lastCoordA[L1A_STAGES]; + MatrixCoord lastCoordB[L1B_STAGES]; + + // The id of current stage + uint32_t l1AListId{0}; + uint32_t l1BListId{0}; + uint32_t l0AListId{0}; + uint32_t l0BListId{0}; + uint32_t l0CListId{0}; + + TileMmad tileMmad; + CopyL1ToL0A copyL1ToL0A; + CopyL1ToL0B copyL1ToL0B; + CopyL1ToBT copyL1ToBT; +}; + +} // namespace Catlass::Gemm::Block + +#endif // CATLASS_GEMM_BLOCK_BLOCK_MMAD_PINGPONG_TLA_MULTI_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/kernel_utils/block/block_mmad_pingpong_tla_preloadA_l1B.hpp b/csrc/chunk_kda_fwd/op_kernel/kernel_utils/block/block_mmad_pingpong_tla_preloadA_l1B.hpp new file mode 100644 index 0000000000..6b1341c9d9 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/kernel_utils/block/block_mmad_pingpong_tla_preloadA_l1B.hpp @@ -0,0 +1,498 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_GEMM_BLOCK_BLOCK_MMAD_PINGPONG_TLA_PRELOADA_L1B_HPP +#define CATLASS_GEMM_BLOCK_BLOCK_MMAD_PINGPONG_TLA_PRELOADA_L1B_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/kda_coord.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" +#include "catlass/gemm/kda_gemm_helper.hpp" +#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" +#include "catlass/gemm/tile/kda_tile_mmad.hpp" +#include "tla/kda_tla_layout.hpp" +#include "tla/kda_tensor.hpp" + +namespace Catlass::Gemm { + +template +struct MmadPingpongTlaPreloadAL1B : public MmadBase { + static constexpr uint32_t L1A_STAGES = L1A_STAGES_; + static constexpr uint32_t L1B_STAGES = L1B_STAGES_; + static constexpr uint32_t L0A_STAGES = L0A_STAGES_; + static constexpr uint32_t L0B_STAGES = L0B_STAGES_; + static constexpr uint32_t L0C_STAGES = L0C_STAGES_; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; +}; + +} // namespace Catlass::Gemm + +namespace Catlass::Gemm::Block { + +template +struct BlockMmadTla, + L1TileShape_, L0TileShape_, ElementA_, ElementB_, ElementC_, ElementBias_, TileCopy_, TileMmad_> { +public: + // Type Aliases + using DispatchPolicy = MmadPingpongTlaPreloadAL1B; + using ArchTag = typename DispatchPolicy::ArchTag; + using TileCopy = TileCopy_; + using L1TileShape = L1TileShape_; + using L0TileShape = L0TileShape_; + using ElementA = ElementA_; + using LayoutA = typename TileCopy::LayoutA; + using ElementB = ElementB_; + using LayoutB = typename TileCopy::LayoutB; + using ElementC = ElementC_; + using LayoutC = typename TileCopy::LayoutC; + + using TileMmad = TileMmad_; + + using CopyL1ToL0A = typename TileCopy::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopy::CopyL1ToL0B; + using CopyL1ToBT = typename TileCopy::CopyL1ToBT; + + using ElementAccumulator = typename TileCopy::ElementAccumulator; + + using LayoutTagL1A = typename TileCopy::LayoutTagL1A; + using LayoutTagL1B = typename TileCopy::LayoutTagL1B; + using LayoutTagL0A = typename TileCopy::LayoutTagL0A; + using LayoutTagL0B = typename TileCopy::LayoutTagL0B; + + using L1AAlignHelper = typename TileCopy_::L1AAlignHelper; + using L1BAlignHelper = typename TileCopy_::L1BAlignHelper; + + static_assert(tla::is_tuple::value && tla::is_static::value, + "L1TileShape must be tla::tuple and static!"); + static_assert(tla::is_tuple::value && tla::is_static::value, + "L0TileShape must be tla::tuple and static!"); + + static constexpr bool ENABLE_UNIT_FLAG = DispatchPolicy::ENABLE_UNIT_FLAG; + static constexpr uint32_t L1A_STAGES = DispatchPolicy::L1A_STAGES; + static constexpr uint32_t L1B_STAGES = DispatchPolicy::L1B_STAGES; + static constexpr uint32_t L0A_STAGES = DispatchPolicy::L0A_STAGES; + static constexpr uint32_t L0B_STAGES = DispatchPolicy::L0B_STAGES; + static constexpr uint32_t L0C_STAGES = DispatchPolicy::L0C_STAGES; + static constexpr uint32_t L1_TILE_M = tla::get<0>(L1TileShape{}); + static constexpr uint32_t L1_TILE_N = tla::get<1>(L1TileShape{}); + static constexpr uint32_t L1_TILE_K = tla::get<2>(L1TileShape{}); + static constexpr uint32_t L0_TILE_M = tla::get<0>(L0TileShape{}); + static constexpr uint32_t L0_TILE_N = tla::get<1>(L0TileShape{}); + static constexpr uint32_t L0_TILE_K = tla::get<2>(L0TileShape{}); + + // L1 tile size + static constexpr uint32_t L1A_TILE_SIZE = L1_TILE_M * L1_TILE_K * sizeof(ElementA); + static constexpr uint32_t L1B_TILE_SIZE = L1_TILE_N * L1_TILE_K * sizeof(ElementB); + // L0 tile size + static constexpr uint32_t L0A_TILE_SIZE = L0_TILE_M * L0_TILE_K * sizeof(ElementA); + static constexpr uint32_t L0B_TILE_SIZE = L0_TILE_K * L0_TILE_N * sizeof(ElementB); + static constexpr uint32_t L0C_TILE_SIZE = L1_TILE_M * L1_TILE_N * sizeof(ElementAccumulator); + + // Check L0C_STAGES + static_assert(!(ENABLE_UNIT_FLAG && L0C_STAGES != 1), "L0C_STAGES must be 1 when UnitFlag is true!"); + + // Check LayoutC + static_assert(tla::detail::isRowMajor::value || + ((std::is_same_v || std::is_same_v || + std::is_same_v) && + tla::detail::iszN::value), + "LayoutC only supports zN in half or bfloat16 or float, RowMajor in all dtype yet!"); + + // Check L1TileShape + static_assert(L1A_TILE_SIZE * L1A_STAGES + L1B_TILE_SIZE * L1B_STAGES <= ArchTag::L1_SIZE, + "L1TileShape exceeding the L1 space!"); + + // Check L0TileShape + static_assert(L0A_TILE_SIZE * L0A_STAGES <= ArchTag::L0A_SIZE, "L0TileShape exceeding the L0A space!"); + static_assert(L0B_TILE_SIZE * L0B_STAGES <= ArchTag::L0B_SIZE, "L0TileShape exceeding the L0B space!"); + static_assert(L0C_TILE_SIZE * L0C_STAGES <= ArchTag::L0C_SIZE, "L0TileShape exceeding the L0C space!"); + + static constexpr uint32_t _32B = 32 * 8; // in bits + static_assert(L1_TILE_M == L0_TILE_M && L1_TILE_N == L0_TILE_N, + "The situation where the basic blocks of L1 and L0 differ on the m and n axes is not supported yet"); + static_assert(L0_TILE_K <= L1_TILE_K, "L0TileShape::K cannot exceed L1TileShape::K"); +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 2201) + static_assert(L1_TILE_M * SizeOfBits::value % _32B == 0, "L1TileShape::M must be 32B aligned."); + static_assert(L1_TILE_K * SizeOfBits::value % _32B == 0, "L1TileShape::K must be 32B aligned."); + static_assert(L1_TILE_K * SizeOfBits::value % _32B == 0, "L1TileShape::K must be 32B aligned."); + static_assert(L1_TILE_N * SizeOfBits::value % _32B == 0, "L1TileShape::N must be 32B aligned."); + static_assert(L0_TILE_K * SizeOfBits::value % _32B == 0, "L0TileShape::K must be 32B aligned."); +#endif + + static constexpr auto L1A_LAYOUT = + tla::MakeLayout(tla::Int{}, tla::Int{}); + static constexpr auto L1B_LAYOUT = + tla::MakeLayout(tla::Int{}, tla::Int{}); + static constexpr auto L1BIAS_LAYOUT = tla::MakeLayout(tla::Int{}); + static constexpr auto L0BIAS_LAYOUT = tla::MakeLayout(tla::Int{}); + + // When enabling L1 resident mode, restore the pointer and coordinates that record the last state + // to the initial state. if two blockmmad instances need to be consecutively invoked at the kernel layer, + // RestoreStatus() must be inserted between them. + CATLASS_DEVICE + void RestoreStatus() + { + for (int i = 0; i < L1A_STAGES; ++i) { + lastAddrA[i] = nullptr; + lastCoordA[i] = MatrixCoord{0U, 0U}; + } + for (int i = 0; i < L1B_STAGES; ++i) { + lastAddrB[i] = nullptr; + lastCoordB[i] = MatrixCoord{0U, 0U}; + } + } + + /// Construct + CATLASS_DEVICE + BlockMmadTla(Arch::Resource &resource, uint32_t l1BufAddrStart = 0) + { + if ASCEND_IS_AIC { + uint32_t l1AOffset = l1BufAddrStart; + uint32_t l1BOffset = l1BufAddrStart + L1A_TILE_SIZE * L1A_STAGES; + // Init buffers + for (uint32_t i = 0; i < L1A_STAGES; i++) { + // Assign L1/L0A/L0B space for each stages + l1ATensorList[i] = resource.l1Buf.template GetBufferByByte(l1AOffset + L1A_TILE_SIZE * i); + // Assign event ID for each stages + l1AEventList[i] = i; + } + for (uint32_t i = 0; i < L1B_STAGES; i++) { + // Assign L1/L0A/L0B space for each stages + l1BTensorList[i] = resource.l1Buf.template GetBufferByByte(l1BOffset + L1B_TILE_SIZE * i); + // Assign event ID for each stages + l1BEventList[i] = i + L1A_STAGES; + } + for (uint32_t i = 0; i < L0A_STAGES; i++) { + // Assign L1/L0A/L0B space for each stages + l0ATensorList[i] = resource.l0ABuf.template GetBufferByByte(L0A_TILE_SIZE * i); + // Assign event ID for each stages + l0AEventList[i] = i; + } + for (uint32_t i = 0; i < L0B_STAGES; i++) { + // Assign L1/L0A/L0B space for each stages + l0BTensorList[i] = resource.l0BBuf.template GetBufferByByte(L0B_TILE_SIZE * i); + // Assign event ID for each stages + l0BEventList[i] = i + L0A_STAGES; + } + if constexpr (!ENABLE_UNIT_FLAG) { + for (uint32_t i = 0; i < L0C_STAGES; i++) { + l0CTensorList[i] = resource.l0CBuf.template GetBufferByByte(L0C_TILE_SIZE * i); + l0CEventList[i] = i; + } + } else { + l0CTensorList[0] = resource.l0CBuf.template GetBufferByByte(0); + } + } + } + + /// Destructor + CATLASS_DEVICE + ~BlockMmadTla() {} + + CATLASS_DEVICE + void preSetFlags() + { + if ASCEND_IS_AIC { + if constexpr (ENABLE_UNIT_FLAG && tla::detail::isRowMajor::value) { + AscendC::SetMMLayoutTransform(true); + } + for (uint32_t i = 0; i < L1A_STAGES; i++) { + AscendC::SetFlag(l1AEventList[i]); + } + for (uint32_t i = 0; i < L1B_STAGES; i++) { + AscendC::SetFlag(l1BEventList[i]); + } + for (uint32_t i = 0; i < L0A_STAGES; i++) { + AscendC::SetFlag(l0AEventList[i]); + } + for (uint32_t i = 0; i < L0B_STAGES; i++) { + AscendC::SetFlag(l0BEventList[i]); + } + if constexpr (!ENABLE_UNIT_FLAG) { + for (uint32_t i = 0; i < L0C_STAGES; i++) { + AscendC::SetFlag(l0CEventList[i]); + } + } + } + } + + CATLASS_DEVICE + void finalWaitFlags() + { + if ASCEND_IS_AIC { + if constexpr (ENABLE_UNIT_FLAG && tla::detail::isRowMajor::value) { + AscendC::SetMMLayoutTransform(false); + } + for (uint32_t i = 0; i < L1A_STAGES; i++) { + AscendC::WaitFlag(l1AEventList[i]); + } + for (uint32_t i = 0; i < L1B_STAGES; i++) { + AscendC::WaitFlag(l1BEventList[i]); + } + for (uint32_t i = 0; i < L0A_STAGES; i++) { + AscendC::WaitFlag(l0AEventList[i]); + } + for (uint32_t i = 0; i < L0B_STAGES; i++) { + AscendC::WaitFlag(l0BEventList[i]); + } + if constexpr (!ENABLE_UNIT_FLAG) { + for (uint32_t i = 0; i < L0C_STAGES; i++) { + AscendC::WaitFlag(l0CEventList[i]); + } + } + } + } + + /// Perform a block-scoped matrix multiply-accumulate + template + CATLASS_DEVICE void operator()(TensorA &tensorA, TensorB &tensorB, TensorC &tensorC, GemmCoord const &actualShape) + { + using CopyGmToL1A = typename TileCopy_::template CopyGmToL1A; + CopyGmToL1A copyGmToL1A; +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 2201) + using CopyL0CToGm = typename TileCopy_::template CopyL0CToGm; + CopyL0CToGm copyL0CToDst; +#endif +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 3510) + using CopyL0CToDst = typename TileCopy_::template CopyL0CToDst; + CopyL0CToDst copyL0CToDst; +#endif + + uint32_t mBlockActual = actualShape.m(); + uint32_t kBlockActual = actualShape.k(); + uint32_t nBlockActual = actualShape.n(); + + uint32_t mL1Actual = mBlockActual; + if constexpr (std::is_same_v) { + // Avoid using the gemv mode in mmad + if (mL1Actual == 1) { + mL1Actual = 16; + } + } + uint32_t nL1Actual = nBlockActual; + + auto layoutInL0C = tla::MakeLayoutL0C(mL1Actual, nL1Actual); + auto tensorL0C = tla::MakeTensor(l0CTensorList[l0CListId], layoutInL0C, Arch::PositionL0C{}); + auto tensorL0Bias = tla::MakeTensor(l0BiasTensor, L0BIAS_LAYOUT, Arch::PositionBias{}); + + uint32_t kL1Actual = min(kBlockActual, L1_TILE_K); + // load first matrix A tile from GM to L1 + AscendC::WaitFlag(l1AEventList[l1AListId]); + auto tensorL1A = tla::MakeTensor(l1ATensorList[l1AListId], L1A_LAYOUT, Arch::PositionL1{}); + auto tensorTileA = GetTileA(tensorA, 0, 0, mBlockActual, kL1Actual); + copyGmToL1A(tensorL1A, tensorTileA); + AscendC::SetFlag(l1AEventList[l1AListId]); + + // load first matrix B tile from GM to L1 + AscendC::WaitFlag(l1BEventList[l1BListId]); + AscendC::SetFlag(l1BEventList[l1BListId]); + + if constexpr (!ENABLE_UNIT_FLAG) { + AscendC::WaitFlag(l0CEventList[l0CListId]); + } + + uint32_t mL0Loop = CeilDiv(mL1Actual); + uint32_t nL0Loop = CeilDiv(nL1Actual); + + // main loop + uint32_t kL1Loop = CeilDiv(kBlockActual); + for (uint32_t kL1Idx = 0; kL1Idx < kL1Loop; kL1Idx++) { + uint32_t l1AListIdNext = (l1AListId + 1 < L1A_STAGES) ? (l1AListId + 1) : 0; + uint32_t l1BListIdNext = (l1BListId + 1 < L1B_STAGES) ? (l1BListId + 1) : 0; + uint32_t kL1ActualNext{0}; + // preload next tile from GM to L1 + if (kL1Idx < kL1Loop - 1) { + uint32_t kL1IdxNext = kL1Idx + 1; + kL1ActualNext = (kL1IdxNext < kL1Loop - 1) ? L1_TILE_K : (kBlockActual - kL1IdxNext * L1_TILE_K); + + // Get L1 tensor for next stage + auto l1ATensor = l1ATensorList[l1AListIdNext]; + auto l1BTensor = l1BTensorList[l1BListIdNext]; + auto tensorL1A = tla::MakeTensor(l1ATensor, L1A_LAYOUT, Arch::PositionL1{}); + auto tensorL1B = tla::MakeTensor(l1BTensor, L1B_LAYOUT, Arch::PositionL1{}); + // Get GM tile for next stage + auto tensorTileA = GetTileA(tensorA, 0, kL1IdxNext * L1_TILE_K, mBlockActual, kL1ActualNext); + auto tensorTileB = GetTile(tensorB, tla::MakeCoord(kL1IdxNext * L1_TILE_K, 0), + tla::MakeShape(kL1ActualNext, nBlockActual)); + + // load next matrix A tile from GM to L1 + AscendC::WaitFlag(l1AEventList[l1AListIdNext]); + copyGmToL1A(tensorL1A, tensorTileA); + AscendC::SetFlag(l1AEventList[l1AListIdNext]); + + // load next matrix B tile from GM to L1 + AscendC::WaitFlag(l1BEventList[l1BListIdNext]); + // copyGmToL1B(tensorL1B, tensorTileB); + AscendC::SetFlag(l1BEventList[l1BListIdNext]); + } + + // Get L1 tensor for current stage + auto l1ATensor = l1ATensorList[l1AListId]; + tensorL1A = tla::MakeTensor(l1ATensor, L1A_LAYOUT, Arch::PositionL1{}); + // Get the loop nums on L0 + uint32_t kL0Loop = CeilDiv(kL1Actual); + + for (int mL0Idx = 0; mL0Idx < mL0Loop; mL0Idx++) { + uint32_t mL0Actual = (mL0Idx < mL0Loop - 1) ? L0_TILE_M : (mL1Actual - mL0Idx * L0_TILE_M); + + for (int kL0Idx = 0; kL0Idx < kL0Loop; kL0Idx++) { + uint32_t kL0Actual = (kL0Idx < kL0Loop - 1) ? L0_TILE_K : (kL1Actual - kL0Idx * L0_TILE_K); + + // Locate the current tile on L0A + auto l0ATile = l0ATensorList[l0AListId]; + auto layoutAInL0 = tla::MakeLayout(mL0Actual, kL0Actual); + auto tensorL0A = tla::MakeTensor(l0ATile, layoutAInL0, Arch::PositionL0A{}); + // Locate the current tile of matrix A on L1 + auto tensorTileL1A = + GetTileA(tensorL1A, mL0Idx * L0_TILE_M, kL0Idx * L0_TILE_K, mL0Actual, kL0Actual); + + AscendC::WaitFlag(l0AEventList[l0AListId]); + if ((mL0Idx == 0) && (kL0Idx == 0)) { + AscendC::WaitFlag(l1AEventList[l1AListId]); + } + + // Load current tile from L1 to L0A + copyL1ToL0A(tensorL0A, tensorTileL1A); + + if ((mL0Idx == mL0Loop - 1) && (kL0Idx == kL0Loop - 1)) { + AscendC::SetFlag(l1AEventList[l1AListId]); + } + + bool initC = ((kL1Idx == 0) && (kL0Idx == 0)); + for (int nL0Idx = 0; nL0Idx < nL0Loop; nL0Idx++) { + uint32_t nL0Actual = (nL0Idx < nL0Loop - 1) ? L0_TILE_N : (nL1Actual - nL0Idx * L0_TILE_N); + + // Locate the current tile on L0B + auto l0BTile = l0BTensorList[l0BListId]; + auto layoutBInL0 = tla::MakeLayout(kL0Actual, nL0Actual); + auto tensorL0B = tla::MakeTensor(l0BTile, layoutBInL0, Arch::PositionL0B{}); + // Locate the current tile of matrix B on L1 + auto tensorTileL1B = GetTile(tensorB, tla::MakeCoord(kL0Idx * L0_TILE_K, nL0Idx * L0_TILE_N), + tla::MakeShape(kL0Actual, nL0Actual)); + + // Wait for mmad finished + AscendC::WaitFlag(l0BEventList[l0BListId]); + // If the current tile is the first one on the k&n axis, wait for loading matrix B from GM to L1 + if ((mL0Idx == 0) && (kL0Idx == 0) && (nL0Idx == 0)) { + AscendC::WaitFlag(l1BEventList[l1BListId]); + } + + // Load current tile from L1 to L0B + copyL1ToL0B(tensorL0B, tensorTileL1B); + + // If the current tile is the last one on the k&n axis, notify to load matrix B from GM to L1 + if ((mL0Idx == mL0Loop - 1) && (kL0Idx == kL0Loop - 1) && (nL0Idx == nL0Loop - 1)) { + AscendC::SetFlag(l1BEventList[l1BListId]); + } + + // Notify to do mmad + AscendC::SetFlag(l0CEventList[l0CListId]); + + // Locate the current tile on L0C + auto tensorTileL0C = GetTile(tensorL0C, tla::MakeCoord(mL0Idx * L0_TILE_M, nL0Idx * L0_TILE_N), + tla::MakeShape(mL0Actual, nL0Actual)); + + // Compute the matrix multiplication on L0A and L0B and write the result to the accumulator + // Wait for loading L0B + AscendC::WaitFlag(l0CEventList[l0CListId]); + + // If the unit flag is enabled, the unit flag is set according to the calculation progress + uint8_t unitFlag = 0b00; + if constexpr (ENABLE_UNIT_FLAG) { + if ((kL1Idx == kL1Loop - 1) && (mL0Idx == mL0Loop - 1) && (kL0Idx == kL0Loop - 1) && + (nL0Idx == nL0Loop - 1)) { + unitFlag = 0b11; + } else { + unitFlag = 0b10; + } + } + + tileMmad(tensorTileL0C, tensorL0A, tensorL0B, mL0Actual, nL0Actual, kL0Actual, initC, unitFlag); + + // Notify to move the next L0B tile + AscendC::SetFlag(l0BEventList[l0BListId]); + l0BListId = (l0BListId + 1 < L0B_STAGES) ? (l0BListId + 1) : 0; + } + AscendC::SetFlag(l0AEventList[l0AListId]); + l0AListId = (l0AListId + 1 < L0A_STAGES) ? (l0AListId + 1) : 0; + } + } + l1AListId = l1AListIdNext; + l1BListId = l1BListIdNext; + kL1Actual = kL1ActualNext; + } + + // copy block out + if constexpr (!ENABLE_UNIT_FLAG) { + AscendC::SetFlag(l0CEventList[l0CListId]); + AscendC::WaitFlag(l0CEventList[l0CListId]); + copyL0CToDst(tensorC, tensorL0C); + AscendC::SetFlag(l0CEventList[l0CListId]); + l0CListId = (l0CListId + 1 < L0C_STAGES) ? (l0CListId + 1) : 0; + } else { + copyL0CToDst(tensorC, tensorL0C, 0b11); + } + } + +protected: + template + CATLASS_DEVICE auto GetTileA(TensorA &tensorA, uint32_t mIndex, uint32_t kIndex, uint32_t mSize, uint32_t kSize) + { + if constexpr (tla::detail::isVector::value) { + return GetTile(tensorA, tla::MakeCoord(kIndex), tla::MakeShape(kSize)); + } else { + return GetTile(tensorA, tla::MakeCoord(mIndex, kIndex), tla::MakeShape(mSize, kSize)); + } + } + + // Multi-stage tensors list + AscendC::LocalTensor l1ATensorList[L1A_STAGES]; + AscendC::LocalTensor l1BTensorList[L1B_STAGES]; + AscendC::LocalTensor l0ATensorList[L0A_STAGES]; + AscendC::LocalTensor l0BTensorList[L0B_STAGES]; + AscendC::LocalTensor l0CTensorList[L0C_STAGES]; + AscendC::LocalTensor l1BiasTensor; + AscendC::LocalTensor l0BiasTensor; + + // Multi-stage event id list + int32_t l1AEventList[L1A_STAGES]; + int32_t l1BEventList[L1B_STAGES]; + int32_t l0AEventList[L0A_STAGES]; + int32_t l0BEventList[L0B_STAGES]; + int32_t l0CEventList[L0C_STAGES]; + + __gm__ typename AscendC::GlobalTensor::PrimType *lastAddrA[L1A_STAGES]; + __gm__ typename AscendC::GlobalTensor::PrimType *lastAddrB[L1B_STAGES]; + MatrixCoord lastCoordA[L1A_STAGES]; + MatrixCoord lastCoordB[L1B_STAGES]; + + // The id of current stage + uint32_t l1AListId{0}; + uint32_t l1BListId{0}; + uint32_t l0AListId{0}; + uint32_t l0BListId{0}; + uint32_t l0CListId{0}; + + TileMmad tileMmad; + CopyL1ToL0A copyL1ToL0A; + CopyL1ToL0B copyL1ToL0B; + CopyL1ToBT copyL1ToBT; +}; + +} // namespace Catlass::Gemm::Block + +#endif // CATLASS_GEMM_BLOCK_BLOCK_MMAD_PINGPONG_TLA_PRELOADA_L1B_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/kernel_utils/tile/copy_l0c_to_ub.hpp b/csrc/chunk_kda_fwd/op_kernel/kernel_utils/tile/copy_l0c_to_ub.hpp new file mode 100644 index 0000000000..516e290c9d --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/kernel_utils/tile/copy_l0c_to_ub.hpp @@ -0,0 +1,381 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef COMMON_TILE_ASCEND950_COPY_L0C_TO_UB_950_HPP +#define COMMON_TILE_ASCEND950_COPY_L0C_TO_UB_950_HPP + +#include "catlass/arch/kda_arch.hpp" +#include "catlass/kda_catlass.hpp" +#if defined(CATLASS_ARCH) && CATLASS_ARCH == 3510 +#include "catlass/gemm/tile/ascend950/kda_ascend950_copy_l0c_to_dst.hpp" +#endif +#include "tla/kda_tensor.hpp" +#include "catlass/detail/kda_tag_to_layout.hpp" + +namespace Common::Tile { + +#if defined(CATLASS_ARCH) && CATLASS_ARCH == 3510 + +template +struct CopyL0CToUBTla { + static_assert(DEPENDENT_FALSE, "Unsupported copy l0c to ub, can not find the specialization."); +}; + +#endif + +#if defined(CATLASS_ARCH) && CATLASS_ARCH == 3510 + +template < + /// Tag indicating architecture + class ArchTag, class ElementA_, class LayoutTagA_, class ElementB_, class LayoutTagB_, class ElementC_, + class LayoutTagC_, class ElementBias = void, bool ReluEnable_ = false, + Catlass::Gemm::Tile::ScaleGranularity DEQUANT_GRANULARITY_ = Catlass::Gemm::Tile::ScaleGranularity::NO_QUANT, + class L0CCopyMode = Catlass::Gemm::Tile::CopyToGM> +struct PackedTileCopyTla { + using ElementA = ElementA_; + using ElementB = ElementB_; + using LayoutTagA = LayoutTagA_; + using LayoutTagB = LayoutTagB_; + using LayoutTagC = LayoutTagC_; + using ElementAccumulator = + typename Catlass::Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; + static constexpr bool ReluEnable = ReluEnable_; + static constexpr Catlass::Gemm::Tile::ScaleGranularity DEQUANT_GRANULARITY = DEQUANT_GRANULARITY_; + + static constexpr bool HAS_BIAS = !std::is_void_v; + static constexpr bool HAS_QUANT_TENSOR = + (DEQUANT_GRANULARITY == Catlass::Gemm::Tile::ScaleGranularity::PER_CHANNEL); + + using LayoutTagL1A = + typename Catlass::Gemm::helper::L1ATypeSelector>::L1AType::Layout; + using LayoutTagL1B = + typename Catlass::Gemm::helper::L1BTypeSelector>::L1BType::Layout; + using LayoutTagL0A = typename Catlass::Gemm::helper::L0ALayoutSelector::Layout; + using LayoutTagL0B = Catlass::layout::nZ; + using LayoutTagL0C = Catlass::layout::L0C; + + using LayoutA = Catlass::detail::TagToLayout_t; + using LayoutB = Catlass::detail::TagToLayout_t; + using LayoutC = Catlass::detail::TagToLayout_t; + + using LayoutL1A = Catlass::detail::TagToLayout_t; + using LayoutL1B = Catlass::detail::TagToLayout_t; + using LayoutL0A = Catlass::detail::TagToLayout_t; + using LayoutL0B = Catlass::detail::TagToLayout_t; + using LayoutL0C = typename Catlass::detail::LayoutL0C; + + using TensorL1AVectorLayout = + tla::Tensor, LayoutL1A, tla::Coord, AscendC::TPosition::A1>; + using TensorL1ALayout = + tla::Tensor, LayoutL1A, tla::Coord, AscendC::TPosition::A1>; + + using TensorL1A = + std::conditional_t::value, TensorL1AVectorLayout, TensorL1ALayout>; + using TensorL1B = + tla::Tensor, LayoutL1B, tla::Coord, AscendC::TPosition::A1>; + using TensorL0A = + tla::Tensor, LayoutL0A, tla::Coord, AscendC::TPosition::A2>; + using TensorL0B = + tla::Tensor, LayoutL0B, tla::Coord, AscendC::TPosition::B2>; + using TensorL0C = tla::Tensor, LayoutL0C, tla::Coord, + AscendC::TPosition::CO1>; + using TensorL1Bias = + std::conditional_t, + Catlass::detail::TagToLayout_t, + tla::Coord, AscendC::TPosition::A1>, + Catlass::EmptyClass>; + using TensorL0Bias = tla::Tensor, + Catlass::detail::TagToLayout_t, + tla::Coord, AscendC::TPosition::C2>; + using TensorL1Quant = + std::conditional_t, + Catlass::detail::TagToLayout_t, + tla::Coord, AscendC::TPosition::A1>, + Catlass::EmptyClass>; + + using L1AAlignHelper = Catlass::Gemm::helper::L1AlignHelper; + using L1BAlignHelper = Catlass::Gemm::helper::L1AlignHelper; + + template + using CopyGmToL1A = Catlass::Gemm::Tile::TileCopyTla; + + template + using CopyGmToL1B = Catlass::Gemm::Tile::TileCopyTla; + + template + using CopyGmToL1Bias = + std::conditional_t, + Catlass::EmptyClass>; + + template + using CopyGmToL1Scale = + std::conditional_t, + Catlass::EmptyClass>; + + using CopyL1ToL0A = Catlass::Gemm::Tile::TileCopyTla; + using CopyL1ToL0B = Catlass::Gemm::Tile::TileCopyTla; + using CopyL1ToBT = + std::conditional_t, + Catlass::EmptyClass>; + +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 3510) + template + using CopyL0CToDst = + Catlass::Gemm::Tile::CopyL0CToGmTla; +#endif +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 2201) + template + using CopyL0CToGm = + Catlass::Gemm::Tile::CopyL0CToGmTla; +#endif +}; + +#if defined(CATLASS_ARCH) && CATLASS_ARCH == 3510 + +template +struct CopyL0CToUBTla< + Catlass::Arch::Ascend950, TensorSrc_, + tla::Tensor, LayoutDst_, CoordDst_, AscendC::TPosition::VECCALC>, + Catlass::Gemm::Tile::CopyL0CToUBMode::NO_SPLIT, Catlass::Gemm::Tile::ScaleGranularity::NO_QUANT, ReluEnable_, + std::enable_if_t::value>> { + using ArchTag = Catlass::Arch::Ascend950; + using ElementDst = ElementDst_; + using ElementSrc = typename TensorSrc_::Element; + static constexpr auto quantPre = + Catlass::Gemm::Tile::CopyL0CToDstQuantMode::VALUE; + static constexpr auto reluEn = ReluEnable_; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, uint8_t unitFlag = 0) + { + static_assert( + tla::detail::isRowMajor::value && + TensorSrc::position == AscendC::TPosition::CO1 && TensorDst::position == AscendC::TPosition::VECCALC, + "The input parameters do not match. TensorSrc must be L0C, while TensorDst must be UB and RowMajor"); + + AscendC::FixpipeParamsC310 intriParams; + + // Fixpipe layout information + intriParams.nSize = tla::get<1>(dstTensor.originShape()); + intriParams.mSize = tla::get<0>(dstTensor.originShape()); + intriParams.srcStride = tla::get<1, 1>(srcTensor.stride()) / tla::get<0, 0>(srcTensor.stride()); + intriParams.dstStride = tla::get<0>(dstTensor.stride()); + + // Fixpipe auxiliary arguments + intriParams.quantPre = quantPre; + intriParams.reluEn = reluEn; + intriParams.unitFlag = unitFlag; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + + // Call AscendC Fixpipe + AscendC::Fixpipe(dstTensor.data()[dstOffset], + srcTensor.data()[srcOffset], intriParams); + } + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, uint32_t ubRowNum, + uint8_t beginSubBlockIdx, uint8_t sendVecNum, uint8_t unitFlag = 0) + { + static_assert( + tla::detail::isRowMajor::value && + TensorSrc::position == AscendC::TPosition::CO1 && TensorDst::position == AscendC::TPosition::VECCALC, + "The input parameters do not match. TensorSrc must be L0C, while TensorDst must be UB and RowMajor"); + + AscendC::FixpipeParamsC310 intriParams; + + uint32_t dstTensorRowNum = tla::get<0>(dstTensor.originShape()); + // Fixpipe layout information + intriParams.nSize = tla::get<1>(dstTensor.originShape()); + intriParams.mSize = min(dstTensorRowNum, ubRowNum); + intriParams.srcStride = tla::get<1, 1>(srcTensor.stride()) / tla::get<0, 0>(srcTensor.stride()); + intriParams.dstStride = tla::get<0>(dstTensor.stride()); + + // Fixpipe auxiliary arguments + intriParams.quantPre = quantPre; + intriParams.reluEn = reluEn; + intriParams.unitFlag = unitFlag; + intriParams.dualDstCtl = 0; + intriParams.subBlockId = beginSubBlockIdx; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + + // Call AscendC Fixpipe + AscendC::Fixpipe(dstTensor.data()[dstOffset], + srcTensor.data()[srcOffset], intriParams); + if (sendVecNum > 1) { + if (dstTensorRowNum > ubRowNum) { + intriParams.mSize = dstTensorRowNum - ubRowNum; + intriParams.subBlockId = (beginSubBlockIdx + 1) % 2; + srcOffset += ubRowNum * tla::get<0, 0>(srcTensor.stride()); + AscendC::Fixpipe(dstTensor.data()[dstOffset], + srcTensor.data()[srcOffset], intriParams); + } + } + } + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, uint8_t subBlockIdx, + uint8_t unitFlag = 0) + { + static_assert( + tla::detail::isRowMajor::value && + TensorSrc::position == AscendC::TPosition::CO1 && TensorDst::position == AscendC::TPosition::VECCALC, + "The input parameters do not match. TensorSrc must be L0C, while TensorDst must be UB and RowMajor"); + + AscendC::FixpipeParamsC310 intriParams; + + // Fixpipe layout information + intriParams.nSize = tla::get<1>(dstTensor.shape()); + intriParams.mSize = tla::get<0>(dstTensor.shape()); + intriParams.srcStride = tla::get<1, 1>(srcTensor.stride()) / tla::get<0, 0>(srcTensor.stride()); + intriParams.dstStride = tla::get<0>(dstTensor.stride()); + + // Fixpipe auxiliary arguments + intriParams.quantPre = quantPre; + intriParams.reluEn = reluEn; + intriParams.unitFlag = unitFlag; + intriParams.dualDstCtl = 0; + intriParams.subBlockId = subBlockIdx; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + + // Call AscendC Fixpipe + AscendC::Fixpipe(dstTensor.data()[dstOffset], + srcTensor.data()[srcOffset], intriParams); + } +}; + +template +struct CopyL0CToUBTla< + Catlass::Arch::Ascend950, TensorSrc_, + tla::Tensor, LayoutDst_, CoordDst_, AscendC::TPosition::VECCALC>, + Catlass::Gemm::Tile::CopyL0CToUBMode::SPLIT_M, Catlass::Gemm::Tile::ScaleGranularity::NO_QUANT, ReluEnable_, + std::enable_if_t::value>> { + using ArchTag = Catlass::Arch::Ascend950; + using ElementDst = ElementDst_; + using ElementSrc = typename TensorSrc_::Element; + static constexpr auto quantPre = + Catlass::Gemm::Tile::CopyL0CToDstQuantMode::VALUE; + static constexpr auto reluEn = ReluEnable_; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, uint8_t unitFlag = 0) + { + static_assert( + tla::detail::isRowMajor::value && + TensorSrc::position == AscendC::TPosition::CO1 && TensorDst::position == AscendC::TPosition::VECCALC, + "The input parameters do not match. TensorSrc must be L0C, while TensorDst must be UB and RowMajor"); + + AscendC::FixpipeParamsC310 intriParams; + + // Fixpipe layout information + intriParams.nSize = tla::get<1>(dstTensor.originShape()); + intriParams.mSize = RoundUp(tla::get<0>(dstTensor.originShape()), 2); // m must be even when spilt m + intriParams.srcStride = tla::get<1, 1>(srcTensor.stride()) / tla::get<0, 0>(srcTensor.stride()); + intriParams.dstStride = tla::get<0>(dstTensor.stride()); + + // Fixpipe auxiliary arguments + intriParams.quantPre = quantPre; + intriParams.reluEn = reluEn; + intriParams.unitFlag = unitFlag; + intriParams.dualDstCtl = 1; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + + // Call AscendC Fixpipe + AscendC::Fixpipe(dstTensor.data()[dstOffset], + srcTensor.data()[srcOffset], intriParams); + } +}; + +template +struct CopyL0CToUBTla< + Catlass::Arch::Ascend950, TensorSrc_, + tla::Tensor, LayoutDst_, CoordDst_, AscendC::TPosition::VECCALC>, + Catlass::Gemm::Tile::CopyL0CToUBMode::SPLIT_N, Catlass::Gemm::Tile::ScaleGranularity::NO_QUANT, ReluEnable_, + std::enable_if_t::value>> { + using ArchTag = Catlass::Arch::Ascend950; + using ElementDst = ElementDst_; + using ElementSrc = typename TensorSrc_::Element; + static constexpr auto quantPre = + Catlass::Gemm::Tile::CopyL0CToDstQuantMode::VALUE; + static constexpr auto reluEn = ReluEnable_; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, uint8_t unitFlag = 0) + { + static_assert( + tla::detail::isRowMajor::value && + TensorSrc::position == AscendC::TPosition::CO1 && TensorDst::position == AscendC::TPosition::VECCALC, + "The input parameters do not match. TensorSrc must be L0C, while TensorDst must be UB and RowMajor"); + + AscendC::FixpipeParamsC310 intriParams; + + // Fixpipe layout information + intriParams.nSize = RoundUp(tla::get<1>(dstTensor.originShape()), 32); + intriParams.mSize = tla::get<0>(dstTensor.originShape()); // m must be even when spilt m + intriParams.srcStride = tla::get<1, 1>(srcTensor.stride()) / tla::get<0, 0>(srcTensor.stride()); + intriParams.dstStride = tla::get<0>(dstTensor.stride()); + + // Fixpipe auxiliary arguments + intriParams.quantPre = quantPre; + intriParams.reluEn = reluEn; + intriParams.unitFlag = unitFlag; + intriParams.dualDstCtl = 2; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + + // Call AscendC Fixpipe + AscendC::Fixpipe(dstTensor.data()[dstOffset], + srcTensor.data()[srcOffset], intriParams); + } +}; + +#endif + +template < + /// Tag indicating architecture + class ArchTag, class ElementA_, class LayoutTagA, class ElementB_, class LayoutTagB, class ElementC_, + class LayoutTagC, class ElementBias = void, + Catlass::Gemm::Tile::CopyL0CToUBMode CopyMode_ = Catlass::Gemm::Tile::CopyL0CToUBMode::NO_SPLIT, + bool ReluEnable = false, + Catlass::Gemm::Tile::ScaleGranularity DEQUANT_GRANULARITY = Catlass::Gemm::Tile::ScaleGranularity::NO_QUANT> +struct PackedTileCopyTlaToUB : public PackedTileCopyTla { + static constexpr Catlass::Gemm::Tile::CopyL0CToUBMode CopyMode = CopyMode_; + // 重写 CopyL0CToDst + using TensorL0C = typename PackedTileCopyTla::TensorL0C; + + template + using CopyL0CToDst = CopyL0CToUBTla; +}; + +#endif + +///////////////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace Common::Tile + +#endif // CATLASS_GEMM_TILE_ASCEND950_COPY_L0C_TO_UB_950_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/kernel_utils/vector/regbase.hpp b/csrc/chunk_kda_fwd/op_kernel/kernel_utils/vector/regbase.hpp new file mode 100644 index 0000000000..0bd2391742 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/kernel_utils/vector/regbase.hpp @@ -0,0 +1,150 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +using namespace AscendC::MicroAPI; + +#pragma once + +constexpr static CastTrait ctHalf2Fp32Zero = { + RegLayout::ZERO, + SatMode::SAT, + MaskMergeMode::ZEROING, + AscendC::RoundMode::CAST_NONE, +}; +constexpr static CastTrait ctHalf2Fp32One = { + RegLayout::ONE, + SatMode::SAT, + MaskMergeMode::ZEROING, + AscendC::RoundMode::CAST_NONE, +}; +constexpr static CastTrait ctFp322HalfZero = {RegLayout::ZERO, SatMode::NO_SAT, MaskMergeMode::MERGING, + AscendC::RoundMode::CAST_ROUND}; +constexpr static CastTrait ctFp322HalfOne = {RegLayout::ONE, SatMode::NO_SAT, MaskMergeMode::ZEROING, + AscendC::RoundMode::CAST_ROUND}; + +constexpr uint16_t PRELOAD_NUM = 2; + +template +__simd_callee__ inline void LoadIn(RegTensor &dstReg, __ubuf__ TType *srcUb) +{ + if constexpr (BroatCast) { + if constexpr (!std::is_same()) { + LoadAlign(dstReg, srcUb); + } else { + LoadAlign(dstReg, srcUb); + } + } else { + LoadAlign(dstReg, srcUb); + } +} + +template +__simd_callee__ inline void CastHalf2Float(RegTensor &dstZeroReg, RegTensor &dstOneReg, + RegTensor &srcReg, MaskReg &mask) +{ + Cast(dstZeroReg, srcReg, mask); + Cast(dstOneReg, srcReg, mask); +} + +template +__simd_callee__ inline void CastFloat2Half(RegTensor &dstReg, RegTensor &srcZeroReg, + RegTensor &srcOneReg, MaskReg &mask) +{ + Cast(dstReg, srcOneReg, mask); + Cast(dstReg, srcZeroReg, mask); +} + +template +__simd_callee__ inline void HalfOrFloat2Float(RegTensor &dstReg, RegTensor &srcReg, MaskReg &maskHalf, + MaskReg &maskFloat) +{ + if constexpr (!std::is_same()) { + Cast(dstReg, srcReg, maskHalf); + } else { + Duplicate(dstReg, srcReg, maskFloat); + } +} + +__simd_callee__ inline void MulFloatTwoReg(RegTensor &dstZeroReg, RegTensor &dstOneReg, + RegTensor &src1ZeroReg, RegTensor &src1OneReg, + RegTensor &src2ZeroReg, RegTensor &src2OneReg, + MaskReg &maskFloat) +{ + Mul(dstZeroReg, src1ZeroReg, src2ZeroReg, maskFloat); + Mul(dstOneReg, src1OneReg, src2OneReg, maskFloat); +} + +__simd_callee__ inline void MinsFloatTwoReg(RegTensor &dstZeroReg, RegTensor &dstOneReg, + RegTensor &srcZeroReg, RegTensor &srcOneReg, + float scalarValue, MaskReg &maskFloat) +{ + Mins(dstZeroReg, srcZeroReg, scalarValue, maskFloat); + Mins(dstOneReg, srcOneReg, scalarValue, maskFloat); +} + +__simd_callee__ inline void ExpFloatTwoReg(RegTensor &dstZeroReg, RegTensor &dstOneReg, + RegTensor &srcZeroReg, RegTensor &srcOneReg, + MaskReg &maskFloat) +{ + Exp(dstZeroReg, srcZeroReg, maskFloat); + Exp(dstOneReg, srcOneReg, maskFloat); +} + +__simd_callee__ inline void SubFloatTwoReg(RegTensor &dstZeroReg, RegTensor &dstOneReg, + RegTensor &src1ZeroReg, RegTensor &src1OneReg, + RegTensor &src2ZeroReg, RegTensor &src2OneReg, + MaskReg &maskFloat) +{ + Sub(dstZeroReg, src1ZeroReg, src2ZeroReg, maskFloat); + Sub(dstOneReg, src1OneReg, src2OneReg, maskFloat); +} + +__simd_callee__ inline void AddFloatTwoReg(RegTensor &dstZeroReg, RegTensor &dstOneReg, + RegTensor &src1ZeroReg, RegTensor &src1OneReg, + RegTensor &src2ZeroReg, RegTensor &src2OneReg, + MaskReg &maskFloat) +{ + Add(dstZeroReg, src1ZeroReg, src2ZeroReg, maskFloat); + Add(dstOneReg, src1OneReg, src2OneReg, maskFloat); +} + +template +__simd_callee__ inline void CompareTwoReg(MaskReg &maskZeroReg, MaskReg &maskOneReg, RegTensor &cmp1ZeroReg, + RegTensor &cmp1OneReg, RegTensor &cmp2ZeroReg, + RegTensor &cmp2OneReg, MaskReg &mask) +{ + Compare(maskZeroReg, cmp1ZeroReg, cmp2ZeroReg, mask); + Compare(maskOneReg, cmp1OneReg, cmp2OneReg, mask); +} + +template +__simd_callee__ inline void SelectTwoReg(RegTensor &dstZeroReg, RegTensor &dstOneReg, + RegTensor &src1ZeroReg, RegTensor &src1OneReg, + RegTensor &src2ZeroReg, RegTensor &src2OneReg, MaskReg &mask1, + MaskReg &mask2) +{ + Select(dstZeroReg, src1ZeroReg, src2ZeroReg, mask1); + Select(dstOneReg, src1OneReg, src2OneReg, mask2); +} + +template +__simd_callee__ inline void StoreUnAlignOut(__ubuf__ TType *dstUb, RegTensor &srcReg, MaskReg &mask, + UnalignRegForStore &uStore, uint32_t copyNum) +{ + RegTensor outZeroReg; + if constexpr (!std::is_same()) { + Cast(outZeroReg, srcReg, mask); + StoreUnAlign(dstUb, outZeroReg, uStore, copyNum); + StoreUnAlignPost(dstUb, uStore, 0); + } else { + StoreUnAlign(dstUb, srcReg, uStore, copyNum); + StoreUnAlignPost(dstUb, uStore, 0); + } +} diff --git a/csrc/chunk_kda_fwd/op_kernel/tla/int_tuple.hpp b/csrc/chunk_kda_fwd/op_kernel/tla/int_tuple.hpp new file mode 100644 index 0000000000..b35081452c --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/tla/int_tuple.hpp @@ -0,0 +1,239 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef TLA_INT_TUPLE_HPP +#define TLA_INT_TUPLE_HPP + +#include "tla/type_traits.hpp" +#include "tla/tuple.hpp" +#include "tla/numeric/integral_constant.hpp" +#include "tla/numeric/integer_sequence.hpp" + +namespace tla { +// +// Apply (Unpack) +// (t, f) => f(t_0,t_1,...,t_n) +// + +namespace detail { + +template +CATLASS_HOST_DEVICE constexpr auto apply(T &&t, F &&f, seq) +{ + return f(get(static_cast(t))...); +} + +template +CATLASS_HOST_DEVICE constexpr auto tapply(T &&t, F &&f, G &&g, seq) +{ + return g(f(get(static_cast(t)))...); +} + +template +CATLASS_HOST_DEVICE constexpr auto tapply(T0 &&t0, T1 &&t1, F &&f, G &&g, seq) +{ + return g(f(get(static_cast(t0)), + get(static_cast(t1)))...); +} + +} // end namespace detail + +template +CATLASS_HOST_DEVICE constexpr auto apply(T &&t, F &&f) +{ + return detail::apply(static_cast(t), f, tuple_seq{}); +} + +template +CATLASS_HOST_DEVICE constexpr auto transform_apply(T &&t, F &&f, G &&g) +{ + if constexpr (is_tuple>::value) { + return detail::tapply(static_cast(t), f, g, tuple_seq{}); + } else { + return g(f(static_cast(t))); + } +} + +struct UnpackedMakeTuple { + template + CATLASS_HOST_DEVICE constexpr auto operator()(T const &...a) const + { + return tla::MakeTuple(a...); + } +}; + +template +CATLASS_HOST_DEVICE constexpr auto transform(T0 const &t0, T1 const &t1, F &&f) +{ + if constexpr (is_tuple::value) { + static_assert(tuple_size::value == tuple_size::value, "Mismatched tuple_size"); + return detail::tapply(t0, t1, f, UnpackedMakeTuple{}, tuple_seq{}); + } else { + return f(t0, t1); + } +} + +template >::value)> +CATLASS_HOST_DEVICE constexpr decltype(auto) get(T &&t) noexcept +{ + static_assert(I == 0, "Index out of range"); + return static_cast(t); +} + +template +CATLASS_HOST_DEVICE constexpr decltype(auto) get(T &&t) noexcept +{ + return get(get(static_cast(t))); +} + +// max +template +CATLASS_HOST_DEVICE constexpr auto max(T0 const &t0, Ts const &...ts); + +struct UnpackedMax { + template + CATLASS_HOST_DEVICE constexpr auto operator()(T const &...v) const + { + return tla::max(v...); + } +}; + +template +CATLASS_HOST_DEVICE constexpr auto max(T0 const &t0, Ts const &...ts) +{ + if constexpr (is_tuple::value) { + return tla::max(tla::apply(t0, UnpackedMax{}), ts...); + } else if constexpr (sizeof...(Ts) == 0) { + return t0; + } else { + return tla::max(t0, tla::max(ts...)); + } +} + +// rank +template +CATLASS_HOST_DEVICE constexpr auto rank(Tuple const &t) +{ + if constexpr (sizeof...(Is) == 0) { + if constexpr (is_tuple::value) { + return Int::value>{}; + } else { + return Int<1>{}; + } + } else { + return rank(get(t)); + } +} + +template +using rank_t = decltype(rank(std::declval())); + +template +constexpr auto rank_v = rank_t::value; + +// depth +template +CATLASS_HOST_DEVICE constexpr auto depth(Tuple const &t); + +struct UnpackedDepth { + template + CATLASS_HOST_DEVICE constexpr auto operator()(T const &...v) const + { + return tla::max(depth(v)...); + } +}; + +template +CATLASS_HOST_DEVICE constexpr auto depth(Tuple const &t) +{ + if constexpr (sizeof...(Is) == 0) { + if constexpr (is_tuple::value) { + return Int<1>{} + tla::apply(t, UnpackedDepth{}); + } else { + return Int<0>{}; + } + } else { + return depth(get(t)); + } +} + +template +using depth_t = decltype(depth(std::declval())); + +template +constexpr auto depth_v = depth_t::value; + +struct MultipliesUnaryLfold { + template + CATLASS_HOST_DEVICE constexpr auto operator()(T const &...v) const + { + return (... * v); + } +}; + +// Implementation of product as a function object +struct Product { + template + CATLASS_HOST_DEVICE constexpr auto operator()(IntTuple const &a) const + { + if constexpr (is_tuple::value) { + if constexpr (tuple_size::value == 0) { + return Int<1>{}; + } else { + return tla::transform_apply(a, Product{}, MultipliesUnaryLfold{}); + } + } else if constexpr (tla::is_integral::value) { + return a; + } + } +}; + +namespace detail { + +template +struct MakeZeroTupleImpl; + +template +struct MakeZeroTupleImpl> { + using type = tla::tuple...>; +}; + +template +using MakeZeroTuple = typename MakeZeroTupleImpl>::type; + +} // end namespace detail + +// Add +template +CATLASS_HOST_DEVICE constexpr auto Add(IntTupleA const &a, IntTupleB const &b); + +struct UnpackedAdd { + template + CATLASS_HOST_DEVICE constexpr auto operator()(IntTupleA const &x, IntTupleB const &y) const + { + return Add(x, y); + } +}; + +template +CATLASS_HOST_DEVICE constexpr auto Add(IntTupleA const &a, IntTupleB const &b) +{ + if constexpr (is_tuple::value && is_tuple::value) { + static_assert(tuple_size::value == tuple_size::value, "Mismatched ranks"); + return transform(a, b, UnpackedAdd{}); + } else { + return tla::add(a, b); + } +} + +} // end namespace tla + +#endif // TLA_INT_TUPLE_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/tla/kda_tensor.hpp b/csrc/chunk_kda_fwd/op_kernel/tla/kda_tensor.hpp new file mode 100644 index 0000000000..3214d1da02 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/tla/kda_tensor.hpp @@ -0,0 +1,413 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef TLA_KDA_TENSOR_HPP +#define TLA_KDA_TENSOR_HPP + +#include "catlass/arch/kda_arch.hpp" +#include "tla/kda_tla_layout.hpp" // tla::Shape +#include "tla/numeric/integral_constant.hpp" // tla::is_integral +#include "tla/int_tuple.hpp" + +namespace tla { + +// +// Underscore slicing utilities +// +namespace detail { + +// Safe "element type at index I" helper that avoids instantiating get for invalid indices. +template +struct coord_elem_type { + using type = void; +}; +template +struct coord_elem_type= 0) && (I < (int)tla::tuple_size>::value)>> { + using type = tla::remove_cvref_t(std::declval()))>; +}; + +template +struct coord_elem_is_underscore : tla::is_underscore::type> {}; + +// Count underscores. +template +struct underscore_count_from : tla::integral_constant {}; +template +struct underscore_count_from= 0)>> { + static constexpr int value = + (coord_elem_is_underscore::value ? 1 : 0) + underscore_count_from::value; +}; +template +struct underscore_count + : tla::integral_constant>::value - 1>::value> {}; + +// Build index sequences for underscore dims (stable 0..R-1 recursion). +template +struct underscore_indices_impl; +template +struct underscore_indices_impl { + using type = seq; +}; +template +struct underscore_indices_impl + : std::conditional_t::value, + underscore_indices_impl, + underscore_indices_impl> {}; + +template +using underscore_indices = + typename underscore_indices_impl>::value>::type; + +// Replace every tla::_ with 0 for offset computation. +template +CATLASS_HOST_DEVICE constexpr decltype(auto) underscore_to_zero(T const &x) +{ + if constexpr (tla::is_underscore::value) { + return tla::_0{}; + } else { + return x; + } +} + +template +CATLASS_HOST_DEVICE constexpr auto replace_underscore_with_zero_impl(Coord const &c, seq) +{ + return tla::MakeCoord(underscore_to_zero(tla::get(c))...); +} + +template +CATLASS_HOST_DEVICE constexpr auto replace_underscore_with_zero(Coord const &c) +{ + static_assert(tla::is_tuple>::value, + "Coord must be tla::tuple for underscore slicing."); + return replace_underscore_with_zero_impl(c, tuple_seq{}); +} + +// Build a layout from selected top-level indices of `layout`. +template +CATLASS_HOST_DEVICE constexpr auto select_layout(Layout const &layout, seq) +{ + auto shape_new = tla::MakeTuple(tla::get(layout.shape())...); + auto stride_new = tla::MakeTuple(tla::get(layout.stride())...); + auto origin_new = tla::MakeTuple(tla::get(layout.originShape())...); + return tla::MakeLayout(shape_new, stride_new, origin_new); +} + +} // namespace detail + +// +// slice_and_offset +// +// A lightweight helper that factors underscore slicing into: +// - a projected layout (keeping only underscored dimensions, in-order), and +// - a base offset computed at the fixed indices (underscored dims treated as 0). +// +// Notes: +// - `coord_arg` must be a one-level `tla::tuple` and contain at least one `tla::_`. +// - This function does not perform runtime bounds checks against originShape(); out-of-bounds is undefined behavior. +// - Returned `offset` is an element offset intended for `BuiltinTensor::operator[](offset)` view creation. +template +CATLASS_HOST_DEVICE constexpr auto slice_and_offset(CoordArg const &coord_arg, + Layout const &layout, + BaseCoord const &base_coord) +{ + static_assert(tla::is_tuple>::value, "slice_and_offset expects a tuple CoordArg."); + static_assert(depth_v == 1, "slice_and_offset only supports one-level CoordArg (no nested tuples)."); + static_assert((int)tla::tuple_size>::value == (int)Layout::rank, + "slice_and_offset requires CoordArg rank == Layout::rank."); + static_assert(tla::is_tuple>::value, "slice_and_offset expects a tuple BaseCoord."); + static_assert((int)tla::tuple_size>::value == (int)Layout::rank, + "slice_and_offset requires BaseCoord rank == Layout::rank."); + + constexpr int k = detail::underscore_count::value; + static_assert(k > 0, "slice_and_offset requires at least one underscore."); + static_assert(k <= Layout::rank, "Invalid underscore count."); + + // Compute base offset using zeros for underscores + auto coord0 = detail::replace_underscore_with_zero(coord_arg); + auto full0 = Add(base_coord, coord0); + auto offset = (int64_t)layout(full0); + + // Determine output dims (underscored dims) and build projected layout. + using Us = detail::underscore_indices; + auto layout_proj = detail::select_layout(layout, Us{}); + + return tla::MakeTuple(layout_proj, offset); +} + +// Convenience overload (no base_coord): assume base_coord == 0. +template +CATLASS_HOST_DEVICE constexpr auto slice_and_offset(CoordArg const &coord_arg, Layout const &layout) +{ + using Z = detail::MakeZeroTuple; + return slice_and_offset(coord_arg, layout, Z{}); +} + +// +// Tensor +// + +namespace detail { + +template +CATLASS_DEVICE constexpr auto HadamardU32(A const &a, B const &b, seq) +{ + return MakeCoord((static_cast(get(a)) * static_cast(get(b)))...); +} + +template +CATLASS_DEVICE constexpr auto GetTileImpl(TensorT const &tensor, CoordT const &coord, ShapeT const &shape, Int) +{ + static_assert(is_tuple::value && depth_v == 1 && rank_v == R, "Coord rank mismatch."); + static_assert(is_tuple::value && depth_v == 1 && rank_v == R, "Shape rank mismatch."); + + auto layoutNew = GetTileLayout(tensor.layout(), shape, coord); + auto coordNew = Add(tensor.coord(), coord); + return MakeTensor(tensor.data(), layoutNew, coordNew, Catlass::Arch::PositionType{}); +} + +template +CATLASS_DEVICE constexpr auto TileViewImpl(TensorT const &tensor, TileCoord const &tileCoord, TileShape const &tileShape, Int) +{ + static_assert(is_tuple::value && depth_v == 1 && rank_v == R, "TileCoord rank mismatch."); + static_assert(is_tuple::value && depth_v == 1 && rank_v == R, "TileShape rank mismatch."); + + auto elementOffset = HadamardU32(tileCoord, tileShape, tuple_seq{}); + auto layoutNew = GetTileLayout(tensor.layout(), tileShape, elementOffset); + auto coordNew = Add(tensor.coord(), elementOffset); + return MakeTensor(tensor.data(), layoutNew, coordNew, Catlass::Arch::PositionType{}); +} + +} // namespace detail + +template +struct Tensor { + using Element = typename BuiltinTensor::PrimType; + using Layout = Layout_; + using Coord = Coord_; + static constexpr AscendC::TPosition position = Position; + + CATLASS_HOST_DEVICE constexpr Tensor() {} + + CATLASS_HOST_DEVICE constexpr Tensor(BuiltinTensor const &builtinTensor, Layout const &layout, Coord const &coord = {}) + : rep_(builtinTensor, layout, coord) + {} + + // + // Accessors + // + + static constexpr int rank = Layout::rank; + + CATLASS_HOST_DEVICE constexpr decltype(auto) tensor() const + { + return *this; + } + + CATLASS_HOST_DEVICE constexpr decltype(auto) data() const + { + return get<0>(rep_); + } + + CATLASS_HOST_DEVICE constexpr decltype(auto) data() + { + return get<0>(rep_); + } + + CATLASS_HOST_DEVICE constexpr decltype(auto) layout() const + { + return get<1>(rep_); + } + + CATLASS_HOST_DEVICE constexpr decltype(auto) coord() const + { + return get<2>(rep_); + } + + CATLASS_HOST_DEVICE constexpr decltype(auto) shape() const + { + return layout().shape(); + } + + CATLASS_HOST_DEVICE constexpr decltype(auto) stride() const + { + return layout().stride(); + } + + CATLASS_HOST_DEVICE constexpr decltype(auto) originShape() const + { + return layout().originShape(); + } + + // + // Indexing / slicing + // + // - No underscore: returns `data()[layout()(coord()+coord_arg)]` + // - Underscores (0..rank, one-level coord): returns a subtensor view over the underscored dimensions (kept in-order) + // Notes: + // - Coord must be one-level (no nested tuples in coord elements). + // - Fixed (non-underscore) indices are expected to be within originShape(). This implementation does not + // generally perform runtime bounds checks/cropping; out-of-bounds indices result in undefined behavior. + template + CATLASS_HOST_DEVICE constexpr decltype(auto) operator()(CoordArg const &coord_arg) const + { + if constexpr (tla::is_tuple>::value) { + static_assert(depth_v == 1, "Underscore slicing only supports one-level Coord (no nested tuples)."); + static_assert(tla::tuple_size>::value == Layout::rank, + "Tensor::operator()(coord): Coord rank must equal tensor rank (Layout::rank)."); + + constexpr int k = detail::underscore_count::value; + if constexpr (k > 0) { + static_assert(k <= Layout::rank, "Invalid underscore count."); + + auto sliced = tla::slice_and_offset(coord_arg, layout(), coord()); + auto layout_proj = tla::get<0>(sliced); + auto offset = (int64_t)tla::get<1>(sliced); + + using CoordZ = detail::MakeZeroTuple<(size_t)k>; + auto data_new = data()[static_cast(offset)]; + return Tensor(data_new, layout_proj, CoordZ{}); + } else { + // No underscore: point view at coord() + coord_arg + auto full = Add(coord(), coord_arg); + return data()[layout()(full)]; + } + } else { + // Scalar coordinate convenience (rank-1): treat it as a 1D coord tuple. + static_assert(Layout::rank == 1, "Tensor::operator()(scalar) is only supported for rank-1 tensors."); + auto full = Add(coord(), MakeCoord(coord_arg)); + return data()[layout()(full)]; + } + } + + template + CATLASS_HOST_DEVICE constexpr decltype(auto) operator()(Coord0 const &c0, Coord1 const &c1, Coords const &...cs) const + { + return operator()(MakeCoord(c0, c1, cs...)); + } + + tla::tuple rep_; +}; + +template +CATLASS_HOST_DEVICE constexpr auto MakeTensor(BuiltinTensor const &builtinTensor, Layout const &layout, PositionType) +{ + using Coord = detail::MakeZeroTuple; + return Tensor(builtinTensor, layout); +} + +template +CATLASS_HOST_DEVICE constexpr auto MakeTensor(BuiltinTensor const &builtinTensor, Layout const &layout, Coord const &coord, PositionType) +{ + return Tensor(builtinTensor, layout, coord); +} + +// Get a tile from tensor, automatically handling boundary cases (tail tile). +// coord is element coordinate; shape is tile size for memory layout. +// The returned Tensor's layout.shape() is used for memory layout calculation, +// and layout.originShape() is the actual logical size (may be smaller than shape). +// Supports tensors of any rank (rank >= 1). +template +CATLASS_DEVICE constexpr auto GetTile(Tensor const &tensor, Coord const &coord, Shape const &shape) +{ + static_assert(Tensor::rank >= 1, "GetTile requires tensor rank >= 1."); + static_assert(Tensor::rank == rank_v && Tensor::rank == rank_v, + "GetTile: coord and shape must have the same rank as the tensor."); + return detail::GetTileImpl(tensor, coord, shape, Int{}); +} + +// 从 tensor 中获取一个 tile，自动处理边界情况（tail tile）。 +// tileCoord 是 tile 单位坐标；tileShape 是用于内存布局的 tile 尺寸。 +// 返回的 Tensor 的 layout.shape() 用于内存布局计算，layout.originShape() 是实际逻辑尺寸（可能小于 tileShape）。 +// TileView(tensor, tileCoord, tileShape) = GetTile(tensor, tileCoord ⊙ tileShape, tileShape) +// Supports tensors of any rank (rank >= 1). +template +CATLASS_DEVICE constexpr auto TileView(TensorT const &tensor, TileCoord const &tileCoord, TileShape const &tileShape) +{ + static_assert(TensorT::rank >= 1, "TileView requires tensor rank >= 1."); + static_assert(TensorT::rank == rank_v && TensorT::rank == rank_v, + "TileView: tileCoord and tileShape must have the same rank as the tensor."); + return detail::TileViewImpl(tensor, tileCoord, tileShape, Int{}); +} + +// 创建一个与另一个 Tensor 类似的 Tensor： +// 目标 layout 根据 LayoutTagDst 构造，从 LikeTensor::Element 推断 ElementDst，从 likeTensor 的 originShape 提取尺寸。 +template +CATLASS_HOST_DEVICE constexpr auto MakeTensorLike(BuiltinTensor const &builtinTensor, + LikeTensor const &likeTensor, + PositionType) +{ + using ElementDst = typename LikeTensor::Element; + static_assert(std::is_same_v, + "BuiltinTensor element type must match LikeTensor element type"); + return MakeTensorLike(builtinTensor, likeTensor, PositionType{}); +} + +// 创建一个与另一个 Tensor 类似的 Tensor： +// 使用 layoutBase 的 shape/stride，但继承 likeTensor 的 originShape，从 LikeTensor::Element 推断 ElementDst。 +template +CATLASS_HOST_DEVICE constexpr auto MakeTensorLike(BuiltinTensor const &builtinTensor, + LikeTensor const &likeTensor, + PositionType, + LayoutBase const &layoutBase) +{ + using ElementDst = typename LikeTensor::Element; + static_assert(std::is_same_v, + "BuiltinTensor element type must match LikeTensor element type"); + return MakeTensorLike(builtinTensor, likeTensor, PositionType{}, layoutBase); +} + +// 创建一个与另一个 Tensor 类似的 Tensor： +// 目标 layout 根据 LayoutTagDst 构造，从 LikeTensor::Element 推断 ElementDst，从 likeTensor 的 originShape 提取尺寸。(调用MakeLayout，可能会因分型布局合法要求对shape进行以分型为粒度的向上取整) +// 允许 BuiltinTensor 的元素类型与 LikeTensor 的元素类型不同（例如 L0C 使用 ElementAccumulator 而不是 ElementC）。 +template +CATLASS_HOST_DEVICE constexpr auto MakeTensorLike(BuiltinTensor const &builtinTensor, + LikeTensor const &likeTensor, + PositionType) +{ + static_assert(LikeTensor::rank == 1 || LikeTensor::rank == 2, + "MakeTensorLike(..., likeTensor) expects rank-1 or rank-2 likeTensor."); + static_assert(std::is_same_v, + "BuiltinTensor element type must match specified ElementDst type"); + // 根据目标布局格式（LayoutTagDst）和指定的元素类型构造 layout + if constexpr (LikeTensor::rank == 1) { + auto layoutNominal = MakeLayout(get<0>(likeTensor.layout().originShape())); + using Coord0 = detail::MakeZeroTuple; + return Tensor(builtinTensor, layoutNominal); + } else { + static_assert(LikeTensor::rank == 2, "MakeTensorLike(..., likeTensor) expects rank-1 or rank-2 likeTensor."); + auto layoutNominal = MakeLayout(get<0>(likeTensor.layout().originShape()), get<1>(likeTensor.layout().originShape())); + using Coord0 = detail::MakeZeroTuple; + return Tensor(builtinTensor, layoutNominal); + } +} + +// 创建一个与另一个 Tensor 类似的 Tensor： +// 使用 layoutBase 的 shape/stride，但继承 likeTensor 的 originShape。允许 BuiltinTensor 的元素类型与 LikeTensor 的元素类型不同。 +template +CATLASS_HOST_DEVICE constexpr auto MakeTensorLike(BuiltinTensor const &builtinTensor, + LikeTensor const &likeTensor, + PositionType, + LayoutBase const &layoutBase) +{ + static_assert(LikeTensor::rank == 1 || LikeTensor::rank == 2, "MakeTensorLike(..., likeTensor, layoutBase) expects rank-1 or rank-2 likeTensor."); + static_assert(std::is_same_v, + "BuiltinTensor element type must match specified ElementDst type"); + + auto layoutFixedStride = MakeLayout(layoutBase.shape(), layoutBase.stride(), likeTensor.originShape()); + using Coord0 = detail::MakeZeroTuple; + return Tensor(builtinTensor, layoutFixedStride); +} + +} // end namespace tla + +#endif // TLA_KDA_TENSOR_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/tla/kda_tla_layout.hpp b/csrc/chunk_kda_fwd/op_kernel/tla/kda_tla_layout.hpp new file mode 100644 index 0000000000..c68c52bfe4 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/tla/kda_tla_layout.hpp @@ -0,0 +1,580 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef KDA_TLA_LAYOUT_HPP +#define KDA_TLA_LAYOUT_HPP + +#include "catlass/kda_catlass.hpp" +#include "tla/numeric/integral_constant.hpp" +#include "tla/numeric/math.hpp" +#include "tla/tuple.hpp" +#include "tla/int_tuple.hpp" +#include "catlass/layout/kda_layout.hpp" + +namespace tla { + +// Aliases + +template +using Shape = tla::tuple; + +template +using Stride = tla::tuple; + +template +using Coord = tla::tuple; + +template +CATLASS_HOST_DEVICE constexpr Shape MakeShape(Ts const &...t) +{ + return {t...}; +} +template +CATLASS_HOST_DEVICE constexpr Stride MakeStride(Ts const &...t) +{ + return {t...}; +} +template +CATLASS_HOST_DEVICE constexpr Coord MakeCoord(Ts const &...t) +{ + return {t...}; +} + +namespace detail { + +// Type trait to generate OriginShape type with rank=R, depth=1, element type=uint32_t +template +struct MakeOriginShapeTypeImpl; + +template +struct MakeOriginShapeTypeImpl> { + template + using repeat_type = uint32_t; + using type = Shape...>; +}; + +template +using MakeOriginShapeType = typename MakeOriginShapeTypeImpl, tla::make_index_sequence>>::type; + +struct UnpackedMakeOriginShapeU32 { + template + CATLASS_HOST_DEVICE constexpr auto operator()(T const &...a) const + { + return MakeShape(static_cast(a)...); + } +}; + +} // namespace detail + +// +// Layout +// +// 自动推导 OriginShape 类型为 Shape，rank=rank_v +template > +struct Layout : private tla::tuple { + // NOTE: This defaults static Shapes/Strides correctly, but not dynamic + CATLASS_HOST_DEVICE constexpr Layout(Shape const &shape = {}, Stride const &stride = {}, OriginShape const &originShape = {}) + : tla::tuple(shape, stride, originShape) + {} + + // + // Accessors + // + + static constexpr int rank = rank_v; + static constexpr int depth = depth_v; + + template + CATLASS_HOST_DEVICE constexpr decltype(auto) shape() + { + return get<0, I...>(static_cast &>(*this)); + } + + template + CATLASS_HOST_DEVICE constexpr decltype(auto) shape() const + { + return get<0, I...>(static_cast const &>(*this)); + } + + template + CATLASS_HOST_DEVICE constexpr decltype(auto) stride() + { + return get<1, I...>(static_cast &>(*this)); + } + + template + CATLASS_HOST_DEVICE constexpr decltype(auto) stride() const + { + return get<1, I...>(static_cast const &>(*this)); + } + + template + CATLASS_HOST_DEVICE constexpr decltype(auto) originShape() + { + return get<2, I...>(static_cast &>(*this)); + } + + template + CATLASS_HOST_DEVICE constexpr decltype(auto) originShape() const + { + return get<2, I...>(static_cast const &>(*this)); + } + + template + CATLASS_HOST_DEVICE constexpr auto operator()(Coord const &coord) const + { + return crd2offset(coord, shape(), stride()); + } +}; + +// Layout construction + +template +CATLASS_HOST_DEVICE constexpr auto MakeLayout(Shape const &shape, Stride const &stride, OriginShape const &originShape) +{ + static_assert(is_tuple::value || is_integral::value); + static_assert(is_tuple::value || is_integral::value); + static_assert(is_tuple::value || is_integral::value); + return Layout(shape, stride, originShape); +} + +template +CATLASS_HOST_DEVICE constexpr auto MakeLayout(Shape const &shape, Stride const &stride) +{ + static_assert(is_tuple::value || is_integral::value); + static_assert(is_tuple::value || is_integral::value); + // 计算默认的 originShape：将 shape 扁平化为 depth=1，并将每个维度归一化为 uint32_t + return MakeLayout(shape, stride, tla::transform_apply(shape, Product{}, detail::UnpackedMakeOriginShapeU32{})); +} + +// Convenience tags for common layouts + +template +CATLASS_HOST_DEVICE constexpr auto MakeLayoutFromTag(LayoutTag const &tag) +{ + static_assert(std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v, + "Unsupported LayoutTag for MakeLayoutFromTag, only support Catlass::layout::RowMajor or" + "Catlass::layout::ColumnMajor or Catlass::layout::VectorLayout or Catlass::layout::zN or Catlass::layout::nZ or Catlass::layout::L0C"); + + if constexpr (std::is_same_v) { + return MakeLayout(MakeShape(tag.shape(0)), + MakeStride(tag.stride(0)), + MakeShape(tag.shape(0))); + } else if constexpr (std::is_same_v) { + return MakeLayout(MakeShape(tag.shape(0), tag.shape(1)), + MakeStride(tag.stride(0), Int<1>{}), + MakeShape(tag.shape(0), tag.shape(1))); + } else if constexpr (std::is_same_v) { + return MakeLayout(MakeShape(tag.shape(0), tag.shape(1)), + MakeStride(Int<1>{}, tag.stride(1)), + MakeShape(tag.shape(0), tag.shape(1))); + } else { // zN or nZ or L0C + return MakeLayout(MakeShape(MakeShape(tag.shape(0), tag.shape(1)), MakeShape(tag.shape(2), tag.shape(3))), + MakeStride(MakeStride(tag.stride(0), tag.stride(1)), MakeStride(tag.stride(2), tag.stride(3))), + MakeShape(tag.orgShape(0), tag.orgShape(1))); + } +} + +// Return the shape of a mode +template +CATLASS_HOST_DEVICE constexpr decltype(auto) shape(Layout &layout) +{ + return layout.template shape(); +} + +template +CATLASS_HOST_DEVICE constexpr decltype(auto) shape(Layout const &layout) +{ + return layout.template shape(); +} + +// Return the stride of a mode +template +CATLASS_HOST_DEVICE constexpr decltype(auto) stride(Layout &layout) +{ + return layout.template stride(); +} + +template +CATLASS_HOST_DEVICE constexpr decltype(auto) stride(Layout const &layout) +{ + return layout.template stride(); +} + +template +CATLASS_HOST_DEVICE constexpr decltype(auto) originShape(Layout &layout) +{ + return layout.template originShape(); +} + +template +CATLASS_HOST_DEVICE constexpr decltype(auto) originShape(Layout const &layout) +{ + return layout.template originShape(); +} + +// Return the rank of layout +template +CATLASS_HOST_DEVICE constexpr auto rank(Layout const &layout) +{ + return rank(shape(layout)); +} + +// Return the depth of the layout +template +CATLASS_HOST_DEVICE constexpr auto depth(Layout const &layout) +{ + return depth(shape(layout)); +} + +// Return the offset of coord +template +CATLASS_HOST_DEVICE constexpr auto crd2offset(Coord const &coord, Shape const &shape, Stride const &stride); + +namespace detail { + +template +CATLASS_HOST_DEVICE constexpr auto crd2offset_ttt(Coord const &coord, Shape const &shape, Stride const &stride, seq) +{ + return (... + crd2offset(get(coord), get(shape), get(stride))); +} + +template +CATLASS_HOST_DEVICE constexpr auto crd2offset_itt(CInt const &coord, STuple const &shape, DTuple const &stride, seq) +{ + if constexpr (sizeof...(Is) == 0) { // Avoid recursion and mod on single/last iter + return crd2offset(coord, get(shape), get(stride)); + } else if constexpr (is_constant<0, CInt>::value) { + return crd2offset(_0{}, get(shape), get(stride)) + + (_0{} + ... + crd2offset(_0{}, get(shape), get(stride))); + } else { // General case + return crd2offset(coord % Product{}(get(shape)), get(shape), get(stride)) + + crd2offset_itt(coord / Product{}(get(shape)), shape, stride, seq{}); + } +} + +} // end namespace detail + +template +CATLASS_HOST_DEVICE constexpr auto crd2offset(Coord const &coord, Shape const &shape, Stride const &stride) +{ + if constexpr (is_tuple::value) { + if constexpr (is_tuple::value) { // tuple tuple tuple + static_assert(tuple_size::value == tuple_size::value, "Mismatched Ranks"); + static_assert(tuple_size::value == tuple_size::value, "Mismatched Ranks"); + return detail::crd2offset_ttt(coord, shape, stride, tuple_seq{}); + } else { // tuple "int" "int" + static_assert(sizeof(Coord) == 0, "Invalid parameters"); + } + } else { + if constexpr (is_tuple::value) { // "int" tuple tuple + static_assert(tuple_size::value == tuple_size::value, "Mismatched Ranks"); + return detail::crd2offset_itt(coord, shape, stride, tuple_seq{}); + } else { // "int" "int" "int" + return coord * stride; + } + } +} + +template +struct is_layout : false_type {}; +template +struct is_layout> : true_type {}; + +// Layout Check +namespace detail { + +template +struct isVector { + static bool const value = false; +}; + +template +struct isVector> { + static bool const value = (stride<0>(Layout{}) == 1); +}; + +template +struct isRowMajor { + static bool const value = false; +}; + +template +struct isRowMajor> { + static bool const value = (stride<1>(Layout{}) == 1); +}; + +template +struct isColumnMajor { + static bool const value = false; +}; + +template +struct isColumnMajor> { + static bool const value = (stride<0>(Layout{}) == 1); +}; + +template +struct iszN { + static bool const value = false; +}; + +template +struct iszN> { + static constexpr uint32_t ELE_NUM_PER_C0 = Catlass::BYTE_PER_C0 / sizeof(Element); + static constexpr uint32_t ELE_NUM_PER_FRACTAL = Catlass::BYTE_PER_FRACTAL / sizeof(Element); + static bool const value = (shape<0, 0>(Layout{}) == Catlass::C0_NUM_PER_FRACTAL && + shape<1, 0>(Layout{}) == ELE_NUM_PER_C0 && + stride<1, 0>(Layout{}) == 1 && + stride<0, 1>(Layout{}) == ELE_NUM_PER_FRACTAL); +}; + +/* +For matmul m axis is not c0 Aligned. +Exp: oriShape(m, k) : (127, 256) +zNUnAlign shape:((127, 1), (16, 256/16)) zN shape: ((16, Ceil(127/16)), (16, 256/16)) +*/ +template +struct iszNUnAlign { + static bool const value = false; +}; + +template +struct iszNUnAlign, std::enable_if_t(Layout{}))> == 2 && rank_v(Layout{}))> == 2>> { + static constexpr uint32_t ELE_NUM_PER_C0 = Catlass::BYTE_PER_C0 / sizeof(Element); + static bool const value = (shape<0, 1>(Layout{}) == 1 && + shape<1, 0>(Layout{}) == ELE_NUM_PER_C0 && + stride<0, 0>(Layout{}) == ELE_NUM_PER_C0 && + stride<1, 0>(Layout{}) == 1); +}; + +template +struct iszZ { + static bool const value = false; +}; + +template +struct iszZ> { + static constexpr uint32_t ELE_NUM_PER_C0 = Catlass::BYTE_PER_C0 / sizeof(Element); + static constexpr uint32_t ELE_NUM_PER_FRACTAL = Catlass::BYTE_PER_FRACTAL / sizeof(Element); + static bool const value = (shape<0, 0>(Layout{}) == Catlass::C0_NUM_PER_FRACTAL && + shape<1, 0>(Layout{}) == ELE_NUM_PER_C0 && + stride<1, 0>(Layout{}) == 1 && + stride<1, 1>(Layout{}) == ELE_NUM_PER_FRACTAL); +}; + +template +struct isnZ { + static bool const value = false; +}; + +template +struct isnZ> { + static constexpr uint32_t ELE_NUM_PER_C0 = Catlass::BYTE_PER_C0 / sizeof(Element); + static constexpr uint32_t ELE_NUM_PER_FRACTAL = Catlass::BYTE_PER_FRACTAL / sizeof(Element); + static bool const value = (shape<0, 0>(Layout{}) == ELE_NUM_PER_C0 && + shape<1, 0>(Layout{}) == Catlass::C0_NUM_PER_FRACTAL && + stride<0, 0>(Layout{}) == 1 && + stride<1, 1>(Layout{}) == ELE_NUM_PER_FRACTAL); +}; + +} // end namespace detail + +// Advanced Layout constructions + +// Make a vector layout. +template +CATLASS_HOST_DEVICE constexpr auto MakeLayout(T const &len) +{ + return MakeLayout(MakeShape(len), MakeStride(Int<1>{}), MakeShape(len)); +} + +// Make a inner layout with Rows and Cols. +template +CATLASS_HOST_DEVICE constexpr auto MakeLayout(T const &rows, U const &cols) +{ + static_assert(std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v, + "Unsupported LayoutTag for MakeLayoutFromTag, only support Catlass::layout::RowMajor or" + "Catlass::layout::ColumnMajor or Catlass::layout::zN or Catlass::layout::nZ or Catlass::layout::zZ or Catlass::layout::L0C"); + + constexpr uint32_t ELE_NUM_PER_C0 = Catlass::BYTE_PER_C0 / sizeof(Element); + constexpr uint32_t ELE_NUM_PER_FRACTAL = Catlass::BYTE_PER_FRACTAL / sizeof(Element); + + if constexpr (std::is_same_v) { + return MakeLayout(MakeShape(cols), MakeStride(Int<1>{}), MakeShape(cols)); + } else if constexpr (std::is_same_v) { + return MakeLayout(MakeShape(rows, cols), + MakeStride((int64_t)cols, Int<1>{}), + MakeShape(rows, cols)); + } else if constexpr (std::is_same_v) { + return MakeLayout(MakeShape(rows, cols), + MakeStride(Int<1>{}, (int64_t)rows), + MakeShape(rows, cols)); + } else if constexpr (std::is_same_v) { + return MakeLayout( + MakeShape(MakeShape(Int{}, CeilDiv(rows, Int{})), + MakeShape(Int{}, CeilDiv(cols, Int{}))), + MakeStride(MakeStride(Int{}, Int{}), + MakeStride(Int<1>{}, RoundUp((int64_t)rows, Int{}) * ELE_NUM_PER_C0)), + MakeShape(rows, cols)); + } else if constexpr (std::is_same_v) { + return MakeLayout( + MakeShape(MakeShape(Int{}, CeilDiv(rows, Int{})), + MakeShape(Int{}, CeilDiv(cols, Int{}))), + MakeStride(MakeStride(Int{}, + RoundUp((int64_t)cols, Int{}) * Catlass::C0_NUM_PER_FRACTAL), + MakeStride(Int<1>{}, Int{})), + MakeShape(rows, cols)); + } else if constexpr (std::is_same_v) { + constexpr uint32_t ELE_NUM_PER_FRACTAL = 256; + return MakeLayout( + MakeShape(MakeShape(Int{}, CeilDiv(rows, Int{})), + MakeShape(Int{}, CeilDiv(cols, Int{}))), + MakeStride(MakeStride(Int{}, Int{}), + MakeStride(Int<1>{}, RoundUp((int64_t)rows, Int{}) * Catlass::C0_NUM_PER_FRACTAL)), + MakeShape(rows, cols)); + } else { + return MakeLayout( + MakeShape(MakeShape(Int{}, CeilDiv(rows, Int{})), + MakeShape(Int{}, CeilDiv(cols, Int{}))), + MakeStride( + MakeStride(Int<1>{}, RoundUp((int64_t)cols, Int{}) * ELE_NUM_PER_C0), + MakeStride(Int{}, Int{})), + MakeShape(rows, cols)); + } +} + +namespace detail { + +template +CATLASS_HOST_DEVICE constexpr auto CropOriginShape(OriginBase const &originBase, TileShape const &tileShape, Coord const &coord, seq) +{ + return MakeShape( + tla::min( + static_cast(get(tileShape)), + (static_cast(get(coord)) < static_cast(get(originBase))) ? (static_cast(get(originBase)) - static_cast(get(coord))) : 0u)...); +} + +} // namespace detail + +/// 创建 tile layout：使用指定的 tile 尺寸用于内存布局计算，同时携带实际逻辑尺寸（origin_shape）。 +/// coord 是元素坐标，用于计算实际的 originShape（处理边界情况）。 +/// Supports layouts of any rank (rank >= 1) for depth==1 layouts. +/// For depth>1 (fractal) layouts, currently only rank-2 is supported. +template +CATLASS_HOST_DEVICE constexpr auto GetTileLayout(Layout const &layout, TileShape const &tileShape, Coord const &coord) +{ + static_assert(is_tuple::value && depth_v == 1 && rank_v >= 1, + "GetTileLayout: TileShape must be a flat tuple with rank >= 1."); + static_assert(is_tuple::value && depth_v == 1 && rank_v == rank_v, + "GetTileLayout: Coord must have the same rank as TileShape."); + + // 统一计算 tail tile 的逻辑尺寸（originShape 裁剪） + auto tileOriginShape = detail::CropOriginShape(layout.originShape(), tileShape, coord, tuple_seq{}); + + // depth==1 的布局（vector/matrix/tensor）：tile shape 直接作为 memory-layout shape + // 支持任意 rank >= 1（但必须与 layout.rank 匹配） + if constexpr (Layout::depth == 1) { + static_assert(Layout::rank == rank_v, + "GetTileLayout: for depth==1 layouts, TileShape rank must match layout rank."); + return MakeLayout(tileShape, layout.stride(), tileOriginShape); + } else { + // depth>1 的布局（fractal layout）：目前只支持 rank=2 + // 因为 fractal layout 通常用于矩阵（rank-2），需要把 (rows, cols) 转为同结构嵌套 shape + static_assert(rank_v == 2, + "GetTileLayout: for depth>1 (fractal) layouts, TileShape must be rank-2 (rows, cols)."); + + if constexpr (is_static(layout))>::value && + is_static(layout))>::value) { + const uint32_t rows = get<0>(tileShape); + const uint32_t cols = get<1>(tileShape); + constexpr uint32_t dstInnerShapeRow = decltype(shape<0, 0>(layout))::value; + constexpr uint32_t dstInnerShapeCol = decltype(shape<1, 0>(layout))::value; + return MakeLayout( + MakeShape(MakeShape(Int{}, CeilDiv(rows)), + MakeShape(Int{}, CeilDiv(cols))), + layout.stride(), + tileOriginShape); + } else { + const uint32_t rows = get<0>(tileShape); + const uint32_t cols = get<1>(tileShape); + const uint32_t dstInnerShapeRow = shape<0, 0>(layout); + const uint32_t dstInnerShapeCol = shape<1, 0>(layout); + return MakeLayout( + MakeShape(MakeShape(dstInnerShapeRow, CeilDiv(rows, dstInnerShapeRow)), + MakeShape(dstInnerShapeCol, CeilDiv(cols, dstInnerShapeCol))), + layout.stride(), + tileOriginShape); + } + } +} + +template +CATLASS_HOST_DEVICE constexpr auto MakeLayoutL0C(T const &rows, U const &cols) +{ + constexpr uint32_t ELE_NUM_PER_FRACTAL = 256; + return MakeLayout( + MakeShape(MakeShape(Int{}, CeilDiv(rows, Int{})), + MakeShape(Int{}, CeilDiv(cols, Int{}))), + MakeStride(MakeStride(Int{}, Int{}), + MakeStride(Int<1>{}, RoundUp((int64_t)rows, Int{}) * Catlass::C0_NUM_PER_FRACTAL)), + MakeShape(rows, cols)); +} + +// +// Layout transforms +// +namespace detail { + +// Prepend one leading dimension to a layout type +// This is the general form of "make batched layout": +// - Given a base layout type of rank R, create a new layout of rank R+1 by +// prefixing (batchShape, batchStride, batchOrigin) to the existing shape/stride/originShape. +// intentionally preserve the *types* of each stride element from the original layout +template > +struct PrependDimLayout; + +template +struct PrependDimLayout> { + using ShapeOld = tla::remove_cvref_t().shape())>; + using StrideOld = tla::remove_cvref_t().stride())>; + using OriginOld = tla::remove_cvref_t().originShape())>; + + using ShapeNew = tla::Shape(std::declval()))>...>; + using StrideNew = tla::Stride(std::declval()))>...>; + using OriginNew = tla::Shape(std::declval()))>...>; + + using type = tla::Layout; +}; + +} // namespace detail + +template +using PrependDimLayout_t = typename detail::PrependDimLayout::type; + +template +using MakeBatchedLayout_t = PrependDimLayout_t; + +} // end namespace tla + +#endif // KDA_TLA_LAYOUT_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/tla/numeric/integer_sequence.hpp b/csrc/chunk_kda_fwd/op_kernel/tla/numeric/integer_sequence.hpp new file mode 100644 index 0000000000..2bd81f0fce --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/tla/numeric/integer_sequence.hpp @@ -0,0 +1,67 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef TLA_NUMERIC_INTEGER_SEQUENCE_HPP +#define TLA_NUMERIC_INTEGER_SEQUENCE_HPP + +#include "tla/numeric/integral_constant.hpp" +#include "tla/type_traits.hpp" + +namespace tla { + +template +struct IntegerSequence { + using value_type = T; + static constexpr size_t size() { return sizeof...(Ns); } +}; + +template +struct MakeIntegerSequenceImpl; + +template +struct MakeIntegerSequenceImpl, T, 0> { + typedef IntegerSequence type; +}; + +template +struct MakeIntegerSequenceImpl, T, N> { + typedef typename MakeIntegerSequenceImpl, T, N - 1>::type type; +}; + +template +using MakeIntegerSequence = typename MakeIntegerSequenceImpl, T, N>::type; + +// index_sequence +template +using index_sequence = IntegerSequence; + +template +using make_index_sequence = MakeIntegerSequence; + +// int_sequence +template +using int_sequence = IntegerSequence; + +template +using make_int_sequence = MakeIntegerSequence; + +// Shortcuts +template +using seq = int_sequence; + +template +using make_seq = make_int_sequence; + +template +using tuple_seq = make_seq>::value>; + +} // end namespace tla + +#endif // TLA_NUMERIC_INTEGER_SEQUENCE_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/tla/numeric/integral_constant.hpp b/csrc/chunk_kda_fwd/op_kernel/tla/numeric/integral_constant.hpp new file mode 100644 index 0000000000..ea18b263d4 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/tla/numeric/integral_constant.hpp @@ -0,0 +1,196 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef TLA_NUMERIC_INTEGER_CONSTANT_HPP +#define TLA_NUMERIC_INTEGER_CONSTANT_HPP + +#include "catlass/detail/kda_macros.hpp" +#include "tla/type_traits.hpp" +#include "tla/numeric/math.hpp" + +namespace tla { + +// A constant value: short name and type-deduction for fast compilation +template +struct C { + using type = C; + static constexpr auto value = v; + using value_type = decltype(v); + CATLASS_HOST_DEVICE constexpr operator value_type() const noexcept { return value; } + CATLASS_HOST_DEVICE constexpr value_type operator()() const noexcept { return value; } +}; + +// Deprecate +template +using constant = C; + +template +using bool_constant = C; + +using true_type = bool_constant; +using false_type = bool_constant; + +template +using is_std_integral = std::is_integral; + +// A more std:: conforming integral_constant that enforces type but interops with C +template +struct integral_constant : C { + using type = integral_constant; + static constexpr T value = v; + using value_type = T; + CATLASS_HOST_DEVICE constexpr value_type operator()() const noexcept { return value; } +}; + +// Use tla::is_std_integral to match built-in integral types (int, int64_t, unsigned, etc) +// Use tla::is_integral to match both built-in integral types AND static integral types. + +template +struct is_integral : bool_constant::value> {}; +template +struct is_integral> : true_type {}; +template +struct is_integral> : true_type {}; + +// is_static detects if an (abstract) value is defined completely by its type (no members) +template +struct is_static : bool_constant>::value> {}; + +// is_constant detects if a type is a static integral type and if v is equal to a value + +template +struct is_constant : false_type {}; +template +struct is_constant : is_constant {}; +template +struct is_constant : is_constant {}; +template +struct is_constant : is_constant {}; +template +struct is_constant : is_constant {}; +template +struct is_constant> : bool_constant {}; +template +struct is_constant> : bool_constant {}; + +// +// Specializations +// + +template +using Int = C; +using _0 = Int<0>; +using _64 = Int<64>; +using _128 = Int<128>; +using _256 = Int<256>; +using _512 = Int<512>; + +// +// Underscore placeholder (for slicing semantics) +// +// Usage: +// - `tla::_` is an empty tag value that can be used inside `tla::Coord` / tensor indexing +// to indicate "take the whole dimension" (full slice). +struct Underscore { + using type = Underscore; +}; + +CATLASS_HOST_DEVICE constexpr Underscore _{}; + +template +struct is_underscore : false_type {}; +template <> +struct is_underscore : true_type {}; +template +struct is_underscore : is_underscore {}; +template +struct is_underscore : is_underscore {}; +template +struct is_underscore : is_underscore {}; +template +struct is_underscore : is_underscore {}; + +/***************/ +/** Operators **/ +/***************/ + +#define TLA_LEFT_UNARY_OP(OP) \ + template \ + CATLASS_HOST_DEVICE constexpr C<(OP t)> operator OP(C) \ + { \ + return {}; \ + } +#define TLA_BINARY_OP(OP) \ + template \ + CATLASS_HOST_DEVICE constexpr C<(t OP u)> operator OP(C, C) \ + { \ + return {}; \ + } + +TLA_LEFT_UNARY_OP(+); +TLA_LEFT_UNARY_OP(-); +TLA_LEFT_UNARY_OP(~); +TLA_LEFT_UNARY_OP(!); +TLA_LEFT_UNARY_OP(*); + +TLA_BINARY_OP(+); +TLA_BINARY_OP(-); +TLA_BINARY_OP(*); +TLA_BINARY_OP(/); +TLA_BINARY_OP(%); +TLA_BINARY_OP(&); +TLA_BINARY_OP(|); +TLA_BINARY_OP(^); +TLA_BINARY_OP(<<); +TLA_BINARY_OP(>>); + +#undef TLA_BINARY_OP +#undef TLA_LEFT_UNARY_OP +#undef TLA_RIGHT_UNARY_OP + +// +// Named functions from math.hpp +// + +#define TLA_NAMED_UNARY_FN(OP) \ + template \ + CATLASS_HOST_DEVICE constexpr auto OP(C) \ + { \ + return C{}; \ + } +#define TLA_NAMED_BINARY_FN(OP) \ + template \ + CATLASS_HOST_DEVICE constexpr auto OP(C, C) \ + { \ + return C{}; \ + } \ + template ::value)> \ + CATLASS_HOST_DEVICE constexpr auto OP(C, U u) \ + { \ + return OP(t, u); \ + } \ + template ::value)> \ + CATLASS_HOST_DEVICE constexpr auto OP(T t, C) \ + { \ + return OP(t, u); \ + } + +TLA_NAMED_BINARY_FN(max); +TLA_NAMED_BINARY_FN(min); +TLA_NAMED_BINARY_FN(add); + +#undef TLA_NAMED_UNARY_FN +#undef TLA_NAMED_BINARY_FN + +} // end namespace tla + +#endif // TLA_NUMERIC_INTEGER_CONSTANT_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/tla/numeric/math.hpp b/csrc/chunk_kda_fwd/op_kernel/tla/numeric/math.hpp new file mode 100644 index 0000000000..00ff91957b --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/tla/numeric/math.hpp @@ -0,0 +1,52 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef TLA_NUMERIC_MATH_HPP +#define TLA_NUMERIC_MATH_HPP + +#include "catlass/detail/kda_macros.hpp" +#include "tla/type_traits.hpp" + +namespace tla { + +// +// Common Operations +// + +template ::value && + std::is_arithmetic::value)> +CATLASS_HOST_DEVICE constexpr auto +max(T const &t, U const &u) +{ + return t < u ? u : t; +} + +template ::value && + std::is_arithmetic::value)> +CATLASS_HOST_DEVICE constexpr auto +min(T const &t, U const &u) +{ + return t < u ? t : u; +} + +template ::value && + std::is_arithmetic::value)> +CATLASS_HOST_DEVICE constexpr auto +add(T const &t, U const &u) +{ + return t + u; +} + +} // namespace tla + +#endif // TLA_NUMERIC_MATH_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/tla/tuple.hpp b/csrc/chunk_kda_fwd/op_kernel/tla/tuple.hpp new file mode 100644 index 0000000000..db8a39ec0a --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/tla/tuple.hpp @@ -0,0 +1,141 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef TLA_TUPLE_HPP +#define TLA_TUPLE_HPP + +#include "tla/numeric/integral_constant.hpp" +#include "tla/numeric/integer_sequence.hpp" + +namespace tla { + +namespace detail { + +// EBO stands for "empty base optimization." +template ::value> +struct EBO; + +// Specialization for types T that are empty; +template +struct EBO { + CATLASS_HOST_DEVICE constexpr EBO() {} + + CATLASS_HOST_DEVICE constexpr EBO(T const &) {} +}; + +template +CATLASS_HOST_DEVICE constexpr T getv(EBO const &) +{ + return {}; +} + +// Specialization for types T that are not empty; +template +struct EBO { + CATLASS_HOST_DEVICE constexpr EBO() + : t_{} + {} + + CATLASS_HOST_DEVICE constexpr EBO(T const &t) + : t_{t} + {} + + T t_; +}; + +template +CATLASS_HOST_DEVICE constexpr T const &getv(EBO const &x) +{ + return x.t_; +} + +template +CATLASS_HOST_DEVICE constexpr T &getv(EBO &x) +{ + return x.t_; +} + +// TupleBase +template +struct TupleBase; + +template +struct TupleBase, T...> : EBO... { + CATLASS_HOST_DEVICE constexpr TupleBase() {} + + CATLASS_HOST_DEVICE constexpr TupleBase(T const &...t) + : EBO(t)... + {} +}; + +} // end namespace detail + +// tla::tuple class. +template +struct tuple : detail::TupleBase, T...> { + CATLASS_HOST_DEVICE constexpr tuple() {} + + CATLASS_HOST_DEVICE constexpr tuple(T const &...t) + : detail::TupleBase, T...>(t...) + {} +}; + +// get for tla::tuple +template +CATLASS_HOST_DEVICE constexpr decltype(auto) get(tuple const &t) noexcept +{ + static_assert(I < sizeof...(T), "Index out of range"); + return detail::getv(t); +} + +template +CATLASS_HOST_DEVICE constexpr decltype(auto) get(tuple &t) noexcept +{ + static_assert(I < sizeof...(T), "Index out of range"); + return detail::getv(t); +} + +template +CATLASS_HOST_DEVICE constexpr decltype(auto) get(tuple &&t) noexcept +{ + static_assert(I < sizeof...(T), "Index out of range"); + return detail::getv(static_cast &&>(t)); +} + +namespace detail { + +template +auto has_tuple_size(T *) -> bool_constant<(0 <= tuple_size::value)>; +auto has_tuple_size(...) -> false_type; + +} // end namespace detail + +template +struct is_tuple : decltype(detail::has_tuple_size((T *)0)){}; + +template +struct tuple_size> + : std::integral_constant {}; + +template +struct tuple_size> + : std::integral_constant {}; + +// make_tuple +template +CATLASS_HOST_DEVICE constexpr tuple +MakeTuple(T const &...t) +{ + return {t...}; +} + +} // end namespace tla + +#endif // TLA_TUPLE_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/tla/type_traits.hpp b/csrc/chunk_kda_fwd/op_kernel/tla/type_traits.hpp new file mode 100644 index 0000000000..81ac1427c2 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/tla/type_traits.hpp @@ -0,0 +1,43 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef TLA_UTIL_TYPE_TRAITS_HPP +#define TLA_UTIL_TYPE_TRAITS_HPP + +#include + +#define TLA_REQUIRES(...) typename std::enable_if<(__VA_ARGS__)>::type * = nullptr + +namespace tla { + +// using std::remove_cvref; +template +struct remove_cvref { + using type = std::remove_cv_t>; +}; + +// using std::remove_cvref_t; +template +using remove_cvref_t = typename remove_cvref::type; + +// tuple_size, tuple_element +template +struct tuple_size; + +template +struct tuple_size::type>> + : std::integral_constant::value> {}; + +template +constexpr size_t tuple_size_v = tuple_size::value; + +} // end namespace tla + +#endif // TLA_UTIL_TYPE_TRAITS_HPP diff --git a/csrc/pytorch_extensions.cpp b/csrc/pytorch_extensions.cpp index 29033c0b53..3a971de949 100644 --- a/csrc/pytorch_extensions.cpp +++ b/csrc/pytorch_extensions.cpp @@ -193,6 +193,17 @@ TORCH_LIBRARY_FRAGMENT(npu, m) "bool has_ori_kv=True, bool has_cmp_kv=True) -> Tensor"); #endif + m.def( + "chunk_kda_fwd(Tensor q, Tensor k, Tensor v, Tensor g, Tensor beta, " + "Tensor? a_log=None, Tensor? dt_bias=None, Tensor? initial_state=None, " + "Tensor? cu_seqlens=None, Tensor? chunk_indices=None, " + "str layout='BSND', float scale=1.0, int chunk_size=64, " + "bool safe_gate=False, float lower_bound=-5.0, bool use_gate_in_kernel=False, " + "bool state_v_first=False, bool output_final_state=True, bool output_gk=False, " + "bool output_w=False, bool output_u=False, bool output_qg=False, " + "bool output_kg=False, bool output_v_new=False, bool output_h=False) " + "-> (Tensor, Tensor?, Tensor?, Tensor, Tensor, Tensor?, Tensor?, Tensor?, Tensor?, Tensor?, Tensor?)"); + #ifdef BUILD_CATLASS_MODULE m.def("catlass_matmul_basic(Tensor tensor_a, Tensor tensor_b, Tensor(a!) tensor_c, str? format_mode=None) -> ()"); @@ -309,6 +320,8 @@ TORCH_LIBRARY_IMPL(npu, PrivateUse1, m) }); #endif + m.impl("chunk_kda_fwd", TORCH_FN(sglang::npu_kernel::chunk_kda_fwd)); + #ifdef BUILD_CATLASS_MODULE m.impl("catlass_matmul_basic", TORCH_FN(sglang::npu_kernel::catlass_matmul_basic)); diff --git a/include/sgl_kenel_npu_ops.h b/include/sgl_kenel_npu_ops.h index 498e3f26a7..6accd5ef49 100644 --- a/include/sgl_kenel_npu_ops.h +++ b/include/sgl_kenel_npu_ops.h @@ -239,6 +239,18 @@ at::Tensor sparse_attn_sharedkv_metadata_host( int64_t cmp_mask_mode, int64_t ori_win_left, int64_t ori_win_right, bool has_ori_kv, bool has_cmp_kv); +std::tuple, c10::optional, at::Tensor, at::Tensor, + c10::optional, c10::optional, c10::optional, + c10::optional, c10::optional, c10::optional> +chunk_kda_fwd(const at::Tensor &q, const at::Tensor &k, const at::Tensor &v, const at::Tensor &g, + const at::Tensor &beta, const c10::optional &aLog, + const c10::optional &dtBias, const c10::optional &initialState, + const c10::optional &cuSeqlens, + const c10::optional &chunkIndices, const std::string &layout, double scale, + int64_t chunkSize, bool safeGate, double lowerBound, bool useGateInKernel, + bool stateVFirst, bool outputFinalState, bool outputGk, bool outputW, bool outputU, + bool outputQG, bool outputKg, bool outputVNew, bool outputH); + } // namespace npu_kernel } // namespace sglang diff --git a/python/sgl_kernel_npu/sgl_kernel_npu/fla/chunk_kda_fwd.py b/python/sgl_kernel_npu/sgl_kernel_npu/fla/chunk_kda_fwd.py new file mode 100644 index 0000000000..37dbefae01 --- /dev/null +++ b/python/sgl_kernel_npu/sgl_kernel_npu/fla/chunk_kda_fwd.py @@ -0,0 +1,93 @@ +from typing import Optional + +import torch + + +def chunk_kda_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + *, + a_log: Optional[torch.Tensor] = None, + dt_bias: Optional[torch.Tensor] = None, + initial_state: Optional[torch.Tensor] = None, + cu_seqlens: Optional[torch.Tensor] = None, + chunk_indices: Optional[torch.Tensor] = None, + layout: str = "BSND", + scale: float = 1.0, + chunk_size: int = 64, + safe_gate: bool = False, + lower_bound: float = -5.0, + use_gate_in_kernel: bool = False, + state_v_first: bool = False, + output_final_state: bool = True, + output_gk: bool = False, + output_w: bool = False, + output_u: bool = False, + output_qg: bool = False, + output_kg: bool = False, + output_v_new: bool = False, + output_h: bool = False, +) -> tuple: + r"""Fused chunk KDA forward kernel (direct-launch NPU implementation). + + Args: + q: (B, S, H, K) bfloat16 (layout-dependent, see ``layout``). + k: (B, S, H, K) bfloat16, same shape as q. + v: (B, S, HV, V) bfloat16. + g: (B, S, HV, K) float32 or bfloat16 raw gate. + beta: (B, S, HV) float32 or bfloat16 delta coefficient. + a_log: (HV,) float32 gate decay parameter, required when use_gate_in_kernel. + dt_bias: (HV*K,) float32 gate bias. + initial_state: (N, HV, K, V) [or (N, HV, V, K) when state_v_first] float32. + cu_seqlens: (N+1,) int64 varlen sequence boundaries. + chunk_indices: (2*NC,) int64 canonical (seq_id, chunk_id) pairs. + layout: one of "BSND", "BNSD", "TND", "NTD". + scale: query scaling factor, usually K^(-0.5). + chunk_size: 64 or 128. + safe_gate: whether to use a bounded gate. + lower_bound: safe gate lower bound in [-5, 0). + use_gate_in_kernel: whether to compute the activated gate in-kernel. + state_v_first: whether the state tensors have (V, K) last two dims. + output_*: whether to materialize the corresponding intermediate output. + + Returns: + A tuple of length 11: + (attn_out, final_state?, gk?, aqk_out, akk_out, w?, u?, qg?, kg?, v_new?, h?). + attn_out: (B, S, HV, V) [or (T, HV, V) for rank-3 layouts]. + final_state: (N, HV, K, V) [or (N, HV, V, K) when state_v_first] float32. + gk: (B, HV, S, K) [or (HV, S, K)] float32. + aqk_out/akk_out: (B, HV, S, C) [or (HV, S, C)]. + w/qg/kg: (B, HV, S, K) [or (HV, S, K)]. + u/v_new: (B, HV, S, V) [or (HV, S, V)]. + h: (B, NC, HV, K, V) [or (NC, HV, K, V)]. + """ + return torch.ops.npu.chunk_kda_fwd( + q, + k, + v, + g, + beta, + a_log, + dt_bias, + initial_state, + cu_seqlens, + chunk_indices, + layout, + scale, + chunk_size, + safe_gate, + lower_bound, + use_gate_in_kernel, + state_v_first, + output_final_state, + output_gk, + output_w, + output_u, + output_qg, + output_kg, + output_v_new, + output_h, + ) From 6d0d5de46b8efb3bb090f8eab9a313a47d95bb26 Mon Sep 17 00:00:00 2001 From: Your Name Date: Sat, 29 Aug 2026 10:19:15 +0800 Subject: [PATCH 02/11] fix chunk kda fwd --- csrc/chunk_kda_fwd/op_host/chunk_kda_fwd.cpp | 28 ++++++-- .../chunk_kda_fwd/op_kernel/chunk_kda_fwd.cpp | 20 +++++- .../op_kernel/chunk_kda_fwd_prepare.h | 1 + csrc/deepep/ops/CMakePresets.json | 67 +++++++++++++++++++ .../attentions/attentions/plugin/__init__.py | 0 .../sgl_kernel_npu/sgl_kernel_npu/config.ini | 2 + 6 files changed, 113 insertions(+), 5 deletions(-) create mode 100644 csrc/deepep/ops/CMakePresets.json create mode 100644 python/attentions/attentions/plugin/__init__.py create mode 100644 python/sgl_kernel_npu/sgl_kernel_npu/config.ini diff --git a/csrc/chunk_kda_fwd/op_host/chunk_kda_fwd.cpp b/csrc/chunk_kda_fwd/op_host/chunk_kda_fwd.cpp index 9b176669c4..f8bb1e305b 100644 --- a/csrc/chunk_kda_fwd/op_host/chunk_kda_fwd.cpp +++ b/csrc/chunk_kda_fwd/op_host/chunk_kda_fwd.cpp @@ -21,6 +21,7 @@ #include #include +#include #include #include #include @@ -61,6 +62,15 @@ constexpr int64_t MAX_KDA_K_DIM = 256; constexpr int64_t MAX_KDA_HEAD_NUM = 128; constexpr int64_t MAX_KDA_VARLEN_SEQUENCES = 1024; +// Diagnostic-only zero initialization. These switches make it possible to +// distinguish an incomplete kernel write/read-before-write from ordinary +// numerical differences without rebuilding for every A/B run. +bool DiagnosticZeroInitEnabled(const char *name) +{ + const char *value = std::getenv(name); + return value != nullptr && value[0] != '\0' && std::strcmp(value, "0") != 0; +} + enum class KdaFwdLayout { BSND, BNSD, @@ -540,8 +550,13 @@ chunk_kda_fwd(const at::Tensor &q, const at::Tensor &k, const at::Tensor &v, con at::Tensor finalStateOut; void *finalStatePtr = nullptr; if (storeFinalState) { - finalStateOut = at::empty({info.seqNum, hv, kd, vd}, - at::TensorOptions().dtype(at::kFloat).device(q.device())); + const auto finalStateOptions = + at::TensorOptions().dtype(at::kFloat).device(q.device()); + const std::vector finalStateShape = {info.seqNum, hv, kd, vd}; + finalStateOut = + DiagnosticZeroInitEnabled("SGLANG_CHUNK_KDA_DIAG_ZERO_FINAL_STATE") + ? at::zeros(finalStateShape, finalStateOptions) + : at::empty(finalStateShape, finalStateOptions); finalStatePtr = finalStateOut.data_ptr(); } @@ -607,10 +622,15 @@ chunk_kda_fwd(const at::Tensor &q, const at::Tensor &k, const at::Tensor &v, con const uint64_t tokenHeads = static_cast(b) * hv * s; const uint64_t totalUserWorkspace = tilingData.outputScratchOffset + 2 * tokenHeads * vd * sizeof(float); + const uint64_t sysytemWorkspaceBytes = static_cast(ascendcPlatform->GetLibApiWorkSpaceSize()); const int64_t totalWorkspaceBytes = - SYS_WORKSPACE_SIZE + static_cast(AlignWorkspace(totalUserWorkspace)); + sysytemWorkspaceBytes + static_cast(AlignWorkspace(totalUserWorkspace)); + const auto workspaceOptions = at::TensorOptions().dtype(at::kByte).device(q.device()); + const std::vector workspaceShape = {totalWorkspaceBytes}; auto workspaceTensor = - at::empty({totalWorkspaceBytes}, at::TensorOptions().dtype(at::kByte).device(q.device())); + DiagnosticZeroInitEnabled("SGLANG_CHUNK_KDA_DIAG_ZERO_WORKSPACE") + ? at::zeros(workspaceShape, workspaceOptions) + : at::empty(workspaceShape, workspaceOptions); EXEC_KERNEL_CMD(chunk_kda_fwd, blockDim, qHead, kHead, vHead, gHead, betaHead, aLogPtr, dtBiasPtr, initStatePtr, cuSeqlensPtr, chunkIndicesPtr, attnOut, finalStatePtr, gkPtr, aqkOut, diff --git a/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd.cpp b/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd.cpp index fc2321360a..e03c8aff33 100644 --- a/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd.cpp +++ b/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd.cpp @@ -109,7 +109,23 @@ __aicore__ inline void DispatchGenericSafeGate( using ChunkKdaFwd::ChunkKdaFwdTilingData; -extern "C" __global__ __aicore__ void chunk_kda_fwd( +// FwdH uses SyncAll() and cross-core flags. The original framework +// operator requests batch scheduling with TilingContext::SetScheduleMode(1). +// Kernel-launch operators must express the same requirement on the entrypoint +// so all requested cores are available together instead of being scheduled in +// partial waves. +// +// ascendc_library's generated device translation unit defines the guard below, +// temporarily rewrites __global__ to inline, and includes this source as the +// implementation body. Schedule-mode attributes are invalid on that inline +// copy, so attach the attribute only to the real kernel declaration processed +// outside the generated inline-body pass. +#if defined(__CHUNK_KDA_FWD__KERNEL_FUN_H__) +#define KDA_FWD_SCHEDULE_MODE +#else +#define KDA_FWD_SCHEDULE_MODE __schedmode__(1) +#endif +extern "C" KDA_FWD_SCHEDULE_MODE __global__ __aicore__ void chunk_kda_fwd( GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR g, GM_ADDR beta, GM_ADDR a_log, GM_ADDR dt_bias, GM_ADDR initial_state, GM_ADDR cu_seqlens, GM_ADDR chunk_indices, GM_ADDR attn_out, @@ -185,4 +201,6 @@ extern "C" __global__ __aicore__ void chunk_kda_fwd( } } +#undef KDA_FWD_SCHEDULE_MODE + #undef KDA_COMPILE_ARCH35_FAST_PATH diff --git a/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_prepare.h b/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_prepare.h index c5f6004802..0bb747f8b6 100644 --- a/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_prepare.h +++ b/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_prepare.h @@ -1893,6 +1893,7 @@ class ChunkKdaFwdPrepareKernel { WaitFlag(mte2ToVEvent_); Muls(aqkLocal, aqkLocal, scale_, static_cast(matrixElems)); + PipeBarrier(); ClampFp32ToOutputType(aqkLocal, static_cast(matrixElems)); ClampFp32ToOutputType(akkLocal, static_cast(matrixElems)); Cast(aqkTyped, aqkLocal, RoundMode::CAST_RINT, static_cast(matrixElems)); diff --git a/csrc/deepep/ops/CMakePresets.json b/csrc/deepep/ops/CMakePresets.json new file mode 100644 index 0000000000..882998d920 --- /dev/null +++ b/csrc/deepep/ops/CMakePresets.json @@ -0,0 +1,67 @@ +{ + "version": 1, + "cmakeMinimumRequired": { + "major": 3, + "minor": 16, + "patch": 0 + }, + "configurePresets": [ + { + "name": "default", + "displayName": "Default Config", + "description": "Default build using Unix Makefiles generator", + "generator": "Unix Makefiles", + "binaryDir": "${sourceDir}/build_out", + "cacheVariables": { + "CMAKE_BUILD_TYPE": { + "type": "STRING", + "value": "Release" + }, + "ENABLE_SOURCE_PACKAGE": { + "type": "BOOL", + "value": "True" + }, + "ENABLE_BINARY_PACKAGE": { + "type": "BOOL", + "value": "True" + }, + "ASCEND_COMPUTE_UNIT": { + "type": "STRING", + "value": "ascend910_93" + }, + "ENABLE_TEST": { + "type": "BOOL", + "value": "True" + }, + "vendor_name": { + "type": "STRING", + "value": "hwcomputing" + }, + "ASCEND_CANN_PACKAGE_PATH": { + "type": "PATH", + "value": "/usr/local/Ascend/cann-9.0.0" + }, + "ASCEND_PYTHON_EXECUTABLE": { + "type": "STRING", + "value": "python3" + }, + "CMAKE_INSTALL_PREFIX": { + "type": "PATH", + "value": "${sourceDir}/build_out" + }, + "ENABLE_CROSS_COMPILE": { + "type": "BOOL", + "value": "False" + }, + "CMAKE_CROSS_PLATFORM_COMPILER": { + "type": "PATH", + "value": "/usr/bin/aarch64-linux-gnu-g++" + }, + "ASCEND_PACK_SHARED_LIBRARY": { + "type": "BOOL", + "value": "False" + } + } + } + ] +} \ No newline at end of file diff --git a/python/attentions/attentions/plugin/__init__.py b/python/attentions/attentions/plugin/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/python/sgl_kernel_npu/sgl_kernel_npu/config.ini b/python/sgl_kernel_npu/sgl_kernel_npu/config.ini new file mode 100644 index 0000000000..ec216af1d6 --- /dev/null +++ b/python/sgl_kernel_npu/sgl_kernel_npu/config.ini @@ -0,0 +1,2 @@ +[global] +version = 2026.6.1 From 79bd713df374fedbc1ac9842190a79d7754a169f Mon Sep 17 00:00:00 2001 From: chenjz Date: Sat, 29 Aug 2026 18:09:28 +0800 Subject: [PATCH 03/11] fix chunk kda fwd & delete unuseful file --- csrc/chunk_kda_fwd/op_host/chunk_kda_fwd.cpp | 25 +------ csrc/deepep/ops/CMakePresets.json | 67 ------------------- .../attentions/attentions/plugin/__init__.py | 0 .../sgl_kernel_npu/sgl_kernel_npu/config.ini | 2 - 4 files changed, 3 insertions(+), 91 deletions(-) delete mode 100644 csrc/deepep/ops/CMakePresets.json delete mode 100644 python/attentions/attentions/plugin/__init__.py delete mode 100644 python/sgl_kernel_npu/sgl_kernel_npu/config.ini diff --git a/csrc/chunk_kda_fwd/op_host/chunk_kda_fwd.cpp b/csrc/chunk_kda_fwd/op_host/chunk_kda_fwd.cpp index f8bb1e305b..60650632fe 100644 --- a/csrc/chunk_kda_fwd/op_host/chunk_kda_fwd.cpp +++ b/csrc/chunk_kda_fwd/op_host/chunk_kda_fwd.cpp @@ -21,7 +21,6 @@ #include #include -#include #include #include #include @@ -48,7 +47,6 @@ namespace sglang { namespace npu_kernel { -constexpr int64_t SYS_WORKSPACE_SIZE = 16777216; // 16 MB reserved system workspace constexpr uint32_t PADDING_BYTE = 32U; constexpr uint64_t KDA_ALIGN = 512; @@ -62,14 +60,6 @@ constexpr int64_t MAX_KDA_K_DIM = 256; constexpr int64_t MAX_KDA_HEAD_NUM = 128; constexpr int64_t MAX_KDA_VARLEN_SEQUENCES = 1024; -// Diagnostic-only zero initialization. These switches make it possible to -// distinguish an incomplete kernel write/read-before-write from ordinary -// numerical differences without rebuilding for every A/B run. -bool DiagnosticZeroInitEnabled(const char *name) -{ - const char *value = std::getenv(name); - return value != nullptr && value[0] != '\0' && std::strcmp(value, "0") != 0; -} enum class KdaFwdLayout { BSND, @@ -550,13 +540,8 @@ chunk_kda_fwd(const at::Tensor &q, const at::Tensor &k, const at::Tensor &v, con at::Tensor finalStateOut; void *finalStatePtr = nullptr; if (storeFinalState) { - const auto finalStateOptions = - at::TensorOptions().dtype(at::kFloat).device(q.device()); - const std::vector finalStateShape = {info.seqNum, hv, kd, vd}; - finalStateOut = - DiagnosticZeroInitEnabled("SGLANG_CHUNK_KDA_DIAG_ZERO_FINAL_STATE") - ? at::zeros(finalStateShape, finalStateOptions) - : at::empty(finalStateShape, finalStateOptions); + finalStateOut = at::empty({info.seqNum, hv, kd, vd}, + at::TensorOptions().dtype(at::kFloat).device(q.device())); finalStatePtr = finalStateOut.data_ptr(); } @@ -625,12 +610,8 @@ chunk_kda_fwd(const at::Tensor &q, const at::Tensor &k, const at::Tensor &v, con const uint64_t sysytemWorkspaceBytes = static_cast(ascendcPlatform->GetLibApiWorkSpaceSize()); const int64_t totalWorkspaceBytes = sysytemWorkspaceBytes + static_cast(AlignWorkspace(totalUserWorkspace)); - const auto workspaceOptions = at::TensorOptions().dtype(at::kByte).device(q.device()); - const std::vector workspaceShape = {totalWorkspaceBytes}; auto workspaceTensor = - DiagnosticZeroInitEnabled("SGLANG_CHUNK_KDA_DIAG_ZERO_WORKSPACE") - ? at::zeros(workspaceShape, workspaceOptions) - : at::empty(workspaceShape, workspaceOptions); + at::empty({totalWorkspaceBytes}, at::TensorOptions().dtype(at::kByte).device(q.device())); EXEC_KERNEL_CMD(chunk_kda_fwd, blockDim, qHead, kHead, vHead, gHead, betaHead, aLogPtr, dtBiasPtr, initStatePtr, cuSeqlensPtr, chunkIndicesPtr, attnOut, finalStatePtr, gkPtr, aqkOut, diff --git a/csrc/deepep/ops/CMakePresets.json b/csrc/deepep/ops/CMakePresets.json deleted file mode 100644 index 882998d920..0000000000 --- a/csrc/deepep/ops/CMakePresets.json +++ /dev/null @@ -1,67 +0,0 @@ -{ - "version": 1, - "cmakeMinimumRequired": { - "major": 3, - "minor": 16, - "patch": 0 - }, - "configurePresets": [ - { - "name": "default", - "displayName": "Default Config", - "description": "Default build using Unix Makefiles generator", - "generator": "Unix Makefiles", - "binaryDir": "${sourceDir}/build_out", - "cacheVariables": { - "CMAKE_BUILD_TYPE": { - "type": "STRING", - "value": "Release" - }, - "ENABLE_SOURCE_PACKAGE": { - "type": "BOOL", - "value": "True" - }, - "ENABLE_BINARY_PACKAGE": { - "type": "BOOL", - "value": "True" - }, - "ASCEND_COMPUTE_UNIT": { - "type": "STRING", - "value": "ascend910_93" - }, - "ENABLE_TEST": { - "type": "BOOL", - "value": "True" - }, - "vendor_name": { - "type": "STRING", - "value": "hwcomputing" - }, - "ASCEND_CANN_PACKAGE_PATH": { - "type": "PATH", - "value": "/usr/local/Ascend/cann-9.0.0" - }, - "ASCEND_PYTHON_EXECUTABLE": { - "type": "STRING", - "value": "python3" - }, - "CMAKE_INSTALL_PREFIX": { - "type": "PATH", - "value": "${sourceDir}/build_out" - }, - "ENABLE_CROSS_COMPILE": { - "type": "BOOL", - "value": "False" - }, - "CMAKE_CROSS_PLATFORM_COMPILER": { - "type": "PATH", - "value": "/usr/bin/aarch64-linux-gnu-g++" - }, - "ASCEND_PACK_SHARED_LIBRARY": { - "type": "BOOL", - "value": "False" - } - } - } - ] -} \ No newline at end of file diff --git a/python/attentions/attentions/plugin/__init__.py b/python/attentions/attentions/plugin/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/python/sgl_kernel_npu/sgl_kernel_npu/config.ini b/python/sgl_kernel_npu/sgl_kernel_npu/config.ini deleted file mode 100644 index ec216af1d6..0000000000 --- a/python/sgl_kernel_npu/sgl_kernel_npu/config.ini +++ /dev/null @@ -1,2 +0,0 @@ -[global] -version = 2026.6.1 From c5fc8d706af0392610a23d9b621b30839e613436 Mon Sep 17 00:00:00 2001 From: chenjz Date: Tue, 1 Sep 2026 09:44:00 +0800 Subject: [PATCH 04/11] fix(npu): fix chunk KDA synchronization issues --- .../block_epilogue_kda_fwdh_vnew_arch22.hpp | 14 ++++-- .../gemm/kernel/kda_fwd_h_kernel_arch22.hpp | 44 ++++++++++++++----- .../arch35/chunk_kda_fwd_finalize_arch35.h | 15 +++---- .../arch35/chunk_kda_fwd_prepare_arch35.h | 2 + .../op_kernel/chunk_kda_fwd_common.h | 13 +++++- csrc/pytorch_extensions.cpp | 2 + 6 files changed, 65 insertions(+), 25 deletions(-) diff --git a/csrc/chunk_kda_fwd/op_kernel/arch22/fwd_h/epilogue/block/block_epilogue_kda_fwdh_vnew_arch22.hpp b/csrc/chunk_kda_fwd/op_kernel/arch22/fwd_h/epilogue/block/block_epilogue_kda_fwdh_vnew_arch22.hpp index f857cc780d..7634c785ad 100644 --- a/csrc/chunk_kda_fwd/op_kernel/arch22/fwd_h/epilogue/block/block_epilogue_kda_fwdh_vnew_arch22.hpp +++ b/csrc/chunk_kda_fwd/op_kernel/arch22/fwd_h/epilogue/block/block_epilogue_kda_fwdh_vnew_arch22.hpp @@ -168,7 +168,7 @@ class BlockEpilogue uInput, AscendC::GlobalTensor wsInput, uint32_t chunkSize, uint32_t vBlockDim, uint32_t vHeadDim, Arch::CrossCoreFlag cube1Done, Arch::CrossCoreFlag vec1Done, bool isInitialState, bool isFinalState, bool storeFinalState, - bool waitWsFromMte3, bool isPing) + bool waitWsFromMte3, bool isPing, bool cube1AlreadyWaited) { static constexpr uint32_t ROW_TILE = 16; uint32_t mActual = chunkSize; @@ -185,7 +185,9 @@ class BlockEpilogue= mActual) { - Arch::CrossCoreWaitFlag(cube1Done); + if (!cube1AlreadyWaited) { + Arch::CrossCoreWaitFlag(cube1Done); + } // A zero-row AIV lane still owns the EVENT0 hand-off consumed by V2. if (waitWsFromMte3) { AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); @@ -233,7 +235,9 @@ class BlockEpilogue(EVENT_ID3 + pingpongFlag); } - Arch::CrossCoreWaitFlag(cube1Done); + if (!cube1AlreadyWaited) { + Arch::CrossCoreWaitFlag(cube1Done); + } if (waitWsFromMte3) { AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); @@ -289,7 +293,9 @@ class BlockEpilogue(EVENT_ID3 + pingpongFlag); } - Arch::CrossCoreWaitFlag(cube1Done); + if (!cube1AlreadyWaited) { + Arch::CrossCoreWaitFlag(cube1Done); + } bool waitWsThisTileFromMte3 = waitWsFromMte3; for (uint32_t rowStart = rowBegin; rowStart < rowEnd;) { diff --git a/csrc/chunk_kda_fwd/op_kernel/arch22/fwd_h/gemm/kernel/kda_fwd_h_kernel_arch22.hpp b/csrc/chunk_kda_fwd/op_kernel/arch22/fwd_h/gemm/kernel/kda_fwd_h_kernel_arch22.hpp index bfd8cb9802..fdc9076ab1 100644 --- a/csrc/chunk_kda_fwd/op_kernel/arch22/fwd_h/gemm/kernel/kda_fwd_h_kernel_arch22.hpp +++ b/csrc/chunk_kda_fwd/op_kernel/arch22/fwd_h/gemm/kernel/kda_fwd_h_kernel_arch22.hpp @@ -274,7 +274,8 @@ class KdaFwdHKernel { return static_cast(value); } - __aicore__ inline void ComputeTailVWorkspace(const KdaFwdHOffsets &offsets) + // Borrow the stream's V_MTE2 free token and restore it before returning. + __aicore__ inline void ComputeTailVWorkspace(const KdaFwdHOffsets &offsets, uint32_t tailEventId) { uint32_t subBlockIdx = AscendC::GetSubBlockIdx(); uint32_t subBlockNum = AscendC::GetSubBlockNum(); @@ -284,6 +285,8 @@ class KdaFwdHKernel { if (rowBegin >= rowEnd) { return; } + AscendC::ResetMask(); + AscendC::WaitFlag(tailEventId); constexpr uint32_t TAIL_INPUT_OFFSET = 166 * 1024; constexpr uint32_t TAIL_FLOAT_OFFSET = 167 * 1024; @@ -297,9 +300,16 @@ class KdaFwdHKernel { AscendC::PipeBarrier(); for (uint32_t kIdx = 0; kIdx < kHeadDim; ++kIdx) { AscendC::DataCopy(inputUb, gmH[offsets.hSrcOffset + kIdx * vHeadDim], offsets.vBlockDim); - AscendC::SetFlag(EVENT_ID7); - AscendC::WaitFlag(EVENT_ID7); + AscendC::SetFlag(tailEventId); + AscendC::WaitFlag(tailEventId); AscendC::Cast(floatUb, inputUb, AscendC::RoundMode::CAST_NONE, offsets.vBlockDim); + // inputUb is reused by the next kIdx iteration. PIPE_V and + // PIPE_MTE2 run independently, so a PIPE_V barrier alone does + // not stop the next DataCopy from overwriting inputUb while + // Cast is still reading it. Queue a V -> MTE2 dependency + // before that next copy, matching ComputeHWorkspaceVector(). + AscendC::SetFlag(tailEventId); + AscendC::WaitFlag(tailEventId); AscendC::PipeBarrier(); float weight = LoadScalarAsFloat(gmW, offsets.wOffset + tokenRow * kHeadDim + kIdx); AscendC::Muls(floatUb, floatUb, weight, offsets.vBlockDim); @@ -307,13 +317,19 @@ class KdaFwdHKernel { AscendC::Add(accumUb, accumUb, floatUb, offsets.vBlockDim); AscendC::PipeBarrier(); } - AscendC::SetFlag(EVENT_ID7); - AscendC::WaitFlag(EVENT_ID7); + AscendC::SetFlag(tailEventId); + AscendC::WaitFlag(tailEventId); AscendC::DataCopy(gmVWorkspace[offsets.vWorkOffset + tokenRow * offsets.vBlockDim], accumUb, offsets.vBlockDim); - AscendC::SetFlag(EVENT_ID7); - AscendC::WaitFlag(EVENT_ID7); + AscendC::SetFlag(tailEventId); + AscendC::WaitFlag(tailEventId); + // The MTE3 copy above still reads accumUb. The next token row + // starts by clearing the same buffer on PIPE_V, so protect that + // reuse explicitly; MTE3_MTE2 does not order MTE3 against PIPE_V. + AscendC::SetFlag(tailEventId); + AscendC::WaitFlag(tailEventId); } + AscendC::SetFlag(tailEventId); } __aicore__ inline void ComputeHWorkspaceVector(const KdaFwdHOffsets &offsets, uint32_t eventId) @@ -681,8 +697,16 @@ class KdaFwdHKernel { continue; } const KdaFwdHOffsets &vec1Offsets = vecBlockScheduler.GetCurTaskOffsets(stream); - if (vec1Offsets.blockTokens < 16) { - ComputeTailVWorkspace(vec1Offsets); + bool tailVectorPath = vec1Offsets.blockTokens < 16; + if (tailVectorPath) { + // For a short tail Cube does not calculate W @ H, + // but cube1Done still orders this task after the + // previous chunk's V2 update of gmH. Wait before + // the software W @ H reads gmH, not afterwards in + // the epilogue. + Arch::CrossCoreWaitFlag(vecBlockScheduler.cube1Done[streamId]); + ComputeTailVWorkspace(vec1Offsets, + EVENT_ID3 + (streamId == 0 ? 0 : pongEventOffset)); } bool waitWsFromMte3 = storeFinalState && std::is_same::value && event0FromMte3[streamId]; @@ -692,7 +716,7 @@ class KdaFwdHKernel { vec1Offsets.vBlockDim, vHeadDim, vecBlockScheduler.cube1Done[streamId], vecBlockScheduler.vec1Done[streamId], vec1Offsets.isInitialState, vec1Offsets.isFinalState, storeFinalState, waitWsFromMte3, - (streamId == 0)); + (streamId == 0), tailVectorPath); AscendC::SetFlag(EVENT_ID1); AscendC::WaitFlag(EVENT_ID1); if (storeFinalState && std::is_same::value) { diff --git a/csrc/chunk_kda_fwd/op_kernel/arch35/chunk_kda_fwd_finalize_arch35.h b/csrc/chunk_kda_fwd/op_kernel/arch35/chunk_kda_fwd_finalize_arch35.h index f77a4caa73..39d021f0da 100644 --- a/csrc/chunk_kda_fwd/op_kernel/arch35/chunk_kda_fwd_finalize_arch35.h +++ b/csrc/chunk_kda_fwd/op_kernel/arch35/chunk_kda_fwd_finalize_arch35.h @@ -330,7 +330,6 @@ class ChunkKdaFwdFinalizeKernel { mte3ToVEvent_ = pipe_->AllocEventID(); mte2ToMte3Event_ = pipe_->AllocEventID(); mte3ToMte2Event_ = pipe_->AllocEventID(); - vToSEvent_ = pipe_->AllocEventID(); sToVEvent_ = pipe_->AllocEventID(); sToMte2Event_ = pipe_->AllocEventID(); vectorEventsAllocated_ = true; @@ -347,7 +346,6 @@ class ChunkKdaFwdFinalizeKernel { pipe_->ReleaseEventID(mte3ToVEvent_); pipe_->ReleaseEventID(mte2ToMte3Event_); pipe_->ReleaseEventID(mte3ToMte2Event_); - pipe_->ReleaseEventID(vToSEvent_); pipe_->ReleaseEventID(sToVEvent_); pipe_->ReleaseEventID(sToMte2Event_); vectorEventsAllocated_ = false; @@ -568,9 +566,9 @@ class ChunkKdaFwdFinalizeKernel { SetFlag(mte2ToVEvent_); WaitFlag(mte2ToVEvent_); Cast(coefficients, coefficientTyped, RoundMode::CAST_NONE, static_cast(curT)); - PipeBarrier(); - SetFlag(vToSEvent_); - WaitFlag(vToSEvent_); + // Earlier megakernel stages reuse V_S event IDs. Drain the cast + // before the first scalar coefficient read in this tail row. + PipeBarrier(); Duplicate(dstRow, 0.0f, static_cast(V_)); PipeBarrier(); for (uint64_t j = 0; j < curT; ++j) { @@ -602,9 +600,9 @@ class ChunkKdaFwdFinalizeKernel { SetFlag(mte2ToVEvent_); WaitFlag(mte2ToVEvent_); Cast(coefficients, coefficientTyped, RoundMode::CAST_NONE, static_cast(K_)); - PipeBarrier(); - SetFlag(vToSEvent_); - WaitFlag(vToSEvent_); + // Earlier megakernel stages reuse V_S event IDs. Drain the cast + // before the first scalar coefficient read in this tail row. + PipeBarrier(); Duplicate(dstRow, 0.0f, static_cast(V_)); PipeBarrier(); for (uint64_t d = 0; d < K_; ++d) { @@ -1457,7 +1455,6 @@ class ChunkKdaFwdFinalizeKernel { TEventID mte3ToVEvent_ = 0; TEventID mte2ToMte3Event_ = 0; TEventID mte3ToMte2Event_ = 0; - TEventID vToSEvent_ = 0; TEventID sToVEvent_ = 0; TEventID sToMte2Event_ = 0; bool vectorEventsAllocated_ = false; diff --git a/csrc/chunk_kda_fwd/op_kernel/arch35/chunk_kda_fwd_prepare_arch35.h b/csrc/chunk_kda_fwd/op_kernel/arch35/chunk_kda_fwd_prepare_arch35.h index e9a34c1b1f..852bad8d21 100644 --- a/csrc/chunk_kda_fwd/op_kernel/arch35/chunk_kda_fwd_prepare_arch35.h +++ b/csrc/chunk_kda_fwd/op_kernel/arch35/chunk_kda_fwd_prepare_arch35.h @@ -3910,9 +3910,11 @@ class ChunkKdaFwdPrepareKernel { #if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 if constexpr (!(SAFE_GATE && COMPILE_BT == 64 && COMPILE_K == 128 && COMPILE_V == 128)) { Muls(aqkLocal, aqkLocal, scale_, static_cast(matrixElems)); + PipeBarrier(); } #else Muls(aqkLocal, aqkLocal, scale_, static_cast(matrixElems)); + PipeBarrier(); #endif #if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 if constexpr (!(SAFE_GATE && COMPILE_BT == 64 && COMPILE_K == 128 && COMPILE_V == 128)) { diff --git a/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_common.h b/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_common.h index 29f5f234a8..2a570ab3cb 100644 --- a/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_common.h +++ b/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_common.h @@ -200,12 +200,21 @@ __aicore__ inline void RunFrontEnd( SyncAll(); } GM_ADDR uSeed = (tiling.fusePostWu || tiling.fusePostWuIntoFwdH) ? addresses.u : addresses.uSeed; + GM_ADDR wSeed = addresses.w; +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if (!tiling.fusePostWu && !tiling.fusePostWuIntoFwdH) { + // Arch35 PostWU reads the prepared W seed while producing W. Keep the + // seed in the already-reserved PostWU workspace so output stores can + // never overwrite seed rows that have not been consumed yet. + wSeed = userWorkspace + tiling.postWuScratchOffset; + } +#endif KdaPrepare::RunChunkKdaPrepare( q, k, v, addresses.gk, g, aLog, dtBias, beta, initialState, cuSeqlens, chunkIndices, aqk, akk, addresses.qg, - addresses.qgScaled, addresses.w, uSeed, addresses.kg, + addresses.qgScaled, wSeed, uSeed, addresses.kg, userWorkspace, tiling, pipe, tiling.storeQG); SyncAll(); pipe.Reset(); @@ -213,7 +222,7 @@ __aicore__ inline void RunFrontEnd( if (!tiling.fusePostWu && !tiling.fusePostWuIntoFwdH) { KdaPostWu::RunChunkKdaPostWu( q, k, v, addresses.gk, beta, initialState, cuSeqlens, - chunkIndices, addresses.w, akk, uSeed, + chunkIndices, wSeed, akk, uSeed, addresses.w, addresses.u, addresses.kg, addresses.vNew, userWorkspace, tiling, pipe); SyncAll(); diff --git a/csrc/pytorch_extensions.cpp b/csrc/pytorch_extensions.cpp index 3a971de949..db5465822a 100644 --- a/csrc/pytorch_extensions.cpp +++ b/csrc/pytorch_extensions.cpp @@ -256,7 +256,9 @@ TORCH_LIBRARY_IMPL(npu, PrivateUse1, m) m.impl("sgemmc_shrink", TORCH_FN(sglang::npu_kernel::sgemmc_shrink)); +#ifdef SGL_KERNEL_ENABLE_A3_ONLY_OPS m.impl("compressor", TORCH_FN(sglang::npu_kernel::compressor)); +#endif m.impl("apply_token_bitmask", [](at::Tensor logits, at::Tensor bitmask, const c10::optional &indices) { auto indices_or_empty = indices.has_value() ? *indices : at::empty({0}, logits.options().dtype(at::kInt)); From 6fc4f2064f7217753072689ffd74bcde10d62249 Mon Sep 17 00:00:00 2001 From: chenjz Date: Thu, 3 Sep 2026 15:07:17 +0800 Subject: [PATCH 05/11] fix(npu): register causal_conv1d on A5 --- csrc/CMakeLists.txt | 2 +- csrc/pytorch_extensions.cpp | 49 ++++++++++++++++++------------------- 2 files changed, 25 insertions(+), 26 deletions(-) diff --git a/csrc/CMakeLists.txt b/csrc/CMakeLists.txt index 5389f9d19e..2cbb0d1648 100644 --- a/csrc/CMakeLists.txt +++ b/csrc/CMakeLists.txt @@ -41,6 +41,7 @@ set(OP_SRCS ${PROJECT_OP_SRC_BASE}/compressor/op_host/compressor.cpp ${PROJECT_OP_SRC_BASE}/compressor/op_host/compressor_tiling.cpp ${PROJECT_OP_SRC_BASE}/causal_conv1d_update/op_host/causal_conv1d_update.cpp + ${PROJECT_OP_SRC_BASE}/causal_conv1d/op_host/causal_conv1d.cpp ${PROJECT_OP_SRC_BASE}/apply_token_bitmask/op_host/apply_token_bitmask.cpp ${PROJECT_OP_SRC_BASE}/sparse_attn_sharedkv_metadata/op_host/sparse_attn_sharedkv_metadata.cpp ${PROJECT_OP_SRC_BASE}/sparse_attention_score/op_host/sparse_attention_score.cpp @@ -59,7 +60,6 @@ if(SGL_KERNEL_ENABLE_A3_ONLY_OPS) ${PROJECT_OP_SRC_BASE}/sparse_attn_sharedkv/op_host/sparse_attn_sharedkv_tiling.cpp ${PROJECT_OP_SRC_BASE}/tri_inv/op_host/tri_inv.cpp ${PROJECT_OP_SRC_BASE}/recurrent_gated_delta_rule/op_host/recurrent_gated_delta_rule.cpp - ${PROJECT_OP_SRC_BASE}/causal_conv1d/op_host/causal_conv1d.cpp ${PROJECT_OP_SRC_BASE}/mega_chunk_gdn/op_host/mega_chunk_gdn.cpp ) endif() diff --git a/csrc/pytorch_extensions.cpp b/csrc/pytorch_extensions.cpp index db5465822a..c913ef504e 100644 --- a/csrc/pytorch_extensions.cpp +++ b/csrc/pytorch_extensions.cpp @@ -16,9 +16,7 @@ #include "torch_helper.h" #include "sgl_kenel_npu_ops.h" #include "causal_conv1d_update/op_host/causal_conv1d_update.h" -#ifdef SGL_KERNEL_ENABLE_A3_ONLY_OPS #include "causal_conv1d/op_host/causal_conv1d.h" -#endif namespace { TORCH_LIBRARY_FRAGMENT(npu, m) @@ -105,6 +103,12 @@ TORCH_LIBRARY_FRAGMENT(npu, m) "Tensor conv_state_indices, Tensor? bias=None, Tensor? num_accepted_tokens=None, " "Tensor? query_start_loc=None, bool activation_mode=False, int pad_slot_id=-1) -> Tensor"); + m.def( + "causal_conv1d(Tensor x, Tensor weight, Tensor(a!) conv_states, Tensor? bias=None, " + "Tensor? query_start_loc=None, Tensor? cache_indices=None, Tensor? has_initial_state=None, " + "Tensor? num_accepted_tokens=None, int activation_mode=0, int pad_slot_id=-1, " + "int run_mode=0) -> Tensor"); + #ifdef SGL_KERNEL_ENABLE_A3_ONLY_OPS m.def( "mla_preprocess(Tensor hiddenState, Tensor gamma0, Tensor beta0, Tensor wdqkv, " @@ -176,12 +180,6 @@ TORCH_LIBRARY_FRAGMENT(npu, m) m.def("triangular_inverse(Tensor x) -> Tensor"); - m.def( - "causal_conv1d(Tensor x, Tensor weight, Tensor conv_states, Tensor? bias=None, " - "Tensor? query_start_loc=None, Tensor? cache_indices=None, Tensor? has_initial_state=None, " - "Tensor? num_accepted_tokens=None, int activation_mode=0, int pad_slot_id=-1, " - "int run_mode=0) -> Tensor"); - m.def( "sparse_attn_sharedkv_metadata_host(" "int num_heads_q, int num_heads_kv, int head_dim, " @@ -282,23 +280,6 @@ TORCH_LIBRARY_IMPL(npu, PrivateUse1, m) query_loc_or_empty, activation_mode, pad_slot_id); }); -#ifdef SGL_KERNEL_ENABLE_A3_ONLY_OPS - m.impl("mla_preprocess", TORCH_FN(sglang::npu_kernel::mla_preprocess)); - - m.impl("batch_matmul_transpose", TORCH_FN(sglang::npu_kernel::batch_matmul_transpose)); - - m.impl("recurrent_gated_delta_rule", TORCH_FN(sglang::npu_kernel::recurrent_gated_delta_rule)); - - m.impl("mega_chunk_gdn", TORCH_FN(sglang::npu_kernel::mega_chunk_gdn)); - - m.impl("lightning_indexer", TORCH_FN(sglang::npu_kernel::lightning_indexer)); - - m.impl("sparse_attn_sharedkv", TORCH_FN(sglang::npu_kernel::sparse_attn_sharedkv)); - - m.impl("npu_sparse_attention_score", TORCH_FN(sglang::npu_kernel::sparse_attention_score)); - - m.impl("triangular_inverse", TORCH_FN(sglang::npu_kernel::tri_inv_col_sweep)); - m.impl("causal_conv1d", [](const at::Tensor &x, const at::Tensor &weight, const at::Tensor &conv_states, const c10::optional &bias, const c10::optional &query_start_loc, const c10::optional &cache_indices, @@ -320,6 +301,24 @@ TORCH_LIBRARY_IMPL(npu, PrivateUse1, m) x, weight, bias_or_empty, conv_states, query_start_loc_or_empty, cache_indices_or_empty, has_initial_state_or_empty, num_accepted_tokens_or_empty, activation_mode, pad_slot_id, run_mode); }); + +#ifdef SGL_KERNEL_ENABLE_A3_ONLY_OPS + m.impl("mla_preprocess", TORCH_FN(sglang::npu_kernel::mla_preprocess)); + + m.impl("batch_matmul_transpose", TORCH_FN(sglang::npu_kernel::batch_matmul_transpose)); + + m.impl("recurrent_gated_delta_rule", TORCH_FN(sglang::npu_kernel::recurrent_gated_delta_rule)); + + m.impl("mega_chunk_gdn", TORCH_FN(sglang::npu_kernel::mega_chunk_gdn)); + + m.impl("lightning_indexer", TORCH_FN(sglang::npu_kernel::lightning_indexer)); + + m.impl("sparse_attn_sharedkv", TORCH_FN(sglang::npu_kernel::sparse_attn_sharedkv)); + + m.impl("npu_sparse_attention_score", TORCH_FN(sglang::npu_kernel::sparse_attention_score)); + + m.impl("triangular_inverse", TORCH_FN(sglang::npu_kernel::tri_inv_col_sweep)); + #endif m.impl("chunk_kda_fwd", TORCH_FN(sglang::npu_kernel::chunk_kda_fwd)); From 38967b3c65470ae8e639544d340b58fdd550eb73 Mon Sep 17 00:00:00 2001 From: chenjz Date: Fri, 4 Sep 2026 09:23:24 +0800 Subject: [PATCH 06/11] refactor(chunk_kda_fwd): drop never-instantiated vendored catlass code The vendored catlass copy under op_kernel/catlass/ was imported wholesale. Its two hub headers pull in every implementation they know about, so the chunk_kda_fwd translation unit was parsing ~26k lines of templates that no KDA code path ever instantiates. Auditing every Catlass symbol referenced from arch22/, arch35/, kernel_utils/ and tla/ shows the KDA kernels use exactly: - Gemm::Block::BlockMmadTla with MmadPingpong (catlass) and with MmadPingpongTlaMulti / MmadPingpongTlaPreloadAL1B (kernel_utils/block/, included directly by their users) - Epilogue::Block::BlockEpilogue with EpilogueAtlasKdaFwdHVnew / Update, both specialized under arch22|arch35/fwd_h/epilogue/block/ Everything else -- 39 BlockMmad implementations, 20 catlass epilogues, the AtlasA2-only BlockMmad/BlockGemm/BlockMmadAiv/BlockMmadSparseTla primaries, the block swizzles (included by both FwdH kernels but never referenced), and the helpers reachable only from those files -- is dead. So kda_block_mmad.hpp and kda_block_epilogue.hpp now keep their primary templates and include only what is selected, and the 70 orphaned files are removed. Preprocessed lines entering the compiler for the A5 build drop from 40943 to 28267 (116 files to 83); the A3 build drops to 23720. Verified for both CATLASS_ARCH configurations that no include dangles and that no surviving file references a removed symbol. Co-Authored-By: Claude Opus 5 --- .../gemm/kernel/kda_fwd_h_kernel_arch22.hpp | 1 - .../gemm/kernel/kda_fwd_h_kernel_arch35.hpp | 1 - .../op_kernel/catlass/detail/kda_callback.hpp | 61 - .../block_epilogue_elemwise_no_source.hpp | 185 --- .../block/block_epilogue_fa_rescale_o.hpp | 337 ------ .../block_epilogue_fa_rescale_o_ascend950.hpp | 216 ---- .../block_epilogue_fa_softmax_ascend950.hpp | 405 ------- .../epilogue/block/block_epilogue_fixpipe.hpp | 83 -- .../block/block_epilogue_mla_fd_rescale_o.hpp | 307 ----- .../block/block_epilogue_mla_rescale_o.hpp | 479 -------- .../block_epilogue_mla_tp1_rescale_o.hpp | 341 ------ .../block_epilogue_online_softmax_no_mask.hpp | 877 -------------- .../block_epilogue_per_group_per_block.hpp | 863 -------------- .../block_epilogue_per_token_dequant_tla.hpp | 784 ------------- .../block_epilogue_rescale_o_no_split_row.hpp | 313 ----- ...gue_w4a4_per_token_per_channel_dequant.hpp | 252 ---- .../epilogue/block/kda_block_epilogue.hpp | 27 +- ...kda_block_epilogue_elemwise_one_source.hpp | 185 --- .../block/kda_block_epilogue_fa_softmax.hpp | 403 ------- .../block/kda_block_epilogue_gemm.hpp | 218 ---- .../block/kda_block_epilogue_gemv.hpp | 222 ---- .../block/kda_block_epilogue_mla_softmax.hpp | 441 ------- .../kda_block_epilogue_mla_tp1_softmax.hpp | 507 -------- .../kda_block_epilogue_per_token_dequant.hpp | 683 ----------- .../catlass/epilogue/kda_dispatch_policy.hpp | 143 --- .../catlass/epilogue/tile/kda_tile_cast.hpp | 46 - .../gemm/block/block_mmad_dynamic_aiv.hpp | 505 -------- .../gemm/block/block_mmad_dynamic_common.hpp | 353 ------ ...ad_dynamic_preload_async_with_callback.hpp | 441 ------- .../block_mmad_dynamic_single_core_splitk.hpp | 391 ------- .../gemm/block/block_mmad_dynamic_small.hpp | 284 ----- .../gemm/block/block_mmad_dynamic_streamk.hpp | 370 ------ .../block/block_mmad_fai_pv_head_tail.hpp | 315 ----- .../block/block_mmad_fai_pv_head_tail_tla.hpp | 350 ------ .../gemm/block/block_mmad_fai_pv_normal.hpp | 305 ----- .../block/block_mmad_fai_pv_normal_tla.hpp | 340 ------ .../gemm/block/block_mmad_fai_pv_tla.hpp | 306 ----- .../block/block_mmad_fai_qk_head_tail.hpp | 306 ----- .../block/block_mmad_fai_qk_head_tail_tla.hpp | 326 ------ .../gemm/block/block_mmad_fai_qk_normal.hpp | 337 ------ .../block/block_mmad_fai_qk_normal_tla.hpp | 347 ------ .../gemm/block/block_mmad_fai_qk_tla.hpp | 349 ------ .../block/block_mmad_pingpong_dequant_tla.hpp | 666 ----------- .../block/block_mmad_pingpong_full_loadA.hpp | 347 ------ .../block_mmad_pingpong_full_loadA_tla.hpp | 582 ---------- ..._mmad_pingpong_per_group_per_block_tla.hpp | 320 ------ ...ck_mmad_pingpong_slice_k_with_prologue.hpp | 806 ------------- .../gemm/block/block_mmad_pingpong_tla_v2.hpp | 394 ------- .../block_mmad_pingpong_with_prologue.hpp | 535 --------- ...k_mmad_preload_async_with_callback_tla.hpp | 650 ----------- .../gemm/block/block_mmad_preload_tla.hpp | 445 ------- .../block/block_mmad_single_core_splitk.hpp | 391 ------- .../catlass/gemm/block/block_mmad_small.hpp | 319 ------ .../gemm/block/block_mmad_sparse_tla.hpp | 397 ------- .../catlass/gemm/block/block_mmad_streamk.hpp | 406 ------- ...mad_w4a4_per_token_per_channel_dequant.hpp | 494 -------- .../catlass/gemm/block/kda_block_mmad.hpp | 108 +- .../gemm/block/kda_block_mmad_fa_pv.hpp | 193 ---- .../gemm/block/kda_block_mmad_fa_qk.hpp | 199 ---- .../gemm/block/kda_block_mmad_gemm.hpp | 348 ------ .../gemm/block/kda_block_mmad_mla_pv.hpp | 207 ---- .../block/kda_block_mmad_mla_pv_tp1_spec.hpp | 221 ---- .../gemm/block/kda_block_mmad_mla_qk.hpp | 241 ---- .../block/kda_block_mmad_mla_qk_tp1_spec.hpp | 259 ----- .../gemm/block/kda_block_mmad_pingpong.hpp | 362 ------ .../block/kda_block_mmad_pingpong_bias.hpp | 421 ------- .../gemm/block/kda_block_mmad_preload.hpp | 402 ------- .../block/kda_block_mmad_preload_async.hpp | 448 -------- ...block_mmad_preload_async_with_callback.hpp | 449 -------- .../catlass/gemm/block/kda_block_swizzle.hpp | 1018 ----------------- .../catlass/gemm/tile/kda_tile_muls.hpp | 49 - .../catlass/gemm/tile/kda_tile_traits.hpp | 48 - .../catlass/gemv/kda_gemv_helper.hpp | 122 -- .../op_kernel/catlass/kda_gemv_coord.hpp | 108 -- 74 files changed, 8 insertions(+), 25952 deletions(-) delete mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/detail/kda_callback.hpp delete mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_elemwise_no_source.hpp delete mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_fa_rescale_o.hpp delete mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_fa_rescale_o_ascend950.hpp delete mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_fa_softmax_ascend950.hpp delete mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_fixpipe.hpp delete mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_mla_fd_rescale_o.hpp delete mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_mla_rescale_o.hpp delete mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_mla_tp1_rescale_o.hpp delete mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_online_softmax_no_mask.hpp delete mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_per_group_per_block.hpp delete mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_per_token_dequant_tla.hpp delete mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_rescale_o_no_split_row.hpp delete mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_w4a4_per_token_per_channel_dequant.hpp delete mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/kda_block_epilogue_elemwise_one_source.hpp delete mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/kda_block_epilogue_fa_softmax.hpp delete mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/kda_block_epilogue_gemm.hpp delete mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/kda_block_epilogue_gemv.hpp delete mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/kda_block_epilogue_mla_softmax.hpp delete mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/kda_block_epilogue_mla_tp1_softmax.hpp delete mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/kda_block_epilogue_per_token_dequant.hpp delete mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/kda_dispatch_policy.hpp delete mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/tile/kda_tile_cast.hpp delete mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_dynamic_aiv.hpp delete mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_dynamic_common.hpp delete mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_dynamic_preload_async_with_callback.hpp delete mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_dynamic_single_core_splitk.hpp delete mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_dynamic_small.hpp delete mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_dynamic_streamk.hpp delete mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_pv_head_tail.hpp delete mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_pv_head_tail_tla.hpp delete mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_pv_normal.hpp delete mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_pv_normal_tla.hpp delete mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_pv_tla.hpp delete mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_qk_head_tail.hpp delete mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_qk_head_tail_tla.hpp delete mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_qk_normal.hpp delete mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_qk_normal_tla.hpp delete mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_qk_tla.hpp delete mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_pingpong_dequant_tla.hpp delete mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_pingpong_full_loadA.hpp delete mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_pingpong_full_loadA_tla.hpp delete mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_pingpong_per_group_per_block_tla.hpp delete mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_pingpong_slice_k_with_prologue.hpp delete mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_pingpong_tla_v2.hpp delete mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_pingpong_with_prologue.hpp delete mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_preload_async_with_callback_tla.hpp delete mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_preload_tla.hpp delete mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_single_core_splitk.hpp delete mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_small.hpp delete mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_sparse_tla.hpp delete mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_streamk.hpp delete mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_w4a4_per_token_per_channel_dequant.hpp delete mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_fa_pv.hpp delete mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_fa_qk.hpp delete mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_gemm.hpp delete mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_mla_pv.hpp delete mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_mla_pv_tp1_spec.hpp delete mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_mla_qk.hpp delete mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_mla_qk_tp1_spec.hpp delete mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_pingpong.hpp delete mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_pingpong_bias.hpp delete mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_preload.hpp delete mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_preload_async.hpp delete mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_preload_async_with_callback.hpp delete mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_swizzle.hpp delete mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_tile_muls.hpp delete mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_tile_traits.hpp delete mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/gemv/kda_gemv_helper.hpp delete mode 100644 csrc/chunk_kda_fwd/op_kernel/catlass/kda_gemv_coord.hpp diff --git a/csrc/chunk_kda_fwd/op_kernel/arch22/fwd_h/gemm/kernel/kda_fwd_h_kernel_arch22.hpp b/csrc/chunk_kda_fwd/op_kernel/arch22/fwd_h/gemm/kernel/kda_fwd_h_kernel_arch22.hpp index fdc9076ab1..3cea7b6167 100644 --- a/csrc/chunk_kda_fwd/op_kernel/arch22/fwd_h/gemm/kernel/kda_fwd_h_kernel_arch22.hpp +++ b/csrc/chunk_kda_fwd/op_kernel/arch22/fwd_h/gemm/kernel/kda_fwd_h_kernel_arch22.hpp @@ -20,7 +20,6 @@ #include "../../epilogue/block/block_epilogue_kda_fwdh_vnew_arch22.hpp" #include "catlass/gemm/block/kda_block_mmad.hpp" #include "../../../../kernel_utils/block/block_mmad_pingpong_tla_multi.hpp" -#include "catlass/gemm/block/kda_block_swizzle.hpp" #include "../block/block_scheduler_kda_fwd_h_arch22.hpp" #include "catlass/gemm/kda_gemm_dispatch_policy.hpp" #include "catlass/gemm/kda_gemm_type.hpp" diff --git a/csrc/chunk_kda_fwd/op_kernel/arch35/fwd_h/gemm/kernel/kda_fwd_h_kernel_arch35.hpp b/csrc/chunk_kda_fwd/op_kernel/arch35/fwd_h/gemm/kernel/kda_fwd_h_kernel_arch35.hpp index 3239a21bd0..6039f9994e 100644 --- a/csrc/chunk_kda_fwd/op_kernel/arch35/fwd_h/gemm/kernel/kda_fwd_h_kernel_arch35.hpp +++ b/csrc/chunk_kda_fwd/op_kernel/arch35/fwd_h/gemm/kernel/kda_fwd_h_kernel_arch35.hpp @@ -23,7 +23,6 @@ #include "../../../../kernel_utils/block/block_mmad_pingpong_tla_kernel_utils.hpp" #include "../../../../kernel_utils/block/block_mmad_pingpong_tla_multi.hpp" #include "../../../../kernel_utils/block/block_mmad_pingpong_tla_preloadA_l1B.hpp" -#include "catlass/gemm/block/kda_block_swizzle.hpp" #include "catlass/gemm/kda_gemm_dispatch_policy.hpp" #include "catlass/gemm/kda_gemm_type.hpp" #include "catlass/layout/kda_layout.hpp" diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/detail/kda_callback.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/detail/kda_callback.hpp deleted file mode 100644 index 91456611c9..0000000000 --- a/csrc/chunk_kda_fwd/op_kernel/catlass/detail/kda_callback.hpp +++ /dev/null @@ -1,61 +0,0 @@ -/** - * Copyright (c) 2025 Huawei Technologies Co., Ltd. - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of - * CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ - -#ifndef CATLASS_DETAIL_KDA_CALLBACK_HPP -#define CATLASS_DETAIL_KDA_CALLBACK_HPP - -#include "catlass/detail/kda_macros.hpp" - -/// @brief Callback is an alternative to std::function, providing a general carrier -/// of callable structure with no parameters and no return value. Compared with function pointers -/// of type void (*)(), Callback can carry lambda expressions with captures, and does not need to -/// pay attention to the captured content. It should be noted that Callback itself does not store -/// the callable structure it carries like std::function, so it is necessary to ensure -/// that it is used within the life cycle of the callable structure. -struct Callback { - void const *func{nullptr}; - void (*caller)(void const *){nullptr}; - - Callback() = default; - - CATLASS_DEVICE - void operator()() const - { - if (func) { - caller(func); - } - } - - CATLASS_DEVICE - operator bool() const - { - return func != nullptr; - } -}; - -template -CATLASS_DEVICE void FuncWrapper(void const *func) -{ - (*static_cast(func))(); -} - -// Use this to make a callback -template -CATLASS_DEVICE - Callback - MakeCallback(Func *func) -{ - Callback callback; - callback.func = func; - callback.caller = &FuncWrapper; - return callback; -} - -#endif // CATLASS_DETAIL_KDA_CALLBACK_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_elemwise_no_source.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_elemwise_no_source.hpp deleted file mode 100644 index 06fe13078a..0000000000 --- a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_elemwise_no_source.hpp +++ /dev/null @@ -1,185 +0,0 @@ -/** - * Copyright (c) 2025 Huawei Technologies Co., Ltd. - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of - * CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ - -#ifndef CATLASS_EPILOGUE_BLOCK_EPILOGUE_ELEMWISE_NO_SOURCE_HPP -#define CATLASS_EPILOGUE_BLOCK_EPILOGUE_ELEMWISE_NO_SOURCE_HPP - -#include "catlass/kda_catlass.hpp" -#include "catlass/arch/kda_resource.hpp" -#include "catlass/epilogue/kda_dispatch_policy.hpp" -#include "catlass/kda_gemm_coord.hpp" -#include "catlass/kda_matrix_coord.hpp" -#include "catlass/layout/kda_layout.hpp" -#include "catlass/epilogue/tile/kda_tile_cast.hpp" - -namespace Catlass::Epilogue::Block { -template < - class CType_, - class DType_, - class TileElemWiseEpilogue_, - class TileCopy_> -class BlockEpilogue< - EpilogueAtlasA2ElemWiseNoSource, - CType_, - DType_, - TileElemWiseEpilogue_, - TileCopy_> { -public: - // Type aliases - using DispatchPolicy = EpilogueAtlasA2ElemWiseNoSource; - using ArchTag = typename DispatchPolicy::ArchTag; - using ElementC = typename CType_::Element; - using LayoutC = typename CType_::Layout; - - using ElementD = typename DType_::Element; - using LayoutD = typename DType_::Layout; - using TileElemWiseEpilogue = TileElemWiseEpilogue_; - using CopyGmToUbC = typename TileCopy_::CopyGmToUbC; - using CopyUbToGmD = typename TileCopy_::CopyUbToGmD; - - static constexpr uint32_t COMPUTE_LENGTH = TileElemWiseEpilogue::COMPUTE_LENGTH; - static constexpr uint32_t OPERANDS_NUM = DispatchPolicy::OPERANDS_NUM; - - using ElementCompute = ElementC; - - using LayoutComputeInUb = layout::RowMajor; - - // Check the element type of C and D - static_assert(std::is_same_v, - "Element type of C must be float"); - // Check the layout type of C and D - static_assert(std::is_same_v && std::is_same_v, - "Layout type of C, D must be RowMajor"); - - // Check if ArchTag is matched - static_assert(std::is_same_v, - "Tile epilogue's ArchTag mismatch"); - // Check if compute length is valid - static_assert(COMPUTE_LENGTH * (OPERANDS_NUM * sizeof(ElementC) + sizeof(ElementD)) <= ArchTag::UB_SIZE, - "UB out of bounds"); - - // Epilogue params definition - struct Params { - GM_ADDR ptrC; - LayoutC layoutC; - GM_ADDR ptrD; - LayoutD layoutD; - - CATLASS_HOST_DEVICE - Params() {} - - CATLASS_HOST_DEVICE - Params(GM_ADDR ptrC_, LayoutD const &layoutC_, GM_ADDR ptrD_, LayoutD const &layoutD_) - : ptrC(ptrC_), - layoutC(layoutC_), - ptrD(ptrD_), - layoutD(layoutD_) - {} - }; - - CATLASS_DEVICE - BlockEpilogue(Arch::Resource &resource, Params const ¶ms) - : params(params) - { - ubC = resource.ubBuf.template GetBufferByByte(0); - ubD = resource.ubBuf.template GetBufferByByte(COMPUTE_LENGTH * sizeof(ElementC)); - if constexpr (!std::is_same_v) { - ubCompute = resource.ubBuf.template GetBufferByByte( - COMPUTE_LENGTH * (sizeof(ElementC) + sizeof(ElementD))); - } else { - // When ElementCompute is same as ElementD, it is not necessary to do Cast from ubCompute to ubD. - // Use the same buffer. - ubCompute = ubD; - } - AscendC::SetFlag(EVENT_ID0); - AscendC::SetFlag(EVENT_ID0); - } - - CATLASS_DEVICE - ~BlockEpilogue() - { - AscendC::WaitFlag(EVENT_ID0); - AscendC::WaitFlag(EVENT_ID0); - } - - CATLASS_DEVICE - void operator()( - GemmCoord const &blockShapeMNK, - GemmCoord const &blockCoordMNK, - GemmCoord const &actualBlockShapeMNK, - AscendC::GlobalTensor const &gmBlockC, - LayoutD const &layoutBlockC) - { - // Calculate the offset of the current block - MatrixCoord blockShape = blockShapeMNK.GetCoordMN(); - MatrixCoord blockCoord = blockCoordMNK.GetCoordMN(); - MatrixCoord actualBlockShape = actualBlockShapeMNK.GetCoordMN(); - MatrixCoord blockOffset = blockCoord * blockShape; - - // Calculate the offset and the shape of the current subblock - MatrixCoord subblockShape{ - CeilDiv(actualBlockShape.row(), static_cast(AscendC::GetSubBlockNum())), - actualBlockShape.column()}; - - MatrixCoord subblockCoord{AscendC::GetSubBlockIdx(), 0}; - MatrixCoord actualSubblockShape = MatrixCoord::Min( - subblockShape, actualBlockShape - subblockCoord * subblockShape); - MatrixCoord subblockOffset = subblockCoord * subblockShape; - - // Get the data and layout of C - auto gmSubblockC = gmBlockC[layoutBlockC.GetOffset(subblockOffset)]; - auto layoutSubblockC = layoutBlockC.GetTileLayout(actualSubblockShape); - - // Get the data and layout of D - AscendC::GlobalTensor gmD; - gmD.SetGlobalBuffer(reinterpret_cast<__gm__ ElementD *>(params.ptrD)); - auto gmSubblockD = gmD[params.layoutD.GetOffset(blockOffset + subblockOffset)]; - auto layoutSubblockD = params.layoutD.GetTileLayout(actualSubblockShape); - - // Get the layout on UB - auto ubTileStride = MakeCoord(static_cast(blockShape.column()), 1L); - LayoutComputeInUb layoutComputeInUb{actualSubblockShape, ubTileStride}; - LayoutComputeInUb layoutComputeOutUb{actualSubblockShape, ubTileStride}; - // Copy the data of C - AscendC::WaitFlag(EVENT_ID0); - copyGmToUbC(ubC, gmSubblockC, layoutComputeInUb, layoutSubblockC); - AscendC::SetFlag(EVENT_ID0); - - // Perform epilogue calculation - AscendC::WaitFlag(EVENT_ID0); - AscendC::WaitFlag(EVENT_ID0); - tileEpilogue(ubCompute, ubC); - if constexpr (!std::is_same_v) { - AscendC::Cast(ubD, ubCompute, AscendC::RoundMode::CAST_NONE, COMPUTE_LENGTH); - } - AscendC::SetFlag(EVENT_ID0); - AscendC::SetFlag(EVENT_ID0); - - // Copy the data of D - AscendC::WaitFlag(EVENT_ID0); - copyUbToGmD(gmSubblockD, ubD, layoutSubblockD, layoutComputeOutUb); - AscendC::SetFlag(EVENT_ID0); - } - -private: - Params params; - - AscendC::LocalTensor ubC; - AscendC::LocalTensor ubCompute; - AscendC::LocalTensor ubD; - - TileElemWiseEpilogue tileEpilogue; - CopyGmToUbC copyGmToUbC; - CopyUbToGmD copyUbToGmD; -}; - -} // namespace Catlass::Epilogue::Block - -#endif // CATLASS_EPILOGUE_BLOCK_EPILOGUE_ELEMWISE_NO_SOURCE_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_fa_rescale_o.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_fa_rescale_o.hpp deleted file mode 100644 index 5f7aba7c2a..0000000000 --- a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_fa_rescale_o.hpp +++ /dev/null @@ -1,337 +0,0 @@ -/** - * Copyright (c) 2025 Huawei Technologies Co., Ltd. - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of - * CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ - -#ifndef CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_FA_RESCALE_O_HPP -#define CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_FA_RESCALE_O_HPP - -#include "catlass/kda_catlass.hpp" -#include "catlass/arch/kda_resource.hpp" -#include "catlass/epilogue/kda_dispatch_policy.hpp" -#include "catlass/kda_gemm_coord.hpp" -#include "catlass/kda_matrix_coord.hpp" -#include "catlass/epilogue/tile/kda_epilogue_tile_copy.hpp" - -namespace Catlass::Epilogue::Block { - -template < - class OutputType_, - class InputType_> -class BlockEpilogue< - EpilogueAtlasA2FARescaleO, - OutputType_, - InputType_> { -public: - // Type aliases - using DispatchPolicy = EpilogueAtlasA2FARescaleO; - using ArchTag = typename DispatchPolicy::ArchTag; - using ElementOutput = typename OutputType_::Element; - using ElementInput = typename InputType_::Element; - - using LayoutOutput = typename OutputType_::Layout; - using LayoutInput = typename InputType_::Layout; - - using CopyGmToUbInput = Tile::CopyGm2Ub; - using CopyUbToGmOutput = Tile::CopyUb2Gm; - - static constexpr uint32_t HALF_ELENUM_PER_BLK = 16; - static constexpr uint32_t FLOAT_ELENUM_PER_BLK = 8; - static constexpr uint32_t HALF_ELENUM_PER_VECCALC = 128; - static constexpr uint32_t FLOAT_ELENUM_PER_VECCALC = 64; - static constexpr uint32_t UB_TILE_SIZE = 16384; // 64 * 128 * 2B - static constexpr uint32_t UB_LINE_SIZE = 512; // 128 * 2 * 2B - static constexpr uint32_t HALF_ELENUM_PER_LINE = 256; // 128 * 2 - static constexpr uint32_t FLOAT_ELENUM_PER_LINE = 128; // 128 - static constexpr uint32_t MULTIPLIER = 2; - - CATLASS_DEVICE - BlockEpilogue(Arch::Resource &resource) - { - constexpr uint32_t LO_UB_TENSOR_OFFSET = 5 * UB_TILE_SIZE; - constexpr uint32_t DM_UB_TENSOR_OFFSET = 7 * UB_TILE_SIZE + 4 * UB_LINE_SIZE; - constexpr uint32_t LL_UB_TENSOR_OFFSET = 7 * UB_TILE_SIZE + 6 * UB_LINE_SIZE; - constexpr uint32_t GL_UB_TENSOR_OFFSET = 7 * UB_TILE_SIZE + 10 * UB_LINE_SIZE; - constexpr uint32_t TV_UB_TENSOR_OFFSET = 7 * UB_TILE_SIZE + 11 * UB_LINE_SIZE; - constexpr uint32_t GO_UB_TENSOR_OFFSET = 8 * UB_TILE_SIZE; - - loUbTensor = resource.ubBuf.template GetBufferByByte(LO_UB_TENSOR_OFFSET); - dmUbTensor = resource.ubBuf.template GetBufferByByte(DM_UB_TENSOR_OFFSET); - llUbTensor = resource.ubBuf.template GetBufferByByte(LL_UB_TENSOR_OFFSET); - glUbTensor = resource.ubBuf.template GetBufferByByte(GL_UB_TENSOR_OFFSET); - tvUbTensor = resource.ubBuf.template GetBufferByByte(TV_UB_TENSOR_OFFSET); - goUbTensor = resource.ubBuf.template GetBufferByByte(GO_UB_TENSOR_OFFSET); - - AscendC::SetFlag(EVENT_ID1); - AscendC::SetFlag(EVENT_ID0); - } - - CATLASS_DEVICE - ~BlockEpilogue() - { - AscendC::WaitFlag(EVENT_ID1); - AscendC::WaitFlag(EVENT_ID0); - } - - CATLASS_DEVICE - void SetMask(int32_t len) - { - const int32_t MAX_MASK_LEN = 128; - const int32_t HALF_MASK_LEN = 64; - if (len >= MAX_MASK_LEN) { - AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); - return; - } - int32_t highMask = len - HALF_MASK_LEN > 0 ? len - HALF_MASK_LEN : 0; - int32_t lowMask = len - HALF_MASK_LEN >= 0 ? HALF_MASK_LEN : len; - if (len < HALF_MASK_LEN) { - AscendC::SetVectorMask(0x0, ((uint64_t)1 << lowMask) - 1); - } else { - AscendC::SetVectorMask(((uint64_t)1 << highMask) - 1, 0xffffffffffffffff); - } - } - - CATLASS_DEVICE - void subCoreCompute( - AscendC::GlobalTensor gOutput, - AscendC::GlobalTensor gInput, - const LayoutOutput &layoutOutput, - const LayoutInput &layoutInput, - uint32_t nIdx, uint32_t isLast) - { - uint32_t subM = layoutInput.shape(0); - uint32_t k = layoutInput.shape(1); - uint32_t kRound = layoutInput.stride(0); - uint32_t strideQO = layoutOutput.stride(0); - uint32_t subMAligned128 = (subM + HALF_ELENUM_PER_VECCALC - 1) / HALF_ELENUM_PER_VECCALC; - uint32_t subMAligned64 = (subM + FLOAT_ELENUM_PER_VECCALC - 1) / FLOAT_ELENUM_PER_VECCALC; - uint32_t subMRound = (subM + HALF_ELENUM_PER_BLK - 1) / HALF_ELENUM_PER_BLK * HALF_ELENUM_PER_BLK; - - // Get the layout on UB - LayoutInput layoutInUb(subM, k, kRound); - - if (subM > 0) { - AscendC::WaitFlag(EVENT_ID1); - // Copy O - copyGmToUbInput(loUbTensor, gInput, layoutInUb, layoutInput); - AscendC::SetFlag(EVENT_ID0); - AscendC::WaitFlag(EVENT_ID0); - // 更新 L 和 O - if (nIdx != 0) { - // dm32 = castfp16to32(dm) - AscendC::Cast( - tvUbTensor, - dmUbTensor[nIdx % MULTIPLIER * HALF_ELENUM_PER_LINE], - AscendC::RoundMode::CAST_NONE, - (uint64_t)0, - subMAligned64, - AscendC::UnaryRepeatParams(1, 1, 8, 4)); - AscendC::PipeBarrier(); - // dm32_block = brcb(dm32) - AscendC::Brcb( - tvUbTensor.ReinterpretCast()[HALF_ELENUM_PER_VECCALC], - tvUbTensor.ReinterpretCast(), - subMRound / FLOAT_ELENUM_PER_BLK, - AscendC::BrcbRepeatParams(1, 8)); - AscendC::PipeBarrier(); - // dm32 = exp(dm32) - AscendC::Exp( - tvUbTensor, - tvUbTensor, - (uint64_t)0, - subMAligned64, - AscendC::UnaryRepeatParams(1, 1, 8, 8)); - AscendC::PipeBarrier(); - // gl = dm * gl - AscendC::Mul( - glUbTensor, - tvUbTensor, - glUbTensor, - (uint64_t)0, - subMAligned64, - AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); - AscendC::PipeBarrier(); - // gl = ll + gl - AscendC::Add( - glUbTensor, - glUbTensor, - llUbTensor[nIdx % MULTIPLIER * FLOAT_ELENUM_PER_LINE], - (uint64_t)0, - subMAligned64, - AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); - AscendC::PipeBarrier(); - // dm32_block = exp(dm32_block) - AscendC::Exp( - tvUbTensor[HALF_ELENUM_PER_VECCALC], - tvUbTensor[HALF_ELENUM_PER_VECCALC], - (uint64_t)0, - (subM * FLOAT_ELENUM_PER_BLK + FLOAT_ELENUM_PER_VECCALC - 1) / FLOAT_ELENUM_PER_VECCALC, - AscendC::UnaryRepeatParams(1, 1, 8, 8)); - AscendC::PipeBarrier(); - if (goFlag == 1) { - AscendC::WaitFlag(EVENT_ID0); - goFlag = 0; - } - // go = go * dm32_block - for (uint32_t vmulIdx = 0; vmulIdx < k / FLOAT_ELENUM_PER_VECCALC; vmulIdx++) { - AscendC::Mul( - goUbTensor[vmulIdx * FLOAT_ELENUM_PER_VECCALC], - goUbTensor[vmulIdx * FLOAT_ELENUM_PER_VECCALC], - tvUbTensor[HALF_ELENUM_PER_VECCALC], - (uint64_t)0, - subM, - AscendC::BinaryRepeatParams( - 1, 1, 0, kRound / FLOAT_ELENUM_PER_BLK, kRound / FLOAT_ELENUM_PER_BLK, 1)); - AscendC::PipeBarrier(); - } - if (k % FLOAT_ELENUM_PER_VECCALC > 0) { - SetMask(k % FLOAT_ELENUM_PER_VECCALC); - AscendC::Mul( - goUbTensor[k / FLOAT_ELENUM_PER_VECCALC * FLOAT_ELENUM_PER_VECCALC], - goUbTensor[k / FLOAT_ELENUM_PER_VECCALC * FLOAT_ELENUM_PER_VECCALC], - tvUbTensor[HALF_ELENUM_PER_VECCALC], - (uint64_t)0, - subM, - AscendC::BinaryRepeatParams( - 1, 1, 0, kRound / FLOAT_ELENUM_PER_BLK, kRound / FLOAT_ELENUM_PER_BLK, 1)); - AscendC::PipeBarrier(); - AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); - } - // go = lo + go - AscendC::Add( - goUbTensor, - goUbTensor, - loUbTensor, - (uint64_t)0, - (subM * kRound + FLOAT_ELENUM_PER_VECCALC - 1) / FLOAT_ELENUM_PER_VECCALC, - AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); - AscendC::PipeBarrier(); - } else { - // gl = ll - AscendC::DataCopy(glUbTensor, llUbTensor[nIdx % MULTIPLIER * FLOAT_ELENUM_PER_LINE], - AscendC::DataCopyParams(1, subMRound / FLOAT_ELENUM_PER_BLK, 0, 0)); - AscendC::PipeBarrier(); - if (goFlag == 1) { - AscendC::WaitFlag(EVENT_ID0); - goFlag = 0; - } - AscendC::DataCopy(goUbTensor, loUbTensor, - AscendC::DataCopyParams(1, subM * kRound / FLOAT_ELENUM_PER_BLK, 0, 0)); - AscendC::PipeBarrier(); - } - AscendC::SetFlag(EVENT_ID1); - if (isLast) { - // gl = castfp32to16(gl) - AscendC::Cast( - glUbTensor.ReinterpretCast(), - glUbTensor, - AscendC::RoundMode::CAST_NONE, - (uint64_t)0, - subMAligned64, - AscendC::UnaryRepeatParams(1, 1, 4, 8)); - AscendC::PipeBarrier(); - // go = castfp32to16(go) - AscendC::Cast( - goUbTensor.ReinterpretCast(), - goUbTensor, - AscendC::RoundMode::CAST_NONE, - (uint64_t)0, - (subM * kRound + FLOAT_ELENUM_PER_VECCALC - 1) / FLOAT_ELENUM_PER_VECCALC, - AscendC::UnaryRepeatParams(1, 1, 4, 8)); - AscendC::PipeBarrier(); - // gl_block = brcb(gl) - AscendC::Brcb( - tvUbTensor.ReinterpretCast(), - glUbTensor.ReinterpretCast(), - subMRound / FLOAT_ELENUM_PER_BLK, - AscendC::BrcbRepeatParams(1, 8)); - AscendC::PipeBarrier(); - // go = go / gl_block - for (uint32_t vdivIdx = 0; vdivIdx < k / HALF_ELENUM_PER_VECCALC; vdivIdx++) { - AscendC::Div( - goUbTensor.ReinterpretCast()[vdivIdx * HALF_ELENUM_PER_VECCALC], - goUbTensor.ReinterpretCast()[vdivIdx * HALF_ELENUM_PER_VECCALC], - tvUbTensor.ReinterpretCast(), - (uint64_t)0, - subM, - AscendC::BinaryRepeatParams( - 1, 1, 0, kRound / HALF_ELENUM_PER_BLK, kRound / HALF_ELENUM_PER_BLK, 1)); - } - if (k % HALF_ELENUM_PER_VECCALC > 0) { - SetMask(k % HALF_ELENUM_PER_VECCALC); - AscendC::Div( - goUbTensor.ReinterpretCast()[k / HALF_ELENUM_PER_VECCALC * HALF_ELENUM_PER_VECCALC], - goUbTensor.ReinterpretCast()[k / HALF_ELENUM_PER_VECCALC * HALF_ELENUM_PER_VECCALC], - tvUbTensor.ReinterpretCast(), - (uint64_t)0, - subM, - AscendC::BinaryRepeatParams( - 1, 1, 0, kRound / HALF_ELENUM_PER_BLK, kRound / HALF_ELENUM_PER_BLK, 1)); - AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); - } - // copy O to GM - AscendC::SetFlag(EVENT_ID1); - AscendC::WaitFlag(EVENT_ID1); - copyUbToGmOutput(gOutput, goUbTensor.ReinterpretCast(), layoutOutput, layoutInUb); - if (goFlag == 0) { - AscendC::SetFlag(EVENT_ID0); - goFlag = 1; - } - } - } - } - - CATLASS_DEVICE - void operator()( - AscendC::GlobalTensor gOutput, - AscendC::GlobalTensor gInput, - const LayoutOutput &layoutOutput, - const LayoutInput &layoutInput, - GemmCoord actualBlockShape, - uint32_t nIdx, uint32_t isLast) - { - uint32_t mActual = actualBlockShape.m(); - uint32_t nActual = actualBlockShape.n(); - - uint32_t subBlockIdx = AscendC::GetSubBlockIdx(); - uint32_t subBlockNum = AscendC::GetSubBlockNum(); - - uint32_t mActualPerSubBlock = CeilDiv(mActual, subBlockNum); - uint32_t mActualThisSubBlock = (subBlockIdx == 0) ? mActualPerSubBlock : (mActual - mActualPerSubBlock); - uint32_t mOffset = subBlockIdx * mActualPerSubBlock; - uint32_t nOffset = 0; - - int64_t offsetOutput = layoutOutput.GetOffset(MatrixCoord(mOffset, nOffset)); - auto gOutputThisSubBlock = gOutput[offsetOutput]; - auto layoutOutputThisSubBlock = layoutOutput.GetTileLayout(MatrixCoord(mActualThisSubBlock, nActual)); - - int64_t offsetInput = layoutInput.GetOffset(MatrixCoord(mOffset, nOffset)); - auto gInputThisSubBlock = gInput[offsetInput]; - auto layoutInputThisSubBlock = layoutInput.GetTileLayout(MatrixCoord(mActualThisSubBlock, nActual)); - - subCoreCompute(gOutputThisSubBlock, gInputThisSubBlock, - layoutOutputThisSubBlock, layoutInputThisSubBlock, - nIdx, isLast); - } - -private: - uint32_t goFlag = 1; - AscendC::LocalTensor loUbTensor; - AscendC::LocalTensor dmUbTensor; - AscendC::LocalTensor llUbTensor; - AscendC::LocalTensor glUbTensor; - AscendC::LocalTensor tvUbTensor; - AscendC::LocalTensor goUbTensor; - - CopyGmToUbInput copyGmToUbInput; - CopyUbToGmOutput copyUbToGmOutput; -}; - -} // namespace Catlass::Epilogue::Block - -#endif // CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_FA_RESCALE_O_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_fa_rescale_o_ascend950.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_fa_rescale_o_ascend950.hpp deleted file mode 100644 index 954f6f4c56..0000000000 --- a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_fa_rescale_o_ascend950.hpp +++ /dev/null @@ -1,216 +0,0 @@ -/** - * Copyright (c) 2026 Huawei Technologies Co., Ltd. - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of - * CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ - -#ifndef CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_RESCALE_O_A5 -#define CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_RESCALE_O_A5 - -#include "catlass/kda_catlass.hpp" -#include "catlass/arch/kda_resource.hpp" -#include "catlass/epilogue/kda_dispatch_policy.hpp" -#include "catlass/kda_gemm_coord.hpp" -#include "catlass/kda_matrix_coord.hpp" -#include "tla/kda_tensor.hpp" -#include "tla/kda_tla_layout.hpp" - -namespace Catlass::Epilogue::Block { - -template < - class L1TileShape_, - class OType_, - class OTmpType_> -class BlockEpilogue< - EpilogueAscend950FARescaleO, - L1TileShape_, - OType_, - OTmpType_> { -public: - using DispatchPolicy = EpilogueAscend950FARescaleO; - using ArchTag = typename DispatchPolicy::ArchTag; - - using ElementO = typename OType_::Element; - using ElementOTmp = typename OTmpType_::Element; - using LayoutTagO = typename OType_::Layout; - using LayoutTagOTmp = typename OTmpType_::Layout; - using L1TileShape = L1TileShape_; - - static constexpr uint32_t S1_BASE_SIZE = tla::get<0>(L1TileShape{}); - static constexpr uint32_t D_BASE_SIZE = tla::get<2>(L1TileShape{}); - static constexpr uint32_t VEC2_UB_SIZE = S1_BASE_SIZE / 2 * D_BASE_SIZE * sizeof(ElementOTmp); - - CATLASS_DEVICE - BlockEpilogue(Arch::Resource &resource, uint32_t &ubBufAddrStart) - { - vf2OutUb = resource.ubBuf.template GetBufferByByte(ubBufAddrStart); - ubBufAddrStart += VEC2_UB_SIZE; - AscendC::SetFlag(eventOMTE3V); - } - - CATLASS_DEVICE - ~BlockEpilogue() - { - AscendC::WaitFlag(eventOMTE3V); - } - - template - CATLASS_DEVICE void operator()( - TensorDst &attenOutGm, - const AscendC::LocalTensor &expMaxUb, - const AscendC::LocalTensor &sumUb, - TensorSrc &bmm2Res, - bool isFirstLoop, - bool isLastUpdate, - uint64_t MM2_RES_INTRA_EVENT) - { - AscendC::WaitFlag(eventOMTE3V); - uint32_t m = tla::get<0>(bmm2Res.shape()); - uint32_t n = tla::get<1>(bmm2Res.shape()); - constexpr int16_t vlSize = static_cast(AscendC::GetVecLen() / sizeof(ElementOTmp)); - int16_t nLoops = AscendC::CeilDivision(n, vlSize) - 1; - uint32_t tailN = (n - 1) % vlSize + 1; - - __ubuf__ float *vec2ResUbAddr = (__ubuf__ ElementOTmp *)vf2OutUb.GetPhyAddr(); - __ubuf__ float *bmm2UbAddr = (__ubuf__ ElementOTmp *)bmm2Res.data().GetPhyAddr(); - __ubuf__ float *expMaxUbAddr = (__ubuf__ ElementOTmp *)expMaxUb.GetPhyAddr(); - __ubuf__ float *sumUbAddr = (__ubuf__ ElementOTmp *)sumUb.GetPhyAddr(); - - if (isFirstLoop) { - DataCopy(vf2OutUb, bmm2Res.data(), m * n); - } else if (!isLastUpdate) { - FlashUpdateNew(vec2ResUbAddr, bmm2UbAddr, expMaxUbAddr, m, nLoops, tailN); - } else { - FlashUpdateLastNew( - vec2ResUbAddr, bmm2UbAddr, expMaxUbAddr, sumUbAddr, m, nLoops, tailN); - } - AscendC::CrossCoreSetFlag(MM2_RES_INTRA_EVENT); - - if (isFirstLoop && isLastUpdate) { - LastDivNew(vec2ResUbAddr, bmm2UbAddr, sumUbAddr, m, nLoops, tailN); - } - if (isLastUpdate) { - AscendC::LocalTensor attenOut; - attenOut.SetAddr(vf2OutUb.address_); - AscendC::Cast(attenOut, vf2OutUb, AscendC::RoundMode::CAST_ROUND, m * D_BASE_SIZE); - AscendC::SetFlag(eventOVMTE3); - AscendC::WaitFlag(eventOVMTE3); - auto layoutUb = tla::MakeLayout( - tla::MakeShape(m, n), - tla::MakeStride(D_BASE_SIZE, tla::Int<1>{})); - auto attenOutUb = tla::MakeTensor(attenOut, layoutUb, Arch::PositionUB{}); - using CopyUbToGmO = Tile::CopyUb2GmTla; - CopyUbToGmO copyUbToGmO; - copyUbToGmO(attenOutGm, attenOutUb); - } - AscendC::SetFlag(eventOMTE3V); - } - -private: - AscendC::LocalTensor vf2OutUb; - static constexpr int32_t SYNC_MODE = 4; - static constexpr uint16_t FLOAT_REP_SIZE = 64; - static constexpr int32_t eventOVMTE3 = 3; - static constexpr int32_t eventOMTE3V = 3; - - template - __simd_vf__ inline void FlashUpdateNew(__ubuf__ T *updateUb, __ubuf__ T *curUb, __ubuf__ T *expMaxUb, - uint16_t m, uint16_t nLoops, uint32_t tailN) - { - using namespace AscendC::MicroAPI; - RegTensor expMaxVreg; - RegTensor preSrcVreg; - RegTensor curSrcVreg; - RegTensor mulVreg; - RegTensor addVreg; - - MaskReg pregFull = CreateMask(); - MaskReg pregTailN = UpdateMask(tailN); - - for (uint16_t i = 0; i < m; ++i) { - LoadAlign(expMaxVreg, expMaxUb + i); - for (uint16_t j = 0; j < nLoops; ++j) { - LoadAlign(preSrcVreg, updateUb + i * DBaseSize + j * FLOAT_REP_SIZE); - LoadAlign(curSrcVreg, curUb + i * DBaseSize + j * FLOAT_REP_SIZE); - Mul(mulVreg, expMaxVreg, preSrcVreg, pregFull); - Add(addVreg, mulVreg, curSrcVreg, pregFull); - StoreAlign( - updateUb + i * DBaseSize + j * FLOAT_REP_SIZE, addVreg, pregFull); - } - LoadAlign(preSrcVreg, updateUb + i * DBaseSize + nLoops * FLOAT_REP_SIZE); - LoadAlign(curSrcVreg, curUb + i * DBaseSize + nLoops * FLOAT_REP_SIZE); - Mul(mulVreg, expMaxVreg, preSrcVreg, pregTailN); - Add(addVreg, mulVreg, curSrcVreg, pregTailN); - StoreAlign( - updateUb + i * DBaseSize + nLoops * FLOAT_REP_SIZE, addVreg, pregTailN); - } - } - template - __simd_vf__ inline void FlashUpdateLastNew(__ubuf__ T *updateUb, __ubuf__ T *curUb, __ubuf__ T *expMaxUb, - __ubuf__ T *expSumUb, uint16_t m, uint16_t nLoops, uint32_t tailN) - { - using namespace AscendC::MicroAPI; - RegTensor expMaxVreg; - RegTensor preSrcVreg; - RegTensor curSrcVreg; - RegTensor mulVreg; - RegTensor addVreg; - RegTensor divDstVreg; - RegTensor expSumVreg; - - MaskReg pregFull = CreateMask(); - MaskReg pregTailN = UpdateMask(tailN); - - for (uint16_t i = 0; i < m; ++i) { - LoadAlign(expMaxVreg, expMaxUb + i); - LoadAlign(expSumVreg, expSumUb + i); - for (uint16_t j = 0; j < nLoops; ++j) { - LoadAlign(preSrcVreg, updateUb + i * DBaseSize + j * FLOAT_REP_SIZE); - LoadAlign(curSrcVreg, curUb + i * DBaseSize + j * FLOAT_REP_SIZE); - Mul(mulVreg, expMaxVreg, preSrcVreg, pregFull); - Add(addVreg, mulVreg, curSrcVreg, pregFull); - Div(divDstVreg, addVreg, expSumVreg, pregFull); - StoreAlign( - updateUb + i * DBaseSize + j * FLOAT_REP_SIZE, divDstVreg, pregFull); - } - LoadAlign(preSrcVreg, updateUb + i * DBaseSize + nLoops * FLOAT_REP_SIZE); - LoadAlign(curSrcVreg, curUb + i * DBaseSize + nLoops * FLOAT_REP_SIZE); - Mul(mulVreg, expMaxVreg, preSrcVreg, pregTailN); - Add(addVreg, mulVreg, curSrcVreg, pregTailN); - Div(divDstVreg, addVreg, expSumVreg, pregTailN); - StoreAlign( - updateUb + i * DBaseSize + nLoops * FLOAT_REP_SIZE, divDstVreg, pregTailN); - } - } - template - __simd_vf__ inline void LastDivNew(__ubuf__ T *updateUb, __ubuf__ T *curUb, __ubuf__ T *expSumUb, - uint16_t m, uint16_t nLoops, uint32_t tailN) - { - using namespace AscendC::MicroAPI; - RegTensor curSrcVreg; - RegTensor divDstVreg; - RegTensor expSumVreg; - MaskReg pregFull = CreateMask(); - MaskReg pregTailN = UpdateMask(tailN); - for (uint16_t i = 0; i < m; ++i) { - LoadAlign(expSumVreg, expSumUb + i); - for (uint16_t j = 0; j < nLoops; ++j) { - LoadAlign(curSrcVreg, curUb + i * DBaseSize + j * FLOAT_REP_SIZE); - Div(divDstVreg, curSrcVreg, expSumVreg, pregFull); - StoreAlign( - updateUb + i * DBaseSize + j * FLOAT_REP_SIZE, divDstVreg, pregFull); - } - LoadAlign(curSrcVreg, curUb + i * DBaseSize + nLoops * FLOAT_REP_SIZE); - Div(divDstVreg, curSrcVreg, expSumVreg, pregTailN); - StoreAlign( - updateUb + i * DBaseSize + nLoops * FLOAT_REP_SIZE, divDstVreg, pregTailN); - } - } -}; -} // namespace Catlass::Epilogue::Block - -#endif // CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_RESCALE_O_A5 diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_fa_softmax_ascend950.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_fa_softmax_ascend950.hpp deleted file mode 100644 index 60c4563a9b..0000000000 --- a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_fa_softmax_ascend950.hpp +++ /dev/null @@ -1,405 +0,0 @@ -/** - * Copyright (c) 2026 Huawei Technologies Co., Ltd. - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of - * CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ - -#ifndef CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_FA_SOFTMAX_A5 -#define CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_FA_SOFTMAX_A5 - -#include "catlass/kda_catlass.hpp" -#include "catlass/arch/kda_cross_core_sync.hpp" -#include "catlass/arch/kda_resource.hpp" -#include "catlass/epilogue/kda_dispatch_policy.hpp" -#include "catlass/kda_gemm_coord.hpp" -#include "catlass/kda_matrix_coord.hpp" -#include "tla/kda_tensor.hpp" -#include "tla/kda_tla_layout.hpp" - -namespace Catlass::Epilogue::Block { - -template < - class L1TileShape_, - class PType_, - class SType_, - class MaskType_, - bool ATTENTION_MASK_FLAG_> -class BlockEpilogue< - EpilogueAscend950FASoftmax, - L1TileShape_, - PType_, - SType_, - MaskType_> { -public: - using DispatchPolicy = EpilogueAscend950FASoftmax; - using ArchTag = typename DispatchPolicy::ArchTag; - using L1TileShape = L1TileShape_; - using ElementP = typename PType_::Element; - using ElementS = typename SType_::Element; - using ElementMask = typename MaskType_::Element; - using LayoutTagMask = typename MaskType_::Layout; - using LayoutTagP = typename PType_::Layout; - - static constexpr uint32_t S1_BASE_SIZE = tla::get<0>(L1TileShape{}); - static constexpr uint32_t S2_BASE_SIZE = tla::get<1>(L1TileShape{}); - static constexpr bool ATTENTION_MASK_FLAG = ATTENTION_MASK_FLAG_; - static constexpr int32_t HALF_S1_BASE_SIZE = S1_BASE_SIZE >> 1; - static constexpr int32_t HALF_VEC_SIZE = HALF_S1_BASE_SIZE * sizeof(ElementS); - static constexpr int32_t HALF_MASK_BLOCK_SIZE = HALF_S1_BASE_SIZE * S2_BASE_SIZE * sizeof(ElementMask); - static constexpr int32_t HALF_SCM_BLOCK_SIZE = HALF_S1_BASE_SIZE * S2_BASE_SIZE * sizeof(ElementP); - static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(ElementP); - - CATLASS_DEVICE - BlockEpilogue(Arch::Resource &resource, ElementS scaleValue_, uint32_t &ubBufAddrStart) - { - int32_t eventVMTE2 = 0; - int32_t eventMTE2V = 0; - int32_t eventMTE3V = 0; - int32_t eventVMTE3 = 0; - scaleValue = scaleValue_; - expSumUb = resource.ubBuf.template GetBufferByByte(ubBufAddrStart); - ubBufAddrStart += HALF_VEC_SIZE; - nowMaxUb = resource.ubBuf.template GetBufferByByte(ubBufAddrStart); - ubBufAddrStart += HALF_VEC_SIZE; - for (uint32_t i = 0; i < TASK_NUM2; i++) { - vf1OutUbList[i] = resource.ubBuf.template GetBufferByByte(ubBufAddrStart); - ubBufAddrStart += HALF_SCM_BLOCK_SIZE; - if constexpr (ATTENTION_MASK_FLAG) { - attenMaskUbList[i] = resource.ubBuf.template GetBufferByByte(ubBufAddrStart); - ubBufAddrStart += HALF_MASK_BLOCK_SIZE; - eventUbMaskVMTE2List[i] = eventVMTE2++; - eventUbMaskMTE2VList[i] = eventMTE2V++; - AscendC::SetFlag(eventUbMaskVMTE2List[i]); - } - } - - for (uint32_t i = 0; i < TASK_NUM3; i++) { - eventUbPMTE3VList[i] = eventMTE3V++; - eventUbPVMTE3List[i] = eventVMTE3++; - AscendC::SetFlag(eventUbPMTE3VList[i]); - } - } - - CATLASS_DEVICE - ~BlockEpilogue() - { - if constexpr (ATTENTION_MASK_FLAG) { - for (uint32_t i = 0; i < TASK_NUM2; i++) { - AscendC::WaitFlag(eventUbMaskVMTE2List[i]); - } - } - for (uint32_t i = 0; i < TASK_NUM3; i++) { - AscendC::WaitFlag(eventUbPMTE3VList[i]); - } - } - - template - CATLASS_DEVICE void operator()( - TensorDst &vf1OutL1, - AscendC::LocalTensor &sumUb, - AscendC::LocalTensor &lastMaxUb, - AscendC::LocalTensor &expMaxUb, - TensorSrc &bmm1Res, - TensorMask attenMaskGm, - bool isUpdate, - int8_t taskIdMod2, - int8_t taskIdMod3, - uint64_t MM1_RES_INTRA_EVENT, - uint64_t SYNC_V1_C2_FLAG) - { - uint32_t m = tla::get<0>(bmm1Res.shape()); - uint32_t n = tla::get<1>(bmm1Res.shape()); - uint32_t blockStride = AscendC::CeilDivision(m, ELE_NUM_PER_C0) * ELE_NUM_PER_C0; - auto vf1UbLayout = tla::MakeLayout( - tla::MakeShape( - tla::MakeShape(m, tla::Int<1>{}), - tla::MakeShape( - tla::Int{}, - AscendC::CeilDivision(S2_BASE_SIZE, tla::Int{}))), - tla::MakeStride( - tla::MakeStride(tla::Int{}, tla::Int{} * m), - tla::MakeStride(tla::Int<1>{}, blockStride * ELE_NUM_PER_C0))); - auto vf1OutUb = tla::MakeTensor( - vf1OutUbList[taskIdMod2], - vf1UbLayout, - Arch::PositionUB{}); - - constexpr int16_t vlSize = static_cast(AscendC::GetVecLen() / sizeof(ElementS)); - uint32_t tailN = (n - 1) % vlSize + 1; - uint32_t tailM = m; - - __ubuf__ ElementP *outputAddr = (__ubuf__ ElementP *)vf1OutUb.data().GetPhyAddr(); - __ubuf__ ElementS *inputAddr = (__ubuf__ ElementS *)bmm1Res.data().GetPhyAddr(); - __ubuf__ ElementMask *maskUbAddr; - __ubuf__ ElementMask *maskUbUnrollAddr; - __ubuf__ ElementS *lastMaxUbAddr; - __ubuf__ ElementS *nowMaxUbAddr; - __ubuf__ ElementS *expSumUbAddr; - if constexpr (ATTENTION_MASK_FLAG) { - AscendC::WaitFlag(taskIdMod2); - auto layoutUb = tla::MakeLayout( - tla::MakeShape(m, n), - tla::MakeStride(S2_BASE_SIZE, tla::Int<1>{})); - auto attenMaskUb = tla::MakeTensor( - attenMaskUbList[taskIdMod2], - layoutUb, - Arch::PositionUB{}); - using CopyGmToUbMask = Tile::CopyGm2UbTla; - CopyGmToUbMask copyGm2UbMask; - copyGm2UbMask(attenMaskUb, attenMaskGm); - AscendC::SetFlag(taskIdMod2); - AscendC::WaitFlag(taskIdMod2); - maskUbAddr = (__ubuf__ ElementMask *)attenMaskUb.data().GetPhyAddr(); - maskUbUnrollAddr = (__ubuf__ ElementMask *)(attenMaskUb.data().GetPhyAddr() + FLOAT_REP_SIZE); - } - - if (likely(isUpdate)) { - lastMaxUbAddr = (__ubuf__ ElementS *)lastMaxUb.GetPhyAddr(); - nowMaxUbAddr = (__ubuf__ ElementS *)nowMaxUb.GetPhyAddr(); - expSumUbAddr = (__ubuf__ ElementS *)expSumUb.GetPhyAddr(); - } else { - nowMaxUbAddr = (__ubuf__ ElementS *)lastMaxUb.GetPhyAddr(); - expSumUbAddr = (__ubuf__ ElementS *)sumUb.GetPhyAddr(); - } - - if (n == 128) { - ComputeMaskandScale( - inputAddr, maskUbAddr, maskUbUnrollAddr, nowMaxUbAddr, m, tailN, scaleValue); - } else if (n <= 64) { - ComputeMaskandScale( - inputAddr, maskUbAddr, maskUbUnrollAddr, nowMaxUbAddr, m, tailN, scaleValue); - } else { - ComputeMaskandScale( - inputAddr, maskUbAddr, maskUbUnrollAddr, nowMaxUbAddr, m, tailN, scaleValue); - } - - if constexpr (ATTENTION_MASK_FLAG) { - AscendC::SetFlag(taskIdMod2); - } - - if (likely(isUpdate)) { - UpdateMax(nowMaxUbAddr, lastMaxUbAddr, tailM); - } - AscendC::WaitFlag(taskIdMod3); - if (unlikely(n > 64)) { - ComputeExpSubSum( - outputAddr, inputAddr, nowMaxUbAddr, expSumUbAddr, m, blockStride); - } else { - ComputeExpSubSum( - outputAddr, inputAddr, nowMaxUbAddr, expSumUbAddr, m, blockStride); - } - - AscendC::SetFlag(taskIdMod2); - AscendC::WaitFlag(taskIdMod2); - AscendC::CrossCoreSetFlag(MM1_RES_INTRA_EVENT); - - if (likely(m != 0)) { - using CopyUbToL1P = Tile::CopyUb2L1Tla; - CopyUbToL1P copyUb2L1P; - copyUb2L1P(vf1OutL1, vf1OutUb); - } - AscendC::SetFlag(taskIdMod3); - AscendC::CrossCoreSetFlag(SYNC_V1_C2_FLAG); - - if (likely(isUpdate)) { - __ubuf__ ElementS *sumUbAddr = (__ubuf__ ElementS *)sumUb.GetPhyAddr(); - __ubuf__ ElementS *expMaxUbAddr = (__ubuf__ ElementS *)expMaxUb.GetPhyAddr(); - UpdateExpSumAndExpMax( - sumUbAddr, expMaxUbAddr, lastMaxUbAddr, expSumUbAddr, nowMaxUbAddr, tailM); - } - } - -private: - static constexpr uint32_t TASK_NUM2 = 2; - static constexpr uint32_t TASK_NUM3 = 3; - static constexpr int32_t SYNC_MODE = 4; - static constexpr int64_t BLOCK_BYTES = 32; - static constexpr uint32_t REPEAT_STRIDE = 1; - static constexpr float MIN_VALUE = -3e38f; - static constexpr uint16_t FLOAT_REP_SIZE = 64; - static constexpr uint16_t DOUBLE_FLOAT_REP_SIZE = 128; - - int32_t eventUbMaskVMTE2List[TASK_NUM2]; - int32_t eventUbMaskMTE2VList[TASK_NUM2]; - int32_t eventUbPMTE3VList[TASK_NUM3]; - int32_t eventUbPVMTE3List[TASK_NUM3]; - - ElementS scaleValue; - AscendC::LocalTensor attenMaskUbList[TASK_NUM2]; - AscendC::LocalTensor vf1OutUbList[TASK_NUM2]; - AscendC::LocalTensor expSumUb; - AscendC::LocalTensor nowMaxUb; - - enum class NRangeIndex { - N0_64 = 0, - N65_127, - N128, - N128_INF - }; - - template - __simd_vf__ inline void ComputeMaskandScale(__ubuf__ ElementS *srcUb, __ubuf__ uint8_t *maskUb, __ubuf__ uint8_t *maskUbUnroll, - __ubuf__ ElementS *newMaxUb, uint16_t m, uint32_t tailN, ElementS dScale) - { - using namespace AscendC::MicroAPI; - RegTensor minVreg; - RegTensor srcVreg0; - RegTensor srcVreg1; - RegTensor maxVreg; - RegTensor maxTmpVreg; - UnalignReg maxUreg; - MaskReg pregCompare0; - MaskReg pregCompare1; - MaskReg pregFull = CreateMask(); - MaskReg pregTailN = UpdateMask(tailN); - - if constexpr (NRange < NRangeIndex::N128 || HasAtten) { - Duplicate(minVreg, MIN_VALUE); - } - for (uint16_t i = 0; i < m; ++i) { - if constexpr (NRange > NRangeIndex::N0_64) { - LoadAlign(srcVreg0, srcUb + i * S2BaseSize); - LoadAlign(srcVreg1, srcUb + i * S2BaseSize + FLOAT_REP_SIZE); - Muls(srcVreg0, srcVreg0, dScale, pregFull); - Muls(srcVreg1, srcVreg1, dScale, pregTailN); - if constexpr (HasAtten) { - LoadAlign( - pregCompare0, (__ubuf__ uint32_t *&)maskUb, S2BaseSize); - LoadAlign( - pregCompare1, (__ubuf__ uint32_t *&)maskUbUnroll, S2BaseSize); - Select(srcVreg0, minVreg, srcVreg0, pregCompare0); - Select(srcVreg1, minVreg, srcVreg1, pregCompare1); - } - if constexpr (NRange < NRangeIndex::N128) { - Select(srcVreg1, srcVreg1, minVreg, pregTailN); - } - StoreAlign(srcUb + i * S2BaseSize, srcVreg0, pregFull); - StoreAlign(srcUb + i * S2BaseSize + FLOAT_REP_SIZE, srcVreg1, pregFull); - Max(maxTmpVreg, srcVreg0, srcVreg1, pregFull); - ReduceMax(maxVreg, maxTmpVreg, pregFull); - } else { - LoadAlign(srcVreg0, srcUb + i * S2BaseSize); - Muls(srcVreg0, srcVreg0, dScale, pregTailN); - if constexpr (HasAtten) { - LoadAlign( - pregCompare0, (__ubuf__ uint32_t *&)maskUb, S2BaseSize); - Select(srcVreg0, minVreg, srcVreg0, pregCompare0); - } - Select(srcVreg0, srcVreg0, minVreg, pregTailN); - StoreAlign(srcUb + i * S2BaseSize, srcVreg0, pregFull); - ReduceMax(maxVreg, srcVreg0, pregFull); - } - StoreUnAlign(newMaxUb, maxVreg, maxUreg, 1); - } - StoreUnAlignPost(newMaxUb, maxUreg, 0); - } - - template - __simd_vf__ inline void UpdateMax( - __ubuf__ ElementS *nowMaxUb, __ubuf__ ElementS *lastMaxUb, uint32_t tailM) - { - using namespace AscendC::MicroAPI; - RegTensor nowMaxVreg; - RegTensor lastMaxVreg; - RegTensor maxVreg; - - MaskReg pregTailM = UpdateMask(tailM); - LoadAlign(lastMaxVreg, lastMaxUb); - LoadAlign(nowMaxVreg, nowMaxUb); - Max(maxVreg, nowMaxVreg, lastMaxVreg, pregTailM); - StoreAlign(nowMaxUb, maxVreg, pregTailM); - } - - template - __simd_vf__ inline void ComputeExpSubSum(__ubuf__ ElementP *expUb, __ubuf__ ElementS *srcUb, - __ubuf__ ElementS *nowMaxUb, __ubuf__ ElementS *expSumUb, uint16_t m, uint32_t blockStride) - { - using namespace AscendC::MicroAPI; - constexpr static CastTrait castTraitZero = { - RegLayout::ZERO, - SatMode::SAT, - MaskMergeMode::ZEROING, - AscendC::RoundMode::CAST_ROUND, - }; - - constexpr static CastTrait castTraitOne = { - RegLayout::ONE, - SatMode::SAT, - MaskMergeMode::ZEROING, - AscendC::RoundMode::CAST_ROUND, - }; - RegTensor srcVreg0; - RegTensor srcVreg1; - RegTensor expVreg; - RegTensor expVreg0; - RegTensor expVreg1; - RegTensor expSumVreg; - RegTensor maxVreg; - - RegTensor expDstVreg0; - RegTensor expDstVreg1; - RegTensor expDstVreg; - - UnalignReg expSumUreg; - - MaskReg pregFull = CreateMask(); - MaskReg pregFull16 = CreateMask(); - for (uint16_t i = 0; i < m; ++i) { - LoadAlign(maxVreg, nowMaxUb + i); - if constexpr (NRange > NRangeIndex::N0_64) { - LoadAlign(srcVreg0, srcVreg1, srcUb + i * S2BaseSize); - ExpSub(expVreg0, srcVreg0, maxVreg, pregFull); - ExpSub(expVreg1, srcVreg1, maxVreg, pregFull); - Add(expVreg, expVreg0, expVreg1, pregFull); - Cast(expDstVreg0, expVreg0, pregFull); - Cast(expDstVreg1, expVreg1, pregFull); - Or((RegTensor &)expDstVreg, - (RegTensor &)expDstVreg0, (RegTensor &)expDstVreg1, - pregFull16); - StoreAlign( - expUb, expDstVreg, blockStride, REPEAT_STRIDE, pregFull16); - } else { - LoadAlign(srcVreg0, srcUb + i * S2BaseSize); - ExpSub(expVreg, srcVreg0, maxVreg, pregFull); - Cast(expDstVreg, expVreg, pregFull); - DeInterleave(expDstVreg0, expDstVreg1, expDstVreg, expDstVreg); - StoreAlign( - expUb, expDstVreg0, blockStride, REPEAT_STRIDE, pregFull16); - } - ReduceSum(expSumVreg, expVreg, pregFull); - StoreUnAlign(expSumUb, expSumVreg, expSumUreg, 1); - } - StoreUnAlignPost(expSumUb, expSumUreg, 0); - } - - template - __simd_vf__ inline void UpdateExpSumAndExpMax(__ubuf__ ElementS *sumUb, __ubuf__ ElementS *expMaxUb, - __ubuf__ ElementS *maxUb, __ubuf__ ElementS *expSumUb, __ubuf__ ElementS *nowMaxUb, uint32_t tailM) - { - using namespace AscendC::MicroAPI; - RegTensor nowMaxVreg; - RegTensor lastMaxVreg; - RegTensor expMaxVreg; - RegTensor lastExpSumVreg; - RegTensor brcExpSumVreg; - MaskReg pregTailM = UpdateMask(tailM); - LoadAlign(lastMaxVreg, maxUb); - LoadAlign(nowMaxVreg, nowMaxUb); - ExpSub(expMaxVreg, lastMaxVreg, nowMaxVreg, pregTailM); - StoreAlign(expMaxUb, expMaxVreg, pregTailM); - StoreAlign(maxUb, nowMaxVreg, pregTailM); - - LoadAlign(lastExpSumVreg, sumUb); - LoadAlign(brcExpSumVreg, expSumUb); - MulDstAdd(expMaxVreg, lastExpSumVreg, brcExpSumVreg, pregTailM); - StoreAlign(sumUb, expMaxVreg, pregTailM); - } -}; - -} // namespace Catlass::Epilogue::Block - -#endif // CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_FA_SOFTMAX_A5 diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_fixpipe.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_fixpipe.hpp deleted file mode 100644 index eaf873b4bd..0000000000 --- a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_fixpipe.hpp +++ /dev/null @@ -1,83 +0,0 @@ -/** - * Copyright (c) 2026 Huawei Technologies Co., Ltd. - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of - * CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ - -#ifndef CATLASS_EPILOGUE_BLOCK_EPILOGUE_FIXPIPE_HPP -#define CATLASS_EPILOGUE_BLOCK_EPILOGUE_FIXPIPE_HPP - -#include "catlass/kda_catlass.hpp" -#include "catlass/arch/kda_resource.hpp" -#include "catlass/epilogue/kda_dispatch_policy.hpp" -#include "catlass/kda_gemm_coord.hpp" -#include "catlass/kda_matrix_coord.hpp" -#include "catlass/epilogue/tile/kda_epilogue_tile_copy.hpp" -#include "catlass/gemm/kda_gemm_helper.hpp" -#include "catlass/epilogue/tile/kda_copy_ub_to_gm_tla.hpp" - -namespace Catlass::Epilogue::Block { - -template < - class L0TileShape_, - class ElementOut_, - class ElementIn_, - bool SPLIT_M_> -class BlockEpilogue< - EpilogueAscend950Fixpipe, - L0TileShape_, - ElementOut_, - ElementIn_> { -public: - using DispatchPolicy = EpilogueAscend950Fixpipe; - using ArchTag = typename DispatchPolicy::ArchTag; - using ElementOut = ElementOut_; - using ElementIn = ElementIn_; - - static constexpr bool SPLIT_M = DispatchPolicy::SPLIT_M; - static constexpr int64_t ML0_ = tla::get<0>(L0TileShape_{}); - static constexpr int64_t NL0_ = tla::get<1>(L0TileShape_{}); - - CATLASS_DEVICE - BlockEpilogue() = default; - - template - CATLASS_DEVICE void operator()(TensorOut &tensorOut, TensorIn &tensorIn) - { - if constexpr (!SPLIT_M && AscendC::GetSubBlockIdx() > 0) { - return; - } - - int64_t blockShapeM = tla::get<0>(tensorOut.shape()); - int64_t blockShapeN = tla::get<1>(tensorOut.shape()); - int64_t halfBlockShapeM = CeilDiv(blockShapeM, AscendC::GetTaskRation()); - if constexpr (SPLIT_M) { - blockShapeM = (static_cast(blockShapeM) & 1UL) > 0UL ? - (halfBlockShapeM - AscendC::GetSubBlockIdx()) : - halfBlockShapeM; - } - - // real copy data size - int64_t copySize = blockShapeM * blockShapeN; - if (copySize <= 0) { - return; - } - - auto tensorTileOut = GetTile( - tensorOut, - tla::MakeCoord(halfBlockShapeM * (AscendC::GetSubBlockIdx() & 0x1), 0), - tla::MakeShape(blockShapeM, blockShapeN)); - - using CopyUbToGm = Epilogue::Tile::CopyUb2GmTla; - CopyUbToGm copyUbToGm; - copyUbToGm(tensorTileOut, tensorIn); - } -}; - -} // namespace Catlass::Epilogue::Block - -#endif // CATLASS_EPILOGUE_BLOCK_EPILOGUE_FIXPIPE_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_mla_fd_rescale_o.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_mla_fd_rescale_o.hpp deleted file mode 100644 index d57ce3cdb7..0000000000 --- a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_mla_fd_rescale_o.hpp +++ /dev/null @@ -1,307 +0,0 @@ -/** - * Copyright (c) 2025 Huawei Technologies Co., Ltd. - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of - * CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ - -#ifndef CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_MLA_FD_RESCALE_O_HPP -#define CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_MLA_FD_RESCALE_O_HPP - -#include "catlass/kda_catlass.hpp" -#include "catlass/arch/kda_resource.hpp" -#include "catlass/epilogue/kda_dispatch_policy.hpp" -#include "catlass/epilogue/tile/kda_epilogue_tile_copy.hpp" -#include "catlass/kda_gemm_coord.hpp" -#include "catlass/kda_matrix_coord.hpp" - -namespace Catlass::Epilogue::Block { - -template < - class OutputType_, - class InputType_, - uint32_t ComputeEleNum_> -class BlockEpilogue< - EpilogueAtlasA2MLAFDRescaleO, - OutputType_, - InputType_> { -public: - // Type aliases - using DispatchPolicy = EpilogueAtlasA2MLAFDRescaleO; - using ArchTag = typename DispatchPolicy::ArchTag; - using ElementOutput = typename OutputType_::Element; - using ElementInput = typename InputType_::Element; - using LayoutOutput = typename OutputType_::Layout; - using LayoutInput = typename InputType_::Layout; - - static constexpr uint32_t KV_SPLIT_MAX = DispatchPolicy::KV_SPLIT_MAX; - static constexpr uint32_t HEADS_PROCESS_MAX = DispatchPolicy::HEADS_PROCESS_MAX; - static constexpr uint32_t COMPUTE_ELE_NUM = DispatchPolicy::COMPUTE_ELE_NUM; - static constexpr uint32_t FLOAT_ELENUM_PER_VECCALC = 64; - static constexpr uint32_t FLOAT_BLOCK_SIZE = 8; - static constexpr uint32_t STAGES = 2; - - CATLASS_DEVICE - BlockEpilogue(Arch::Resource &resource, uint32_t kvSplitCoreNum_) - { - kvSplitCoreNum = kvSplitCoreNum_; - - uint32_t ubOffset = 0; - oIn[0] = resource.ubBuf.template GetBufferByByte(ubOffset); - ubOffset += COMPUTE_ELE_NUM * sizeof(float); - oIn[1] = resource.ubBuf.template GetBufferByByte(ubOffset); - ubOffset += COMPUTE_ELE_NUM * sizeof(float); - oTemp[0] = resource.ubBuf.template GetBufferByByte(ubOffset); - ubOffset += COMPUTE_ELE_NUM * sizeof(float); - oTemp[1] = resource.ubBuf.template GetBufferByByte(ubOffset); - ubOffset += COMPUTE_ELE_NUM * sizeof(float); - oSum = resource.ubBuf.template GetBufferByByte(ubOffset); - ubOffset += COMPUTE_ELE_NUM * sizeof(float); - out = resource.ubBuf.template GetBufferByByte(ubOffset); - ubOffset += COMPUTE_ELE_NUM * sizeof(ElementOutput); - lIn = resource.ubBuf.template GetBufferByByte(ubOffset); - ubOffset += KV_SPLIT_MAX * HEADS_PROCESS_MAX * sizeof(float); - lExp = resource.ubBuf.template GetBufferByByte(ubOffset); - ubOffset += KV_SPLIT_MAX * HEADS_PROCESS_MAX * sizeof(float); - lMax = resource.ubBuf.template GetBufferByByte(ubOffset); - ubOffset += HEADS_PROCESS_MAX * sizeof(float); - lSum = resource.ubBuf.template GetBufferByByte(ubOffset); - ubOffset += HEADS_PROCESS_MAX * sizeof(float); - lBrcb[0] = resource.ubBuf.template GetBufferByByte(ubOffset); - ubOffset += HEADS_PROCESS_MAX * FLOAT_BLOCK_SIZE * sizeof(float); - lBrcb[1] = resource.ubBuf.template GetBufferByByte(ubOffset); - - AscendC::SetFlag(EVENT_ID0); - AscendC::SetFlag(EVENT_ID0); - AscendC::SetFlag(EVENT_ID1); - AscendC::SetFlag(EVENT_ID2); - } - CATLASS_DEVICE - ~BlockEpilogue() - { - AscendC::WaitFlag(EVENT_ID0); - AscendC::WaitFlag(EVENT_ID0); - AscendC::WaitFlag(EVENT_ID1); - AscendC::WaitFlag(EVENT_ID2); - } - - CATLASS_DEVICE - void SetMask(int32_t len) - { - constexpr int32_t MAX_MASK_LEN = 128; - constexpr int32_t HALF_MASK_LEN = 64; - if (len >= MAX_MASK_LEN) { - AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); - return; - } - int32_t highMask = len - HALF_MASK_LEN > 0 ? len - HALF_MASK_LEN : 0; - int32_t lowMask = len - HALF_MASK_LEN >= 0 ? HALF_MASK_LEN : len; - if (len < HALF_MASK_LEN) { - AscendC::SetVectorMask(0x0, ((uint64_t)1 << lowMask) - 1); - } else { - AscendC::SetVectorMask(((uint64_t)1 << highMask) - 1, 0xffffffffffffffff); - } - } - - CATLASS_DEVICE - void operator()( - AscendC::GlobalTensor gOutput, - AscendC::GlobalTensor gOCoreTmp, - AscendC::GlobalTensor gl, - uint32_t actualHeads, uint32_t headsProcess, uint32_t headSize) - { - uint32_t kvSplitRound = (kvSplitCoreNum + FLOAT_BLOCK_SIZE - 1) / FLOAT_BLOCK_SIZE * FLOAT_BLOCK_SIZE; - - AscendC::WaitFlag(EVENT_ID2); - AscendC::DataCopyPad( - lIn, gl, - AscendC::DataCopyExtParams( - actualHeads, kvSplitCoreNum * sizeof(ElementInput), 0, - (KV_SPLIT_MAX - kvSplitCoreNum) / FLOAT_BLOCK_SIZE, 0), - AscendC::DataCopyPadExtParams(false, 0, 0, 0)); - - AscendC::SetFlag(EVENT_ID2); - AscendC::WaitFlag(EVENT_ID2); - - SetMask(kvSplitCoreNum); - AscendC::WholeReduceMax( - lMax, lIn, (int32_t)0, actualHeads, 1, 1, 8, - AscendC::ReduceOrder::ORDER_ONLY_VALUE); - AscendC::PipeBarrier(); - - for (uint32_t i = 0; i < kvSplitRound / FLOAT_BLOCK_SIZE; i++) { - AscendC::Brcb( - lExp[i * FLOAT_BLOCK_SIZE], - lMax, - (headsProcess + FLOAT_BLOCK_SIZE - 1) / FLOAT_BLOCK_SIZE, - AscendC::BrcbRepeatParams(KV_SPLIT_MAX / FLOAT_BLOCK_SIZE, - 8 * KV_SPLIT_MAX / FLOAT_BLOCK_SIZE)); - } - AscendC::PipeBarrier(); - - SetMask(kvSplitCoreNum); - AscendC::Sub( - lExp, - lIn, - lExp, - (uint64_t)0, - actualHeads, - AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); - AscendC::PipeBarrier(); - - AscendC::Exp( - lExp, - lExp, - (uint64_t)0, - actualHeads, - AscendC::UnaryRepeatParams(1, 1, 8, 8)); - AscendC::PipeBarrier(); - - AscendC::RepeatReduceSum(lSum, lExp, actualHeads, 0, 0, 1, 1, 8); - AscendC::PipeBarrier(); - - AscendC::Ln(lSum, lSum, (headsProcess + FLOAT_BLOCK_SIZE - 1) / FLOAT_BLOCK_SIZE * FLOAT_BLOCK_SIZE); - AscendC::PipeBarrier(); - - AscendC::Add(lSum, lSum, lMax, (headsProcess + FLOAT_BLOCK_SIZE - 1) / FLOAT_BLOCK_SIZE * FLOAT_BLOCK_SIZE); - AscendC::PipeBarrier(); - - for (uint32_t i = 0; i < kvSplitRound / FLOAT_BLOCK_SIZE; i++) { - AscendC::Brcb( - lExp[i * FLOAT_BLOCK_SIZE], - lSum, - (headsProcess + FLOAT_BLOCK_SIZE - 1) / FLOAT_BLOCK_SIZE, - AscendC::BrcbRepeatParams(KV_SPLIT_MAX / FLOAT_BLOCK_SIZE, - 8 * KV_SPLIT_MAX / FLOAT_BLOCK_SIZE)); - } - AscendC::PipeBarrier(); - - SetMask(kvSplitCoreNum); - AscendC::Sub( - lExp, - lIn, - lExp, - (uint64_t)0, - actualHeads, - AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); - AscendC::PipeBarrier(); - AscendC::SetFlag(EVENT_ID2); - - AscendC::Exp( - lExp, - lExp, - (uint64_t)0, - actualHeads, - AscendC::UnaryRepeatParams(1, 1, 8, 8)); - AscendC::PipeBarrier(); - - // preload - AscendC::WaitFlag(EVENT_ID0); - AscendC::DataCopyPad( - oIn[0], gOCoreTmp, - AscendC::DataCopyExtParams( - actualHeads, headSize * sizeof(ElementInput), - (kvSplitCoreNum * headSize - headSize) * sizeof(ElementInput), 0, 0), - AscendC::DataCopyPadExtParams(false, 0, 0, 0)); - AscendC::SetFlag(EVENT_ID0); - - AscendC::SetFlag(EVENT_ID2); - AscendC::WaitFlag(EVENT_ID2); - - SetMask(FLOAT_ELENUM_PER_VECCALC); - uint32_t bufferId = 0; - for (uint32_t i = 0; i < kvSplitCoreNum; i++) { - // load next o - if (i < kvSplitCoreNum - 1) { - uint32_t nextBufferId = 1 - bufferId; - AscendC::WaitFlag(oInEventList[nextBufferId]); - AscendC::DataCopyPad( - oIn[nextBufferId], gOCoreTmp[(i + 1) * headSize], - AscendC::DataCopyExtParams( - actualHeads, headSize * sizeof(ElementInput), - (kvSplitCoreNum * headSize - headSize) * sizeof(ElementInput), 0, 0), - AscendC::DataCopyPadExtParams(false, 0, 0, 0)); - AscendC::SetFlag(oInEventList[nextBufferId]); - } - - AscendC::PipeBarrier(); - for (uint32_t j = 0; j < actualHeads; j++) { - float a = lExp[j * KV_SPLIT_MAX + i].GetValue(0); - AscendC::SetFlag(oTempEventList[bufferId]); - AscendC::WaitFlag(oTempEventList[bufferId]); - AscendC::Duplicate( - lBrcb[bufferId][j * FLOAT_BLOCK_SIZE], a, uint64_t(0), 1, 0, 0); - } - AscendC::PipeBarrier(); - - // caculate current o - AscendC::WaitFlag(oInEventList[bufferId]); - uint32_t loops = (headSize + FLOAT_ELENUM_PER_VECCALC - 1) / FLOAT_ELENUM_PER_VECCALC; - if (i > 0) { - for (uint32_t j = 0; j < loops; j++) { - AscendC::Mul( - oTemp[bufferId][j * FLOAT_ELENUM_PER_VECCALC], lBrcb[bufferId], - oIn[bufferId][j * FLOAT_ELENUM_PER_VECCALC], (uint64_t)0, - actualHeads, - AscendC::BinaryRepeatParams( - 1, 0, 1, headSize / FLOAT_BLOCK_SIZE, 1, - headSize / FLOAT_BLOCK_SIZE)); - } - } else { - for (uint32_t j = 0; j < loops; j++) { - AscendC::Mul( - oSum[j * FLOAT_ELENUM_PER_VECCALC], lBrcb[bufferId], - oIn[bufferId][j * FLOAT_ELENUM_PER_VECCALC], (uint64_t)0, - actualHeads, - AscendC::BinaryRepeatParams( - 1, 0, 1, headSize / FLOAT_BLOCK_SIZE, 1, - headSize / FLOAT_BLOCK_SIZE)); - } - } - AscendC::PipeBarrier(); - AscendC::SetFlag(oInEventList[bufferId]); - - if (i > 0) { - AscendC::Add(oSum, oSum, oTemp[bufferId], actualHeads * headSize); - } - AscendC::PipeBarrier(); - bufferId = 1 - bufferId; - } - - AscendC::WaitFlag(EVENT_ID0); - if (std::is_same::value) { - AscendC::Cast(out, oSum, AscendC::RoundMode::CAST_RINT, actualHeads * headSize); - } else { - AscendC::Cast(out, oSum, AscendC::RoundMode::CAST_NONE, actualHeads * headSize); - } - - AscendC::SetFlag(EVENT_ID0); - AscendC::WaitFlag(EVENT_ID0); - AscendC::DataCopyPad( - gOutput, out, - AscendC::DataCopyExtParams( - actualHeads, headSize * sizeof(ElementOutput), 0, 0, 0)); - AscendC::SetFlag(EVENT_ID0); - } - -private: - uint32_t kvSplitCoreNum = 1; - AscendC::LocalTensor out; - AscendC::LocalTensor oIn[STAGES]; - AscendC::LocalTensor oTemp[STAGES]; - AscendC::LocalTensor lBrcb[STAGES]; - AscendC::LocalTensor oSum; - AscendC::LocalTensor lIn; - AscendC::LocalTensor lExp; - AscendC::LocalTensor lTrans; - AscendC::LocalTensor lMax; - AscendC::LocalTensor lSum; - - int32_t oTempEventList[STAGES] = {0, 1}; - int32_t oInEventList[STAGES] = {0, 1}; -}; -} // namespace Catlass::Epilogue::Block -#endif // CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_MLA_FD_RESCALE_O_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_mla_rescale_o.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_mla_rescale_o.hpp deleted file mode 100644 index b1723cc657..0000000000 --- a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_mla_rescale_o.hpp +++ /dev/null @@ -1,479 +0,0 @@ -/** - * Copyright (c) 2025 Huawei Technologies Co., Ltd. - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of - * CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ - -#ifndef CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_MLA_RESCALE_O_HPP -#define CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_MLA_RESCALE_O_HPP - -#include "catlass/kda_catlass.hpp" -#include "catlass/arch/kda_resource.hpp" -#include "catlass/epilogue/kda_dispatch_policy.hpp" -#include "catlass/epilogue/tile/kda_epilogue_tile_copy.hpp" -#include "catlass/kda_gemm_coord.hpp" -#include "catlass/kda_matrix_coord.hpp" - -namespace Catlass::Epilogue::Block { - -template < - class OutputType_, - class UpdateType_, - class InputType_> -class BlockEpilogue< - EpilogueAtlasA2MLARescaleO, - OutputType_, - UpdateType_, - InputType_> { -public: - // Type aliases - using DispatchPolicy = EpilogueAtlasA2MLARescaleO; - using ArchTag = typename DispatchPolicy::ArchTag; - - using ElementOutput = typename OutputType_::Element; - using ElementUpdate = typename UpdateType_::Element; - using ElementInput = typename InputType_::Element; - - using LayoutOutput = typename OutputType_::Layout; - using LayoutUpdate = typename UpdateType_::Layout; - using LayoutInput = typename InputType_::Layout; - - static constexpr uint32_t HALF_ELENUM_PER_BLK = 16; - static constexpr uint32_t HALF_ELENUM_PER_VECCALC = 128; - static constexpr uint32_t FLOAT_ELENUM_PER_VECCALC = 64; - static constexpr uint32_t HALF_ELENUM_PER_LINE = 256; - static constexpr uint32_t FLOAT_ELENUM_PER_LINE = 128; - static constexpr uint32_t MULTIPLIER = 2; - static constexpr uint32_t FLOAT_BLOCK_SIZE = 8; - static constexpr uint32_t FLOAT_VECTOR_SIZE = 64; - static constexpr uint32_t UB_UINT8_LINE_SIZE = 512; - static constexpr uint32_t UB_UINT8_BLOCK_SIZE_MLA = 16384; - static constexpr uint32_t ROW_WISE_CYCLE_TILE = 8; - static constexpr uint32_t HALF_DM_UB_SIZE = 128; - static constexpr uint32_t HALF_LL_UB_SIZE = 256; - static constexpr uint32_t VECTOR_SIZE = 128; - - CATLASS_DEVICE - BlockEpilogue(Arch::Resource &resource, uint32_t kvSplitCoreNum_) - { - // Allocate UB space - constexpr uint32_t LO_UB_TENSOR_OFFSET = 4 * UB_UINT8_BLOCK_SIZE_MLA; - constexpr uint32_t DM_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA + 6 * UB_UINT8_LINE_SIZE; - constexpr uint32_t LL_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA + 9 * UB_UINT8_LINE_SIZE; - constexpr uint32_t GL_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA + 15 * UB_UINT8_LINE_SIZE; - constexpr uint32_t GO_UB_TENSOR_OFFSET = 8 * UB_UINT8_BLOCK_SIZE_MLA; - constexpr uint32_t TV_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE_MLA; - constexpr uint32_t HM_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA + 1 * UB_UINT8_LINE_SIZE; - constexpr uint32_t GM_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA + 13 * UB_UINT8_LINE_SIZE; - - kvSplitCoreNum = kvSplitCoreNum_; - loUbTensor = resource.ubBuf.template GetBufferByByte(LO_UB_TENSOR_OFFSET); - dmUbTensor = resource.ubBuf.template GetBufferByByte(DM_UB_TENSOR_OFFSET); - llUbTensor = resource.ubBuf.template GetBufferByByte(LL_UB_TENSOR_OFFSET); - glUbTensor = resource.ubBuf.template GetBufferByByte(GL_UB_TENSOR_OFFSET); - tvUbTensor = resource.ubBuf.template GetBufferByByte(TV_UB_TENSOR_OFFSET); - goUbTensor32 = resource.ubBuf.template GetBufferByByte(GO_UB_TENSOR_OFFSET); - goUbTensor16 = resource.ubBuf.template GetBufferByByte(GO_UB_TENSOR_OFFSET); - hmUbTensor = resource.ubBuf.template GetBufferByByte(HM_UB_TENSOR_OFFSET); - gmUbTensor = resource.ubBuf.template GetBufferByByte(GM_UB_TENSOR_OFFSET); - } - - CATLASS_DEVICE - ~BlockEpilogue() - { - } - - CATLASS_DEVICE - void SetMask(int32_t len) - { - uint64_t mask = 0; - uint64_t one = 1; - uint64_t temp = len % FLOAT_VECTOR_SIZE; - for (int64_t i = 0; i < temp; i++) { - mask |= one << i; - } - if (len == VECTOR_SIZE) { - AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); - } else if (len >= FLOAT_VECTOR_SIZE) { - AscendC::SetVectorMask(mask, (uint64_t)-1); - } else { - AscendC::SetVectorMask(0x0, mask); - } - } - - CATLASS_DEVICE - void SubCoreCompute( - AscendC::GlobalTensor gInput, - AscendC::GlobalTensor gUpdate, - AscendC::GlobalTensor gOutput, - AscendC::GlobalTensor gOCoreTmp, - AscendC::GlobalTensor gl, - const LayoutInput &layoutInput, - const LayoutOutput &layoutOutput, - const LayoutUpdate &layoutUpdate, - uint32_t nIdx, - uint32_t isLastNTile, - uint32_t needRowLoop, - uint32_t rowLoopIdx, - uint32_t proTokenIdx, - uint32_t proTokenNum, - uint32_t epiTokenNum, - uint32_t integralHeadNum, - uint32_t rescaleOPingPongFlag, - uint32_t &glFlag) - { - uint32_t curRowNum = layoutInput.shape(0); - uint32_t embed = layoutInput.shape(1); - uint32_t embedRound = layoutInput.stride(0); - uint32_t strideQO = layoutOutput.stride(0); - uint32_t tokenNumPerHead = layoutOutput.shape(0); - uint32_t curRowNumAligned64 = (curRowNum + FLOAT_ELENUM_PER_VECCALC - 1) / FLOAT_ELENUM_PER_VECCALC; - uint32_t curRowNumRound = (curRowNum + HALF_ELENUM_PER_BLK - 1) / HALF_ELENUM_PER_BLK * HALF_ELENUM_PER_BLK; - uint64_t dmUbOffsetCurCycle = (uint64_t)(rescaleOPingPongFlag * HALF_DM_UB_SIZE + - rowLoopIdx * ROW_WISE_CYCLE_TILE); - uint64_t llUbOffsetCurCycle = (uint64_t)(rescaleOPingPongFlag * HALF_LL_UB_SIZE + - rowLoopIdx * ROW_WISE_CYCLE_TILE); - uint32_t oUbOffset = oPingPangFlag * ROW_WISE_CYCLE_TILE * embedRound; - AscendC::WaitFlag(oPingPangFlag); - if ((nIdx - 1) != 0) { - AscendC::DataCopy( - loUbTensor[oUbOffset], - gInput, - AscendC::DataCopyParams( - 1, curRowNum * embedRound / FLOAT_BLOCK_SIZE, 0, 0)); - AscendC::SetFlag(EVENT_ID0); - AscendC::WaitFlag(EVENT_ID0); - } - AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); - AscendC::WaitFlag(oPingPangFlag + 4); - if ((nIdx - 1) != 0) { - // *** dm = exp(dm) - if (rowLoopIdx == 0) { - AscendC::Exp( - dmUbTensor[dmUbOffsetCurCycle], - dmUbTensor[dmUbOffsetCurCycle], - (uint64_t)0, - curRowNumAligned64, - AscendC::UnaryRepeatParams(1, 1, 8, 8)); - AscendC::PipeBarrier(); - AscendC::Mul( - glUbTensor, - dmUbTensor[dmUbOffsetCurCycle], - glUbTensor, - (uint64_t)0, - curRowNumAligned64, - AscendC::BinaryRepeatParams( - 1, 1, 1, 8, 8, 8)); - AscendC::PipeBarrier(); - AscendC::Add( - glUbTensor, - glUbTensor, - llUbTensor[llUbOffsetCurCycle], - (uint64_t)0, - curRowNumAligned64, - AscendC::BinaryRepeatParams( - 1, 1, 1, 8, 8, 8)); - AscendC::PipeBarrier(); - } - AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); - AscendC::Brcb( - tvUbTensor.ReinterpretCast(), - dmUbTensor[dmUbOffsetCurCycle].ReinterpretCast(), - curRowNumRound / FLOAT_BLOCK_SIZE, - AscendC::BrcbRepeatParams(1, 8)); - AscendC::PipeBarrier(); - if (needRowLoop) { - AscendC::DataCopy( - goUbTensor32[oUbOffset], - gUpdate, - AscendC::DataCopyParams( - 1, curRowNum * embedRound / FLOAT_BLOCK_SIZE, 0, 0)); - AscendC::SetFlag(EVENT_ID0); - AscendC::WaitFlag(EVENT_ID0); - } - // *** go = go * dm_block - AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); - for (uint32_t mulIdx = 0; mulIdx < embed / FLOAT_VECTOR_SIZE; ++mulIdx) { - AscendC::Mul( - goUbTensor32[oUbOffset + mulIdx * FLOAT_VECTOR_SIZE], - goUbTensor32[oUbOffset + mulIdx * FLOAT_VECTOR_SIZE], - tvUbTensor, - (uint64_t)0, - curRowNum, - AscendC::BinaryRepeatParams( - 1, 1, 0, embedRound / FLOAT_BLOCK_SIZE, embedRound / FLOAT_BLOCK_SIZE, 1)); - } - if (embed % FLOAT_VECTOR_SIZE > 0) { - SetMask(embed % FLOAT_VECTOR_SIZE); - AscendC::Mul( - goUbTensor32[oUbOffset + embed / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - goUbTensor32[oUbOffset + embed / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - tvUbTensor, - (uint64_t)0, - curRowNum, - AscendC::BinaryRepeatParams( - 1, 1, 0, embedRound / FLOAT_BLOCK_SIZE, embedRound / FLOAT_BLOCK_SIZE, 1)); - - AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); - } - AscendC::PipeBarrier(); - // *** go = lo + go - AscendC::Add( - goUbTensor32[oUbOffset], - goUbTensor32[oUbOffset], - loUbTensor[oUbOffset], - (uint64_t)0, - (curRowNum * embedRound + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, - AscendC::BinaryRepeatParams( - 1, 1, 1, 8, 8, 8)); - - AscendC::PipeBarrier(); - } else { - // *** gl = ll - if (rowLoopIdx == 0) { - AscendC::DataCopy( - glUbTensor, - llUbTensor[llUbOffsetCurCycle], - AscendC::DataCopyParams( - 1, 64 / FLOAT_BLOCK_SIZE, 0, 0)); - AscendC::PipeBarrier(); - } - AscendC::DataCopy( - goUbTensor32[oUbOffset], - gInput, - AscendC::DataCopyParams( - 1, curRowNum * embedRound / FLOAT_BLOCK_SIZE, 0, 0)); - AscendC::SetFlag(EVENT_ID0); - AscendC::WaitFlag(EVENT_ID0); - } - AscendC::SetFlag(oPingPangFlag); - - if (isLastNTile) { - AscendC::Brcb( - tvUbTensor.ReinterpretCast(), - glUbTensor.ReinterpretCast()[rowLoopIdx * ROW_WISE_CYCLE_TILE], - curRowNumRound / FLOAT_BLOCK_SIZE, - AscendC::BrcbRepeatParams(1, 8)); - AscendC::PipeBarrier(); - // *** go = go / gl_block - AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); - for (uint32_t divIdx = 0; divIdx < embed / FLOAT_VECTOR_SIZE; ++divIdx) { - AscendC::Div( - goUbTensor32[oUbOffset + divIdx * FLOAT_VECTOR_SIZE], - goUbTensor32[oUbOffset + divIdx * FLOAT_VECTOR_SIZE], - tvUbTensor, - (uint64_t)0, - curRowNum, - AscendC::BinaryRepeatParams( - 1, 1, 0, embedRound / FLOAT_BLOCK_SIZE, embedRound / FLOAT_BLOCK_SIZE, 1)); - } - - if (embed % FLOAT_VECTOR_SIZE > 0) { - SetMask(embed % FLOAT_VECTOR_SIZE); - AscendC::Div( - goUbTensor32[oUbOffset + embed / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - goUbTensor32[oUbOffset + embed / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - tvUbTensor, - (uint64_t)0, - curRowNum, - AscendC::BinaryRepeatParams( - 1, 1, 0, embedRound / FLOAT_BLOCK_SIZE, embedRound / FLOAT_BLOCK_SIZE, 1)); - AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); // fix hidden_size=96 - } - AscendC::PipeBarrier(); - - if (kvSplitCoreNum != 1) { - // log(l) - AscendC::Ln( - tvUbTensor, - tvUbTensor, - (uint64_t)0, - curRowNum, - AscendC::UnaryRepeatParams(1, 1, 8, 8)); - AscendC::PipeBarrier(); - AscendC::Brcb( - hmUbTensor.ReinterpretCast(), - gmUbTensor.ReinterpretCast()[rowLoopIdx * ROW_WISE_CYCLE_TILE], - curRowNumRound / FLOAT_BLOCK_SIZE, - AscendC::BrcbRepeatParams(1, 8)); - AscendC::PipeBarrier(); - // logf(lse_sum) + lse_max - AscendC::Add( - tvUbTensor, - tvUbTensor, - hmUbTensor, - (uint64_t)0, - curRowNum, - AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); - AscendC::PipeBarrier(); - - AscendC::SetFlag(EVENT_ID2); - AscendC::WaitFlag(EVENT_ID2); - AscendC::DataCopyPad(gl, tvUbTensor, - AscendC::DataCopyExtParams(curRowNum, 4, 0, (kvSplitCoreNum - 1) * 4, 0)); - - if (glFlag == 0) { - AscendC::SetFlag(EVENT_ID2); - glFlag = 1; - } - uint32_t srcGap = ((embed % 16 <= 8) && (embed % 16 > 0)) ? 1 : 0; - AscendC::DataCopyPad(gOCoreTmp, goUbTensor32[oUbOffset], - AscendC::DataCopyExtParams(curRowNum, embed * 4, srcGap, (kvSplitCoreNum - 1) * embed * 4, 0)); - AscendC::SetFlag(EVENT_ID3); - AscendC::WaitFlag(EVENT_ID3); - AscendC::SetFlag(EVENT_ID1); - AscendC::WaitFlag(EVENT_ID1); - } else { - // *** go = castfp32to16(go) - if (std::is_same::value) { - AscendC::Cast( - goUbTensor16[oUbOffset * 2], goUbTensor32[oUbOffset], - AscendC::RoundMode::CAST_RINT, (uint64_t)0, - (curRowNum * embedRound + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, - AscendC::UnaryRepeatParams(1, 1, 4, 8)); - } else { - AscendC::Cast( - goUbTensor16[oUbOffset * 2], goUbTensor32[oUbOffset], - AscendC::RoundMode::CAST_NONE, (uint64_t)0, - (curRowNum * embedRound + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, - AscendC::UnaryRepeatParams(1, 1, 4, 8)); - } - AscendC::SetFlag(EVENT_ID0); - AscendC::WaitFlag(EVENT_ID0); - // ********************* move O to GM ************************ - if (tokenNumPerHead == 1) { - AscendC::DataCopyPad(gOutput, goUbTensor16[oUbOffset * 2], - AscendC::DataCopyExtParams(curRowNum, embed * 2, 0, 0, 0)); - } else { - uint32_t innerOGmOffset = 0; - uint32_t inner_go_ubuf_offset = oUbOffset * 2; - if (proTokenNum != 0) { - AscendC::DataCopyPad(gOutput[innerOGmOffset + proTokenIdx * strideQO], - goUbTensor16[inner_go_ubuf_offset], - AscendC::DataCopyExtParams(proTokenNum, - embed * 2, 0, (strideQO - embed) * 2, 0)); - innerOGmOffset += embed; - inner_go_ubuf_offset += proTokenNum * embed; - } - for (uint32_t qN_idx = 0; qN_idx < integralHeadNum; qN_idx++) { - AscendC::DataCopyPad(gOutput[innerOGmOffset], - goUbTensor16[inner_go_ubuf_offset], - AscendC::DataCopyExtParams(tokenNumPerHead, - embed * 2, 0, (strideQO - embed) * 2, 0)); - innerOGmOffset += embed; - inner_go_ubuf_offset += tokenNumPerHead * embed; - } - if (epiTokenNum != 0) { - AscendC::DataCopyPad(gOutput[innerOGmOffset], - goUbTensor16[inner_go_ubuf_offset], - AscendC::DataCopyExtParams(epiTokenNum, - embed * 2, 0, (strideQO - embed) * 2, 0)); - } - } - } - } else if (needRowLoop) { - AscendC::SetFlag(EVENT_ID5); - AscendC::WaitFlag(EVENT_ID5); - AscendC::DataCopy( - gUpdate, - goUbTensor32[oUbOffset], - AscendC::DataCopyParams( - 1, curRowNum * embedRound / FLOAT_BLOCK_SIZE, 0, 0)); - } - AscendC::SetFlag(oPingPangFlag + 4); - if (needRowLoop) { - oPingPangFlag = 1 - oPingPangFlag; - } - } - - CATLASS_DEVICE - void operator()( - AscendC::GlobalTensor gInput, - AscendC::GlobalTensor gUpdate, - AscendC::GlobalTensor gOutput, - AscendC::GlobalTensor gOCoreTmp, - AscendC::GlobalTensor gl, - const LayoutInput &layoutInput, - const LayoutOutput &layoutOutput, - const LayoutUpdate &layoutUpdate, - GemmCoord actualBlockShape, - uint32_t nIdx, - uint32_t isLastNTile, - uint32_t curHeadNum, - uint32_t rescaleOPingPongFlag, - uint32_t &glFlag) - { - uint32_t tokenNumPerHead = layoutOutput.shape(0); - uint32_t embed = layoutInput.shape(1); - uint32_t rowActual = actualBlockShape.m(); // curHeadNum * tokenNumPerHead - uint32_t columnActual = actualBlockShape.n(); // embed - - uint32_t subBlockIdx = AscendC::GetSubBlockIdx(); - uint32_t subBlockNum = AscendC::GetSubBlockNum(); - - uint32_t curHeadSplitSubBlock = curHeadNum / subBlockNum; - uint32_t curHeadThisSubBlock = (subBlockIdx == 0) ? curHeadSplitSubBlock : (curHeadNum - curHeadSplitSubBlock); - - uint32_t rowActualThisSubBlock = curHeadThisSubBlock * tokenNumPerHead; - uint32_t rowOffsetSubBlock = subBlockIdx * curHeadSplitSubBlock * tokenNumPerHead; - uint32_t outOffsetSubBlock = subBlockIdx * curHeadSplitSubBlock * embed; - - if (rowActualThisSubBlock > 0) { - uint32_t rowLoop = (rowActualThisSubBlock + ROW_WISE_CYCLE_TILE - 1) / ROW_WISE_CYCLE_TILE; - uint32_t needRowLoop = (rowLoop > 1) ? 1 : 0; - // The rows of each cycle consist of multiple heads with several tokens. - // There are several integral heads, one prologue head, one epilogue head. - uint32_t proTokenIdx = 0; // the token idx of the start token of the prologue part - uint32_t proTokenNum = 0; // the token num of the prologue part - uint32_t epiTokenNum = 0; // the token num of the epilogue part - uint32_t integralHeadNum = 0; // the number of integral heads within a cycle - for (uint32_t rowLoopIdx = 0; rowLoopIdx < rowLoop; rowLoopIdx++) { - uint32_t rowOffsetLoop = rowLoopIdx * ROW_WISE_CYCLE_TILE; - uint32_t rowOffsetCurCycle = rowOffsetSubBlock + rowOffsetLoop; - uint32_t rowActualCurCycle = (rowLoopIdx == (rowLoop - 1)) ? - rowActualThisSubBlock - rowLoopIdx * ROW_WISE_CYCLE_TILE : - ROW_WISE_CYCLE_TILE; - int64_t offsetInput = layoutInput.GetOffset(MatrixCoord(rowOffsetCurCycle, 0)); - auto gInputThisCurCycle = gInput[offsetInput]; - auto layoutInputCurCycle = layoutInput.GetTileLayout(MatrixCoord(rowActualCurCycle, columnActual)); - int64_t offsetOutput = rowLoopIdx * ROW_WISE_CYCLE_TILE / tokenNumPerHead * embed + outOffsetSubBlock; - auto gOutputCurCycle = gOutput[offsetOutput]; - auto layoutOutputCurCycle = layoutOutput; - int64_t offsetUpdate = layoutUpdate.GetOffset(MatrixCoord(rowOffsetCurCycle, 0)); - auto gUpdateCurCycle = gUpdate[offsetUpdate]; - auto layoutUpdateCurCycle = layoutUpdate.GetTileLayout(MatrixCoord(rowActualCurCycle, columnActual)); - proTokenIdx = epiTokenNum; - proTokenNum = (tokenNumPerHead - epiTokenNum) % tokenNumPerHead; - integralHeadNum = (rowActualCurCycle - proTokenNum) / tokenNumPerHead; - epiTokenNum = rowActualCurCycle - proTokenNum - integralHeadNum * tokenNumPerHead; - SubCoreCompute(gInputThisCurCycle, gUpdateCurCycle, gOutputCurCycle, - gOCoreTmp[rowOffsetLoop * embed * kvSplitCoreNum], - gl[rowOffsetLoop * kvSplitCoreNum], - layoutInputCurCycle, layoutOutputCurCycle, layoutUpdateCurCycle, - nIdx, isLastNTile, needRowLoop, rowLoopIdx, - proTokenIdx, proTokenNum, epiTokenNum, integralHeadNum, rescaleOPingPongFlag, glFlag); - } - } - } - -private: - uint32_t kvSplitCoreNum = 1; - uint32_t oPingPangFlag = 0; - AscendC::LocalTensor goUbTensor16; - AscendC::LocalTensor loUbTensor; - AscendC::LocalTensor dmUbTensor; - AscendC::LocalTensor llUbTensor; - AscendC::LocalTensor glUbTensor; - AscendC::LocalTensor tvUbTensor; - AscendC::LocalTensor goUbTensor32; - AscendC::LocalTensor hmUbTensor; - AscendC::LocalTensor gmUbTensor; -}; - -} // namespace Catlass::Epilogue::Block - -#endif // CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_MLA_RESCALE_O_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_mla_tp1_rescale_o.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_mla_tp1_rescale_o.hpp deleted file mode 100644 index 5ac4da6989..0000000000 --- a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_mla_tp1_rescale_o.hpp +++ /dev/null @@ -1,341 +0,0 @@ -/** - * Copyright (c) 2025 Huawei Technologies Co., Ltd. - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of - * CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ - -#ifndef CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_MLA_TP1_RESCALE_O_HPP -#define CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_MLA_TP1_RESCALE_O_HPP - -#include "catlass/kda_catlass.hpp" -#include "catlass/arch/kda_resource.hpp" -#include "catlass/epilogue/kda_dispatch_policy.hpp" -#include "catlass/epilogue/tile/kda_epilogue_tile_copy.hpp" -#include "catlass/kda_gemm_coord.hpp" -#include "catlass/kda_matrix_coord.hpp" - -namespace Catlass::Epilogue::Block { - -template -class BlockEpilogue { -public: - // Type aliases - using DispatchPolicy = EpilogueAtlasA2MLATP1RescaleO; - using ArchTag = typename DispatchPolicy::ArchTag; - - using ElementOutput = typename OutputType_::Element; - using ElementUpdate = typename UpdateType_::Element; - using ElementInput = typename InputType_::Element; - - using LayoutOutput = typename OutputType_::Layout; - using LayoutUpdate = typename UpdateType_::Layout; - using LayoutInput = typename InputType_::Layout; - - static constexpr uint32_t HALF_ELENUM_PER_BLK = 16; - static constexpr uint32_t HALF_ELENUM_PER_VECCALC = 128; - static constexpr uint32_t FLOAT_ELENUM_PER_VECCALC = 64; - static constexpr uint32_t HALF_ELENUM_PER_LINE = 256; - static constexpr uint32_t FLOAT_ELENUM_PER_LINE = 128; - static constexpr uint32_t MULTIPLIER = 2; - static constexpr uint32_t FLOAT_BLOCK_SIZE = 8; - static constexpr uint32_t FLOAT_VECTOR_SIZE = 64; - static constexpr uint32_t UB_UINT8_LINE_SIZE = 512; - static constexpr uint32_t UB_UINT8_BLOCK_SIZE_MLA = 16384; - static constexpr uint32_t ROW_WISE_CYCLE_TILE = 8; - static constexpr uint32_t HALF_DM_UB_SIZE = 64; - static constexpr uint32_t HALF_LL_UB_SIZE = 256; - static constexpr uint32_t VECTOR_SIZE = 128; - static constexpr uint32_t NUM4 = 4; - - CATLASS_DEVICE - BlockEpilogue(Arch::Resource &resource, uint32_t kvSplitCoreNum_ = 1) - { - // Allocate UB space - constexpr uint32_t LO_UB_TENSOR_OFFSET = 4 * UB_UINT8_BLOCK_SIZE_MLA; - constexpr uint32_t DM_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA + 6 * UB_UINT8_LINE_SIZE; - constexpr uint32_t GL_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA + 16 * UB_UINT8_LINE_SIZE; - constexpr uint32_t GO_UB_TENSOR_OFFSET = 8 * UB_UINT8_BLOCK_SIZE_MLA; - constexpr uint32_t TV_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE_MLA; - constexpr uint32_t HM_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA + 1 * UB_UINT8_LINE_SIZE; - constexpr uint32_t GM_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA + 14 * UB_UINT8_LINE_SIZE; - - kvSplitCoreNum = kvSplitCoreNum_; - loUbTensor = resource.ubBuf.template GetBufferByByte(LO_UB_TENSOR_OFFSET); - dmUbTensor = resource.ubBuf.template GetBufferByByte(DM_UB_TENSOR_OFFSET); - glUbTensor = resource.ubBuf.template GetBufferByByte(GL_UB_TENSOR_OFFSET); - tvUbTensor = resource.ubBuf.template GetBufferByByte(TV_UB_TENSOR_OFFSET); - goUbTensor16 = resource.ubBuf.template GetBufferByByte(GO_UB_TENSOR_OFFSET); - goUbTensor32 = resource.ubBuf.template GetBufferByByte(GO_UB_TENSOR_OFFSET); - hmUbTensor = resource.ubBuf.template GetBufferByByte(HM_UB_TENSOR_OFFSET); - gmUbTensor = resource.ubBuf.template GetBufferByByte(GM_UB_TENSOR_OFFSET); - } - - CATLASS_DEVICE - ~BlockEpilogue() {} - - CATLASS_DEVICE - void SetMask(int32_t len) - { - uint64_t mask = 0; - uint64_t one = 1; - uint64_t temp = len % FLOAT_VECTOR_SIZE; - for (int64_t i = 0; i < temp; i++) { - mask |= one << i; - } - - if (len == VECTOR_SIZE) { - AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); - } else if (len >= FLOAT_VECTOR_SIZE) { - AscendC::SetVectorMask(mask, (uint64_t)-1); - } else { - AscendC::SetVectorMask(0x0, mask); - } - } - - CATLASS_DEVICE - void SetkvSplitCoreNum(uint32_t kvSplitCoreNum_) - { - kvSplitCoreNum = kvSplitCoreNum_; - } - - CATLASS_DEVICE - void SubCoreCompute(AscendC::GlobalTensor gInput, AscendC::GlobalTensor gUpdate, - AscendC::GlobalTensor gOutput, AscendC::GlobalTensor gOCoreTmp, - AscendC::GlobalTensor gl, const LayoutInput &layoutInput, - const LayoutOutput &layoutOutput, const LayoutUpdate &layoutUpdate, uint32_t nIdx, - uint32_t isLastNTile, uint32_t needRowLoop, uint32_t rowLoopIdx, uint32_t rescaleOPingPongFlag, - uint32_t &glFlag) - { - uint32_t curRowNum = layoutInput.shape(0); - uint32_t embed = layoutInput.shape(1); - uint32_t embedRound = layoutInput.stride(0); - uint32_t curRowNumRound = RoundUp(curRowNum); - uint64_t dmUbOffsetCurCycle = - (uint64_t)(rescaleOPingPongFlag * HALF_DM_UB_SIZE + rowLoopIdx * ROW_WISE_CYCLE_TILE); - uint32_t oUbOffset = oPingPangFlag * ROW_WISE_CYCLE_TILE * embedRound; - AscendC::WaitFlag(oPingPangFlag); - if (nIdx != NUM4) { - AscendC::DataCopy(loUbTensor[oUbOffset], gInput, - AscendC::DataCopyParams(1, curRowNum * embedRound / FLOAT_BLOCK_SIZE, 0, 0)); - AscendC::SetFlag(EVENT_ID0); - AscendC::WaitFlag(EVENT_ID0); - } - AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); - AscendC::WaitFlag(oPingPangFlag + 4); - if (nIdx != NUM4) { - AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); - AscendC::Brcb(tvUbTensor.ReinterpretCast(), - dmUbTensor[dmUbOffsetCurCycle].ReinterpretCast(), curRowNumRound / FLOAT_BLOCK_SIZE, - AscendC::BrcbRepeatParams(1, 8)); - AscendC::PipeBarrier(); - if (needRowLoop) { - AscendC::DataCopy(goUbTensor32[oUbOffset], gUpdate, - AscendC::DataCopyParams(1, curRowNum * embedRound / FLOAT_BLOCK_SIZE, 0, 0)); - AscendC::SetFlag(EVENT_ID0); - AscendC::WaitFlag(EVENT_ID0); - } - // *** go = go * dm_block - AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); - for (uint32_t vmul_idx = 0; vmul_idx < embed / FLOAT_VECTOR_SIZE; ++vmul_idx) { - AscendC::Mul(goUbTensor32[oUbOffset + vmul_idx * FLOAT_VECTOR_SIZE], - goUbTensor32[oUbOffset + vmul_idx * FLOAT_VECTOR_SIZE], - tvUbTensor, (uint64_t)0, - curRowNum, - AscendC::BinaryRepeatParams(1, 1, 0, embedRound / FLOAT_BLOCK_SIZE, - embedRound / FLOAT_BLOCK_SIZE, 1)); - } - if (embed % FLOAT_VECTOR_SIZE > 0) { - SetMask(embed % FLOAT_VECTOR_SIZE); - AscendC::Mul(goUbTensor32[oUbOffset + embed / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - goUbTensor32[oUbOffset + embed / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - tvUbTensor, - (uint64_t)0, curRowNum, - AscendC::BinaryRepeatParams(1, 1, 0, embedRound / FLOAT_BLOCK_SIZE, - embedRound / FLOAT_BLOCK_SIZE, 1)); - AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); - } - AscendC::PipeBarrier(); - // *** go = lo + go - AscendC::Add(goUbTensor32[oUbOffset], goUbTensor32[oUbOffset], - loUbTensor[oUbOffset], (uint64_t)0, - (curRowNum * embedRound + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, - AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); - AscendC::PipeBarrier(); - } else { - // *** gl = ll - AscendC::DataCopy(goUbTensor32[oUbOffset], gInput, - AscendC::DataCopyParams(1, curRowNum * embedRound / FLOAT_BLOCK_SIZE, 0, 0)); - AscendC::SetFlag(EVENT_ID0); - AscendC::WaitFlag(EVENT_ID0); - } - AscendC::SetFlag(oPingPangFlag); - - if (isLastNTile) { - // *** gl_block = expand_to_block(gl), 存放于 tv - AscendC::Brcb(tvUbTensor.ReinterpretCast(), - glUbTensor.ReinterpretCast()[rowLoopIdx * ROW_WISE_CYCLE_TILE], - curRowNumRound / FLOAT_BLOCK_SIZE, - AscendC::BrcbRepeatParams(1, 8)); - AscendC::PipeBarrier(); - // *** go = go / gl_block - AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); - for (uint32_t vdiv_idx = 0; vdiv_idx < embed / FLOAT_VECTOR_SIZE; ++vdiv_idx) { - AscendC::Div(goUbTensor32[oUbOffset + vdiv_idx * FLOAT_VECTOR_SIZE], - goUbTensor32[oUbOffset + vdiv_idx * FLOAT_VECTOR_SIZE], - tvUbTensor, (uint64_t)0, - curRowNum, - AscendC::BinaryRepeatParams(1, 1, 0, embedRound / FLOAT_BLOCK_SIZE, - embedRound / FLOAT_BLOCK_SIZE, 1)); - } - if (embed % FLOAT_VECTOR_SIZE > 0) { - SetMask(embed % FLOAT_VECTOR_SIZE); - AscendC::Div(goUbTensor32[oUbOffset + embed / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - goUbTensor32[oUbOffset + embed / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - tvUbTensor, - (uint64_t)0, curRowNum, - AscendC::BinaryRepeatParams(1, 1, 0, embedRound / FLOAT_BLOCK_SIZE, - embedRound / FLOAT_BLOCK_SIZE, 1)); - AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); - } - AscendC::PipeBarrier(); - - if (kvSplitCoreNum != 1) { - // log(l) - AscendC::Ln( - tvUbTensor, - tvUbTensor, - (uint64_t)0, - curRowNum, - AscendC::UnaryRepeatParams(1, 1, 8, 8)); - AscendC::PipeBarrier(); - AscendC::Brcb( - hmUbTensor.ReinterpretCast(), - gmUbTensor.ReinterpretCast()[rowLoopIdx * ROW_WISE_CYCLE_TILE], - curRowNumRound / FLOAT_BLOCK_SIZE, - AscendC::BrcbRepeatParams(1, 8)); - AscendC::PipeBarrier(); - // logf(lse_sum) + lse_max - AscendC::Add( - tvUbTensor, - tvUbTensor, - hmUbTensor, - (uint64_t)0, - curRowNum, - AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); - AscendC::PipeBarrier(); - - AscendC::SetFlag(EVENT_ID2); - AscendC::WaitFlag(EVENT_ID2); - AscendC::DataCopyPad(gl, tvUbTensor, - AscendC::DataCopyExtParams(curRowNum, 4, 0, (kvSplitCoreNum - 1) * 4, 0)); - - if (glFlag == 0) { - AscendC::SetFlag(EVENT_ID2); - glFlag = 1; - } - uint32_t srcGap = ((embed % 16 <= 8) && (embed % 16 > 0)) ? 1 : 0; - AscendC::DataCopyPad(gOCoreTmp, goUbTensor32[oUbOffset], - AscendC::DataCopyExtParams(curRowNum, embed * 4, srcGap, (kvSplitCoreNum - 1) * embed * 4, 0)); - AscendC::SetFlag(EVENT_ID3); - AscendC::WaitFlag(EVENT_ID3); - AscendC::SetFlag(EVENT_ID1); - AscendC::WaitFlag(EVENT_ID1); - } else { - // *** go = castfp32to16(go) - if (std::is_same::value) { - AscendC::Cast( - goUbTensor16[oUbOffset * 2], goUbTensor32[oUbOffset], - AscendC::RoundMode::CAST_RINT, (uint64_t)0, - (curRowNum * embedRound + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, - AscendC::UnaryRepeatParams(1, 1, 4, 8)); - } else { - AscendC::Cast( - goUbTensor16[oUbOffset * 2], goUbTensor32[oUbOffset], - AscendC::RoundMode::CAST_NONE, (uint64_t)0, - (curRowNum * embedRound + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, - AscendC::UnaryRepeatParams(1, 1, 4, 8)); - } - AscendC::SetFlag(EVENT_ID0); - AscendC::WaitFlag(EVENT_ID0); - - // ********************* move O to GM ************************ - AscendC::DataCopyPad(gOutput, goUbTensor16[oUbOffset * 2], - AscendC::DataCopyExtParams(curRowNum, embed * 2, 0, 0, 0)); - } - } else if (needRowLoop) { - AscendC::SetFlag(EVENT_ID5); - AscendC::WaitFlag(EVENT_ID5); - AscendC::DataCopy(gUpdate, goUbTensor32[oUbOffset], - AscendC::DataCopyParams(1, curRowNum * embedRound / FLOAT_BLOCK_SIZE, 0, 0)); - } - AscendC::SetFlag(oPingPangFlag + 4); - oPingPangFlag = 1 - oPingPangFlag; - } - - CATLASS_DEVICE - void operator()(AscendC::GlobalTensor gInput, AscendC::GlobalTensor gUpdate, - AscendC::GlobalTensor gOutput, AscendC::GlobalTensor gOCoreTmp, - AscendC::GlobalTensor gl, const LayoutInput &layoutInput, - const LayoutUpdate &layoutUpdate, const LayoutOutput &layoutOutput, GemmCoord actualBlockShape, - uint32_t nIdx, uint32_t isLastNTile, uint32_t rescaleOPingPongFlag, uint32_t &glFlag) - { - uint32_t embed = layoutInput.shape(1); - uint32_t rowActual = actualBlockShape.m(); - uint32_t columnActual = actualBlockShape.n(); - - uint32_t subBlockIdx = AscendC::GetSubBlockIdx(); - uint32_t subBlockNum = AscendC::GetSubBlockNum(); - - uint32_t curRowSplitSubBlock = rowActual / subBlockNum; - uint32_t rowActualThisSubBlock = (subBlockIdx == 0) ? curRowSplitSubBlock : (rowActual - curRowSplitSubBlock); - uint32_t rowOffsetSubBlock = subBlockIdx * curRowSplitSubBlock; - - if (rowActualThisSubBlock > 0) { - uint32_t rowLoop = (rowActualThisSubBlock + ROW_WISE_CYCLE_TILE - 1) / ROW_WISE_CYCLE_TILE; - uint32_t needRowLoop = (rowLoop > 1) ? 1 : 0; - for (uint32_t rowLoopIdx = 0; rowLoopIdx < rowLoop; rowLoopIdx++) { - uint32_t rowOffsetLoop = rowLoopIdx * ROW_WISE_CYCLE_TILE; - uint32_t rowOffsetCurCycle = rowOffsetSubBlock + rowOffsetLoop; - uint32_t rowActualCurCycle = - (rowLoopIdx == (rowLoop - 1)) ? rowActualThisSubBlock - rowOffsetLoop : ROW_WISE_CYCLE_TILE; - int64_t offsetInput = layoutInput.GetOffset(MatrixCoord(rowOffsetCurCycle, 0)); - auto gInputThisCurCycle = gInput[offsetInput]; - auto layoutInputCurCycle = layoutInput.GetTileLayout(MatrixCoord(rowActualCurCycle, columnActual)); - - int64_t offsetUpdate = layoutUpdate.GetOffset(MatrixCoord(rowOffsetCurCycle, 0)); - auto gUpdateCurCycle = gUpdate[offsetUpdate]; - auto layoutUpdateCurCycle = layoutUpdate.GetTileLayout(MatrixCoord(rowActualCurCycle, columnActual)); - - int64_t offsetOutput = layoutOutput.GetOffset(MatrixCoord(rowOffsetCurCycle, 0)); - auto gOutputCurCycle = gOutput[offsetOutput]; - auto layoutOutputCurCycle = layoutOutput.GetTileLayout(MatrixCoord(rowActualCurCycle, columnActual)); - - SubCoreCompute(gInputThisCurCycle, gUpdateCurCycle, gOutputCurCycle, - gOCoreTmp[rowOffsetLoop * embed * kvSplitCoreNum], - gl[rowOffsetLoop * kvSplitCoreNum], layoutInputCurCycle, - layoutOutputCurCycle, layoutUpdateCurCycle, nIdx, isLastNTile, needRowLoop, rowLoopIdx, - rescaleOPingPongFlag, glFlag); - } - } - } - -private: - uint32_t kvSplitCoreNum = 1; - uint32_t oPingPangFlag = 0; - AscendC::LocalTensor loUbTensor; - AscendC::LocalTensor dmUbTensor; - AscendC::LocalTensor glUbTensor; - AscendC::LocalTensor tvUbTensor; - AscendC::LocalTensor goUbTensor16; - AscendC::LocalTensor goUbTensor32; - AscendC::LocalTensor hmUbTensor; - AscendC::LocalTensor gmUbTensor; -}; - -} // namespace Catlass::Epilogue::Block - -#endif // CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_MLA_TP1_RESCALE_O_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_online_softmax_no_mask.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_online_softmax_no_mask.hpp deleted file mode 100644 index a783da27e3..0000000000 --- a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_online_softmax_no_mask.hpp +++ /dev/null @@ -1,877 +0,0 @@ -/** - * Copyright (c) 2025 Huawei Technologies Co., Ltd. - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of - * CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ - -#ifndef CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_ONLINE_SOFTMAX_NO_MASK_HPP -#define CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_ONLINE_SOFTMAX_NO_MASK_HPP - -#include "catlass/kda_catlass.hpp" -#include "catlass/arch/kda_cross_core_sync.hpp" -#include "catlass/arch/kda_resource.hpp" -#include "catlass/epilogue/kda_dispatch_policy.hpp" -#include "catlass/epilogue/tile/kda_epilogue_tile_copy.hpp" -#include "catlass/kda_gemm_coord.hpp" -#include "catlass/kda_matrix_coord.hpp" - -namespace Catlass::Epilogue::Block { - -template < - class OutputType_, - class InputType_, - class MaskType_> -class BlockEpilogue< - EpilogueAtlasA2OnlineSoftmax, - OutputType_, - InputType_, - MaskType_> { -public: - using DispatchPolicy = EpilogueAtlasA2OnlineSoftmax; - using ArchTag = typename DispatchPolicy::ArchTag; - using ElementOutput = typename OutputType_::Element; - using ElementInput = typename InputType_::Element; - using ElementMask = typename MaskType_::Element; - - using LayoutOutput = typename OutputType_::Layout; - using LayoutInput = typename InputType_::Layout; - using LayoutMask = typename MaskType_::Layout; - - static constexpr uint32_t FLOAT_BLOCK_SIZE = 8; - static constexpr uint32_t FLOAT_VECTOR_SIZE = 64; - static constexpr uint32_t HALF_VECTOR_SIZE = 128; - static constexpr uint32_t BLOCK_SIZE = 16; - static constexpr uint32_t UB_UINT8_VECTOR_SIZE = 1024; - static constexpr uint32_t UB_UINT8_BLOCK_SIZE = 16384; - static constexpr uint32_t VECTOR_SIZE = 128; - static constexpr uint32_t MAX_UB_S_ELEM_NUM = 8192; - - static constexpr uint32_t REDUCE_UB_SIZE = 1024; - static constexpr uint32_t ROW_OPS_SPEC_MASK_32 = 32; - static constexpr uint32_t ROW_OPS_SPEC_MASK_4 = 4; - static constexpr uint32_t MAX_ROW_NUM_SUB_CORE = 128; - static constexpr int64_t UB_FLOAT_LINE_SIZE = 64; - enum class MaskCategory { - NO_MASK = 0, - CAUSAL_MASK = 1 - }; - CATLASS_DEVICE - BlockEpilogue(Arch::Resource &resource, float scaleValue_) - { - // Allocate UB space - constexpr uint32_t LS_UB_TENSOR_OFFSET = 0; - constexpr uint32_t LP_UB_TENSOR_OFFSET = 4 * UB_UINT8_BLOCK_SIZE; - constexpr uint32_t MASK32_UB_TENSOR_OFFSET = 4 * UB_UINT8_BLOCK_SIZE; - - constexpr uint32_t TV_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE; - constexpr uint32_t LM_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 8 * UB_UINT8_VECTOR_SIZE; - - constexpr uint32_t HM_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 9 * UB_UINT8_VECTOR_SIZE; - constexpr uint32_t GM_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 10 * UB_UINT8_VECTOR_SIZE; - constexpr uint32_t LL_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 11 * UB_UINT8_VECTOR_SIZE; - constexpr uint32_t GL_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 12 * UB_UINT8_VECTOR_SIZE; - constexpr uint32_t DM_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 13 * UB_UINT8_VECTOR_SIZE; - - constexpr uint32_t MASK_UB_TENSOR_OFFSET = 11 * UB_UINT8_BLOCK_SIZE; - - scaleValue = scaleValue_; - lsUbTensor = resource.ubBuf.template GetBufferByByte(LS_UB_TENSOR_OFFSET); - lpUbTensor = resource.ubBuf.template GetBufferByByte(LP_UB_TENSOR_OFFSET); - maskUbTensor = resource.ubBuf.template GetBufferByByte(MASK_UB_TENSOR_OFFSET); - maskUbTensor32 = resource.ubBuf.template GetBufferByByte(MASK32_UB_TENSOR_OFFSET); - lmUbTensor = resource.ubBuf.template GetBufferByByte(LM_UB_TENSOR_OFFSET); - hmUbTensor = resource.ubBuf.template GetBufferByByte(HM_UB_TENSOR_OFFSET); - gmUbTensor = resource.ubBuf.template GetBufferByByte(GM_UB_TENSOR_OFFSET); - dmUbTensor = resource.ubBuf.template GetBufferByByte(DM_UB_TENSOR_OFFSET); - llUbTensor = resource.ubBuf.template GetBufferByByte(LL_UB_TENSOR_OFFSET); - tvUbTensor = resource.ubBuf.template GetBufferByByte(TV_UB_TENSOR_OFFSET); - glUbTensor = resource.ubBuf.template GetBufferByByte(GL_UB_TENSOR_OFFSET); - } - - CATLASS_DEVICE - ~BlockEpilogue() {} - - template - CATLASS_DEVICE T Min(T a, T b) - { - return (a > b) ? b : a; - } - - CATLASS_DEVICE - void SetVecMask(int32_t len) - { - uint64_t mask = 0; - uint64_t one = 1; - uint64_t temp = len % FLOAT_VECTOR_SIZE; - for (int64_t i = 0; i < temp; i++) { - mask |= one << i; - } - - if (len == VECTOR_SIZE || len == 0) { - AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); - } else if (len >= FLOAT_VECTOR_SIZE) { - AscendC::SetVectorMask(mask, (uint64_t)-1); - } else { - AscendC::SetVectorMask(0x0, mask); - } - } - - CATLASS_DEVICE - void SetBlockReduceMask(int32_t len) - { - if (len > 8 || len < 1) { - AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); - return; - } - uint64_t subMask = ((uint64_t)1 << len) - 1; - uint64_t maskValue = (subMask << 48) + (subMask << 32) + (subMask << 16) + subMask + (subMask << 56) + - (subMask << 40) + (subMask << 24) + (subMask << 8); - AscendC::SetVectorMask(maskValue, maskValue); - } - - CATLASS_DEVICE - void RowsumSPECTILE512(const AscendC::LocalTensor &srcUb, const AscendC::LocalTensor &rowsumUb, - const AscendC::LocalTensor &tvUbTensor, uint32_t numRowsRound, uint32_t numElems, - uint32_t numElemsAligned) - { - AscendC::BlockReduceSum(tvUbTensor, srcUb, numRowsRound * numElemsAligned / FLOAT_VECTOR_SIZE, 0, - 1, 1, 8); - AscendC::PipeBarrier(); - - AscendC::BlockReduceSum(tvUbTensor[REDUCE_UB_SIZE], tvUbTensor, - numRowsRound * numElemsAligned / FLOAT_BLOCK_SIZE / FLOAT_VECTOR_SIZE, 0, - 1, 1, 8); - AscendC::PipeBarrier(); - AscendC::BlockReduceSum(rowsumUb, tvUbTensor[REDUCE_UB_SIZE], - numRowsRound * numElemsAligned / FLOAT_VECTOR_SIZE / FLOAT_VECTOR_SIZE, 0, - 1, 1, 8); - AscendC::PipeBarrier(); - } - - CATLASS_DEVICE - void RowsumSPECTILE256(const AscendC::LocalTensor &srcUb, const AscendC::LocalTensor &rowsumUb, - const AscendC::LocalTensor &tvUbTensor, uint32_t numRowsRound, uint32_t numElems, - uint32_t numElemsAligned) - { - AscendC::BlockReduceSum(tvUbTensor, srcUb, numRowsRound * numElemsAligned / FLOAT_VECTOR_SIZE, 0, - 1, 1, 8); - AscendC::PipeBarrier(); - SetVecMask(ROW_OPS_SPEC_MASK_32); - AscendC::BlockReduceSum(tvUbTensor[REDUCE_UB_SIZE], tvUbTensor, numRowsRound, 0, 1, 1, 4); - AscendC::PipeBarrier(); - SetBlockReduceMask(ROW_OPS_SPEC_MASK_4); - AscendC::BlockReduceSum( - rowsumUb, tvUbTensor[REDUCE_UB_SIZE], - (numRowsRound * FLOAT_BLOCK_SIZE + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, 0, 1, 1, 8); - AscendC::PipeBarrier(); - AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); - } - - CATLASS_DEVICE - void RowsumTAILTILE(const AscendC::LocalTensor &srcUb, const AscendC::LocalTensor &rowsumUb, - const AscendC::LocalTensor &tvUbTensor, uint32_t numRowsRound, uint32_t numElems, - uint32_t numElemsAligned) - { - if (numElems >= FLOAT_VECTOR_SIZE) { - AscendC::BlockReduceSum(tvUbTensor, srcUb, numRowsRound, 0, 1, 1, - numElemsAligned / FLOAT_BLOCK_SIZE); - AscendC::PipeBarrier(); - AscendC::BlockReduceSum( - rowsumUb, tvUbTensor, (numRowsRound * FLOAT_BLOCK_SIZE + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, 0, - 1, 1, 8); - AscendC::PipeBarrier(); - for (uint64_t rowSumIdx = 1; rowSumIdx < (uint64_t)numElems / FLOAT_VECTOR_SIZE; ++rowSumIdx) { - AscendC::BlockReduceSum(tvUbTensor, srcUb[rowSumIdx * FLOAT_VECTOR_SIZE], numRowsRound, 0, - 1, 1, numElemsAligned / FLOAT_BLOCK_SIZE); - AscendC::PipeBarrier(); - AscendC::BlockReduceSum( - tvUbTensor[REDUCE_UB_SIZE], tvUbTensor, - (numRowsRound * FLOAT_BLOCK_SIZE + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, 0, 1, 1, 8); - AscendC::PipeBarrier(); - SetVecMask(numRowsRound); - AscendC::Add(rowsumUb, rowsumUb, tvUbTensor[REDUCE_UB_SIZE], (uint64_t)0, 1, - AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); - AscendC::PipeBarrier(); - AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); - } - } - if (numElems % FLOAT_VECTOR_SIZE > 0) { - SetVecMask(numElems % FLOAT_VECTOR_SIZE); - AscendC::BlockReduceSum(tvUbTensor, srcUb[numElems / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - numRowsRound, 0, 1, 1, numElemsAligned / FLOAT_BLOCK_SIZE); - AscendC::PipeBarrier(); - SetBlockReduceMask((numElems % FLOAT_VECTOR_SIZE + FLOAT_BLOCK_SIZE - 1) / FLOAT_BLOCK_SIZE); - if (numElems < FLOAT_VECTOR_SIZE) { - AscendC::BlockReduceSum( - rowsumUb, tvUbTensor, (numRowsRound * FLOAT_BLOCK_SIZE + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, - 0, 1, 1, 8); - AscendC::PipeBarrier(); - } else { - AscendC::BlockReduceSum( - tvUbTensor[REDUCE_UB_SIZE], tvUbTensor, - (numRowsRound * FLOAT_BLOCK_SIZE + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, 0, 1, 1, 8); - AscendC::PipeBarrier(); - SetVecMask(numRowsRound); - AscendC::Add(rowsumUb, rowsumUb, tvUbTensor[REDUCE_UB_SIZE], (uint64_t)0, 1, - AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); - AscendC::PipeBarrier(); - } - AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); - } - } - - CATLASS_DEVICE - void RowmaxSPECTILE512(const AscendC::LocalTensor &srcUb, const AscendC::LocalTensor &rowmaxUb, - const AscendC::LocalTensor &tvUbTensor, uint32_t numRowsRound, uint32_t numElems, - uint32_t numElemsAligned) - { - AscendC::BlockReduceMax(tvUbTensor, srcUb, numRowsRound * numElemsAligned / FLOAT_VECTOR_SIZE, 0, - 1, 1, 8); - AscendC::PipeBarrier(); - AscendC::BlockReduceMax(tvUbTensor[REDUCE_UB_SIZE], tvUbTensor, - numRowsRound * numElemsAligned / FLOAT_BLOCK_SIZE / FLOAT_VECTOR_SIZE, 0, - 1, 1, 8); - AscendC::PipeBarrier(); - AscendC::BlockReduceMax(rowmaxUb, tvUbTensor[REDUCE_UB_SIZE], - numRowsRound * numElemsAligned / FLOAT_VECTOR_SIZE / FLOAT_VECTOR_SIZE, 0, - 1, 1, 8); - AscendC::PipeBarrier(); - } - - CATLASS_DEVICE - void RowmaxSPECTILE256(const AscendC::LocalTensor &srcUb, const AscendC::LocalTensor &rowmaxUb, - const AscendC::LocalTensor &tvUbTensor, uint32_t numRowsRound, uint32_t numElems, - uint32_t numElemsAligned) - { - AscendC::BlockReduceMax(tvUbTensor, srcUb, numRowsRound * numElemsAligned / FLOAT_VECTOR_SIZE, 0, - 1, 1, 8); - AscendC::PipeBarrier(); - SetVecMask(ROW_OPS_SPEC_MASK_32); - AscendC::BlockReduceMax(tvUbTensor[REDUCE_UB_SIZE], tvUbTensor, numRowsRound, 0, 1, 1, 4); - AscendC::PipeBarrier(); - SetBlockReduceMask(ROW_OPS_SPEC_MASK_4); - AscendC::BlockReduceMax( - rowmaxUb, tvUbTensor[REDUCE_UB_SIZE], - (numRowsRound * FLOAT_BLOCK_SIZE + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, 0, 1, 1, 8); - AscendC::PipeBarrier(); - AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); - } - - CATLASS_DEVICE - void RowmaxTAILTILE(const AscendC::LocalTensor &srcUb, const AscendC::LocalTensor &rowmaxUb, - const AscendC::LocalTensor &tvUbTensor, uint32_t numRowsRound, uint32_t numElems, - uint32_t numElemsAligned) - { - if (numElems >= FLOAT_VECTOR_SIZE) { - AscendC::BlockReduceMax(tvUbTensor, srcUb, numRowsRound, 0, 1, 1, - numElemsAligned / FLOAT_BLOCK_SIZE); - AscendC::PipeBarrier(); - AscendC::BlockReduceMax( - rowmaxUb, tvUbTensor, (numRowsRound * FLOAT_BLOCK_SIZE + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, 0, - 1, 1, 8); - AscendC::PipeBarrier(); - for (uint64_t rowmax_idx = 1; rowmax_idx < (uint64_t)numElems / FLOAT_VECTOR_SIZE; ++rowmax_idx) { - AscendC::BlockReduceMax(tvUbTensor, srcUb[rowmax_idx * FLOAT_VECTOR_SIZE], numRowsRound, - 0, 1, 1, numElemsAligned / FLOAT_BLOCK_SIZE); - AscendC::PipeBarrier(); - AscendC::BlockReduceMax( - tvUbTensor[REDUCE_UB_SIZE], tvUbTensor, - (numRowsRound * FLOAT_BLOCK_SIZE + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, 0, 1, 1, 8); - AscendC::PipeBarrier(); - SetVecMask(numRowsRound); - AscendC::Max(rowmaxUb, rowmaxUb, tvUbTensor[REDUCE_UB_SIZE], (uint64_t)0, 1, - AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); - AscendC::PipeBarrier(); - AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); - } - } - if (numElems % FLOAT_VECTOR_SIZE > 0) { - SetVecMask(numElems % FLOAT_VECTOR_SIZE); - AscendC::BlockReduceMax(tvUbTensor, srcUb[numElems / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - numRowsRound, 0, 1, 1, numElemsAligned / FLOAT_BLOCK_SIZE); - AscendC::PipeBarrier(); - SetBlockReduceMask((numElems % FLOAT_VECTOR_SIZE + FLOAT_BLOCK_SIZE - 1) / FLOAT_BLOCK_SIZE); - if (numElems < FLOAT_VECTOR_SIZE) { - AscendC::BlockReduceMax( - rowmaxUb, tvUbTensor, (numRowsRound * FLOAT_BLOCK_SIZE + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, - 0, 1, 1, 8); - AscendC::PipeBarrier(); - } else { - AscendC::BlockReduceMax( - tvUbTensor[REDUCE_UB_SIZE], tvUbTensor, - (numRowsRound * FLOAT_BLOCK_SIZE + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, 0, 1, 1, 8); - AscendC::PipeBarrier(); - SetVecMask(numRowsRound); - AscendC::Max(rowmaxUb, rowmaxUb, tvUbTensor[REDUCE_UB_SIZE], (uint64_t)0, 1, - AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); - AscendC::PipeBarrier(); - } - AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); - } - } - - CATLASS_DEVICE - void CopySGmToUb( - AscendC::GlobalTensor gInput, - uint32_t sUbOffset, - uint32_t rowNumCurLoop, - uint32_t columnNumRound, - uint32_t columnNumPad) - { - // input S - AscendC::DataCopy( - lsUbTensor[sUbOffset], gInput, - AscendC::DataCopyParams( - rowNumCurLoop, columnNumRound / FLOAT_BLOCK_SIZE, - (columnNumPad - columnNumRound) / FLOAT_BLOCK_SIZE, 0)); - } - - CATLASS_DEVICE - void CopyMaskGmToUb( - AscendC::GlobalTensor gMask, - uint32_t columnNum, - uint32_t columnNumRound, - uint32_t maskStride, - uint32_t qSBlockSize, - uint32_t proTokenIdx, - uint32_t proTokenNum, - uint32_t integralHeadNum, - uint32_t epiTokenNum) - { - uint32_t innerUbRowOffset = 0; - if (proTokenNum != 0) { - AscendC::DataCopyPad( - maskUbTensor[innerUbRowOffset], gMask[proTokenIdx * maskStride], - AscendC::DataCopyExtParams( - proTokenNum, columnNum * 2, (maskStride - columnNum) * 2, 0, 0), - AscendC::DataCopyPadExtParams(false, 0, 0, 0)); - innerUbRowOffset += proTokenNum * columnNumRound; - } - for (uint32_t headIdx = 0; headIdx < integralHeadNum; headIdx++) { - AscendC::DataCopyPad( - maskUbTensor[innerUbRowOffset], gMask, - AscendC::DataCopyExtParams( - qSBlockSize, columnNum * 2, (maskStride - columnNum) * 2, 0, 0), - AscendC::DataCopyPadExtParams(false, 0, 0, 0)); - innerUbRowOffset += qSBlockSize * columnNumRound; - } - if (epiTokenNum != 0) { - AscendC::DataCopyPad( - maskUbTensor[innerUbRowOffset], gMask, - AscendC::DataCopyExtParams( - epiTokenNum, columnNum * 2, (maskStride - columnNum) * 2, 0, 0), - AscendC::DataCopyPadExtParams(false, 0, 0, 0)); - } - } - - CATLASS_DEVICE - void ScaleS( - uint32_t sUbOffset, - uint32_t rowNumCurLoop, - uint32_t columnNumRound) - { - // *** ls = scaleValue * ls - AscendC::Muls( - lsUbTensor[sUbOffset], lsUbTensor[sUbOffset], scaleValue, (uint64_t)0, - (rowNumCurLoop * columnNumRound + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, - AscendC::UnaryRepeatParams(1, 1, 8, 8)); - - AscendC::PipeBarrier(); - } - - CATLASS_DEVICE - void UpCastMask( - uint32_t rowNumCurLoop, - uint32_t columnNumRound) - { - AscendC::Cast( - maskUbTensor32, maskUbTensor, AscendC::RoundMode::CAST_NONE, (uint64_t)0, - (rowNumCurLoop * columnNumRound + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, - AscendC::UnaryRepeatParams(1, 1, 8, 4)); - AscendC::PipeBarrier(); - } - - CATLASS_DEVICE - void ApplyMask( - uint32_t sUbOffset, - uint32_t rowNumCurLoop, - uint32_t columnNumRound) - { - AscendC::Muls( - maskUbTensor32, maskUbTensor32, (float)-3e38, (uint64_t)0, - (rowNumCurLoop * columnNumRound + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, - AscendC::UnaryRepeatParams(1, 1, 8, 8)); - AscendC::PipeBarrier(); - AscendC::Add( - lsUbTensor[sUbOffset], lsUbTensor[sUbOffset], maskUbTensor32, - (uint64_t)0, (rowNumCurLoop * columnNumRound + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, - AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); - AscendC::PipeBarrier(); - } - - CATLASS_DEVICE - void CalcLocalRowMax( - uint32_t sUbOffset, - uint32_t rowNumCurLoopRound, - uint32_t columnNum, - uint32_t columnNumRound, - uint32_t rowOffset) - { - if (columnNum == 512) { - RowmaxSPECTILE512( - lsUbTensor[sUbOffset], lmUbTensor[rowOffset], - tvUbTensor, rowNumCurLoopRound, columnNum, columnNumRound); - } else if (columnNum == 256) { - RowmaxSPECTILE256( - lsUbTensor[sUbOffset], lmUbTensor[rowOffset], - tvUbTensor, rowNumCurLoopRound, columnNum, columnNumRound); - } else { - RowmaxTAILTILE( - lsUbTensor[sUbOffset], lmUbTensor[rowOffset], - tvUbTensor, rowNumCurLoopRound, columnNum, columnNumRound); - } - } - - CATLASS_DEVICE - void UpdateGlobalRowMax( - uint32_t rowNumCurLoop, - uint32_t rowNumCurLoopRound, - uint32_t columnNum, - uint32_t columnNumRound, - uint32_t dmUbOffsetCurCycle, - uint32_t rowOffset, - uint32_t isFirstStackTile) - { - if (isFirstStackTile) { - AscendC::DataCopy( - hmUbTensor[rowOffset], lmUbTensor[rowOffset], - AscendC::DataCopyParams(1, rowNumCurLoopRound / FLOAT_BLOCK_SIZE, 0, 0)); - AscendC::PipeBarrier(); - } else { - SetVecMask(rowNumCurLoop); - // *** hm = vmax(lm, gm) - AscendC::Max( - hmUbTensor[rowOffset], lmUbTensor[rowOffset], gmUbTensor[rowOffset], (uint64_t)0, - 1, AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); - - AscendC::PipeBarrier(); - // *** dm = gm - hm - AscendC::Sub( - dmUbTensor[dmUbOffsetCurCycle], - gmUbTensor[rowOffset], hmUbTensor[rowOffset], (uint64_t)0, 1, - AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); - - AscendC::PipeBarrier(); - // *** dm = exp(dm) - AscendC::Exp( - dmUbTensor[dmUbOffsetCurCycle], - dmUbTensor[dmUbOffsetCurCycle], - (uint64_t)0, 1, AscendC::UnaryRepeatParams(1, 1, 8, 8)); - } - AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); - AscendC::PipeBarrier(); - // *** gm = hm - AscendC::DataCopy(gmUbTensor[rowOffset], hmUbTensor[rowOffset], - AscendC::DataCopyParams(1, rowNumCurLoopRound / FLOAT_BLOCK_SIZE, 0, 0)); - AscendC::PipeBarrier(); - } - - CATLASS_DEVICE - void CalcExp( - uint32_t sUbOffset, - uint32_t rowNumCurLoop, - uint32_t rowNumCurLoopRound, - uint32_t columnNum, - uint32_t columnNumRound, - uint32_t rowOffset) - { - // *** hm_block = expand_to_block(hm), 存放于 tv - AscendC::Brcb( - tvUbTensor.template ReinterpretCast(), - hmUbTensor[rowOffset].template ReinterpretCast(), rowNumCurLoopRound / FLOAT_BLOCK_SIZE, - AscendC::BrcbRepeatParams(1, 8)); - AscendC::PipeBarrier(); - // *** ls = ls - hm_block - for (uint32_t subIdx = 0; subIdx < columnNum / FLOAT_VECTOR_SIZE; ++subIdx) { - AscendC::Sub( - lsUbTensor[sUbOffset][subIdx * FLOAT_VECTOR_SIZE], lsUbTensor[sUbOffset][subIdx * FLOAT_VECTOR_SIZE], - tvUbTensor, (uint64_t)0, rowNumCurLoop, - AscendC::BinaryRepeatParams(1, 1, 0, columnNumRound / FLOAT_BLOCK_SIZE, columnNumRound / FLOAT_BLOCK_SIZE, 1)); - } - if (columnNum % FLOAT_VECTOR_SIZE > 0) { - SetVecMask(columnNum % FLOAT_VECTOR_SIZE); - AscendC::Sub( - lsUbTensor[sUbOffset][columnNum / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - lsUbTensor[sUbOffset][columnNum / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], tvUbTensor, (uint64_t)0, rowNumCurLoop, - AscendC::BinaryRepeatParams(1, 1, 0, columnNumRound / FLOAT_BLOCK_SIZE, columnNumRound / FLOAT_BLOCK_SIZE, 1)); - AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); - } - AscendC::PipeBarrier(); - // *** ls = exp(ls) - AscendC::Exp(lsUbTensor[sUbOffset], lsUbTensor[sUbOffset], (uint64_t)0, - (rowNumCurLoop * columnNumRound + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, - AscendC::UnaryRepeatParams(1, 1, 8, 8)); - AscendC::PipeBarrier(); - } - - CATLASS_DEVICE - void CalcLocalRowSum( - uint32_t sUbOffset, - uint32_t rowNumCurLoopRound, - uint32_t columnNum, - uint32_t columnNumRound, - uint32_t rowOffset) - { - // *** ll = rowsum(ls32) - if (columnNum == 512) { - RowsumSPECTILE512( - lsUbTensor[sUbOffset], llUbTensor[rowOffset], tvUbTensor, - rowNumCurLoopRound, columnNum, columnNumRound); - } else if (columnNum == 256) { - RowsumSPECTILE256( - lsUbTensor[sUbOffset], llUbTensor[rowOffset], tvUbTensor, - rowNumCurLoopRound, columnNum, columnNumRound); - } else { - RowsumTAILTILE( - lsUbTensor[sUbOffset], llUbTensor[rowOffset], tvUbTensor, - rowNumCurLoopRound, columnNum, columnNumRound); - } - } - - CATLASS_DEVICE - void UpdateGlobalRowSum( - uint32_t sUbOffset, - uint32_t rowNumCurLoop, - uint32_t rowNumCurLoopRound, - uint32_t dmUbOffsetCurCycle, - uint32_t rowOffset, - uint32_t isFirstStackTile) - { - if (isFirstStackTile) { - // *** gl = ll - AscendC::DataCopy( - glUbTensor[rowOffset], llUbTensor[rowOffset], - AscendC::DataCopyParams(1, rowNumCurLoopRound / FLOAT_BLOCK_SIZE, 0, 0)); - AscendC::PipeBarrier(); - } else { - SetVecMask(rowNumCurLoop); - // *** gl = dm * gl - AscendC::Mul( - glUbTensor[rowOffset], dmUbTensor[dmUbOffsetCurCycle], - glUbTensor[rowOffset], (uint64_t)0, 1, AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); - AscendC::PipeBarrier(); - // *** gl = ll + gl - AscendC::Add(glUbTensor[rowOffset], glUbTensor[rowOffset], llUbTensor[rowOffset], (uint64_t)0, - 1, AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); - AscendC::PipeBarrier(); - AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); - } - } - - CATLASS_DEVICE - void DownCastP( - uint32_t sUbOffset, - uint32_t rowNumCurLoop, - uint32_t columnNumRound) - { - // *** lp = castfp32to16(ls) - if (std::is_same::value) { - AscendC::Cast( - lpUbTensor[sUbOffset], lsUbTensor[sUbOffset], AscendC::RoundMode::CAST_RINT, (uint64_t)0, - (rowNumCurLoop * columnNumRound + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, AscendC::UnaryRepeatParams(1, 1, 4, 8)); - } else { - AscendC::Cast( - lpUbTensor[sUbOffset], lsUbTensor[sUbOffset], AscendC::RoundMode::CAST_NONE, (uint64_t)0, - (rowNumCurLoop * columnNumRound + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, AscendC::UnaryRepeatParams(1, 1, 4, 8)); - } - } - - CATLASS_DEVICE - void CopyPUbToGm( - AscendC::GlobalTensor gOutput, - uint32_t sUbOffset, - uint32_t rowNumCurLoop, - uint32_t columnNumRound, - uint32_t columnNumPad) - { - if (columnNumRound == columnNumPad) { - AscendC::DataCopy( - gOutput, lpUbTensor[sUbOffset], - AscendC::DataCopyParams( - 1, rowNumCurLoop * columnNumRound / BLOCK_SIZE, - 0, 0)); - } else { - AscendC::DataCopy( - gOutput, lpUbTensor[sUbOffset], - AscendC::DataCopyParams( - rowNumCurLoop, columnNumRound / BLOCK_SIZE, - 0, (columnNumPad - columnNumRound) / BLOCK_SIZE)); - } - } - - template - CATLASS_DEVICE void SubCoreCompute(AscendC::GlobalTensor gOutput, - const LayoutOutput &layoutOutput, - uint32_t rowOffset, uint32_t isFirstStackTile, - uint32_t columnNumRound, uint32_t pingpongFlag, - uint32_t curStackTileMod) - { - uint32_t rowNumCurLoop = layoutOutput.shape(0); - uint32_t rowNumCurLoopRound = RoundUp(rowNumCurLoop, FLOAT_BLOCK_SIZE); - uint32_t columnNum = layoutOutput.shape(1); - // Align colNum to 16, for both float&half compute© - uint32_t columnNumPad = layoutOutput.stride(0); - uint32_t sUbOffset = pingpongFlag * MAX_UB_S_ELEM_NUM; - uint32_t dmUbOffsetCurCycle = curStackTileMod * MAX_ROW_NUM_SUB_CORE + rowOffset; - - CalcLocalRowMax(sUbOffset, rowNumCurLoopRound, columnNum, columnNumRound, rowOffset); - UpdateGlobalRowMax(rowNumCurLoop, rowNumCurLoopRound, columnNum, columnNumRound, dmUbOffsetCurCycle, rowOffset, isFirstStackTile); - - CalcExp(sUbOffset, rowNumCurLoop, rowNumCurLoopRound, columnNum, columnNumRound, rowOffset); - if constexpr (maskCat == MaskCategory::NO_MASK) { - AscendC::WaitFlag(pingpongFlag); - } - - DownCastP(sUbOffset, rowNumCurLoop, columnNumRound); - AscendC::SetFlag(pingpongFlag); - - CalcLocalRowSum(sUbOffset, rowNumCurLoopRound, columnNum, columnNumRound, rowOffset); - AscendC::SetFlag(pingpongFlag); - - AscendC::WaitFlag(pingpongFlag); - CopyPUbToGm(gOutput, sUbOffset, rowNumCurLoop, columnNumRound, columnNumPad); - if constexpr (maskCat == MaskCategory::NO_MASK) { - AscendC::SetFlag(pingpongFlag); - } else if constexpr (maskCat == MaskCategory::CAUSAL_MASK) { - AscendC::SetFlag(EVENT_ID0); - AscendC::SetFlag(EVENT_ID1); - } - UpdateGlobalRowSum(sUbOffset, rowNumCurLoop, rowNumCurLoopRound, dmUbOffsetCurCycle, rowOffset, isFirstStackTile); - } - - CATLASS_DEVICE - void operator()(AscendC::GlobalTensor gOutput, AscendC::GlobalTensor gInput, - const LayoutOutput &layoutOutput, const LayoutInput &layoutInput, GemmCoord actualBlockShape, - uint32_t isFirstStackTile, uint32_t qSBlockSize, uint32_t qNBlockSize, uint32_t curStackTileMod) - { - uint32_t rowNum = actualBlockShape.m(); - uint32_t columnNum = actualBlockShape.n(); - uint32_t columnNumRound = RoundUp(columnNum, BLOCK_SIZE); - uint32_t columnNumPad = layoutInput.stride(0); - - uint32_t subBlockIdx = AscendC::GetSubBlockIdx(); - uint32_t subBlockNum = AscendC::GetSubBlockNum(); - - uint32_t qNSplitSubBlock = qNBlockSize / subBlockNum; - uint32_t qNThisSubBlock = (qNBlockSize == 1) ? - 0 : - (subBlockIdx == 1) ? (qNBlockSize - qNSplitSubBlock) : - qNSplitSubBlock; - uint32_t rowSplitSubBlock = (qNBlockSize == 1) ? - (qSBlockSize / 2) : - (qSBlockSize * qNSplitSubBlock); - uint32_t rowActualThisSubBlock = (subBlockIdx == 1) ? - (rowNum - rowSplitSubBlock) : - rowSplitSubBlock; - uint32_t rowOffsetThisSubBlock = subBlockIdx * rowSplitSubBlock; - uint32_t maxRowNumPerLoop = MAX_UB_S_ELEM_NUM / columnNumRound; - uint32_t rowNumTile = RoundDown(maxRowNumPerLoop, FLOAT_BLOCK_SIZE); - uint32_t rowLoopNum = CeilDiv(rowActualThisSubBlock, rowNumTile); - uint32_t preLoad = 1; - - for (uint32_t rowLoopIdx = 0; rowLoopIdx < rowLoopNum + preLoad; rowLoopIdx++) { - if (rowLoopIdx < rowLoopNum) { - uint32_t pingpongFlag = rowLoopIdx % 2; - uint32_t rowOffsetCurLoop = rowLoopIdx * rowNumTile; - uint32_t rowOffsetIoGm = rowOffsetCurLoop + rowOffsetThisSubBlock; - uint32_t rowNumCurLoop = (rowLoopIdx == rowLoopNum - 1) ? - (rowActualThisSubBlock - rowOffsetCurLoop) : - rowNumTile; - // loop 0 mask load before cross core sync - - int64_t offsetInput = layoutInput.GetOffset(MatrixCoord(rowOffsetIoGm, 0)); - auto gInputCurLoop = gInput[offsetInput]; - - AscendC::WaitFlag(pingpongFlag); - CopySGmToUb(gInputCurLoop, (pingpongFlag * MAX_UB_S_ELEM_NUM), rowNumCurLoop, columnNumRound, columnNumPad); - AscendC::SetFlag(pingpongFlag); - } - if (rowLoopIdx >= preLoad) { - uint32_t delayedRowLoopIdx = rowLoopIdx - preLoad; - uint32_t pingpongFlag = delayedRowLoopIdx % 2; - uint32_t rowOffsetCurLoop = delayedRowLoopIdx * rowNumTile; - uint32_t rowOffsetIoGm = rowOffsetCurLoop + rowOffsetThisSubBlock; - uint32_t rowNumCurLoop = (delayedRowLoopIdx == rowLoopNum - 1) ? - (rowActualThisSubBlock - rowOffsetCurLoop) : - rowNumTile; - - int64_t offsetOutput = layoutOutput.GetOffset(MatrixCoord(rowOffsetIoGm, 0)); - auto gOutputCurLoop = gOutput[offsetOutput]; - auto layoutOutputCurLoop = layoutOutput.GetTileLayout(MatrixCoord(rowNumCurLoop, columnNum)); - AscendC::WaitFlag(pingpongFlag); - ScaleS((pingpongFlag * MAX_UB_S_ELEM_NUM), rowNumCurLoop, columnNumRound); - SubCoreCompute( - gOutputCurLoop, - layoutOutputCurLoop, - rowOffsetCurLoop, - isFirstStackTile, - columnNumRound, - pingpongFlag, - curStackTileMod); - } - } - } - - CATLASS_DEVICE - void operator()(AscendC::GlobalTensor gOutput, AscendC::GlobalTensor gInput, - AscendC::GlobalTensor gMask, - const LayoutOutput &layoutOutput, const LayoutInput &layoutInput, const LayoutInput &layoutMask, - GemmCoord actualBlockShape, - uint32_t isFirstStackTile, uint32_t qSBlockSize, uint32_t qNBlockSize, uint32_t curStackTileMod, - Arch::CrossCoreFlag qkReady) - { - uint32_t rowNum = actualBlockShape.m(); - uint32_t columnNum = actualBlockShape.n(); - uint32_t columnNumRound = RoundUp(columnNum, BLOCK_SIZE); - uint32_t columnNumPad = layoutInput.stride(0); - uint32_t maskStride = layoutMask.stride(0); - uint32_t subBlockIdx = AscendC::GetSubBlockIdx(); - uint32_t subBlockNum = AscendC::GetSubBlockNum(); - - uint32_t qNSplitSubBlock = qNBlockSize / subBlockNum; - uint32_t qNThisSubBlock = (qNBlockSize == 1) ? - 0 : - (subBlockIdx == 1) ? (qNBlockSize - qNSplitSubBlock) : - qNSplitSubBlock; - uint32_t rowSplitSubBlock = (qNBlockSize == 1) ? - (qSBlockSize / 2) : - (qSBlockSize * qNSplitSubBlock); - uint32_t rowActualThisSubBlock = (subBlockIdx == 1) ? - (rowNum - rowSplitSubBlock) : - rowSplitSubBlock; - uint32_t rowOffsetThisSubBlock = subBlockIdx * rowSplitSubBlock; - - uint32_t tokenNumPerHeadThisSubBlock = Min(qSBlockSize, rowActualThisSubBlock); - - uint32_t maskOffsetThisSubBlock = (qNBlockSize == 1) ? - rowOffsetThisSubBlock : - 0; - int64_t offsetMask = layoutMask.GetOffset(MatrixCoord(maskOffsetThisSubBlock, 0)); - auto gMaskThisSubBlock = gMask[offsetMask]; - auto layoutMaskThisSubBlock = layoutMask; - - uint32_t maxRowNumPerLoop = MAX_UB_S_ELEM_NUM / columnNumRound; - uint32_t rowNumTile = RoundDown(maxRowNumPerLoop, FLOAT_BLOCK_SIZE); - uint32_t rowLoopNum = CeilDiv(rowActualThisSubBlock, rowNumTile); - uint32_t preLoad = 1; - - if (rowActualThisSubBlock == 0) { - Arch::CrossCoreWaitFlag(qkReady); - return; - } - - for (uint32_t rowLoopIdx = 0; rowLoopIdx < rowLoopNum + preLoad; rowLoopIdx++) { - if (rowLoopIdx < rowLoopNum) { - uint32_t pingpongFlag = rowLoopIdx % 2; - uint32_t rowOffsetCurLoop = rowLoopIdx * rowNumTile; - uint32_t rowOffsetIoGm = rowOffsetCurLoop + rowOffsetThisSubBlock; - uint32_t rowNumCurLoop = (rowLoopIdx == rowLoopNum - 1) ? - (rowActualThisSubBlock - rowOffsetCurLoop) : - rowNumTile; - // loop 0 mask load before cross core sync - if (rowLoopIdx == 0) { - // the token idx of the start token of the prologue part - uint32_t proTokenIdx = rowOffsetCurLoop % tokenNumPerHeadThisSubBlock; - // the token num of the prologue part - uint32_t proTokenNum = Min(rowNumCurLoop, (tokenNumPerHeadThisSubBlock - proTokenIdx)) % - tokenNumPerHeadThisSubBlock; - // the token num of the epilogue part - uint32_t integralHeadNum = (rowNumCurLoop - proTokenNum) / tokenNumPerHeadThisSubBlock; - // the number of integral heads within a cycle - uint32_t epiTokenNum = rowNumCurLoop - proTokenNum - integralHeadNum * tokenNumPerHeadThisSubBlock; - AscendC::WaitFlag(EVENT_ID2); - CopyMaskGmToUb(gMaskThisSubBlock, columnNum, columnNumRound, maskStride, - tokenNumPerHeadThisSubBlock, proTokenIdx, proTokenNum, integralHeadNum, epiTokenNum); - AscendC::SetFlag(EVENT_ID2); - Arch::CrossCoreWaitFlag(qkReady); - } - int64_t offsetInput = layoutInput.GetOffset(MatrixCoord(rowOffsetIoGm, 0)); - auto gInputCurLoop = gInput[offsetInput]; - AscendC::WaitFlag(pingpongFlag); - CopySGmToUb(gInputCurLoop, (pingpongFlag * MAX_UB_S_ELEM_NUM), rowNumCurLoop, columnNumRound, columnNumPad); - AscendC::SetFlag(pingpongFlag); - } - if (rowLoopIdx >= preLoad) { - uint32_t delayedRowLoopIdx = rowLoopIdx - preLoad; - uint32_t pingpongFlag = delayedRowLoopIdx % 2; - uint32_t rowOffsetCurLoop = delayedRowLoopIdx * rowNumTile; - uint32_t rowNumCurLoop = (delayedRowLoopIdx == rowLoopNum - 1) ? - (rowActualThisSubBlock - rowOffsetCurLoop) : - rowNumTile; - - AscendC::WaitFlag(EVENT_ID0); - AscendC::WaitFlag(EVENT_ID1); - AscendC::WaitFlag(EVENT_ID2); - UpCastMask(rowNumCurLoop, columnNumRound); - AscendC::WaitFlag(pingpongFlag); - ScaleS((pingpongFlag * MAX_UB_S_ELEM_NUM), rowNumCurLoop, columnNumRound); - ApplyMask((pingpongFlag * MAX_UB_S_ELEM_NUM), rowNumCurLoop, columnNumRound); - AscendC::SetFlag(EVENT_ID2); - // next loop mask load - if (rowLoopIdx < rowLoopNum) { - uint32_t rowOffsetCurLoop = rowLoopIdx * rowNumTile; - uint32_t rowNumCurLoop = (rowLoopIdx == rowLoopNum - 1) ? - (rowActualThisSubBlock - rowOffsetCurLoop) : - rowNumTile; - // the token idx of the start token of the prologue part - uint32_t proTokenIdx = rowOffsetCurLoop % tokenNumPerHeadThisSubBlock; - // the token num of the prologue part - uint32_t proTokenNum = Min(rowNumCurLoop, (tokenNumPerHeadThisSubBlock - proTokenIdx)) % - tokenNumPerHeadThisSubBlock; - // the number of integral heads within a cycle - uint32_t integralHeadNum = (rowNumCurLoop - proTokenNum) / tokenNumPerHeadThisSubBlock; - // the token num of the epilogue part - uint32_t epiTokenNum = rowNumCurLoop - proTokenNum - integralHeadNum * tokenNumPerHeadThisSubBlock; - AscendC::WaitFlag(EVENT_ID2); - CopyMaskGmToUb(gMaskThisSubBlock, columnNum, columnNumRound, maskStride, - tokenNumPerHeadThisSubBlock, proTokenIdx, proTokenNum, integralHeadNum, epiTokenNum); - AscendC::SetFlag(EVENT_ID2); - } - // online softmax vectorized compute - uint32_t rowOffsetIoGm = rowOffsetCurLoop + rowOffsetThisSubBlock; - int64_t offsetOutput = layoutOutput.GetOffset(MatrixCoord(rowOffsetIoGm, 0)); - auto gOutputCurLoop = gOutput[offsetOutput]; - auto layoutOutputCurLoop = layoutOutput.GetTileLayout(MatrixCoord(rowNumCurLoop, columnNum)); - SubCoreCompute( - gOutputCurLoop, - layoutOutputCurLoop, - rowOffsetCurLoop, - isFirstStackTile, - columnNumRound, - pingpongFlag, - curStackTileMod); - } - } - } - -private: - float scaleValue; - AscendC::LocalTensor lsUbTensor; - AscendC::LocalTensor lpUbTensor; - AscendC::LocalTensor maskUbTensor; - AscendC::LocalTensor maskUbTensor32; - AscendC::LocalTensor lmUbTensor; - AscendC::LocalTensor hmUbTensor; - AscendC::LocalTensor gmUbTensor; - AscendC::LocalTensor dmUbTensor; - AscendC::LocalTensor llUbTensor; - AscendC::LocalTensor tvUbTensor; - AscendC::LocalTensor glUbTensor; -}; -} // namespace Catlass::Epilogue::Block - -#endif // CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_ONLINE_SOFTMAX_NO_MASK_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_per_group_per_block.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_per_group_per_block.hpp deleted file mode 100644 index e84dcbd790..0000000000 --- a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_per_group_per_block.hpp +++ /dev/null @@ -1,863 +0,0 @@ -/** - * Copyright (c) 2026 Huawei Technologies Co., Ltd. - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of - * CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ - -#ifndef CATLASS_EPILOGUE_BLOCK_EPILOGUE_PER_GROUP_PER_BLOCK_HPP -#define CATLASS_EPILOGUE_BLOCK_EPILOGUE_PER_GROUP_PER_BLOCK_HPP - -#include "catlass/kda_catlass.hpp" -#include "catlass/arch/kda_resource.hpp" -#include "catlass/epilogue/kda_dispatch_policy.hpp" -#include "catlass/kda_gemm_coord.hpp" -#include "catlass/kda_matrix_coord.hpp" -#include "tla/kda_tensor.hpp" -#include "tla/kda_tla_layout.hpp" - -namespace Catlass::Epilogue::Block { - -#define QBMM_BLOCK_EPILOGUE_PERTILE_CLASS_LOCAL_PARAMS \ - template < \ - class L0TileShape_, class DataTypeOut_, class DataTypeIn_, class DataTypeBias_, \ - class DataTypeX1Scale_, class DataTypeX2Scale_, class LayoutX1Scale_, class LayoutX2Scale_> -#define QBMM_BLOCK_EPILOGUE_PERTILE_FUNC_LOCAL_PARAMS \ - BlockEpiloguePertile, L0TileShape_, DataTypeOut_, DataTypeIn_, DataTypeBias_, DataTypeX1Scale_, DataTypeX2Scale_, LayoutX1Scale_, \ - LayoutX2Scale_ - -template < - class L0TileShape_, - class DataTypeOut_, - class DataTypeIn_, - class DataTypeBias_, - class DataTypeX1Scale_, - class DataTypeX2Scale_, - class LayoutX1Scale_, - class LayoutX2Scale_> -class BlockEpilogue< - BlockEpiloguePertile, - L0TileShape_, - DataTypeOut_, - DataTypeIn_, - DataTypeBias_, - DataTypeX1Scale_, - DataTypeX2Scale_, - LayoutX1Scale_, - LayoutX2Scale_> { -public: - using DispatchPolicy = BlockEpiloguePertile; - using ArchTag = typename DispatchPolicy::ArchTag; - CATLASS_DEVICE BlockEpilogue() - { - if ASCEND_IS_AIV { - AscendC::SetFlag(0); - AscendC::SetFlag(1); - if constexpr (!AscendC::IsSameType::value) { - AscendC::SetFlag(0); - AscendC::SetFlag(1); - } - } - } - - CATLASS_DEVICE ~BlockEpilogue() - { - AscendC::WaitFlag(0); - AscendC::WaitFlag(1); - if ASCEND_IS_AIV { - if constexpr (!AscendC::IsSameType::value) { - AscendC::WaitFlag(0); - AscendC::WaitFlag(1); - } - if (isBias_) { - AscendC::WaitFlag(biasPingPongID_); - AscendC::WaitFlag(biasPingPongID_ ^ 1); - } - } - } - - struct Arguments { - GM_ADDR outGmAddr{nullptr}; - GM_ADDR x2ScaleGmAddr{nullptr}; - GM_ADDR x1ScaleGmAddr{nullptr}; - GM_ADDR biasGmAddr{nullptr}; - uint32_t baseM; - uint32_t baseN; - uint32_t baseK; - uint32_t groupSizeM = 1U; - uint32_t groupSizeN = 128U; - uint32_t groupSizeK = 128U; - uint32_t isBias = 0; - Arguments() = default; - }; - - // params - using Params = Arguments; - using YType = DataTypeOut_; - using CType = DataTypeIn_; - using BiasType = float; - using X2ScaleType = DataTypeX2Scale_; - using X1ScaleType = DataTypeX1Scale_; - using LayoutX1Scale = LayoutX1Scale_; - using LayoutX2Scale = LayoutX2Scale_; - using CalcType = float; - using BlockCoord = tla::Coord; - - constexpr static uint32_t Y_IDX = 0; - constexpr static uint32_t X2SCALE_IDX = 1; - constexpr static uint32_t X1SCALE_IDX = 2; - constexpr static uint32_t BIAS_IDX = 3; - constexpr static uint32_t QBMM_BUFFER_NUM = 2; - constexpr static uint16_t QBMM_AIV_SYNC_AIC_FLAG = 6; - constexpr static uint16_t QBMM_AIC_SYNC_AIV_FLAG = 8; - constexpr static uint8_t QBMM_AIC_SYNC_AIV_MODE = 4; - constexpr static uint64_t QBMM_MAX_STEP_SCALEA_K = 16; - constexpr static uint32_t QBMM_UB_ALIGN_SIZE = 32; - constexpr static uint32_t UB_TWO_BANK_ELEMS_B32 = 128U; - constexpr static int64_t PER_BLOCK_SIZE = 128LL; - constexpr static uint32_t UB_SUB_BANK_NUM = 2U; - constexpr static uint32_t UB_SUB_BANK_ELEMS_B32 = 64U; - constexpr static uint32_t UB_ALIGN_SIZE = 32U; - constexpr static uint32_t QBMM_BMM_BLOCK_NUM = 16; - - struct PertileUBParam { - bool CopyOutWithSplitN = false; - uint16_t ndNum; - uint64_t singleM; - uint64_t singleN; - uint64_t validM; - uint32_t validN[UB_SUB_BANK_NUM]; - uint64_t offsetScaleM; - uint64_t offsetScaleN[UB_SUB_BANK_NUM]; - uint64_t offsetY[UB_SUB_BANK_NUM]; - uint64_t offsetBias[UB_SUB_BANK_NUM]; - }; - -public: - CATLASS_DEVICE void Init(const Params *params); - template - CATLASS_DEVICE void operator()(const TensorUb &tensorUbPing); - CATLASS_DEVICE void UpdateGlobalAddr(const BlockCoord &baseOffset); - CATLASS_DEVICE void UpdateParamsForNextProblem(const GemmCoord &problemShape); - CATLASS_DEVICE auto GetL0c2UbPingTensor(); - CATLASS_DEVICE auto GetL0c2UbPongTensor(); - -private: - CATLASS_DEVICE void ProcessAivSingleKPertile(int64_t x1ScaleOffset, - __gm__ X2ScaleType *x2ScaleAddr[UB_SUB_BANK_NUM]); - - CATLASS_DEVICE void ProcessAivSingleKPerblock(int64_t x1ScaleOffset, - __gm__ X2ScaleType *x2ScaleAddr[UB_SUB_BANK_NUM]); - CATLASS_DEVICE void CopyInBias(); - template - CATLASS_DEVICE __ubuf__ T *CopyInX1Scale(uint64_t srcOffset, uint64_t m, uint64_t k); - template - CATLASS_DEVICE T CopyInX1ScalePertile(__gm__ T *src, uint64_t offset); - template - CATLASS_DEVICE void CopyInX2Scale(T x2Scale[UB_SUB_BANK_NUM], __gm__ T *src[UB_SUB_BANK_NUM], uint64_t offset); - CATLASS_DEVICE int64_t CalcX1OffsetPerGroup(); - CATLASS_DEVICE void CalcX2OffsetPerGroup(int64_t x2ScaleOffset[UB_SUB_BANK_NUM]); - template - CATLASS_DEVICE __ubuf__ T *GetX1ScaleUbAddrPerGroup(int64_t x1ScaleOffset, uint64_t kOffset, uint64_t kElem); - template - __simd_vf__ void AivPerTensor(__ubuf__ CalcType *dst, __ubuf__ CType *l0cOut, __ubuf__ X1ScaleType *x1Scale, - uint16_t mSize, uint32_t nSize0, uint32_t nSize1, uint16_t kSize, - X2ScaleType x2Scale0, X2ScaleType x2Scale1, uint64_t x1ScaleKIdxInCache); - template - __simd_vf__ void AivPerTensor(__ubuf__ CType *dst, __ubuf__ CType *l0cOut, X1ScaleType x1Scale, - uint16_t mSize, uint32_t nSize0, uint32_t nSize1, X2ScaleType x2Scale0, - X2ScaleType x2Scale1); - template - __simd_vf__ void AddBias( - __ubuf__ CalcType *mmAdd, __ubuf__ BiasType *bias, uint16_t mSize, uint32_t nSize0, uint32_t nSize1); - CATLASS_DEVICE void AivPostProcess(const AscendC::LocalTensor &mmAddUb); - CATLASS_DEVICE void CopyOut(const AscendC::LocalTensor &ubRes, uint16_t eventId, uint16_t blkCount, - uint32_t blkLen, uint32_t srcStride, uint32_t dstStride, uint64_t yOffset); - CATLASS_DEVICE void CastAndCopyOut(const AscendC::LocalTensor &mmAddUb); - CATLASS_DEVICE void UpdatePertileUBValidMN(); - CATLASS_DEVICE void UpdatePertileUBParam(); - CATLASS_DEVICE void WaitForCube(uint16_t crossPingPongID) - { - AscendC::CrossCoreWaitFlag(QBMM_AIV_SYNC_AIC_FLAG + crossPingPongID); - } - CATLASS_DEVICE void NotifyCube(uint16_t crossPingPongID) - { - AscendC::CrossCoreSetFlag(QBMM_AIC_SYNC_AIV_FLAG + crossPingPongID); - } - - AscendC::GlobalTensor biasGlobal_; - AscendC::GlobalTensor cGlobal_; - AscendC::GlobalTensor x1ScaleGlobal_; - __gm__ X1ScaleType *x1ScaleGlobalPerblock_; - __gm__ X2ScaleType *x2ScaleGlobal_; - AscendC::LocalTensor mmResPing_; - AscendC::LocalTensor mmResPong_; - AscendC::LocalTensor ubResPing_; - AscendC::LocalTensor ubResPong_; - AscendC::LocalTensor mmAddUb_; - AscendC::LocalTensor x1ScaleUbPing_; - AscendC::LocalTensor x1ScaleUbPong_; - AscendC::LocalTensor biasUbPing_; - AscendC::LocalTensor biasUbPong_; - -private: - const Params *params_; - PertileUBParam ubParams_; - GemmCoord problemShape_{}; - GemmCoord actualSingleShape_{}; - BlockCoord baseOffset_{0, 0, 0, 0}; - BlockCoord blockCoord_{0, 0, 0, 0}; - - uint64_t scaleM_ = 0; - uint64_t scaleN_ = 0; - uint64_t scaleK_ = 0; - uint32_t subBlockIdx_; - uint16_t crossPingPongID_ = 0; - uint16_t x1ScalePingPongID_ = 0; - uint16_t biasPingPongID_ = 2; - bool isPergroup_ = false; - bool isBias_ = false; - bool needAivSet = false; -}; - -QBMM_BLOCK_EPILOGUE_PERTILE_CLASS_LOCAL_PARAMS -CATLASS_DEVICE void BlockEpilogue::Init(const Params *params) -{ - if ASCEND_IS_AIC { - return; - } - - params_ = params; - subBlockIdx_ = AscendC::GetSubBlockIdx(); - constexpr uint32_t elems = UB_TWO_BANK_ELEMS_B32 * PER_BLOCK_SIZE; - constexpr uint32_t addUbOffset = elems * UB_SUB_BANK_NUM * sizeof(CType); // l0c res 128KB - mmAddUb_ = AscendC::LocalTensor(AscendC::TPosition::VECCALC, addUbOffset, elems); - constexpr uint32_t afterAddOffset = addUbOffset + elems * sizeof(CalcType); // ub add res 64KB - if constexpr (!AscendC::IsSameType::value) { - ubResPing_ = AscendC::LocalTensor(AscendC::TPosition::VECCALC, afterAddOffset, elems); - ubResPong_ = ubResPing_[elems / QBMM_BUFFER_NUM]; - } - - isPergroup_ = params_->groupSizeM == 1; - if (isPergroup_) { - constexpr uint32_t x1ScaleUbOffset = - (AscendC::IsSameType::value) ? afterAddOffset : afterAddOffset + elems * sizeof(YType); - x1ScaleUbPing_ = AscendC::LocalTensor(AscendC::TPosition::VECCALC, x1ScaleUbOffset, - PER_BLOCK_SIZE * QBMM_MAX_STEP_SCALEA_K); - x1ScaleUbPong_ = AscendC::LocalTensor( - AscendC::TPosition::VECCALC, x1ScaleUbOffset + PER_BLOCK_SIZE * QBMM_MAX_STEP_SCALEA_K * sizeof(X1ScaleType), - PER_BLOCK_SIZE * QBMM_MAX_STEP_SCALEA_K); - if (params->isBias == 1) { - AscendC::SetFlag( - biasPingPongID_); - AscendC::SetFlag(biasPingPongID_ ^ 1); - isBias_ = true; - constexpr uint32_t afterx1ScaleUbOffset = - x1ScaleUbOffset + 2 * PER_BLOCK_SIZE * QBMM_MAX_STEP_SCALEA_K * sizeof(X1ScaleType); - biasUbPing_ = - AscendC::LocalTensor(AscendC::TPosition::VECCALC, afterx1ScaleUbOffset, PER_BLOCK_SIZE); - biasUbPong_ = AscendC::LocalTensor( - AscendC::TPosition::VECCALC, afterx1ScaleUbOffset + PER_BLOCK_SIZE * sizeof(BiasType), PER_BLOCK_SIZE); - } - } -} - -QBMM_BLOCK_EPILOGUE_PERTILE_CLASS_LOCAL_PARAMS -CATLASS_DEVICE void BlockEpilogue::UpdateParamsForNextProblem( - const GemmCoord &problemShape) -{ - problemShape_ = problemShape; - - scaleM_ = CeilDiv(problemShape_.m(), params_->groupSizeM); - scaleN_ = CeilDiv(problemShape_.n(), params_->groupSizeN); - scaleK_ = CeilDiv(problemShape_.k(), params_->groupSizeK); -} - -QBMM_BLOCK_EPILOGUE_PERTILE_CLASS_LOCAL_PARAMS -CATLASS_DEVICE void BlockEpilogue::UpdateGlobalAddr( - const BlockCoord &baseOffset) -{ - if ASCEND_IS_AIV { - x1ScaleGlobal_.SetGlobalBuffer((__gm__ X1ScaleType *)params_->x1ScaleGmAddr + tla::get(baseOffset)); - x1ScaleGlobalPerblock_ = (__gm__ X1ScaleType *)params_->x1ScaleGmAddr + tla::get(baseOffset); - x2ScaleGlobal_ = (__gm__ X2ScaleType *)params_->x2ScaleGmAddr + tla::get(baseOffset); - cGlobal_.SetGlobalBuffer((__gm__ YType *)params_->outGmAddr + tla::get(baseOffset)); - biasGlobal_.SetGlobalBuffer((__gm__ BiasType *)params_->biasGmAddr + tla::get(baseOffset)); - } -} - -QBMM_BLOCK_EPILOGUE_PERTILE_CLASS_LOCAL_PARAMS -CATLASS_DEVICE int64_t BlockEpilogue::CalcX1OffsetPerGroup() -{ - int64_t x1ScaleOffset = tla::get(blockCoord_) * scaleK_; - if (subBlockIdx_ == 1) { - x1ScaleOffset += (ubParams_.offsetScaleM * scaleK_); - } - return x1ScaleOffset; -} - -QBMM_BLOCK_EPILOGUE_PERTILE_CLASS_LOCAL_PARAMS -CATLASS_DEVICE void BlockEpilogue::CalcX2OffsetPerGroup( - int64_t x2ScaleOffset[UB_SUB_BANK_NUM]) -{ - x2ScaleOffset[0] = tla::get(blockCoord_) / params_->groupSizeN + ubParams_.offsetScaleN[0]; - x2ScaleOffset[1] = tla::get(blockCoord_) / params_->groupSizeN + ubParams_.offsetScaleN[1]; -} - -QBMM_BLOCK_EPILOGUE_PERTILE_CLASS_LOCAL_PARAMS -template -CATLASS_DEVICE __ubuf__ T *BlockEpilogue::CopyInX1Scale( - uint64_t srcOffset, uint64_t m, uint64_t k) -{ - AscendC::DataCopyExtParams x1ScaleGm2UbParams; - AscendC::DataCopyPadExtParams padParams; - x1ScaleGm2UbParams.blockCount = m; - x1ScaleGm2UbParams.blockLen = k * sizeof(T); - x1ScaleGm2UbParams.srcStride = (scaleK_ - k) * sizeof(T); - - auto x1ScaleUb = x1ScalePingPongID_ == 0 ? &x1ScaleUbPing_ : &x1ScaleUbPong_; - AscendC::WaitFlag(x1ScalePingPongID_); - AscendC::DataCopyPad(*x1ScaleUb, x1ScaleGlobal_[srcOffset], x1ScaleGm2UbParams, padParams); - AscendC::SetFlag(x1ScalePingPongID_); - AscendC::WaitFlag(x1ScalePingPongID_); - return reinterpret_cast<__ubuf__ T *>(x1ScaleUb->GetPhyAddr()); -} -QBMM_BLOCK_EPILOGUE_PERTILE_CLASS_LOCAL_PARAMS -CATLASS_DEVICE void BlockEpilogue::CopyInBias() -{ - AscendC::DataCopyExtParams bias4N0Gm2UbParams; - AscendC::DataCopyExtParams bias4N1Gm2UbParams; - AscendC::DataCopyPadExtParams padParams; - uint32_t validN0Align = ubParams_.validN[0] * sizeof(BiasType); - bias4N0Gm2UbParams.blockCount = 1; - bias4N0Gm2UbParams.blockLen = validN0Align; - bias4N0Gm2UbParams.srcStride = 0; - bias4N0Gm2UbParams.dstStride = 0; - AscendC::WaitFlag(biasPingPongID_); - auto biasUb = biasPingPongID_ == 2 ? biasUbPing_ : biasUbPong_; - AscendC::DataCopyPad(biasUb, biasGlobal_[ubParams_.offsetBias[0]], bias4N0Gm2UbParams, padParams); - if (ubParams_.validN[1] != 0) { - uint32_t validN1Align = ubParams_.validN[1] * sizeof(BiasType); - bias4N1Gm2UbParams.blockCount = 1; - bias4N1Gm2UbParams.blockLen = validN1Align; - bias4N1Gm2UbParams.srcStride = 0; - bias4N1Gm2UbParams.dstStride = 0; - AscendC::DataCopyPad( - biasUb[RoundUp(ubParams_.validN[0], UB_ALIGN_SIZE / sizeof(BiasType))], biasGlobal_[ubParams_.offsetBias[1]], - bias4N1Gm2UbParams, padParams); - } - AscendC::SetFlag(biasPingPongID_); -} -QBMM_BLOCK_EPILOGUE_PERTILE_CLASS_LOCAL_PARAMS -template -CATLASS_DEVICE T BlockEpilogue::CopyInX1ScalePertile( - __gm__ T *src, uint64_t offset) -{ - return src[offset]; -} - -QBMM_BLOCK_EPILOGUE_PERTILE_CLASS_LOCAL_PARAMS -template -CATLASS_DEVICE void BlockEpilogue::CopyInX2Scale( - T x2Scale[UB_SUB_BANK_NUM], __gm__ T *src[UB_SUB_BANK_NUM], uint64_t offset) -{ - x2Scale[0] = src[0][offset * scaleN_]; - x2Scale[1] = src[1][offset * scaleN_]; -} - -QBMM_BLOCK_EPILOGUE_PERTILE_CLASS_LOCAL_PARAMS -CATLASS_DEVICE void BlockEpilogue::UpdatePertileUBValidMN() -{ - int64_t actualN = actualSingleShape_.n(); - if (ubParams_.CopyOutWithSplitN) { - ubParams_.validM = ubParams_.singleM; - uint64_t subBlockIdxOffset = AscendC::GetSubBlockIdx() * ubParams_.singleN; - uint64_t ndNumN = 2 * ubParams_.singleN + subBlockIdxOffset; // 2: the nSize of 2 ND, base is 2 ND - ubParams_.validN[0] = actualN < subBlockIdxOffset ? 0 : Min(ubParams_.singleN, actualN - subBlockIdxOffset); - ubParams_.validN[1] = actualN < ndNumN ? 0 : Min(ubParams_.singleN, actualN - ndNumN); - } else { - if (AscendC::GetSubBlockIdx() == 0) { - ubParams_.validM = ubParams_.singleM; - } else { - ubParams_.validM = actualSingleShape_.m() - ubParams_.singleM; - } - ubParams_.validN[0] = Min(ubParams_.singleN, static_cast(actualN)); - ubParams_.validN[1] = actualN < ubParams_.singleN ? 0 : Min(ubParams_.singleN, actualN - ubParams_.singleN); - } -} - -QBMM_BLOCK_EPILOGUE_PERTILE_CLASS_LOCAL_PARAMS -CATLASS_DEVICE void BlockEpilogue::UpdatePertileUBParam() -{ - ubParams_.CopyOutWithSplitN = - actualSingleShape_.n() > params_->groupSizeN || actualSingleShape_.m() == 1; - uint32_t fixpipeN = 0; - if (ubParams_.CopyOutWithSplitN) { - // (m * n/2) is written to 2 UB, n must be multiples of 32 - // | AIV0 singleN | AIV1 singleN | AIV0 singleN | AIV1 singleN |, max(singleN) = 64 - ubParams_.ndNum = actualSingleShape_.n() > UB_TWO_BANK_ELEMS_B32 ? 2 : 1; // 2: 2 ND - int64_t alignedNBase = - actualSingleShape_.n() > PER_BLOCK_SIZE ? PER_BLOCK_SIZE : AscendC::ONE_BLK_SIZE * ubParams_.ndNum; - fixpipeN = RoundUp(actualSingleShape_.n(), static_cast(alignedNBase)) / ubParams_.ndNum; - ubParams_.singleN = fixpipeN / static_cast(AscendC::GetTaskRation()); - ubParams_.singleM = actualSingleShape_.m(); - } else { - // (m/2 * n) is written to 2 UB, m must be multiples of 2 - // | AIV0 singleN | AIV0 singleN | - // | AIV1 singleN | AIV1 singleN | - ubParams_.ndNum = actualSingleShape_.n() > UB_SUB_BANK_ELEMS_B32 ? 2 : 1; // 2: 2 ND - fixpipeN = RoundUp(actualSingleShape_.n(), static_cast(AscendC::BLOCK_CUBE) * ubParams_.ndNum) / - ubParams_.ndNum; - ubParams_.singleN = fixpipeN; - ubParams_.singleM = CeilDiv(actualSingleShape_.m(), AscendC::GetTaskRation()); - } - - UpdatePertileUBValidMN(); - int64_t offsetM = 0; - int64_t offsetN0 = 0; - int64_t offsetN1 = 0; - - if (ubParams_.CopyOutWithSplitN) { - offsetN0 = ubParams_.validN[0] == 0 ? 0 : AscendC::GetSubBlockIdx() * ubParams_.singleN; - offsetN1 = ubParams_.validN[1] == 0 ? offsetN0 : offsetN0 + UB_SUB_BANK_NUM * ubParams_.singleN; - } else { - if (AscendC::GetSubBlockIdx() == 1) { - offsetM += ubParams_.singleM; - } - offsetN1 = ubParams_.validN[1] == 0 ? 0 : ubParams_.singleN; - } - - ubParams_.offsetScaleM = offsetM / params_->groupSizeM; - ubParams_.offsetScaleN[0] = offsetN0 / params_->groupSizeN; - ubParams_.offsetScaleN[1] = offsetN1 / params_->groupSizeN; - ubParams_.offsetY[0] = tla::get(blockCoord_) + offsetM * problemShape_.n() + offsetN0; - ubParams_.offsetY[1] = tla::get(blockCoord_) + offsetM * problemShape_.n() + offsetN1; - ubParams_.offsetBias[0] = tla::get(blockCoord_) + offsetN0; - ubParams_.offsetBias[1] = tla::get(blockCoord_) + offsetN1; -} - -QBMM_BLOCK_EPILOGUE_PERTILE_CLASS_LOCAL_PARAMS -template -CATLASS_DEVICE void BlockEpilogue::operator()( - const TensorUb &tensorUb) -{ - actualSingleShape_ = GemmCoord{tla::get<0>(tensorUb.shape()), - tla::get<1>(tensorUb.shape()), - problemShape_.k()}; - blockCoord_ = BlockCoord{tla::get<0>(tensorUb.coord()) * problemShape_.n() + tla::get<1>(tensorUb.coord()), - tla::get<1>(tensorUb.coord()), - tla::get<0>(tensorUb.coord()), 0}; - - constexpr uint32_t elems = UB_TWO_BANK_ELEMS_B32 * PER_BLOCK_SIZE; - mmResPing_ = tensorUb.data(); - mmResPong_ = tensorUb.data()[elems]; - UpdatePertileUBParam(); - int64_t x1ScaleOffset = CalcX1OffsetPerGroup(); // one block same x1Scale - int64_t x2ScaleOffset[UB_SUB_BANK_NUM] = {0}; // maybe 2 different x2Scale - CalcX2OffsetPerGroup(x2ScaleOffset); - __gm__ X2ScaleType *x2ScaleAddr[UB_SUB_BANK_NUM] = {x2ScaleGlobal_ + x2ScaleOffset[0], - x2ScaleGlobal_ + x2ScaleOffset[1]}; - - if (isPergroup_) { - ProcessAivSingleKPertile(x1ScaleOffset, x2ScaleAddr); - } else { - ProcessAivSingleKPerblock(x1ScaleOffset, x2ScaleAddr); - } -} - -QBMM_BLOCK_EPILOGUE_PERTILE_CLASS_LOCAL_PARAMS -template -CATLASS_DEVICE __ubuf__ T * -BlockEpilogue::GetX1ScaleUbAddrPerGroup( - int64_t x1ScaleOffset, uint64_t kOffset, uint64_t kElem) -{ - uint64_t scaleX1GmOffset; - scaleX1GmOffset = x1ScaleOffset + kOffset; - return CopyInX1Scale(scaleX1GmOffset, ubParams_.validM, kElem); -} - -QBMM_BLOCK_EPILOGUE_PERTILE_CLASS_LOCAL_PARAMS -CATLASS_DEVICE void BlockEpilogue::ProcessAivSingleKPertile( - int64_t x1ScaleOffset, __gm__ X2ScaleType *x2ScaleAddr[UB_SUB_BANK_NUM]) -{ - auto mmAddUbAddr = reinterpret_cast<__ubuf__ CType *>(mmAddUb_.GetPhyAddr()); - const uint16_t x1ScaleKElem = Min(QBMM_MAX_STEP_SCALEA_K, scaleK_); - uint64_t kElem; - __ubuf__ X1ScaleType *x1ScaleUbAddr; - X2ScaleType x2Scale[UB_SUB_BANK_NUM]; - - for (uint64_t kb = 0, kOffset = 0; kb < problemShape_.k(); kb += params_->baseK, kOffset++) { - CopyInX2Scale(x2Scale, x2ScaleAddr, kOffset); - uint64_t x1ScaleKRem = kOffset % x1ScaleKElem; - if (x1ScaleKRem == 0) { - kElem = Min(static_cast(x1ScaleKElem), scaleK_ - kOffset); - x1ScaleUbAddr = GetX1ScaleUbAddrPerGroup(x1ScaleOffset, kOffset, kElem); - } - - WaitForCube(crossPingPongID_); - auto mmUbInputAddr = crossPingPongID_ == 0 ? reinterpret_cast(mmResPing_.GetPhyAddr()) : - reinterpret_cast(mmResPong_.GetPhyAddr()); - if (kb == 0) { - if (ubParams_.ndNum == 1) { - AivPerTensor((__ubuf__ CalcType *)mmAddUbAddr, (__ubuf__ CType *)mmUbInputAddr, x1ScaleUbAddr, - ubParams_.validM, ubParams_.validN[0], ubParams_.validN[1], kElem, x2Scale[0], - x2Scale[1], x1ScaleKRem); - } else { - AivPerTensor((__ubuf__ CalcType *)mmAddUbAddr, (__ubuf__ CType *)mmUbInputAddr, x1ScaleUbAddr, - ubParams_.validM, ubParams_.validN[0], ubParams_.validN[1], kElem, x2Scale[0], - x2Scale[1], x1ScaleKRem); - } - } else { - if (ubParams_.ndNum == 1) { - AivPerTensor((__ubuf__ CalcType *)mmAddUbAddr, (__ubuf__ CType *)mmUbInputAddr, x1ScaleUbAddr, - ubParams_.validM, ubParams_.validN[0], ubParams_.validN[1], kElem, x2Scale[0], - x2Scale[1], x1ScaleKRem); - } else { - AivPerTensor((__ubuf__ CalcType *)mmAddUbAddr, (__ubuf__ CType *)mmUbInputAddr, x1ScaleUbAddr, - ubParams_.validM, ubParams_.validN[0], ubParams_.validN[1], kElem, x2Scale[0], - x2Scale[1], x1ScaleKRem); - } - } - if (x1ScaleKRem == x1ScaleKElem - 1 || kOffset == scaleK_ - 1) { - AscendC::SetFlag(x1ScalePingPongID_); - x1ScalePingPongID_ = (x1ScalePingPongID_ + 1) & 1; - } - NotifyCube(crossPingPongID_); - needAivSet = needAivSet || crossPingPongID_ == 1; - crossPingPongID_ = (crossPingPongID_ + 1) & 1; - } - if (isBias_) { - CopyInBias(); - } - AivPostProcess(mmAddUb_); -} - -QBMM_BLOCK_EPILOGUE_PERTILE_CLASS_LOCAL_PARAMS -CATLASS_DEVICE void BlockEpilogue::ProcessAivSingleKPerblock( - int64_t x1ScaleOffset, __gm__ X2ScaleType *x2ScaleAddr[UB_SUB_BANK_NUM]) -{ - auto mmAddUbAddr = reinterpret_cast<__ubuf__ CType *>(mmAddUb_.GetPhyAddr()); - auto x1ScaleAddr = x1ScaleGlobalPerblock_ + x1ScaleOffset; - X2ScaleType x2Scale[UB_SUB_BANK_NUM]; - for (uint64_t kb = 0, kOffset = 0; kb < problemShape_.k(); kb += params_->baseK, kOffset++) { - CopyInX2Scale(x2Scale, x2ScaleAddr, kOffset); - X1ScaleType x1Scale = CopyInX1ScalePertile(x1ScaleAddr, kOffset); - WaitForCube(crossPingPongID_); - auto mmUbInputAddr = crossPingPongID_ == 0 ? reinterpret_cast(mmResPing_.GetPhyAddr()) : - reinterpret_cast(mmResPong_.GetPhyAddr()); - if (kb == 0) { - if (ubParams_.ndNum == 1) { - AivPerTensor((__ubuf__ CType *)mmAddUbAddr, (__ubuf__ CType *)mmUbInputAddr, x1Scale, - ubParams_.validM, ubParams_.validN[0], ubParams_.validN[1], x2Scale[0], - x2Scale[1]); - } else { - AivPerTensor((__ubuf__ CType *)mmAddUbAddr, (__ubuf__ CType *)mmUbInputAddr, x1Scale, - ubParams_.validM, ubParams_.validN[0], ubParams_.validN[1], x2Scale[0], - x2Scale[1]); - } - } else { - if (ubParams_.ndNum == 1) { - AivPerTensor((__ubuf__ CType *)mmAddUbAddr, (__ubuf__ CType *)mmUbInputAddr, x1Scale, - ubParams_.validM, ubParams_.validN[0], ubParams_.validN[1], x2Scale[0], - x2Scale[1]); - } else { - AivPerTensor((__ubuf__ CType *)mmAddUbAddr, (__ubuf__ CType *)mmUbInputAddr, x1Scale, - ubParams_.validM, ubParams_.validN[0], ubParams_.validN[1], x2Scale[0], - x2Scale[1]); - } - } - NotifyCube(crossPingPongID_); - crossPingPongID_ = (crossPingPongID_ + 1) & 1; - } - AivPostProcess(mmAddUb_); -} - -QBMM_BLOCK_EPILOGUE_PERTILE_CLASS_LOCAL_PARAMS -template -__simd_vf__ void BlockEpilogue::AivPerTensor( - __ubuf__ CalcType *dst, __ubuf__ CType *l0cOut, __ubuf__ X1ScaleType *x1Scale, uint16_t mSize, uint32_t nSize0, - uint32_t nSize1, uint16_t kSize, X2ScaleType x2Scale0, X2ScaleType x2Scale1, uint64_t x1ScaleKIdxInCache) -{ - uint16_t alignM = RoundUp(mSize, QBMM_UB_ALIGN_SIZE / sizeof(X1ScaleType)); - uint16_t alignK = RoundUp(kSize, QBMM_UB_ALIGN_SIZE / sizeof(X1ScaleType)); - AscendC::MicroAPI::RegTensor x1ScaleReg, muledScaleReg; - AscendC::MicroAPI::RegTensor l0cOutReg; - AscendC::MicroAPI::RegTensor addReg; - AscendC::MicroAPI::RegTensor ResReg, mulScaleOutReg; - constexpr static AscendC::MicroAPI::CastTrait ctInt322Fp32 = { - AscendC::MicroAPI::RegLayout::UNKNOWN, AscendC::MicroAPI::SatMode::UNKNOWN, - AscendC::MicroAPI::MaskMergeMode::ZEROING, AscendC::RoundMode::CAST_RINT}; - { - for (uint16_t mIdx = 0; mIdx < mSize; mIdx++) { - AscendC::MicroAPI::LoadAlign( - x1ScaleReg, x1Scale + mIdx * alignK + x1ScaleKIdxInCache); - - // 1 ND - uint32_t elementNum = nSize0; - AscendC::MicroAPI::MaskReg maskN = AscendC::MicroAPI::UpdateMask(elementNum); - // copy input from ub to register, addr of ub should align to 32B - uint32_t offset = mIdx * UB_TWO_BANK_ELEMS_B32; - uint32_t l0cOutOffset = mIdx * (nSize0 + nSize1); - AscendC::MicroAPI::LoadAlign(l0cOutReg, l0cOut + l0cOutOffset); - // l0c_out * scale - AscendC::MicroAPI::Muls(muledScaleReg, x1ScaleReg, x2Scale0, maskN); - if constexpr (AscendC::IsSameType::value) { - AscendC::MicroAPI::RegTensor l0cOutRegAfterCast{}; - AscendC::MicroAPI::Cast(l0cOutRegAfterCast, l0cOutReg, maskN); - AscendC::MicroAPI::Mul(mulScaleOutReg, l0cOutRegAfterCast, muledScaleReg, maskN); - } else { - AscendC::MicroAPI::Mul(mulScaleOutReg, l0cOutReg, muledScaleReg, maskN); - } - uint32_t dstUbOffset = offset; - if constexpr (isFirstKLoop) { - AscendC::MicroAPI::StoreAlign(dst + dstUbOffset, - mulScaleOutReg, maskN); - } else { - AscendC::MicroAPI::LoadAlign(addReg, dst + dstUbOffset); - AscendC::MicroAPI::Add(ResReg, mulScaleOutReg, addReg, maskN); - // copy out from register to ub - AscendC::MicroAPI::StoreAlign(dst + dstUbOffset, - ResReg, maskN); - } - // 2 ND - if constexpr (ndNum == 1) { - continue; - } - elementNum = nSize1; - maskN = AscendC::MicroAPI::UpdateMask(elementNum); - // copy input from ub to register, addr of ub should align to 32B - l0cOutOffset = l0cOutOffset + nSize0; // + distance of 2 NDs - AscendC::MicroAPI::LoadAlign(l0cOutReg, l0cOut + l0cOutOffset); - // l0c_out * scale - AscendC::MicroAPI::Muls(muledScaleReg, x1ScaleReg, x2Scale1, maskN); - if constexpr (AscendC::IsSameType::value) { - AscendC::MicroAPI::RegTensor l0cOutRegAfterCast{}; - AscendC::MicroAPI::Cast(l0cOutRegAfterCast, l0cOutReg, maskN); - AscendC::MicroAPI::Mul(mulScaleOutReg, l0cOutRegAfterCast, muledScaleReg, maskN); - } else { - AscendC::MicroAPI::Mul(mulScaleOutReg, l0cOutReg, muledScaleReg, maskN); - } - dstUbOffset = offset + nSize0; - if constexpr (isFirstKLoop) { - AscendC::MicroAPI::StoreAlign(dst + dstUbOffset, - mulScaleOutReg, maskN); - } else { - AscendC::MicroAPI::LoadAlign(addReg, dst + dstUbOffset); - AscendC::MicroAPI::Add(ResReg, mulScaleOutReg, addReg, maskN); - // copy out from register to ub - AscendC::MicroAPI::StoreAlign(dst + dstUbOffset, - ResReg, maskN); - } - } - } -} - -QBMM_BLOCK_EPILOGUE_PERTILE_CLASS_LOCAL_PARAMS -template -__simd_vf__ void BlockEpilogue::AivPerTensor( - __ubuf__ CType *dst, __ubuf__ CType *l0cOut, X1ScaleType x1Scale, uint16_t mSize, uint32_t nSize0, uint32_t nSize1, - X2ScaleType x2Scale0, X2ScaleType x2Scale1) -{ - { - for (uint16_t mIdx = 0; mIdx < mSize; mIdx++) { - // 1 ND - X2ScaleType scaleMul = x1Scale * x2Scale0; - uint32_t elementNum = nSize0; - AscendC::MicroAPI::MaskReg maskN = AscendC::MicroAPI::UpdateMask(elementNum); - AscendC::MicroAPI::RegTensor l0cOutReg; - AscendC::MicroAPI::RegTensor addReg; - AscendC::MicroAPI::RegTensor ResReg, mulScaleOutReg; - // copy input from ub to register, addr of ub should align to 32B - uint32_t offset = mIdx * UB_TWO_BANK_ELEMS_B32; - uint32_t l0cOutOffset = offset; - AscendC::MicroAPI::LoadAlign(l0cOutReg, l0cOut + offset); - // l0c_out * scale - AscendC::MicroAPI::Muls(mulScaleOutReg, l0cOutReg, scaleMul, maskN); - uint32_t dstUbOffset = offset; - if constexpr (isFirstKLoop) { - AscendC::MicroAPI::StoreAlign(dst + dstUbOffset, - mulScaleOutReg, maskN); - } else { - AscendC::MicroAPI::LoadAlign(addReg, dst + dstUbOffset); - AscendC::MicroAPI::Add(ResReg, mulScaleOutReg, addReg, maskN); - // copy out from register to ub - AscendC::MicroAPI::StoreAlign(dst + dstUbOffset, - ResReg, maskN); - } - // 2 ND - if constexpr (ndNum == 1) { - continue; - } - scaleMul = x1Scale * x2Scale1; - elementNum = nSize1; - maskN = AscendC::MicroAPI::UpdateMask(elementNum); - // copy input from ub to register, addr of ub should align to 32B - l0cOutOffset = offset + nSize0; // + distance of 2 NDs - AscendC::MicroAPI::LoadAlign(l0cOutReg, l0cOut + l0cOutOffset); - // l0c_out * scale - AscendC::MicroAPI::Muls(mulScaleOutReg, l0cOutReg, scaleMul, maskN); - dstUbOffset = offset + nSize0; - if constexpr (isFirstKLoop) { - AscendC::MicroAPI::StoreAlign(dst + dstUbOffset, - mulScaleOutReg, maskN); - } else { - AscendC::MicroAPI::LoadAlign(addReg, dst + dstUbOffset); - AscendC::MicroAPI::Add(ResReg, mulScaleOutReg, addReg, maskN); - // copy out from register to ub - AscendC::MicroAPI::StoreAlign(dst + dstUbOffset, - ResReg, maskN); - } - } - } -} - -QBMM_BLOCK_EPILOGUE_PERTILE_CLASS_LOCAL_PARAMS -template -__simd_vf__ void BlockEpilogue::AddBias( - __ubuf__ CalcType *mmAdd, __ubuf__ BiasType *bias, uint16_t mSize, uint32_t nSize0, uint32_t nSize1) -{ - { - uint32_t mmAddOffset{}; - for (uint16_t mIdx = 0; mIdx < mSize; mIdx++) { - AscendC::MicroAPI::RegTensor biasReg{}; - AscendC::MicroAPI::RegTensor mmAddReg{}; - AscendC::MicroAPI::DataCopy(biasReg, bias); - mmAddOffset = mIdx * PER_BLOCK_SIZE; - uint32_t elementNum = nSize0; - AscendC::MicroAPI::MaskReg maskN = AscendC::MicroAPI::UpdateMask(elementNum); - AscendC::MicroAPI::DataCopy(mmAddReg, mmAdd + mmAddOffset); - AscendC::MicroAPI::Add(mmAddReg, mmAddReg, biasReg, maskN); - AscendC::MicroAPI::DataCopy( - mmAdd + mmAddOffset, mmAddReg, maskN); - if constexpr (ndNum == 1) { - continue; - } - elementNum = nSize1; - maskN = AscendC::MicroAPI::UpdateMask(elementNum); - mmAddOffset += nSize0; - AscendC::MicroAPI::DataCopy(biasReg, bias + nSize0); - AscendC::MicroAPI::DataCopy(mmAddReg, mmAdd + mmAddOffset); - AscendC::MicroAPI::Add(mmAddReg, mmAddReg, biasReg, maskN); - AscendC::MicroAPI::DataCopy( - mmAdd + mmAddOffset, mmAddReg, maskN); - } - } -} - -QBMM_BLOCK_EPILOGUE_PERTILE_CLASS_LOCAL_PARAMS -CATLASS_DEVICE void BlockEpilogue::AivPostProcess( - const AscendC::LocalTensor &mmAddUb) -{ - if (ubParams_.validM == 0) { - return; - } - if constexpr (AscendC::IsSameType::value) { - if (isBias_) { - auto biasUb = biasPingPongID_ == 2 ? biasUbPing_ : biasUbPong_; - AscendC::WaitFlag(biasPingPongID_); - auto mmAddUbAddr = reinterpret_cast<__ubuf__ CalcType *>(mmAddUb.GetPhyAddr()); - auto biasUbAddr = reinterpret_cast<__ubuf__ BiasType *>(biasUb.GetPhyAddr()); - if (ubParams_.ndNum == 1) { - AddBias<1U>( - (__ubuf__ CalcType *)mmAddUbAddr, (__ubuf__ BiasType *)biasUbAddr, ubParams_.validM, - ubParams_.validN[0], ubParams_.validN[1]); - } else if (ubParams_.ndNum == 2) { - AddBias<2U>( - (__ubuf__ CalcType *)mmAddUbAddr, (__ubuf__ BiasType *)biasUbAddr, ubParams_.validM, - ubParams_.validN[0], ubParams_.validN[1]); - } - AscendC::SetFlag(biasPingPongID_); - biasPingPongID_ = biasPingPongID_ ^ 1; - } - } - if constexpr (AscendC::IsSameType::value) { - // mov optimize in splitM, 0~63 + 64 ~127 -> 0~127 - AscendC::SetFlag(0); - AscendC::WaitFlag(0); - if (ubParams_.ndNum == 2 && !ubParams_.CopyOutWithSplitN) { // 2: 2ND, opt branch with splitM - uint32_t sumN = ubParams_.validN[0] + ubParams_.validN[1]; - CopyOut(mmAddUb, 0, ubParams_.validM, sumN, UB_TWO_BANK_ELEMS_B32 - sumN, problemShape_.n() - sumN, - ubParams_.offsetY[0]); - } else { - for (uint64_t ndIdx = 0; ndIdx < ubParams_.ndNum; ndIdx++) { - if (ubParams_.validN[ndIdx] > 0) { - CopyOut(mmAddUb[ndIdx * ubParams_.validN[0]], 0, ubParams_.validM, ubParams_.validN[ndIdx], - UB_TWO_BANK_ELEMS_B32 - ubParams_.validN[ndIdx], - problemShape_.n() - ubParams_.validN[ndIdx], ubParams_.offsetY[ndIdx]); - } - } - } - AscendC::SetFlag(0); - AscendC::WaitFlag(0); - } else { - AscendC::PipeBarrier(); - CastAndCopyOut(mmAddUb); - } -} - -QBMM_BLOCK_EPILOGUE_PERTILE_CLASS_LOCAL_PARAMS -CATLASS_DEVICE void BlockEpilogue::CopyOut( - const AscendC::LocalTensor &ubRes, uint16_t eventId, uint16_t blkCount, uint32_t blkLen, uint32_t srcStride, - uint32_t dstStride, uint64_t yOffset) -{ - AscendC::DataCopyExtParams copyParams{ - blkCount, static_cast(blkLen * sizeof(YType)), - static_cast(srcStride * sizeof(YType) / AscendC::ONE_BLK_SIZE), - static_cast(dstStride * sizeof(YType)), 0}; - AscendC::DataCopyPad(cGlobal_[yOffset], ubRes, copyParams); -} - -QBMM_BLOCK_EPILOGUE_PERTILE_CLASS_LOCAL_PARAMS -CATLASS_DEVICE void BlockEpilogue::CastAndCopyOut( - const AscendC::LocalTensor &mmAddUb) -{ - if (ubParams_.ndNum == 2 && !ubParams_.CopyOutWithSplitN) { // 2: 2ND, opt branch with splitM - uint32_t sumN = ubParams_.validN[0] + ubParams_.validN[1]; - uint32_t mSizePing = CeilDiv(ubParams_.validM, static_cast(QBMM_BUFFER_NUM)); - uint32_t mSize[QBMM_BUFFER_NUM] = {mSizePing, static_cast(ubParams_.validM - mSizePing)}; - for (uint32_t mDbIdx = 0; mDbIdx < QBMM_BUFFER_NUM; ++mDbIdx) { - if (mSize[mDbIdx] > 0 && sumN > 0) { - auto ubRes = mDbIdx == 0 ? &ubResPing_ : &ubResPong_; - AscendC::WaitFlag(mDbIdx); - AscendC::Cast( - *ubRes, mmAddUb[mDbIdx * mSizePing * UB_TWO_BANK_ELEMS_B32], AscendC::RoundMode::CAST_RINT, - mSize[mDbIdx] * UB_TWO_BANK_ELEMS_B32); - AscendC::SetFlag(mDbIdx); - AscendC::WaitFlag(mDbIdx); - CopyOut( - *ubRes, mDbIdx, mSize[mDbIdx], sumN, UB_TWO_BANK_ELEMS_B32 - sumN, problemShape_.n() - sumN, - ubParams_.offsetY[0] + mDbIdx * mSizePing * problemShape_.n()); - AscendC::SetFlag(mDbIdx); - } - } - } else { - for (uint64_t ndIdx = 0; ndIdx < ubParams_.ndNum; ndIdx++) { - auto ubRes = ndIdx == 0 ? &ubResPing_ : &ubResPong_; - if (ubParams_.validN[ndIdx] > 0) { - AscendC::UnaryRepeatParams repeatParam; - repeatParam.srcBlkStride = 1; - repeatParam.dstBlkStride = 1; - // write continuously - repeatParam.dstRepStride = CeilDiv(ubParams_.singleN, AscendC::ONE_BLK_SIZE / sizeof(YType)); - // srcStride is 16(512B / 32B), subBank0 256B one repeat - repeatParam.srcRepStride = QBMM_BMM_BLOCK_NUM; - AscendC::WaitFlag(ndIdx); - AscendC::Cast(*ubRes, mmAddUb[ubParams_.validN[0] * ndIdx], AscendC::RoundMode::CAST_RINT, - ubParams_.validN[ndIdx], ubParams_.validM, repeatParam); - AscendC::SetFlag(ndIdx); - AscendC::WaitFlag(ndIdx); - CopyOut(*ubRes, ndIdx, ubParams_.validM, ubParams_.validN[ndIdx], - ubParams_.singleN - ubParams_.validN[ndIdx], - problemShape_.n() - ubParams_.validN[ndIdx], ubParams_.offsetY[ndIdx]); - AscendC::SetFlag(ndIdx); - } - } - } -} - -} // namespace Catlass::Epilogue::Block - -#endif // CATLASS_EPILOGUE_BLOCK_EPILOGUE_PER_GROUP_PER_BLOCK_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_per_token_dequant_tla.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_per_token_dequant_tla.hpp deleted file mode 100644 index cdf6e9847a..0000000000 --- a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_per_token_dequant_tla.hpp +++ /dev/null @@ -1,784 +0,0 @@ -/** - * Copyright (c) 2026 Huawei Technologies Co., Ltd. - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of - * CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ - -#ifndef CATLASS_EPILOGUE_BLOCK_EPILOGUE_PER_TOKEN_DEQUANT_TLA_HPP -#define CATLASS_EPILOGUE_BLOCK_EPILOGUE_PER_TOKEN_DEQUANT_TLA_HPP - -#include "catlass/kda_catlass.hpp" -#include "catlass/arch/kda_resource.hpp" -#include "catlass/epilogue/kda_dispatch_policy.hpp" -#include "catlass/kda_gemm_coord.hpp" -#include "catlass/gemm/kda_gemm_type.hpp" -#include "catlass/kda_matrix_coord.hpp" -#include "catlass/layout/kda_layout.hpp" -#include "catlass/detail/kda_callback.hpp" -#include "tla/kda_tla_layout.hpp" -#include "tla/kda_tensor.hpp" - -namespace Catlass::Epilogue::Block { - -template < - uint32_t UB_STAGES_, - class ElementC_, - class ElementScale_, - class ElementPerTokenScale_, - class ElementD_, - class TileRowBroadcastMul_, - class TileBroadcastOneBlk_, - class TileOneBlkColumnBroadcastMul_, - class TileCopy_, - class EpilogueTileSwizzle_> -class BlockEpilogue< - EpilogueAtlasA2PerTokenDequantTla, - ElementC_, - ElementScale_, - ElementPerTokenScale_, - ElementD_, - TileRowBroadcastMul_, - TileBroadcastOneBlk_, - TileOneBlkColumnBroadcastMul_, - TileCopy_, - EpilogueTileSwizzle_> { -public: - using DispatchPolicy = EpilogueAtlasA2PerTokenDequantTla; - using ArchTag = typename DispatchPolicy::ArchTag; - static constexpr uint32_t UB_STAGES = UB_STAGES_; - using TileCopy = TileCopy_; - - // Data infos - using ElementC = ElementC_; - using LayoutC = typename TileCopy::LayoutC; - using ElementScale = ElementScale_; - using LayoutScale = typename TileCopy::LayoutX; - using ElementPerTokenScale = ElementPerTokenScale_; - using LayoutPerTokenScale = typename TileCopy::LayoutY; - using ElementD = ElementD_; - using LayoutD = typename TileCopy::LayoutD; - - // Check data infos - static_assert( - std::is_same_v && (std::is_same_v || std::is_same_v) && - std::is_same_v && std::is_same_v, - "The element type template parameters of BlockEpilogue are wrong"); - static_assert( - tla::detail::isRowMajor::value && tla::detail::isRowMajor::value, - "The layout template parameters of BlockEpilogue are wrong"); - - // Tile compute ops - using TileRowBroadcastMul = TileRowBroadcastMul_; - using TileBroadcastOneBlk = TileBroadcastOneBlk_; - using TileOneBlkColumnBroadcastMul = TileOneBlkColumnBroadcastMul_; - - using EpilogueTileSwizzle = EpilogueTileSwizzle_; - - using TileShape = typename TileRowBroadcastMul::TileShape; - - static constexpr uint32_t ELE_NUM_PER_BLK = BYTE_PER_BLK / sizeof(float); - - static_assert( - TileShape::ROW == TileBroadcastOneBlk::COMPUTE_LENGTH && - std::is_same_v, - "TileShape must be consistent for all tile compute ops"); - - static_assert( - (UB_STAGES * (TileShape::COUNT * sizeof(ElementC) + TileShape::COLUMN * sizeof(ElementScale) + TileShape::ROW * sizeof(ElementPerTokenScale) + TileShape::COUNT * sizeof(ElementD)) + (TileShape::COUNT + TileShape::COLUMN + TileShape::COUNT + TileShape::ROW) * sizeof(float) + TileShape::ROW * BYTE_PER_BLK) <= ArchTag::UB_SIZE, - "TileShape is too large to fit in UB"); - - struct Params { - GM_ADDR ptrScale{nullptr}; - LayoutScale layoutScale{}; - GM_ADDR ptrPerTokenScale{nullptr}; - LayoutPerTokenScale layoutPerTokenScale{}; - GM_ADDR ptrD{nullptr}; - LayoutD layoutD{}; - Params() = default; - }; - - CATLASS_DEVICE - BlockEpilogue(Arch::Resource const &resource, Params const ¶ms = Params{}, uint32_t sharedUbSize = 0) - : params(params) - { - uint32_t ubOffset = sharedUbSize; - int32_t eventVMTE2 = 0; - int32_t eventMTE2V = 0; - int32_t eventMTE3V = 0; - int32_t eventVMTE3 = 0; - for (uint32_t i = 0; i < UB_STAGES; ++i) { - ubCList[i] = resource.ubBuf.template GetBufferByByte(ubOffset); - ubOffset += TileShape::COUNT * sizeof(ElementC); - ubScaleList[i] = resource.ubBuf.template GetBufferByByte(ubOffset); - ubOffset += TileShape::COLUMN * sizeof(ElementScale); - ubPerTokenScaleList[i] = resource.ubBuf.template GetBufferByByte(ubOffset); - ubOffset += TileShape::ROW * sizeof(ElementPerTokenScale); - ubDList[i] = resource.ubBuf.template GetBufferByByte(ubOffset); - ubOffset += TileShape::COUNT * sizeof(ElementD); - - eventUbCVMTE2List[i] = eventVMTE2++; - eventUbCMTE2VList[i] = eventMTE2V++; - eventUbScaleVMTE2List[i] = eventVMTE2++; - eventUbScaleMTE2VList[i] = eventMTE2V++; - eventUbPerTokenScaleVMTE2List[i] = eventVMTE2++; - eventUbPerTokenScaleMTE2VList[i] = eventMTE2V++; - eventUbDMTE3VList[i] = eventMTE3V++; - eventUbDVMTE3List[i] = eventVMTE3++; - - AscendC::SetFlag(eventUbCVMTE2List[i]); - AscendC::SetFlag(eventUbScaleVMTE2List[i]); - AscendC::SetFlag(eventUbPerTokenScaleVMTE2List[i]); - AscendC::SetFlag(eventUbDMTE3VList[i]); - } - ubCFp32 = resource.ubBuf.template GetBufferByByte(ubOffset); - ubOffset += TileShape::COUNT * sizeof(float); - ubScaleFp32 = resource.ubBuf.template GetBufferByByte(ubOffset); - ubOffset += TileShape::COLUMN * sizeof(float); - ubMul = resource.ubBuf.template GetBufferByByte(ubOffset); - ubOffset += TileShape::COUNT * sizeof(float); - ubPerTokenScaleFp32 = resource.ubBuf.template GetBufferByByte(ubOffset); - ubOffset += TileShape::ROW * sizeof(float); - ubPerTokenScaleFp32Brcb = resource.ubBuf.template GetBufferByByte(ubOffset); - ubOffset += TileShape::ROW * BYTE_PER_BLK; - ubPerTokenMul = ubMul; - } - - CATLASS_DEVICE - ~BlockEpilogue() - { - for (uint32_t i = 0; i < UB_STAGES; ++i) { - AscendC::WaitFlag(eventUbCVMTE2List[i]); - AscendC::WaitFlag(eventUbScaleVMTE2List[i]); - AscendC::WaitFlag(eventUbPerTokenScaleVMTE2List[i]); - AscendC::WaitFlag(eventUbDMTE3VList[i]); - } - } - - CATLASS_DEVICE - void UpdateParams(Params const ¶ms_) - { - params = params_; - } - - template - CATLASS_DEVICE void operator()( - TensorC &tensorBlockC, TensorScale &tensorBlockScale, TensorPerTokenScaleScale &tensorBlockPerTokenScaleScale, - TensorD &tensorBlockD, GemmCoord const &actualBlockShapeMNK, Callback &&callback = Callback{}) - { - if (actualBlockShapeMNK.k() == 0) { - return; - } - callback(); - - MatrixCoord actualBlockShape = actualBlockShapeMNK.GetCoordMN(); - - using CopyGmToUbC = typename TileCopy::template CopyGmToUbC; - using CopyGmToUbScale = typename TileCopy::template CopyGmToUbX; - using CopyGmToUbPerTokenScale = typename TileCopy::template CopyGmToUbY; - using CopyUbToGmD = typename TileCopy::template CopyUbToGmD; - CopyGmToUbC copyGmToUbC; - CopyGmToUbScale copyGmToUbScale; - CopyGmToUbPerTokenScale copyGmToUbPerTokenScale; - CopyUbToGmD copyUbToGmD; - - auto ubTileStride = static_cast(TileShape::COLUMN); - auto ubTileStrideRow = static_cast(TileShape::ROW); - auto tileShape = TileShape::ToCoord(); - EpilogueTileSwizzle epilogueTileSwizzle(actualBlockShape, tileShape); - uint32_t tileLoops = epilogueTileSwizzle.GetLoops(); - uint32_t subblockIdx = AscendC::GetSubBlockIdx(); - uint32_t subblockNum = AscendC::GetSubBlockNum(); - for (uint32_t loopIdx = subblockIdx; loopIdx < tileLoops; loopIdx += subblockNum) { - auto tileCoord = epilogueTileSwizzle.GetTileCoord(loopIdx); - auto actualTileShape = epilogueTileSwizzle.GetActualTileShape(tileCoord); - MatrixCoord tileOffsetInBlock = tileCoord * tileShape; - auto tileOffsetInBlockRow = tileOffsetInBlock.row(); - auto tileOffsetInBlockColumn = tileOffsetInBlock.column(); - - // build tensor C block in GM - auto tensorSubBlockC = GetTile( - tensorBlockC, tla::MakeCoord(tileOffsetInBlockRow, tileOffsetInBlockColumn), - tla::MakeShape(actualTileShape.row(), actualTileShape.column())); - // build tensor C block in UB - auto &ubC = ubCList[ubListId]; - auto layoutUbC = tla::MakeLayout( - tla::MakeShape(actualTileShape.row(), actualTileShape.column()), tla::MakeStride(ubTileStride, tla::Int<1>{})); - auto tensorUbC = tla::MakeTensor(ubC, layoutUbC, Arch::PositionUB{}); - // copy tensor C from GM to UB - AscendC::WaitFlag(eventUbCVMTE2List[ubListId]); - copyGmToUbC(tensorUbC, tensorSubBlockC); - AscendC::SetFlag(eventUbCMTE2VList[ubListId]); - - /////////////////////////////////////////////////// - // build tensor Scale block in GM - auto tensorSubBlockScale = GetTile( - tensorBlockScale, tla::MakeCoord(0, tileOffsetInBlockColumn), - tla::MakeShape(tla::Int<1>{}, actualTileShape.column())); - // build tensor Scale block in UB - auto &ubScale = ubScaleList[ubListId]; - auto layoutUbScale = tla::MakeLayout( - tla::MakeShape(tla::Int<1>{}, actualTileShape.column()), tla::MakeStride(ubTileStride, tla::Int<1>{})); - auto tensorUbScale = tla::MakeTensor(ubScale, layoutUbScale, Arch::PositionUB{}); - // copy tensor Scale from GM to UB - AscendC::WaitFlag(eventUbScaleVMTE2List[ubListId]); - copyGmToUbScale(tensorUbScale, tensorSubBlockScale); - AscendC::SetFlag(eventUbScaleMTE2VList[ubListId]); - - /////////////////////////////////////////////////// - // build tensor PerTokenScale block in GM - auto tensorSubBlockPerTokenScale = GetTile( - tensorBlockPerTokenScaleScale, tla::MakeCoord(0, tileOffsetInBlockRow), - tla::MakeShape(1, actualTileShape.row())); - // build tensor PerTokenScale block in UB - auto &ubPerTokenScale = ubPerTokenScaleList[ubListId]; - auto layoutUbPerTokenScale = tla::MakeLayout( - tla::MakeShape(tla::Int<1>{}, actualTileShape.row()), tla::MakeStride(ubTileStrideRow, tla::Int<1>{})); - auto tensorUbPerTokenScale = tla::MakeTensor(ubPerTokenScale, layoutUbPerTokenScale, Arch::PositionUB{}); - // copy tensor PerTokenScale from GM to UB - AscendC::WaitFlag(eventUbPerTokenScaleVMTE2List[ubListId]); - copyGmToUbPerTokenScale(tensorUbPerTokenScale, tensorSubBlockPerTokenScale); - AscendC::SetFlag(eventUbPerTokenScaleMTE2VList[ubListId]); - - /////////////////////////////////////////////////// - auto tensorUbScaleFp32 = tla::MakeTensor(ubScaleFp32, layoutUbScale, Arch::PositionUB{}); - auto tensorUbPerTokenScaleFp32 = tla::MakeTensor(ubPerTokenScaleFp32, layoutUbPerTokenScale, Arch::PositionUB{}); - auto tensorUbCFp32 = tla::MakeTensor(ubCFp32, layoutUbC, Arch::PositionUB{}); - auto tensorUbMul = tla::MakeTensor(ubMul, layoutUbC, Arch::PositionUB{}); - - auto layoutUbPerTokenScaleBrcb = tla::MakeLayout( - tla::MakeShape(actualTileShape.row(), ELE_NUM_PER_BLK), tla::MakeStride(ELE_NUM_PER_BLK, tla::Int<1>{})); - auto tensorUbPerTokenScaleFp32Brcb = tla::MakeTensor(ubPerTokenScaleFp32Brcb, layoutUbPerTokenScaleBrcb, Arch::PositionUB{}); - auto tensorUbPerTokenMul = tla::MakeTensor(ubPerTokenMul, layoutUbC, Arch::PositionUB{}); - - AscendC::WaitFlag(eventUbCMTE2VList[ubListId]); - AscendC::Cast(ubCFp32, ubC, AscendC::RoundMode::CAST_RINT, TileShape::COUNT); - AscendC::SetFlag(eventUbCVMTE2List[ubListId]); - - AscendC::WaitFlag(eventUbScaleMTE2VList[ubListId]); - AscendC::Cast(ubScaleFp32, ubScale, AscendC::RoundMode::CAST_NONE, TileShape::COLUMN); - AscendC::SetFlag(eventUbScaleVMTE2List[ubListId]); - - AscendC::WaitFlag(eventUbPerTokenScaleMTE2VList[ubListId]); - AscendC::Cast(ubPerTokenScaleFp32, ubPerTokenScale, AscendC::RoundMode::CAST_NONE, TileShape::ROW); - AscendC::SetFlag(eventUbPerTokenScaleVMTE2List[ubListId]); - - AscendC::PipeBarrier(); - tileRowBroadcastMul(tensorUbMul, tensorUbCFp32, tensorUbScaleFp32); - tileBroadcastOneBlk(tensorUbPerTokenScaleFp32Brcb, tensorUbPerTokenScaleFp32); - AscendC::PipeBarrier(); - tileOneBlkColumnBroadcastMul(tensorUbPerTokenMul, tensorUbMul, tensorUbPerTokenScaleFp32Brcb); - AscendC::PipeBarrier(); - - auto &ubD = ubDList[ubListId]; - - AscendC::WaitFlag(eventUbDMTE3VList[ubListId]); - AscendC::Cast(ubD, ubPerTokenMul, AscendC::RoundMode::CAST_RINT, TileShape::COUNT); - AscendC::SetFlag(eventUbDVMTE3List[ubListId]); - - // build tensor D block in GM - auto tensorSubBlockD = GetTile( - tensorBlockD, tla::MakeCoord(tileOffsetInBlockRow, tileOffsetInBlockColumn), - tla::MakeShape(actualTileShape.row(), actualTileShape.column())); - // build tensor D block in UB - auto tensorUbD = tla::MakeTensor(ubD, layoutUbC, Arch::PositionUB{}); - - AscendC::WaitFlag(eventUbDVMTE3List[ubListId]); - copyUbToGmD(tensorSubBlockD, tensorUbD); - AscendC::SetFlag(eventUbDMTE3VList[ubListId]); - - ubListId = (ubListId + 1 < UB_STAGES) ? (ubListId + 1) : 0; - } - } - -private: - Params params; - - AscendC::LocalTensor ubCList[UB_STAGES]; - AscendC::LocalTensor ubScaleList[UB_STAGES]; - AscendC::LocalTensor ubPerTokenScaleList[UB_STAGES]; - AscendC::LocalTensor ubDList[UB_STAGES]; - - int32_t eventUbCVMTE2List[UB_STAGES]; - int32_t eventUbCMTE2VList[UB_STAGES]; - int32_t eventUbScaleVMTE2List[UB_STAGES]; - int32_t eventUbScaleMTE2VList[UB_STAGES]; - int32_t eventUbPerTokenScaleVMTE2List[UB_STAGES]; - int32_t eventUbPerTokenScaleMTE2VList[UB_STAGES]; - int32_t eventUbDMTE3VList[UB_STAGES]; - int32_t eventUbDVMTE3List[UB_STAGES]; - - uint32_t ubListId{0}; - - AscendC::LocalTensor ubCFp32; - AscendC::LocalTensor ubScaleFp32; - AscendC::LocalTensor ubMul; - AscendC::LocalTensor ubPerTokenScaleFp32; - AscendC::LocalTensor ubPerTokenScaleFp32Brcb; - AscendC::LocalTensor ubPerTokenMul; - - TileRowBroadcastMul tileRowBroadcastMul; - TileBroadcastOneBlk tileBroadcastOneBlk; - TileOneBlkColumnBroadcastMul tileOneBlkColumnBroadcastMul; -}; - -template < - uint32_t UB_STAGES_, - class ElementC_, - class ElementD_, - class TileRowBroadcastMul_, - class TileBroadcastOneBlk_, - class TileOneBlkColumnBroadcastMul_, - class TileCopy_, - class EpilogueTileSwizzle_> -class BlockEpilogue< - EpilogueAtlasA2PerTokenDequantTla, - ElementC_, - float, - float, - ElementD_, - TileRowBroadcastMul_, - TileBroadcastOneBlk_, - TileOneBlkColumnBroadcastMul_, - TileCopy_, - EpilogueTileSwizzle_> { -public: - using DispatchPolicy = EpilogueAtlasA2PerTokenDequantTla; - using ArchTag = typename DispatchPolicy::ArchTag; - static constexpr uint32_t UB_STAGES = UB_STAGES_; - using TileCopy = TileCopy_; - - // Data infos - using ElementC = ElementC_; - using LayoutC = typename TileCopy::LayoutC; - using ElementScale = float; - using LayoutScale = typename TileCopy::LayoutX; - using ElementPerTokenScale = float; - using LayoutPerTokenScale = typename TileCopy::LayoutY; - using ElementD = ElementD_; - using LayoutD = typename TileCopy::LayoutD; - - // Check data infos - static_assert( - std::is_same_v && (std::is_same_v || std::is_same_v || - std::is_same_v), - "The element type template parameters of BlockEpilogue are wrong"); - static_assert( - tla::detail::isRowMajor::value && tla::detail::isRowMajor::value, - "The layout template parameters of BlockEpilogue are wrong"); - - // Tile compute ops - using TileRowBroadcastMul = TileRowBroadcastMul_; - using TileBroadcastOneBlk = TileBroadcastOneBlk_; - using TileOneBlkColumnBroadcastMul = TileOneBlkColumnBroadcastMul_; - - using EpilogueTileSwizzle = EpilogueTileSwizzle_; - - using TileShape = typename TileRowBroadcastMul::TileShape; - - static constexpr uint32_t ELE_NUM_PER_BLK = BYTE_PER_BLK / sizeof(float); - - static_assert( - TileShape::ROW == TileBroadcastOneBlk::COMPUTE_LENGTH && - std::is_same_v, - "TileShape must be consistent for all tile compute ops"); - - static_assert( - (UB_STAGES * (TileShape::COUNT * sizeof(ElementC) + TileShape::COLUMN * sizeof(ElementScale) + TileShape::ROW * sizeof(ElementPerTokenScale) + TileShape::COUNT * sizeof(ElementD)) + (TileShape::COUNT + TileShape::COUNT) * sizeof(float) + TileShape::ROW * BYTE_PER_BLK) <= ArchTag::UB_SIZE, - "TileShape is too large to fit in UB"); - - struct Params { - GM_ADDR ptrScale{nullptr}; - LayoutScale layoutScale{}; - GM_ADDR ptrPerTokenScale{nullptr}; - LayoutPerTokenScale layoutPerTokenScale{}; - GM_ADDR ptrD{nullptr}; - LayoutD layoutD{}; - Params() = default; - }; - - CATLASS_DEVICE - BlockEpilogue(Arch::Resource const &resource, Params const ¶ms = Params{}, uint32_t sharedUbSize = 0) - : params(params) - { - uint32_t ubOffset = sharedUbSize; - int32_t eventVMTE2 = 0; - int32_t eventMTE2V = 0; - int32_t eventMTE3V = 0; - int32_t eventVMTE3 = 0; - for (uint32_t i = 0; i < UB_STAGES; ++i) { - ubCList[i] = resource.ubBuf.template GetBufferByByte(ubOffset); - ubOffset += TileShape::COUNT * sizeof(ElementC); - ubScaleList[i] = resource.ubBuf.template GetBufferByByte(ubOffset); - ubOffset += TileShape::COLUMN * sizeof(ElementScale); - ubPerTokenScaleList[i] = resource.ubBuf.template GetBufferByByte(ubOffset); - ubOffset += TileShape::ROW * sizeof(ElementPerTokenScale); - ubDList[i] = resource.ubBuf.template GetBufferByByte(ubOffset); - ubOffset += TileShape::COUNT * sizeof(ElementD); - - eventUbCVMTE2List[i] = eventVMTE2++; - eventUbCMTE2VList[i] = eventMTE2V++; - eventUbScaleVMTE2List[i] = eventVMTE2++; - eventUbScaleMTE2VList[i] = eventMTE2V++; - eventUbPerTokenScaleVMTE2List[i] = eventVMTE2++; - eventUbPerTokenScaleMTE2VList[i] = eventMTE2V++; - eventUbDMTE3VList[i] = eventMTE3V++; - eventUbDVMTE3List[i] = eventVMTE3++; - - AscendC::SetFlag(eventUbCVMTE2List[i]); - AscendC::SetFlag(eventUbScaleVMTE2List[i]); - AscendC::SetFlag(eventUbPerTokenScaleVMTE2List[i]); - AscendC::SetFlag(eventUbDMTE3VList[i]); - } - ubCFp32 = resource.ubBuf.template GetBufferByByte(ubOffset); - ubOffset += TileShape::COUNT * sizeof(float); - ubMul = resource.ubBuf.template GetBufferByByte(ubOffset); - ubOffset += TileShape::COUNT * sizeof(float); - ubPerTokenScaleBrcb = resource.ubBuf.template GetBufferByByte(ubOffset); - ubOffset += TileShape::ROW * BYTE_PER_BLK; - ubPerTokenMul = ubCFp32; - } - - CATLASS_DEVICE - ~BlockEpilogue() - { - for (uint32_t i = 0; i < UB_STAGES; ++i) { - AscendC::WaitFlag(eventUbCVMTE2List[i]); - AscendC::WaitFlag(eventUbScaleVMTE2List[i]); - AscendC::WaitFlag(eventUbPerTokenScaleVMTE2List[i]); - AscendC::WaitFlag(eventUbDMTE3VList[i]); - } - } - - CATLASS_DEVICE - void UpdateParams(Params const ¶ms_) - { - params = params_; - } - - template - CATLASS_DEVICE void operator()( - TensorC &tensorBlockC, TensorScale &tensorBlockScale, TensorPerTokenScaleScale &tensorBlockPerTokenScaleScale, - TensorD &tensorBlockD, GemmCoord const &actualBlockShapeMNK, Callback &&callback = Callback{}) - { - if (actualBlockShapeMNK.k() == 0) { - return; - } - callback(); - - MatrixCoord actualBlockShape = actualBlockShapeMNK.GetCoordMN(); - - using CopyGmToUbC = typename TileCopy::template CopyGmToUbC; - using CopyGmToUbScale = typename TileCopy::template CopyGmToUbX; - using CopyGmToUbPerTokenScale = typename TileCopy::template CopyGmToUbY; - using CopyUbToGmD = typename TileCopy::template CopyUbToGmD; - CopyGmToUbC copyGmToUbC; - CopyGmToUbScale copyGmToUbScale; - CopyGmToUbPerTokenScale copyGmToUbPerTokenScale; - CopyUbToGmD copyUbToGmD; - - auto ubTileStride = static_cast(TileShape::COLUMN); - auto ubTileStrideRow = static_cast(TileShape::ROW); - auto tileShape = TileShape::ToCoord(); - EpilogueTileSwizzle epilogueTileSwizzle(actualBlockShape, tileShape); - uint32_t tileLoops = epilogueTileSwizzle.GetLoops(); - uint32_t subblockIdx = AscendC::GetSubBlockIdx(); - uint32_t subblockNum = AscendC::GetSubBlockNum(); - for (uint32_t loopIdx = subblockIdx; loopIdx < tileLoops; loopIdx += subblockNum) { - auto tileCoord = epilogueTileSwizzle.GetTileCoord(loopIdx); - auto actualTileShape = epilogueTileSwizzle.GetActualTileShape(tileCoord); - MatrixCoord tileOffsetInBlock = tileCoord * tileShape; - auto tileOffsetInBlockRow = tileOffsetInBlock.row(); - auto tileOffsetInBlockColumn = tileOffsetInBlock.column(); - - // build tensor C block in GM - auto tensorSubBlockC = GetTile( - tensorBlockC, tla::MakeCoord(tileOffsetInBlockRow, tileOffsetInBlockColumn), - tla::MakeShape(actualTileShape.row(), actualTileShape.column())); - // build tensor C block in UB - auto &ubC = ubCList[ubListId]; - auto layoutUbC = tla::MakeLayout( - tla::MakeShape(actualTileShape.row(), actualTileShape.column()), tla::MakeStride(ubTileStride, tla::Int<1>{})); - auto tensorUbC = tla::MakeTensor(ubC, layoutUbC, Arch::PositionUB{}); - // copy tensor C from GM to UB - AscendC::WaitFlag(eventUbCVMTE2List[ubListId]); - copyGmToUbC(tensorUbC, tensorSubBlockC); - AscendC::SetFlag(eventUbCMTE2VList[ubListId]); - - /////////////////////////////////////////////////// - // build tensor Scale block in GM - auto tensorSubBlockScale = GetTile( - tensorBlockScale, tla::MakeCoord(0, tileOffsetInBlockColumn), - tla::MakeShape(tla::Int<1>{}, actualTileShape.column())); - // build tensor Scale block in UB - auto &ubScale = ubScaleList[ubListId]; - auto layoutUbScale = tla::MakeLayout( - tla::MakeShape(tla::Int<1>{}, actualTileShape.column()), tla::MakeStride(ubTileStride, tla::Int<1>{})); - auto tensorUbScale = tla::MakeTensor(ubScale, layoutUbScale, Arch::PositionUB{}); - // copy tensor Scale from GM to UB - AscendC::WaitFlag(eventUbScaleVMTE2List[ubListId]); - copyGmToUbScale(tensorUbScale, tensorSubBlockScale); - AscendC::SetFlag(eventUbScaleMTE2VList[ubListId]); - - /////////////////////////////////////////////////// - // build tensor PerTokenScale block in GM - auto tensorSubBlockPerTokenScale = GetTile( - tensorBlockPerTokenScaleScale, tla::MakeCoord(0, tileOffsetInBlockRow), - tla::MakeShape(tla::Int<1>{}, actualTileShape.row())); - // build tensor PerTokenScale block in UB - auto &ubPerTokenScale = ubPerTokenScaleList[ubListId]; - auto layoutUbPerTokenScale = tla::MakeLayout( - tla::MakeShape(tla::Int<1>{}, actualTileShape.row()), tla::MakeStride(ubTileStrideRow, tla::Int<1>{})); - auto tensorUbPerTokenScale = tla::MakeTensor(ubPerTokenScale, layoutUbPerTokenScale, Arch::PositionUB{}); - // copy tensor PerTokenScale from GM to UB - AscendC::WaitFlag(eventUbPerTokenScaleVMTE2List[ubListId]); - copyGmToUbPerTokenScale(tensorUbPerTokenScale, tensorSubBlockPerTokenScale); - AscendC::SetFlag(eventUbPerTokenScaleMTE2VList[ubListId]); - - /////////////////////////////////////////////////// - auto tensorUbCFp32 = tla::MakeTensor(ubCFp32, layoutUbC, Arch::PositionUB{}); - auto tensorUbMul = tla::MakeTensor(ubMul, layoutUbC, Arch::PositionUB{}); - - auto layoutUbPerTokenScaleBrcb = tla::MakeLayout( - tla::MakeShape(actualTileShape.row(), ELE_NUM_PER_BLK), tla::MakeStride(ELE_NUM_PER_BLK, tla::Int<1>{})); - auto tensorUbPerTokenScaleBrcb = tla::MakeTensor(ubPerTokenScaleBrcb, layoutUbPerTokenScaleBrcb, Arch::PositionUB{}); - auto tensorUbPerTokenMul = tla::MakeTensor(ubPerTokenMul, layoutUbC, Arch::PositionUB{}); - - AscendC::WaitFlag(eventUbCMTE2VList[ubListId]); - AscendC::Cast(ubCFp32, ubC, AscendC::RoundMode::CAST_RINT, TileShape::COUNT); - AscendC::SetFlag(eventUbCVMTE2List[ubListId]); - - AscendC::PipeBarrier(); - AscendC::WaitFlag(eventUbScaleMTE2VList[ubListId]); - tileRowBroadcastMul(tensorUbMul, tensorUbCFp32, tensorUbScale); - AscendC::SetFlag(eventUbScaleVMTE2List[ubListId]); - - AscendC::PipeBarrier(); - AscendC::WaitFlag(eventUbPerTokenScaleMTE2VList[ubListId]); - tileBroadcastOneBlk(tensorUbPerTokenScaleBrcb, tensorUbPerTokenScale); - AscendC::SetFlag(eventUbPerTokenScaleVMTE2List[ubListId]); - - // build tensor D block in UB - auto &ubD = ubDList[ubListId]; - auto tensorUbD = tla::MakeTensor(ubD, layoutUbC, Arch::PositionUB{}); - - if constexpr (std::is_same_v) { - AscendC::PipeBarrier(); - AscendC::WaitFlag(eventUbDMTE3VList[ubListId]); - tileOneBlkColumnBroadcastMul(tensorUbD, tensorUbMul, tensorUbPerTokenScaleBrcb); - AscendC::SetFlag(eventUbDVMTE3List[ubListId]); - } else if constexpr (std::is_same_v) { - AscendC::PipeBarrier(); - tileOneBlkColumnBroadcastMul(tensorUbPerTokenMul, tensorUbMul, tensorUbPerTokenScaleBrcb); - AscendC::PipeBarrier(); - - AscendC::WaitFlag(eventUbDMTE3VList[ubListId]); - AscendC::Cast(ubD, ubPerTokenMul, AscendC::RoundMode::CAST_NONE, TileShape::COUNT); - AscendC::SetFlag(eventUbDVMTE3List[ubListId]); - } else { // std::is_same_v - AscendC::PipeBarrier(); - tileOneBlkColumnBroadcastMul(tensorUbPerTokenMul, tensorUbMul, tensorUbPerTokenScaleBrcb); - AscendC::PipeBarrier(); - - AscendC::WaitFlag(eventUbDMTE3VList[ubListId]); - AscendC::Cast(ubD, ubPerTokenMul, AscendC::RoundMode::CAST_RINT, TileShape::COUNT); - AscendC::SetFlag(eventUbDVMTE3List[ubListId]); - } - - // build tensor D block in GM - auto tensorSubBlockD = GetTile( - tensorBlockD, tla::MakeCoord(tileOffsetInBlockRow, tileOffsetInBlockColumn), - tla::MakeShape(actualTileShape.row(), actualTileShape.column())); - - AscendC::WaitFlag(eventUbDVMTE3List[ubListId]); - copyUbToGmD(tensorSubBlockD, tensorUbD); - AscendC::SetFlag(eventUbDMTE3VList[ubListId]); - - ubListId = (ubListId + 1 < UB_STAGES) ? (ubListId + 1) : 0; - } - } - -private: - Params params; - - AscendC::LocalTensor ubCList[UB_STAGES]; - AscendC::LocalTensor ubScaleList[UB_STAGES]; - AscendC::LocalTensor ubPerTokenScaleList[UB_STAGES]; - AscendC::LocalTensor ubDList[UB_STAGES]; - - int32_t eventUbCVMTE2List[UB_STAGES]; - int32_t eventUbCMTE2VList[UB_STAGES]; - int32_t eventUbScaleVMTE2List[UB_STAGES]; - int32_t eventUbScaleMTE2VList[UB_STAGES]; - int32_t eventUbPerTokenScaleVMTE2List[UB_STAGES]; - int32_t eventUbPerTokenScaleMTE2VList[UB_STAGES]; - int32_t eventUbDMTE3VList[UB_STAGES]; - int32_t eventUbDVMTE3List[UB_STAGES]; - - uint32_t ubListId{0}; - - AscendC::LocalTensor ubCFp32; - AscendC::LocalTensor ubMul; - AscendC::LocalTensor ubPerTokenScaleBrcb; - AscendC::LocalTensor ubPerTokenMul; - - TileRowBroadcastMul tileRowBroadcastMul; - TileBroadcastOneBlk tileBroadcastOneBlk; - TileOneBlkColumnBroadcastMul tileOneBlkColumnBroadcastMul; -}; - -template < - uint32_t UB_STAGES_, - class TileShape_, - class ElementSrc_, - class ElementScale_, - class ElementPerToken_, - class ElementDst_, - class TilePerTokenDequant_, - class TileCopy_> -class BlockEpilogue< - EpilogueAscend950PerTokenDequantTla, - TileShape_, - ElementSrc_, - ElementScale_, - ElementPerToken_, - ElementDst_, - TilePerTokenDequant_, - TileCopy_> { -public: - using DispatchPolicy = EpilogueAscend950PerTokenDequantTla; - using ArchTag = typename DispatchPolicy::ArchTag; - using TileCopy = TileCopy_; - - using ElementSrc = ElementSrc_; - using ElementScale = ElementScale_; - using LayoutScale = typename TileCopy::LayoutTagX; - using ElementPerToken = ElementPerToken_; - using LayoutPerToken = typename TileCopy::LayoutTagY; - using ElementDst = ElementDst_; - using LayoutDst = typename TileCopy::LayoutTagD; - - using TilePerTokenDequant = TilePerTokenDequant_; - using TileShape = TileShape_; - - static constexpr uint32_t UB_STAGES = UB_STAGES_; - static constexpr int16_t N_BASE_SIZE = static_cast(TileShape::COLUMN); - static constexpr uint16_t CV_RATIO = 2; - static constexpr uint32_t BLOCK_SIZE = TileShape::COUNT / CV_RATIO; - - static_assert( - UB_STAGES * (BLOCK_SIZE * sizeof(ElementSrc) + TileShape::COLUMN / CV_RATIO * sizeof(ElementScale) + TileShape::ROW / CV_RATIO * sizeof(ElementPerToken)) <= ArchTag::UB_SIZE, - "TileShape is too large to fit in UB"); - - CATLASS_DEVICE - BlockEpilogue(Arch::Resource const &resource, uint32_t &ubOffset) - { - int32_t eventVMTE2 = 0; - int32_t eventMTE2V = 0; - int32_t eventMTE3V = 0; - int32_t eventVMTE3 = 0; - for (uint32_t i = 0; i < UB_STAGES; ++i) { - ubScaleList[i] = resource.ubBuf.template GetBufferByByte(ubOffset); - ubOffset += TileShape::COLUMN * sizeof(ElementScale); - ubPerTokenList[i] = resource.ubBuf.template GetBufferByByte(ubOffset); - ubOffset += TileShape::ROW * sizeof(ElementPerToken); - - eventUbScaleVMTE2List[i] = eventVMTE2++; - eventUbScaleMTE2VList[i] = eventMTE2V++; - eventUbPerTokenVMTE2List[i] = eventVMTE2++; - eventUbPerTokenMTE2VList[i] = eventMTE2V++; - eventUbDMTE3VList[i] = eventMTE3V++; - eventUbDVMTE3List[i] = eventVMTE3++; - - AscendC::SetFlag(eventUbScaleVMTE2List[i]); - AscendC::SetFlag(eventUbPerTokenVMTE2List[i]); - AscendC::SetFlag(eventUbDMTE3VList[i]); - } - } - - CATLASS_DEVICE - ~BlockEpilogue() - { - for (uint32_t i = 0; i < UB_STAGES; ++i) { - AscendC::WaitFlag(eventUbScaleVMTE2List[i]); - AscendC::WaitFlag(eventUbPerTokenVMTE2List[i]); - AscendC::WaitFlag(eventUbDMTE3VList[i]); - } - } - - template - CATLASS_DEVICE void operator()( - TensorDst gmDequantOut, - TensorSrc ubGmmRes, - TensorScale gmScale, - TensorPerToken gmPerToken, - uint32_t ubListId) - { - using CopyGmToUbScale = typename TileCopy::template CopyGmToUbX; - using CopyGmToUbPerToken = typename TileCopy::template CopyGmToUbY; - using CopyUbToGmDequant = typename TileCopy::template CopyUbToGmD; - - CopyGmToUbScale copyGmToUbScale; - CopyGmToUbPerToken copyGmToUbPerToken; - CopyUbToGmDequant copyUbToGmDequant; - TilePerTokenDequant tilePerTokenDequant; - - uint32_t m = tla::get<0>(ubGmmRes.shape()); - uint32_t n = tla::get<1>(ubGmmRes.shape()); - - auto scaleLayout = tla::MakeLayout(n); - auto ubScale = tla::MakeTensor( - ubScaleList[ubListId], - scaleLayout, - Arch::PositionUB{}); - - auto perTokenLayout = tla::MakeLayout(m); - auto ubPerToken = tla::MakeTensor( - ubPerTokenList[ubListId], - perTokenLayout, - Arch::PositionUB{}); - - auto ubDequantOut = tla::MakeTensor( - ubGmmRes.data().template ReinterpretCast(), - ubGmmRes.layout(), - Arch::PositionUB{}); - - AscendC::WaitFlag(eventUbScaleVMTE2List[ubListId]); - copyGmToUbScale(ubScale, gmScale); - AscendC::SetFlag(eventUbScaleMTE2VList[ubListId]); - - AscendC::WaitFlag(eventUbPerTokenVMTE2List[ubListId]); - copyGmToUbPerToken(ubPerToken, gmPerToken); - AscendC::SetFlag(eventUbPerTokenMTE2VList[ubListId]); - - AscendC::WaitFlag(eventUbScaleMTE2VList[ubListId]); - AscendC::WaitFlag(eventUbPerTokenMTE2VList[ubListId]); - AscendC::WaitFlag(eventUbDMTE3VList[ubListId]); - tilePerTokenDequant(ubDequantOut, ubGmmRes, ubScale, ubPerToken); - AscendC::SetFlag(eventUbDVMTE3List[ubListId]); - AscendC::SetFlag(eventUbScaleVMTE2List[ubListId]); - AscendC::SetFlag(eventUbPerTokenVMTE2List[ubListId]); - - AscendC::WaitFlag(eventUbDVMTE3List[ubListId]); - copyUbToGmDequant(gmDequantOut, ubDequantOut); - AscendC::SetFlag(eventUbDMTE3VList[ubListId]); - } - -private: - AscendC::LocalTensor ubScaleList[UB_STAGES]; - AscendC::LocalTensor ubPerTokenList[UB_STAGES]; - - int32_t eventUbScaleVMTE2List[UB_STAGES]; - int32_t eventUbScaleMTE2VList[UB_STAGES]; - int32_t eventUbPerTokenVMTE2List[UB_STAGES]; - int32_t eventUbPerTokenMTE2VList[UB_STAGES]; - int32_t eventUbDMTE3VList[UB_STAGES]; - int32_t eventUbDVMTE3List[UB_STAGES]; -}; - -} // namespace Catlass::Epilogue::Block - -#endif // CATLASS_EPILOGUE_BLOCK_EPILOGUE_PER_TOKEN_DEQUANT_TLA_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_rescale_o_no_split_row.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_rescale_o_no_split_row.hpp deleted file mode 100644 index 657ace9526..0000000000 --- a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_rescale_o_no_split_row.hpp +++ /dev/null @@ -1,313 +0,0 @@ -/** - * Copyright (c) 2025 Huawei Technologies Co., Ltd. - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of - * CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ - -#ifndef CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_RESCALE_O_NO_SPLIT_ROW_HPP -#define CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_RESCALE_O_NO_SPLIT_ROW_HPP - -#include "catlass/kda_catlass.hpp" -#include "catlass/arch/kda_resource.hpp" -#include "catlass/epilogue/kda_dispatch_policy.hpp" -#include "catlass/epilogue/tile/kda_epilogue_tile_copy.hpp" -#include "catlass/kda_gemm_coord.hpp" -#include "catlass/kda_matrix_coord.hpp" - -namespace Catlass::Epilogue::Block { - -template < - class OutputType_, - class InputType_, - class UpdateType_> -class BlockEpilogue< - EpilogueAtlasA2RescaleO, - OutputType_, - InputType_, - UpdateType_> { -public: - // Type aliases - using DispatchPolicy = EpilogueAtlasA2RescaleO; - using ArchTag = typename DispatchPolicy::ArchTag; - - using ElementOutput = typename OutputType_::Element; - using ElementInput = typename InputType_::Element; - using ElementUpdate = typename UpdateType_::Element; - - using LayoutOutput = typename OutputType_::Layout; - using LayoutInput = typename InputType_::Layout; - using LayoutUpdate = typename UpdateType_::Layout; - - static constexpr uint32_t HALF_ELENUM_PER_BLK = 16; - static constexpr uint32_t BLOCK_SIZE = 16; - static constexpr uint32_t HALF_ELENUM_PER_VECCALC = 128; - static constexpr uint32_t FLOAT_ELENUM_PER_VECCALC = 64; - static constexpr uint32_t HALF_ELENUM_PER_LINE = 256; - static constexpr uint32_t FLOAT_ELENUM_PER_LINE = 128; - static constexpr uint32_t MULTIPLIER = 2; - static constexpr uint32_t FLOAT_BLOCK_SIZE = 8; - static constexpr uint32_t FLOAT_VECTOR_SIZE = 64; - static constexpr uint32_t UB_UINT8_VECTOR_SIZE = 1024; - static constexpr uint32_t UB_UINT8_BLOCK_SIZE = 16384; - static constexpr uint32_t HALF_DM_UB_SIZE = 64; - static constexpr uint32_t HALF_LL_UB_SIZE = 256; - static constexpr uint32_t VECTOR_SIZE = 128; - static constexpr uint32_t NUM4 = 4; - static constexpr uint32_t MAX_UB_O_ELEM_NUM = 4096; - static constexpr uint32_t MAX_ROW_NUM_SUB_CORE = 128; - - CATLASS_DEVICE - BlockEpilogue(Arch::Resource &resource) - { - // Allocate UB space - constexpr uint32_t LO_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE; - constexpr uint32_t GO_UB_TENSOR_OFFSET = 8 * UB_UINT8_BLOCK_SIZE; - constexpr uint32_t TV_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE; - - constexpr uint32_t HM_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 9 * UB_UINT8_VECTOR_SIZE; - constexpr uint32_t GL_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 12 * UB_UINT8_VECTOR_SIZE; - constexpr uint32_t DM_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 13 * UB_UINT8_VECTOR_SIZE; - - loUbTensor = resource.ubBuf.template GetBufferByByte(LO_UB_TENSOR_OFFSET); - dmUbTensor = resource.ubBuf.template GetBufferByByte(DM_UB_TENSOR_OFFSET); - glUbTensor = resource.ubBuf.template GetBufferByByte(GL_UB_TENSOR_OFFSET); - tvUbTensor = resource.ubBuf.template GetBufferByByte(TV_UB_TENSOR_OFFSET); - goUbTensor16 = resource.ubBuf.template GetBufferByByte(GO_UB_TENSOR_OFFSET); - goUbTensor32 = resource.ubBuf.template GetBufferByByte(GO_UB_TENSOR_OFFSET); - hmUbTensor = resource.ubBuf.template GetBufferByByte(HM_UB_TENSOR_OFFSET); - } - - CATLASS_DEVICE - ~BlockEpilogue() - { - } - - CATLASS_DEVICE - void SetMask(int32_t len) - { - uint64_t mask = 0; - uint64_t one = 1; - uint64_t temp = len % FLOAT_VECTOR_SIZE; - for (int64_t i = 0; i < temp; i++) { - mask |= one << i; - } - - if (len == VECTOR_SIZE) { - AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); - } else if (len >= FLOAT_VECTOR_SIZE) { - AscendC::SetVectorMask(mask, (uint64_t)-1); - } else { - AscendC::SetVectorMask(0x0, mask); - } - } - - CATLASS_DEVICE - void CopyOToGm( - AscendC::GlobalTensor gOutput, - uint32_t curRowNum, uint32_t qSBlockSize, uint32_t embed, - uint32_t embedRound, uint32_t qNThisSubBlock, uint32_t oHiddenSize) - { - if (qNThisSubBlock == 0) { - AscendC::DataCopyPad( - gOutput, - goUbTensor16, - AscendC::DataCopyExtParams(curRowNum, embed * 2, 0, (oHiddenSize - embed) * 2, 0)); - } else { - for (uint32_t qNIdx = 0; qNIdx < qNThisSubBlock; qNIdx++) { - AscendC::DataCopyPad( - gOutput[qNIdx * embed], - goUbTensor16[qNIdx * embedRound * qSBlockSize], - AscendC::DataCopyExtParams(qSBlockSize, embed * 2, 0, (oHiddenSize - embed) * 2, 0)); - } - } - } - - CATLASS_DEVICE - void SubCoreCompute( - AscendC::GlobalTensor gOutput, - AscendC::GlobalTensor gInput, - const LayoutOutput &layoutOutput, - const LayoutInput &layoutInput, - uint32_t qNThisSubBlock, - uint32_t isFirstStackTile, uint32_t isLastStackTile, uint32_t curStackTileMod) - { - uint32_t curRowNum = layoutInput.shape(0); - uint32_t embed = layoutInput.shape(1); - uint32_t embedRound = layoutInput.stride(0); - uint32_t curRowNumRound = RoundUp(curRowNum, FLOAT_BLOCK_SIZE); - uint32_t qSBlockSize = layoutOutput.shape(0); - uint32_t oHiddenSize = layoutOutput.shape(1); - uint32_t dmUbOffsetCurStackTile = - curStackTileMod * MAX_ROW_NUM_SUB_CORE; - - AscendC::WaitFlag(EVENT_ID3); - if (!isFirstStackTile) { - AscendC::DataCopy(loUbTensor, gInput, - AscendC::DataCopyParams(1, curRowNum * embedRound / FLOAT_BLOCK_SIZE, 0, 0)); - AscendC::SetFlag(EVENT_ID0); - AscendC::WaitFlag(EVENT_ID0); - - AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); - AscendC::Brcb(tvUbTensor.ReinterpretCast(), - dmUbTensor[dmUbOffsetCurStackTile].ReinterpretCast(), curRowNumRound / FLOAT_BLOCK_SIZE, - AscendC::BrcbRepeatParams(1, 8)); - AscendC::PipeBarrier(); - // *** go = go * dm_block - AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); - for (uint32_t vmul_idx = 0; vmul_idx < embed / FLOAT_VECTOR_SIZE; ++vmul_idx) { - AscendC::Mul(goUbTensor32[vmul_idx * FLOAT_VECTOR_SIZE], - goUbTensor32[vmul_idx * FLOAT_VECTOR_SIZE], - tvUbTensor, (uint64_t)0, - curRowNum, - AscendC::BinaryRepeatParams(1, 1, 0, embedRound / FLOAT_BLOCK_SIZE, - embedRound / FLOAT_BLOCK_SIZE, 1)); - } - if (embed % FLOAT_VECTOR_SIZE > 0) { - SetMask(embed % FLOAT_VECTOR_SIZE); - AscendC::Mul(goUbTensor32[embed / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - goUbTensor32[embed / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - tvUbTensor, - (uint64_t)0, curRowNum, - AscendC::BinaryRepeatParams(1, 1, 0, embedRound / FLOAT_BLOCK_SIZE, - embedRound / FLOAT_BLOCK_SIZE, 1)); - AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); - } - AscendC::PipeBarrier(); - // *** go = lo + go - AscendC::Add(goUbTensor32, goUbTensor32, - loUbTensor, (uint64_t)0, - (curRowNum * embedRound + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, - AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); - AscendC::PipeBarrier(); - } else { - // *** go = lo - AscendC::DataCopy(goUbTensor32, gInput, - AscendC::DataCopyParams(1, curRowNum * embedRound / FLOAT_BLOCK_SIZE, 0, 0)); - AscendC::SetFlag(EVENT_ID0); - AscendC::WaitFlag(EVENT_ID0); - } - AscendC::SetFlag(EVENT_ID3); - - if (isLastStackTile) { - // *** gl_block = expand_to_block(gl), 存放于 tv - AscendC::Brcb(tvUbTensor.ReinterpretCast(), - glUbTensor.ReinterpretCast(), - curRowNumRound / FLOAT_BLOCK_SIZE, - AscendC::BrcbRepeatParams(1, 8)); - AscendC::PipeBarrier(); - // *** go = go / gl_block - AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); - for (uint32_t vdiv_idx = 0; vdiv_idx < embed / FLOAT_VECTOR_SIZE; ++vdiv_idx) { - AscendC::Div(goUbTensor32[vdiv_idx * FLOAT_VECTOR_SIZE], - goUbTensor32[vdiv_idx * FLOAT_VECTOR_SIZE], - tvUbTensor, (uint64_t)0, - curRowNum, - AscendC::BinaryRepeatParams(1, 1, 0, embedRound / FLOAT_BLOCK_SIZE, - embedRound / FLOAT_BLOCK_SIZE, 1)); - } - if (embed % FLOAT_VECTOR_SIZE > 0) { - SetMask(embed % FLOAT_VECTOR_SIZE); - AscendC::Div(goUbTensor32[embed / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - goUbTensor32[embed / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - tvUbTensor, - (uint64_t)0, curRowNum, - AscendC::BinaryRepeatParams(1, 1, 0, embedRound / FLOAT_BLOCK_SIZE, - embedRound / FLOAT_BLOCK_SIZE, 1)); - AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); - } - AscendC::PipeBarrier(); - - // *** go = castfp32to16(go) - if (std::is_same::value) { - AscendC::Cast( - goUbTensor16, goUbTensor32, - AscendC::RoundMode::CAST_RINT, (uint64_t)0, - (curRowNum * embedRound + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, - AscendC::UnaryRepeatParams(1, 1, 4, 8)); - } else { - AscendC::Cast( - goUbTensor16, goUbTensor32, - AscendC::RoundMode::CAST_NONE, (uint64_t)0, - (curRowNum * embedRound + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, - AscendC::UnaryRepeatParams(1, 1, 4, 8)); - } - AscendC::SetFlag(EVENT_ID0); - AscendC::WaitFlag(EVENT_ID0); - - // ***move O to GM - CopyOToGm( - gOutput, curRowNum, qSBlockSize, embed, - embedRound, qNThisSubBlock, oHiddenSize); - } - } - - CATLASS_DEVICE - void operator()( - AscendC::GlobalTensor gOutput, - AscendC::GlobalTensor gInput, - const LayoutOutput &layoutOutput, - const LayoutInput &layoutInput, - GemmCoord actualBlockShape, - uint32_t qSBlockSize, uint32_t qNBlockSize, - uint32_t isFirstStackTile, uint32_t isLastStackTile, uint32_t curStackTileMod) - { - uint32_t rowNum = actualBlockShape.m(); - uint32_t embed = actualBlockShape.n(); - - uint32_t subBlockIdx = AscendC::GetSubBlockIdx(); - uint32_t subBlockNum = AscendC::GetSubBlockNum(); - - uint32_t qNSplitSubBlock = qNBlockSize / subBlockNum; - uint32_t qNThisSubBlock = (qNBlockSize == 1) ? - 0 : - (subBlockIdx == 1) ? - (qNBlockSize - qNSplitSubBlock) : - qNSplitSubBlock; - uint32_t inRowSplitSubBlock = (qNBlockSize == 1) ? - (qSBlockSize / subBlockNum) : - (qSBlockSize * qNSplitSubBlock); - uint32_t inRowActualThisSubBlock = (subBlockIdx == 1) ? - (rowNum - inRowSplitSubBlock) : - inRowSplitSubBlock; - uint32_t inRowOffsetThisSubBlock = subBlockIdx * inRowSplitSubBlock; - uint32_t outRowOffsetThisSubBlock = (qNBlockSize == 1) ? - inRowOffsetThisSubBlock : - 0; - uint32_t outColOffsetThisSubBlock = (qNBlockSize == 1) ? - 0 : - subBlockIdx * qNSplitSubBlock * embed; - - if (inRowActualThisSubBlock > 0) { - int64_t offsetOutput = - layoutOutput.GetOffset(MatrixCoord(outRowOffsetThisSubBlock, outColOffsetThisSubBlock)); - auto gOutputThisSubBlock = gOutput[offsetOutput]; - auto layoutOutputThisSubBlock = layoutOutput; - - int64_t offsetInput = layoutInput.GetOffset(MatrixCoord(inRowOffsetThisSubBlock, 0)); - auto gInputThisSubBlock = gInput[offsetInput]; - auto layoutInputThisSubBlock = layoutInput.GetTileLayout(MatrixCoord(inRowActualThisSubBlock, embed)); - SubCoreCompute( - gOutputThisSubBlock, - gInputThisSubBlock, - layoutOutputThisSubBlock, - layoutInputThisSubBlock, - qNThisSubBlock, isFirstStackTile, isLastStackTile, curStackTileMod); - } - } - -private: - AscendC::LocalTensor loUbTensor; - AscendC::LocalTensor dmUbTensor; - AscendC::LocalTensor hmUbTensor; - AscendC::LocalTensor glUbTensor; - AscendC::LocalTensor tvUbTensor; - AscendC::LocalTensor goUbTensor16; - AscendC::LocalTensor goUbTensor32; -}; -} // namespace Catlass::Epilogue::Block - -#endif // CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_RESCALE_O_NO_SPLIT_ROW_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_w4a4_per_token_per_channel_dequant.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_w4a4_per_token_per_channel_dequant.hpp deleted file mode 100644 index 2d2e768c55..0000000000 --- a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/block_epilogue_w4a4_per_token_per_channel_dequant.hpp +++ /dev/null @@ -1,252 +0,0 @@ -/** - * Copyright (c) 2025 Huawei Technologies Co., Ltd. - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of - * CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ -#ifndef CATLASS_EPILOGUE_BLOCK_EPILOGUE_W4A4_HPP -#define CATLASS_EPILOGUE_BLOCK_EPILOGUE_W4A4_HPP - -#include "catlass/arch/kda_resource.hpp" -#include "catlass/kda_catlass.hpp" -#include "catlass/detail/kda_callback.hpp" -#include "catlass/epilogue/kda_dispatch_policy.hpp" -#include "catlass/gemm/kda_gemm_type.hpp" -#include "catlass/kda_gemm_coord.hpp" -#include "catlass/layout/kda_layout.hpp" -#include "catlass/kda_matrix_coord.hpp" - -namespace Catlass::Epilogue::Block { - -template < - class CType_, - class PerTokenScaleType_, - class DType_, - class TileBroadcastOneBlk_, - class TileOneBlkColumnBroadcastMul_, - class TileCopy_, - class EpilogueTileSwizzle_> -class BlockEpilogue< - EpilogueAtlasA2W4A4PerTokenPerChannelDequant, - CType_, - PerTokenScaleType_, - DType_, - TileBroadcastOneBlk_, - TileOneBlkColumnBroadcastMul_, - TileCopy_, - EpilogueTileSwizzle_> { -public: - using DispatchPolicy = EpilogueAtlasA2W4A4PerTokenPerChannelDequant; - using ArchTag = typename DispatchPolicy::ArchTag; - - // Data infos - using ElementC = typename CType_::Element; - using LayoutC = typename CType_::Layout; - using ElementPerTokenScale = typename PerTokenScaleType_::Element; - using LayoutPerTokenScale = typename PerTokenScaleType_::Layout; - using ElementD = typename DType_::Element; - using LayoutD = typename DType_::Layout; - - // Tile compute ops - using TileBroadcastOneBlk = TileBroadcastOneBlk_; - using TileOneBlkColumnBroadcastMul = TileOneBlkColumnBroadcastMul_; - - // Tile copy - using CopyGmToUbC = typename TileCopy_::CopyGmToUbC; - using CopyGmToUbPerTokenScale = typename TileCopy_::CopyGmToUbPerTokenScale; - using CopyUbToGmD = typename TileCopy_::CopyUbToGmD; - - using EpilogueTileSwizzle = EpilogueTileSwizzle_; - using TileShape = typename TileOneBlkColumnBroadcastMul::TileShape; - - // Check - static_assert( - std::is_same_v && std::is_same_v && std::is_same_v, - "Accumulator type (`ElementC`) should be half, output type (`ElementD`) should be bfloat16_t while " - "perTokenScale type should be float"); - static_assert( - std::is_same_v && std::is_same_v && std::is_same_v, - "The layout template parameters of BlockEpilogue are wrong"); - - static_assert( - TileShape::ROW == TileBroadcastOneBlk::COMPUTE_LENGTH, - "TileShape must be consistent for all tile compute ops"); - - static_assert(TileShape::ROW % 16 == 0, "Epilogue TileShape::ROW must be divisible by 16"); - - static_assert( - /*ubC*/ TileShape::COUNT * sizeof(ElementC) + /*ubPerTokenScale*/ TileShape::ROW * sizeof(ElementPerTokenScale) + /*ubD*/ TileShape::COUNT * sizeof(ElementD) + /*ubCFp32 & ubPerTokenScaleBrcb*/ 2 * TileShape::COUNT * sizeof(float) <= ArchTag::UB_SIZE, - "TileShape is too large to fit in UB"); - - struct Params { - __gm__ ElementPerTokenScale *ptrPerTokenScale{nullptr}; - LayoutPerTokenScale layoutPerTokenScale{}; - __gm__ ElementD *ptrD{nullptr}; - LayoutD layoutD{}; - - CATLASS_DEVICE - Params() {}; - - CATLASS_DEVICE - Params( - __gm__ ElementPerTokenScale *ptrPerTokenScale_, - LayoutPerTokenScale const &layoutPerTokenScale_, - __gm__ ElementD *ptrD_, - LayoutD const &layoutD_) - : ptrPerTokenScale(ptrPerTokenScale_), - layoutPerTokenScale(layoutPerTokenScale_), - ptrD(ptrD_), - layoutD(layoutD_) - { - } - }; - - CATLASS_DEVICE - BlockEpilogue(Arch::Resource const &resource, Params const ¶ms = Params{}) - : params(params) - { - size_t ubCOffset = 0; - size_t ubPerTokenScaleOffset = ubCOffset + TileShape::COUNT * sizeof(ElementC); - size_t ubDOffset = ubPerTokenScaleOffset + TileShape::ROW * sizeof(ElementPerTokenScale); - size_t ubOffset1 = ubDOffset + TileShape::COUNT * sizeof(ElementD); // for `ubCFp32` & `ubDFp32` - size_t ubOffset2 = ubOffset1 + TileShape::COUNT * sizeof(float); // for `ubPerTokenScaleBrcb` - // for `ubPerTokenScaleBrcb` - - // Load from GM - ubC = resource.ubBuf.template GetBufferByByte(ubCOffset); - ubPerTokenScale = resource.ubBuf.template GetBufferByByte(ubPerTokenScaleOffset); - ubD = resource.ubBuf.template GetBufferByByte(ubDOffset); - - // Caculation (with perToken scale) - ubCFp32 = resource.ubBuf.template GetBufferByByte(ubOffset1); - ubPerTokenScaleBrcb = resource.ubBuf.template GetBufferByByte(ubOffset2); - ubDFp32 = resource.ubBuf.template GetBufferByByte(ubOffset1); - - AscendC::SetFlag(eventID); - } - - CATLASS_DEVICE - ~BlockEpilogue() - { - AscendC::WaitFlag(eventID); - } - - CATLASS_DEVICE - void UpdateParams(Params const ¶ms_) - { - params = params_; - } - - CATLASS_DEVICE - void operator()( - GemmCoord const &blockShapeMNK, - GemmCoord const &blockCoordMNK, - GemmCoord const &actualBlockShapeMNK, - AscendC::GlobalTensor const &gmBlockC, - LayoutC const &layoutBlockC, - Callback &&callback = Callback{}) - { - if (actualBlockShapeMNK.k() == 0) { - return; - } - callback(); - - // Calculate the offset of the current block - MatrixCoord blockShape = blockShapeMNK.GetCoordMN(); - MatrixCoord blockCoord = blockCoordMNK.GetCoordMN(); - MatrixCoord actualBlockShape = actualBlockShapeMNK.GetCoordMN(); - MatrixCoord blockOffset = blockCoord * blockShape; - - AscendC::GlobalTensor gmPerTokenScale; - gmPerTokenScale.SetGlobalBuffer(params.ptrPerTokenScale); - AscendC::GlobalTensor gmD; - gmD.SetGlobalBuffer(params.ptrD); - - auto ubTileStride = MakeCoord(static_cast(TileShape::COLUMN), 1L); - auto tileShape = TileShape::ToCoord(); - EpilogueTileSwizzle epilogueTileSwizzle(actualBlockShape, tileShape); - uint32_t tileLoops = epilogueTileSwizzle.GetLoops(); - uint32_t subblockIdx = AscendC::GetSubBlockIdx(); - uint32_t subblockNum = AscendC::GetSubBlockNum(); - for (uint32_t loopIdx = subblockIdx; loopIdx < tileLoops; loopIdx += subblockNum) { - auto tileCoord = epilogueTileSwizzle.GetTileCoord(loopIdx); - auto actualTileShape = epilogueTileSwizzle.GetActualTileShape(tileCoord); - auto tileOffsetInBlock = tileCoord * tileShape; - auto tileOffset = blockOffset + tileOffsetInBlock; - - // for GMc - auto gmTileC = gmBlockC[layoutBlockC.GetOffset(tileOffsetInBlock)]; - auto layoutGmTileC = layoutBlockC.GetTileLayout(actualTileShape); - - LayoutC layoutUbC{actualTileShape, ubTileStride}; - - auto calcNum = actualTileShape.row() * TileShape::COLUMN; - - AscendC::WaitFlag(eventID); - copyGmToUbC(ubC, gmTileC, layoutUbC, layoutGmTileC); - - auto perTokenScaleTileOffset = tileOffset.template GetCoordByAxis<0>(); - auto perTokenScaleTileShape = actualTileShape.template GetCoordByAxis<0>(); - - auto gmTilePerTokenScale = gmPerTokenScale[params.layoutPerTokenScale.GetOffset(perTokenScaleTileOffset)]; - auto layoutGmTilePerTokenScale = params.layoutPerTokenScale.GetTileLayout(perTokenScaleTileShape); - - auto layoutUbPerTokenScale = LayoutPerTokenScale::template MakeLayoutInUb( - perTokenScaleTileShape); - - copyGmToUbPerTokenScale( - ubPerTokenScale, gmTilePerTokenScale, layoutUbPerTokenScale, layoutGmTilePerTokenScale); - AscendC::SetFlag(eventID); - - AscendC::WaitFlag(eventID); - AscendC::Cast(/*float*/ ubCFp32, /*half*/ ubC, AscendC::RoundMode::CAST_NONE, calcNum); - AscendC::PipeBarrier(); - - tileBroadcastOneBlk(ubPerTokenScaleBrcb, ubPerTokenScale); - AscendC::PipeBarrier(); - - tileOneBlkColumnBroadcastMul(ubDFp32, ubCFp32, ubPerTokenScaleBrcb); - AscendC::PipeBarrier(); - - AscendC::Cast(ubD, ubDFp32, AscendC::RoundMode::CAST_RINT, calcNum); - AscendC::SetFlag(eventID); - - auto resTileShape = MatrixCoord{actualTileShape.row(), actualTileShape.column()}; - LayoutD layoutUbD{resTileShape, ubTileStride}; - - auto gmTileD = gmD[params.layoutD.GetOffset(tileOffset)]; - auto layoutGmTileD = params.layoutD.GetTileLayout(resTileShape); - - AscendC::WaitFlag(eventID); - copyUbToGmD(gmTileD, ubD, layoutGmTileD, layoutUbD); - AscendC::SetFlag(eventID); - } - } - -private: - Params params; - - AscendC::LocalTensor ubC; - AscendC::LocalTensor ubPerTokenScale; - AscendC::LocalTensor ubD; - AscendC::LocalTensor ubCFp32; - AscendC::LocalTensor ubPerTokenScaleBrcb; - AscendC::LocalTensor ubDFp32; - // [gmTileC] -> ubC --(Cast:float)--> ubCFp32 --(Muls)--> ubDFp32 --(Cast:bf16)--> ubD -> [gmTileD] - - int32_t eventID{0}; - - TileBroadcastOneBlk tileBroadcastOneBlk; - TileOneBlkColumnBroadcastMul tileOneBlkColumnBroadcastMul; - - CopyGmToUbC copyGmToUbC; - CopyGmToUbPerTokenScale copyGmToUbPerTokenScale; - CopyUbToGmD copyUbToGmD; -}; - -} // namespace Catlass::Epilogue::Block - -#endif // CATLASS_EPILOGUE_BLOCK_EPILOGUE_W4A4_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/kda_block_epilogue.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/kda_block_epilogue.hpp index 47185bcf5c..cc072c2173 100644 --- a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/kda_block_epilogue.hpp +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/kda_block_epilogue.hpp @@ -24,28 +24,9 @@ class BlockEpilogue { } // namespace Catlass::Epilogue::Block -#include "catlass/epilogue/block/block_epilogue_elemwise_no_source.hpp" -#include "catlass/epilogue/block/kda_block_epilogue_elemwise_one_source.hpp" -#include "catlass/epilogue/block/kda_block_epilogue_fa_softmax.hpp" -#include "catlass/epilogue/block/block_epilogue_fa_rescale_o.hpp" -#include "catlass/epilogue/block/kda_block_epilogue_mla_softmax.hpp" -#include "catlass/epilogue/block/block_epilogue_mla_rescale_o.hpp" -#include "catlass/epilogue/block/block_epilogue_mla_fd_rescale_o.hpp" -#include "catlass/epilogue/block/kda_block_epilogue_per_token_dequant.hpp" -#include "catlass/epilogue/block/block_epilogue_per_token_dequant_tla.hpp" -#include "catlass/epilogue/block/kda_block_epilogue_gemm.hpp" -#include "catlass/epilogue/block/kda_block_epilogue_gemv.hpp" -#include "catlass/epilogue/block/kda_block_epilogue_mla_tp1_softmax.hpp" -#include "catlass/epilogue/block/block_epilogue_mla_tp1_rescale_o.hpp" -#include "catlass/epilogue/block/block_epilogue_online_softmax_no_mask.hpp" -#include "catlass/epilogue/block/block_epilogue_rescale_o_no_split_row.hpp" -#include "catlass/epilogue/block/block_epilogue_w4a4_per_token_per_channel_dequant.hpp" - -#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 3510) -#include "catlass/epilogue/block/block_epilogue_fa_softmax_ascend950.hpp" -#include "catlass/epilogue/block/block_epilogue_fa_rescale_o_ascend950.hpp" -#include "catlass/epilogue/block/block_epilogue_fixpipe.hpp" -#include "catlass/epilogue/block/block_epilogue_per_group_per_block.hpp" -#endif +// chunk_kda_fwd only instantiates BlockEpilogue for its own FwdH policies +// (EpilogueAtlasKdaFwdHVnew / EpilogueAtlasKdaFwdHUpdate), whose specializations +// live under arch22|arch35/fwd_h/epilogue/block/ and are included by their users. +// None of the upstream catlass epilogues were ever selected here. #endif // CATLASS_KDA_BLOCK_EPILOGUE_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/kda_block_epilogue_elemwise_one_source.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/kda_block_epilogue_elemwise_one_source.hpp deleted file mode 100644 index 1f2b6e4283..0000000000 --- a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/kda_block_epilogue_elemwise_one_source.hpp +++ /dev/null @@ -1,185 +0,0 @@ -/** - * Copyright (c) 2025 Huawei Technologies Co., Ltd. - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of - * CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ - -#ifndef CATLASS_KDA_BLOCK_EPILOGUE_ELEMWISE_ONE_SOURCE_HPP -#define CATLASS_KDA_BLOCK_EPILOGUE_ELEMWISE_ONE_SOURCE_HPP - -#include "catlass/kda_catlass.hpp" -#include "catlass/arch/kda_resource.hpp" -#include "catlass/epilogue/kda_dispatch_policy.hpp" -#include "catlass/kda_gemm_coord.hpp" -#include "catlass/kda_matrix_coord.hpp" -#include "catlass/layout/kda_layout.hpp" - -namespace Catlass::Epilogue::Block { - -template < - class CType_, - class XType_, - class DType_, - class TileElemWiseEpilogue_, - class TileCopy_> -class BlockEpilogue< - EpilogueAtlasA2ElemWiseOneSource, - CType_, - XType_, - DType_, - TileElemWiseEpilogue_, - TileCopy_> { -public: - // Type aliases - using DispatchPolicy = EpilogueAtlasA2ElemWiseOneSource; - using ArchTag = typename DispatchPolicy::ArchTag; - using ElementC = typename CType_::Element; - using LayoutC = typename CType_::Layout; - using ElementX = typename XType_::Element; - using LayoutX = typename XType_::Layout; - using ElementD = typename DType_::Element; - using LayoutD = typename DType_::Layout; - using TileElemWiseEpilogue = TileElemWiseEpilogue_; - using CopyGmToUbC = typename TileCopy_::CopyGmToUbC; - using CopyGmToUbX = typename TileCopy_::CopyGmToUbX; - using CopyUbToGmD = typename TileCopy_::CopyUbToGmD; - - static constexpr uint32_t COMPUTE_LENGTH = TileElemWiseEpilogue::COMPUTE_LENGTH; - static constexpr uint32_t OPERANDS_NUM = DispatchPolicy::OPERANDS_NUM; - - // Check the element type of C, X and D - static_assert(std::is_same_v && std::is_same_v, - "Element type of C, X and D must be the same"); - using ElementCompute = ElementD; - - // Check the layout type of C, X and D - static_assert(std::is_same_v && std::is_same_v && - std::is_same_v, - "Layout type of C, X and D must be RowMajor"); - using LayoutComputeInUb = layout::RowMajor; - - // Check if ArchTag is matched - static_assert(std::is_same_v, "Tile epilogue's ArchTag mismatch"); - // Check if compute length is valid - static_assert(COMPUTE_LENGTH * OPERANDS_NUM * sizeof(ElementCompute) <= ArchTag::UB_SIZE, "UB out of bounds"); - - // Epilogue params definition - struct Params { - GM_ADDR ptrX; - LayoutX layoutX; - GM_ADDR ptrD; - LayoutD layoutD; - - CATLASS_HOST_DEVICE - Params() {} - - CATLASS_HOST_DEVICE - Params(GM_ADDR ptrX_, LayoutX const &layoutX_, GM_ADDR ptrD_, LayoutD const &layoutD_) - : ptrX(ptrX_), - layoutX(layoutX_), - ptrD(ptrD_), - layoutD(layoutD_) - {} - }; - - CATLASS_DEVICE - BlockEpilogue(Arch::Resource &resource, Params const ¶ms) - : params(params) - { - ubC = resource.ubBuf.template GetBufferByByte(0); - ubX = resource.ubBuf.template GetBufferByByte(COMPUTE_LENGTH * sizeof(ElementC)); - ubD = resource.ubBuf.template GetBufferByByte( - COMPUTE_LENGTH * sizeof(ElementC) + COMPUTE_LENGTH * sizeof(ElementX)); - - AscendC::SetFlag(EVENT_ID0); - AscendC::SetFlag(EVENT_ID0); - } - - CATLASS_DEVICE - ~BlockEpilogue() - { - AscendC::WaitFlag(EVENT_ID0); - AscendC::WaitFlag(EVENT_ID0); - } - - CATLASS_DEVICE - void operator()( - GemmCoord const &blockShapeMNK, - GemmCoord const &blockCoordMNK, - GemmCoord const &actualBlockShapeMNK, - AscendC::GlobalTensor const &gmBlockC, - LayoutX const &layoutBlockC) - { - // Calculate the offset of the current block - MatrixCoord blockShape = blockShapeMNK.GetCoordMN(); - MatrixCoord blockCoord = blockCoordMNK.GetCoordMN(); - MatrixCoord actualBlockShape = actualBlockShapeMNK.GetCoordMN(); - MatrixCoord blockOffset = blockCoord * blockShape; - - // Calculate the offset and the shape of the current subblock - MatrixCoord subblockShape{ - CeilDiv(actualBlockShape.row(), static_cast(AscendC::GetSubBlockNum())), - actualBlockShape.column()}; - MatrixCoord subblockCoord{AscendC::GetSubBlockIdx(), 0}; - MatrixCoord actualSubblockShape = MatrixCoord::Min( - subblockShape, actualBlockShape - subblockCoord * subblockShape); - MatrixCoord subblockOffset = subblockCoord * subblockShape; - - // Get the data and layout of C - auto gmSubblockC = gmBlockC[layoutBlockC.GetOffset(subblockOffset)]; - auto layoutSubblockC = layoutBlockC.GetTileLayout(actualSubblockShape); - - // Get the data and layout of X - AscendC::GlobalTensor gmX; - gmX.SetGlobalBuffer(reinterpret_cast<__gm__ ElementX *>(params.ptrX)); - auto gmSubblockX = gmX[params.layoutX.GetOffset(blockOffset + subblockOffset)]; - auto layoutSubblockX = params.layoutX.GetTileLayout(actualSubblockShape); - - // Get the data and layout of D - AscendC::GlobalTensor gmD; - gmD.SetGlobalBuffer(reinterpret_cast<__gm__ ElementD *>(params.ptrD)); - auto gmSubblockD = gmD[params.layoutD.GetOffset(blockOffset + subblockOffset)]; - auto layoutSubblockD = params.layoutD.GetTileLayout(actualSubblockShape); - - // Get the layout on UB - auto layoutComputeInUb = LayoutComputeInUb::template MakeLayoutInUb(actualSubblockShape); - - // Copy the data of C and X - AscendC::WaitFlag(EVENT_ID0); - copyGmToUbC(ubC, gmSubblockC, layoutComputeInUb, layoutSubblockC); - copyGmToUbX(ubX, gmSubblockX, layoutComputeInUb, layoutSubblockX); - AscendC::SetFlag(EVENT_ID0); - - // Perform epilogue calculation - AscendC::WaitFlag(EVENT_ID0); - AscendC::WaitFlag(EVENT_ID0); - tileEpilogue(ubD, ubC, ubX); - AscendC::SetFlag(EVENT_ID0); - AscendC::SetFlag(EVENT_ID0); - - // Copy the data of D - AscendC::WaitFlag(EVENT_ID0); - copyUbToGmD(gmSubblockD, ubD, layoutSubblockD, layoutComputeInUb); - AscendC::SetFlag(EVENT_ID0); - } - -private: - Params params; - - AscendC::LocalTensor ubC; - AscendC::LocalTensor ubX; - AscendC::LocalTensor ubD; - - TileElemWiseEpilogue tileEpilogue; - CopyGmToUbC copyGmToUbC; - CopyGmToUbX copyGmToUbX; - CopyUbToGmD copyUbToGmD; -}; - -} // namespace Catlass::Epilogue::Block - -#endif // CATLASS_KDA_BLOCK_EPILOGUE_ELEMWISE_ONE_SOURCE_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/kda_block_epilogue_fa_softmax.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/kda_block_epilogue_fa_softmax.hpp deleted file mode 100644 index 8b1ecfa270..0000000000 --- a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/kda_block_epilogue_fa_softmax.hpp +++ /dev/null @@ -1,403 +0,0 @@ -/** - * Copyright (c) 2025 Huawei Technologies Co., Ltd. - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of - * CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ - -#ifndef CATLASS_KDA_BLOCK_EPILOGUE_FA_SOFTMAX_HPP -#define CATLASS_KDA_BLOCK_EPILOGUE_FA_SOFTMAX_HPP - -#include "catlass/kda_catlass.hpp" -#include "catlass/arch/kda_cross_core_sync.hpp" -#include "catlass/arch/kda_resource.hpp" -#include "catlass/epilogue/kda_dispatch_policy.hpp" -#include "catlass/kda_gemm_coord.hpp" -#include "catlass/kda_matrix_coord.hpp" -#include "catlass/epilogue/tile/kda_epilogue_tile_copy.hpp" - -namespace Catlass::Epilogue::Block { - -template < - class OutputType_, - class InputType_, - class MaskType_> -class BlockEpilogue< - EpilogueAtlasA2FASoftmax, - OutputType_, - InputType_, - MaskType_> { -public: - // Type aliases - using DispatchPolicy = EpilogueAtlasA2FASoftmax; - using ArchTag = typename DispatchPolicy::ArchTag; - using ElementOutput = typename OutputType_::Element; - using ElementInput = typename InputType_::Element; - using ElementMask = typename MaskType_::Element; - - using LayoutOutput = typename OutputType_::Layout; - using LayoutInput = typename InputType_::Layout; - using LayoutMask = typename MaskType_::Layout; - - using CopyGmToUbInput = Tile::CopyGm2Ub; - using CopyGmToUbMask = Tile::CopyGm2Ub; - using CopyUbToGmOutput = Tile::CopyUb2Gm; - - static constexpr uint32_t HALF_ELENUM_PER_BLK = 16; - static constexpr uint32_t FLOAT_ELENUM_PER_BLK = 8; - static constexpr uint32_t HALF_ELENUM_PER_VECCALC = 128; - static constexpr uint32_t FLOAT_ELENUM_PER_VECCALC = 64; - static constexpr uint32_t UB_TILE_SIZE = 16384; // 64 * 128 * 2B - static constexpr uint32_t UB_LINE_SIZE = 512; // 128 * 2 * 2B - static constexpr uint32_t HALF_ELENUM_PER_LINE = 256; // 128 * 2 - static constexpr uint32_t FLOAT_ELENUM_PER_LINE = 128; // 128 - static constexpr uint32_t MULTIPLIER = 2; - - CATLASS_DEVICE - BlockEpilogue(Arch::Resource &resource, half tor_) - { - constexpr uint32_t LS32_UB_TENSOR_OFFSET = 2 * UB_TILE_SIZE; - constexpr uint32_t MASK_UB_TENSOR_OFFSET = 4 * UB_TILE_SIZE; - constexpr uint32_t LM_UB_TENSOR_OFFSET = 7 * UB_TILE_SIZE; - constexpr uint32_t HM_UB_TENSOR_OFFSET = 7 * UB_TILE_SIZE + 1 * UB_LINE_SIZE; - constexpr uint32_t GM_UB_TENSOR_OFFSET = 7 * UB_TILE_SIZE + 2 * UB_LINE_SIZE; - constexpr uint32_t DM_UB_TENSOR_OFFSET = 7 * UB_TILE_SIZE + 4 * UB_LINE_SIZE; - constexpr uint32_t LL_UB_TENSOR_OFFSET = 7 * UB_TILE_SIZE + 6 * UB_LINE_SIZE; - constexpr uint32_t TV_UB_TENSOR_OFFSET = 7 * UB_TILE_SIZE + 11 * UB_LINE_SIZE; - - tor = tor_; - lsUbTensor = resource.ubBuf.template GetBufferByByte(0); - lpUbTensor = resource.ubBuf.template GetBufferByByte(0); - ls32UbTensor = resource.ubBuf.template GetBufferByByte(LS32_UB_TENSOR_OFFSET); - maskUbTensor = resource.ubBuf.template GetBufferByByte(MASK_UB_TENSOR_OFFSET); - lmUbTensor = resource.ubBuf.template GetBufferByByte(LM_UB_TENSOR_OFFSET); - hmUbTensor = resource.ubBuf.template GetBufferByByte(HM_UB_TENSOR_OFFSET); - gmUbTensor = resource.ubBuf.template GetBufferByByte(GM_UB_TENSOR_OFFSET); - dmUbTensor = resource.ubBuf.template GetBufferByByte(DM_UB_TENSOR_OFFSET); - llUbTensor = resource.ubBuf.template GetBufferByByte(LL_UB_TENSOR_OFFSET); - tvUbTensor = resource.ubBuf.template GetBufferByByte(TV_UB_TENSOR_OFFSET); - - AscendC::SetFlag(EVENT_ID0); - AscendC::SetFlag(EVENT_ID0); - AscendC::SetFlag(EVENT_ID1); - } - - CATLASS_DEVICE - ~BlockEpilogue() - { - AscendC::WaitFlag(EVENT_ID0); - AscendC::WaitFlag(EVENT_ID0); - AscendC::WaitFlag(EVENT_ID1); - } - - CATLASS_DEVICE - void SetMask(int32_t len) - { - const int32_t MAX_MASK_LEN = 128; - const int32_t HALF_MASK_LEN = 64; - if (len >= MAX_MASK_LEN) { - AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); - return; - } - int32_t highMask = len - HALF_MASK_LEN > 0 ? len - HALF_MASK_LEN : 0; - int32_t lowMask = len - HALF_MASK_LEN >= 0 ? HALF_MASK_LEN : len; - if (len < HALF_MASK_LEN) { - AscendC::SetVectorMask(0x0, ((uint64_t)1 << lowMask) - 1); - } else { - AscendC::SetVectorMask(((uint64_t)1 << highMask) - 1, 0xffffffffffffffff); - } - } - - CATLASS_DEVICE - void SetVcgMask(int32_t len) - { - const int32_t MAX_LEN = 16; - if (len > MAX_LEN) { - AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); - return; - } - uint64_t subMask = ((uint64_t)1 << len) - 1; - uint64_t maskValue = (subMask << 48) + (subMask << 32) + (subMask << 16) + subMask; - AscendC::SetVectorMask(maskValue, maskValue); - } - - CATLASS_DEVICE - void subCoreCompute( - AscendC::GlobalTensor gOutput, - AscendC::GlobalTensor gInput, - AscendC::GlobalTensor gMask, - const LayoutOutput &layoutOutput, - const LayoutInput &layoutInput, - const LayoutMask &layoutMask, - uint32_t nIdx, Arch::CrossCoreFlag qkReady) - { - uint32_t subM = layoutInput.shape(0); - uint32_t qkN = layoutInput.shape(1); - uint32_t qkNRound = layoutInput.stride(0); - uint32_t maxSeqlen = layoutMask.stride(0); - uint32_t offset = pingpongFlag * UB_TILE_SIZE / sizeof(ElementInput); - uint32_t subMAligned128 = (subM + HALF_ELENUM_PER_VECCALC - 1) / HALF_ELENUM_PER_VECCALC; - uint32_t subMRound = (subM + HALF_ELENUM_PER_BLK - 1) / HALF_ELENUM_PER_BLK * HALF_ELENUM_PER_BLK; - - // Get the layout on UB - auto layoutInUb = LayoutInput::template MakeLayoutInUb(MatrixCoord{subM, qkN}); - - if (subM > 0) { - // Copy mask - AscendC::WaitFlag(EVENT_ID0); - copyGmToUbMask(maskUbTensor, gMask, layoutInUb, layoutMask); - } - Arch::CrossCoreWaitFlag(qkReady); - if (subM > 0) { - AscendC::WaitFlag(pingpongFlag); - // Copy QK - copyGmToUbInput(lsUbTensor[offset], gInput, layoutInUb, layoutInput); - AscendC::SetFlag(EVENT_ID0); - AscendC::WaitFlag(EVENT_ID0); - - // ls = tor * ls - AscendC::Muls( - lsUbTensor[offset], - lsUbTensor[offset], - tor, - (uint64_t)0, - (subM * qkNRound + HALF_ELENUM_PER_VECCALC - 1) / HALF_ELENUM_PER_VECCALC, - AscendC::UnaryRepeatParams(1, 1, 8, 8)); - AscendC::PipeBarrier(); - // ls = ls + mask - AscendC::Add( - lsUbTensor[offset], - lsUbTensor[offset], - maskUbTensor, - (uint64_t)0, - (subM * qkNRound + HALF_ELENUM_PER_VECCALC - 1) / HALF_ELENUM_PER_VECCALC, - AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); - AscendC::PipeBarrier(); - AscendC::SetFlag(EVENT_ID0); - // lm = rowmax(ls) - if (qkN <= HALF_ELENUM_PER_VECCALC) { - SetMask(qkN); - AscendC::BlockReduceMax( - tvUbTensor.ReinterpretCast(), - lsUbTensor[offset], - subM, - 0, - 2, - 1, - qkNRound / HALF_ELENUM_PER_BLK); - AscendC::PipeBarrier(); - SetVcgMask(qkNRound / HALF_ELENUM_PER_BLK); - AscendC::BlockReduceMax( - lmUbTensor, - tvUbTensor.ReinterpretCast(), - (subM * HALF_ELENUM_PER_BLK + HALF_ELENUM_PER_VECCALC - 1) / HALF_ELENUM_PER_VECCALC, - 0, - 1, - 1, - 8); - AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); - } - AscendC::PipeBarrier(); - if (nIdx == 0) { - // hm = lm - AscendC::DataCopy( - hmUbTensor, lmUbTensor, AscendC::DataCopyParams(1, subMRound / HALF_ELENUM_PER_BLK, 0, 0)); - AscendC::PipeBarrier(); - } else { - // hm = vmax(lm, gm) - AscendC::Max( - hmUbTensor, - lmUbTensor, - gmUbTensor, - (uint64_t)0, - subMAligned128, - AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); - AscendC::PipeBarrier(); - // dm = gm - hm - AscendC::Sub( - dmUbTensor[nIdx % MULTIPLIER * HALF_ELENUM_PER_LINE], - gmUbTensor, - hmUbTensor, - (uint64_t)0, - subMAligned128, - AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); - AscendC::PipeBarrier(); - } - // gm = hm - AscendC::DataCopy( - gmUbTensor, hmUbTensor, AscendC::DataCopyParams(1, subMRound / HALF_ELENUM_PER_BLK, 0, 0)); - AscendC::PipeBarrier(); - // hm_block = brcb(hm), 存放于tv - AscendC::Brcb( - tvUbTensor.ReinterpretCast(), - hmUbTensor.ReinterpretCast(), - subMRound / FLOAT_ELENUM_PER_BLK, - AscendC::BrcbRepeatParams(1, 8)); - AscendC::PipeBarrier(); - // ls = ls - hm_block - for (uint32_t vsubIdx = 0; vsubIdx < qkN / HALF_ELENUM_PER_VECCALC; vsubIdx++) { - AscendC::Sub( - lsUbTensor[offset + vsubIdx * HALF_ELENUM_PER_VECCALC], - lsUbTensor[offset + vsubIdx * HALF_ELENUM_PER_VECCALC], - tvUbTensor.ReinterpretCast(), - (uint64_t)0, - subM, - AscendC::BinaryRepeatParams( - 1, 1, 0, qkNRound / HALF_ELENUM_PER_BLK, qkNRound / HALF_ELENUM_PER_BLK, 1)); - } - if (qkN % HALF_ELENUM_PER_VECCALC > 0) { - SetMask(qkN % HALF_ELENUM_PER_VECCALC); - AscendC::Sub( - lsUbTensor[offset + qkN / HALF_ELENUM_PER_VECCALC * HALF_ELENUM_PER_VECCALC], - lsUbTensor[offset + qkN / HALF_ELENUM_PER_VECCALC * HALF_ELENUM_PER_VECCALC], - tvUbTensor.ReinterpretCast(), - (uint64_t)0, - subM, - AscendC::BinaryRepeatParams( - 1, 1, 0, qkNRound / HALF_ELENUM_PER_BLK, qkNRound / HALF_ELENUM_PER_BLK, 1)); - AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); - } - AscendC::PipeBarrier(); - // ls32 = castfp16to32(ls) - AscendC::Cast( - ls32UbTensor, - lsUbTensor[offset], - AscendC::RoundMode::CAST_NONE, - (uint64_t)0, - (subM * qkNRound + FLOAT_ELENUM_PER_VECCALC - 1) / FLOAT_ELENUM_PER_VECCALC, - AscendC::UnaryRepeatParams(1, 1, 8, 4)); - AscendC::PipeBarrier(); - // ls32 = exp(ls32) - AscendC::Exp( - ls32UbTensor, - ls32UbTensor, - (uint64_t)0, - (subM * qkNRound + FLOAT_ELENUM_PER_VECCALC - 1) / FLOAT_ELENUM_PER_VECCALC, - AscendC::UnaryRepeatParams(1, 1, 8, 8)); - AscendC::PipeBarrier(); - // lp = castfp32to16(ls) - AscendC::Cast( - lpUbTensor[offset], - ls32UbTensor, - AscendC::RoundMode::CAST_NONE, - (uint64_t)0, - (subM * qkNRound + FLOAT_ELENUM_PER_VECCALC - 1) / FLOAT_ELENUM_PER_VECCALC, - AscendC::UnaryRepeatParams(1, 1, 4, 8)); - AscendC::PipeBarrier(); - AscendC::SetFlag(EVENT_ID0); - // ll = rowsum(ls32) - if (qkN <= FLOAT_ELENUM_PER_VECCALC) { - SetMask(qkN); - AscendC::RepeatReduceSum( - llUbTensor[nIdx % MULTIPLIER * FLOAT_ELENUM_PER_LINE], - ls32UbTensor, - subM, 0, 0, 1, 1, qkNRound / FLOAT_ELENUM_PER_BLK); - AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); - } else { - for (uint32_t vaddIdx = 1; vaddIdx < qkN / FLOAT_ELENUM_PER_VECCALC; vaddIdx++) { - AscendC::Add( - ls32UbTensor, - ls32UbTensor, - ls32UbTensor[vaddIdx * FLOAT_ELENUM_PER_VECCALC], - (uint64_t)0, - subM, - AscendC::BinaryRepeatParams( - 1, 1, 1, - qkNRound / FLOAT_ELENUM_PER_BLK, - qkNRound / FLOAT_ELENUM_PER_BLK, - qkNRound / FLOAT_ELENUM_PER_BLK)); - AscendC::PipeBarrier(); - } - if (qkN % FLOAT_ELENUM_PER_VECCALC > 0) { - SetMask(qkN % FLOAT_ELENUM_PER_VECCALC); - AscendC::Add( - ls32UbTensor, - ls32UbTensor, - ls32UbTensor[qkN / FLOAT_ELENUM_PER_VECCALC * FLOAT_ELENUM_PER_VECCALC], - (uint64_t)0, - subM, - AscendC::BinaryRepeatParams( - 1, 1, 1, - qkNRound / FLOAT_ELENUM_PER_BLK, - qkNRound / FLOAT_ELENUM_PER_BLK, - qkNRound / FLOAT_ELENUM_PER_BLK)); - AscendC::PipeBarrier(); - AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); - } - AscendC::RepeatReduceSum( - llUbTensor[nIdx % MULTIPLIER * FLOAT_ELENUM_PER_LINE], - ls32UbTensor, - subM, 0, 0, 1, 1, qkNRound / FLOAT_ELENUM_PER_BLK); - } - AscendC::PipeBarrier(); - AscendC::WaitFlag(EVENT_ID0); - copyUbToGmOutput(gOutput, lpUbTensor[offset], layoutOutput, layoutInUb); - AscendC::SetFlag(pingpongFlag); - pingpongFlag = 1 - pingpongFlag; - } - } - - CATLASS_DEVICE - void operator()( - AscendC::GlobalTensor gOutput, - AscendC::GlobalTensor gInput, - AscendC::GlobalTensor gMask, - const LayoutOutput &layoutOutput, - const LayoutInput &layoutInput, - const LayoutMask &layoutMask, - GemmCoord actualBlockShape, - uint32_t nIdx, Arch::CrossCoreFlag qkReady) - { - uint32_t mActual = actualBlockShape.m(); - uint32_t nActual = actualBlockShape.n(); - - uint32_t subBlockIdx = AscendC::GetSubBlockIdx(); - uint32_t subBlockNum = AscendC::GetSubBlockNum(); - - uint32_t mActualPerSubBlock = CeilDiv(mActual, subBlockNum); - uint32_t mActualThisSubBlock = (subBlockIdx == 0) ? mActualPerSubBlock : (mActual - mActualPerSubBlock); - uint32_t mOffset = subBlockIdx * mActualPerSubBlock; - uint32_t nOffset = 0; - - int64_t offsetOutput = layoutOutput.GetOffset(MatrixCoord(mOffset, nOffset)); - auto gOutputThisSubBlock = gOutput[offsetOutput]; - auto layoutOutputThisSubBlock = layoutOutput.GetTileLayout(MatrixCoord(mActualThisSubBlock, nActual)); - - int64_t offsetInput = layoutInput.GetOffset(MatrixCoord(mOffset, nOffset)); - auto gInputThisSubBlock = gInput[offsetInput]; - auto layoutInputThisSubBlock = layoutInput.GetTileLayout(MatrixCoord(mActualThisSubBlock, nActual)); - - int64_t offsetMask = layoutMask.GetOffset(MatrixCoord(mOffset, nOffset)); - auto gMaskThisSubBlock = gMask[offsetMask]; - auto layoutMaskThisSubBlock = layoutMask.GetTileLayout(MatrixCoord(mActualThisSubBlock, nActual)); - - subCoreCompute(gOutputThisSubBlock, gInputThisSubBlock, gMaskThisSubBlock, - layoutOutputThisSubBlock, layoutInputThisSubBlock, layoutMaskThisSubBlock, - nIdx, qkReady); - } - -private: - half tor; - uint32_t pingpongFlag = 0; - AscendC::LocalTensor lsUbTensor; - AscendC::LocalTensor lpUbTensor; - AscendC::LocalTensor ls32UbTensor; - AscendC::LocalTensor maskUbTensor; - AscendC::LocalTensor lmUbTensor; - AscendC::LocalTensor hmUbTensor; - AscendC::LocalTensor gmUbTensor; - AscendC::LocalTensor dmUbTensor; - AscendC::LocalTensor llUbTensor; - AscendC::LocalTensor tvUbTensor; - - CopyGmToUbInput copyGmToUbInput; - CopyGmToUbMask copyGmToUbMask; - CopyUbToGmOutput copyUbToGmOutput; -}; - -} // namespace Catlass::Epilogue::Block - -#endif // CATLASS_KDA_BLOCK_EPILOGUE_FA_SOFTMAX_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/kda_block_epilogue_gemm.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/kda_block_epilogue_gemm.hpp deleted file mode 100644 index 7e7d043ec1..0000000000 --- a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/kda_block_epilogue_gemm.hpp +++ /dev/null @@ -1,218 +0,0 @@ -/** - * Copyright (c) 2025 Huawei Technologies Co., Ltd. - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of - * CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ - -#ifndef CATLASS_KDA_BLOCK_EPILOGUE_GEMM_HPP -#define CATLASS_KDA_BLOCK_EPILOGUE_GEMM_HPP - -#include "catlass/kda_catlass.hpp" -#include "catlass/arch/kda_resource.hpp" -#include "catlass/epilogue/kda_dispatch_policy.hpp" -#include "catlass/kda_gemm_coord.hpp" -#include "catlass/kda_matrix_coord.hpp" -#include "catlass/epilogue/tile/kda_epilogue_tile_copy.hpp" -#include "catlass/gemm/kda_gemm_helper.hpp" - -namespace Catlass::Epilogue::Block { - -template < - class CType_, - class XType_, - class DType_, - class TileElemWiseEpilogueAdd_, - class TileElemWiseEpilogueMuls_, - class TileElemWiseCastD_, - class TileCopy_> -class BlockEpilogue< - EpilogueAtlasA2Gemm, - CType_, - XType_, - DType_, - TileElemWiseEpilogueAdd_, - TileElemWiseEpilogueMuls_, - TileElemWiseCastD_, - TileCopy_> { -public: - // Type aliases - using DispatchPolicy = EpilogueAtlasA2Gemm; - using ArchTag = typename DispatchPolicy::ArchTag; - using ElementC = typename CType_::Element; - using LayoutC = typename CType_::Layout; - using ElementX = typename XType_::Element; - using LayoutX = typename XType_::Layout; - using ElementD = typename DType_::Element; - using LayoutD = typename DType_::Layout; - using TileElemWiseEpilogueAdd = TileElemWiseEpilogueAdd_; - using TileElemWiseEpilogueMuls = TileElemWiseEpilogueMuls_; - using TileElemWiseCastD = TileElemWiseCastD_; - using CopyGmToUbC = typename TileCopy_::CopyGmToUbC; - using CopyGmToUbX = typename TileCopy_::CopyGmToUbX; - using CopyUbToGmD = typename TileCopy_::CopyUbToGmD; - - const uint32_t SubNum = AscendC::GetSubBlockNum(); - const uint32_t UBSize = ArchTag::UB_SIZE; - static constexpr bool isNeedCast = !std::is_same::value; - static constexpr uint32_t COMPUTE_LENGTH = TileElemWiseEpilogueAdd::COMPUTE_LENGTH; - - using ElementCompute = typename Catlass::Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; - using ElementScalar = ElementCompute; - - // Check if ArchTag is matched - static_assert(std::is_same_v, "Tile epilogue's ArchTag mismatch"); - static_assert(std::is_same_v, "Tile epilogue's ArchTag mismatch"); - static_assert(std::is_same_v, "LayoutC only support RowMajor yet!"); - - struct Params { - ElementScalar alpha; - ElementScalar beta; - GM_ADDR ptrX; - LayoutX layoutX; - GM_ADDR ptrD; - LayoutD layoutD; - - CATLASS_HOST_DEVICE - Params() {} - - CATLASS_HOST_DEVICE - Params(ElementScalar alpha_, ElementScalar beta_, GM_ADDR ptrX_, LayoutX layoutX_, GM_ADDR ptrD_, LayoutD layoutD_) - : alpha(alpha_), - beta(beta_), - ptrX(ptrX_), - layoutX(layoutX_), - ptrD(ptrD_), - layoutD(layoutD_) - {} - }; - - CATLASS_DEVICE - BlockEpilogue(Arch::Resource &resource, GemmCoord blockShape_, Params const ¶ms_, uint32_t ubByteStart = 0) - : blockShapeMNK(blockShape_), - params(params_) - { - uint32_t maxMPerBlock = blockShapeMNK.m(); - uint32_t maxNPerBlock = blockShapeMNK.n(); - uint32_t tileSize = maxMPerBlock * maxNPerBlock / SubNum; - uint32_t ubCSize = tileSize * sizeof(ElementC); - uint32_t ubXSize = tileSize * sizeof(ElementX); - uint32_t ubDSize = tileSize * sizeof(ElementD); - uint32_t ubXCastSize = tileSize * sizeof(ElementCompute); - uint32_t ubDCastSize = tileSize * sizeof(ElementCompute); - ubCTensor = resource.ubBuf.template GetBufferByByte(ubByteStart); - ubByteStart += ubCSize; - ubXTensor = resource.ubBuf.template GetBufferByByte(ubByteStart); - ubByteStart += ubXSize; - ubDTensor = resource.ubBuf.template GetBufferByByte(ubByteStart); - ubByteStart += ubDSize; - if constexpr (isNeedCast) { - ubXTensorCast = resource.ubBuf.template GetBufferByByte(ubByteStart); - ubByteStart += ubXCastSize; - ubDTensorCast = resource.ubBuf.template GetBufferByByte(ubByteStart); - ; - ubByteStart += ubDCastSize; - } - AscendC::SetFlag(EVENT_ID0); - AscendC::SetFlag(EVENT_ID1); - AscendC::SetFlag(EVENT_ID0); - } - - CATLASS_DEVICE - ~BlockEpilogue() - { - AscendC::WaitFlag(EVENT_ID0); - AscendC::WaitFlag(EVENT_ID1); - AscendC::WaitFlag(EVENT_ID0); - } - - CATLASS_DEVICE - void operator()( - GemmCoord const &actualShapeMNK, GemmCoord const &blockCoordMNK, - AscendC::GlobalTensor const &gmBlockC, LayoutC const &layoutC, - uint64_t const &offset) - { - AscendC::GlobalTensor gmBlockX; - gmBlockX.SetGlobalBuffer(reinterpret_cast<__gm__ ElementX *>(params.ptrX)); - AscendC::GlobalTensor gmBlockD; - gmBlockD.SetGlobalBuffer(reinterpret_cast<__gm__ ElementD *>(params.ptrD)); - MatrixCoord blockShapeMN = blockShapeMNK.GetCoordMN(); - MatrixCoord blockCoordMN = blockCoordMNK.GetCoordMN(); - MatrixCoord actualShapeMN = actualShapeMNK.GetCoordMN(); - MatrixCoord blockOffset = blockCoordMN * blockShapeMN; - MatrixCoord subblockShape{ - CeilDiv(actualShapeMN.row(), SubNum), - actualShapeMN.column()}; - MatrixCoord subblockCoord{AscendC::GetSubBlockIdx(), 0}; - MatrixCoord actualSubblockShape = MatrixCoord::Min( - subblockShape, actualShapeMN - subblockCoord * subblockShape); - MatrixCoord subblockOffset = subblockCoord * subblockShape; - LayoutC layoutInUb{blockShapeMN.row() / SubNum, blockShapeMN.column()}; - AscendC::WaitFlag(EVENT_ID1); - auto layoutTileX = params.layoutX.GetTileLayout(actualSubblockShape); - auto layoutXInUb = layoutInUb.GetTileLayout(actualSubblockShape); - auto gmTileX = gmBlockX[offset + params.layoutX.GetOffset(blockOffset + subblockOffset)]; - copyGmToUbX(ubXTensor, gmTileX, layoutXInUb, layoutTileX); - AscendC::SetFlag(EVENT_ID1); - AscendC::WaitFlag(EVENT_ID1); - if constexpr (isNeedCast) { - AscendC::Cast(ubXTensorCast, ubXTensor, AscendC::RoundMode::CAST_NONE, COMPUTE_LENGTH); - AscendC::PipeBarrier(); - tileElemWiseEpilogueMuls(ubXTensorCast, ubXTensorCast, (ElementCompute)params.beta); - } else { - tileElemWiseEpilogueMuls(ubXTensor, ubXTensor, (ElementX)params.beta); - } - AscendC::WaitFlag(EVENT_ID0); - auto layoutTileC = layoutC.GetTileLayout(actualSubblockShape); - auto layoutCInUb = layoutInUb.GetTileLayout(actualSubblockShape); - auto gmTileC = gmBlockC[offset + layoutC.GetOffset(blockOffset + subblockOffset)]; - copyGmToUbC(ubCTensor, gmTileC, layoutCInUb, layoutTileC); - AscendC::SetFlag(EVENT_ID0); - AscendC::WaitFlag(EVENT_ID0); - tileElemWiseEpilogueMuls(ubCTensor, ubCTensor, (ElementC)params.alpha); - AscendC::WaitFlag(EVENT_ID0); - AscendC::PipeBarrier(); - if constexpr (isNeedCast) { - tileElemWiseEpilogueAdd(ubDTensorCast, ubCTensor, ubXTensorCast); - } else { - tileElemWiseEpilogueAdd(ubDTensor, ubCTensor, ubXTensor); - } - AscendC::SetFlag(EVENT_ID0); - AscendC::SetFlag(EVENT_ID1); - AscendC::PipeBarrier(); - if constexpr (isNeedCast) { - tileElemWiseCastD(ubDTensor, ubDTensorCast); - } - AscendC::SetFlag(EVENT_ID0); - AscendC::WaitFlag(EVENT_ID0); - auto layoutDInGm = params.layoutD.GetTileLayout(actualSubblockShape); - auto layoutTileD = layoutInUb.GetTileLayout(actualSubblockShape); - auto gmTileD = gmBlockD[offset + params.layoutD.GetOffset(blockOffset + subblockOffset)]; - copyUbToGmD(gmTileD, ubDTensor, layoutDInGm, layoutTileD); - AscendC::SetFlag(EVENT_ID0); - } - -private: - GemmCoord blockShapeMNK; - Params params; - - AscendC::LocalTensor ubCTensor; - AscendC::LocalTensor ubXTensor; - AscendC::LocalTensor ubDTensor; - AscendC::LocalTensor ubXTensorCast; - AscendC::LocalTensor ubDTensorCast; - - CopyGmToUbC copyGmToUbC; - CopyGmToUbX copyGmToUbX; - CopyUbToGmD copyUbToGmD; - - TileElemWiseEpilogueAdd tileElemWiseEpilogueAdd; - TileElemWiseEpilogueMuls tileElemWiseEpilogueMuls; - TileElemWiseCastD tileElemWiseCastD; -}; -} // namespace Catlass::Epilogue::Block - -#endif // CATLASS_KDA_BLOCK_EPILOGUE_GEMM_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/kda_block_epilogue_gemv.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/kda_block_epilogue_gemv.hpp deleted file mode 100644 index 4f58adc934..0000000000 --- a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/kda_block_epilogue_gemv.hpp +++ /dev/null @@ -1,222 +0,0 @@ -/** - * Copyright (c) 2025 Huawei Technologies Co., Ltd. - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of - * CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ - -#ifndef CATLASS_KDA_BLOCK_EPILOGUE_GEMV_HPP -#define CATLASS_KDA_BLOCK_EPILOGUE_GEMV_HPP - -#include "catlass/kda_catlass.hpp" -#include "catlass/arch/kda_resource.hpp" -#include "catlass/epilogue/kda_dispatch_policy.hpp" -#include "catlass/gemm/kda_gemm_helper.hpp" -#include "catlass/gemv/kda_gemv_helper.hpp" -#include "catlass/kda_gemv_coord.hpp" -#include "catlass/layout/kda_layout.hpp" -#include "catlass/kda_matrix_coord.hpp" - -namespace Catlass::Epilogue::Block { - -template < - class CType_, - class YType_, - class ZType_, - class TileElemWiseEpilogueAdd_, - class TileElemWiseEpilogueMuls_, - class TileCopy_> -class BlockEpilogue< - EpilogueAtlasA2Gemv, - CType_, - YType_, - ZType_, - TileElemWiseEpilogueAdd_, - TileElemWiseEpilogueMuls_, - TileCopy_> { -public: - using DispatchPolicy = EpilogueAtlasA2Gemv; - using ArchTag = typename DispatchPolicy::ArchTag; - - using ElementC = typename CType_::Element; - using LayoutC = typename CType_::Layout; - using ElementY = typename YType_::Element; - using LayoutY = typename YType_::Layout; - using ElementZ = typename ZType_::Element; - using LayoutZ = typename ZType_::Layout; - - using TileElemWiseEpilogueAdd = TileElemWiseEpilogueAdd_; - using TileElemWiseEpilogueMuls = TileElemWiseEpilogueMuls_; - - using CopyGmToUbY = typename TileCopy_::CopyGmToUbC; - using CopyGmToubC = typename TileCopy_::CopyGmToUbX; - using CopyUbToGmZ = typename TileCopy_::CopyUbToGmD; - - static constexpr uint32_t COMPUTE_LENGTH = TileElemWiseEpilogueMuls::COMPUTE_LENGTH; - - static constexpr bool isNeedCast = !std::is_same::value; - - using ElementCompute = typename Catlass::Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; - using ElementScalar = ElementCompute; - using TensorCoord = layout::VectorLayout::TensorCoord; - - // check the layout of Y, C and Z - static_assert(std::is_same_v && std::is_same_v && - std::is_same_v, - "Layout type of Y, C and Z must be VectorLayout"); - - using LayoutComputeInUb = layout::VectorLayout; - - // Check if ArchTag is matched - static_assert(std::is_same_v, "Tile epilogue's ArchTag mismatch"); - - struct Params { - ElementScalar alpha; - ElementScalar beta; - GM_ADDR ptrY; - LayoutY layoutY; - GM_ADDR ptrZ; - LayoutZ layoutZ; - - // Methods - CATLASS_HOST_DEVICE - Params() {} - - CATLASS_HOST_DEVICE - Params(ElementScalar alpha_, ElementScalar beta_, GM_ADDR ptrY_, LayoutC layoutY_, GM_ADDR ptrZ_, LayoutZ layoutZ_) - : alpha(alpha_), - beta(beta_), - ptrY(ptrY_), - layoutY(layoutY_), - ptrZ(ptrZ_), - layoutZ(layoutZ_) - {} - }; - - CATLASS_DEVICE - BlockEpilogue(Arch::Resource &resource, Params const ¶ms) - : params(params) - { - ubC = resource.ubBuf.template GetBufferByByte(0); - ubY = resource.ubBuf.template GetBufferByByte(COMPUTE_LENGTH * sizeof(ElementC)); - ubYCast = resource.ubBuf.template GetBufferByByte(COMPUTE_LENGTH * sizeof(ElementC)); - ubZ = resource.ubBuf.template GetBufferByByte( - COMPUTE_LENGTH * sizeof(ElementY) + COMPUTE_LENGTH * sizeof(ElementC)); - ubZCast = resource.ubBuf.template GetBufferByByte( - COMPUTE_LENGTH * sizeof(ElementY) + COMPUTE_LENGTH * sizeof(ElementC)); - AscendC::SetFlag(EVENT_ID0); - } - - CATLASS_DEVICE - ~BlockEpilogue() - { - AscendC::WaitFlag(EVENT_ID0); - } - - CATLASS_DEVICE - void operator()( - TensorCoord const &blockOffsetMN, - TensorCoord const &actualBlockShapeMN, - AscendC::GlobalTensor const &gmBlockC, - LayoutC const &layoutBlockC) - { - TensorCoord actualBlockShape = actualBlockShapeMN; - TensorCoord blockOffset = blockOffsetMN; - - TensorCoord subblockShape{ - CeilDiv(actualBlockShape[0], static_cast(AscendC::GetSubBlockNum()))}; - TensorCoord subblockCoord{static_cast(AscendC::GetSubBlockIdx())}; - - TensorCoord actualSubblockShape = TensorCoord::Min(subblockShape, actualBlockShape - subblockCoord * subblockShape); - TensorCoord subblockOffset = subblockCoord * subblockShape; - - // Get the data and layout of C - auto gmSubblockC = gmBlockC[layoutBlockC.GetOffset(subblockOffset)]; - auto layoutSubblockC = layoutBlockC.GetTileLayout(actualSubblockShape); - - // Get the data and layout of y - AscendC::GlobalTensor gmY; - gmY.SetGlobalBuffer(reinterpret_cast<__gm__ ElementY *>(params.ptrY)); - auto gmSubblockY = gmY[params.layoutY.GetOffset(blockOffset + subblockOffset)]; - auto layoutSubblockY = params.layoutY.GetTileLayout(actualSubblockShape); - - // Get the data and layout of Z - AscendC::GlobalTensor gmZ; - gmZ.SetGlobalBuffer(reinterpret_cast<__gm__ ElementZ *>(params.ptrZ)); - auto gmSubblockZ = gmZ[params.layoutZ.GetOffset(blockOffset + subblockOffset)]; - auto layoutSubblockZ = params.layoutZ.GetTileLayout(actualSubblockShape); - - // get the layout on UB - auto layoutComputeInUb = LayoutComputeInUb::template MakeLayoutInUb(actualSubblockShape); - - // load C(A*x) from gm to ub - AscendC::WaitFlag(EVENT_ID0); - copyGmToubC(ubC, gmSubblockC, layoutComputeInUb, layoutSubblockC); - AscendC::SetFlag(EVENT_ID0); - AscendC::WaitFlag(EVENT_ID0); - - // compute C * alpha - tileEpilogueMul(ubC, ubC, params.alpha); - - AscendC::SetFlag(EVENT_ID0); - AscendC::WaitFlag(EVENT_ID0); - - // load Y from gm to ub - copyGmToUbY(ubY, gmSubblockY, layoutComputeInUb, layoutSubblockY); - - AscendC::SetFlag(EVENT_ID0); - AscendC::WaitFlag(EVENT_ID0); - - // compute Y * beta - if constexpr (isNeedCast) { - AscendC::Cast(ubYCast, ubY, AscendC::RoundMode::CAST_NONE, COMPUTE_LENGTH); - AscendC::PipeBarrier(); - tileEpilogueMul(ubYCast, ubYCast, params.beta); - AscendC::PipeBarrier(); - } else { - tileEpilogueMul(ubY, ubY, params.beta); - AscendC::PipeBarrier(); - } - - if constexpr (isNeedCast) { - tileEpilogueAdd(ubZCast, ubC, ubYCast); - } else { - tileEpilogueAdd(ubZ, ubC, ubY); - } - - if constexpr (isNeedCast) { - AscendC::PipeBarrier(); - AscendC::Cast(ubZ, ubZCast, AscendC::RoundMode::CAST_RINT, COMPUTE_LENGTH); - AscendC::PipeBarrier(); - } - - AscendC::SetFlag(EVENT_ID0); - AscendC::WaitFlag(EVENT_ID0); - - copyUbToGmZ(gmSubblockZ, ubZ, layoutSubblockZ, layoutComputeInUb); - AscendC::SetFlag(EVENT_ID0); - }; - -private: - Params params; - - AscendC::LocalTensor ubY; - AscendC::LocalTensor ubYCast; - AscendC::LocalTensor ubC; - AscendC::LocalTensor ubZ; - AscendC::LocalTensor ubZCast; - - TileElemWiseEpilogueAdd tileEpilogueAdd; - TileElemWiseEpilogueMuls tileEpilogueMul; - - CopyGmToUbY copyGmToUbY; - CopyGmToubC copyGmToubC; - CopyUbToGmZ copyUbToGmZ; -}; - -} // namespace Catlass::Epilogue::Block - -#endif // CATLASS_KDA_BLOCK_EPILOGUE_GEMV_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/kda_block_epilogue_mla_softmax.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/kda_block_epilogue_mla_softmax.hpp deleted file mode 100644 index 46d939296b..0000000000 --- a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/kda_block_epilogue_mla_softmax.hpp +++ /dev/null @@ -1,441 +0,0 @@ -/** - * Copyright (c) 2025 Huawei Technologies Co., Ltd. - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of - * CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ - -#ifndef CATLASS_KDA_BLOCK_EPILOGUE_MLA_SOFTMAX_HPP -#define CATLASS_KDA_BLOCK_EPILOGUE_MLA_SOFTMAX_HPP - -#include "catlass/kda_catlass.hpp" -#include "catlass/arch/kda_cross_core_sync.hpp" -#include "catlass/arch/kda_resource.hpp" -#include "catlass/epilogue/kda_dispatch_policy.hpp" -#include "catlass/epilogue/tile/kda_epilogue_tile_copy.hpp" -#include "catlass/kda_gemm_coord.hpp" -#include "catlass/kda_matrix_coord.hpp" - -namespace Catlass::Epilogue::Block { - -template < - class OutputType_, - class InputType_, - class MaskType_> -class BlockEpilogue< - EpilogueAtlasA2MLASoftmax, - OutputType_, - InputType_, - MaskType_> { -public: - // Type aliases - using DispatchPolicy = EpilogueAtlasA2MLASoftmax; - using ArchTag = typename DispatchPolicy::ArchTag; - using ElementOutput = typename OutputType_::Element; - using ElementInput = typename InputType_::Element; - using ElementMask = typename MaskType_::Element; - - using LayoutOutput = typename OutputType_::Layout; - using LayoutInput = typename InputType_::Layout; - using LayoutMask = typename MaskType_::Layout; - - using CopyGmToUbInput = Tile::CopyGm2Ub; - using CopyGmToUbMask = Tile::CopyGm2Ub; - using CopyUbToGmOutput = Tile::CopyUb2Gm; - - static constexpr uint32_t HALF_ELENUM_PER_BLK = 16; - static constexpr uint32_t FLOAT_BLOCK_SIZE = 8; - static constexpr uint32_t HALF_ELENUM_PER_VECCALC = 128; - static constexpr uint32_t FLOAT_VECTOR_SIZE = 64; - static constexpr uint32_t UB_TILE_SIZE = 16384; // 64 * 128 * 2B - static constexpr uint32_t UB_LINE_SIZE = 512; // 128 * 2 * 2B - static constexpr uint32_t HALF_ELENUM_PER_LINE = 256; // 128 * 2 - static constexpr uint32_t FLOAT_ELENUM_PER_LINE = 128; // 128 - static constexpr uint32_t MULTIPLIER = 2; - static constexpr uint32_t HALF_VECTOR_SIZE = 128; - static constexpr uint32_t BLOCK_SIZE = 16; - static constexpr uint32_t T_BLOCK_SIZE = 32 / 2; - static constexpr uint32_t UB_UINT8_LINE_SIZE = 512; - static constexpr uint32_t UB_UINT8_BLOCK_SIZE_MLA = 16384; - static constexpr uint32_t HALF_DM_UB_SIZE = 128; - static constexpr uint32_t VECTOR_SIZE = 128; - static constexpr uint32_t HALF_LL_UB_SIZE = 256; - - CATLASS_DEVICE - BlockEpilogue(Arch::Resource &resource, half tor_, uint32_t kvSplitCoreNum_) - { - // Allocate UB space - constexpr uint32_t LS_UB_TENSOR_OFFSET = 0; - constexpr uint32_t LP_UB_TENSOR_OFFSET = 2 * UB_UINT8_BLOCK_SIZE_MLA; - constexpr uint32_t LM_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA; - constexpr uint32_t HM_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA + 1 * UB_UINT8_LINE_SIZE; - constexpr uint32_t DM_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA + 6 * UB_UINT8_LINE_SIZE; - constexpr uint32_t LL_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA + 9 * UB_UINT8_LINE_SIZE; - constexpr uint32_t GM_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA + 13 * UB_UINT8_LINE_SIZE; - constexpr uint32_t TV_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE_MLA; - - tor = tor_; - kvSplitCoreNum = kvSplitCoreNum_; - tvUbTensor16 = resource.ubBuf.template GetBufferByByte(LP_UB_TENSOR_OFFSET); - lpUbTensor32 = resource.ubBuf.template GetBufferByByte(LP_UB_TENSOR_OFFSET); - lsUbTensor = resource.ubBuf.template GetBufferByByte(LS_UB_TENSOR_OFFSET); - lmUbTensor = resource.ubBuf.template GetBufferByByte(LM_UB_TENSOR_OFFSET); - hmUbTensor = resource.ubBuf.template GetBufferByByte(HM_UB_TENSOR_OFFSET); - gmUbTensor = resource.ubBuf.template GetBufferByByte(GM_UB_TENSOR_OFFSET); - dmUbTensor = resource.ubBuf.template GetBufferByByte(DM_UB_TENSOR_OFFSET); - llUbTensor = resource.ubBuf.template GetBufferByByte(LL_UB_TENSOR_OFFSET); - tvUbTensor = resource.ubBuf.template GetBufferByByte(TV_UB_TENSOR_OFFSET); - } - - CATLASS_DEVICE - ~BlockEpilogue() - { - } - - CATLASS_DEVICE - void SetMask(int32_t len) - { - uint64_t mask = 0; - uint64_t one = 1; - uint64_t temp = len % FLOAT_VECTOR_SIZE; - for (int64_t i = 0; i < temp; i++) { - mask |= one << i; - } - if (len == VECTOR_SIZE) { - AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); - } else if (len >= FLOAT_VECTOR_SIZE) { - AscendC::SetVectorMask(mask, (uint64_t)-1); - } else { - AscendC::SetVectorMask(0x0, mask); - } - } - - CATLASS_DEVICE - void ReduceSumRepeatM( - const AscendC::LocalTensor &dst, - const AscendC::LocalTensor &src, - uint32_t curRowNum, - uint32_t kSeqTile, - uint32_t kSeqTileRound) - { - if (kSeqTile <= FLOAT_VECTOR_SIZE) { - SetMask(kSeqTile); - AscendC::RepeatReduceSum(dst, src, curRowNum, 0, 0, 1, 1, kSeqTileRound / FLOAT_BLOCK_SIZE); - AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); - } else { - for (uint32_t rowsum_idx = 1; rowsum_idx < kSeqTile / FLOAT_VECTOR_SIZE; ++rowsum_idx) { - AscendC::Add( - src, - src, - src[rowsum_idx * FLOAT_VECTOR_SIZE], - (uint64_t)0, - curRowNum, - AscendC::BinaryRepeatParams( - 1, 1, 1, kSeqTileRound / FLOAT_BLOCK_SIZE, - kSeqTileRound / FLOAT_BLOCK_SIZE, kSeqTileRound / FLOAT_BLOCK_SIZE)); - AscendC::PipeBarrier(); - } - if (kSeqTile % FLOAT_VECTOR_SIZE > 0) { - SetMask(kSeqTile % FLOAT_VECTOR_SIZE); - AscendC::Add( - src, - src, - src[kSeqTile / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - (uint64_t)0, - curRowNum, - AscendC::BinaryRepeatParams( - 1, 1, 1, kSeqTileRound / FLOAT_BLOCK_SIZE, - kSeqTileRound / FLOAT_BLOCK_SIZE, kSeqTileRound / FLOAT_BLOCK_SIZE)); - AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); - } - AscendC::PipeBarrier(); - AscendC::RepeatReduceSum(dst, src, curRowNum, 0, 0, 1, 1, kSeqTileRound / FLOAT_BLOCK_SIZE); - } - } - - CATLASS_DEVICE - void TensorSubValueRepeatM( - const AscendC::LocalTensor &dst, - const AscendC::LocalTensor &src, - const AscendC::LocalTensor &MaxTensor, - const AscendC::LocalTensor &tempMaxTensor, - uint32_t curRowNum, - uint32_t subMRound, - uint32_t kSeqTile, - uint32_t kSeqTileRound) - { - AscendC::Brcb(tempMaxTensor.ReinterpretCast(), MaxTensor.ReinterpretCast(), - subMRound / FLOAT_BLOCK_SIZE, AscendC::BrcbRepeatParams(1, 8)); - AscendC::PipeBarrier(); - for (uint32_t subIdx = 0; subIdx < kSeqTile / FLOAT_VECTOR_SIZE; ++subIdx) { - AscendC::Sub( - dst[subIdx * FLOAT_VECTOR_SIZE], - src[subIdx * FLOAT_VECTOR_SIZE], - tempMaxTensor, - (uint64_t)0, - curRowNum, - AscendC::BinaryRepeatParams( - 1, 1, 0, kSeqTileRound / FLOAT_BLOCK_SIZE, kSeqTileRound / FLOAT_BLOCK_SIZE, 1)); - } - if (kSeqTile % FLOAT_VECTOR_SIZE > 0) { - SetMask(kSeqTile % FLOAT_VECTOR_SIZE); - AscendC::Sub( - dst[kSeqTile / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - src[kSeqTile / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - tempMaxTensor, - (uint64_t)0, - curRowNum, - AscendC::BinaryRepeatParams( - 1, 1, 0, kSeqTileRound / FLOAT_BLOCK_SIZE, kSeqTileRound / FLOAT_BLOCK_SIZE, 1)); - AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); - } - AscendC::PipeBarrier(); - } - - CATLASS_DEVICE - void ReduceMaxRepeatM( - const AscendC::LocalTensor &dst, - const AscendC::LocalTensor &src, - const AscendC::LocalTensor &tempTensor, - uint32_t curRowNum, - uint32_t kSeqTile, - uint32_t kSeqTileRound) - { - if (kSeqTile <= FLOAT_VECTOR_SIZE) { - SetMask(kSeqTile); - AscendC::WholeReduceMax( - dst, src, (int32_t)0, curRowNum, 1, 1, - kSeqTileRound / FLOAT_BLOCK_SIZE, AscendC::ReduceOrder::ORDER_ONLY_VALUE); - } else { - AscendC::DataCopy(tempTensor, src, AscendC::DataCopyParams(curRowNum, HALF_VECTOR_SIZE / BLOCK_SIZE, (kSeqTileRound - FLOAT_VECTOR_SIZE) / FLOAT_BLOCK_SIZE, 0)); - AscendC::PipeBarrier(); - for (uint32_t rowmaxIdx = 1; rowmaxIdx < kSeqTile / FLOAT_VECTOR_SIZE; ++rowmaxIdx) { - AscendC::Max( - tempTensor, - tempTensor, - src[rowmaxIdx * FLOAT_VECTOR_SIZE], - (uint64_t)0, - curRowNum, - AscendC::BinaryRepeatParams( - 1, 1, 1, 8, 8, kSeqTileRound / FLOAT_BLOCK_SIZE)); - AscendC::PipeBarrier(); - } - if (kSeqTile % FLOAT_VECTOR_SIZE > 0) { - SetMask(kSeqTile % FLOAT_VECTOR_SIZE); - AscendC::Max( - tempTensor, - tempTensor, - src[kSeqTile / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - (uint64_t)0, - curRowNum, - AscendC::BinaryRepeatParams( - 1, 1, 1, 8, 8, kSeqTileRound / FLOAT_BLOCK_SIZE)); - } - AscendC::PipeBarrier(); - AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); - AscendC::WholeReduceMax( - dst, tempTensor, (int32_t)0, curRowNum, 1, 1, 8, AscendC::ReduceOrder::ORDER_ONLY_VALUE); - } - AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); - AscendC::PipeBarrier(); - } - - CATLASS_DEVICE - void SubCoreCompute( - AscendC::GlobalTensor gOutput, - AscendC::GlobalTensor gInput, - const LayoutOutput &layoutOutput, - const LayoutInput &layoutInput, - uint32_t nIdx, - uint32_t softmaxPingPongFlag, - uint32_t &glFlag) - { - uint32_t curRowNum = layoutInput.shape(0); - uint32_t kSeqTile = layoutInput.shape(1); - uint32_t kSeqTileRound = layoutInput.stride(0); - uint32_t subMRound = (curRowNum + 16 - 1) / 16 * 16; - uint32_t sub_m_d64 = (curRowNum + 63) / 64; // up aligned to 128 - uint64_t dmUbOffsetCurCycle = (uint64_t)(softmaxPingPongFlag * HALF_DM_UB_SIZE); - uint64_t llUbOffsetCurCycle = (uint64_t)(softmaxPingPongFlag * HALF_LL_UB_SIZE); - AscendC::WaitFlag(EVENT_ID2); - AscendC::DataCopy(lsUbTensor, gInput, - AscendC::DataCopyParams(1, curRowNum * kSeqTileRound / FLOAT_BLOCK_SIZE, 0, 0)); - - AscendC::SetFlag(EVENT_ID0); - AscendC::WaitFlag(EVENT_ID0); - - // muls scale_value - for (uint32_t mulsIdx = 0; mulsIdx < kSeqTile / FLOAT_VECTOR_SIZE; ++mulsIdx) { - AscendC::Muls( - lsUbTensor[mulsIdx * FLOAT_VECTOR_SIZE], - lsUbTensor[mulsIdx * FLOAT_VECTOR_SIZE], - tor, - (uint64_t)0, - curRowNum, - AscendC::UnaryRepeatParams(1, 1, kSeqTileRound / FLOAT_BLOCK_SIZE, kSeqTileRound / FLOAT_BLOCK_SIZE)); - } - if (kSeqTile % FLOAT_VECTOR_SIZE > 0) { - SetMask(kSeqTile % FLOAT_VECTOR_SIZE); - AscendC::Muls( - lsUbTensor[kSeqTile / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - lsUbTensor[kSeqTile / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - tor, - (uint64_t)0, - curRowNum, - AscendC::UnaryRepeatParams(1, 1, kSeqTileRound / FLOAT_BLOCK_SIZE, kSeqTileRound / FLOAT_BLOCK_SIZE)); - - AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); - } - AscendC::PipeBarrier(); - - // *** lm = rowmax(ls) - ReduceMaxRepeatM(lmUbTensor, lsUbTensor, lpUbTensor32, curRowNum, kSeqTile, kSeqTileRound); - - if (nIdx != 0) { - AscendC::Max( - hmUbTensor, - lmUbTensor, - gmUbTensor, - (uint64_t)0, - sub_m_d64, - AscendC::BinaryRepeatParams( - 1, 1, 1, 8, 8, 8)); - AscendC::PipeBarrier(); - AscendC::Sub( - dmUbTensor[dmUbOffsetCurCycle], - gmUbTensor, - hmUbTensor, - (uint64_t)0, - sub_m_d64, - AscendC::BinaryRepeatParams( - 1, 1, 1, 8, 8, 8)); - AscendC::PipeBarrier(); - } else { - AscendC::DataCopy(hmUbTensor, lmUbTensor, AscendC::DataCopyParams(1, subMRound / FLOAT_BLOCK_SIZE, 0, 0)); - AscendC::PipeBarrier(); - } - // *** gm = hm - AscendC::DataCopy(gmUbTensor, hmUbTensor, AscendC::DataCopyParams(1, subMRound / FLOAT_BLOCK_SIZE, 0, 0)); - AscendC::PipeBarrier(); - - if (kvSplitCoreNum != 1) { - if (nIdx == 0) { - if (glFlag == 1) { - AscendC::WaitFlag(EVENT_ID2); - glFlag = 0; - } - } - } - - // *** ls = ls - hm_block - TensorSubValueRepeatM(lsUbTensor, lsUbTensor, - hmUbTensor, tvUbTensor, - curRowNum, subMRound, kSeqTile, kSeqTileRound); - - AscendC::Exp( - lsUbTensor, - lsUbTensor, - (uint64_t)0, - (curRowNum * kSeqTileRound + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, - AscendC::UnaryRepeatParams(1, 1, 8, 8)); - - AscendC::PipeBarrier(); - // *** lp = castfp32to16(ls) - if (std::is_same::value) { - AscendC::Cast( - tvUbTensor16, lsUbTensor, AscendC::RoundMode::CAST_RINT, (uint64_t)0, - (curRowNum * kSeqTileRound + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, - AscendC::UnaryRepeatParams(1, 1, 4, 8)); - } else { - AscendC::Cast( - tvUbTensor16, lsUbTensor, AscendC::RoundMode::CAST_NONE, (uint64_t)0, - (curRowNum * kSeqTileRound + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, - AscendC::UnaryRepeatParams(1, 1, 4, 8)); - } - - AscendC::PipeBarrier(); - - AscendC::SetFlag(EVENT_ID0); - AscendC::WaitFlag(EVENT_ID0); - - uint16_t blockCount = 1; - uint16_t blockLen = curRowNum * kSeqTileRound / T_BLOCK_SIZE; - uint16_t srcStride = 0; - uint16_t dstStride = 0; - - AscendC::DataCopy(gOutput, - tvUbTensor16, - AscendC::DataCopyParams( - blockCount, // blockCount - blockLen, // blockLen - srcStride, // srcGap - dstStride)); - - // *** ll = rowsum(ls32) - ReduceSumRepeatM(llUbTensor[llUbOffsetCurCycle], lsUbTensor, curRowNum, kSeqTile, kSeqTileRound); - AscendC::SetFlag(EVENT_ID2); - AscendC::PipeBarrier(); - } - - CATLASS_DEVICE - void operator()( - AscendC::GlobalTensor gOutput, - AscendC::GlobalTensor gInput, - const LayoutOutput &layoutOutput, - const LayoutInput &layoutInput, - GemmCoord actualBlockShape, - uint32_t nIdx, - uint32_t curHeadNum, - uint32_t softmaxPingPongFlag, - uint32_t &glFlag) - { - uint32_t rowActual = actualBlockShape.m(); - uint32_t nActual = actualBlockShape.n(); - uint32_t tokenNumPerHead = rowActual / curHeadNum; - - uint32_t subBlockIdx = AscendC::GetSubBlockIdx(); - uint32_t subBlockNum = AscendC::GetSubBlockNum(); - - uint32_t curHeadSplitSubBlock = curHeadNum / subBlockNum; - uint32_t curHeadThisSubBlock = (subBlockIdx == 0) ? curHeadSplitSubBlock : (curHeadNum - curHeadSplitSubBlock); - - uint32_t rowActualThisSubBlock = curHeadThisSubBlock * tokenNumPerHead; - uint32_t rowOffsetSubBlock = subBlockIdx * curHeadSplitSubBlock * tokenNumPerHead; - - if (rowActualThisSubBlock > 0) { - int64_t offsetInput = layoutInput.GetOffset(MatrixCoord(rowOffsetSubBlock, 0)); - auto gInputThisSubBlock = gInput[offsetInput]; - auto layoutInputThisSubBlock = layoutInput.GetTileLayout(MatrixCoord(rowActualThisSubBlock, nActual)); - int64_t offsetOutput = layoutOutput.GetOffset(MatrixCoord(rowOffsetSubBlock, 0)); - auto gOutputThisSubBlock = gOutput[offsetOutput]; - auto layoutOutputThisSubBlock = layoutOutput.GetTileLayout(MatrixCoord(rowActualThisSubBlock, nActual)); - SubCoreCompute(gOutputThisSubBlock, gInputThisSubBlock, layoutOutputThisSubBlock, layoutInputThisSubBlock, - nIdx, softmaxPingPongFlag, glFlag); - } - } - -private: - float tor; - uint32_t pingpongFlag = 0; - uint32_t kvSplitCoreNum = 1; - AscendC::LocalTensor tvUbTensor16; - AscendC::LocalTensor lpUbTensor32; - AscendC::LocalTensor lsUbTensor; - AscendC::LocalTensor lmUbTensor; - AscendC::LocalTensor hmUbTensor; - AscendC::LocalTensor gmUbTensor; - AscendC::LocalTensor dmUbTensor; - AscendC::LocalTensor llUbTensor; - AscendC::LocalTensor tvUbTensor; - - CopyGmToUbInput copyGmToUbInput; - CopyGmToUbMask copyGmToUbMask; - CopyUbToGmOutput copyUbToGmOutput; -}; - -} // namespace Catlass::Epilogue::Block - -#endif // CATLASS_KDA_BLOCK_EPILOGUE_MLA_SOFTMAX_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/kda_block_epilogue_mla_tp1_softmax.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/kda_block_epilogue_mla_tp1_softmax.hpp deleted file mode 100644 index 1ade25fd59..0000000000 --- a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/kda_block_epilogue_mla_tp1_softmax.hpp +++ /dev/null @@ -1,507 +0,0 @@ -/** - * Copyright (c) 2025 Huawei Technologies Co., Ltd. - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of - * CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ - -#ifndef CATLASS_KDA_BLOCK_EPILOGUE_MLA_TP1_SOFTMAX_HPP -#define CATLASS_KDA_BLOCK_EPILOGUE_MLA_TP1_SOFTMAX_HPP - -#include "catlass/kda_catlass.hpp" -#include "catlass/arch/kda_cross_core_sync.hpp" -#include "catlass/arch/kda_resource.hpp" -#include "catlass/epilogue/kda_dispatch_policy.hpp" -#include "catlass/epilogue/tile/kda_epilogue_tile_copy.hpp" -#include "catlass/kda_gemm_coord.hpp" -#include "catlass/kda_matrix_coord.hpp" - -namespace Catlass::Epilogue::Block { - -template < - class OutputType_, - class InputType_, - class MaskType_> -class BlockEpilogue< - EpilogueAtlasA2MLATP1Softmax, - OutputType_, - InputType_, - MaskType_> { -public: - // Type aliases - using DispatchPolicy = EpilogueAtlasA2MLATP1Softmax; - using ArchTag = typename DispatchPolicy::ArchTag; - using ElementOutput = typename OutputType_::Element; - using ElementInput = typename InputType_::Element; - using ElementMask = typename MaskType_::Element; - - using LayoutOutput = typename OutputType_::Layout; - using LayoutInput = typename InputType_::Layout; - using LayoutMask = typename MaskType_::Layout; - - using CopyGmToUbInput = Tile::CopyGm2Ub; - using CopyGmToUbMask = Tile::CopyGm2Ub; - using CopyUbToGmOutput = Tile::CopyUb2Gm; - - static constexpr uint32_t FLOAT_BLOCK_SIZE = 8; - static constexpr uint32_t FLOAT_VECTOR_SIZE = 64; - static constexpr uint32_t HALF_VECTOR_SIZE = 128; - static constexpr uint32_t BLOCK_SIZE = 16; - static constexpr uint32_t UB_UINT8_LINE_SIZE = 512; - static constexpr uint32_t UB_UINT8_BLOCK_SIZE_MLA = 16384; - static constexpr uint32_t VECTOR_SIZE = 128; - - static constexpr uint32_t REDUCE_UB_SIZE = 1024; - static constexpr uint32_t ROW_OPS_SPEC_MASK_32 = 32; - static constexpr uint32_t ROW_OPS_SPEC_MASK_4 = 4; - static constexpr uint32_t S_BLOCK_STACK = 4; - static constexpr int64_t UB_FLOAT_LINE_SIZE = 64; - static constexpr uint32_t M_SLICE = 16; - static constexpr uint32_t QK_READY_ID = 1; - - CATLASS_DEVICE - BlockEpilogue(Arch::Resource &resource, half tor_, uint32_t kvSplitCoreNum_ = 1) - { - // Allocate UB space - constexpr uint32_t LS_UB_TENSOR_OFFSET = 0; - constexpr uint32_t LP_UB_TENSOR_OFFSET = 0; - constexpr uint32_t LM_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA; - constexpr uint32_t HM_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA + 1 * UB_UINT8_LINE_SIZE; - constexpr uint32_t DM_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA + 6 * UB_UINT8_LINE_SIZE; - constexpr uint32_t LL_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA + 10 * UB_UINT8_LINE_SIZE; - constexpr uint32_t GM_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA + 14 * UB_UINT8_LINE_SIZE; - constexpr uint32_t GL_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA + 16 * UB_UINT8_LINE_SIZE; - constexpr uint32_t TV_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE_MLA; - - tor = tor_; - kvSplitCoreNum = kvSplitCoreNum_; - lsUbTensor = resource.ubBuf.template GetBufferByByte(LS_UB_TENSOR_OFFSET); - lpUbTensor = resource.ubBuf.template GetBufferByByte(LP_UB_TENSOR_OFFSET); - lmUbTensor = resource.ubBuf.template GetBufferByByte(LM_UB_TENSOR_OFFSET); - hmUbTensor = resource.ubBuf.template GetBufferByByte(HM_UB_TENSOR_OFFSET); - gmUbTensor = resource.ubBuf.template GetBufferByByte(GM_UB_TENSOR_OFFSET); - dmUbTensor = resource.ubBuf.template GetBufferByByte(DM_UB_TENSOR_OFFSET); - llUbTensor = resource.ubBuf.template GetBufferByByte(LL_UB_TENSOR_OFFSET); - tvUbTensor = resource.ubBuf.template GetBufferByByte(TV_UB_TENSOR_OFFSET); - glUbTensor = resource.ubBuf.template GetBufferByByte(GL_UB_TENSOR_OFFSET); - } - - CATLASS_DEVICE - ~BlockEpilogue() {} - - CATLASS_DEVICE - void SetVecMask(int32_t len) - { - uint64_t mask = 0; - uint64_t one = 1; - uint64_t temp = len % FLOAT_VECTOR_SIZE; - for (int64_t i = 0; i < temp; i++) { - mask |= one << i; - } - - if (len == VECTOR_SIZE || len == 0) { - AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); - } else if (len >= FLOAT_VECTOR_SIZE) { - AscendC::SetVectorMask(mask, (uint64_t)-1); - } else { - AscendC::SetVectorMask(0x0, mask); - } - } - - CATLASS_DEVICE - void SetkvSplitCoreNum(uint32_t kvSplitCoreNum_) - { - kvSplitCoreNum = kvSplitCoreNum_; - } - - CATLASS_DEVICE - void SetBlockReduceMask(int32_t len) - { - if (len > 8 || len < 1) { - AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); - return; - } - uint64_t subMask = ((uint64_t)1 << len) - 1; - uint64_t maskValue = (subMask << 48) + (subMask << 32) + (subMask << 16) + subMask + (subMask << 56) + - (subMask << 40) + (subMask << 24) + (subMask << 8); - AscendC::SetVectorMask(maskValue, maskValue); - } - - CATLASS_DEVICE - void RowsumSPECTILE512(const AscendC::LocalTensor &srcUb, const AscendC::LocalTensor &rowsumUb, - const AscendC::LocalTensor &tvUbTensor, uint32_t numRowsRound, uint32_t numElems, - uint32_t numElemsAligned) - { - AscendC::BlockReduceSum(tvUbTensor, srcUb, numRowsRound * numElemsAligned / FLOAT_VECTOR_SIZE, 0, - 1, 1, 8); - AscendC::PipeBarrier(); - - AscendC::BlockReduceSum(tvUbTensor[REDUCE_UB_SIZE], tvUbTensor, - numRowsRound * numElemsAligned / FLOAT_BLOCK_SIZE / FLOAT_VECTOR_SIZE, 0, - 1, 1, 8); - AscendC::PipeBarrier(); - AscendC::BlockReduceSum(rowsumUb, tvUbTensor[REDUCE_UB_SIZE], - numRowsRound * numElemsAligned / FLOAT_VECTOR_SIZE / FLOAT_VECTOR_SIZE, 0, - 1, 1, 8); - AscendC::PipeBarrier(); - } - - CATLASS_DEVICE - void RowsumSPECTILE256(const AscendC::LocalTensor &srcUb, const AscendC::LocalTensor &rowsumUb, - const AscendC::LocalTensor &tvUbTensor, uint32_t numRowsRound, uint32_t numElems, - uint32_t numElemsAligned) - { - AscendC::BlockReduceSum(tvUbTensor, srcUb, numRowsRound * numElemsAligned / FLOAT_VECTOR_SIZE, 0, - 1, 1, 8); - AscendC::PipeBarrier(); - SetVecMask(ROW_OPS_SPEC_MASK_32); - AscendC::BlockReduceSum(tvUbTensor[REDUCE_UB_SIZE], tvUbTensor, numRowsRound, 0, 1, 1, 4); - AscendC::PipeBarrier(); - SetBlockReduceMask(ROW_OPS_SPEC_MASK_4); - AscendC::BlockReduceSum( - rowsumUb, tvUbTensor[REDUCE_UB_SIZE], - (numRowsRound * FLOAT_BLOCK_SIZE + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, 0, 1, 1, 8); - AscendC::PipeBarrier(); - AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); - } - - CATLASS_DEVICE - void RowsumTAILTILE(const AscendC::LocalTensor &srcUb, const AscendC::LocalTensor &rowsumUb, - const AscendC::LocalTensor &tvUbTensor, uint32_t numRowsRound, uint32_t numElems, - uint32_t numElemsAligned) - { - if (numElems >= FLOAT_VECTOR_SIZE) { - AscendC::BlockReduceSum(tvUbTensor, srcUb, numRowsRound, 0, 1, 1, - numElemsAligned / FLOAT_BLOCK_SIZE); - AscendC::PipeBarrier(); - AscendC::BlockReduceSum( - rowsumUb, tvUbTensor, (numRowsRound * FLOAT_BLOCK_SIZE + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, 0, - 1, 1, 8); - AscendC::PipeBarrier(); - for (uint64_t rowSumIdx = 1; rowSumIdx < (uint64_t)numElems / FLOAT_VECTOR_SIZE; ++rowSumIdx) { - AscendC::BlockReduceSum(tvUbTensor, srcUb[rowSumIdx * FLOAT_VECTOR_SIZE], numRowsRound, 0, - 1, 1, numElemsAligned / FLOAT_BLOCK_SIZE); - AscendC::PipeBarrier(); - AscendC::BlockReduceSum( - tvUbTensor[REDUCE_UB_SIZE], tvUbTensor, - (numRowsRound * FLOAT_BLOCK_SIZE + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, 0, 1, 1, 8); - AscendC::PipeBarrier(); - SetVecMask(numRowsRound); - AscendC::Add(rowsumUb, rowsumUb, tvUbTensor[REDUCE_UB_SIZE], (uint64_t)0, 1, - AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); - AscendC::PipeBarrier(); - AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); - } - } - if (numElems % FLOAT_VECTOR_SIZE > 0) { - SetVecMask(numElems % FLOAT_VECTOR_SIZE); - AscendC::BlockReduceSum(tvUbTensor, srcUb[numElems / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - numRowsRound, 0, 1, 1, numElemsAligned / FLOAT_BLOCK_SIZE); - AscendC::PipeBarrier(); - SetBlockReduceMask((numElems % FLOAT_VECTOR_SIZE + FLOAT_BLOCK_SIZE - 1) / FLOAT_BLOCK_SIZE); - if (numElems < FLOAT_VECTOR_SIZE) { - AscendC::BlockReduceSum( - rowsumUb, tvUbTensor, (numRowsRound * FLOAT_BLOCK_SIZE + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, - 0, 1, 1, 8); - AscendC::PipeBarrier(); - } else { - AscendC::BlockReduceSum( - tvUbTensor[REDUCE_UB_SIZE], tvUbTensor, - (numRowsRound * FLOAT_BLOCK_SIZE + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, 0, 1, 1, 8); - AscendC::PipeBarrier(); - SetVecMask(numRowsRound); - AscendC::Add(rowsumUb, rowsumUb, tvUbTensor[REDUCE_UB_SIZE], (uint64_t)0, 1, - AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); - AscendC::PipeBarrier(); - } - AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); - } - } - - CATLASS_DEVICE - void RowmaxSPECTILE512(const AscendC::LocalTensor &srcUb, const AscendC::LocalTensor &rowmaxUb, - const AscendC::LocalTensor &tvUbTensor, uint32_t numRowsRound, uint32_t numElems, - uint32_t numElemsAligned) - { - AscendC::BlockReduceMax(tvUbTensor, srcUb, numRowsRound * numElemsAligned / FLOAT_VECTOR_SIZE, 0, - 1, 1, 8); - AscendC::PipeBarrier(); - AscendC::BlockReduceMax(tvUbTensor[REDUCE_UB_SIZE], tvUbTensor, - numRowsRound * numElemsAligned / FLOAT_BLOCK_SIZE / FLOAT_VECTOR_SIZE, 0, - 1, 1, 8); - AscendC::PipeBarrier(); - AscendC::BlockReduceMax(rowmaxUb, tvUbTensor[REDUCE_UB_SIZE], - numRowsRound * numElemsAligned / FLOAT_VECTOR_SIZE / FLOAT_VECTOR_SIZE, 0, - 1, 1, 8); - AscendC::PipeBarrier(); - } - - CATLASS_DEVICE - void RowmaxSPECTILE256(const AscendC::LocalTensor &srcUb, const AscendC::LocalTensor &rowmaxUb, - const AscendC::LocalTensor &tvUbTensor, uint32_t numRowsRound, uint32_t numElems, - uint32_t numElemsAligned) - { - AscendC::BlockReduceMax(tvUbTensor, srcUb, numRowsRound * numElemsAligned / FLOAT_VECTOR_SIZE, 0, - 1, 1, 8); - AscendC::PipeBarrier(); - SetVecMask(ROW_OPS_SPEC_MASK_32); - AscendC::BlockReduceMax(tvUbTensor[REDUCE_UB_SIZE], tvUbTensor, numRowsRound, 0, 1, 1, 4); - AscendC::PipeBarrier(); - SetBlockReduceMask(ROW_OPS_SPEC_MASK_4); - AscendC::BlockReduceMax( - rowmaxUb, tvUbTensor[REDUCE_UB_SIZE], - (numRowsRound * FLOAT_BLOCK_SIZE + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, 0, 1, 1, 8); - AscendC::PipeBarrier(); - AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); - } - - CATLASS_DEVICE - void RowmaxTAILTILE(const AscendC::LocalTensor &srcUb, const AscendC::LocalTensor &rowmaxUb, - const AscendC::LocalTensor &tvUbTensor, uint32_t numRowsRound, uint32_t numElems, - uint32_t numElemsAligned) - { - if (numElems >= FLOAT_VECTOR_SIZE) { - AscendC::BlockReduceMax(tvUbTensor, srcUb, numRowsRound, 0, 1, 1, - numElemsAligned / FLOAT_BLOCK_SIZE); - AscendC::PipeBarrier(); - AscendC::BlockReduceMax( - rowmaxUb, tvUbTensor, (numRowsRound * FLOAT_BLOCK_SIZE + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, 0, - 1, 1, 8); - AscendC::PipeBarrier(); - for (uint64_t rowmax_idx = 1; rowmax_idx < (uint64_t)numElems / FLOAT_VECTOR_SIZE; ++rowmax_idx) { - AscendC::BlockReduceMax(tvUbTensor, srcUb[rowmax_idx * FLOAT_VECTOR_SIZE], numRowsRound, - 0, 1, 1, numElemsAligned / FLOAT_BLOCK_SIZE); - AscendC::PipeBarrier(); - AscendC::BlockReduceMax( - tvUbTensor[REDUCE_UB_SIZE], tvUbTensor, - (numRowsRound * FLOAT_BLOCK_SIZE + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, 0, 1, 1, 8); - AscendC::PipeBarrier(); - SetVecMask(numRowsRound); - AscendC::Max(rowmaxUb, rowmaxUb, tvUbTensor[REDUCE_UB_SIZE], (uint64_t)0, 1, - AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); - AscendC::PipeBarrier(); - AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); - } - } - if (numElems % FLOAT_VECTOR_SIZE > 0) { - SetVecMask(numElems % FLOAT_VECTOR_SIZE); - AscendC::BlockReduceMax(tvUbTensor, srcUb[numElems / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - numRowsRound, 0, 1, 1, numElemsAligned / FLOAT_BLOCK_SIZE); - AscendC::PipeBarrier(); - SetBlockReduceMask((numElems % FLOAT_VECTOR_SIZE + FLOAT_BLOCK_SIZE - 1) / FLOAT_BLOCK_SIZE); - if (numElems < FLOAT_VECTOR_SIZE) { - AscendC::BlockReduceMax( - rowmaxUb, tvUbTensor, (numRowsRound * FLOAT_BLOCK_SIZE + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, - 0, 1, 1, 8); - AscendC::PipeBarrier(); - } else { - AscendC::BlockReduceMax( - tvUbTensor[REDUCE_UB_SIZE], tvUbTensor, - (numRowsRound * FLOAT_BLOCK_SIZE + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, 0, 1, 1, 8); - AscendC::PipeBarrier(); - SetVecMask(numRowsRound); - AscendC::Max(rowmaxUb, rowmaxUb, tvUbTensor[REDUCE_UB_SIZE], (uint64_t)0, 1, - AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); - AscendC::PipeBarrier(); - } - AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); - } - } - - CATLASS_DEVICE - void SubCoreCompute(AscendC::GlobalTensor gOutput, AscendC::GlobalTensor gInput, - uint32_t m, uint32_t nReal, uint32_t nStride, uint32_t pingpongFlag, uint32_t rowOffset, - uint32_t sUbOffset, uint32_t nIdx, uint32_t &glFlag) - { - uint32_t round_m = (m + FLOAT_BLOCK_SIZE - 1) / FLOAT_BLOCK_SIZE * FLOAT_BLOCK_SIZE; - AscendC::WaitFlag(pingpongFlag); - // input QK - AscendC::DataCopy(lsUbTensor[sUbOffset], gInput, AscendC::DataCopyParams(m, nStride / FLOAT_BLOCK_SIZE, 0, 0)); - - AscendC::SetFlag(pingpongFlag); - AscendC::WaitFlag(pingpongFlag); - - // *** ls = tor * ls - AscendC::Muls(lsUbTensor[sUbOffset], lsUbTensor[sUbOffset], tor, (uint64_t)0, - (m * nStride + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, - AscendC::UnaryRepeatParams(1, 1, 8, 8)); - - AscendC::PipeBarrier(); - - if (kvSplitCoreNum != 1) { - if (nIdx == 0) { - if (glFlag == 1) { - AscendC::WaitFlag(EVENT_ID2); - glFlag = 0; - } - } - } - - if (nReal == 512) { - RowmaxSPECTILE512(lsUbTensor[sUbOffset], lmUbTensor[rowOffset], tvUbTensor, round_m, nReal, nStride); - } else if (nReal == 256) { - RowmaxSPECTILE256(lsUbTensor[sUbOffset], lmUbTensor[rowOffset], tvUbTensor, round_m, nReal, nStride); - } else { - RowmaxTAILTILE(lsUbTensor[sUbOffset], lmUbTensor[rowOffset], tvUbTensor, round_m, nReal, nStride); - } - - if (nIdx == 0) { - AscendC::DataCopy(hmUbTensor[rowOffset], lmUbTensor[rowOffset], - AscendC::DataCopyParams(1, round_m / FLOAT_BLOCK_SIZE, 0, 0)); - AscendC::PipeBarrier(); - } else { - SetVecMask(m); - // *** hm = vmax(lm, gm) - AscendC::Max(hmUbTensor[rowOffset], lmUbTensor[rowOffset], gmUbTensor[rowOffset], (uint64_t)0, - 1, AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); - - AscendC::PipeBarrier(); - // *** dm = gm - hm - AscendC::Sub(dmUbTensor[((nIdx / S_BLOCK_STACK) % 2) * UB_FLOAT_LINE_SIZE + rowOffset], - gmUbTensor[rowOffset], hmUbTensor[rowOffset], (uint64_t)0, 1, - AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); - - AscendC::PipeBarrier(); - // *** dm = exp(dm) - AscendC::Exp(dmUbTensor[((nIdx / S_BLOCK_STACK) % 2) * UB_FLOAT_LINE_SIZE + rowOffset], - dmUbTensor[((nIdx / S_BLOCK_STACK) % 2) * UB_FLOAT_LINE_SIZE + rowOffset], - (uint64_t)0, 1, AscendC::UnaryRepeatParams(1, 1, 8, 8)); - } - AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); - AscendC::PipeBarrier(); - // *** gm = hm - AscendC::DataCopy(gmUbTensor[rowOffset], hmUbTensor[rowOffset], - AscendC::DataCopyParams(1, round_m / FLOAT_BLOCK_SIZE, 0, 0)); - AscendC::PipeBarrier(); - // *** hm_block = expand_to_block(hm), 存放于 tv - AscendC::Brcb(tvUbTensor.template ReinterpretCast(), - hmUbTensor[rowOffset].template ReinterpretCast(), round_m / FLOAT_BLOCK_SIZE, - AscendC::BrcbRepeatParams(1, 8)); - AscendC::PipeBarrier(); - // *** ls = ls - hm_block - for (uint32_t subIdx = 0; subIdx < nReal / FLOAT_VECTOR_SIZE; ++subIdx) { - AscendC::Sub( - lsUbTensor[sUbOffset][subIdx * FLOAT_VECTOR_SIZE], lsUbTensor[sUbOffset][subIdx * FLOAT_VECTOR_SIZE], - tvUbTensor, (uint64_t)0, m, - AscendC::BinaryRepeatParams(1, 1, 0, nStride / FLOAT_BLOCK_SIZE, nStride / FLOAT_BLOCK_SIZE, 1)); - } - if (nReal % FLOAT_VECTOR_SIZE > 0) { - SetVecMask(nReal % FLOAT_VECTOR_SIZE); - AscendC::Sub( - lsUbTensor[sUbOffset][nReal / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - lsUbTensor[sUbOffset][nReal / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], tvUbTensor, (uint64_t)0, m, - AscendC::BinaryRepeatParams(1, 1, 0, nStride / FLOAT_BLOCK_SIZE, nStride / FLOAT_BLOCK_SIZE, 1)); - AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); - } - AscendC::PipeBarrier(); - - // *** ls = exp(ls) - AscendC::Exp(lsUbTensor[sUbOffset], lsUbTensor[sUbOffset], (uint64_t)0, - (m * nStride + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, - AscendC::UnaryRepeatParams(1, 1, 8, 8)); - AscendC::PipeBarrier(); - // *** ll = rowsum(ls32) - if (nReal == 512) { - RowsumSPECTILE512(lsUbTensor[sUbOffset], llUbTensor[rowOffset], tvUbTensor, round_m, nReal, nStride); - } else if (nReal == 256) { - RowsumSPECTILE256(lsUbTensor[sUbOffset], llUbTensor[rowOffset], tvUbTensor, round_m, nReal, nStride); - } else { - RowsumTAILTILE(lsUbTensor[sUbOffset], llUbTensor[rowOffset], tvUbTensor, round_m, nReal, nStride); - } - - // *** lp = castfp32to16(ls) - if (std::is_same::value) { - AscendC::Cast( - lpUbTensor[sUbOffset * 2], lsUbTensor[sUbOffset], AscendC::RoundMode::CAST_RINT, (uint64_t)0, - (m * nStride + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, AscendC::UnaryRepeatParams(1, 1, 4, 8)); - } else { - AscendC::Cast( - lpUbTensor[sUbOffset * 2], lsUbTensor[sUbOffset], AscendC::RoundMode::CAST_NONE, (uint64_t)0, - (m * nStride + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, AscendC::UnaryRepeatParams(1, 1, 4, 8)); - } - - AscendC::SetFlag(pingpongFlag); - AscendC::WaitFlag(pingpongFlag); - AscendC::DataCopy(gOutput, lpUbTensor[sUbOffset * 2], AscendC::DataCopyParams(m, nStride * 2 / 32, 0, 0)); - AscendC::SetFlag(pingpongFlag); - if (nIdx == 0) { - // *** gl = ll - AscendC::DataCopy(glUbTensor[rowOffset], llUbTensor[rowOffset], - AscendC::DataCopyParams(1, round_m / FLOAT_BLOCK_SIZE, 0, 0)); - AscendC::PipeBarrier(); - } else { - SetVecMask(m); - // *** gl = dm * gl - AscendC::Mul( - glUbTensor[rowOffset], dmUbTensor[((nIdx / S_BLOCK_STACK) % 2) * UB_FLOAT_LINE_SIZE + rowOffset], - glUbTensor[rowOffset], (uint64_t)0, 1, AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); - AscendC::PipeBarrier(); - // *** gl = ll + gl - AscendC::Add(glUbTensor[rowOffset], glUbTensor[rowOffset], llUbTensor[rowOffset], (uint64_t)0, - 1, AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); - AscendC::PipeBarrier(); - AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); - } - } - - CATLASS_DEVICE - void operator()(AscendC::GlobalTensor gOutput, AscendC::GlobalTensor gInput, - const LayoutOutput &layoutOutput, const LayoutInput &layoutInput, GemmCoord actualBlockShape, - uint32_t nIdx, uint32_t &glFlag) - { - uint32_t cur_head_num = actualBlockShape.m(); - uint32_t qkN = actualBlockShape.n(); - uint32_t qkRoundN = layoutInput.stride(0); - uint32_t pingpongFlag = 0; - - uint32_t subBlockIdx = AscendC::GetSubBlockIdx(); - uint32_t subBlockNum = AscendC::GetSubBlockNum(); - - uint32_t subM = (subBlockIdx == 1) ? (cur_head_num - cur_head_num / 2) : cur_head_num / 2; - - uint32_t mEnd = (subM + M_SLICE - 1) / M_SLICE; - - for (uint32_t mInd = 0; mInd < mEnd; mInd++) { - uint32_t rowOffset = mInd * M_SLICE; - uint32_t currM = mInd == mEnd - 1 ? subM - rowOffset : M_SLICE; - uint32_t sUbOffset = pingpongFlag * 8192; - int64_t offsetOutput = rowOffset * qkRoundN; - auto gOutputThisSubBlock = gOutput[offsetOutput]; - int64_t offsetInput = rowOffset * qkRoundN; - auto gInputThisSubBlock = gInput[offsetInput]; - if (mInd == 0) { - Arch::CrossCoreWaitFlag(qkReady); - } - if (currM == 0) { - continue; - } - SubCoreCompute(gOutputThisSubBlock, gInputThisSubBlock, currM, qkN, qkRoundN, pingpongFlag, rowOffset, - sUbOffset, nIdx, glFlag); - pingpongFlag = 1 - pingpongFlag; - } - } - -private: - float tor; - uint32_t pingpongFlag = 0; - uint32_t kvSplitCoreNum = 1; - - AscendC::LocalTensor lsUbTensor; - AscendC::LocalTensor lpUbTensor; - AscendC::LocalTensor lmUbTensor; - AscendC::LocalTensor hmUbTensor; - AscendC::LocalTensor gmUbTensor; - AscendC::LocalTensor dmUbTensor; - AscendC::LocalTensor llUbTensor; - AscendC::LocalTensor tvUbTensor; - AscendC::LocalTensor glUbTensor; - - Arch::CrossCoreFlag qkReady{QK_READY_ID}; -}; - -} // namespace Catlass::Epilogue::Block - -#endif // CATLASS_KDA_BLOCK_EPILOGUE_MLA_TP1_SOFTMAX_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/kda_block_epilogue_per_token_dequant.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/kda_block_epilogue_per_token_dequant.hpp deleted file mode 100644 index 53c7cd6636..0000000000 --- a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/kda_block_epilogue_per_token_dequant.hpp +++ /dev/null @@ -1,683 +0,0 @@ -/** - * Copyright (c) 2025 Huawei Technologies Co., Ltd. - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of - * CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ - -#ifndef CATLASS_KDA_BLOCK_EPILOGUE_PER_TOKEN_DEQUANT_HPP -#define CATLASS_KDA_BLOCK_EPILOGUE_PER_TOKEN_DEQUANT_HPP - -#include "catlass/kda_catlass.hpp" -#include "catlass/arch/kda_resource.hpp" -#include "catlass/epilogue/kda_dispatch_policy.hpp" -#include "catlass/kda_gemm_coord.hpp" -#include "catlass/gemm/kda_gemm_type.hpp" -#include "catlass/kda_matrix_coord.hpp" -#include "catlass/layout/kda_layout.hpp" -#include "catlass/detail/kda_callback.hpp" - -namespace Catlass::Epilogue::Block { - -template < - uint32_t UB_STAGES_, - class CType_, - class ScaleType_, - class PerTokenScaleType_, - class DType_, - class TileRowBroadcastMul_, - class TileBroadcastOneBlk_, - class TileOneBlkColumnBroadcastMul_, - class TileCopy_, - class EpilogueTileSwizzle_> -class BlockEpilogue< - EpilogueAtlasA2PerTokenDequant, - CType_, - ScaleType_, - PerTokenScaleType_, - DType_, - TileRowBroadcastMul_, - TileBroadcastOneBlk_, - TileOneBlkColumnBroadcastMul_, - TileCopy_, - EpilogueTileSwizzle_> { -public: - using DispatchPolicy = EpilogueAtlasA2PerTokenDequant; - using ArchTag = typename DispatchPolicy::ArchTag; - static constexpr uint32_t UB_STAGES = UB_STAGES_; - - // Data infos - using ElementC = typename CType_::Element; - using LayoutC = typename CType_::Layout; - using ElementScale = typename ScaleType_::Element; - using LayoutScale = typename ScaleType_::Layout; - using ElementPerTokenScale = typename PerTokenScaleType_::Element; - using LayoutPerTokenScale = typename PerTokenScaleType_::Layout; - using ElementD = typename DType_::Element; - using LayoutD = typename DType_::Layout; - - // Check data infos - static_assert( - std::is_same_v && (std::is_same_v || std::is_same_v) && - std::is_same_v && std::is_same_v, - "The element type template parameters of BlockEpilogue are wrong"); - static_assert( - std::is_same_v && std::is_same_v && - std::is_same_v && std::is_same_v, - "The layout template parameters of BlockEpilogue are wrong"); - - // Tile compute ops - using TileRowBroadcastMul = TileRowBroadcastMul_; - using TileBroadcastOneBlk = TileBroadcastOneBlk_; - using TileOneBlkColumnBroadcastMul = TileOneBlkColumnBroadcastMul_; - - // Tile copy - using CopyGmToUbC = typename TileCopy_::CopyGmToUbC; - using CopyGmToUbScale = typename TileCopy_::CopyGmToUbX; - using CopyGmToUbPerTokenScale = typename TileCopy_::CopyGmToUbY; - using CopyUbToGmD = typename TileCopy_::CopyUbToGmD; - - using EpilogueTileSwizzle = EpilogueTileSwizzle_; - - using TileShape = typename TileRowBroadcastMul::TileShape; - - static_assert( - TileShape::ROW == TileBroadcastOneBlk::COMPUTE_LENGTH && - std::is_same_v, - "TileShape must be consistent for all tile compute ops"); - - static_assert( - (UB_STAGES * (TileShape::COUNT * sizeof(ElementC) + TileShape::COLUMN * sizeof(ElementScale) + TileShape::ROW * sizeof(ElementPerTokenScale) + TileShape::COUNT * sizeof(ElementD)) + (TileShape::COUNT + TileShape::COLUMN + TileShape::COUNT + TileShape::ROW) * sizeof(float) + TileShape::ROW * BYTE_PER_BLK) <= ArchTag::UB_SIZE, - "TileShape is too large to fit in UB"); - - struct Params { - __gm__ ElementScale *ptrScale{nullptr}; - LayoutScale layoutScale{}; - __gm__ ElementPerTokenScale *ptrPerTokenScale{nullptr}; - LayoutPerTokenScale layoutPerTokenScale{}; - __gm__ ElementD *ptrD{nullptr}; - LayoutD layoutD{}; - - CATLASS_DEVICE - Params() {}; - - CATLASS_DEVICE - Params( - __gm__ ElementScale *ptrScale_, LayoutScale const &layoutScale_, - __gm__ ElementPerTokenScale *ptrPerTokenScale_, LayoutPerTokenScale const &layoutPerTokenScale_, - __gm__ ElementD *ptrD_, LayoutD const &layoutD_) - : ptrScale(ptrScale_), - layoutScale(layoutScale_), - ptrPerTokenScale(ptrPerTokenScale_), - layoutPerTokenScale(layoutPerTokenScale_), - ptrD(ptrD_), - layoutD(layoutD_) - {} - }; - - CATLASS_DEVICE - BlockEpilogue(Arch::Resource const &resource, Params const ¶ms = Params{}) - : params(params) - { - size_t ubOffset = 0; - int32_t eventVMTE2 = 0; - int32_t eventMTE2V = 0; - int32_t eventMTE3V = 0; - int32_t eventVMTE3 = 0; - for (uint32_t i = 0; i < UB_STAGES; ++i) { - ubCList[i] = resource.ubBuf.template GetBufferByByte(ubOffset); - ubOffset += TileShape::COUNT * sizeof(ElementC); - ubScaleList[i] = resource.ubBuf.template GetBufferByByte(ubOffset); - ubOffset += TileShape::COLUMN * sizeof(ElementScale); - ubPerTokenScaleList[i] = resource.ubBuf.template GetBufferByByte(ubOffset); - ubOffset += TileShape::ROW * sizeof(ElementPerTokenScale); - ubDList[i] = resource.ubBuf.template GetBufferByByte(ubOffset); - ubOffset += TileShape::COUNT * sizeof(ElementD); - - eventUbCVMTE2List[i] = eventVMTE2++; - eventUbCMTE2VList[i] = eventMTE2V++; - eventUbScaleVMTE2List[i] = eventVMTE2++; - eventUbScaleMTE2VList[i] = eventMTE2V++; - eventUbPerTokenScaleVMTE2List[i] = eventVMTE2++; - eventUbPerTokenScaleMTE2VList[i] = eventMTE2V++; - eventUbDMTE3VList[i] = eventMTE3V++; - eventUbDVMTE3List[i] = eventVMTE3++; - - AscendC::SetFlag(eventUbCVMTE2List[i]); - AscendC::SetFlag(eventUbScaleVMTE2List[i]); - AscendC::SetFlag(eventUbPerTokenScaleVMTE2List[i]); - AscendC::SetFlag(eventUbDMTE3VList[i]); - } - ubCFp32 = resource.ubBuf.template GetBufferByByte(ubOffset); - ubOffset += TileShape::COUNT * sizeof(float); - ubScaleFp32 = resource.ubBuf.template GetBufferByByte(ubOffset); - ubOffset += TileShape::COLUMN * sizeof(float); - ubMul = resource.ubBuf.template GetBufferByByte(ubOffset); - ubOffset += TileShape::COUNT * sizeof(float); - ubPerTokenScaleFp32 = resource.ubBuf.template GetBufferByByte(ubOffset); - ubOffset += TileShape::ROW * sizeof(float); - ubPerTokenScaleFp32Brcb = resource.ubBuf.template GetBufferByByte(ubOffset); - ubOffset += TileShape::ROW * BYTE_PER_BLK; - ubPerTokenMul = ubMul; - } - - CATLASS_DEVICE - ~BlockEpilogue() - { - for (uint32_t i = 0; i < UB_STAGES; ++i) { - AscendC::WaitFlag(eventUbCVMTE2List[i]); - AscendC::WaitFlag(eventUbScaleVMTE2List[i]); - AscendC::WaitFlag(eventUbPerTokenScaleVMTE2List[i]); - AscendC::WaitFlag(eventUbDMTE3VList[i]); - } - } - - CATLASS_DEVICE - void UpdateParams(Params const ¶ms_) - { - params = params_; - } - - CATLASS_DEVICE - void operator()( - GemmCoord const &blockShapeMNK, - GemmCoord const &blockCoordMNK, - GemmCoord const &actualBlockShapeMNK, - AscendC::GlobalTensor const &gmBlockC, - LayoutC const &layoutBlockC, Callback &&callback = Callback{}) - { - MatrixCoord blockShape = blockShapeMNK.GetCoordMN(); - MatrixCoord blockCoord = blockCoordMNK.GetCoordMN(); - MatrixCoord actualBlockShape = actualBlockShapeMNK.GetCoordMN(); - MatrixCoord blockOffset = blockCoord * blockShape; - - auto ubTileStride = MakeCoord(static_cast(TileShape::COLUMN), 1L); - auto tileShape = TileShape::ToCoord(); - EpilogueTileSwizzle epilogueTileSwizzle(actualBlockShape, tileShape); - uint32_t tileLoops = epilogueTileSwizzle.GetLoops(); - uint32_t subblockIdx = AscendC::GetSubBlockIdx(); - uint32_t subblockNum = AscendC::GetSubBlockNum(); - - if (actualBlockShapeMNK.k() == 0) { - AscendC::GlobalTensor gmD; - gmD.SetGlobalBuffer(params.ptrD); - - ElementD zeroD{}; - for (uint32_t loopIdx = subblockIdx; loopIdx < tileLoops; loopIdx += subblockNum) { - auto tileCoord = epilogueTileSwizzle.GetTileCoord(loopIdx); - auto actualTileShape = epilogueTileSwizzle.GetActualTileShape(tileCoord); - auto tileOffset = blockOffset + tileCoord * tileShape; - - auto &ubD = ubDList[ubListId]; - LayoutD layoutUbD{actualTileShape, ubTileStride}; - - AscendC::WaitFlag(eventUbDMTE3VList[ubListId]); - AscendC::Duplicate(ubD, zeroD, TileShape::COUNT); - AscendC::SetFlag(eventUbDVMTE3List[ubListId]); - - auto gmTileD = gmD[params.layoutD.GetOffset(tileOffset)]; - auto layoutGmTileD = params.layoutD.GetTileLayout(actualTileShape); - - AscendC::WaitFlag(eventUbDVMTE3List[ubListId]); - copyUbToGmD(gmTileD, ubD, layoutGmTileD, layoutUbD); - AscendC::SetFlag(eventUbDMTE3VList[ubListId]); - - ubListId = (ubListId + 1 < UB_STAGES) ? (ubListId + 1) : 0; - } - return; - } - callback(); - - AscendC::GlobalTensor gmScale; - gmScale.SetGlobalBuffer(params.ptrScale); - AscendC::GlobalTensor gmPerTokenScale; - gmPerTokenScale.SetGlobalBuffer(params.ptrPerTokenScale); - AscendC::GlobalTensor gmD; - gmD.SetGlobalBuffer(params.ptrD); - - for (uint32_t loopIdx = subblockIdx; loopIdx < tileLoops; loopIdx += subblockNum) { - auto tileCoord = epilogueTileSwizzle.GetTileCoord(loopIdx); - auto actualTileShape = epilogueTileSwizzle.GetActualTileShape(tileCoord); - auto tileOffsetInBlock = tileCoord * tileShape; - auto tileOffset = blockOffset + tileOffsetInBlock; - - auto gmTileC = gmBlockC[layoutBlockC.GetOffset(tileOffsetInBlock)]; - auto layoutGmTileC = layoutBlockC.GetTileLayout(actualTileShape); - - auto &ubC = ubCList[ubListId]; - LayoutC layoutUbC{actualTileShape, ubTileStride}; - - AscendC::WaitFlag(eventUbCVMTE2List[ubListId]); - copyGmToUbC(ubC, gmTileC, layoutUbC, layoutGmTileC); - AscendC::SetFlag(eventUbCMTE2VList[ubListId]); - - auto scaleTileOffset = tileOffset.template GetCoordByAxis<1>(); - auto scaleTileShape = actualTileShape.template GetCoordByAxis<1>(); - - auto gmTileScale = gmScale[params.layoutScale.GetOffset(scaleTileOffset)]; - auto layoutGmTileScale = params.layoutScale.GetTileLayout(scaleTileShape); - - auto &ubScale = ubScaleList[ubListId]; - auto layoutUbScale = LayoutScale::template MakeLayoutInUb(scaleTileShape); - - AscendC::WaitFlag(eventUbScaleVMTE2List[ubListId]); - copyGmToUbScale(ubScale, gmTileScale, layoutUbScale, layoutGmTileScale); - AscendC::SetFlag(eventUbScaleMTE2VList[ubListId]); - - auto perTokenScaleTileOffset = tileOffset.template GetCoordByAxis<0>(); - auto perTokenScaleTileShape = actualTileShape.template GetCoordByAxis<0>(); - - auto gmTilePerTokenScale = gmPerTokenScale[params.layoutPerTokenScale.GetOffset(perTokenScaleTileOffset)]; - auto layoutGmTilePerTokenScale = params.layoutPerTokenScale.GetTileLayout(perTokenScaleTileShape); - - auto &ubPerTokenScale = ubPerTokenScaleList[ubListId]; - auto layoutUbPerTokenScale = LayoutScale::template MakeLayoutInUb( - perTokenScaleTileShape); - - AscendC::WaitFlag(eventUbPerTokenScaleVMTE2List[ubListId]); - copyGmToUbPerTokenScale(ubPerTokenScale, gmTilePerTokenScale, layoutUbPerTokenScale, - layoutGmTilePerTokenScale); - AscendC::SetFlag(eventUbPerTokenScaleMTE2VList[ubListId]); - - AscendC::WaitFlag(eventUbCMTE2VList[ubListId]); - AscendC::Cast(ubCFp32, ubC, AscendC::RoundMode::CAST_RINT, TileShape::COUNT); - AscendC::SetFlag(eventUbCVMTE2List[ubListId]); - - AscendC::WaitFlag(eventUbScaleMTE2VList[ubListId]); - AscendC::Cast(ubScaleFp32, ubScale, AscendC::RoundMode::CAST_NONE, TileShape::COLUMN); - AscendC::SetFlag(eventUbScaleVMTE2List[ubListId]); - - AscendC::WaitFlag(eventUbPerTokenScaleMTE2VList[ubListId]); - AscendC::Cast(ubPerTokenScaleFp32, ubPerTokenScale, AscendC::RoundMode::CAST_NONE, TileShape::ROW); - AscendC::SetFlag(eventUbPerTokenScaleVMTE2List[ubListId]); - - AscendC::PipeBarrier(); - tileRowBroadcastMul(ubMul, ubCFp32, ubScaleFp32); - tileBroadcastOneBlk(ubPerTokenScaleFp32Brcb, ubPerTokenScaleFp32); - AscendC::PipeBarrier(); - tileOneBlkColumnBroadcastMul(ubPerTokenMul, ubMul, ubPerTokenScaleFp32Brcb); - AscendC::PipeBarrier(); - - auto &ubD = ubDList[ubListId]; - LayoutD layoutUbD{actualTileShape, ubTileStride}; - - AscendC::WaitFlag(eventUbDMTE3VList[ubListId]); - AscendC::Cast(ubD, ubPerTokenMul, AscendC::RoundMode::CAST_RINT, TileShape::COUNT); - AscendC::SetFlag(eventUbDVMTE3List[ubListId]); - - auto gmTileD = gmD[params.layoutD.GetOffset(tileOffset)]; - auto layoutGmTileD = params.layoutD.GetTileLayout(actualTileShape); - - AscendC::WaitFlag(eventUbDVMTE3List[ubListId]); - copyUbToGmD(gmTileD, ubD, layoutGmTileD, layoutUbD); - AscendC::SetFlag(eventUbDMTE3VList[ubListId]); - - ubListId = (ubListId + 1 < UB_STAGES) ? (ubListId + 1) : 0; - } - } - -private: - Params params; - - AscendC::LocalTensor ubCList[UB_STAGES]; - AscendC::LocalTensor ubScaleList[UB_STAGES]; - AscendC::LocalTensor ubPerTokenScaleList[UB_STAGES]; - AscendC::LocalTensor ubDList[UB_STAGES]; - - int32_t eventUbCVMTE2List[UB_STAGES]; - int32_t eventUbCMTE2VList[UB_STAGES]; - int32_t eventUbScaleVMTE2List[UB_STAGES]; - int32_t eventUbScaleMTE2VList[UB_STAGES]; - int32_t eventUbPerTokenScaleVMTE2List[UB_STAGES]; - int32_t eventUbPerTokenScaleMTE2VList[UB_STAGES]; - int32_t eventUbDMTE3VList[UB_STAGES]; - int32_t eventUbDVMTE3List[UB_STAGES]; - - uint32_t ubListId{0}; - - AscendC::LocalTensor ubCFp32; - AscendC::LocalTensor ubScaleFp32; - AscendC::LocalTensor ubMul; - AscendC::LocalTensor ubPerTokenScaleFp32; - AscendC::LocalTensor ubPerTokenScaleFp32Brcb; - AscendC::LocalTensor ubPerTokenMul; - - TileRowBroadcastMul tileRowBroadcastMul; - TileBroadcastOneBlk tileBroadcastOneBlk; - TileOneBlkColumnBroadcastMul tileOneBlkColumnBroadcastMul; - - CopyGmToUbC copyGmToUbC; - CopyGmToUbScale copyGmToUbScale; - CopyGmToUbPerTokenScale copyGmToUbPerTokenScale; - CopyUbToGmD copyUbToGmD; -}; - -template < - uint32_t UB_STAGES_, - class CType_, - class LayoutScale_, - class LayoutPerTokenScale_, - class DType_, - class TileRowBroadcastMul_, - class TileBroadcastOneBlk_, - class TileOneBlkColumnBroadcastMul_, - class TileCopy_, - class EpilogueTileSwizzle_> -class BlockEpilogue< - EpilogueAtlasA2PerTokenDequant, - CType_, - Gemm::GemmType, - Gemm::GemmType, - DType_, - TileRowBroadcastMul_, - TileBroadcastOneBlk_, - TileOneBlkColumnBroadcastMul_, - TileCopy_, - EpilogueTileSwizzle_> { -public: - using DispatchPolicy = EpilogueAtlasA2PerTokenDequant; - using ArchTag = typename DispatchPolicy::ArchTag; - static constexpr uint32_t UB_STAGES = UB_STAGES_; - - // Data infos - using ElementC = typename CType_::Element; - using LayoutC = typename CType_::Layout; - using ElementScale = float; - using LayoutScale = LayoutScale_; - using ElementPerTokenScale = float; - using LayoutPerTokenScale = LayoutPerTokenScale_; - using ElementD = typename DType_::Element; - using LayoutD = typename DType_::Layout; - - // Check data infos - static_assert( - std::is_same_v && (std::is_same_v || std::is_same_v), - "The element type template parameters of BlockEpilogue are wrong"); - static_assert( - std::is_same_v && std::is_same_v && - std::is_same_v && std::is_same_v, - "The layout template parameters of BlockEpilogue are wrong"); - - // Tile compute ops - using TileRowBroadcastMul = TileRowBroadcastMul_; - using TileBroadcastOneBlk = TileBroadcastOneBlk_; - using TileOneBlkColumnBroadcastMul = TileOneBlkColumnBroadcastMul_; - - // Tile copy - using CopyGmToUbC = typename TileCopy_::CopyGmToUbC; - using CopyGmToUbScale = typename TileCopy_::CopyGmToUbX; - using CopyGmToUbPerTokenScale = typename TileCopy_::CopyGmToUbY; - using CopyUbToGmD = typename TileCopy_::CopyUbToGmD; - - using EpilogueTileSwizzle = EpilogueTileSwizzle_; - - using TileShape = typename TileRowBroadcastMul::TileShape; - - static_assert( - TileShape::ROW == TileBroadcastOneBlk::COMPUTE_LENGTH && - std::is_same_v, - "TileShape must be consistent for all tile compute ops"); - - static_assert( - (UB_STAGES * (TileShape::COUNT * sizeof(ElementC) + TileShape::COLUMN * sizeof(ElementScale) + TileShape::ROW * sizeof(ElementPerTokenScale) + TileShape::COUNT * sizeof(ElementD)) + (TileShape::COUNT + TileShape::COUNT) * sizeof(float) + TileShape::ROW * BYTE_PER_BLK) <= ArchTag::UB_SIZE, - "TileShape is too large to fit in UB"); - - struct Params { - __gm__ ElementScale *ptrScale{nullptr}; - LayoutScale layoutScale{}; - __gm__ ElementPerTokenScale *ptrPerTokenScale{nullptr}; - LayoutPerTokenScale layoutPerTokenScale{}; - __gm__ ElementD *ptrD{nullptr}; - LayoutD layoutD{}; - - CATLASS_DEVICE - Params() {}; - - CATLASS_DEVICE - Params( - __gm__ ElementScale *ptrScale_, LayoutScale const &layoutScale_, - __gm__ ElementPerTokenScale *ptrPerTokenScale_, LayoutPerTokenScale const &layoutPerTokenScale_, - __gm__ ElementD *ptrD_, LayoutD const &layoutD_) - : ptrScale(ptrScale_), - layoutScale(layoutScale_), - ptrPerTokenScale(ptrPerTokenScale_), - layoutPerTokenScale(layoutPerTokenScale_), - ptrD(ptrD_), - layoutD(layoutD_) - {} - }; - - CATLASS_DEVICE - BlockEpilogue(Arch::Resource const &resource, Params const ¶ms = Params{}) - : params(params) - { - size_t ubOffset = 0; - int32_t eventVMTE2 = 0; - int32_t eventMTE2V = 0; - int32_t eventMTE3V = 0; - int32_t eventVMTE3 = 0; - for (uint32_t i = 0; i < UB_STAGES; ++i) { - ubCList[i] = resource.ubBuf.template GetBufferByByte(ubOffset); - ubOffset += TileShape::COUNT * sizeof(ElementC); - ubScaleList[i] = resource.ubBuf.template GetBufferByByte(ubOffset); - ubOffset += TileShape::COLUMN * sizeof(ElementScale); - ubPerTokenScaleList[i] = resource.ubBuf.template GetBufferByByte(ubOffset); - ubOffset += TileShape::ROW * sizeof(ElementPerTokenScale); - ubDList[i] = resource.ubBuf.template GetBufferByByte(ubOffset); - ubOffset += TileShape::COUNT * sizeof(ElementD); - - eventUbCVMTE2List[i] = eventVMTE2++; - eventUbCMTE2VList[i] = eventMTE2V++; - eventUbScaleVMTE2List[i] = eventVMTE2++; - eventUbScaleMTE2VList[i] = eventMTE2V++; - eventUbPerTokenScaleVMTE2List[i] = eventVMTE2++; - eventUbPerTokenScaleMTE2VList[i] = eventMTE2V++; - eventUbDMTE3VList[i] = eventMTE3V++; - eventUbDVMTE3List[i] = eventVMTE3++; - - AscendC::SetFlag(eventUbCVMTE2List[i]); - AscendC::SetFlag(eventUbScaleVMTE2List[i]); - AscendC::SetFlag(eventUbPerTokenScaleVMTE2List[i]); - AscendC::SetFlag(eventUbDMTE3VList[i]); - } - ubCFp32 = resource.ubBuf.template GetBufferByByte(ubOffset); - ubOffset += TileShape::COUNT * sizeof(float); - ubMul = resource.ubBuf.template GetBufferByByte(ubOffset); - ubOffset += TileShape::COUNT * sizeof(float); - ubPerTokenScaleBrcb = resource.ubBuf.template GetBufferByByte(ubOffset); - ubOffset += TileShape::ROW * BYTE_PER_BLK; - ubPerTokenMul = ubCFp32; - } - - CATLASS_DEVICE - ~BlockEpilogue() - { - for (uint32_t i = 0; i < UB_STAGES; ++i) { - AscendC::WaitFlag(eventUbCVMTE2List[i]); - AscendC::WaitFlag(eventUbScaleVMTE2List[i]); - AscendC::WaitFlag(eventUbPerTokenScaleVMTE2List[i]); - AscendC::WaitFlag(eventUbDMTE3VList[i]); - } - } - - CATLASS_DEVICE - void UpdateParams(Params const ¶ms_) - { - params = params_; - } - - CATLASS_DEVICE - void operator()( - GemmCoord const &blockShapeMNK, - GemmCoord const &blockCoordMNK, - GemmCoord const &actualBlockShapeMNK, - AscendC::GlobalTensor const &gmBlockC, - LayoutC const &layoutBlockC, Callback &&callback = Callback{}) - { - MatrixCoord blockShape = blockShapeMNK.GetCoordMN(); - MatrixCoord blockCoord = blockCoordMNK.GetCoordMN(); - MatrixCoord actualBlockShape = actualBlockShapeMNK.GetCoordMN(); - MatrixCoord blockOffset = blockCoord * blockShape; - - auto ubTileStride = MakeCoord(static_cast(TileShape::COLUMN), 1L); - auto tileShape = TileShape::ToCoord(); - EpilogueTileSwizzle epilogueTileSwizzle(actualBlockShape, tileShape); - uint32_t tileLoops = epilogueTileSwizzle.GetLoops(); - uint32_t subblockIdx = AscendC::GetSubBlockIdx(); - uint32_t subblockNum = AscendC::GetSubBlockNum(); - - if (actualBlockShapeMNK.k() == 0) { - AscendC::GlobalTensor gmD; - gmD.SetGlobalBuffer(params.ptrD); - - ElementD zeroD{}; - for (uint32_t loopIdx = subblockIdx; loopIdx < tileLoops; loopIdx += subblockNum) { - auto tileCoord = epilogueTileSwizzle.GetTileCoord(loopIdx); - auto actualTileShape = epilogueTileSwizzle.GetActualTileShape(tileCoord); - auto tileOffset = blockOffset + tileCoord * tileShape; - - auto &ubD = ubDList[ubListId]; - LayoutD layoutUbD{actualTileShape, ubTileStride}; - - AscendC::WaitFlag(eventUbDMTE3VList[ubListId]); - AscendC::Duplicate(ubD, zeroD, TileShape::COUNT); - AscendC::SetFlag(eventUbDVMTE3List[ubListId]); - - auto gmTileD = gmD[params.layoutD.GetOffset(tileOffset)]; - auto layoutGmTileD = params.layoutD.GetTileLayout(actualTileShape); - - AscendC::WaitFlag(eventUbDVMTE3List[ubListId]); - copyUbToGmD(gmTileD, ubD, layoutGmTileD, layoutUbD); - AscendC::SetFlag(eventUbDMTE3VList[ubListId]); - - ubListId = (ubListId + 1 < UB_STAGES) ? (ubListId + 1) : 0; - } - return; - } - callback(); - - AscendC::GlobalTensor gmScale; - gmScale.SetGlobalBuffer(params.ptrScale); - AscendC::GlobalTensor gmPerTokenScale; - gmPerTokenScale.SetGlobalBuffer(params.ptrPerTokenScale); - AscendC::GlobalTensor gmD; - gmD.SetGlobalBuffer(params.ptrD); - - for (uint32_t loopIdx = subblockIdx; loopIdx < tileLoops; loopIdx += subblockNum) { - auto tileCoord = epilogueTileSwizzle.GetTileCoord(loopIdx); - auto actualTileShape = epilogueTileSwizzle.GetActualTileShape(tileCoord); - auto tileOffsetInBlock = tileCoord * tileShape; - auto tileOffset = blockOffset + tileOffsetInBlock; - - auto gmTileC = gmBlockC[layoutBlockC.GetOffset(tileOffsetInBlock)]; - auto layoutGmTileC = layoutBlockC.GetTileLayout(actualTileShape); - - auto &ubC = ubCList[ubListId]; - LayoutC layoutUbC{actualTileShape, ubTileStride}; - - AscendC::WaitFlag(eventUbCVMTE2List[ubListId]); - copyGmToUbC(ubC, gmTileC, layoutUbC, layoutGmTileC); - AscendC::SetFlag(eventUbCMTE2VList[ubListId]); - - auto scaleTileOffset = tileOffset.template GetCoordByAxis<1>(); - auto scaleTileShape = actualTileShape.template GetCoordByAxis<1>(); - - auto gmTileScale = gmScale[params.layoutScale.GetOffset(scaleTileOffset)]; - auto layoutGmTileScale = params.layoutScale.GetTileLayout(scaleTileShape); - - auto &ubScale = ubScaleList[ubListId]; - auto layoutUbScale = LayoutScale::template MakeLayoutInUb(scaleTileShape); - - AscendC::WaitFlag(eventUbScaleVMTE2List[ubListId]); - copyGmToUbScale(ubScale, gmTileScale, layoutUbScale, layoutGmTileScale); - AscendC::SetFlag(eventUbScaleMTE2VList[ubListId]); - - auto perTokenScaleTileOffset = tileOffset.template GetCoordByAxis<0>(); - auto perTokenScaleTileShape = actualTileShape.template GetCoordByAxis<0>(); - - auto gmTilePerTokenScale = gmPerTokenScale[params.layoutPerTokenScale.GetOffset(perTokenScaleTileOffset)]; - auto layoutGmTilePerTokenScale = params.layoutPerTokenScale.GetTileLayout(perTokenScaleTileShape); - - auto &ubPerTokenScale = ubPerTokenScaleList[ubListId]; - auto layoutUbPerTokenScale = LayoutScale::template MakeLayoutInUb( - perTokenScaleTileShape); - - AscendC::WaitFlag(eventUbPerTokenScaleVMTE2List[ubListId]); - copyGmToUbPerTokenScale(ubPerTokenScale, gmTilePerTokenScale, layoutUbPerTokenScale, - layoutGmTilePerTokenScale); - AscendC::SetFlag(eventUbPerTokenScaleMTE2VList[ubListId]); - - AscendC::WaitFlag(eventUbCMTE2VList[ubListId]); - AscendC::Cast(ubCFp32, ubC, AscendC::RoundMode::CAST_RINT, TileShape::COUNT); - AscendC::SetFlag(eventUbCVMTE2List[ubListId]); - - AscendC::WaitFlag(eventUbScaleMTE2VList[ubListId]); - tileRowBroadcastMul(ubMul, ubCFp32, ubScale); - AscendC::SetFlag(eventUbScaleVMTE2List[ubListId]); - - AscendC::WaitFlag(eventUbPerTokenScaleMTE2VList[ubListId]); - tileBroadcastOneBlk(ubPerTokenScaleBrcb, ubPerTokenScale); - AscendC::SetFlag(eventUbPerTokenScaleVMTE2List[ubListId]); - - AscendC::PipeBarrier(); - tileOneBlkColumnBroadcastMul(ubPerTokenMul, ubMul, ubPerTokenScaleBrcb); - AscendC::PipeBarrier(); - - auto &ubD = ubDList[ubListId]; - LayoutD layoutUbD{actualTileShape, ubTileStride}; - - AscendC::WaitFlag(eventUbDMTE3VList[ubListId]); - AscendC::Cast(ubD, ubPerTokenMul, AscendC::RoundMode::CAST_RINT, TileShape::COUNT); - AscendC::SetFlag(eventUbDVMTE3List[ubListId]); - - auto gmTileD = gmD[params.layoutD.GetOffset(tileOffset)]; - auto layoutGmTileD = params.layoutD.GetTileLayout(actualTileShape); - - AscendC::WaitFlag(eventUbDVMTE3List[ubListId]); - copyUbToGmD(gmTileD, ubD, layoutGmTileD, layoutUbD); - AscendC::SetFlag(eventUbDMTE3VList[ubListId]); - - ubListId = (ubListId + 1 < UB_STAGES) ? (ubListId + 1) : 0; - } - } - -private: - Params params; - - AscendC::LocalTensor ubCList[UB_STAGES]; - AscendC::LocalTensor ubScaleList[UB_STAGES]; - AscendC::LocalTensor ubPerTokenScaleList[UB_STAGES]; - AscendC::LocalTensor ubDList[UB_STAGES]; - - int32_t eventUbCVMTE2List[UB_STAGES]; - int32_t eventUbCMTE2VList[UB_STAGES]; - int32_t eventUbScaleVMTE2List[UB_STAGES]; - int32_t eventUbScaleMTE2VList[UB_STAGES]; - int32_t eventUbPerTokenScaleVMTE2List[UB_STAGES]; - int32_t eventUbPerTokenScaleMTE2VList[UB_STAGES]; - int32_t eventUbDMTE3VList[UB_STAGES]; - int32_t eventUbDVMTE3List[UB_STAGES]; - - uint32_t ubListId{0}; - - AscendC::LocalTensor ubCFp32; - AscendC::LocalTensor ubMul; - AscendC::LocalTensor ubPerTokenScaleBrcb; - AscendC::LocalTensor ubPerTokenMul; - - TileRowBroadcastMul tileRowBroadcastMul; - TileBroadcastOneBlk tileBroadcastOneBlk; - TileOneBlkColumnBroadcastMul tileOneBlkColumnBroadcastMul; - - CopyGmToUbC copyGmToUbC; - CopyGmToUbScale copyGmToUbScale; - CopyGmToUbPerTokenScale copyGmToUbPerTokenScale; - CopyUbToGmD copyUbToGmD; -}; - -} // namespace Catlass::Epilogue::Block - -#endif // CATLASS_KDA_BLOCK_EPILOGUE_PER_TOKEN_DEQUANT_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/kda_dispatch_policy.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/kda_dispatch_policy.hpp deleted file mode 100644 index f417f91865..0000000000 --- a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/kda_dispatch_policy.hpp +++ /dev/null @@ -1,143 +0,0 @@ -/** - * Copyright (c) 2025 Huawei Technologies Co., Ltd. - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of - * CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ - -#ifndef CATLASS_KDA_DISPATCH_POLICY_HPP -#define CATLASS_KDA_DISPATCH_POLICY_HPP - -#include "catlass/kda_catlass.hpp" -#include "catlass/arch/kda_arch.hpp" - -namespace Catlass::Epilogue { - -// For AtlasA2, an element wise epilogue of the form D = C + X, where X is an additional source -struct EpilogueAtlasA2ElemWiseOneSource { - using ArchTag = Arch::AtlasA2; - // Number of operands. Including C, X, and D 3 operands - static constexpr uint32_t OPERANDS_NUM = 3; -}; - -struct EpilogueAtlasA2ElemWiseNoSource { - using ArchTag = Arch::AtlasA2; - // Number of operands. Including C, D 2 operands - static constexpr uint32_t OPERANDS_NUM = 2; -}; - -// For AtlasA2, FA Softmax -struct EpilogueAtlasA2FASoftmax { - using ArchTag = Arch::AtlasA2; -}; - -// For AtlasA2, FA RescaleO -struct EpilogueAtlasA2FARescaleO { - using ArchTag = Arch::AtlasA2; -}; - -// For AtlasA2, MLA Softmax -struct EpilogueAtlasA2MLASoftmax { - using ArchTag = Arch::AtlasA2; -}; - -// For AtlasA2, FA Infer online Softmax no mask -struct EpilogueAtlasA2OnlineSoftmax { - using ArchTag = Arch::AtlasA2; -}; - -// For AtlasA2, FA Infer RescaleO no split row -struct EpilogueAtlasA2RescaleO { - using ArchTag = Arch::AtlasA2; -}; - -// For Ascend950, FA Infer online Softmax -template -struct EpilogueAscend950FASoftmax { - using ArchTag = Arch::Ascend950; - static constexpr bool ATTENTION_MASK_FLAG = ATTENTION_MASK_FLAG_; -}; - -// For Ascend950, FA Infer RescaleO -struct EpilogueAscend950FARescaleO { - using ArchTag = Arch::Ascend950; -}; - -// For AtlasA2, MLA RescaleO -struct EpilogueAtlasA2MLARescaleO { - using ArchTag = Arch::AtlasA2; -}; - -// For AtlasA2, MLA FD RescaleO -template -struct EpilogueAtlasA2MLAFDRescaleO { - using ArchTag = Arch::AtlasA2; - static constexpr uint32_t KV_SPLIT_MAX = 64; - static constexpr uint32_t HEADS_PROCESS_MAX = 16; - static constexpr uint32_t COMPUTE_ELE_NUM = COMPUTE_ELE_NUM_; -}; - -// For AtlasA2, MLA TP1 Softmax -struct EpilogueAtlasA2MLATP1Softmax { - using ArchTag = Arch::AtlasA2; -}; - -// For AtlasA2, MLA TP1 RescaleO -struct EpilogueAtlasA2MLATP1RescaleO { - using ArchTag = Arch::AtlasA2; -}; - -// For AtlasA2, per token dequant -template -struct EpilogueAtlasA2PerTokenDequant { - using ArchTag = Arch::AtlasA2; - static constexpr uint32_t UB_STAGES = UB_STAGES_; -}; - -// For AtlasA2, per token dequant tla version -template -struct EpilogueAtlasA2PerTokenDequantTla { - using ArchTag = Arch::AtlasA2; - static constexpr uint32_t UB_STAGES = UB_STAGES_; -}; - -// For Ascend950, per token dequant -template -struct EpilogueAscend950PerTokenDequantTla { - using ArchTag = Arch::Ascend950; - static constexpr uint32_t UB_STAGES = UB_STAGES_; -}; - -// For Ascend950, perGroup + perBlock dequant -struct BlockEpiloguePertile { - using ArchTag = Arch::Ascend950; -}; - -// For AtlasA2, W4A4 epilogue process -struct EpilogueAtlasA2W4A4PerTokenPerChannelDequant { - using ArchTag = Arch::AtlasA2; -}; -//////////////////////////// -/// new add -// For AtlasA2, GEMM -struct EpilogueAtlasA2Gemm { - using ArchTag = Arch::AtlasA2; -}; - -// For AtlasA2, GEMV -struct EpilogueAtlasA2Gemv { - using ArchTag = Arch::AtlasA2; -}; -/////////////////////////// -// For Ascend950, fixpipe-opti -template -struct EpilogueAscend950Fixpipe { - using ArchTag = Arch::Ascend950; - static constexpr bool SPLIT_M = SPLIT_M_; -}; -} // namespace Catlass::Epilogue - -#endif // CATLASS_KDA_DISPATCH_POLICY_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/tile/kda_tile_cast.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/tile/kda_tile_cast.hpp deleted file mode 100644 index 38e2e43c0e..0000000000 --- a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/tile/kda_tile_cast.hpp +++ /dev/null @@ -1,46 +0,0 @@ -/** - * Copyright (c) 2025 Huawei Technologies Co., Ltd. - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of - * CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ - -#ifndef CATLASS_KDA_TILE_CAST_HPP -#define CATLASS_KDA_TILE_CAST_HPP - -#include "catlass/kda_catlass.hpp" - -namespace Catlass::Epilogue::Tile { - -template < - /// Tag indicating architecture - class ArchTag_, - /// Compute data type - class DstType_, - class SrcType_, - /// Length of the compute buffer - class TileShape_> -struct TileCast { - using ArchTag = ArchTag_; - using ElementDst = typename DstType_::Element; - using ElementSrc = typename SrcType_::Element; - using TileShape = TileShape_; - - CATLASS_DEVICE - TileCast() {} - - CATLASS_DEVICE - void operator()( - AscendC::LocalTensor const &ubOut, - AscendC::LocalTensor const &ubIn) - { - AscendC::Cast(ubOut, ubIn, AscendC::RoundMode::CAST_RINT, TileShape::COUNT); - } -}; - -} // namespace Catlass::Epilogue::Tile - -#endif diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_dynamic_aiv.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_dynamic_aiv.hpp deleted file mode 100644 index 0817c0772c..0000000000 --- a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_dynamic_aiv.hpp +++ /dev/null @@ -1,505 +0,0 @@ -/** - * Copyright (c) 2025 Huawei Technologies Co., Ltd. - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of - * CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ - -#ifndef CATLASS_GEMM_BLOCK_MMAD_DYNAMIC_AIV_HPP -#define CATLASS_GEMM_BLOCK_MMAD_DYNAMIC_AIV_HPP - -#include "catlass/kda_catlass.hpp" -#include "catlass/arch/kda_resource.hpp" -#include "catlass/kda_coord.hpp" -#include "catlass/kda_gemm_coord.hpp" -#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" -#include "catlass/gemm/kda_gemm_helper.hpp" -#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" -#include "catlass/gemm/tile/kda_tile_muls.hpp" - -namespace Catlass::Gemm::Block { -template < - uint32_t SCALAR_BUFFER_ELE_NUM_, - uint32_t STAGES_, - class AType_, - class BType_, - class CType_, - class BiasType_, - class TileCopy_, - class TileVmuls_> -struct BlockMmadAiv< - MmadAtlasA2DynamicAiv, - AType_, - BType_, - CType_, - BiasType_, - TileCopy_, - TileVmuls_> { -public: - // Type Aliases - using DispatchPolicy = MmadAtlasA2DynamicAiv; - using ArchTag = typename DispatchPolicy::ArchTag; - using ElementA = typename AType_::Element; - using LayoutA = typename AType_::Layout; - using ElementB = typename BType_::Element; - using LayoutB = typename BType_::Layout; - using ElementC = typename CType_::Element; - using LayoutC = typename CType_::Layout; - using CopyGmToUbA = typename TileCopy_::CopyGmToUbA; - using CopyGmToUbB = typename TileCopy_::CopyGmToUbB; - using CopyUbToGmC = typename TileCopy_::CopyUbToGmC; - using TileVmuls = TileVmuls_; - using TensorCoord = layout::VectorLayout::TensorCoord; - - static constexpr uint32_t UBA_ALIGN = BYTE_PER_BLK / sizeof(ElementA); - static constexpr uint32_t UBB_ALIGN = BYTE_PER_BLK / sizeof(ElementB); - static constexpr uint32_t STAGES = DispatchPolicy::STAGES; - static constexpr uint32_t SCALAR_BUFFER_ELE_NUM = DispatchPolicy::SCALAR_BUFFER_ELE_NUM; - - // Check LayoutC - static_assert(std::is_same_v, "LayoutC only support RowMajor yet!"); - - /// Construct - CATLASS_DEVICE - BlockMmadAiv(Arch::Resource &resource, MatrixCoord taskShape, uint32_t ubBufAddrStart = 0) - { - uint32_t ubASize = SCALAR_BUFFER_ELE_NUM * sizeof(ElementA); - uint32_t ubBSize = taskShape.column() * sizeof(ElementB); - uint32_t ubCSize = ArchTag::UB_SIZE / STAGES - ubASize - ubBSize; - - ubTileNum = ubCSize / ubBSize; - ubTileSize = ubBSize / sizeof(ElementB); - - uint32_t ubAOffset = ubBufAddrStart; - uint32_t ubBOffset = ubBufAddrStart + ubASize * STAGES; - uint32_t ubCOffset = ubBOffset + ubBSize * STAGES; - // Init buffers - for (uint32_t i = 0; i < STAGES; i++) { - ubATensorList[i] = resource.ubBuf.template GetBufferByByte(ubAOffset + ubASize * i); - ubBTensorList[i] = resource.ubBuf.template GetBufferByByte(ubBOffset + ubBSize * i); - ubCTensorList[i] = resource.ubBuf.template GetBufferByByte(ubCOffset + ubCSize * i); - - inEventIds[i] = i; - outEventIds[i] = i; - AscendC::SetFlag(inEventIds[i]); - AscendC::SetFlag(outEventIds[i]); - } - } - - /// Destructor - CATLASS_DEVICE - ~BlockMmadAiv() - { - for (uint32_t i = 0; i < STAGES; i++) { - AscendC::WaitFlag(inEventIds[i]); - AscendC::WaitFlag(outEventIds[i]); - } - } - - /// Perform a block-scoped matrix multiply-accumulate - CATLASS_DEVICE - void operator()( - AscendC::GlobalTensor const &gmBlockA, LayoutA const &layoutA, - AscendC::GlobalTensor const &gmBlockB, LayoutB const &layoutB, - AscendC::GlobalTensor const &gmBlockC, LayoutC const &layoutC, - GemmCoord const &actualShape) - { - AscendC::WaitFlag((event_t)(inEventIds[UbInStageId])); - - uint32_t mRound = RoundUp(actualShape.m()); - ElementA scalarA[SCALAR_BUFFER_ELE_NUM]; - LayoutA layoutAInUb{mRound}; - LayoutB layoutBInUb{ubTileSize}; - LayoutC layoutCInUb{ubTileNum, ubTileSize}; - - // load B - auto layoutTileB = layoutB.GetTileLayout(TensorCoord(actualShape.n())); - copyGmToUbB(ubBTensorList[UbInStageId], gmBlockB[0], layoutBInUb, layoutTileB); - AscendC::SetFlag((event_t)(inEventIds[UbInStageId])); - - // load A - auto layoutTileA = layoutA.GetTileLayout(TensorCoord(actualShape.m())); - copyGmToUbA(ubATensorList[UbInStageId], gmBlockA[0], layoutAInUb, layoutTileA); - AscendC::SetFlag((event_t)(inEventIds[UbInStageId])); - AscendC::WaitFlag((event_t)(inEventIds[UbInStageId])); - for (uint32_t i = 0; i < actualShape.m(); i++) { - scalarA[i] = ubATensorList[UbInStageId].GetValue(i); - } - AscendC::SetFlag((event_t)(inEventIds[UbInStageId])); - AscendC::WaitFlag((event_t)(inEventIds[UbInStageId])); - AscendC::WaitFlag((event_t)(inEventIds[UbInStageId])); - - auto mTiles = CeilDiv(actualShape.m(), ubTileNum); - uint32_t lastTileM = actualShape.m() - (mTiles - 1) * ubTileNum; - // tile vmuls && write c to gm - for (uint32_t i = 0; i < mTiles; i++) { - auto actualTileM = (i == mTiles - 1) ? lastTileM : ubTileNum; - AscendC::WaitFlag((event_t)(outEventIds[UbOutStageId])); - for (uint32_t j = 0; j < actualTileM; j++) { - tileVmuls(ubCTensorList[UbOutStageId][j * ubTileSize], ubBTensorList[UbInStageId], - scalarA[i * ubTileNum + j], actualShape.n()); - } - AscendC::PipeBarrier(); - AscendC::SetFlag((event_t)(outEventIds[UbOutStageId])); - AscendC::WaitFlag((event_t)(outEventIds[UbOutStageId])); - MatrixCoord gmTileCOffset{i * ubTileNum, 0}; - auto gmTileC = gmBlockC[layoutC.GetOffset(gmTileCOffset)]; - LayoutC layoutCBlock = layoutC.GetTileLayout(MakeCoord(actualTileM, actualShape.n())); - LayoutC lastLayoutCInUb = layoutCInUb.GetTileLayout(MakeCoord(actualTileM, actualShape.n())); - copyUbToGmC(gmTileC, ubCTensorList[UbOutStageId], layoutCBlock, lastLayoutCInUb); - AscendC::SetFlag((event_t)(outEventIds[UbOutStageId])); - UbOutStageId = (UbOutStageId + 1 < STAGES) ? (UbOutStageId + 1) : 0; - } - AscendC::SetFlag((event_t)(inEventIds[UbInStageId])); - UbInStageId = (UbInStageId + 1 < STAGES) ? (UbInStageId + 1) : 0; - } - -protected: - /// Data members - AscendC::LocalTensor ubATensorList[STAGES]; - AscendC::LocalTensor ubBTensorList[STAGES]; - AscendC::LocalTensor ubCTensorList[STAGES]; - - CopyGmToUbA copyGmToUbA; - CopyGmToUbB copyGmToUbB; - CopyUbToGmC copyUbToGmC; - TileVmuls tileVmuls; - - int32_t inEventIds[STAGES]; - int32_t outEventIds[STAGES]; - - uint32_t UbOutStageId{0}; - uint32_t UbInStageId{0}; - - uint32_t ubTileNum{0}; - uint32_t ubTileSize{0}; -}; - -template < - uint32_t SCALAR_BUFFER_ELE_NUM_, - bool IS_TILE_M_, - class AType_, - class BType_, - class CType_, - class BiasType_, - class TileCopy_, - class TileVmuls_> -struct BlockMmadAiv< - MmadAtlasA2DynamicAivSimple, - AType_, - BType_, - CType_, - BiasType_, - TileCopy_, - TileVmuls_> { -public: - // Type Aliases - using DispatchPolicy = MmadAtlasA2DynamicAivSimple; - using ArchTag = typename DispatchPolicy::ArchTag; - using ElementA = typename AType_::Element; - using LayoutA = typename AType_::Layout; - using ElementB = typename BType_::Element; - using LayoutB = typename BType_::Layout; - using ElementC = typename CType_::Element; - using LayoutC = typename CType_::Layout; - - using CopyGmToUbA = typename TileCopy_::CopyGmToUbA; - using CopyGmToUbB = typename TileCopy_::CopyGmToUbB; - using CopyUbToGmC = typename TileCopy_::CopyUbToGmC; - using TileVmuls = TileVmuls_; - using TensorCoord = layout::VectorLayout::TensorCoord; - - static constexpr uint32_t UBA_ALIGN = BYTE_PER_BLK / sizeof(ElementA); - static constexpr uint32_t UBB_ALIGN = BYTE_PER_BLK / sizeof(ElementB); - static constexpr uint32_t SCALAR_BUFFER_ELE_NUM = DispatchPolicy::SCALAR_BUFFER_ELE_NUM; - static constexpr bool IS_TILE_M = DispatchPolicy::IS_TILE_M; - - // Check LayoutC - static_assert(std::is_same_v, "LayoutC only support RowMajor yet!"); - - /// Construct - CATLASS_DEVICE - BlockMmadAiv(Arch::Resource &resource, MatrixCoord taskShape, uint32_t ubBufAddrStart = 0) - { - uint32_t ubASize = SCALAR_BUFFER_ELE_NUM * sizeof(ElementA); - uint32_t ubBSize = taskShape.column() * sizeof(ElementB); - uint32_t ubCSize = ArchTag::UB_SIZE - ubASize - ubBSize; - - ubTileNum = ubCSize / ubBSize; - ubTileSize = ubBSize / sizeof(ElementB); - - uint32_t ubAOffset = ubBufAddrStart; - uint32_t ubBOffset = ubBufAddrStart + ubASize; - uint32_t ubCOffset = ubBOffset + ubBSize; - // Init buffers - ubATensor = resource.ubBuf.template GetBufferByByte(ubAOffset); - ubBTensor = resource.ubBuf.template GetBufferByByte(ubBOffset); - ubCTensor = resource.ubBuf.template GetBufferByByte(ubCOffset); - - AscendC::SetFlag(EVENT_ID0); - AscendC::SetFlag(EVENT_ID1); - AscendC::SetFlag(EVENT_ID0); - } - - /// Destructor - CATLASS_DEVICE - ~BlockMmadAiv() - { - AscendC::WaitFlag(EVENT_ID0); - AscendC::WaitFlag(EVENT_ID1); - AscendC::WaitFlag(EVENT_ID0); - } - - /// Perform a block-scoped matrix multiply-accumulate - CATLASS_DEVICE - void operator()( - AscendC::GlobalTensor const &gmBlockA, LayoutA const &layoutA, - AscendC::GlobalTensor const &gmBlockB, LayoutB const &layoutB, - AscendC::GlobalTensor const &gmBlockC, LayoutC const &layoutC, - GemmCoord const &actualShape) - { - uint32_t mRound = RoundUp(actualShape.m()); - - AscendC::WaitFlag((event_t)(EVENT_ID0)); - AscendC::WaitFlag((event_t)(EVENT_ID1)); - AscendC::WaitFlag((event_t)(EVENT_ID0)); - - ElementA scalarA[SCALAR_BUFFER_ELE_NUM]; - LayoutA layoutAInUb{mRound}; - LayoutB layoutBInUb{ubTileSize}; - LayoutC layoutCInUb{ubTileNum, ubTileSize}; - - // load B - auto layoutTileB = layoutB.GetTileLayout(TensorCoord(actualShape.n())); - copyGmToUbB(ubBTensor, gmBlockB[0], layoutBInUb, layoutTileB); - AscendC::SetFlag((event_t)(EVENT_ID1)); - // load A - auto layoutTileA = layoutA.GetTileLayout(TensorCoord(actualShape.m())); - copyGmToUbA(ubATensor, gmBlockA[0], layoutAInUb, layoutTileA); - AscendC::SetFlag((event_t)(EVENT_ID0)); - AscendC::SetFlag((event_t)(EVENT_ID0)); - AscendC::WaitFlag((event_t)(EVENT_ID0)); - if constexpr (IS_TILE_M) { - for (uint32_t i = 0; i < actualShape.m(); i++) { - scalarA[i] = ubATensor.GetValue(i); - } - } else { - scalarA[0] = ubATensor.GetValue(0); - } - AscendC::WaitFlag((event_t)(EVENT_ID0)); - AscendC::WaitFlag((event_t)(EVENT_ID1)); - AscendC::SetFlag((event_t)(EVENT_ID0)); - AscendC::WaitFlag((event_t)(EVENT_ID0)); - - // compute && write c to gm - LayoutC layoutCBlock = layoutC.GetTileLayout(MakeCoord(actualShape.m(), actualShape.n())); - LayoutC lastLayoutCInUb = layoutCInUb.GetTileLayout(MakeCoord(actualShape.m(), actualShape.n())); - if constexpr (IS_TILE_M) { - for (uint32_t j = 0; j < actualShape.m(); j++) { - tileVmuls(ubCTensor[j * ubTileSize], ubBTensor, scalarA[j], actualShape.n()); - } - } else { - tileVmuls(ubCTensor[0], ubBTensor, scalarA[0], actualShape.n()); - } - - AscendC::SetFlag((event_t)(EVENT_ID0)); - AscendC::SetFlag((event_t)(EVENT_ID1)); - AscendC::SetFlag((event_t)(EVENT_ID0)); - AscendC::WaitFlag((event_t)(EVENT_ID0)); - copyUbToGmC(gmBlockC[0], ubCTensor, layoutCBlock, lastLayoutCInUb); - AscendC::SetFlag((event_t)(EVENT_ID0)); - } - -protected: - /// Data members - AscendC::LocalTensor ubATensor; - AscendC::LocalTensor ubBTensor; - AscendC::LocalTensor ubCTensor; - - CopyGmToUbA copyGmToUbA; - CopyGmToUbB copyGmToUbB; - CopyUbToGmC copyUbToGmC; - TileVmuls tileVmuls; - - uint32_t ubTileNum{0}; - uint32_t ubTileSize{0}; -}; - -template < - uint32_t SCALAR_BUFFER_ELE_NUM_, - class AType_, - class BType_, - class CType_, - class BiasType_, - class TileCopy_, - class TileVmuls_> -struct BlockMmadAiv< - MmadAtlasA2DynamicAivTrans, - AType_, - BType_, - CType_, - BiasType_, - TileCopy_, - TileVmuls_> { -public: - // Type Aliases - using DispatchPolicy = MmadAtlasA2DynamicAivTrans; - using ArchTag = typename DispatchPolicy::ArchTag; - using ElementA = typename AType_::Element; - using LayoutA = typename AType_::Layout; - using ElementB = typename BType_::Element; - using LayoutB = typename BType_::Layout; - using ElementC = typename CType_::Element; - using LayoutC = typename CType_::Layout; - using TileVmuls = TileVmuls_; - using TensorCoord = layout::VectorLayout::TensorCoord; - - using CopyGmToUbB = typename TileCopy_::CopyGmToUbB; - using CopyGmToUbA = typename TileCopy_::CopyGmToUbA; - using CopyUbToGmC = typename TileCopy_::CopyUbToGmC; - - static constexpr uint32_t UBA_ALIGN = BYTE_PER_BLK / sizeof(ElementA); - static constexpr uint32_t UBB_ALIGN = BYTE_PER_BLK / sizeof(ElementB); - static constexpr uint32_t SCALAR_BUFFER_ELE_NUM = DispatchPolicy::SCALAR_BUFFER_ELE_NUM; - static constexpr uint16_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(ElementC); - static constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(ElementC); - static constexpr uint32_t VNCHW_SIZE = 16; - - // Check LayoutC - static_assert(std::is_same_v, "LayoutC only support RowMajor yet!"); - - /// Construct - CATLASS_DEVICE - BlockMmadAiv(Arch::Resource &resource, MatrixCoord taskShape, uint32_t ubBufAddrStart = 0) - { - ubTileSize = RoundUp(taskShape.row()); - ubTileNum = RoundUp(taskShape.column()); - uint32_t ubBSize = SCALAR_BUFFER_ELE_NUM * sizeof(ElementB); - uint32_t ubASize = ubTileSize * sizeof(ElementA); - uint32_t ubCSize = ubTileNum * ubTileSize * sizeof(ElementA); - uint32_t ubTransposeCSize = ubTileNum * ubTileSize * sizeof(ElementA); - - uint32_t ubBOffset = ubBufAddrStart; - uint32_t ubAOffset = ubBufAddrStart + ubBSize; - uint32_t ubTransposeCOffset = ubAOffset + ubASize; - uint32_t ubCOffset = ubTransposeCOffset + ubTransposeCSize; - // Init buffers - ubBTensor = resource.ubBuf.template GetBufferByByte(ubBOffset); - ubATensor = resource.ubBuf.template GetBufferByByte(ubAOffset); - ubTransposeCTensor = resource.ubBuf.template GetBufferByByte(ubTransposeCOffset); - ubCTensor = resource.ubBuf.template GetBufferByByte(ubCOffset); - - AscendC::SetFlag(EVENT_ID0); - AscendC::SetFlag(EVENT_ID1); - AscendC::SetFlag(EVENT_ID2); - } - - /// Destructor - CATLASS_DEVICE - ~BlockMmadAiv() - { - AscendC::WaitFlag(EVENT_ID2); - AscendC::WaitFlag(EVENT_ID1); - AscendC::WaitFlag(EVENT_ID0); - } - - /// Perform a block-scoped matrix multiply-accumulate - CATLASS_DEVICE - void operator()( - AscendC::GlobalTensor const &gmBlockA, LayoutA const &layoutA, - AscendC::GlobalTensor const &gmBlockB, LayoutB const &layoutB, - AscendC::GlobalTensor const &gmBlockC, LayoutC const &layoutC, - GemmCoord const &actualShape) - { - AscendC::WaitFlag((event_t)(EVENT_ID0)); - AscendC::WaitFlag((event_t)(EVENT_ID1)); - - ElementA scalarB[SCALAR_BUFFER_ELE_NUM]; - LayoutA layoutBInUb{SCALAR_BUFFER_ELE_NUM}; - LayoutB layoutAInUb{ubTileSize}; - - // load A - auto layoutTileA = layoutA.GetTileLayout(TensorCoord(actualShape.m())); - copyGmToUbA(ubATensor, gmBlockA[0], layoutAInUb, layoutTileA); - AscendC::SetFlag((event_t)(EVENT_ID1)); - // load B - auto layoutTileB = layoutB.GetTileLayout(TensorCoord(actualShape.n())); - copyGmToUbB(ubBTensor, gmBlockB[0], layoutBInUb, layoutTileB); - AscendC::SetFlag((event_t)(EVENT_ID0)); - AscendC::SetFlag((event_t)(EVENT_ID0)); - AscendC::WaitFlag((event_t)(EVENT_ID0)); - for (uint32_t i = 0; i < actualShape.n(); i++) { - scalarB[i] = ubBTensor.GetValue(i); - } - AscendC::WaitFlag((event_t)(EVENT_ID0)); - AscendC::WaitFlag((event_t)(EVENT_ID1)); - AscendC::SetFlag((event_t)(EVENT_ID0)); - AscendC::WaitFlag((event_t)(EVENT_ID0)); - - // compute && write c to ub - for (uint32_t i = 0; i < actualShape.n(); i++) { - tileVmuls(ubCTensor[i * ubTileSize], ubATensor, scalarB[i], actualShape.m()); - } - AscendC::PipeBarrier(); - AscendC::SetFlag((event_t)(EVENT_ID0)); - AscendC::SetFlag((event_t)(EVENT_ID1)); - - LayoutC layoutCInUb{ubTileNum, ubTileSize}; - LayoutC layoutTranposeCInUb{ubTileSize, ubTileNum}; - auto actualLayoutC = layoutC.GetTileLayout(MakeCoord(actualShape.m(), actualShape.n())); - auto actualTransLayoutC = layoutTranposeCInUb.GetTileLayout(MakeCoord(actualShape.m(), actualShape.n())); - - // transpose 16*16 block - auto mTiles = CeilDiv(actualShape.n()); - auto usedBufferNum = CeilDiv(actualShape.m()); - AscendC::WaitFlag((event_t)(EVENT_ID2)); - uint64_t dstLocalList[VNCHW_SIZE]; - uint64_t srcLocalList[VNCHW_SIZE]; - for (uint32_t i = 0; i < mTiles; ++i) { - for (uint32_t j = 0; j < VNCHW_SIZE; ++j) { - srcLocalList[j] = ubCTensor.GetPhyAddr() + i * C0_NUM_PER_FRACTAL * ubTileSize * sizeof(ElementC) + j * ubTileSize * sizeof(ElementC); - if constexpr (std::is_same_v) { - dstLocalList[j] = ubTransposeCTensor.GetPhyAddr() + i * C0_NUM_PER_FRACTAL * sizeof(ElementC) + j * ubTileNum * sizeof(ElementC); - } - if constexpr (std::is_same_v) { - dstLocalList[j] = ubTransposeCTensor.GetPhyAddr() + i * C0_NUM_PER_FRACTAL * sizeof(ElementC) + j / 2 * ubTileNum * sizeof(ElementC) + j % 2 * BYTE_PER_C0; - } - } - if (usedBufferNum > 1) { - AscendC::TransDataTo5HDParams params(false, false, usedBufferNum, ubTileNum, 1); - AscendC::TransDataTo5HD(dstLocalList, srcLocalList, params); - } else { - AscendC::TransDataTo5HDParams paramsRepeatOne{false, false, 1, 0, 0}; - AscendC::TransDataTo5HD(dstLocalList, srcLocalList, paramsRepeatOne); - } - } - - AscendC::SetFlag((event_t)(EVENT_ID0)); - AscendC::WaitFlag((event_t)(EVENT_ID0)); - // write to gmC - copyUbToGmC(gmBlockC[0], ubTransposeCTensor, actualLayoutC, actualTransLayoutC); - AscendC::SetFlag((event_t)(EVENT_ID2)); - } - -protected: - /// Data members - AscendC::LocalTensor ubATensor; - AscendC::LocalTensor ubBTensor; - AscendC::LocalTensor ubCTensor; - AscendC::LocalTensor ubTransposeCTensor; - - CopyGmToUbA copyGmToUbA; - CopyGmToUbB copyGmToUbB; - CopyUbToGmC copyUbToGmC; - TileVmuls tileVmuls; - - uint32_t stageId{0}; - uint32_t ubTileNum{0}; - uint32_t ubTileSize{0}; -}; - -} // namespace Catlass::Gemm::Block -#endif // CATLASS_GEMM_BLOCK_MMAD_DYNAMIC_AIV_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_dynamic_common.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_dynamic_common.hpp deleted file mode 100644 index 6d61d88158..0000000000 --- a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_dynamic_common.hpp +++ /dev/null @@ -1,353 +0,0 @@ -/** - * Copyright (c) 2025 Huawei Technologies Co., Ltd. - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of - * CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ - -#ifndef CATLASS_GEMM_BLOCK_BLOCK_MMAD_DYNAMIC_COMMON_HPP -#define CATLASS_GEMM_BLOCK_BLOCK_MMAD_DYNAMIC_COMMON_HPP - -#include "catlass/kda_catlass.hpp" -#include "catlass/arch/kda_resource.hpp" -#include "catlass/kda_coord.hpp" -#include "catlass/kda_gemm_coord.hpp" -#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" -#include "catlass/gemm/kda_gemm_helper.hpp" -#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" -#include "catlass/gemm/tile/kda_tile_mmad.hpp" - -namespace Catlass::Gemm::Block { - -template -struct BlockMmad, L1TileShape_, L0TileShape_, AType_, - BType_, CType_, BiasType_, TileCopy_, TileMmad_> { -public: - // Type Aliases - using DispatchPolicy = MmadAtlasA2DynamicCommon; - using ArchTag = typename DispatchPolicy::ArchTag; - using ElementA = typename AType_::Element; - using LayoutA = typename AType_::Layout; - using ElementB = typename BType_::Element; - using LayoutB = typename BType_::Layout; - using ElementC = typename CType_::Element; - using LayoutC = typename CType_::Layout; - using TileMmad = TileMmad_; - using CopyGmToL1A = typename TileCopy_::CopyGmToL1A; - using CopyGmToL1B = typename TileCopy_::CopyGmToL1B; - using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; - using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; - using CopyL0CToGm = typename TileCopy_::CopyL0CToGm; - using ElementAccumulator = - typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; - using LayoutAInL1 = typename CopyL1ToL0A::LayoutSrc; - using LayoutBInL1 = typename CopyL1ToL0B::LayoutSrc; - using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst; - using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst; - using LayoutCInL0 = layout::zN; - - using L1AAlignHelper = Gemm::helper::L1AlignHelper; - using L1BAlignHelper = Gemm::helper::L1AlignHelper; - - static constexpr bool ENABLE_UNIT_FLAG = DispatchPolicy::ENABLE_UNIT_FLAG; - static constexpr bool ENABLE_SHUFFLE_K = DispatchPolicy::ENABLE_SHUFFLE_K; - static constexpr uint32_t STAGES = DispatchPolicy::STAGES; - - static constexpr uint32_t L0A_PINGPONG_BUF_SIZE = ArchTag::L0A_SIZE / STAGES; - static constexpr uint32_t L0B_PINGPONG_BUF_SIZE = ArchTag::L0B_SIZE / STAGES; - - // Check LayoutC - static_assert(std::is_same_v, "LayoutC only support RowMajor yet!"); - - /// Construct - CATLASS_DEVICE - BlockMmad(GemmCoord const &l1TileShape_, Arch::Resource &resource, uint32_t l1BufAddrStart = 0) - : l1TileShape(l1TileShape_) - { - uint32_t l1ASize = l1TileShape.m() * l1TileShape.k() * sizeof(ElementA); - uint32_t l1BSize = l1TileShape.k() * l1TileShape.n() * sizeof(ElementB); - - kPartLenMax = min(L0A_PINGPONG_BUF_SIZE / sizeof(ElementA) / l1TileShape.m() / L1AAlignHelper::ELE_NUM_PER_C0 * - L1AAlignHelper::ELE_NUM_PER_C0, - L0B_PINGPONG_BUF_SIZE / sizeof(ElementB) / l1TileShape.n() / L1BAlignHelper::ELE_NUM_PER_C0 * - L1BAlignHelper::ELE_NUM_PER_C0); - - if constexpr (std::is_same_v && std::is_same_v) { - kPartLenMax = RoundDown(kPartLenMax); - } - - uint32_t l1AOffset = l1BufAddrStart; - uint32_t l1BOffset = l1BufAddrStart + l1ASize * STAGES; - // Init buffers - for (uint32_t i = 0; i < STAGES; i++) { - l1ATensorList[i] = resource.l1Buf.template GetBufferByByte(l1AOffset + l1ASize * i); - l1BTensorList[i] = resource.l1Buf.template GetBufferByByte(l1BOffset + l1BSize * i); - l0ATensorList[i] = resource.l0ABuf.template GetBufferByByte(L0A_PINGPONG_BUF_SIZE * i); - l0BTensorList[i] = resource.l0BBuf.template GetBufferByByte(L0B_PINGPONG_BUF_SIZE * i); - - l1AEventList[i] = i; - l1BEventList[i] = i + STAGES; - l0AEventList[i] = i; - l0BEventList[i] = i + STAGES; - AscendC::SetFlag(l1AEventList[i]); - AscendC::SetFlag(l1BEventList[i]); - AscendC::SetFlag(l0AEventList[i]); - AscendC::SetFlag(l0BEventList[i]); - } - l0CTensor = resource.l0CBuf.template GetBufferByByte(0); - AscendC::SetFlag(EVENT_ID0); - } - - /// Destructor - CATLASS_DEVICE - ~BlockMmad() - { - for (uint32_t i = 0; i < STAGES; i++) { - AscendC::WaitFlag(l1AEventList[i]); - AscendC::WaitFlag(l1BEventList[i]); - AscendC::WaitFlag(l0AEventList[i]); - AscendC::WaitFlag(l0BEventList[i]); - } - AscendC::WaitFlag(EVENT_ID0); - } - - /// Perform a block-scoped matrix multiply-accumulate - CATLASS_DEVICE - void operator()(AscendC::GlobalTensor const &gmBlockA, LayoutA const &layoutA, - AscendC::GlobalTensor const &gmBlockB, LayoutB const &layoutB, - AscendC::GlobalTensor const &gmBlockC, LayoutC const &layoutC, - AscendC::GlobalTensor const &gmNextBlockA, AscendC::GlobalTensor const &gmNextBlockB, - GemmCoord const &actualShape, GemmCoord const &actualShapeNext, bool isFirstBlock, bool hasNextBlock) - { - uint32_t mRound = RoundUp(actualShape.m()); - uint32_t nRound = RoundUp(actualShape.n()); - - auto layoutAInL1 = LayoutAInL1::template MakeLayout(l1TileShape.m(), l1TileShape.k()); - auto layoutBInL1 = LayoutBInL1::template MakeLayout(l1TileShape.k(), l1TileShape.n()); - auto layoutInL0C = LayoutCInL0::MakeLayoutInL0C(MakeCoord(mRound, nRound)); - - uint32_t kTileCount = CeilDiv(actualShape.k(), l1TileShape.k()); - - if constexpr (!ENABLE_UNIT_FLAG) { - AscendC::WaitFlag(EVENT_ID0); - } - uint32_t startTileIdx = 0; - if constexpr (ENABLE_SHUFFLE_K) { - startTileIdx = AscendC::GetBlockIdx(); - } - uint32_t firstTileIdx = startTileIdx % kTileCount; - uint32_t lastTileIdx = (startTileIdx + kTileCount - 1) % kTileCount; - uint32_t kActual = - (firstTileIdx < kTileCount - 1) ? l1TileShape.k() : (actualShape.k() - firstTileIdx * l1TileShape.k()); - - // k loop - for (uint32_t kLoopIdx = 0; kLoopIdx < kTileCount; kLoopIdx++) { - uint32_t shuffleKIdx = (startTileIdx + kLoopIdx) % kTileCount; - // Load first matrix A tile in total kernel loop from GM to L1 - if (shuffleKIdx == firstTileIdx && isFirstBlock) { - MatrixCoord gmTileACoord{0, shuffleKIdx * l1TileShape.k()}; - MatrixCoord gmTileBCoord{shuffleKIdx * l1TileShape.k(), 0}; - auto gmTileA = gmBlockA[layoutA.GetOffset(gmTileACoord)]; - auto gmTileB = gmBlockB[layoutB.GetOffset(gmTileBCoord)]; - // Load first matrix A tile from GM to L1 - AscendC::WaitFlag(l1AEventList[l1ListId]); - auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShape.m(), kActual)); - copyGmToL1A(l1ATensorList[l1ListId], gmTileA, layoutAInL1, layoutTileA); - AscendC::SetFlag(l1AEventList[l1ListId]); - - // Load first matrix B tile from GM to L1 - AscendC::WaitFlag(l1BEventList[l1ListId]); - auto layoutTileB = layoutB.GetTileLayout(MakeCoord(kActual, actualShape.n())); - copyGmToL1B(l1BTensorList[l1ListId], gmTileB, layoutBInL1, layoutTileB); - AscendC::SetFlag(l1BEventList[l1ListId]); - } - - uint32_t l1ListIdNext = (l1ListId + 1 < STAGES) ? (l1ListId + 1) : 0; - uint32_t kActualNext{0}; - - // preload next tile from GM to L1 - if (shuffleKIdx != lastTileIdx) { - uint32_t shuffleKIdxNext = (startTileIdx + kLoopIdx + 1) % kTileCount; - // Get L1 tensor for next stage - auto l1ATensor = l1ATensorList[l1ListIdNext]; - auto l1BTensor = l1BTensorList[l1ListIdNext]; - // Get GM tensor for next stage - kActualNext = (shuffleKIdxNext < kTileCount - 1) ? l1TileShape.k() : (actualShape.k() - shuffleKIdxNext * l1TileShape.k()); - MatrixCoord gmTileACoord{0, shuffleKIdxNext * l1TileShape.k()}; - MatrixCoord gmTileBCoord{shuffleKIdxNext * l1TileShape.k(), 0}; - auto gmTileA = gmBlockA[layoutA.GetOffset(gmTileACoord)]; - auto gmTileB = gmBlockB[layoutB.GetOffset(gmTileBCoord)]; - - // load next matrix A tile from GM to L1 - AscendC::WaitFlag(l1AEventList[l1ListIdNext]); - auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShape.m(), kActualNext)); - copyGmToL1A(l1ATensor, gmTileA, layoutAInL1, layoutTileA); - AscendC::SetFlag(l1AEventList[l1ListIdNext]); - - // load next matrix B tile from GM to L1 - AscendC::WaitFlag(l1BEventList[l1ListIdNext]); - auto layoutTileB = layoutB.GetTileLayout(MakeCoord(kActualNext, actualShape.n())); - copyGmToL1B(l1BTensor, gmTileB, layoutBInL1, layoutTileB); - AscendC::SetFlag(l1BEventList[l1ListIdNext]); - } - if (shuffleKIdx == lastTileIdx && hasNextBlock) { - uint32_t kTileCount = CeilDiv(actualShapeNext.k(), l1TileShape.k()); - uint32_t firstTileIdx = startTileIdx % kTileCount; - // Get L1 tensor for next stage - auto l1ATensor = l1ATensorList[l1ListIdNext]; - auto l1BTensor = l1BTensorList[l1ListIdNext]; - // Get GM tensor for next stage - kActualNext = (firstTileIdx < kTileCount - 1) ? l1TileShape.k() : (actualShapeNext.k() - firstTileIdx * l1TileShape.k()); - MatrixCoord gmTileACoord{0, firstTileIdx * l1TileShape.k()}; - MatrixCoord gmTileBCoord{firstTileIdx * l1TileShape.k(), 0}; - auto gmTileA = gmNextBlockA[layoutA.GetOffset(gmTileACoord)]; - auto gmTileB = gmNextBlockB[layoutB.GetOffset(gmTileBCoord)]; - // load next matrix A tile from GM to L1 - AscendC::WaitFlag(l1AEventList[l1ListIdNext]); - auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShapeNext.m(), kActualNext)); - copyGmToL1A(l1ATensor, gmTileA, layoutAInL1, layoutTileA); - AscendC::SetFlag(l1AEventList[l1ListIdNext]); - - // load next matrix B tile from GM to L1 - AscendC::WaitFlag(l1BEventList[l1ListIdNext]); - auto layoutTileB = layoutB.GetTileLayout(MakeCoord(kActualNext, actualShapeNext.n())); - copyGmToL1B(l1BTensor, gmTileB, layoutBInL1, layoutTileB); - AscendC::SetFlag(l1BEventList[l1ListIdNext]); - } - - // Get L1 tensor for current stage - auto l1ATensor = l1ATensorList[l1ListId]; - auto l1BTensor = l1BTensorList[l1ListId]; - - // Get the loop nums on L0 - uint32_t kPartLoop = CeilDiv(kActual, kPartLenMax); - - uint32_t l0ABufId = 0; - uint32_t l0BBufId = 0; - - for (int kPartIdx = 0; kPartIdx < kPartLoop; kPartIdx++) { - uint32_t kPartActual = - (kPartIdx < kPartLoop - 1) ? kPartLenMax : (kActual - kPartIdx * kPartLenMax); - - // Locate the current tile on L0A - auto l0ATile = l0ATensorList[l0ABufId]; - LayoutAInL0 layoutAInL0 = LayoutAInL0::template MakeLayout(mRound, kPartActual); - // Locate the current tile of matrix A on L1 - MatrixCoord l1ACoord{0, kPartIdx * kPartLenMax}; - auto l1ATile = l1ATensor[layoutAInL1.GetOffset(l1ACoord)]; - - AscendC::WaitFlag(l0AEventList[l0ABufId]); - if (kPartIdx == 0) { - AscendC::WaitFlag(l1AEventList[l1ListId]); - } - - copyL1ToL0A(l0ATile, l1ATile, layoutAInL0, layoutAInL1); - - if (kPartIdx == kPartLoop - 1) { - AscendC::SetFlag(l1AEventList[l1ListId]); - } - - // Locate the current tile on L0B - auto l0BTile = l0BTensorList[l0BBufId]; - LayoutBInL0 layoutBInL0 = LayoutBInL0::template MakeLayout(kPartActual, nRound); - // Locate the current tile of matrix B on L1 - MatrixCoord l1BCoord{kPartIdx * kPartLenMax, 0}; - auto l1BTile = l1BTensor[layoutBInL1.GetOffset(l1BCoord)]; - - // Wait for mmad finished - AscendC::WaitFlag(l0BEventList[l0BBufId]); - // If the current tile is the first one on the k&n axis, wait for loading matrix B from GM to L1 - if (kPartIdx == 0) { - AscendC::WaitFlag(l1BEventList[l1ListId]); - } - - // Load current tile from L1 to L0B - copyL1ToL0B(l0BTile, l1BTile, layoutBInL0, layoutBInL1); - - // If the current tile is the last one on the k&n axis, notify to load matrix B from GM to L1 - if (kPartIdx == kPartLoop - 1) { - AscendC::SetFlag(l1BEventList[l1ListId]); - } - // Notify to do mmad - AscendC::SetFlag(EVENT_ID0); - - // Locate the current tile on L0C - MatrixCoord l0CCoord{static_cast(0U), static_cast(0U)}; - auto l0CTile = l0CTensor[layoutInL0C.GetOffset(l0CCoord)]; - - // Compute the matrix multiplication on L0A and L0B and write the result to the accumulator - // Wait for loading L0B - AscendC::WaitFlag(EVENT_ID0); - - // If the current tile is the first tile on the k axis, the accumulator needs to be reset to 0 - bool initC = ((kLoopIdx == 0) && (kPartIdx == 0)); - // If the unit flag is enabled, the unit flag is set according to the calculation progress - uint8_t unitFlag = 0b00; - if constexpr (ENABLE_UNIT_FLAG) { - if ((kLoopIdx == kTileCount - 1) && (kPartIdx == kPartLoop - 1)) { - unitFlag = 0b11; - } else { - unitFlag = 0b10; - } - } - // Perform calculation operations - tileMmad(l0CTile, l0ATile, l0BTile, mRound, nRound, kPartActual, initC, unitFlag); - - // Notify to move the next L0B tile - AscendC::SetFlag(l0BEventList[l0BBufId]); - - l0BBufId = (l0BBufId + 1 < STAGES) ? (l0BBufId + 1) : 0; - - AscendC::SetFlag(l0AEventList[l0ABufId]); - l0ABufId = (l0ABufId + 1 < STAGES) ? (l0ABufId + 1) : 0; - } - - l1ListId = l1ListIdNext; - kActual = kActualNext; - } - - // copy block out - LayoutC layoutBlock = layoutC.GetTileLayout(actualShape.GetCoordMN()); - - if constexpr (!ENABLE_UNIT_FLAG) { - AscendC::SetFlag(EVENT_ID0); - AscendC::WaitFlag(EVENT_ID0); - copyL0CToGm(gmBlockC, l0CTensor, layoutBlock, layoutInL0C); - AscendC::SetFlag(EVENT_ID0); - } else { - copyL0CToGm(gmBlockC, l0CTensor, layoutBlock, layoutInL0C, 0b11); - } - } - -protected: - /// Data members - AscendC::LocalTensor l1ATensorList[STAGES]; - AscendC::LocalTensor l1BTensorList[STAGES]; - AscendC::LocalTensor l0ATensorList[STAGES]; - AscendC::LocalTensor l0BTensorList[STAGES]; - AscendC::LocalTensor l0CTensor; - - int32_t l1AEventList[STAGES]; - int32_t l1BEventList[STAGES]; - int32_t l0AEventList[STAGES]; - int32_t l0BEventList[STAGES]; - - uint32_t l1ListId{0}; - GemmCoord l1TileShape; - uint32_t kPartLenMax; - - TileMmad tileMmad; - CopyGmToL1A copyGmToL1A; - CopyGmToL1B copyGmToL1B; - CopyL1ToL0A copyL1ToL0A; - CopyL1ToL0B copyL1ToL0B; - CopyL0CToGm copyL0CToGm; -}; - -} // namespace Catlass::Gemm::Block - -#endif // CATLASS_GEMM_BLOCK_BLOCK_MMAD_DYNAMIC_COMMON_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_dynamic_preload_async_with_callback.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_dynamic_preload_async_with_callback.hpp deleted file mode 100644 index c0d7093aa5..0000000000 --- a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_dynamic_preload_async_with_callback.hpp +++ /dev/null @@ -1,441 +0,0 @@ -/** - * Copyright (c) 2026 Huawei Technologies Co., Ltd. - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of - * CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ - -#ifndef CATLASS_GEMM_BLOCK_BLOCK_MMAD_DYNAMIC_PRELOAD_ASYNC_WITH_CALLBACK_HPP -#define CATLASS_GEMM_BLOCK_BLOCK_MMAD_DYNAMIC_PRELOAD_ASYNC_WITH_CALLBACK_HPP - -#include "catlass/kda_catlass.hpp" -#include "catlass/arch/kda_resource.hpp" -#include "catlass/kda_coord.hpp" -#include "catlass/detail/kda_callback.hpp" -#include "catlass/kda_gemm_coord.hpp" -#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" -#include "catlass/gemm/kda_gemm_helper.hpp" -#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" -#include "catlass/gemm/tile/kda_tile_mmad.hpp" - -namespace Catlass::Gemm::Block { - -template < - uint32_t PRELOAD_STAGES_, - uint32_t L1_STAGES_, - uint32_t L0A_STAGES_, - uint32_t L0B_STAGES_, - uint32_t L0C_STAGES_, - bool ENABLE_UNIT_FLAG_, - bool ENABLE_SHUFFLE_K_, - class L1TileShape_, - class L0TileShape_, - class AType_, - class BType_, - class CType_, - class BiasType_, - class TileCopy_, - class TileMmad_> -struct BlockMmad< - MmadAtlasA2DynamicPreloadAsyncWithCallback< - PRELOAD_STAGES_, - L1_STAGES_, - L0A_STAGES_, - L0B_STAGES_, - L0C_STAGES_, - ENABLE_UNIT_FLAG_, - ENABLE_SHUFFLE_K_>, - L1TileShape_, - L0TileShape_, - AType_, - BType_, - CType_, - BiasType_, - TileCopy_, - TileMmad_> { -public: - // Type Aliases - using DispatchPolicy = MmadAtlasA2DynamicPreloadAsyncWithCallback< - PRELOAD_STAGES_, - L1_STAGES_, - L0A_STAGES_, - L0B_STAGES_, - L0C_STAGES_, - ENABLE_UNIT_FLAG_, - ENABLE_SHUFFLE_K_>; - using ArchTag = typename DispatchPolicy::ArchTag; - using ElementA = typename AType_::Element; - using LayoutA = typename AType_::Layout; - using ElementB = typename BType_::Element; - using LayoutB = typename BType_::Layout; - using ElementC = typename CType_::Element; - using LayoutC = typename CType_::Layout; - using TileMmad = TileMmad_; - using CopyGmToL1A = typename TileCopy_::CopyGmToL1A; - using CopyGmToL1B = typename TileCopy_::CopyGmToL1B; - using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; - using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; - using CopyL0CToGm = typename TileCopy_::CopyL0CToGm; - using ElementAccumulator = - typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; - using LayoutAInL1 = typename CopyL1ToL0A::LayoutSrc; - using LayoutBInL1 = typename CopyL1ToL0B::LayoutSrc; - using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst; - using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst; - using LayoutCInL0 = layout::zN; - - using L1AAlignHelper = Gemm::helper::L1AlignHelper; - using L1BAlignHelper = Gemm::helper::L1AlignHelper; - - static constexpr uint32_t PRELOAD_STAGES = DispatchPolicy::PRELOAD_STAGES; - static constexpr uint32_t L1_STAGES = DispatchPolicy::L1_STAGES; - static constexpr uint32_t L0A_STAGES = DispatchPolicy::L0A_STAGES; - static constexpr uint32_t L0B_STAGES = DispatchPolicy::L0B_STAGES; - static constexpr uint32_t L0C_STAGES = DispatchPolicy::L0C_STAGES; - - static constexpr bool ENABLE_UNIT_FLAG = DispatchPolicy::ENABLE_UNIT_FLAG; - static constexpr bool ENABLE_SHUFFLE_K = DispatchPolicy::ENABLE_SHUFFLE_K; - - // Check LayoutC - static_assert(std::is_same_v, "LayoutC only support RowMajor yet!"); - - CATLASS_DEVICE - BlockMmad(GemmCoord const &l1TileShape_, GemmCoord const &l0TileShape_, Arch::Resource &resource, uint32_t l1BufAddrStart = 0) - : l1TileShape(l1TileShape_), - l0TileShape(l0TileShape_) - { - if ASCEND_IS_AIV { - return; - } - - uint32_t l1ASize = l1TileShape.m() * l1TileShape.k() * sizeof(ElementA); - uint32_t l1BSize = l1TileShape.k() * l1TileShape.n() * sizeof(ElementB); - - uint32_t l0ASize = l0TileShape.m() * l0TileShape.k() * sizeof(ElementA); - uint32_t l0BSize = l0TileShape.k() * l0TileShape.n() * sizeof(ElementB); - uint32_t l0CSize = l0TileShape.m() * l0TileShape.n() * sizeof(ElementAccumulator); - - uint32_t l1AOffset = l1BufAddrStart; - uint32_t l1BOffset = l1BufAddrStart + l1ASize * L1_STAGES; - for (uint32_t i = 0; i < L1_STAGES; ++i) { - l1ATensorList[i] = resource.l1Buf.template GetBufferByByte(l1AOffset + l1ASize * i); - l1BTensorList[i] = resource.l1Buf.template GetBufferByByte(l1BOffset + l1BSize * i); - l1AEventList[i] = i; - l1BEventList[i] = i + L1_STAGES; - AscendC::SetFlag(l1AEventList[i]); - AscendC::SetFlag(l1BEventList[i]); - } - - for (uint32_t i = 0; i < L0A_STAGES; ++i) { - l0ATensorList[i] = resource.l0ABuf.template GetBufferByByte(l0ASize * i); - l0AEventList[i] = i; - AscendC::SetFlag(l0AEventList[i]); - } - - for (uint32_t i = 0; i < L0B_STAGES; ++i) { - l0BTensorList[i] = resource.l0BBuf.template GetBufferByByte(l0BSize * i); - l0BEventList[i] = i + L0A_STAGES; - AscendC::SetFlag(l0BEventList[i]); - } - - for (uint32_t i = 0; i < L0C_STAGES; ++i) { - l0CTensorList[i] = resource.l0CBuf.template GetBufferByByte(l0CSize * i); - l0CEventList[i] = i; - AscendC::SetFlag(l0CEventList[i]); - } - } - - CATLASS_DEVICE - ~BlockMmad() - { - if ASCEND_IS_AIV { - return; - } - - SynchronizeBlock(); - for (uint32_t i = 0; i < L1_STAGES; ++i) { - AscendC::WaitFlag(l1AEventList[i]); - AscendC::WaitFlag(l1BEventList[i]); - } - for (uint32_t i = 0; i < L0A_STAGES; ++i) { - AscendC::WaitFlag(l0AEventList[i]); - } - for (uint32_t i = 0; i < L0B_STAGES; ++i) { - AscendC::WaitFlag(l0BEventList[i]); - } - for (uint32_t i = 0; i < L0C_STAGES; ++i) { - AscendC::WaitFlag(l0CEventList[i]); - } - } - - CATLASS_DEVICE - void operator()( - AscendC::GlobalTensor const &gmBlockA, LayoutA const &layoutA, - AscendC::GlobalTensor const &gmBlockB, LayoutB const &layoutB, - AscendC::GlobalTensor const &gmBlockC, LayoutC const &layoutC, - GemmCoord const &actualShape, - Callback const &callbackBeforeFixpipe, Callback const &callbackAfterFixpipe) - { - uint32_t kTileCount = CeilDiv(actualShape.k(), l1TileShape.k()); - - uint32_t mRound = RoundUp(actualShape.m()); - uint32_t nRound = RoundUp(actualShape.n()); - - uint32_t startTileIdx = 0; - if constexpr (ENABLE_SHUFFLE_K) { - startTileIdx = AscendC::GetBlockIdx() % kTileCount; - } - - for (uint32_t kLoopIdx = 0; kLoopIdx < kTileCount; ++kLoopIdx) { - uint32_t kTileIdx = (startTileIdx + kLoopIdx < kTileCount) ? - (startTileIdx + kLoopIdx) : - (startTileIdx + kLoopIdx - kTileCount); - - uint32_t kActual = (kTileIdx < kTileCount - 1) ? - l1TileShape.k() : - (actualShape.k() - kTileIdx * l1TileShape.k()); - - // Emission load instruction from GM to L1 - MatrixCoord gmTileAOffset{0, kTileIdx * l1TileShape.k()}; - MatrixCoord gmTileBOffset{kTileIdx * l1TileShape.k(), 0}; - auto gmTileA = gmBlockA[layoutA.GetOffset(gmTileAOffset)]; - auto gmTileB = gmBlockB[layoutB.GetOffset(gmTileBOffset)]; - // Load first matrix A tile from GM to L1 - AscendC::WaitFlag(l1AEventList[l1ListId]); - auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShape.m(), kActual)); - auto l1a_layout = LayoutAInL1::template MakeLayout(l1TileShape.m(), l1TileShape.k()); - copyGmToL1A(l1ATensorList[l1ListId], gmTileA, l1a_layout, layoutTileA); - AscendC::SetFlag(l1AEventList[l1ListId]); - // Load first matrix B tile from GM to L1 - AscendC::WaitFlag(l1BEventList[l1ListId]); - auto layoutTileB = layoutB.GetTileLayout(MakeCoord(kActual, actualShape.n())); - auto l1b_layout = LayoutBInL1::template MakeLayout(l1TileShape.k(), l1TileShape.n()); - copyGmToL1B(l1BTensorList[l1ListId], gmTileB, l1b_layout, layoutTileB); - AscendC::SetFlag(l1BEventList[l1ListId]); - - // If the number of preload instructions reaches the upper limit, perform an mmad calculation on L1 tile - if (preloadCount == PRELOAD_STAGES) { - L1TileMmad(l1TileMmadParamsList[l1TileMmadParamsId]); - } - - // Store the current load status - uint32_t preloadL1TileMmadParamsId = (l1TileMmadParamsId + preloadCount < PRELOAD_STAGES) ? - (l1TileMmadParamsId + preloadCount) : - (l1TileMmadParamsId + preloadCount - PRELOAD_STAGES); - auto &l1TileMmadParams = l1TileMmadParamsList[preloadL1TileMmadParamsId]; - l1TileMmadParams.l1ListId = l1ListId; - l1TileMmadParams.mRound = mRound; - l1TileMmadParams.nRound = nRound; - l1TileMmadParams.kActual = kActual; - l1TileMmadParams.isKLoopFirst = (kLoopIdx == 0); - l1TileMmadParams.isKLoopLast = (kLoopIdx == kTileCount - 1); - if (kLoopIdx == kTileCount - 1) { - l1TileMmadParams.gmBlockC = gmBlockC; - l1TileMmadParams.layoutCInGm = layoutC.GetTileLayout(actualShape.GetCoordMN()); - l1TileMmadParams.callbackBeforeFixpipe = callbackBeforeFixpipe; - l1TileMmadParams.callbackAfterFixpipe = callbackAfterFixpipe; - } - - if (preloadCount < PRELOAD_STAGES) { - ++preloadCount; - } else { - l1TileMmadParamsId = (l1TileMmadParamsId + 1 < PRELOAD_STAGES) ? (l1TileMmadParamsId + 1) : 0; - } - l1ListId = (l1ListId + 1 < L1_STAGES) ? (l1ListId + 1) : 0; - } - } - - CATLASS_DEVICE - void SynchronizeBlock() - { - while (preloadCount > 0) { - L1TileMmad(l1TileMmadParamsList[l1TileMmadParamsId]); - l1TileMmadParamsId = (l1TileMmadParamsId + 1 < PRELOAD_STAGES) ? (l1TileMmadParamsId + 1) : 0; - --preloadCount; - } - } - -private: - struct L1TileMmadParams { - uint32_t l1ListId; - uint32_t mRound; - uint32_t nRound; - uint32_t kActual; - bool isKLoopFirst; - bool isKLoopLast; - AscendC::GlobalTensor gmBlockC; - LayoutC layoutCInGm; - Callback callbackBeforeFixpipe; - Callback callbackAfterFixpipe; - - CATLASS_DEVICE - L1TileMmadParams() = default; - }; - - CATLASS_DEVICE - void InitL1(Arch::Resource &resource, uint32_t l1BufAddrStart) - { - uint32_t l1ASize = l1TileShape.m() * l1TileShape.k() * sizeof(ElementA); - uint32_t l1BSize = l1TileShape.k() * l1TileShape.n() * sizeof(ElementB); - uint32_t l1AOffset = l1BufAddrStart; - uint32_t l1BOffset = l1BufAddrStart + l1ASize * L1_STAGES; - for (uint32_t i = 0; i < L1_STAGES; ++i) { - l1ATensorList[i] = resource.l1Buf.template GetBufferByByte(l1AOffset + l1ASize * i); - l1BTensorList[i] = resource.l1Buf.template GetBufferByByte(l1BOffset + l1BSize * i); - l1AEventList[i] = i; - l1BEventList[i] = i + L1_STAGES; - AscendC::SetFlag(l1AEventList[i]); - AscendC::SetFlag(l1BEventList[i]); - } - } - - CATLASS_DEVICE - void L1TileMmad(L1TileMmadParams const ¶ms) - { - uint32_t mPartLoop = CeilDiv(params.mRound, l0TileShape.m()); - uint32_t nPartLoop = CeilDiv(params.nRound, l0TileShape.n()); - uint32_t kPartLoop = CeilDiv(params.kActual, l0TileShape.k()); - - auto &l1ATensor = l1ATensorList[params.l1ListId]; - auto &l1BTensor = l1BTensorList[params.l1ListId]; - - auto &l0CTensor = l0CTensorList[l0CListId]; - LayoutCInL0 layoutCInL0 = LayoutCInL0::MakeLayoutInL0C(MakeCoord(params.mRound, params.nRound)); - - if constexpr (!ENABLE_UNIT_FLAG) { - if (params.isKLoopFirst) { - AscendC::WaitFlag(l0CEventList[l0CListId]); - } - } - - for (uint32_t mPartIdx = 0; mPartIdx < mPartLoop; ++mPartIdx) { - uint32_t mPartActual = (mPartIdx < mPartLoop - 1) ? - l0TileShape.m() : - (params.mRound - mPartIdx * l0TileShape.m()); - - for (uint32_t kPartIdx = 0; kPartIdx < kPartLoop; ++kPartIdx) { - uint32_t kPartActual = (kPartIdx < kPartLoop - 1) ? - l0TileShape.k() : - (params.kActual - kPartIdx * l0TileShape.k()); - - auto &l0ATile = l0ATensorList[l0AListId]; - auto layoutAInL0 = LayoutAInL0::template MakeLayout(mPartActual, kPartActual); - auto l1AOffset = MakeCoord(mPartIdx, kPartIdx) * MakeCoord(l0TileShape.m(), l0TileShape.k()); - auto l1a_layout = LayoutAInL1::template MakeLayout(l1TileShape.m(), l1TileShape.k()); - auto l1ATile = l1ATensor[l1a_layout.GetOffset(l1AOffset)]; - - AscendC::WaitFlag(l0AEventList[l0AListId]); - if ((mPartIdx == 0) && (kPartIdx == 0)) { - AscendC::WaitFlag(l1AEventList[params.l1ListId]); - } - copyL1ToL0A(l0ATile, l1ATile, layoutAInL0, l1a_layout); - if ((mPartIdx == mPartLoop - 1) && (kPartIdx == kPartLoop - 1)) { - AscendC::SetFlag(l1AEventList[params.l1ListId]); - } - - for (uint32_t nPartIdx = 0; nPartIdx < nPartLoop; ++nPartIdx) { - uint32_t nPartActual = (nPartIdx < nPartLoop - 1) ? - l0TileShape.n() : - (params.nRound - nPartIdx * l0TileShape.n()); - - auto &l0BTile = l0BTensorList[l0BListId]; - auto layoutBInL0 = LayoutBInL0::template MakeLayout(kPartActual, nPartActual); - auto l1BOffset = MakeCoord(kPartIdx, nPartIdx) * MakeCoord(l0TileShape.k(), l0TileShape.n()); - auto l1b_layout = LayoutBInL1::template MakeLayout(l1TileShape.k(), l1TileShape.n()); - auto l1BTile = l1BTensor[l1b_layout.GetOffset(l1BOffset)]; - - AscendC::WaitFlag(l0BEventList[l0BListId]); - if ((kPartIdx == 0) && (nPartIdx == 0)) { - AscendC::WaitFlag(l1BEventList[params.l1ListId]); - } - copyL1ToL0B(l0BTile, l1BTile, layoutBInL0, l1b_layout); - if ((kPartIdx == kPartLoop - 1) && (nPartIdx == nPartLoop - 1)) { - AscendC::SetFlag(l1BEventList[params.l1ListId]); - } - - AscendC::SetFlag(EVENT_ID0); - - auto l0COffset = MakeCoord(mPartIdx, nPartIdx) * MakeCoord(l0TileShape.m(), l0TileShape.n()); - auto l0CTile = l0CTensor[layoutCInL0.GetOffset(l0COffset)]; - - AscendC::WaitFlag(EVENT_ID0); - // If the current tile is the first tile on the k axis, the accumulator needs to be reset to 0 - bool initC = (params.isKLoopFirst && (kPartIdx == 0)); - // If the unit flag is enabled, the unit flag is set according to the calculation progress - uint8_t unitFlag = 0b00; - if constexpr (ENABLE_UNIT_FLAG) { - if (params.isKLoopLast && - (mPartIdx == mPartLoop - 1) && (kPartIdx == kPartLoop - 1) && (nPartIdx == nPartLoop - 1)) { - unitFlag = 0b11; - } else { - unitFlag = 0b10; - } - } - tileMmad(l0CTile, l0ATile, l0BTile, mPartActual, nPartActual, kPartActual, initC, unitFlag); - - AscendC::SetFlag(l0BEventList[l0BListId]); - l0BListId = (l0BListId + 1 < L0B_STAGES) ? (l0BListId + 1) : 0; - } - AscendC::SetFlag(l0AEventList[l0AListId]); - l0AListId = (l0AListId + 1 < L0A_STAGES) ? (l0AListId + 1) : 0; - } - } - - if (params.isKLoopLast) { - auto layoutCInGm = params.layoutCInGm; - - params.callbackBeforeFixpipe(); - - if constexpr (!ENABLE_UNIT_FLAG) { - AscendC::SetFlag(l0CEventList[l0CListId]); - AscendC::WaitFlag(l0CEventList[l0CListId]); - copyL0CToGm(params.gmBlockC, l0CTensor, layoutCInGm, layoutCInL0); - AscendC::SetFlag(l0CEventList[l0CListId]); - } else { - copyL0CToGm(params.gmBlockC, l0CTensor, layoutCInGm, layoutCInL0, 0b11); - } - l0CListId = (l0CListId + 1 < L0C_STAGES) ? (l0CListId + 1) : 0; - - params.callbackAfterFixpipe(); - } - } - -protected: - AscendC::LocalTensor l1ATensorList[L1_STAGES]; - AscendC::LocalTensor l1BTensorList[L1_STAGES]; - int32_t l1AEventList[L1_STAGES]; - int32_t l1BEventList[L1_STAGES]; - uint32_t l1ListId{0}; - - AscendC::LocalTensor l0ATensorList[L0A_STAGES]; - int32_t l0AEventList[L0A_STAGES]; - uint32_t l0AListId{0}; - - AscendC::LocalTensor l0BTensorList[L0B_STAGES]; - int32_t l0BEventList[L0B_STAGES]; - uint32_t l0BListId{0}; - - AscendC::LocalTensor l0CTensorList[L0C_STAGES_]; - int32_t l0CEventList[L0C_STAGES_]; - uint32_t l0CListId{0}; - - L1TileMmadParams l1TileMmadParamsList[PRELOAD_STAGES]; - uint32_t l1TileMmadParamsId{0}; - uint32_t preloadCount{0}; - GemmCoord l1TileShape; - GemmCoord l0TileShape; - - TileMmad tileMmad; - CopyGmToL1A copyGmToL1A; - CopyGmToL1B copyGmToL1B; - CopyL1ToL0A copyL1ToL0A; - CopyL1ToL0B copyL1ToL0B; - CopyL0CToGm copyL0CToGm; -}; - -} // namespace Catlass::Gemm::Block - -#endif // CATLASS_GEMM_BLOCK_BLOCK_MMAD_DYNAMIC_PRELOAD_ASYNC_WITH_CALLBACK_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_dynamic_single_core_splitk.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_dynamic_single_core_splitk.hpp deleted file mode 100644 index bacb903288..0000000000 --- a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_dynamic_single_core_splitk.hpp +++ /dev/null @@ -1,391 +0,0 @@ -/** - * Copyright (c) 2025 Huawei Technologies Co., Ltd. - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of - * CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ - -#ifndef CATLASS_GEMM_BLOCK_MMAD_DYNAMIC_SINGLE_CORE_SPLITK_HPP -#define CATLASS_GEMM_BLOCK_MMAD_DYNAMIC_SINGLE_CORE_SPLITK_HPP - -#include "catlass/arch/kda_resource.hpp" -#include "catlass/kda_catlass.hpp" -#include "catlass/kda_coord.hpp" -#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" -#include "catlass/gemm/kda_gemm_helper.hpp" -#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" -#include "catlass/gemm/tile/kda_tile_mmad.hpp" -#include "catlass/kda_gemm_coord.hpp" - -namespace Catlass::Gemm::Block { - -template < - uint32_t L1A_STAGES_, - uint32_t L1B_STAGES_, - uint32_t L0C_STAGES_, - bool ENABLE_UNIT_FLAG_, - class L1TileShape_, - class L0TileShape_, - class AType_, - class BType_, - class CType_, - class BiasType_, - class TileCopy_, - class TileMmad_> -struct BlockMmad< - MmadAtlasA2DynamicSingleCoreSplitk, - L1TileShape_, - L0TileShape_, - AType_, - BType_, - CType_, - BiasType_, - TileCopy_, - TileMmad_> { -public: - // Type Aliases - using DispatchPolicy = MmadAtlasA2DynamicSingleCoreSplitk; - using ArchTag = typename DispatchPolicy::ArchTag; - using ElementA = typename AType_::Element; - using LayoutA = typename AType_::Layout; - using ElementB = typename BType_::Element; - using LayoutB = typename BType_::Layout; - using ElementC = typename CType_::Element; - using LayoutC = typename CType_::Layout; - using TileMmad = TileMmad_; - using CopyGmToL1A = typename TileCopy_::CopyGmToL1A; - using CopyGmToL1B = typename TileCopy_::CopyGmToL1B; - using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; - using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; - using CopyL0CToGm = typename TileCopy_::CopyL0CToGm; - using ElementAccumulator = typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; - using LayoutAInL1 = typename CopyL1ToL0A::LayoutSrc; - using LayoutBInL1 = typename CopyL1ToL0B::LayoutSrc; - using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst; - using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst; - using LayoutCInL0 = layout::zN; - - using L1AAlignHelper = Gemm::helper::L1AlignHelper; - using L1BAlignHelper = Gemm::helper::L1AlignHelper; - - static constexpr uint32_t L1A_STAGES = DispatchPolicy::L1A_STAGES; - static constexpr uint32_t L1B_STAGES = DispatchPolicy::L1B_STAGES; - static constexpr uint32_t L0AB_STAGES = DispatchPolicy::L0AB_STAGES; - static constexpr uint32_t L0C_STAGES = DispatchPolicy::L0C_STAGES; - static constexpr bool ENABLE_UNIT_FLAG = DispatchPolicy::ENABLE_UNIT_FLAG; - - static constexpr uint32_t L0A_PINGPONG_BUF_SIZE = ArchTag::L0A_SIZE / L0AB_STAGES; - static constexpr uint32_t L0B_PINGPONG_BUF_SIZE = ArchTag::L0B_SIZE / L0AB_STAGES; - static constexpr uint32_t L0C_PINGPONG_BUF_SIZE = ArchTag::L0C_SIZE / L0C_STAGES; - - // Check LayoutC - static_assert(std::is_same_v, "LayoutC only support RowMajor yet!"); - - /// Construct - CATLASS_DEVICE - BlockMmad( - GemmCoord const &l1TileShape_, - GemmCoord const &l0TileShape_, - Arch::Resource &resource, - uint32_t l1BufAddrStart = 0) - : l1TileShape(l1TileShape_), - l0TileShape(l0TileShape_) - { - if ASCEND_IS_AIV { - return; - } - - uint32_t l1ASize = l1TileShape.m() * l1TileShape.k() * sizeof(ElementA); - uint32_t l1BSize = l1TileShape.k() * l1TileShape.n() * sizeof(ElementB); - - uint32_t l1AOffset = l1BufAddrStart; - uint32_t l1BOffset = l1BufAddrStart + l1ASize * L1A_STAGES; - // Init buffers - for (uint32_t i = 0; i < L1A_STAGES; i++) { - l1ATensorList[i] = resource.l1Buf.template GetBufferByByte(l1AOffset + l1ASize * i); - l1AEventList[i] = i; - AscendC::SetFlag(l1AEventList[i]); - } - for (uint32_t i = 0; i < L1B_STAGES; i++) { - l1BTensorList[i] = resource.l1Buf.template GetBufferByByte(l1BOffset + l1BSize * i); - l1BEventList[i] = i + L1A_STAGES; - AscendC::SetFlag(l1BEventList[i]); - } - - for (uint32_t i = 0; i < L0AB_STAGES; i++) { - l0ATensorList[i] = resource.l0ABuf.template GetBufferByByte(L0A_PINGPONG_BUF_SIZE * i); - l0BTensorList[i] = resource.l0BBuf.template GetBufferByByte(L0B_PINGPONG_BUF_SIZE * i); - l0AEventList[i] = i; - l0BEventList[i] = i + L0AB_STAGES; - AscendC::SetFlag(l0AEventList[i]); - AscendC::SetFlag(l0BEventList[i]); - } - for (uint32_t i = 0; i < L0C_STAGES; i++) { - l0CTensorList[i] = resource.l0CBuf.template GetBufferByByte(L0C_PINGPONG_BUF_SIZE * i); - l0CEventList[i] = i; - AscendC::SetFlag(l0CEventList[i]); - } - } - - /// Destructor - CATLASS_DEVICE - ~BlockMmad() - { - if ASCEND_IS_AIV { - return; - } - - for (uint32_t i = 0; i < L1A_STAGES; i++) { - AscendC::WaitFlag(l1AEventList[i]); - } - for (uint32_t i = 0; i < L1B_STAGES; i++) { - AscendC::WaitFlag(l1BEventList[i]); - } - - for (uint32_t i = 0; i < L0AB_STAGES; i++) { - AscendC::WaitFlag(l0AEventList[i]); - AscendC::WaitFlag(l0BEventList[i]); - } - for (uint32_t i = 0; i < L0C_STAGES; i++) { - AscendC::WaitFlag(l0CEventList[i]); - } - } - - /// Perform a block-scoped matrix multiply-accumulate - CATLASS_DEVICE - void operator()( - AscendC::GlobalTensor const &gmBlockA, - LayoutA const &layoutA, - AscendC::GlobalTensor const &gmBlockB, - LayoutB const &layoutB, - AscendC::GlobalTensor const &gmBlockC, - LayoutC const &layoutC, - AscendC::GlobalTensor const &gmNextBlockA, - AscendC::GlobalTensor const &gmNextBlockB, - GemmCoord const &actualShape, - GemmCoord const &actualShapeNext, - bool needLoadNextA, - bool needLoadNextB, - bool atomicAdd) - { - auto layoutAInL1 = LayoutAInL1::template MakeLayout(l1TileShape.m(), l1TileShape.k()); - auto layoutBInL1 = LayoutBInL1::template MakeLayout(l1TileShape.k(), l1TileShape.n()); - - bool isFirstBlock = !preLoadedA && !preLoadedB; - bool isLastBlock = !needLoadNextA && !needLoadNextB; - uint32_t l1AListIdNext = (l1AListId + 1 < L1A_STAGES) ? (l1AListId + 1) : 0; - uint32_t l1BListIdNext = (l1BListId + 1 < L1B_STAGES) ? (l1BListId + 1) : 0; - - if (isFirstBlock) { - // Load matrix A tile from GM to L1 - AscendC::WaitFlag(l1AEventList[l1AListId]); - auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShape.m(), actualShape.k())); - copyGmToL1A(l1ATensorList[l1AListId], gmBlockA, layoutAInL1, layoutTileA); - AscendC::SetFlag(l1AEventList[l1AListId]); - - // Load matrix B tile from GM to L1 - AscendC::WaitFlag(l1BEventList[l1BListId]); - auto layoutTileB = layoutB.GetTileLayout(MakeCoord(actualShape.k(), actualShape.n())); - copyGmToL1B(l1BTensorList[l1BListId], gmBlockB, layoutBInL1, layoutTileB); - AscendC::SetFlag(l1BEventList[l1BListId]); - } - - if (needLoadNextA && (L1A_STAGES > 1)) { - AscendC::WaitFlag(l1AEventList[l1AListIdNext]); - auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShapeNext.m(), actualShapeNext.k())); - copyGmToL1A(l1ATensorList[l1AListIdNext], gmNextBlockA, layoutAInL1, layoutTileA); - AscendC::SetFlag(l1AEventList[l1AListIdNext]); - } - if (needLoadNextB && (L1B_STAGES > 1)) { - AscendC::WaitFlag(l1BEventList[l1BListIdNext]); - auto layoutTileB = layoutB.GetTileLayout(MakeCoord(actualShapeNext.k(), actualShapeNext.n())); - copyGmToL1B(l1BTensorList[l1BListIdNext], gmNextBlockB, layoutBInL1, layoutTileB); - AscendC::SetFlag(l1BEventList[l1BListIdNext]); - } - - // Get L1 tensor for current stage - auto l1ATensor = l1ATensorList[l1AListId]; - auto l1BTensor = l1BTensorList[l1BListId]; - - // Atomic set once in mmad level - if (atomicAdd) { - AscendC::SetAtomicAdd(); - } else { - AscendC::SetAtomicNone(); - } - - uint32_t l1mLoops = CeilDiv(actualShape.m(), l0TileShape.m()); - uint32_t l1nLoops = CeilDiv(actualShape.n(), l0TileShape.n()); - for (uint32_t l1mLoopIdx = 0; l1mLoopIdx < l1mLoops; l1mLoopIdx++) { - for (uint32_t l1nLoopIdx = 0; l1nLoopIdx < l1nLoops; l1nLoopIdx++) { - // layoutAInL0 - uint32_t mActual = (l1mLoopIdx < l1mLoops - 1) ? l0TileShape.m() : (actualShape.m() - l1mLoopIdx * l0TileShape.m()); - uint32_t nActual = (l1nLoopIdx < l1nLoops - 1) ? l0TileShape.n() : (actualShape.n() - l1nLoopIdx * l0TileShape.n()); - uint32_t mRound = RoundUp(mActual); - uint32_t nRound = RoundUp(nActual); - auto layoutInL0C = LayoutCInL0::MakeLayoutInL0C(MakeCoord(mRound, nRound)); - - uint32_t l0ABufId = 0; - uint32_t l0BBufId = 0; - - // wait/set conditions in block/l1mLoops/l1nLoops level - bool mnLoopStart = (l1mLoopIdx == 0 && l1nLoopIdx == 0); - bool mnLoopEnd = (l1mLoopIdx == l1mLoops - 1) && (l1nLoopIdx == l1nLoops - 1); - bool waitFlagA = (isFirstBlock || preLoadedA) && mnLoopStart; - bool setFlagA = (needLoadNextA || isLastBlock) && mnLoopEnd; - bool waitFlagB = (isFirstBlock || preLoadedB) && mnLoopStart; - bool setFlagB = (needLoadNextB || isLastBlock) && mnLoopEnd; - - if constexpr (!ENABLE_UNIT_FLAG) { - AscendC::WaitFlag(l0CEventList[l0CListId]); - } - - // Get the loop nums on L0 - uint32_t l1kLoops = CeilDiv(actualShape.k(), l0TileShape.k()); - for (uint32_t l1kLoopIdx = 0; l1kLoopIdx < l1kLoops; l1kLoopIdx++) { - uint32_t kActual = (l1kLoopIdx < l1kLoops - 1) ? l0TileShape.k() : (actualShape.k() - l1kLoopIdx * l0TileShape.k()); - - // Locate the current tile on L0A - auto l0ATile = l0ATensorList[l0ABufId]; - LayoutAInL0 layoutAInL0 = LayoutAInL0::template MakeLayout(mRound, kActual); - // Locate the current tile of matrix A on L1 - MatrixCoord l1ACoord{l1mLoopIdx * l0TileShape.m(), l1kLoopIdx * l0TileShape.k()}; - auto l1ATile = l1ATensor[layoutAInL1.GetOffset(l1ACoord)]; - - AscendC::WaitFlag(l0AEventList[l0ABufId]); - if (waitFlagA && (l1kLoopIdx == 0)) { - AscendC::WaitFlag(l1AEventList[l1AListId]); - } - - copyL1ToL0A(l0ATile, l1ATile, layoutAInL0, layoutAInL1); - - if (setFlagA && (l1kLoopIdx == l1kLoops - 1)) { - AscendC::SetFlag(l1AEventList[l1AListId]); - } - - // Locate the current tile on L0B - auto l0BTile = l0BTensorList[l0BBufId]; - LayoutBInL0 layoutBInL0 = LayoutBInL0::template MakeLayout(kActual, nRound); - // Locate the current tile of matrix B on L1 - MatrixCoord l1BCoord{l1kLoopIdx * l0TileShape.k(), l1nLoopIdx * l0TileShape.n()}; - auto l1BTile = l1BTensor[layoutBInL1.GetOffset(l1BCoord)]; - - // Wait for mmad finished - AscendC::WaitFlag(l0BEventList[l0BBufId]); - // If the current tile is the first one on the k&n axis, wait for loading matrix B from GM to L1 - if (waitFlagB && (l1kLoopIdx == 0)) { - AscendC::WaitFlag(l1BEventList[l1BListId]); - } - - // Load current tile from L1 to L0B - copyL1ToL0B(l0BTile, l1BTile, layoutBInL0, layoutBInL1); - - // If the current tile is the last one on the k&n axis, notify to load matrix B from GM to L1 - if (setFlagB && (l1kLoopIdx == l1kLoops - 1)) { - AscendC::SetFlag(l1BEventList[l1BListId]); - } - // Notify to do mmad - AscendC::SetFlag(l0CEventList[l0CListId]); - - // Locate the current tile on L0C - auto l0CTile = l0CTensorList[l0CListId]; - - // Compute the matrix multiplication on L0A and L0B and write the result to the accumulator - // Wait for loading L0B - AscendC::WaitFlag(l0CEventList[l0CListId]); - - // If the current tile is the first tile on the k axis, the accumulator needs to be reset to 0 - bool initC = (l1kLoopIdx == 0); - // If the unit flag is enabled, the unit flag is set according to the calculation progress - uint8_t unitFlag = 0b00; - if constexpr (ENABLE_UNIT_FLAG) { - if (l1kLoopIdx == l1kLoops - 1) { - unitFlag = 0b11; - } else { - unitFlag = 0b10; - } - } - // Perform calculation operations - tileMmad(l0CTile, l0ATile, l0BTile, mRound, nRound, kActual, initC, unitFlag); - - // Notify to move the next L0B tile - AscendC::SetFlag(l0BEventList[l0BBufId]); - l0BBufId = (l0BBufId + 1 < L0AB_STAGES) ? (l0BBufId + 1) : 0; - - AscendC::SetFlag(l0AEventList[l0ABufId]); - l0ABufId = (l0ABufId + 1 < L0AB_STAGES) ? (l0ABufId + 1) : 0; - } - // copy block out - LayoutC layoutBlock = layoutC.GetTileLayout(MakeCoord(mActual, nActual)); - MatrixCoord gmCoordC{l1mLoopIdx * l0TileShape.m(), l1nLoopIdx * l0TileShape.n()}; - auto gmTileC = gmBlockC[layoutC.GetOffset(gmCoordC)]; - if constexpr (!ENABLE_UNIT_FLAG) { - AscendC::SetFlag(l0CEventList[l0CListId]); - AscendC::WaitFlag(l0CEventList[l0CListId]); - copyL0CToGm(gmTileC, l0CTensorList[l0CListId], layoutBlock, layoutInL0C); - AscendC::SetFlag(l0CEventList[l0CListId]); - } else { - copyL0CToGm(gmTileC, l0CTensorList[l0CListId], layoutBlock, layoutInL0C, 0b11); - } - l0CListId = (l0CListId + 1 < L0C_STAGES) ? (l0CListId + 1) : 0; - } - } - - // if L1A_STAGES/L1B_STAGES == 1, load next block after compute - if (needLoadNextA && (L1A_STAGES == 1)) { - // Load first matrix A tile from GM to L1 - AscendC::WaitFlag(l1AEventList[l1AListIdNext]); - auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShapeNext.m(), actualShapeNext.k())); - copyGmToL1A(l1ATensorList[l1AListIdNext], gmNextBlockA, layoutAInL1, layoutTileA); - AscendC::SetFlag(l1AEventList[l1AListIdNext]); - } - if (needLoadNextB && (L1B_STAGES == 1)) { - // Load first matrix B tile from GM to L1 - AscendC::WaitFlag(l1BEventList[l1BListIdNext]); - auto layoutTileB = layoutB.GetTileLayout(MakeCoord(actualShapeNext.k(), actualShapeNext.n())); - copyGmToL1B(l1BTensorList[l1BListIdNext], gmNextBlockB, layoutBInL1, layoutTileB); - AscendC::SetFlag(l1BEventList[l1BListIdNext]); - } - - preLoadedA = needLoadNextA; - preLoadedB = needLoadNextB; - l1AListId = needLoadNextA ? l1AListIdNext : l1AListId; - l1BListId = needLoadNextB ? l1BListIdNext : l1BListId; - } - -protected: - /// Data members - AscendC::LocalTensor l1ATensorList[L1A_STAGES]; - AscendC::LocalTensor l1BTensorList[L1B_STAGES]; - AscendC::LocalTensor l0ATensorList[L0AB_STAGES]; - AscendC::LocalTensor l0BTensorList[L0AB_STAGES]; - AscendC::LocalTensor l0CTensorList[L0C_STAGES]; - - int32_t l1AEventList[L1A_STAGES]; - int32_t l1BEventList[L1B_STAGES]; - int32_t l0AEventList[L0AB_STAGES]; - int32_t l0BEventList[L0AB_STAGES]; - int32_t l0CEventList[L0C_STAGES]; - - bool preLoadedA{false}; - bool preLoadedB{false}; - uint32_t l1AListId{0}; - uint32_t l1BListId{0}; - uint32_t l0CListId{0}; - GemmCoord l1TileShape; - GemmCoord l0TileShape; - - TileMmad tileMmad; - CopyGmToL1A copyGmToL1A; - CopyGmToL1B copyGmToL1B; - CopyL1ToL0A copyL1ToL0A; - CopyL1ToL0B copyL1ToL0B; - CopyL0CToGm copyL0CToGm; -}; - -} // namespace Catlass::Gemm::Block - -#endif // CATLASS_GEMM_BLOCK_MMAD_DYNAMIC_SINGLE_CORE_SPLITK_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_dynamic_small.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_dynamic_small.hpp deleted file mode 100644 index 37449a8dae..0000000000 --- a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_dynamic_small.hpp +++ /dev/null @@ -1,284 +0,0 @@ -/** - * Copyright (c) 2025 Huawei Technologies Co., Ltd. - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of - * CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ - -#ifndef CATLASS_GEMM_BLOCK_BLOCK_MMAD_DYNAMIC_SMALL_HPP -#define CATLASS_GEMM_BLOCK_BLOCK_MMAD_DYNAMIC_SMALL_HPP - -#include "catlass/kda_catlass.hpp" -#include "catlass/arch/kda_resource.hpp" -#include "catlass/kda_coord.hpp" -#include "catlass/kda_gemm_coord.hpp" -#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" -#include "catlass/gemm/kda_gemm_helper.hpp" -#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" -#include "catlass/gemm/tile/kda_tile_mmad.hpp" - -namespace Catlass::Gemm::Block { - -template -struct BlockMmad, L1TileShape_, L0TileShape_, - AType_, BType_, CType_, BiasType_, TileCopy_, TileMmad_> { -public: - // Type Aliases - using DispatchPolicy = MmadAtlasA2DynamicSmall; - using ArchTag = typename DispatchPolicy::ArchTag; - - using ElementA = typename AType_::Element; - using LayoutA = typename AType_::Layout; - using ElementB = typename BType_::Element; - using LayoutB = typename BType_::Layout; - using ElementC = typename CType_::Element; - using LayoutC = typename CType_::Layout; - using TileMmad = TileMmad_; - using CopyGmToL1A = typename TileCopy_::CopyGmToL1A; - using CopyGmToL1B = typename TileCopy_::CopyGmToL1B; - using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; - using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; - using CopyL0CToGm = typename TileCopy_::CopyL0CToGm; - using ElementAccumulator = - typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; - using LayoutAInL1 = typename CopyL1ToL0A::LayoutSrc; - using LayoutBInL1 = typename CopyL1ToL0B::LayoutSrc; - using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst; - using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst; - using LayoutCInL0 = layout::zN; - - using L1AAlignHelper = Gemm::helper::L1AlignHelper; - using L1BAlignHelper = Gemm::helper::L1AlignHelper; - - static constexpr bool ENABLE_UNIT_FLAG = DispatchPolicy::ENABLE_UNIT_FLAG; - static constexpr uint32_t STAGES = DispatchPolicy::STAGES; - - static constexpr uint32_t L0A_PINGPONG_BUF_SIZE = ArchTag::L0A_SIZE / STAGES; - static constexpr uint32_t L0B_PINGPONG_BUF_SIZE = ArchTag::L0B_SIZE / STAGES; - - // Check LayoutC - static_assert(std::is_same_v, "LayoutC only support RowMajor yet!"); - - /// Construct - CATLASS_DEVICE - BlockMmad(GemmCoord const &l1TileShape_, Arch::Resource &resource, uint32_t l1BufAddrStart = 0) - : l1TileShape(l1TileShape_) - { - uint32_t l1ASize = l1TileShape.m() * l1TileShape.k() * sizeof(ElementA); - uint32_t l1BSize = l1TileShape.k() * l1TileShape.n() * sizeof(ElementB); - - kPartLenMax = min(L0A_PINGPONG_BUF_SIZE / sizeof(ElementA) / l1TileShape.m() / L1AAlignHelper::ELE_NUM_PER_C0 * - L1AAlignHelper::ELE_NUM_PER_C0, - L0B_PINGPONG_BUF_SIZE / sizeof(ElementB) / l1TileShape.n() / L1BAlignHelper::ELE_NUM_PER_C0 * - L1BAlignHelper::ELE_NUM_PER_C0); - - if constexpr (std::is_same_v && std::is_same_v) { - kPartLenMax = RoundDown(kPartLenMax); - } - - uint32_t l1AOffset = l1BufAddrStart; - uint32_t l1BOffset = l1BufAddrStart + l1ASize * STAGES; - // Init buffers - for (uint32_t i = 0; i < STAGES; i++) { - l1ATensorList[i] = resource.l1Buf.template GetBufferByByte(l1AOffset + l1ASize * i); - l1BTensorList[i] = resource.l1Buf.template GetBufferByByte(l1BOffset + l1BSize * i); - l0ATensorList[i] = resource.l0ABuf.template GetBufferByByte(L0A_PINGPONG_BUF_SIZE * i); - l0BTensorList[i] = resource.l0BBuf.template GetBufferByByte(L0B_PINGPONG_BUF_SIZE * i); - - l1AEventList[i] = i; - l1BEventList[i] = i + STAGES; - l0AEventList[i] = i; - l0BEventList[i] = i + STAGES; - AscendC::SetFlag(l1AEventList[i]); - AscendC::SetFlag(l1BEventList[i]); - AscendC::SetFlag(l0AEventList[i]); - AscendC::SetFlag(l0BEventList[i]); - } - l0CTensor = resource.l0CBuf.template GetBufferByByte(0); - if constexpr (!ENABLE_UNIT_FLAG) { - AscendC::SetFlag(EVENT_ID0); - } - } - - /// Destructor - CATLASS_DEVICE - ~BlockMmad() - { - for (uint32_t i = 0; i < STAGES; i++) { - AscendC::WaitFlag(l1AEventList[i]); - AscendC::WaitFlag(l1BEventList[i]); - AscendC::WaitFlag(l0AEventList[i]); - AscendC::WaitFlag(l0BEventList[i]); - } - if constexpr (!ENABLE_UNIT_FLAG) { - AscendC::WaitFlag(EVENT_ID0); - } - } - - /// Perform a block-scoped matrix multiply-accumulate - CATLASS_DEVICE - void operator()(AscendC::GlobalTensor const &gmBlockA, LayoutA const &layoutA, - AscendC::GlobalTensor const &gmBlockB, LayoutB const &layoutB, - AscendC::GlobalTensor const &gmBlockC, LayoutC const &layoutC, - GemmCoord const &actualShape) - { - uint32_t mRound = RoundUp(actualShape.m()); - uint32_t nRound = RoundUp(actualShape.n()); - - auto layoutAInL1 = LayoutAInL1::template MakeLayout(l1TileShape.m(), l1TileShape.k()); - auto layoutBInL1 = LayoutBInL1::template MakeLayout(l1TileShape.k(), l1TileShape.n()); - auto layoutInL0C = LayoutCInL0::MakeLayoutInL0C(MakeCoord(mRound, nRound)); - - if constexpr (!ENABLE_UNIT_FLAG) { - AscendC::WaitFlag(EVENT_ID0); - } - - uint32_t kActual = actualShape.k(); - - // k loop - { - // For small-shape matmul, the k-loop executes only once. - - // load first matrix A tile from GM to L1 - AscendC::WaitFlag(l1AEventList[l1ListId]); - auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShape.m(), kActual)); - copyGmToL1A(l1ATensorList[l1ListId], gmBlockA, layoutAInL1, layoutTileA); - AscendC::SetFlag(l1AEventList[l1ListId]); - - // load first matrix B tile from GM to L1 - AscendC::WaitFlag(l1BEventList[l1ListId]); - auto layoutTileB = layoutB.GetTileLayout(MakeCoord(kActual, actualShape.n())); - copyGmToL1B(l1BTensorList[l1ListId], gmBlockB, layoutBInL1, layoutTileB); - AscendC::SetFlag(l1BEventList[l1ListId]); - - // Get L1 tensor for current stage - auto l1ATensor = l1ATensorList[l1ListId]; - auto l1BTensor = l1BTensorList[l1ListId]; - - // Get the loop nums on L0 - uint32_t kPartLoop = CeilDiv(kActual, kPartLenMax); - - for (int kPartIdx = 0; kPartIdx < kPartLoop; kPartIdx++) { - uint32_t kPartActual = (kPartIdx < kPartLoop - 1) ? kPartLenMax : (kActual - kPartIdx * kPartLenMax); - - // Locate the current tile on L0A - auto l0ATile = l0ATensorList[l0AListId]; - LayoutAInL0 layoutAInL0 = LayoutAInL0::template MakeLayout(mRound, kPartActual); - // Locate the current tile of matrix A on L1 - MatrixCoord l1AOffset{0, kPartIdx * kPartLenMax}; - auto l1ATile = l1ATensor[layoutAInL1.GetOffset(l1AOffset)]; - - AscendC::WaitFlag(l0AEventList[l0AListId]); - if (kPartIdx == 0) { - AscendC::WaitFlag(l1AEventList[l1ListId]); - } - - // Load current tile from L1 to L0A - copyL1ToL0A(l0ATile, l1ATile, layoutAInL0, layoutAInL1); - - if (kPartIdx == kPartLoop - 1) { - AscendC::SetFlag(l1AEventList[l1ListId]); - } - - // Locate the current tile on L0B - auto l0BTile = l0BTensorList[l0BListId]; - LayoutBInL0 layoutBInL0 = LayoutBInL0::template MakeLayout(kPartActual, nRound); - // Locate the current tile of matrix B on L1 - MatrixCoord l1BOffset{kPartIdx * kPartLenMax, 0}; - auto l1BTile = l1BTensor[layoutBInL1.GetOffset(l1BOffset)]; - - // Wait for mmad finished - AscendC::WaitFlag(l0BEventList[l0BListId]); - // If the current tile is the first one on the k&n axis, wait for loading matrix B from GM to L1 - if (kPartIdx == 0) { - AscendC::WaitFlag(l1BEventList[l1ListId]); - } - - // Load current tile from L1 to L0B - copyL1ToL0B(l0BTile, l1BTile, layoutBInL0, layoutBInL1); - - // If the current tile is the last one on the k&n axis, notify to load matrix B from GM to L1 - if (kPartIdx == kPartLoop - 1) { - AscendC::SetFlag(l1BEventList[l1ListId]); - } - // Notify to do mmad - AscendC::SetFlag(EVENT_ID0); - - // Locate the current tile on L0C - auto l0CTile = l0CTensor; - - // Compute the matrix multiplication on L0A and L0B and write the result to the accumulator - // Wait for loading L0B - AscendC::WaitFlag(EVENT_ID0); - - // If the current tile is the first tile on the k axis, the accumulator needs to be reset to 0 - bool initC = (kPartIdx == 0); - // If the unit flag is enabled, the unit flag is set according to the calculation progress - uint8_t unitFlag = 0b00; - if constexpr (ENABLE_UNIT_FLAG) { - if (kPartIdx == kPartLoop - 1) { - unitFlag = 0b11; - } else { - unitFlag = 0b10; - } - } - // Perform calculation operations - tileMmad(l0CTile, l0ATile, l0BTile, mRound, nRound, kPartActual, initC, unitFlag); - - // Notify to move the next L0B tile - AscendC::SetFlag(l0BEventList[l0BListId]); - l0BListId = (l0BListId + 1 < STAGES) ? (l0BListId + 1) : 0; - - AscendC::SetFlag(l0AEventList[l0AListId]); - l0AListId = (l0AListId + 1 < STAGES) ? (l0AListId + 1) : 0; - } - } - - // copy block out - LayoutC layoutBlock = layoutC.GetTileLayout(actualShape.GetCoordMN()); - - if constexpr (!ENABLE_UNIT_FLAG) { - AscendC::SetFlag(EVENT_ID0); - AscendC::WaitFlag(EVENT_ID0); - copyL0CToGm(gmBlockC, l0CTensor, layoutBlock, layoutInL0C); - AscendC::SetFlag(EVENT_ID0); - } else { - copyL0CToGm(gmBlockC, l0CTensor, layoutBlock, layoutInL0C, 0b11); - } - } - -protected: - /// Data members - AscendC::LocalTensor l1ATensorList[STAGES]; - AscendC::LocalTensor l1BTensorList[STAGES]; - AscendC::LocalTensor l0ATensorList[STAGES]; - AscendC::LocalTensor l0BTensorList[STAGES]; - AscendC::LocalTensor l0CTensor; - - int32_t l1AEventList[STAGES]; - int32_t l1BEventList[STAGES]; - int32_t l0AEventList[STAGES]; - int32_t l0BEventList[STAGES]; - - uint32_t l1ListId{0}; - uint32_t l0AListId{0}; - uint32_t l0BListId{0}; - - GemmCoord l1TileShape; - uint32_t kPartLenMax; - - TileMmad tileMmad; - CopyGmToL1A copyGmToL1A; - CopyGmToL1B copyGmToL1B; - CopyL1ToL0A copyL1ToL0A; - CopyL1ToL0B copyL1ToL0B; - CopyL0CToGm copyL0CToGm; -}; - -} // namespace Catlass::Gemm::Block - -#endif // CATLASS_GEMM_BLOCK_BLOCK_MMAD_DYNAMIC_SMALL_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_dynamic_streamk.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_dynamic_streamk.hpp deleted file mode 100644 index e5e6933ef4..0000000000 --- a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_dynamic_streamk.hpp +++ /dev/null @@ -1,370 +0,0 @@ -/** - * Copyright (c) 2025 Huawei Technologies Co., Ltd. - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of - * CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ - -#ifndef CATLASS_GEMM_BLOCK_BLOCK_MMAD_DYNAMIC_STREAMK_HPP -#define CATLASS_GEMM_BLOCK_BLOCK_MMAD_DYNAMIC_STREAMK_HPP - -#include "catlass/kda_catlass.hpp" -#include "catlass/arch/kda_resource.hpp" -#include "catlass/kda_coord.hpp" -#include "catlass/kda_gemm_coord.hpp" -#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" -#include "catlass/gemm/kda_gemm_helper.hpp" -#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" -#include "catlass/gemm/tile/kda_tile_mmad.hpp" - -namespace Catlass::Gemm::Block { - -template -struct BlockMmad, L1TileShape_, L0TileShape_, AType_, - BType_, CType_, BiasType_, TileCopy_, TileMmad_> { -public: - // Type Aliases - using DispatchPolicy = MmadAtlasA2DynamicStreamk; - using ArchTag = typename DispatchPolicy::ArchTag; - using ElementA = typename AType_::Element; - using LayoutA = typename AType_::Layout; - using ElementB = typename BType_::Element; - using LayoutB = typename BType_::Layout; - using ElementC = typename CType_::Element; - using LayoutC = typename CType_::Layout; - using TileMmad = TileMmad_; - using CopyGmToL1A = typename TileCopy_::CopyGmToL1A; - using CopyGmToL1B = typename TileCopy_::CopyGmToL1B; - using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; - using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; - using CopyL0CToGmNormalBlock = typename TileCopy_::CopyL0CToGm; - using CopyL0CToGmStreamkBlock = Gemm::Tile::CopyL0CToGm>; - using ElementAccumulator = - typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; - using LayoutAInL1 = typename CopyL1ToL0A::LayoutSrc; - using LayoutBInL1 = typename CopyL1ToL0B::LayoutSrc; - using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst; - using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst; - using LayoutCInL0 = layout::zN; - - using L1AAlignHelper = Gemm::helper::L1AlignHelper; - using L1BAlignHelper = Gemm::helper::L1AlignHelper; - - static constexpr bool ENABLE_UNIT_FLAG = DispatchPolicy::ENABLE_UNIT_FLAG; - static constexpr bool ENABLE_SHUFFLE_K = DispatchPolicy::ENABLE_SHUFFLE_K; - static constexpr uint32_t STAGES = DispatchPolicy::STAGES; - - static constexpr uint32_t L0A_PINGPONG_BUF_SIZE = ArchTag::L0A_SIZE / STAGES; - static constexpr uint32_t L0B_PINGPONG_BUF_SIZE = ArchTag::L0B_SIZE / STAGES; - - // Check LayoutC - static_assert(std::is_same_v, "LayoutC only support RowMajor yet!"); - - /// Construct - CATLASS_DEVICE - BlockMmad(GemmCoord const &l1TileShape_, Arch::Resource &resource, uint32_t l1BufAddrStart = 0) - : l1TileShape(l1TileShape_) - { - uint32_t l1ASize = l1TileShape.m() * l1TileShape.k() * sizeof(ElementA); - uint32_t l1BSize = l1TileShape.k() * l1TileShape.n() * sizeof(ElementB); - - kPartLenMax = min(L0A_PINGPONG_BUF_SIZE / sizeof(ElementA) / l1TileShape.m() / L1AAlignHelper::ELE_NUM_PER_C0 * - L1AAlignHelper::ELE_NUM_PER_C0, - L0B_PINGPONG_BUF_SIZE / sizeof(ElementB) / l1TileShape.n() / L1BAlignHelper::ELE_NUM_PER_C0 * - L1BAlignHelper::ELE_NUM_PER_C0); - - if constexpr (std::is_same_v && std::is_same_v) { - kPartLenMax = RoundDown(kPartLenMax); - } - - uint32_t l1AOffset = l1BufAddrStart; - uint32_t l1BOffset = l1BufAddrStart + l1ASize * STAGES; - // Init buffers - for (uint32_t i = 0; i < STAGES; i++) { - l1ATensorList[i] = resource.l1Buf.template GetBufferByByte(l1AOffset + l1ASize * i); - l1BTensorList[i] = resource.l1Buf.template GetBufferByByte(l1BOffset + l1BSize * i); - l0ATensorList[i] = resource.l0ABuf.template GetBufferByByte(L0A_PINGPONG_BUF_SIZE * i); - l0BTensorList[i] = resource.l0BBuf.template GetBufferByByte(L0B_PINGPONG_BUF_SIZE * i); - - l1AEventList[i] = i; - l1BEventList[i] = i + STAGES; - l0AEventList[i] = i; - l0BEventList[i] = i + STAGES; - AscendC::SetFlag(l1AEventList[i]); - AscendC::SetFlag(l1BEventList[i]); - AscendC::SetFlag(l0AEventList[i]); - AscendC::SetFlag(l0BEventList[i]); - } - l0CTensor = resource.l0CBuf.template GetBufferByByte(0); - if constexpr (!ENABLE_UNIT_FLAG) { - AscendC::SetFlag(EVENT_ID0); - } - } - - /// Destructor - CATLASS_DEVICE - ~BlockMmad() - { - for (uint32_t i = 0; i < STAGES; i++) { - AscendC::WaitFlag(l1AEventList[i]); - AscendC::WaitFlag(l1BEventList[i]); - AscendC::WaitFlag(l0AEventList[i]); - AscendC::WaitFlag(l0BEventList[i]); - } - if constexpr (!ENABLE_UNIT_FLAG) { - AscendC::WaitFlag(EVENT_ID0); - } - } - - /// Perform a block-scoped matrix multiply-accumulate - CATLASS_DEVICE - void operator()( - AscendC::GlobalTensor const &gmBlockA, LayoutA const &layoutA, - AscendC::GlobalTensor const &gmBlockB, LayoutB const &layoutB, - AscendC::GlobalTensor const &gmBlockC, LayoutC const &layoutC, - AscendC::GlobalTensor const &gmW, LayoutC const &layoutW, - AscendC::GlobalTensor const &gmNextBlockA, AscendC::GlobalTensor const &gmNextBlockB, - GemmCoord const &actualShape, GemmCoord const &actualShapeNext, - bool isFirstBlock, bool hasNextBlock, bool isSkBlock) - { - uint32_t mRound = RoundUp(actualShape.m()); - uint32_t nRound = RoundUp(actualShape.n()); - - auto layoutAInL1 = LayoutAInL1::template MakeLayout(l1TileShape.m(), l1TileShape.k()); - auto layoutBInL1 = LayoutBInL1::template MakeLayout(l1TileShape.k(), l1TileShape.n()); - auto layoutInL0C = LayoutCInL0::MakeLayoutInL0C(MakeCoord(mRound, nRound)); - - uint32_t kTileCount = CeilDiv(actualShape.k(), l1TileShape.k()); - - if constexpr (!ENABLE_UNIT_FLAG) { - AscendC::WaitFlag(EVENT_ID0); - } - uint32_t startTileIdx = 0; - if constexpr (ENABLE_SHUFFLE_K) { - startTileIdx = AscendC::GetBlockIdx(); - } - uint32_t firstTileIdx = startTileIdx % kTileCount; - uint32_t lastTileIdx = (startTileIdx + kTileCount - 1) % kTileCount; - uint32_t kActual = - (firstTileIdx < kTileCount - 1) ? l1TileShape.k() : (actualShape.k() - firstTileIdx * l1TileShape.k()); - - // k loop - for (uint32_t kLoopIdx = 0; kLoopIdx < kTileCount; kLoopIdx++) { - uint32_t shuffleKIdx = (startTileIdx + kLoopIdx) % kTileCount; - // Load first matrix A tile in total kernel loop from GM to L1 - if (shuffleKIdx == firstTileIdx && isFirstBlock) { - MatrixCoord gmTileACoord{0, shuffleKIdx * l1TileShape.k()}; - MatrixCoord gmTileBCoord{shuffleKIdx * l1TileShape.k(), 0}; - auto gmTileA = gmBlockA[layoutA.GetOffset(gmTileACoord)]; - auto gmTileB = gmBlockB[layoutB.GetOffset(gmTileBCoord)]; - // Load first matrix A tile from GM to L1 - AscendC::WaitFlag(l1AEventList[l1ListId]); - auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShape.m(), kActual)); - copyGmToL1A(l1ATensorList[l1ListId], gmTileA, layoutAInL1, layoutTileA); - AscendC::SetFlag(l1AEventList[l1ListId]); - - // Load first matrix B tile from GM to L1 - AscendC::WaitFlag(l1BEventList[l1ListId]); - auto layoutTileB = layoutB.GetTileLayout(MakeCoord(kActual, actualShape.n())); - copyGmToL1B(l1BTensorList[l1ListId], gmTileB, layoutBInL1, layoutTileB); - AscendC::SetFlag(l1BEventList[l1ListId]); - } - - uint32_t l1ListIdNext = (l1ListId + 1 < STAGES) ? (l1ListId + 1) : 0; - uint32_t kActualNext{0}; - - // preload next tile from GM to L1 - if (shuffleKIdx != lastTileIdx) { - uint32_t shuffleKIdxNext = (startTileIdx + kLoopIdx + 1) % kTileCount; - // Get L1 tensor for next stage - auto l1ATensor = l1ATensorList[l1ListIdNext]; - auto l1BTensor = l1BTensorList[l1ListIdNext]; - // Get GM tensor for next stage - kActualNext = (shuffleKIdxNext < kTileCount - 1) ? l1TileShape.k() : (actualShape.k() - shuffleKIdxNext * l1TileShape.k()); - MatrixCoord gmTileACoord{0, shuffleKIdxNext * l1TileShape.k()}; - MatrixCoord gmTileBCoord{shuffleKIdxNext * l1TileShape.k(), 0}; - auto gmTileA = gmBlockA[layoutA.GetOffset(gmTileACoord)]; - auto gmTileB = gmBlockB[layoutB.GetOffset(gmTileBCoord)]; - - // load next matrix A tile from GM to L1 - AscendC::WaitFlag(l1AEventList[l1ListIdNext]); - auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShape.m(), kActualNext)); - copyGmToL1A(l1ATensor, gmTileA, layoutAInL1, layoutTileA); - AscendC::SetFlag(l1AEventList[l1ListIdNext]); - - // load next matrix B tile from GM to L1 - AscendC::WaitFlag(l1BEventList[l1ListIdNext]); - auto layoutTileB = layoutB.GetTileLayout(MakeCoord(kActualNext, actualShape.n())); - copyGmToL1B(l1BTensor, gmTileB, layoutBInL1, layoutTileB); - AscendC::SetFlag(l1BEventList[l1ListIdNext]); - } - if (shuffleKIdx == lastTileIdx && hasNextBlock) { - uint32_t kTileCount = CeilDiv(actualShapeNext.k(), l1TileShape.k()); - uint32_t firstTileIdx = startTileIdx % kTileCount; - // Get L1 tensor for next stage - auto l1ATensor = l1ATensorList[l1ListIdNext]; - auto l1BTensor = l1BTensorList[l1ListIdNext]; - // Get GM tensor for next stage - kActualNext = (firstTileIdx < kTileCount - 1) ? l1TileShape.k() : (actualShapeNext.k() - firstTileIdx * l1TileShape.k()); - MatrixCoord gmTileACoord{0, firstTileIdx * l1TileShape.k()}; - MatrixCoord gmTileBCoord{firstTileIdx * l1TileShape.k(), 0}; - auto gmTileA = gmNextBlockA[layoutA.GetOffset(gmTileACoord)]; - auto gmTileB = gmNextBlockB[layoutB.GetOffset(gmTileBCoord)]; - // load next matrix A tile from GM to L1 - AscendC::WaitFlag(l1AEventList[l1ListIdNext]); - auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShapeNext.m(), kActualNext)); - copyGmToL1A(l1ATensor, gmTileA, layoutAInL1, layoutTileA); - AscendC::SetFlag(l1AEventList[l1ListIdNext]); - - // load next matrix B tile from GM to L1 - AscendC::WaitFlag(l1BEventList[l1ListIdNext]); - auto layoutTileB = layoutB.GetTileLayout(MakeCoord(kActualNext, actualShapeNext.n())); - copyGmToL1B(l1BTensor, gmTileB, layoutBInL1, layoutTileB); - AscendC::SetFlag(l1BEventList[l1ListIdNext]); - } - - // Get L1 tensor for current stage - auto l1ATensor = l1ATensorList[l1ListId]; - auto l1BTensor = l1BTensorList[l1ListId]; - - // Get the loop nums on L0 - uint32_t kPartLoop = CeilDiv(kActual, kPartLenMax); - - uint32_t l0ABufId = 0; - uint32_t l0BBufId = 0; - - for (int kPartIdx = 0; kPartIdx < kPartLoop; kPartIdx++) { - uint32_t kPartActual = - (kPartIdx < kPartLoop - 1) ? kPartLenMax : (kActual - kPartIdx * kPartLenMax); - - // Locate the current tile on L0A - auto l0ATile = l0ATensorList[l0ABufId]; - LayoutAInL0 layoutAInL0 = LayoutAInL0::template MakeLayout(mRound, kPartActual); - // Locate the current tile of matrix A on L1 - MatrixCoord l1ACoord{0, kPartIdx * kPartLenMax}; - auto l1ATile = l1ATensor[layoutAInL1.GetOffset(l1ACoord)]; - - AscendC::WaitFlag(l0AEventList[l0ABufId]); - if (kPartIdx == 0) { - AscendC::WaitFlag(l1AEventList[l1ListId]); - } - - copyL1ToL0A(l0ATile, l1ATile, layoutAInL0, layoutAInL1); - - if (kPartIdx == kPartLoop - 1) { - AscendC::SetFlag(l1AEventList[l1ListId]); - } - - // Locate the current tile on L0B - auto l0BTile = l0BTensorList[l0BBufId]; - LayoutBInL0 layoutBInL0 = LayoutBInL0::template MakeLayout(kPartActual, nRound); - // Locate the current tile of matrix B on L1 - MatrixCoord l1BCoord{kPartIdx * kPartLenMax, 0}; - auto l1BTile = l1BTensor[layoutBInL1.GetOffset(l1BCoord)]; - - // Wait for mmad finished - AscendC::WaitFlag(l0BEventList[l0BBufId]); - // If the current tile is the first one on the k&n axis, wait for loading matrix B from GM to L1 - if (kPartIdx == 0) { - AscendC::WaitFlag(l1BEventList[l1ListId]); - } - - // Load current tile from L1 to L0B - copyL1ToL0B(l0BTile, l1BTile, layoutBInL0, layoutBInL1); - - // If the current tile is the last one on the k&n axis, notify to load matrix B from GM to L1 - if (kPartIdx == kPartLoop - 1) { - AscendC::SetFlag(l1BEventList[l1ListId]); - } - // Notify to do mmad - AscendC::SetFlag(EVENT_ID0); - - // Locate the current tile on L0C - MatrixCoord l0CCoord{static_cast(0U), static_cast(0U)}; - auto l0CTile = l0CTensor[layoutInL0C.GetOffset(l0CCoord)]; - - // Compute the matrix multiplication on L0A and L0B and write the result to the accumulator - // Wait for loading L0B - AscendC::WaitFlag(EVENT_ID0); - - // If the current tile is the first tile on the k axis, the accumulator needs to be reset to 0 - bool initC = ((kLoopIdx == 0) && (kPartIdx == 0)); - // If the unit flag is enabled, the unit flag is set according to the calculation progress - uint8_t unitFlag = 0b00; - if constexpr (ENABLE_UNIT_FLAG) { - if ((kLoopIdx == kTileCount - 1) && (kPartIdx == kPartLoop - 1)) { - unitFlag = 0b11; - } else { - unitFlag = 0b10; - } - } - // Perform calculation operations - tileMmad(l0CTile, l0ATile, l0BTile, mRound, nRound, kPartActual, initC, unitFlag); - - // Notify to move the next L0B tile - AscendC::SetFlag(l0BEventList[l0BBufId]); - - l0BBufId = (l0BBufId + 1 < STAGES) ? (l0BBufId + 1) : 0; - - AscendC::SetFlag(l0AEventList[l0ABufId]); - l0ABufId = (l0ABufId + 1 < STAGES) ? (l0ABufId + 1) : 0; - } - - l1ListId = l1ListIdNext; - kActual = kActualNext; - } - - // copy block out - LayoutC layoutBlock = layoutC.GetTileLayout(actualShape.GetCoordMN()); - - if constexpr (!ENABLE_UNIT_FLAG) { - AscendC::SetFlag(EVENT_ID0); - AscendC::WaitFlag(EVENT_ID0); - if (isSkBlock) { - copyL0CToGmStreamkBlock(gmW, l0CTensor, layoutW, layoutInL0C); - } else { - copyL0CToGmNormalBlock(gmBlockC, l0CTensor, layoutBlock, layoutInL0C); - } - AscendC::SetFlag(EVENT_ID0); - } else { - if (isSkBlock) { - copyL0CToGmStreamkBlock(gmW, l0CTensor, layoutW, layoutInL0C, 0b11); - } else { - copyL0CToGmNormalBlock(gmBlockC, l0CTensor, layoutBlock, layoutInL0C, 0b11); - } - } - } - -protected: - /// Data members - AscendC::LocalTensor l1ATensorList[STAGES]; - AscendC::LocalTensor l1BTensorList[STAGES]; - AscendC::LocalTensor l0ATensorList[STAGES]; - AscendC::LocalTensor l0BTensorList[STAGES]; - AscendC::LocalTensor l0CTensor; - - int32_t l1AEventList[STAGES]; - int32_t l1BEventList[STAGES]; - int32_t l0AEventList[STAGES]; - int32_t l0BEventList[STAGES]; - - uint32_t l1ListId{0}; - GemmCoord l1TileShape; - uint32_t kPartLenMax; - - TileMmad tileMmad; - CopyGmToL1A copyGmToL1A; - CopyGmToL1B copyGmToL1B; - CopyL1ToL0A copyL1ToL0A; - CopyL1ToL0B copyL1ToL0B; - CopyL0CToGmNormalBlock copyL0CToGmNormalBlock; - CopyL0CToGmStreamkBlock copyL0CToGmStreamkBlock; -}; - -} // namespace Catlass::Gemm::Block - -#endif // CATLASS_GEMM_BLOCK_BLOCK_MMAD_DYNAMIC_STREAMK_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_pv_head_tail.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_pv_head_tail.hpp deleted file mode 100644 index 63fcde0508..0000000000 --- a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_pv_head_tail.hpp +++ /dev/null @@ -1,315 +0,0 @@ -/** - * Copyright (c) 2025 Huawei Technologies Co., Ltd. - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of - * CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ - -#ifndef CATLASS_GEMM_BLOCK_MMAD_PV_TAIL_HPP -#define CATLASS_GEMM_BLOCK_MMAD_PV_TAIL_HPP - -#include "catlass/kda_catlass.hpp" -#include "catlass/arch/kda_resource.hpp" -#include "catlass/kda_coord.hpp" -#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" -#include "catlass/gemm/kda_gemm_helper.hpp" -#include "catlass/kda_gemm_coord.hpp" -#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" -#include "catlass/gemm/tile/kda_tile_mmad.hpp" - -//////////////////////////////////////////////////////////////////// - -namespace Catlass::Gemm::Block { -//////////////////////////////////////////////////////////////////// - -template < - bool PAGED_CACHE_FLAG_, - bool ENABLE_UNIT_FLAG_, - class L1TileShape_, - class L0TileShape_, - class AType_, - class BType_, - class CType_, - class BiasType_, - class TileCopy_, - class TileMmad_> -struct BlockMmad< - MmadAtlasA2FAITailPV, - L1TileShape_, - L0TileShape_, - AType_, - BType_, - CType_, - BiasType_, - TileCopy_, - TileMmad_> { -public: - // Type Aliases - using DispatchPolicy = MmadAtlasA2FAITailPV; - using ArchTag = typename DispatchPolicy::ArchTag; - using L1TileShape = L1TileShape_; - using L0TileShape = L0TileShape_; - using ElementA = typename AType_::Element; - using LayoutA = typename AType_::Layout; - using ElementB = typename BType_::Element; - using LayoutB = typename BType_::Layout; - using ElementC = typename CType_::Element; - using LayoutC = typename CType_::Layout; - using TileMmad = TileMmad_; - using CopyGmToL1A = typename TileCopy_::CopyGmToL1A; - using CopyGmToL1B = typename TileCopy_::CopyGmToL1B; - using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; - using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; - using CopyL0CToGm = typename TileCopy_::CopyL0CToGm; - using ElementAccumulator = - typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; - using LayoutAInL1 = typename CopyL1ToL0A::LayoutSrc; - using LayoutBInL1 = typename CopyL1ToL0B::LayoutSrc; - using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst; - using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst; - using LayoutCInL0 = layout::zN; - - using L1AAlignHelper = Gemm::helper::L1AlignHelper; - using L1BAlignHelper = Gemm::helper::L1AlignHelper; - - static constexpr uint32_t STAGES = DispatchPolicy::STAGES; - static constexpr uint32_t L1A_SIZE = L1TileShape::M * L1TileShape::K * sizeof(ElementA); - static constexpr uint32_t L1B_SIZE = L1TileShape::N * L1TileShape::K * sizeof(ElementB); - static constexpr uint32_t L0A_SIZE = ArchTag::L0A_SIZE; - static constexpr uint32_t L0B_SIZE = ArchTag::L0B_SIZE; - static constexpr uint32_t L0C_SIZE = ArchTag::L0C_SIZE; - static constexpr uint32_t L0A_PINGPONG_BUF_SIZE = L0A_SIZE / STAGES; - static constexpr uint32_t L0B_PINGPONG_BUF_SIZE = L0B_SIZE / STAGES; - static constexpr uint32_t L0C_PINGPONG_BUF_SIZE = L0C_SIZE / STAGES; - static constexpr uint32_t BLOCK_SIZE = 16; - static constexpr uint32_t EMBED_SPLIT_SIZE = 128; - static constexpr uint32_t UNIT_BLOCK_STACK_NUM = 4; - static constexpr uint32_t KV_BASE_BLOCK = 512; - static constexpr uint32_t KV_SPLIT_SIZE = 128; - static constexpr uint32_t LOAB_BLOCK = 1; - - // Check LayoutC - static_assert(std::is_same_v, "LayoutC only support RowMajor yet!"); - - /// Construct - CATLASS_DEVICE - BlockMmad(Arch::Resource &resource, uint32_t l1BufAddrStart = 0) - { - // Allocate L1 memory space - l0ATensor = resource.l0ABuf.template GetBufferByByte(0); - l0BTensor = resource.l0BBuf.template GetBufferByByte(0); - for (uint32_t i = 0; i < STAGES; i++) { - l1ATensor[i] = resource.l1Buf.template GetBufferByByte(l1BufAddrStart + L1A_SIZE * LOAB_BLOCK * i); - l0CTensor[i] = resource.l0CBuf.template GetBufferByByte(L0C_PINGPONG_BUF_SIZE * i); - l1BTensor[i] = - resource.l1Buf.template GetBufferByByte(l1BufAddrStart + L1A_SIZE * 2 * 2 + L1B_SIZE * i); - } - } - - /// Destructor - CATLASS_DEVICE - ~BlockMmad() {} - - CATLASS_DEVICE - void getBlockShape(GemmCoord &actualShape, uint32_t &nowNIdx, uint32_t &kIdx, - uint32_t &nLoop, uint32_t &kLoop, uint32_t &kvSeqlen, uint32_t &embed, bool firstBlock, uint32_t maskTailS = 0) - { - uint32_t nSplitSize = KV_SPLIT_SIZE * LOAB_BLOCK; - uint32_t embedSplitSize = EMBED_SPLIT_SIZE; - if (nowNIdx + LOAB_BLOCK > nLoop - 1) { - nSplitSize = kvSeqlen - nowNIdx * KV_SPLIT_SIZE; - } - if (firstBlock && maskTailS != 0) { - nSplitSize = nSplitSize - maskTailS; - } - if (kIdx == kLoop - 1) { - embedSplitSize = embed - kIdx * EMBED_SPLIT_SIZE; - } - actualShape[1] = embedSplitSize; - actualShape[2] = nSplitSize; - } - - CATLASS_DEVICE - void getKVOffset(AscendC::GlobalTensor &gBlockTable, uint32_t &kOffset, uint32_t &nowNIdx, uint32_t &kIdx, - uint32_t &nLoop, uint32_t &kLoop, uint32_t &strideKV, uint32_t &blockSize, uint32_t maskTailS = 0) - { - if (nowNIdx >= nLoop || kIdx >= kLoop) { - kOffset = 0; - } - if constexpr (PAGED_CACHE_FLAG_) { - uint32_t blockTableId = gBlockTable.GetValue(nowNIdx); - kOffset = blockTableId * blockSize * strideKV + maskTailS * strideKV + kIdx * EMBED_SPLIT_SIZE; - } else { - kOffset = nowNIdx * KV_SPLIT_SIZE * strideKV + kIdx * EMBED_SPLIT_SIZE; - } - } - - CATLASS_DEVICE - void operator()(AscendC::GlobalTensor gA, - AscendC::GlobalTensor gB, - AscendC::GlobalTensor gC, - AscendC::GlobalTensor gBlockTable, - LayoutA layoutA, LayoutB layoutB, GemmCoord actualOriShape, - uint32_t &nIdx, uint32_t &nLoop, uint32_t &blockSize, uint32_t kvSeqlen, uint32_t strideKV, Arch::CrossCoreFlag softmaxFlag, uint32_t maskTailS, bool preloadFlag) - { - uint32_t embed = actualOriShape[1]; - uint32_t kLoop = CeilDiv(embed); - uint32_t rowNum = layoutA.shape(0); - uint32_t blockN = layoutA.shape(1); - GemmCoord actualShape{rowNum, 0, 0}; - GemmCoord actualNextShape{rowNum, 0, 0}; - uint32_t nkBlockLoop = (nLoop + LOAB_BLOCK - 1) / LOAB_BLOCK * kLoop; // gap - uint32_t nkBlockNextIdx = (nIdx + LOAB_BLOCK - 1) / LOAB_BLOCK * kLoop + 1; // gap - uint32_t gBOffset = 0; - uint32_t gBNextOffset = 0; - uint32_t nowMaskTailS = 0; - uint32_t gPOffset = 0; - for (uint32_t kIdx = 0; kIdx < kLoop; kIdx++) { - nowMaskTailS = maskTailS; - gPOffset = 0; - for (uint32_t blockStackIdx = 0; (blockStackIdx < UNIT_BLOCK_STACK_NUM) && ((nIdx + blockStackIdx) < nLoop); - blockStackIdx += LOAB_BLOCK) { - uint32_t nowNIdx = nIdx + blockStackIdx; - uint32_t kLoopNextIdx = (nkBlockNextIdx % (kLoop * UNIT_BLOCK_STACK_NUM)) / (UNIT_BLOCK_STACK_NUM / LOAB_BLOCK); - uint32_t nLoopNextIdx = (nkBlockNextIdx % (kLoop * UNIT_BLOCK_STACK_NUM)) % (UNIT_BLOCK_STACK_NUM / LOAB_BLOCK) + nkBlockNextIdx / (kLoop * UNIT_BLOCK_STACK_NUM) * UNIT_BLOCK_STACK_NUM; - uint32_t startSeqOffset = nowNIdx == nIdx ? maskTailS : 0; - uint32_t startSeqNxtOffset = nLoopNextIdx == nIdx ? maskTailS : 0; - getBlockShape(actualShape, nowNIdx, kIdx, nLoop, kLoop, kvSeqlen, embed, nowNIdx == nIdx, nowMaskTailS); - getBlockShape(actualNextShape, nLoopNextIdx, kLoopNextIdx, nLoop, kLoop, kvSeqlen, embed, nLoopNextIdx == nIdx, nowMaskTailS); - getKVOffset(gBlockTable, gBOffset, nowNIdx, kIdx, nLoop, kLoop, strideKV, blockSize, startSeqOffset); - getKVOffset(gBlockTable, gBNextOffset, nLoopNextIdx, kLoopNextIdx, nLoop, kLoop, strideKV, blockSize, startSeqNxtOffset); - bool firstItr = blockStackIdx == 0; - bool endItr = (blockStackIdx + LOAB_BLOCK > UNIT_BLOCK_STACK_NUM - 1) || (nowNIdx + LOAB_BLOCK > nLoop - 1); - bool initMmad = blockStackIdx == 0; - bool pvCVItr = firstItr && kIdx == 0; - LayoutC layoutOTmpTemp(rowNum, embed, embed); - computePV(gA[gPOffset], gB[gBOffset], gC, gB[gBNextOffset], layoutA, layoutB, layoutOTmpTemp, - actualShape, actualNextShape, nowNIdx, nkBlockNextIdx, nkBlockLoop, firstItr, endItr, initMmad, pvCVItr, softmaxFlag, preloadFlag); - gPOffset += actualShape.k(); - ++nkBlockNextIdx; - nowMaskTailS = 0; - preloadFlag = false; - } - } - } - - CATLASS_DEVICE - void computePV( - AscendC::GlobalTensor const &gA, - AscendC::GlobalTensor const &gB, - AscendC::GlobalTensor const &gC, - AscendC::GlobalTensor const &gmNextBlockB, - LayoutA layoutA, LayoutB layoutB, LayoutC layoutC, - GemmCoord actualShape, GemmCoord actualNextShape, uint32_t nowIdx, uint32_t &nkblockIdx, - uint32_t &nkblockLoop, bool firstItr, bool endItr, bool initMmad, bool pvCVItr, Arch::CrossCoreFlag softmaxFlag, bool preloadFlag = false) - { - uint32_t MActual = actualShape.m(); - uint32_t kActual = actualShape.k(); - uint32_t nActual = actualShape.n(); - uint32_t mRound = RoundUp(MActual); - uint32_t kRound = RoundUp(kActual); - uint32_t nRound = RoundUp(nActual); - LayoutAInL1 layoutAInL1 = LayoutAInL1::template MakeLayout(mRound, kActual); - LayoutAInL0 layoutAInL0 = LayoutAInL0::template MakeLayout(mRound, kActual); - LayoutBInL1 layoutBInL1 = LayoutBInL1::template MakeLayout(kActual, nActual); - LayoutBInL0 layoutBInL0 = LayoutBInL0::template MakeLayout(kActual, nActual); - uint32_t l1KvPingPongFlag = nkblockIdx % 2; - uint32_t l0ABPingPongFlag = nkblockIdx % 2; - if (nkblockIdx == 1 || preloadFlag) { - auto layoutBTile = layoutB.GetTileLayout(MakeCoord(kActual, nActual)); - AscendC::WaitFlag(l1KvPingPongFlag + 2); - copyGmToL1B(l1BTensor[l1KvPingPongFlag], gB, layoutBInL1, layoutBTile); - AscendC::SetFlag(l1KvPingPongFlag + 2); - } - - AscendC::WaitFlag(l1KvPingPongFlag + 2); - AscendC::WaitFlag(2); - AscendC::WaitFlag(3); - copyL1ToL0B(l0BTensor, l1BTensor[l1KvPingPongFlag], layoutBInL0, layoutBInL1); - AscendC::SetFlag(l1KvPingPongFlag + 2); - - if (pvCVItr) { - Arch::CrossCoreWaitFlag(softmaxFlag); - } - AscendC::WaitFlag(l1KvPingPongFlag + 4); - auto layoutATile = layoutA.GetTileLayout(MakeCoord(MActual, kActual)); - copyGmToL1A(l1ATensor[l1KvPingPongFlag], gA, layoutAInL1, layoutATile); - AscendC::SetFlag(EVENT_ID5); - - if (nkblockIdx != nkblockLoop) { - uint32_t nNextActual = actualNextShape.n(); - uint32_t kNextActual = actualNextShape.k(); - LayoutBInL1 layoutBNextInL1 = LayoutBInL1::template MakeLayout(kNextActual, nNextActual); - auto layoutNextBTile = layoutB.GetTileLayout(MakeCoord(kNextActual, nNextActual)); - AscendC::WaitFlag(1 - l1KvPingPongFlag + 2); - copyGmToL1B(l1BTensor[1 - l1KvPingPongFlag], gmNextBlockB, layoutBNextInL1, layoutNextBTile); - AscendC::SetFlag(1 - l1KvPingPongFlag + 2); - } - - AscendC::WaitFlag(EVENT_ID5); - AscendC::WaitFlag(0); - AscendC::WaitFlag(1); - copyL1ToL0A(l0ATensor, l1ATensor[l1KvPingPongFlag], layoutAInL0, layoutAInL1); - AscendC::SetFlag(l1KvPingPongFlag + 4); - - AscendC::SetFlag(l0ABPingPongFlag); - AscendC::WaitFlag(l0ABPingPongFlag); - uint8_t unitFlag = 0b00; - if constexpr (!ENABLE_UNIT_FLAG_) { - if (firstItr) { - AscendC::WaitFlag(0); - AscendC::WaitFlag(1); - } - } else { - if (endItr) { - unitFlag = 0b11; - } else { - unitFlag = 0b10; - } - } - tileMmad( - l0CTensor[0], l0ATensor, l0BTensor, - mRound, nActual, kActual, initMmad, unitFlag); - // AscendC::PipeBarrier(); - AscendC::SetFlag(0); - AscendC::SetFlag(1); - AscendC::SetFlag(2); - AscendC::SetFlag(3); - auto layoutInL0C = LayoutCInL0::MakeLayoutInL0C(MakeCoord(MActual, nActual)); - if (endItr) { - if constexpr (!ENABLE_UNIT_FLAG_) { - AscendC::SetFlag(0); - AscendC::WaitFlag(0); - copyL0CToGm(gC, l0CTensor[0], layoutC, layoutInL0C); - AscendC::SetFlag(0); - AscendC::SetFlag(1); - } else { - copyL0CToGm(gC, l0CTensor[0], layoutC, layoutInL0C, 0b11); - } - } - } - -protected: - /// Data members - AscendC::LocalTensor l1ATensor[STAGES]; - AscendC::LocalTensor l1BTensor[STAGES]; - AscendC::LocalTensor l0ATensor; - AscendC::LocalTensor l0BTensor; - AscendC::LocalTensor l0CTensor[STAGES]; - - TileMmad tileMmad; - CopyGmToL1A copyGmToL1A; - CopyGmToL1B copyGmToL1B; - CopyL1ToL0A copyL1ToL0A; - CopyL1ToL0B copyL1ToL0B; - CopyL0CToGm copyL0CToGm; -}; - -//////////////////////////////////////////////////////////////////// - -} // namespace Catlass::Gemm::Block - -#endif // CATLASS_GEMM_BLOCK_MMAD_PV_TAIL_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_pv_head_tail_tla.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_pv_head_tail_tla.hpp deleted file mode 100644 index fc17f5fda4..0000000000 --- a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_pv_head_tail_tla.hpp +++ /dev/null @@ -1,350 +0,0 @@ -/** - * Copyright (c) 2025 Huawei Technologies Co., Ltd. - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of - * CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ - -#ifndef CATLASS_GEMM_BLOCK_MMAD_PV_TAIL_TLA_HPP -#define CATLASS_GEMM_BLOCK_MMAD_PV_TAIL_TLA_HPP - -#include "catlass/kda_catlass.hpp" -#include "catlass/arch/kda_resource.hpp" -#include "catlass/kda_coord.hpp" -#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" -#include "catlass/gemm/kda_gemm_helper.hpp" -#include "catlass/kda_gemm_coord.hpp" -#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" -#include "catlass/gemm/tile/kda_tile_mmad.hpp" -#include "tla/kda_tla_layout.hpp" -#include "tla/kda_tensor.hpp" - -//////////////////////////////////////////////////////////////////// - -namespace Catlass::Gemm::Block { -//////////////////////////////////////////////////////////////////// - -template < - class ArchTag_, - bool PAGED_CACHE_FLAG_, - bool ENABLE_UNIT_FLAG_, - class L1TileShape_, - class L0TileShape_, - class ElementA_, - class ElementB_, - class ElementC_, - class ElementBias_, - class TileCopy_, - class TileMmad_> -struct BlockMmadTla< - MmadFAITailPV, - L1TileShape_, - L0TileShape_, - ElementA_, - ElementB_, - ElementC_, - ElementBias_, - TileCopy_, - TileMmad_> { -public: - // Type Aliases - using DispatchPolicy = MmadFAITailPV; - using ArchTag = typename DispatchPolicy::ArchTag; - using L1TileShape = L1TileShape_; - using L0TileShape = L0TileShape_; - using ElementA = ElementA_; - using LayoutA = typename TileCopy_::LayoutA; - using ElementB = ElementB_; - using LayoutB = typename TileCopy_::LayoutB; - using ElementC = ElementC_; - using LayoutC = typename TileCopy_::LayoutC; - - using TileMmad = TileMmad_; - - using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; - using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; - - using ElementAccumulator = typename TileCopy_::ElementAccumulator; - - using LayoutTagL1A = typename TileCopy_::LayoutTagL1A; - using LayoutTagL1B = typename TileCopy_::LayoutTagL1B; - using LayoutTagL0A = typename TileCopy_::LayoutTagL0A; - using LayoutTagL0B = typename TileCopy_::LayoutTagL0B; - - using L1AAlignHelper = typename TileCopy_::L1AAlignHelper; - using L1BAlignHelper = typename TileCopy_::L1BAlignHelper; - - static_assert(tla::is_tuple::value && tla::is_static::value, - "L1TileShape must be tla::tuple and static!"); - static_assert(tla::is_tuple::value && tla::is_static::value, - "L0TileShape must be tla::tuple and static!"); - - static constexpr bool ENABLE_UNIT_FLAG = DispatchPolicy::ENABLE_UNIT_FLAG; - static constexpr uint32_t STAGES = DispatchPolicy::STAGES; - static constexpr uint32_t L1_TILE_M = tla::get<0>(L1TileShape{}); - static constexpr uint32_t L1_TILE_N = tla::get<1>(L1TileShape{}); - static constexpr uint32_t L1_TILE_K = tla::get<2>(L1TileShape{}); - static constexpr uint32_t L0_TILE_M = tla::get<0>(L0TileShape{}); - static constexpr uint32_t L0_TILE_N = tla::get<1>(L0TileShape{}); - static constexpr uint32_t L0_TILE_K = tla::get<2>(L0TileShape{}); - - // L1 tile size - static constexpr uint32_t L1A_TILE_SIZE = L1_TILE_M * L1_TILE_K * sizeof(ElementA); - static constexpr uint32_t L1B_TILE_SIZE = L1_TILE_N * L1_TILE_K * sizeof(ElementB); - // L0 tile size - static constexpr uint32_t L0A_TILE_SIZE = L0_TILE_M * L0_TILE_K * sizeof(ElementA); - static constexpr uint32_t L0B_TILE_SIZE = L0_TILE_K * L0_TILE_N * sizeof(ElementB); - static constexpr uint32_t L0C_TILE_SIZE = L0_TILE_M * L0_TILE_N * sizeof(ElementAccumulator); - static constexpr uint32_t L0A_PINGPONG_BUF_SIZE = ArchTag::L0A_SIZE / STAGES; - static constexpr uint32_t L0B_PINGPONG_BUF_SIZE = ArchTag::L0B_SIZE / STAGES; - static constexpr uint32_t L0C_PINGPONG_BUF_SIZE = ArchTag::L0C_SIZE / STAGES; - - // Check LayoutC - static_assert(tla::detail::isRowMajor::value || - ((std::is_same_v || std::is_same_v || - std::is_same_v) && - tla::detail::iszN::value), - "LayoutC only supports zN in half or bfloat16 or float, RowMajor in all dtype yet!"); - - // Check L1TileShape - static_assert((L1A_TILE_SIZE + L1B_TILE_SIZE) * STAGES <= ArchTag::L1_SIZE, - "L1TileShape exceeding the L1 space!"); - - // Check L0TileShape - static_assert(L0A_TILE_SIZE <= ArchTag::L0A_SIZE, "L0TileShape exceeding the L0A space!"); - static_assert(L0B_TILE_SIZE <= ArchTag::L0B_SIZE, "L0TileShape exceeding the L0B space!"); - static_assert(L0C_TILE_SIZE * STAGES <= ArchTag::L0C_SIZE, "L0TileShape exceeding the L0C space!"); - - static constexpr uint32_t BLOCK_SIZE = 16; - static constexpr uint32_t EMBED_SPLIT_SIZE = 128; - static constexpr uint32_t UNIT_BLOCK_STACK_NUM = 4; - static constexpr uint32_t KV_BASE_BLOCK = 512; - static constexpr uint32_t KV_SPLIT_SIZE = 128; - static constexpr uint32_t LOAB_BLOCK = 1; - - /// Construct - CATLASS_DEVICE - BlockMmadTla(Arch::Resource &resource, uint32_t l1BufAddrStart = 0) - { - // Allocate L1 memory space - l0ATensor = resource.l0ABuf.template GetBufferByByte(0); - l0BTensor = resource.l0BBuf.template GetBufferByByte(0); - for (uint32_t i = 0; i < STAGES; i++) { - l1ATensor[i] = resource.l1Buf.template GetBufferByByte(l1BufAddrStart + L1A_TILE_SIZE * LOAB_BLOCK * i); - l0CTensor[i] = resource.l0CBuf.template GetBufferByByte(L0C_PINGPONG_BUF_SIZE * i); - l1BTensor[i] = - resource.l1Buf.template GetBufferByByte(l1BufAddrStart + L1A_TILE_SIZE * 2 * 2 + L1B_TILE_SIZE * i); - } - } - - /// Destructor - CATLASS_DEVICE - ~BlockMmadTla() {} - - CATLASS_DEVICE - void getBlockShape(GemmCoord &actualShape, uint32_t &nowNIdx, uint32_t &kIdx, - uint32_t &nLoop, uint32_t &kLoop, uint32_t &kvSeqlen, uint32_t &embed, bool firstBlock, uint32_t maskTailS = 0) - { - uint32_t nSplitSize = KV_SPLIT_SIZE * LOAB_BLOCK; - uint32_t embedSplitSize = EMBED_SPLIT_SIZE; - if (nowNIdx + LOAB_BLOCK > nLoop - 1) { - nSplitSize = kvSeqlen - nowNIdx * KV_SPLIT_SIZE; - } - if (firstBlock && maskTailS != 0) { - nSplitSize = nSplitSize - maskTailS; - } - if (kIdx == kLoop - 1) { - embedSplitSize = embed - kIdx * EMBED_SPLIT_SIZE; - } - actualShape[1] = embedSplitSize; - actualShape[2] = nSplitSize; - } - - CATLASS_DEVICE - void getKVOffset(AscendC::GlobalTensor &gBlockTable, uint32_t &kOffset, uint32_t &nowNIdx, uint32_t &kIdx, - uint32_t &nLoop, uint32_t &kLoop, uint32_t &strideKV, uint32_t &blockSize, uint32_t maskTailS = 0) - { - if (nowNIdx >= nLoop || kIdx >= kLoop) { - kOffset = 0; - } - if constexpr (PAGED_CACHE_FLAG_) { - uint32_t blockTableId = gBlockTable.GetValue(nowNIdx); - kOffset = blockTableId * blockSize * strideKV + maskTailS * strideKV + kIdx * EMBED_SPLIT_SIZE; - } else { - kOffset = nowNIdx * KV_SPLIT_SIZE * strideKV + kIdx * EMBED_SPLIT_SIZE; - } - } - - template - CATLASS_DEVICE void operator()(TensorA &tensorA, TensorB &tensorB, TensorC &tensorC, - AscendC::GlobalTensor gBlockTable, GemmCoord actualOriShape, - uint32_t &nIdx, uint32_t &nLoop, uint32_t &blockSize, uint32_t kvSeqlen, uint32_t strideKV, Arch::CrossCoreFlag softmaxFlag, uint32_t maskTailS, bool preloadFlag) - { - uint32_t embed = actualOriShape[1]; - uint32_t kLoop = CeilDiv(embed); - uint32_t rowNum = tla::get<0>(tensorA.layout().shape()); - uint32_t blockN = tla::get<1>(tensorA.layout().shape()); - GemmCoord actualShape{rowNum, 0, 0}; - GemmCoord actualNextShape{rowNum, 0, 0}; - uint32_t nkBlockLoop = (nLoop + LOAB_BLOCK - 1) / LOAB_BLOCK * kLoop; // gap - uint32_t nkBlockNextIdx = (nIdx + LOAB_BLOCK - 1) / LOAB_BLOCK * kLoop + 1; // gap - uint32_t gBOffset = 0; - uint32_t gBNextOffset = 0; - uint32_t nowMaskTailS = 0; - uint32_t gPOffset = 0; - for (uint32_t kIdx = 0; kIdx < kLoop; kIdx++) { - nowMaskTailS = maskTailS; - gPOffset = 0; - for (uint32_t blockStackIdx = 0; (blockStackIdx < UNIT_BLOCK_STACK_NUM) && ((nIdx + blockStackIdx) < nLoop); - blockStackIdx += LOAB_BLOCK) { - uint32_t nowNIdx = nIdx + blockStackIdx; - uint32_t kLoopNextIdx = (nkBlockNextIdx % (kLoop * UNIT_BLOCK_STACK_NUM)) / (UNIT_BLOCK_STACK_NUM / LOAB_BLOCK); - uint32_t nLoopNextIdx = (nkBlockNextIdx % (kLoop * UNIT_BLOCK_STACK_NUM)) % (UNIT_BLOCK_STACK_NUM / LOAB_BLOCK) + nkBlockNextIdx / (kLoop * UNIT_BLOCK_STACK_NUM) * UNIT_BLOCK_STACK_NUM; - uint32_t startSeqOffset = nowNIdx == nIdx ? maskTailS : 0; - uint32_t startSeqNxtOffset = nLoopNextIdx == nIdx ? maskTailS : 0; - getBlockShape(actualShape, nowNIdx, kIdx, nLoop, kLoop, kvSeqlen, embed, nowNIdx == nIdx, nowMaskTailS); - getBlockShape(actualNextShape, nLoopNextIdx, kLoopNextIdx, nLoop, kLoop, kvSeqlen, embed, nLoopNextIdx == nIdx, nowMaskTailS); - getKVOffset(gBlockTable, gBOffset, nowNIdx, kIdx, nLoop, kLoop, strideKV, blockSize, startSeqOffset); - getKVOffset(gBlockTable, gBNextOffset, nLoopNextIdx, kLoopNextIdx, nLoop, kLoop, strideKV, blockSize, startSeqNxtOffset); - bool firstItr = blockStackIdx == 0; - bool endItr = (blockStackIdx + LOAB_BLOCK > UNIT_BLOCK_STACK_NUM - 1) || (nowNIdx + LOAB_BLOCK > nLoop - 1); - bool initMmad = blockStackIdx == 0; - bool pvCVItr = firstItr && kIdx == 0; - auto tensorAOffset = GetTile(tensorA, tla::MakeCoord(gPOffset / rowNum, gPOffset % rowNum), tensorA.layout().shape()); - computePV(tensorAOffset, tensorB, tensorC, gBOffset, gBNextOffset, - actualShape, actualNextShape, nowNIdx, nkBlockNextIdx, nkBlockLoop, strideKV, firstItr, endItr, initMmad, pvCVItr, softmaxFlag, preloadFlag); - gPOffset += actualShape.k(); - ++nkBlockNextIdx; - nowMaskTailS = 0; - preloadFlag = false; - } - } - } - - template - CATLASS_DEVICE void computePV( - TensorA &tensorA, TensorB &tensorB, TensorC &tensorC, uint32_t gBOffset, uint32_t gBNextOffset, - GemmCoord actualShape, GemmCoord actualNextShape, uint32_t nowIdx, uint32_t &nkblockIdx, - uint32_t &nkblockLoop, uint32_t strideKV, bool firstItr, bool endItr, bool initMmad, bool pvCVItr, Arch::CrossCoreFlag softmaxFlag, bool preloadFlag = false) - { - using CopyGmToL1A = typename TileCopy_::template CopyGmToL1A; - using CopyGmToL1B = typename TileCopy_::template CopyGmToL1B; - using CopyL0CToGm = typename TileCopy_::template CopyL0CToGm; - CopyGmToL1A copyGmToL1A; - CopyGmToL1B copyGmToL1B; - CopyL0CToGm copyL0CToGm; - - uint32_t MActual = actualShape.m(); - uint32_t kActual = actualShape.k(); - uint32_t nActual = actualShape.n(); - uint32_t mRound = RoundUp(MActual); - uint32_t kRound = RoundUp(kActual); - uint32_t nRound = RoundUp(nActual); - auto layoutAInL1 = tla::MakeLayout(mRound, kActual); - auto layoutAInL0 = tla::MakeLayout(mRound, kActual); - auto layoutBInL1 = tla::MakeLayout(kActual, nActual); - auto layoutBInL0 = tla::MakeLayout(kActual, nActual); - uint32_t l1KvPingPongFlag = nkblockIdx % 2; - uint32_t l0ABPingPongFlag = nkblockIdx % 2; - auto tensorTileA = GetTile(tensorA, tla::MakeCoord(0, 0), tla::MakeShape(MActual, kActual)); - auto tensorL1A = tla::MakeTensor(l1ATensor[l1KvPingPongFlag], layoutAInL1, Arch::PositionL1{}); - auto tensorL0A = tla::MakeTensor(l0ATensor, layoutAInL0, Arch::PositionL0A{}); - auto tensorTileB = GetTile(tensorB, tla::MakeCoord(gBOffset / strideKV, gBOffset % strideKV), tla::MakeShape(kActual, nActual)); - auto tensorL1B = tla::MakeTensor(l1BTensor[l1KvPingPongFlag], layoutBInL1, Arch::PositionL1{}); - auto tensorL0B = tla::MakeTensor(l0BTensor, layoutBInL0, Arch::PositionL0B{}); - if (nkblockIdx == 1 || preloadFlag) { - AscendC::WaitFlag(l1KvPingPongFlag + 2); - copyGmToL1B(tensorL1B, tensorTileB); - AscendC::SetFlag(l1KvPingPongFlag + 2); - } - - AscendC::WaitFlag(l1KvPingPongFlag + 2); - AscendC::WaitFlag(2); - AscendC::WaitFlag(3); - copyL1ToL0B(tensorL0B, tensorL1B); - AscendC::SetFlag(l1KvPingPongFlag + 2); - - if (pvCVItr) { - Arch::CrossCoreWaitFlag(softmaxFlag); - } - AscendC::WaitFlag(l1KvPingPongFlag + 4); - copyGmToL1A(tensorL1A, tensorTileA); - AscendC::SetFlag(EVENT_ID5); - - if (nkblockIdx != nkblockLoop) { - uint32_t nNextActual = actualNextShape.n(); - uint32_t kNextActual = actualNextShape.k(); - auto tensorTileBNext = GetTile(tensorB, tla::MakeCoord(gBNextOffset / strideKV, gBNextOffset % strideKV), tla::MakeShape(kNextActual, nNextActual)); - auto layoutBNextInL1 = tla::MakeLayout(kNextActual, nNextActual); - auto tensorL1BNext = tla::MakeTensor(l1BTensor[1 - l1KvPingPongFlag], layoutBNextInL1, Arch::PositionL1{}); - AscendC::WaitFlag(1 - l1KvPingPongFlag + 2); - copyGmToL1B(tensorL1BNext, tensorTileBNext); - AscendC::SetFlag(1 - l1KvPingPongFlag + 2); - } - - AscendC::WaitFlag(EVENT_ID5); - AscendC::WaitFlag(0); - AscendC::WaitFlag(1); - copyL1ToL0A(tensorL0A, tensorL1A); - AscendC::SetFlag(l1KvPingPongFlag + 4); - - AscendC::SetFlag(l0ABPingPongFlag); - AscendC::WaitFlag(l0ABPingPongFlag); - uint8_t unitFlag = 0b00; - if constexpr (!ENABLE_UNIT_FLAG_) { - if (firstItr) { - AscendC::WaitFlag(0); - AscendC::WaitFlag(1); - } - } else { - if (endItr) { - unitFlag = 0b11; - } else { - unitFlag = 0b10; - } - } - auto layoutCTileInL0 = tla::MakeLayoutL0C(mRound, nActual); - auto tensorL0CTile = tla::MakeTensor(l0CTensor[0], layoutCTileInL0, Arch::PositionL0C{}); - tileMmad( - tensorL0CTile, tensorL0A, tensorL0B, - mRound, nActual, kActual, initMmad, unitFlag); - AscendC::SetFlag(0); - AscendC::SetFlag(1); - AscendC::SetFlag(2); - AscendC::SetFlag(3); - auto layoutCInL0 = tla::MakeLayoutL0C(MActual, nActual); - auto tensorL0C = tla::MakeTensor(l0CTensor[0], layoutCInL0, Arch::PositionL0C{}); - if (endItr) { - if constexpr (!ENABLE_UNIT_FLAG_) { - AscendC::SetFlag(0); - AscendC::WaitFlag(0); - copyL0CToGm(tensorC, tensorL0C); - AscendC::SetFlag(0); - AscendC::SetFlag(1); - } else { - copyL0CToGm(tensorC, tensorL0C, 0b11); - } - } - } - -protected: - /// Data members - AscendC::LocalTensor l1ATensor[STAGES]; - AscendC::LocalTensor l1BTensor[STAGES]; - AscendC::LocalTensor l0ATensor; - AscendC::LocalTensor l0BTensor; - AscendC::LocalTensor l0CTensor[STAGES]; - - TileMmad tileMmad; - CopyL1ToL0A copyL1ToL0A; - CopyL1ToL0B copyL1ToL0B; -}; - -//////////////////////////////////////////////////////////////////// - -} // namespace Catlass::Gemm::Block - -#endif // CATLASS_GEMM_BLOCK_MMAD_PV_TAIL_TLA_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_pv_normal.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_pv_normal.hpp deleted file mode 100644 index d8d42f9489..0000000000 --- a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_pv_normal.hpp +++ /dev/null @@ -1,305 +0,0 @@ -/** - * Copyright (c) 2025 Huawei Technologies Co., Ltd. - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of - * CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ - -#ifndef CATLASS_GEMM_BLOCK_MMAD_PV_HPP -#define CATLASS_GEMM_BLOCK_MMAD_PV_HPP - -#include "catlass/kda_catlass.hpp" -#include "catlass/arch/kda_resource.hpp" -#include "catlass/kda_coord.hpp" -#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" -#include "catlass/gemm/kda_gemm_helper.hpp" -#include "catlass/kda_gemm_coord.hpp" -#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" -#include "catlass/gemm/tile/kda_tile_mmad.hpp" - -//////////////////////////////////////////////////////////////////// - -namespace Catlass::Gemm::Block { -//////////////////////////////////////////////////////////////////// - -template < - bool PAGED_CACHE_FLAG_, - bool ENABLE_UNIT_FLAG_, - class L1TileShape_, - class L0TileShape_, - class AType_, - class BType_, - class CType_, - class BiasType_, - class TileCopy_, - class TileMmad_> -struct BlockMmad< - MmadAtlasA2FAIPV, - L1TileShape_, - L0TileShape_, - AType_, - BType_, - CType_, - BiasType_, - TileCopy_, - TileMmad_> { -public: - // Type Aliases - using DispatchPolicy = MmadAtlasA2FAIPV; - using ArchTag = typename DispatchPolicy::ArchTag; - using L1TileShape = L1TileShape_; - using L0TileShape = L0TileShape_; - using ElementA = typename AType_::Element; - using LayoutA = typename AType_::Layout; - using ElementB = typename BType_::Element; - using LayoutB = typename BType_::Layout; - using ElementC = typename CType_::Element; - using LayoutC = typename CType_::Layout; - using TileMmad = TileMmad_; - using CopyGmToL1A = typename TileCopy_::CopyGmToL1A; - using CopyGmToL1B = typename TileCopy_::CopyGmToL1B; - using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; - using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; - using CopyL0CToGm = typename TileCopy_::CopyL0CToGm; - using ElementAccumulator = - typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; - using LayoutAInL1 = typename CopyL1ToL0A::LayoutSrc; - using LayoutBInL1 = typename CopyL1ToL0B::LayoutSrc; - using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst; - using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst; - using LayoutCInL0 = layout::zN; - - using L1AAlignHelper = Gemm::helper::L1AlignHelper; - using L1BAlignHelper = Gemm::helper::L1AlignHelper; - - static constexpr uint32_t STAGES = DispatchPolicy::STAGES; - static constexpr uint32_t L1A_SIZE = L1TileShape::M * L1TileShape::K * sizeof(ElementA); - static constexpr uint32_t L1B_SIZE = L1TileShape::N * L1TileShape::K * sizeof(ElementB); - static constexpr uint32_t L0A_SIZE = ArchTag::L0A_SIZE; - static constexpr uint32_t L0B_SIZE = ArchTag::L0B_SIZE; - static constexpr uint32_t L0C_SIZE = ArchTag::L0C_SIZE; - static constexpr uint32_t L0A_PINGPONG_BUF_SIZE = L0A_SIZE / STAGES; - static constexpr uint32_t L0B_PINGPONG_BUF_SIZE = L0B_SIZE / STAGES; - static constexpr uint32_t L0C_PINGPONG_BUF_SIZE = L0C_SIZE / STAGES; - static constexpr uint32_t BLOCK_SIZE = 16; - static constexpr uint32_t EMBED_SPLIT_SIZE = 128; - static constexpr uint32_t UNIT_BLOCK_STACK_NUM = 4; - static constexpr uint32_t KV_BASE_BLOCK = 512; - static constexpr uint32_t KV_SPLIT_SIZE = 128; - static constexpr uint32_t LOAB_BLOCK = 1; - - // Check LayoutC - static_assert(std::is_same_v, "LayoutC only support RowMajor yet!"); - - /// Construct - CATLASS_DEVICE - BlockMmad(Arch::Resource &resource, uint32_t l1BufAddrStart = 0) - { - // Allocate L1 memory space - l1ATensor = resource.l1Buf.template GetBufferByByte(l1BufAddrStart); - for (uint32_t i = 0; i < STAGES; i++) { - l0ATensor[i] = resource.l0ABuf.template GetBufferByByte(L0A_PINGPONG_BUF_SIZE * i); - l0BTensor[i] = resource.l0BBuf.template GetBufferByByte(L0B_PINGPONG_BUF_SIZE * i); - l0CTensor[i] = resource.l0CBuf.template GetBufferByByte(L0C_PINGPONG_BUF_SIZE * i); - l1BTensor[i] = - resource.l1Buf.template GetBufferByByte(l1BufAddrStart + L1A_SIZE * 2 * 2 + L1B_SIZE * i); - } - } - - /// Destructor - CATLASS_DEVICE - ~BlockMmad() {} - - CATLASS_DEVICE - void getBlockShape(GemmCoord &actualShape, uint32_t &nowNIdx, uint32_t &kIdx, - uint32_t &nLoop, uint32_t &kLoop, uint32_t &kvSeqlen, uint32_t &embed, bool firstBlock, uint32_t maskTailS = 0) - { - uint32_t nSplitSize = KV_SPLIT_SIZE * LOAB_BLOCK; - uint32_t embedSplitSize = EMBED_SPLIT_SIZE; - if (nowNIdx + LOAB_BLOCK > nLoop - 1) { - nSplitSize = kvSeqlen - nowNIdx * KV_SPLIT_SIZE; - } - if (firstBlock && maskTailS != 0) { - nSplitSize = nSplitSize - maskTailS; - } - if (kIdx == kLoop - 1) { - embedSplitSize = embed - kIdx * EMBED_SPLIT_SIZE; - } - actualShape[1] = embedSplitSize; - actualShape[2] = nSplitSize; - } - - CATLASS_DEVICE - void getKVOffset(AscendC::GlobalTensor &gBlockTable, uint32_t &kOffset, uint32_t &nowNIdx, uint32_t &kIdx, - uint32_t &nLoop, uint32_t &kLoop, uint32_t &strideKV, uint32_t &blockSize, uint32_t maskTailS = 0) - { - if (nowNIdx >= nLoop || kIdx >= kLoop) { - kOffset = 0; - } - if constexpr (PAGED_CACHE_FLAG_) { - uint32_t blockTableId = gBlockTable.GetValue(nowNIdx); - kOffset = blockTableId * blockSize * strideKV + maskTailS * strideKV + kIdx * EMBED_SPLIT_SIZE; - } else { - kOffset = nowNIdx * KV_SPLIT_SIZE * strideKV + kIdx * EMBED_SPLIT_SIZE; - } - } - - CATLASS_DEVICE - void operator()(AscendC::GlobalTensor gA, - AscendC::GlobalTensor gB, - AscendC::GlobalTensor gC, - AscendC::GlobalTensor gBlockTable, - LayoutA layoutA, LayoutB layoutB, GemmCoord actualOriShape, - uint32_t &nIdx, uint32_t &nLoop, uint32_t &blockSize, uint32_t kvSeqlen, uint32_t strideKV, Arch::CrossCoreFlag softmaxFlag) - { - // Arch::CrossCoreWaitFlag(softmaxFlag); - uint32_t embed = actualOriShape[1]; - uint32_t kLoop = CeilDiv(embed); - uint32_t rowNum = layoutA.shape(0); - uint32_t blockN = layoutA.shape(1); - GemmCoord actualShape{rowNum, 0, 0}; - GemmCoord actualNextShape{rowNum, 0, 0}; - uint32_t nkBlockLoop = (nLoop + LOAB_BLOCK - 1) / LOAB_BLOCK * kLoop; // gap - uint32_t nkBlockNextIdx = (nIdx + LOAB_BLOCK - 1) / LOAB_BLOCK * kLoop + 1; // gap - uint32_t gBOffset = 0; - uint32_t gBNextOffset = 0; - for (uint32_t kIdx = 0; kIdx < kLoop; kIdx++) { - for (uint32_t blockStackIdx = 0; (blockStackIdx < UNIT_BLOCK_STACK_NUM) && ((nIdx + blockStackIdx) < nLoop); - blockStackIdx += LOAB_BLOCK) { - uint32_t nowNIdx = nIdx + blockStackIdx; - uint32_t kLoopNextIdx = (nkBlockNextIdx % (kLoop * UNIT_BLOCK_STACK_NUM)) / (UNIT_BLOCK_STACK_NUM / LOAB_BLOCK); - uint32_t nLoopNextIdx = (nkBlockNextIdx % (kLoop * UNIT_BLOCK_STACK_NUM)) % (UNIT_BLOCK_STACK_NUM / LOAB_BLOCK) + nkBlockNextIdx / (kLoop * UNIT_BLOCK_STACK_NUM) * UNIT_BLOCK_STACK_NUM; - getBlockShape(actualShape, nowNIdx, kIdx, nLoop, kLoop, kvSeqlen, embed, nowNIdx == nIdx); - getBlockShape(actualNextShape, nLoopNextIdx, kLoopNextIdx, nLoop, kLoop, kvSeqlen, embed, nowNIdx == nIdx); - getKVOffset(gBlockTable, gBOffset, nowNIdx, kIdx, nLoop, kLoop, strideKV, blockSize); - getKVOffset(gBlockTable, gBNextOffset, nLoopNextIdx, kLoopNextIdx, nLoop, kLoop, strideKV, blockSize); - bool firstItr = blockStackIdx == 0; - bool endItr = (blockStackIdx + LOAB_BLOCK > UNIT_BLOCK_STACK_NUM - 1) || (nowNIdx + LOAB_BLOCK > nLoop - 1); - bool initMmad = blockStackIdx == 0; - bool pvCVItr = firstItr && kIdx == 0; - LayoutC layoutOTmpTemp(rowNum, embed, embed); - computePV(gA, gB[gBOffset], gC, gB[gBNextOffset], layoutA, layoutB, layoutOTmpTemp, - actualShape, actualNextShape, blockStackIdx, nkBlockNextIdx, nkBlockLoop, firstItr, endItr, initMmad, pvCVItr, softmaxFlag); - ++nkBlockNextIdx; - } - } - } - - CATLASS_DEVICE - void computePV( - AscendC::GlobalTensor const &gA, - AscendC::GlobalTensor const &gB, - AscendC::GlobalTensor const &gC, - AscendC::GlobalTensor const &gmNextBlockB, - LayoutA layoutA, LayoutB layoutB, LayoutC layoutC, - GemmCoord actualShape, GemmCoord actualNextShape, uint32_t nowIdx, uint32_t &nkblockIdx, - uint32_t &nkblockLoop, bool firstItr, bool endItr, bool initMmad, bool pvCVItr, Arch::CrossCoreFlag softmaxFlag, bool preloadFlag = false) - { - uint32_t MActual = actualShape.m(); - uint32_t kActual = actualShape.k(); - uint32_t nActual = actualShape.n(); - uint32_t mRound = RoundUp(MActual); - uint32_t kRound = RoundUp(kActual); - uint32_t nRound = RoundUp(nActual); - LayoutAInL1 layoutAInL1 = LayoutAInL1::template MakeLayout(mRound, (uint32_t)512); - LayoutAInL0 layoutAInL0 = LayoutAInL0::template MakeLayout(mRound, kActual); - LayoutBInL1 layoutBInL1 = LayoutBInL1::template MakeLayout(kActual, nActual); - LayoutBInL0 layoutBInL0 = LayoutBInL0::template MakeLayout(kActual, nActual); - uint32_t l1KvPingPongFlag = nkblockIdx % 2; - uint32_t l0ABPingPongFlag = nkblockIdx % 2; - if (nkblockIdx == 1 || preloadFlag) { - auto layoutBTile = layoutB.GetTileLayout(MakeCoord(kActual, nActual)); - AscendC::WaitFlag(l1KvPingPongFlag + 2); - copyGmToL1B(l1BTensor[l1KvPingPongFlag], gB, layoutBInL1, layoutBTile); - AscendC::SetFlag(l1KvPingPongFlag + 2); - } - - AscendC::WaitFlag(l1KvPingPongFlag + 2); - AscendC::WaitFlag(l0ABPingPongFlag + 2); - copyL1ToL0B(l0BTensor[l0ABPingPongFlag], l1BTensor[l1KvPingPongFlag], layoutBInL0, layoutBInL1); - AscendC::SetFlag(l1KvPingPongFlag + 2); - - if (pvCVItr) { - Arch::CrossCoreWaitFlag(softmaxFlag); - AscendC::WaitFlag(4); - AscendC::WaitFlag(5); - auto layoutATile = layoutA.GetTileLayout(MakeCoord(MActual, (uint32_t)512)); - copyGmToL1A(l1ATensor, gA, layoutAInL1, layoutATile); - AscendC::SetFlag(4); - AscendC::WaitFlag(4); - } - - if (nkblockIdx != nkblockLoop) { - uint32_t nNextActual = actualNextShape.n(); - uint32_t kNextActual = actualNextShape.k(); - LayoutBInL1 layoutBNextInL1 = LayoutBInL1::template MakeLayout(kNextActual, nNextActual); - auto layoutNextBTile = layoutB.GetTileLayout(MakeCoord(kNextActual, nNextActual)); - AscendC::WaitFlag(1 - l1KvPingPongFlag + 2); - copyGmToL1B(l1BTensor[1 - l1KvPingPongFlag], gmNextBlockB, layoutBNextInL1, layoutNextBTile); - AscendC::SetFlag(1 - l1KvPingPongFlag + 2); - } - - AscendC::WaitFlag(l0ABPingPongFlag); - copyL1ToL0A(l0ATensor[l0ABPingPongFlag], l1ATensor[mRound * 128 * nowIdx], layoutAInL0, layoutAInL1); - - AscendC::SetFlag(l0ABPingPongFlag); - AscendC::WaitFlag(l0ABPingPongFlag); - uint8_t unitFlag = 0b00; - if constexpr (!ENABLE_UNIT_FLAG_) { - if (firstItr) { - AscendC::WaitFlag(0); - AscendC::WaitFlag(1); - } - } else { - if (endItr) { - unitFlag = 0b11; - } else { - unitFlag = 0b10; - } - } - tileMmad( - l0CTensor[0], l0ATensor[l0ABPingPongFlag], l0BTensor[l0ABPingPongFlag], - mRound, nActual, kActual, initMmad, unitFlag); - // AscendC::PipeBarrier(); - AscendC::SetFlag(l0ABPingPongFlag); - AscendC::SetFlag(l0ABPingPongFlag + 2); - auto layoutInL0C = LayoutCInL0::MakeLayoutInL0C(MakeCoord(MActual, nActual)); - if (endItr) { - AscendC::SetFlag(4); - AscendC::SetFlag(5); - if constexpr (!ENABLE_UNIT_FLAG_) { - AscendC::SetFlag(0); - AscendC::WaitFlag(0); - copyL0CToGm(gC, l0CTensor[0], layoutC, layoutInL0C); - AscendC::SetFlag(0); - AscendC::SetFlag(1); - } else { - copyL0CToGm(gC, l0CTensor[0], layoutC, layoutInL0C, 0b11); - } - } - } - -protected: - /// Data members - AscendC::LocalTensor l1ATensor; - AscendC::LocalTensor l1BTensor[STAGES]; - AscendC::LocalTensor l0ATensor[STAGES]; - AscendC::LocalTensor l0BTensor[STAGES]; - AscendC::LocalTensor l0CTensor[STAGES]; - - TileMmad tileMmad; - CopyGmToL1A copyGmToL1A; - CopyGmToL1B copyGmToL1B; - CopyL1ToL0A copyL1ToL0A; - CopyL1ToL0B copyL1ToL0B; - CopyL0CToGm copyL0CToGm; -}; - -//////////////////////////////////////////////////////////////////// - -} // namespace Catlass::Gemm::Block - -#endif // CATLASS_GEMM_BLOCK_MMAD_PV_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_pv_normal_tla.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_pv_normal_tla.hpp deleted file mode 100644 index 7f98bea467..0000000000 --- a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_pv_normal_tla.hpp +++ /dev/null @@ -1,340 +0,0 @@ -/** - * Copyright (c) 2025 Huawei Technologies Co., Ltd. - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of - * CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ - -#ifndef CATLASS_GEMM_BLOCK_MMAD_PV_TLA_HPP -#define CATLASS_GEMM_BLOCK_MMAD_PV_TLA_HPP - -#include "catlass/kda_catlass.hpp" -#include "catlass/arch/kda_resource.hpp" -#include "catlass/kda_coord.hpp" -#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" -#include "catlass/gemm/kda_gemm_helper.hpp" -#include "catlass/kda_gemm_coord.hpp" -#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" -#include "catlass/gemm/tile/kda_tile_mmad.hpp" -#include "tla/kda_tla_layout.hpp" -#include "tla/kda_tensor.hpp" - -//////////////////////////////////////////////////////////////////// - -namespace Catlass::Gemm::Block { -//////////////////////////////////////////////////////////////////// - -template < - class ArchTag_, - bool PAGED_CACHE_FLAG_, - bool ENABLE_UNIT_FLAG_, - class L1TileShape_, - class L0TileShape_, - class ElementA_, - class ElementB_, - class ElementC_, - class ElementBias_, - class TileCopy_, - class TileMmad_> -struct BlockMmadTla< - MmadFAIPV, - L1TileShape_, - L0TileShape_, - ElementA_, - ElementB_, - ElementC_, - ElementBias_, - TileCopy_, - TileMmad_> { -public: - // Type Aliases - using DispatchPolicy = MmadFAIPV; - using ArchTag = typename DispatchPolicy::ArchTag; - using L1TileShape = L1TileShape_; - using L0TileShape = L0TileShape_; - using ElementA = ElementA_; - using LayoutA = typename TileCopy_::LayoutA; - using ElementB = ElementB_; - using LayoutB = typename TileCopy_::LayoutB; - using ElementC = ElementC_; - using LayoutC = typename TileCopy_::LayoutC; - - using TileMmad = TileMmad_; - - using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; - using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; - - using ElementAccumulator = typename TileCopy_::ElementAccumulator; - - using LayoutTagL1A = typename TileCopy_::LayoutTagL1A; - using LayoutTagL1B = typename TileCopy_::LayoutTagL1B; - using LayoutTagL0A = typename TileCopy_::LayoutTagL0A; - using LayoutTagL0B = typename TileCopy_::LayoutTagL0B; - - using L1AAlignHelper = typename TileCopy_::L1AAlignHelper; - using L1BAlignHelper = typename TileCopy_::L1BAlignHelper; - - static_assert(tla::is_tuple::value && tla::is_static::value, - "L1TileShape must be tla::tuple and static!"); - static_assert(tla::is_tuple::value && tla::is_static::value, - "L0TileShape must be tla::tuple and static!"); - - static constexpr bool ENABLE_UNIT_FLAG = DispatchPolicy::ENABLE_UNIT_FLAG; - static constexpr uint32_t STAGES = DispatchPolicy::STAGES; - static constexpr uint32_t L1_TILE_M = tla::get<0>(L1TileShape{}); - static constexpr uint32_t L1_TILE_N = tla::get<1>(L1TileShape{}); - static constexpr uint32_t L1_TILE_K = tla::get<2>(L1TileShape{}); - static constexpr uint32_t L0_TILE_M = tla::get<0>(L0TileShape{}); - static constexpr uint32_t L0_TILE_N = tla::get<1>(L0TileShape{}); - static constexpr uint32_t L0_TILE_K = tla::get<2>(L0TileShape{}); - - // L1 tile size - static constexpr uint32_t L1A_TILE_SIZE = L1_TILE_M * L1_TILE_K * sizeof(ElementA); - static constexpr uint32_t L1B_TILE_SIZE = L1_TILE_N * L1_TILE_K * sizeof(ElementB); - // L0 tile size - static constexpr uint32_t L0A_TILE_SIZE = L0_TILE_M * L0_TILE_K * sizeof(ElementA); - static constexpr uint32_t L0B_TILE_SIZE = L0_TILE_K * L0_TILE_N * sizeof(ElementB); - static constexpr uint32_t L0C_TILE_SIZE = L0_TILE_M * L0_TILE_N * sizeof(ElementAccumulator); - static constexpr uint32_t L0A_PINGPONG_BUF_SIZE = ArchTag::L0A_SIZE / STAGES; - static constexpr uint32_t L0B_PINGPONG_BUF_SIZE = ArchTag::L0B_SIZE / STAGES; - static constexpr uint32_t L0C_PINGPONG_BUF_SIZE = ArchTag::L0C_SIZE / STAGES; - - // Check LayoutC - static_assert(tla::detail::isRowMajor::value || - ((std::is_same_v || std::is_same_v || - std::is_same_v) && - tla::detail::iszN::value), - "LayoutC only supports zN in half or bfloat16 or float, RowMajor in all dtype yet!"); - - // Check L1TileShape - static_assert(L1A_TILE_SIZE + L1B_TILE_SIZE * STAGES <= ArchTag::L1_SIZE, - "L1TileShape exceeding the L1 space!"); - - // Check L0TileShape - static_assert(L0A_TILE_SIZE * STAGES <= ArchTag::L0A_SIZE, "L0TileShape exceeding the L0A space!"); - static_assert(L0B_TILE_SIZE * STAGES <= ArchTag::L0B_SIZE, "L0TileShape exceeding the L0B space!"); - static_assert(L0C_TILE_SIZE * STAGES <= ArchTag::L0C_SIZE, "L0TileShape exceeding the L0C space!"); - - static constexpr uint32_t BLOCK_SIZE = 16; - static constexpr uint32_t EMBED_SPLIT_SIZE = 128; - static constexpr uint32_t UNIT_BLOCK_STACK_NUM = 4; - static constexpr uint32_t KV_BASE_BLOCK = 512; - static constexpr uint32_t KV_SPLIT_SIZE = 128; - static constexpr uint32_t LOAB_BLOCK = 1; - - /// Construct - CATLASS_DEVICE - BlockMmadTla(Arch::Resource &resource, uint32_t l1BufAddrStart = 0) - { - // Allocate L1 memory space - l1ATensor = resource.l1Buf.template GetBufferByByte(l1BufAddrStart); - for (uint32_t i = 0; i < STAGES; i++) { - l0ATensor[i] = resource.l0ABuf.template GetBufferByByte(L0A_PINGPONG_BUF_SIZE * i); - l0BTensor[i] = resource.l0BBuf.template GetBufferByByte(L0B_PINGPONG_BUF_SIZE * i); - l0CTensor[i] = resource.l0CBuf.template GetBufferByByte(L0C_PINGPONG_BUF_SIZE * i); - l1BTensor[i] = - resource.l1Buf.template GetBufferByByte(l1BufAddrStart + L1A_TILE_SIZE * 2 * 2 + L1B_TILE_SIZE * i); - } - } - - /// Destructor - CATLASS_DEVICE - ~BlockMmadTla() {} - - CATLASS_DEVICE - void getBlockShape(GemmCoord &actualShape, uint32_t &nowNIdx, uint32_t &kIdx, - uint32_t &nLoop, uint32_t &kLoop, uint32_t &kvSeqlen, uint32_t &embed, bool firstBlock, uint32_t maskTailS = 0) - { - uint32_t nSplitSize = KV_SPLIT_SIZE * LOAB_BLOCK; - uint32_t embedSplitSize = EMBED_SPLIT_SIZE; - if (nowNIdx + LOAB_BLOCK > nLoop - 1) { - nSplitSize = kvSeqlen - nowNIdx * KV_SPLIT_SIZE; - } - if (firstBlock && maskTailS != 0) { - nSplitSize = nSplitSize - maskTailS; - } - if (kIdx == kLoop - 1) { - embedSplitSize = embed - kIdx * EMBED_SPLIT_SIZE; - } - actualShape[1] = embedSplitSize; - actualShape[2] = nSplitSize; - } - - CATLASS_DEVICE - void getKVOffset(AscendC::GlobalTensor &gBlockTable, uint32_t &kOffset, uint32_t &nowNIdx, uint32_t &kIdx, - uint32_t &nLoop, uint32_t &kLoop, uint32_t &strideKV, uint32_t &blockSize, uint32_t maskTailS = 0) - { - if (nowNIdx >= nLoop || kIdx >= kLoop) { - kOffset = 0; - } - if constexpr (PAGED_CACHE_FLAG_) { - uint32_t blockTableId = gBlockTable.GetValue(nowNIdx); - kOffset = blockTableId * blockSize * strideKV + maskTailS * strideKV + kIdx * EMBED_SPLIT_SIZE; - } else { - kOffset = nowNIdx * KV_SPLIT_SIZE * strideKV + kIdx * EMBED_SPLIT_SIZE; - } - } - - template - CATLASS_DEVICE void operator()(TensorA &tensorA, TensorB &tensorB, TensorC &tensorC, - AscendC::GlobalTensor gBlockTable, GemmCoord actualOriShape, - uint32_t &nIdx, uint32_t &nLoop, uint32_t &blockSize, uint32_t kvSeqlen, uint32_t strideKV, Arch::CrossCoreFlag softmaxFlag) - { - // Arch::CrossCoreWaitFlag(softmaxFlag); - uint32_t embed = actualOriShape[1]; - uint32_t kLoop = CeilDiv(embed); - uint32_t rowNum = tla::get<0>(tensorA.layout().shape()); - uint32_t blockN = tla::get<1>(tensorA.layout().shape()); - GemmCoord actualShape{rowNum, 0, 0}; - GemmCoord actualNextShape{rowNum, 0, 0}; - uint32_t nkBlockLoop = (nLoop + LOAB_BLOCK - 1) / LOAB_BLOCK * kLoop; // gap - uint32_t nkBlockNextIdx = (nIdx + LOAB_BLOCK - 1) / LOAB_BLOCK * kLoop + 1; // gap - uint32_t gBOffset = 0; - uint32_t gBNextOffset = 0; - for (uint32_t kIdx = 0; kIdx < kLoop; kIdx++) { - for (uint32_t blockStackIdx = 0; (blockStackIdx < UNIT_BLOCK_STACK_NUM) && ((nIdx + blockStackIdx) < nLoop); - blockStackIdx += LOAB_BLOCK) { - uint32_t nowNIdx = nIdx + blockStackIdx; - uint32_t kLoopNextIdx = (nkBlockNextIdx % (kLoop * UNIT_BLOCK_STACK_NUM)) / (UNIT_BLOCK_STACK_NUM / LOAB_BLOCK); - uint32_t nLoopNextIdx = (nkBlockNextIdx % (kLoop * UNIT_BLOCK_STACK_NUM)) % (UNIT_BLOCK_STACK_NUM / LOAB_BLOCK) + nkBlockNextIdx / (kLoop * UNIT_BLOCK_STACK_NUM) * UNIT_BLOCK_STACK_NUM; - getBlockShape(actualShape, nowNIdx, kIdx, nLoop, kLoop, kvSeqlen, embed, nowNIdx == nIdx); - getBlockShape(actualNextShape, nLoopNextIdx, kLoopNextIdx, nLoop, kLoop, kvSeqlen, embed, nowNIdx == nIdx); - getKVOffset(gBlockTable, gBOffset, nowNIdx, kIdx, nLoop, kLoop, strideKV, blockSize); - getKVOffset(gBlockTable, gBNextOffset, nLoopNextIdx, kLoopNextIdx, nLoop, kLoop, strideKV, blockSize); - bool firstItr = blockStackIdx == 0; - bool endItr = (blockStackIdx + LOAB_BLOCK > UNIT_BLOCK_STACK_NUM - 1) || (nowNIdx + LOAB_BLOCK > nLoop - 1); - bool initMmad = blockStackIdx == 0; - bool pvCVItr = firstItr && kIdx == 0; - computePV(tensorA, tensorB, tensorC, gBOffset, gBNextOffset, - actualShape, actualNextShape, blockStackIdx, nkBlockNextIdx, nkBlockLoop, strideKV, firstItr, endItr, initMmad, pvCVItr, softmaxFlag); - ++nkBlockNextIdx; - } - } - } - - template - CATLASS_DEVICE void computePV( - TensorA &tensorA, TensorB &tensorB, TensorC &tensorC, uint32_t gBOffset, uint32_t gBNextOffset, - GemmCoord actualShape, GemmCoord actualNextShape, uint32_t nowIdx, uint32_t &nkblockIdx, - uint32_t &nkblockLoop, uint32_t strideKV, bool firstItr, bool endItr, bool initMmad, bool pvCVItr, Arch::CrossCoreFlag softmaxFlag, bool preloadFlag = false) - { - using CopyGmToL1A = typename TileCopy_::template CopyGmToL1A; - using CopyGmToL1B = typename TileCopy_::template CopyGmToL1B; - using CopyL0CToGm = typename TileCopy_::template CopyL0CToGm; - CopyGmToL1A copyGmToL1A; - CopyGmToL1B copyGmToL1B; - CopyL0CToGm copyL0CToGm; - - uint32_t MActual = actualShape.m(); - uint32_t kActual = actualShape.k(); - uint32_t nActual = actualShape.n(); - uint32_t mRound = RoundUp(MActual); - uint32_t kRound = RoundUp(kActual); - uint32_t nRound = RoundUp(nActual); - auto layoutAInL1 = tla::MakeLayout(mRound, (uint32_t)512); - auto layoutAInL0 = tla::MakeLayout(mRound, kActual); - auto layoutBInL1 = tla::MakeLayout(kActual, nActual); - auto layoutBInL0 = tla::MakeLayout(kActual, nActual); - uint32_t l1KvPingPongFlag = nkblockIdx % 2; - uint32_t l0ABPingPongFlag = nkblockIdx % 2; - auto tensorTileA = GetTile(tensorA, tla::MakeCoord(0, 0), tla::MakeShape(MActual, (uint32_t)512)); - auto tensorL1A = tla::MakeTensor(l1ATensor, layoutAInL1, Arch::PositionL1{}); - auto tensorL0A = tla::MakeTensor(l0ATensor[l0ABPingPongFlag], layoutAInL0, Arch::PositionL0A{}); - auto tensorTileB = GetTile(tensorB, tla::MakeCoord(gBOffset / strideKV, gBOffset % strideKV), tla::MakeShape(kActual, nActual)); - auto tensorL1B = tla::MakeTensor(l1BTensor[l1KvPingPongFlag], layoutBInL1, Arch::PositionL1{}); - auto tensorL0B = tla::MakeTensor(l0BTensor[l0ABPingPongFlag], layoutBInL0, Arch::PositionL0B{}); - if (nkblockIdx == 1 || preloadFlag) { - AscendC::WaitFlag(l1KvPingPongFlag + 2); - copyGmToL1B(tensorL1B, tensorTileB); - AscendC::SetFlag(l1KvPingPongFlag + 2); - } - - AscendC::WaitFlag(l1KvPingPongFlag + 2); - AscendC::WaitFlag(l0ABPingPongFlag + 2); - copyL1ToL0B(tensorL0B, tensorL1B); - AscendC::SetFlag(l1KvPingPongFlag + 2); - - if (pvCVItr) { - Arch::CrossCoreWaitFlag(softmaxFlag); - AscendC::WaitFlag(4); - AscendC::WaitFlag(5); - copyGmToL1A(tensorL1A, tensorTileA); - AscendC::SetFlag(4); - AscendC::WaitFlag(4); - } - - if (nkblockIdx != nkblockLoop) { - uint32_t nNextActual = actualNextShape.n(); - uint32_t kNextActual = actualNextShape.k(); - auto tensorTileBNext = GetTile(tensorB, tla::MakeCoord(gBNextOffset / strideKV, gBNextOffset % strideKV), tla::MakeShape(kNextActual, nNextActual)); - auto layoutBNextInL1 = tla::MakeLayout(kNextActual, nNextActual); - auto tensorL1BNext = tla::MakeTensor(l1BTensor[1 - l1KvPingPongFlag], layoutBNextInL1, Arch::PositionL1{}); - AscendC::WaitFlag(1 - l1KvPingPongFlag + 2); - copyGmToL1B(tensorL1BNext, tensorTileBNext); - AscendC::SetFlag(1 - l1KvPingPongFlag + 2); - } - - AscendC::WaitFlag(l0ABPingPongFlag); - auto tensorL1AOffset = tla::MakeTensor(l1ATensor[mRound * 128 * nowIdx], layoutAInL1, Arch::PositionL1{}); - copyL1ToL0A(tensorL0A, tensorL1AOffset); - - AscendC::SetFlag(l0ABPingPongFlag); - AscendC::WaitFlag(l0ABPingPongFlag); - uint8_t unitFlag = 0b00; - if constexpr (!ENABLE_UNIT_FLAG_) { - if (firstItr) { - AscendC::WaitFlag(0); - AscendC::WaitFlag(1); - } - } else { - if (endItr) { - unitFlag = 0b11; - } else { - unitFlag = 0b10; - } - } - auto layoutCTileInL0 = tla::MakeLayoutL0C(mRound, nActual); - auto tensorL0CTile = tla::MakeTensor(l0CTensor[0], layoutCTileInL0, Arch::PositionL0C{}); - tileMmad( - tensorL0CTile, tensorL0A, tensorL0B, - mRound, nActual, kActual, initMmad, unitFlag); - AscendC::SetFlag(l0ABPingPongFlag); - AscendC::SetFlag(l0ABPingPongFlag + 2); - auto layoutCInL0 = tla::MakeLayoutL0C(MActual, nActual); - auto tensorL0C = tla::MakeTensor(l0CTensor[0], layoutCInL0, Arch::PositionL0C{}); - if (endItr) { - AscendC::SetFlag(4); - AscendC::SetFlag(5); - if constexpr (!ENABLE_UNIT_FLAG_) { - AscendC::SetFlag(0); - AscendC::WaitFlag(0); - copyL0CToGm(tensorC, tensorL0C); - AscendC::SetFlag(0); - AscendC::SetFlag(1); - } else { - copyL0CToGm(tensorC, tensorL0C, 0b11); - } - } - } - -protected: - /// Data members - AscendC::LocalTensor l1ATensor; - AscendC::LocalTensor l1BTensor[STAGES]; - AscendC::LocalTensor l0ATensor[STAGES]; - AscendC::LocalTensor l0BTensor[STAGES]; - AscendC::LocalTensor l0CTensor[STAGES]; - - TileMmad tileMmad; - CopyL1ToL0A copyL1ToL0A; - CopyL1ToL0B copyL1ToL0B; -}; - -//////////////////////////////////////////////////////////////////// - -} // namespace Catlass::Gemm::Block - -#endif // CATLASS_GEMM_BLOCK_MMAD_PV_TLA_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_pv_tla.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_pv_tla.hpp deleted file mode 100644 index 63066b8283..0000000000 --- a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_pv_tla.hpp +++ /dev/null @@ -1,306 +0,0 @@ -/** - * Copyright (c) 2026 Huawei Technologies Co., Ltd. - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of - * CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ -#ifndef CATLASS_GEMM_BLOCK_BLOCK_MMAD_FAI_PV_TLA_HPP -#define CATLASS_GEMM_BLOCK_BLOCK_MMAD_FAI_PV_TLA_HPP - -#include "catlass/kda_catlass.hpp" -#include "catlass/arch/kda_resource.hpp" -#include "catlass/kda_coord.hpp" -#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" -#include "catlass/gemm/kda_gemm_helper.hpp" -#include "catlass/kda_gemm_coord.hpp" -#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" -#include "catlass/gemm/tile/kda_tile_mmad.hpp" -#include "tla/kda_tensor.hpp" -#include "tla/kda_tla_layout.hpp" - -//////////////////////////////////////////////////////////////////// - -namespace Catlass::Gemm::Block { -//////////////////////////////////////////////////////////////////// - -template < - bool PAGED_CACHE_FLAG_, - bool ENABLE_UNIT_FLAG_, - class L1TileShape_, - class L0TileShape_, - class ElementA_, - class ElementB_, - class ElementC_, - class ElementBias_, - class TileCopy_, - class TileMmad_> -struct BlockMmadTla< - MmadFAIPV, - L1TileShape_, - L0TileShape_, - ElementA_, - ElementB_, - ElementC_, - ElementBias_, - TileCopy_, - TileMmad_> { -public: - using DispatchPolicy = MmadFAIPV; - using ArchTag = typename DispatchPolicy::ArchTag; - using L1TileShape = L1TileShape_; - using L0TileShape = L0TileShape_; - using TileCopy = TileCopy_; - using TileMmad = TileMmad_; - - using ElementA = ElementA_; - using ElementB = ElementB_; - using ElementC = ElementC_; - using ElementAccumulator = typename TileCopy::ElementAccumulator; - - using CopyL1ToL0A = typename TileCopy::CopyL1ToL0A; - using CopyL1ToL0B = typename TileCopy::CopyL1ToL0B; - - using LayoutTagA = typename TileCopy::LayoutTagA; - using LayoutTagB = typename TileCopy::LayoutTagB; - using LayoutTagC = typename TileCopy::LayoutTagC; - - using LayoutTagL1A = typename TileCopy::LayoutTagL1A; - using LayoutTagL1B = typename TileCopy::LayoutTagL1B; - using LayoutTagL0A = typename TileCopy::LayoutTagL0A; - using LayoutTagL0B = typename TileCopy::LayoutTagL0B; - - static constexpr uint32_t STAGES = DispatchPolicy::STAGES; - // L1 tile shape - static constexpr uint32_t L1_TILE_M = tla::get<0>(L1TileShape{}); // s1 - static constexpr uint32_t L1_TILE_K = tla::get<1>(L1TileShape{}); // s2 - static constexpr uint32_t L1_TILE_N = tla::get<2>(L1TileShape{}); // d - // L0 tile shape - static constexpr uint32_t L0_TILE_M = tla::get<0>(L0TileShape{}); // s1 - static constexpr uint32_t L0_TILE_K = tla::get<1>(L0TileShape{}); // s2 - static constexpr uint32_t L0_TILE_N = tla::get<2>(L0TileShape{}); // d - // L1 tile size - static constexpr uint32_t L1A_TILE_SIZE = L1_TILE_M * L1_TILE_K * sizeof(ElementA); - static constexpr uint32_t L1B_TILE_SIZE = L1_TILE_N * L1_TILE_K * sizeof(ElementB); - // L0 tile size - static constexpr uint32_t L0A_TILE_SIZE = L0_TILE_M * L0_TILE_K * sizeof(ElementA); - static constexpr uint32_t L0B_TILE_SIZE = L0_TILE_K * L0_TILE_N * sizeof(ElementB); - static constexpr uint32_t L0C_TILE_SIZE = L0_TILE_M * L0_TILE_N * sizeof(ElementAccumulator); - - static constexpr uint32_t BLOCK_L1_SIZE = L1B_TILE_SIZE * STAGES; - static constexpr uint32_t BLOCK_L0C_SIZE = L0C_TILE_SIZE * STAGES; - - // Check L1/L0TileShape - static_assert(L1_TILE_M == L0_TILE_M && L1_TILE_N == L0_TILE_N && L1_TILE_K == L0_TILE_K, - "The situation where the basic blocks of L1 and L0 differ on the m, n, k axes is not supported yet"); - static_assert(L0A_TILE_SIZE * STAGES <= ArchTag::L0A_SIZE, "L0TileShape exceeding the L0A space!"); - static_assert(L0B_TILE_SIZE * STAGES <= ArchTag::L0B_SIZE, "L0TileShape exceeding the L0B space!"); - -public: - /// Construct - CATLASS_DEVICE - BlockMmadTla(Arch::Resource &resource, uint32_t &l1BufAddrStart, uint32_t &l0CBufAddrStart) - { - uint32_t l1BOffset = l1BufAddrStart; - uint32_t l0COffset = l0CBufAddrStart; - - for (uint32_t i = 0; i < STAGES; i++) { - // Assign L1/L0A/L0B space for each stages - l1BTensorList_[i] = resource.l1Buf.template GetBufferByByte(l1BOffset + L1B_TILE_SIZE * i); - l0ATensorList_[i] = resource.l0ABuf.template GetBufferByByte(L0A_TILE_SIZE * i); - l0BTensorList_[i] = resource.l0BBuf.template GetBufferByByte(L0B_TILE_SIZE * i); - l0CTensorList_[i] = resource.l0CBuf.template GetBufferByByte(l0COffset + L0C_TILE_SIZE * i); - - // Assign event ID for each stages - l1BEventList_[i] = BLOCK_EVENT_ID + i + STAGES; - l0AEventList_[i] = i; - l0BEventList_[i] = i + STAGES; - l0CEventList_[i] = BLOCK_EVENT_ID + i; - - // The event id that needs to be set before the loop - AscendC::SetFlag(l1BEventList_[i]); - AscendC::SetFlag(l0CEventList_[i]); - } - - l1BufAddrStart += BLOCK_L1_SIZE; - l0CBufAddrStart += BLOCK_L0C_SIZE; - } - - /// Destructor - CATLASS_DEVICE - ~BlockMmadTla() - { - for (uint32_t i = 0; i < STAGES; i++) { - AscendC::WaitFlag(l1BEventList_[i]); - AscendC::WaitFlag(l0CEventList_[i]); - } - } - -public: - template - CATLASS_DEVICE void operator()( - TensorA &tensorA, - TensorB &tensorB, - TensorC &tensorC, - AscendC::GlobalTensor &blockTable, - Shape &actualShape, - int64_t taskId, int32_t blockSize) - { - int32_t blockM = tla::get<0>(actualShape); - int32_t blockN = tla::get<1>(actualShape); - int32_t blockK = tla::get<2>(actualShape); - - // load matrix B tile from GM to L1 - auto layoutBInL1 = tla::MakeLayout(blockK, blockN); - auto tensorL1B = tla::MakeTensor(l1BTensorList_[l1BListId_], layoutBInL1, Arch::PositionL1{}); - AscendC::WaitFlag(l1BEventList_[l1BListId_]); - CopyInL1B(tensorL1B, tensorB, blockTable, blockK, blockN, blockSize); - AscendC::SetFlag(l1BEventList_[l1BListId_]); - AscendC::WaitFlag(l1BEventList_[l1BListId_]); - - AscendC::WaitFlag(l0CEventList_[l0CListId_]); - auto layoutInL0C = tla::MakeLayoutL0C(blockM, blockN); - auto tensorL0C = tla::MakeTensor(l0CTensorList_[l0CListId_], layoutInL0C, Arch::PositionL0C{}); - - uint32_t nLoops = (blockN + L0_TILE_N - 1) / L0_TILE_N; - uint32_t tailSize = blockN % L0_TILE_N; - uint32_t tailN = tailSize ? tailSize : L0_TILE_N; - - AscendC::WaitFlag(l0AEventList_[l0AListId_]); - for (uint32_t nIdx = 0; nIdx < nLoops; nIdx++) { - uint32_t tileN = (nIdx == (nLoops - 1)) ? tailN : L0_TILE_N; - - // load matrix A tile from L1 to L0A - auto layoutAInL0 = tla::MakeLayout(blockM, blockK); - auto tensorL0A = tla::MakeTensor(l0ATensorList_[l0AListId_], layoutAInL0, Arch::PositionL0A{}); - if (nIdx == 0) { - auto tensorL1TileA = GetTile( - tensorA, - tla::MakeCoord(0, 0), - tla::MakeShape(blockM, blockK)); - copyL1ToL0A(tensorL0A, tensorL1TileA); - } - - // load matrix B tile from L1 to L0B - auto layoutBInL0 = tla::MakeLayout(blockK, tileN); - auto tensorL0B = tla::MakeTensor(l0BTensorList_[l0BListId_], layoutBInL0, Arch::PositionL0B{}); - auto tensorL1TileB = GetTile( - tensorL1B, - tla::MakeCoord(0, nIdx * L0_TILE_N), - tla::MakeShape(blockK, tileN)); - AscendC::WaitFlag(l0BEventList_[l0BListId_]); - copyL1ToL0B(tensorL0B, tensorL1TileB); - - AscendC::SetFlag(l0CEventList_[l0CListId_]); - AscendC::WaitFlag(l0CEventList_[l0CListId_]); - - // Mmad - bool initC = true; - uint32_t tileM = (blockM == 1 ? M_ALIGN : blockM); - uint32_t tileK = blockK; - auto tensorTileL0C = GetTile( - tensorL0C, - tla::MakeCoord(0, nIdx * L0_TILE_N), - tla::MakeShape(tileM, tileN)); - tileMmad(tensorTileL0C, tensorL0A, tensorL0B, tileM, tileN, tileK, initC); - - AscendC::SetFlag(l0BEventList_[l0BListId_]); - l0BListId_ = (l0BListId_ + 1 < STAGES) ? (l0BListId_ + 1) : 0; - } - AscendC::SetFlag(l0AEventList_[l0AListId_]); - l0AListId_ = (l0AListId_ + 1 < STAGES) ? (l0AListId_ + 1) : 0; - - AscendC::SetFlag(l1BEventList_[l1BListId_]); - l1BListId_ = (l1BListId_ + 1 < STAGES) ? (l1BListId_ + 1) : 0; - - AscendC::SetFlag(l0CEventList_[l0CListId_]); - AscendC::WaitFlag(l0CEventList_[l0CListId_]); - - AscendC::CrossCoreWaitFlag(MM2_RES_INTRA_EVENT[taskId]); - AscendC::CrossCoreWaitFlag(16 + MM2_RES_INTRA_EVENT[taskId]); - - // copy block out - using CopyL0CToDst = typename TileCopy::template CopyL0CToDst; - CopyL0CToDst copyL0CToDst; - copyL0CToDst(tensorC, tensorL0C); - - AscendC::SetFlag(l0CEventList_[l0CListId_]); - l0CListId_ = (l0CListId_ + 1 < STAGES) ? (l0CListId_ + 1) : 0; - } - -private: - template - CATLASS_DEVICE void CopyInL1B( - TensorL1B &tensorL1B, - TensorB &tensorB, - AscendC::GlobalTensor &blockTable, - int32_t tileK, int32_t tileN, int32_t blockSize) - { - using CopyGmToL1B = typename TileCopy::template CopyGmToL1B; - CopyGmToL1B copyGmToL1B; - - if constexpr (PAGED_CACHE_FLAG_) { - int32_t copyRowCnt = 0; - int32_t blockLoops = (tileK + blockSize - 1) / blockSize; - - for (int32_t blockIdx = 0; blockIdx < blockLoops; ++blockIdx) { - int32_t curCopyRow = blockIdx == blockLoops - 1 ? tileK - copyRowCnt : blockSize; - int32_t idInBlockTable = blockTable.GetValue(blockIdx); - - auto tensorL1TileB = GetTile( - tensorL1B, - tla::MakeCoord(copyRowCnt, 0), - tla::MakeShape(curCopyRow, tileN)); - - auto tensorTileB = GetTile( - tensorB, - tla::MakeCoord(idInBlockTable * blockSize, 0), - tla::MakeShape(curCopyRow, tileN)); - - copyGmToL1B(tensorL1TileB, tensorTileB); - copyRowCnt += curCopyRow; - } - } else { - auto tensorTileB = GetTile( - tensorB, - tla::MakeCoord(0, 0), - tla::MakeShape(tileK, tileN)); - - copyGmToL1B(tensorL1B, tensorTileB); - } - } - -private: - static constexpr uint16_t M_ALIGN = 16; - static constexpr uint32_t SYNC_MODE = 4; - static constexpr uint32_t MM2_RES_INTRA_EVENT[2] = {7, 8}; - static constexpr uint32_t BLOCK_EVENT_ID = 4; - - // Multi-stage tensors list - AscendC::LocalTensor l1ATensorList_[STAGES]; - AscendC::LocalTensor l1BTensorList_[STAGES]; - AscendC::LocalTensor l0ATensorList_[STAGES]; - AscendC::LocalTensor l0BTensorList_[STAGES]; - AscendC::LocalTensor l0CTensorList_[STAGES]; - - // Multi-stage event id list - int32_t l1AEventList_[STAGES]; - int32_t l1BEventList_[STAGES]; - int32_t l0AEventList_[STAGES]; - int32_t l0BEventList_[STAGES]; - int32_t l0CEventList_[STAGES]; - - uint32_t l1AListId_{0}; - uint32_t l1BListId_{0}; - uint32_t l0AListId_{0}; - uint32_t l0BListId_{0}; - uint32_t l0CListId_{0}; - - CopyL1ToL0A copyL1ToL0A; - CopyL1ToL0B copyL1ToL0B; - TileMmad tileMmad; -}; -} // namespace Catlass::Gemm::Block -#endif diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_qk_head_tail.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_qk_head_tail.hpp deleted file mode 100644 index 08971b249f..0000000000 --- a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_qk_head_tail.hpp +++ /dev/null @@ -1,306 +0,0 @@ -/** - * Copyright (c) 2025 Huawei Technologies Co., Ltd. - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of - * CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ - -#ifndef CATLASS_GEMM_BLOCK_MMAD_QK_TAIL_HPP -#define CATLASS_GEMM_BLOCK_MMAD_QK_TAIL_HPP - -#include "catlass/kda_catlass.hpp" -#include "catlass/arch/kda_resource.hpp" -#include "catlass/kda_coord.hpp" -#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" -#include "catlass/gemm/kda_gemm_helper.hpp" -#include "catlass/kda_gemm_coord.hpp" -#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" -#include "catlass/gemm/tile/kda_tile_mmad.hpp" - -//////////////////////////////////////////////////////////////////// - -namespace Catlass::Gemm::Block { -//////////////////////////////////////////////////////////////////// - -template < - bool PAGED_CACHE_FLAG_, - bool ENABLE_UNIT_FLAG_, - class L1TileShape_, - class L0TileShape_, - class AType_, - class BType_, - class CType_, - class BiasType_, - class TileCopy_, - class TileMmad_> -struct BlockMmad< - MmadAtlasA2FAITailQK, - L1TileShape_, - L0TileShape_, - AType_, - BType_, - CType_, - BiasType_, - TileCopy_, - TileMmad_> { -public: - // Type Aliases - using DispatchPolicy = MmadAtlasA2FAITailQK; - using ArchTag = typename DispatchPolicy::ArchTag; - using L1TileShape = L1TileShape_; - using L0TileShape = L0TileShape_; - using ElementA = typename AType_::Element; - using LayoutA = typename AType_::Layout; - using ElementB = typename BType_::Element; - using LayoutB = typename BType_::Layout; - using ElementC = typename CType_::Element; - using LayoutC = typename CType_::Layout; - using TileMmad = TileMmad_; - using CopyGmToL1A = typename TileCopy_::CopyGmToL1A; - using CopyGmToL1B = typename TileCopy_::CopyGmToL1B; - using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; - using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; - using CopyL0CToGm = typename TileCopy_::CopyL0CToGm; - using ElementAccumulator = - typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; - using LayoutAInL1 = typename CopyL1ToL0A::LayoutSrc; - using LayoutBInL1 = typename CopyL1ToL0B::LayoutSrc; - using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst; - using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst; - using LayoutCInL0 = layout::zN; - - using L1AAlignHelper = Gemm::helper::L1AlignHelper; - using L1BAlignHelper = Gemm::helper::L1AlignHelper; - - static constexpr uint32_t STAGES = DispatchPolicy::STAGES; - static constexpr uint32_t L1A_SIZE = L1TileShape::M * L1TileShape::K * sizeof(ElementA); - static constexpr uint32_t L1B_SIZE = L1TileShape::N * L1TileShape::K * sizeof(ElementB); - static constexpr uint32_t L0A_SIZE = ArchTag::L0A_SIZE; - static constexpr uint32_t L0B_SIZE = ArchTag::L0B_SIZE; - static constexpr uint32_t L0C_SIZE = ArchTag::L0C_SIZE; - static constexpr uint32_t L0A_PINGPONG_BUF_SIZE = L0A_SIZE / STAGES; - static constexpr uint32_t L0B_PINGPONG_BUF_SIZE = L0B_SIZE / STAGES; - static constexpr uint32_t L0C_PINGPONG_BUF_SIZE = L0C_SIZE / STAGES; - static constexpr uint32_t BLOCK_SIZE = 16; - static constexpr uint32_t EMBED_SPLIT_SIZE = 128; - static constexpr uint32_t UNIT_BLOCK_STACK_NUM = 4; - static constexpr uint32_t KV_BASE_BLOCK = 512; - static constexpr uint32_t KV_SPLIT_SIZE = 128; - - static_assert(std::is_same_v, "LayoutC only support RowMajor yet!"); - - CATLASS_DEVICE - BlockMmad(Arch::Resource &resource, uint32_t l1BufAddrStart = 0) - { - // Allocate L1 memory space - l1ATensor = resource.l1Buf.template GetBufferByByte(l1BufAddrStart); - for (uint32_t i = 0; i < STAGES; i++) { - l1BTensor[i] = - resource.l1Buf.template GetBufferByByte(l1BufAddrStart + L1A_SIZE + L1B_SIZE * i); - l0ATensor[i] = resource.l0ABuf.template GetBufferByByte(L0A_PINGPONG_BUF_SIZE * i); - l0BTensor[i] = resource.l0BBuf.template GetBufferByByte(L0B_PINGPONG_BUF_SIZE * i); - } - l0CTensor = resource.l0CBuf.template GetBufferByByte(0); - } - - /// Destructor - CATLASS_DEVICE - ~BlockMmad() {} - - CATLASS_DEVICE - void loadQGM(AscendC::GlobalTensor gA, LayoutA layoutA, uint32_t rowNum, uint32_t &singleGroupHeads, uint32_t &qHeads) - { - uint32_t embed = layoutA.shape(1); - uint32_t rowNumRound = RoundUp(rowNum); - uint32_t tokenNumPerGroup = rowNum / singleGroupHeads; - auto layoutSingleANd = layoutA.GetTileLayout(MakeCoord(singleGroupHeads, embed)); - LayoutAInL1 layoutAInL1 = LayoutAInL1::template MakeLayout(rowNum, embed); - copyGmToL1A( - l1ATensor, gA, - layoutAInL1, layoutSingleANd, - tokenNumPerGroup, qHeads * embed, tokenNumPerGroup, BLOCK_SIZE, rowNumRound); - AscendC::SetFlag(EVENT_ID3); - AscendC::WaitFlag(EVENT_ID3); - } - - CATLASS_DEVICE - void getBlockShape(GemmCoord &actualShape, uint32_t &nowNIdx, uint32_t &kIdx, - uint32_t &nLoop, uint32_t &kLoop, uint32_t &kvSeqlen, uint32_t &embed, bool firstBlock, uint32_t maskTailS = 0) - { - uint32_t nSplitSize = KV_SPLIT_SIZE; - uint32_t embedSplitSize = EMBED_SPLIT_SIZE; - if (nowNIdx == nLoop - 1) { - nSplitSize = kvSeqlen - nowNIdx * KV_SPLIT_SIZE; - } - if (firstBlock && maskTailS != 0) { - nSplitSize = nSplitSize - maskTailS; - } - if (kIdx == kLoop - 1) { - embedSplitSize = embed - kIdx * EMBED_SPLIT_SIZE; - } - actualShape[1] = nSplitSize; - actualShape[2] = embedSplitSize; - } - - CATLASS_DEVICE - void getKVOffset(AscendC::GlobalTensor &gBlockTable, uint32_t &kOffset, uint32_t &nowNIdx, uint32_t &kIdx, - uint32_t &nLoop, uint32_t &kLoop, uint32_t &strideKV, uint32_t &blockSize, uint32_t maskTailS = 0) - { - if (nowNIdx >= nLoop || kIdx >= kLoop) { - kOffset = 0; - } - if constexpr (PAGED_CACHE_FLAG_) { - uint32_t blockTableId = gBlockTable.GetValue(nowNIdx); - kOffset = blockTableId * blockSize * strideKV + maskTailS * strideKV + kIdx * EMBED_SPLIT_SIZE; - } else { - kOffset = nowNIdx * KV_SPLIT_SIZE * strideKV + kIdx * EMBED_SPLIT_SIZE; - } - } - - CATLASS_DEVICE - void operator()(AscendC::GlobalTensor gA, - AscendC::GlobalTensor gB, - AscendC::GlobalTensor gC, - AscendC::GlobalTensor gBlockTable, - LayoutA layoutA, LayoutB layoutB, GemmCoord actualOriShape, - uint32_t &nIdx, uint32_t &nLoop, uint32_t &blockSize, uint32_t kvSeqlen, uint32_t strideKV, uint32_t maskTailS, bool preloadFlag) - { - uint32_t rowNum = actualOriShape[0]; - uint32_t embed = actualOriShape[2]; - uint32_t kLoop = CeilDiv(embed); - uint32_t nkBlockLoop = nLoop * kLoop; - GemmCoord actualShape{rowNum, 0, 0}; - GemmCoord actualNextShape{rowNum, 0, 0}; - uint32_t nkBlockNextIdx = nIdx * kLoop + 1; - uint32_t gBOffset = 0; - uint32_t gBNextOffset = 0; - uint32_t stackTile = 0; - uint32_t preStackTile = 0; - for (uint32_t blockStackIdx = 0; (blockStackIdx < UNIT_BLOCK_STACK_NUM) && ((nIdx + blockStackIdx) < nLoop); - ++blockStackIdx) { - for (uint32_t kIdx = 0; kIdx < kLoop; kIdx++) { - uint32_t nowNIdx = nIdx + blockStackIdx; - uint32_t nLoopNextIdx = nkBlockNextIdx / kLoop; - uint32_t kLoopNextIdx = nkBlockNextIdx % kLoop; - uint32_t gCOffset = stackTile; - uint32_t startSeqOffset = nowNIdx == nIdx ? maskTailS : 0; - uint32_t startSeqNxtOffset = nLoopNextIdx == nIdx ? maskTailS : 0; - getBlockShape(actualShape, nowNIdx, kIdx, nLoop, kLoop, kvSeqlen, embed, nowNIdx == nIdx, maskTailS); - getBlockShape(actualNextShape, nLoopNextIdx, kLoopNextIdx, nLoop, kLoop, kvSeqlen, embed, nLoopNextIdx == nIdx, maskTailS); - getKVOffset(gBlockTable, gBOffset, nowNIdx, kIdx, nLoop, kLoop, strideKV, blockSize, startSeqOffset); - getKVOffset(gBlockTable, gBNextOffset, nLoopNextIdx, kLoopNextIdx, nLoop, kLoop, strideKV, blockSize, startSeqNxtOffset); - bool firstItr = kIdx == 0; - bool endItr = kIdx == kLoop - 1; - bool firstQtr = blockStackIdx == 0; - bool endQItr = ((nowNIdx == nLoop - 1) || (blockStackIdx == UNIT_BLOCK_STACK_NUM - 1)) && (kIdx == kLoop - 1); - bool initMmad = kIdx == 0; - preStackTile = stackTile; - stackTile += actualShape[1]; - LayoutC layOutSTemp(rowNum, stackTile, 512); - computeQK(gA, gB[gBOffset], gC[gCOffset], gB[gBNextOffset], layoutA, layoutB, layOutSTemp, - actualShape, actualNextShape, blockStackIdx, nkBlockNextIdx, nkBlockLoop, firstItr, endItr, initMmad, firstQtr, endQItr, preloadFlag); - ++nkBlockNextIdx; - preloadFlag = false; - } - maskTailS = 0; - } - } - - CATLASS_DEVICE void computeQK( - AscendC::GlobalTensor const &gA, - AscendC::GlobalTensor const &gB, - AscendC::GlobalTensor const &gC, - AscendC::GlobalTensor const &gmNextBlockB, - LayoutA layoutA, LayoutB layoutB, LayoutC layoutC, - GemmCoord actualShape, GemmCoord actualNextShape, uint32_t nowIdx, uint32_t &nkblockIdx, - uint32_t &nkblockLoop, bool firstItr, bool endItr, bool initMmad, bool firstQItr, bool endQItr, bool preloadFlag) - { - uint32_t mActual = actualShape.m(); - uint32_t kActual = actualShape.k(); - uint32_t nActual = actualShape.n(); - uint32_t mRound = RoundUp(mActual); - uint32_t kRound = RoundUp(kActual); - uint32_t nRound = RoundUp(nActual); - LayoutAInL1 layoutAInL1 = LayoutAInL1::template MakeLayout(mRound, kActual); // embed - LayoutAInL0 layoutAInL0 = LayoutAInL0::template MakeLayout(mRound, kActual); - LayoutBInL1 layoutBInL1 = LayoutBInL1::template MakeLayout(kActual, nActual); - LayoutBInL0 layoutBInL0 = LayoutBInL0::template MakeLayout(kActual, nRound); - uint32_t locPingPongFlag = nkblockIdx % 2; - uint32_t l1KvPingPongFlag = nkblockIdx % 2; - uint32_t l0ABPingPongFlag = nkblockIdx % 2; - if (preloadFlag) { - auto layoutBTile = layoutB.GetTileLayout(MakeCoord(kActual, nActual)); - AscendC::WaitFlag(l1KvPingPongFlag); - copyGmToL1B(l1BTensor[l1KvPingPongFlag], gB, layoutBInL1, layoutBTile); - AscendC::SetFlag(l1KvPingPongFlag); - } - if (nkblockIdx != nkblockLoop) { - uint32_t nNextActual = actualNextShape.n(); - uint32_t kNextActual = actualNextShape.k(); - LayoutBInL1 layoutBNextInL1 = LayoutBInL1::template MakeLayout(kNextActual, nNextActual); - auto layoutNextBTile = layoutB.GetTileLayout(MakeCoord(kNextActual, nNextActual)); - AscendC::WaitFlag(1 - l1KvPingPongFlag); - copyGmToL1B(l1BTensor[1 - l1KvPingPongFlag], gmNextBlockB, layoutBNextInL1, layoutNextBTile); - AscendC::SetFlag(1 - l1KvPingPongFlag); - } - AscendC::WaitFlag(l0ABPingPongFlag); - copyL1ToL0A(l0ATensor[l0ABPingPongFlag], l1ATensor, layoutAInL0, layoutAInL1); - - AscendC::WaitFlag(l1KvPingPongFlag); - AscendC::WaitFlag(l0ABPingPongFlag + 2); - copyL1ToL0B(l0BTensor[l0ABPingPongFlag], l1BTensor[l1KvPingPongFlag], layoutBInL0, layoutBInL1); - AscendC::SetFlag(l1KvPingPongFlag); - - AscendC::SetFlag(l0ABPingPongFlag); - AscendC::WaitFlag(l0ABPingPongFlag); - uint8_t unitFlag = 0b00; - if constexpr (!ENABLE_UNIT_FLAG_) { - if (firstItr) { - AscendC::WaitFlag(locPingPongFlag); - } - } else { - unitFlag = 0b11; - } - tileMmad( - l0CTensor[locPingPongFlag * mRound * 128], l0ATensor[l0ABPingPongFlag], l0BTensor[l0ABPingPongFlag], - mRound, nActual, kActual, initMmad, unitFlag); - // AscendC::PipeBarrier(); - AscendC::SetFlag(l0ABPingPongFlag); - AscendC::SetFlag(l0ABPingPongFlag + 2); - auto layoutInL0C = LayoutCInL0::MakeLayoutInL0C(MakeCoord(mActual, nActual)); - if (endItr) { - if constexpr (!ENABLE_UNIT_FLAG_) { - AscendC::SetFlag(locPingPongFlag); - AscendC::WaitFlag(locPingPongFlag); - copyL0CToGm(gC, l0CTensor[locPingPongFlag * mRound * 128], layoutC, layoutInL0C); - AscendC::SetFlag(locPingPongFlag); - } else { - copyL0CToGm(gC, l0CTensor[locPingPongFlag * mRound * 128], layoutC, layoutInL0C, unitFlag); - } - } - } - -protected: - /// Data members - AscendC::LocalTensor l1ATensor; - AscendC::LocalTensor l1BTensor[STAGES]; - AscendC::LocalTensor l0ATensor[STAGES]; - AscendC::LocalTensor l0BTensor[STAGES]; - AscendC::LocalTensor l0CTensor; - - TileMmad tileMmad; - CopyGmToL1A copyGmToL1A; - CopyGmToL1B copyGmToL1B; - CopyL1ToL0A copyL1ToL0A; - CopyL1ToL0B copyL1ToL0B; - CopyL0CToGm copyL0CToGm; -}; - -//////////////////////////////////////////////////////////////////// - -} // namespace Catlass::Gemm::Block - -#endif // CATLASS_GEMM_BLOCK_MMAD_QK_TAIL_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_qk_head_tail_tla.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_qk_head_tail_tla.hpp deleted file mode 100644 index 75b24d6608..0000000000 --- a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_qk_head_tail_tla.hpp +++ /dev/null @@ -1,326 +0,0 @@ -/** - * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of - * CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ - -#ifndef CATLASS_GEMM_BLOCK_MMAD_QK_TAIL_TLA_HPP -#define CATLASS_GEMM_BLOCK_MMAD_QK_TAIL_TLA_HPP - -#include "catlass/kda_catlass.hpp" -#include "catlass/arch/kda_resource.hpp" -#include "catlass/kda_coord.hpp" -#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" -#include "catlass/gemm/kda_gemm_helper.hpp" -#include "catlass/kda_gemm_coord.hpp" -#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" -#include "catlass/gemm/tile/kda_tile_mmad.hpp" -#include "tla/kda_tla_layout.hpp" -#include "tla/kda_tensor.hpp" - -//////////////////////////////////////////////////////////////////// - -namespace Catlass::Gemm::Block { -//////////////////////////////////////////////////////////////////// - -template < - class ArchTag_, - bool PAGED_CACHE_FLAG_, - bool ENABLE_UNIT_FLAG_, - class L1TileShape_, - class L0TileShape_, - class ElementA_, - class ElementB_, - class ElementC_, - class ElementBias_, - class TileCopy_, - class TileMmad_> -struct BlockMmadTla< - MmadFAITailQK, - L1TileShape_, - L0TileShape_, - ElementA_, - ElementB_, - ElementC_, - ElementBias_, - TileCopy_, - TileMmad_> { -public: - // Type Aliases - using DispatchPolicy = MmadFAITailQK; - using ArchTag = typename DispatchPolicy::ArchTag; - using L1TileShape = L1TileShape_; - using L0TileShape = L0TileShape_; - using ElementA = ElementA_; - using LayoutA = typename TileCopy_::LayoutA; - using ElementB = ElementB_; - using LayoutB = typename TileCopy_::LayoutB; - using ElementC = ElementC_; - using LayoutC = typename TileCopy_::LayoutC; - - using TileMmad = TileMmad_; - - using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; - using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; - - using ElementAccumulator = typename TileCopy_::ElementAccumulator; - - using LayoutTagL1A = typename TileCopy_::LayoutTagL1A; - using LayoutTagL1B = typename TileCopy_::LayoutTagL1B; - using LayoutTagL0A = typename TileCopy_::LayoutTagL0A; - using LayoutTagL0B = typename TileCopy_::LayoutTagL0B; - - using L1AAlignHelper = typename TileCopy_::L1AAlignHelper; - using L1BAlignHelper = typename TileCopy_::L1BAlignHelper; - - static_assert(tla::is_tuple::value && tla::is_static::value, - "L1TileShape must be tla::tuple and static!"); - static_assert(tla::is_tuple::value && tla::is_static::value, - "L0TileShape must be tla::tuple and static!"); - - static constexpr bool ENABLE_UNIT_FLAG = DispatchPolicy::ENABLE_UNIT_FLAG; - static constexpr uint32_t STAGES = DispatchPolicy::STAGES; - static constexpr uint32_t L1_TILE_M = tla::get<0>(L1TileShape{}); - static constexpr uint32_t L1_TILE_N = tla::get<1>(L1TileShape{}); - static constexpr uint32_t L1_TILE_K = tla::get<2>(L1TileShape{}); - static constexpr uint32_t L0_TILE_M = tla::get<0>(L0TileShape{}); - static constexpr uint32_t L0_TILE_N = tla::get<1>(L0TileShape{}); - static constexpr uint32_t L0_TILE_K = tla::get<2>(L0TileShape{}); - - // L1 tile size - static constexpr uint32_t L1A_TILE_SIZE = L1_TILE_M * L1_TILE_K * sizeof(ElementA); - static constexpr uint32_t L1B_TILE_SIZE = L1_TILE_N * L1_TILE_K * sizeof(ElementB); - // L0 tile size - static constexpr uint32_t L0A_TILE_SIZE = L0_TILE_M * L0_TILE_K * sizeof(ElementA); - static constexpr uint32_t L0B_TILE_SIZE = L0_TILE_K * L0_TILE_N * sizeof(ElementB); - static constexpr uint32_t L0C_TILE_SIZE = L0_TILE_M * L0_TILE_N * sizeof(ElementAccumulator); - static constexpr uint32_t L0A_PINGPONG_BUF_SIZE = ArchTag::L0A_SIZE / STAGES; - static constexpr uint32_t L0B_PINGPONG_BUF_SIZE = ArchTag::L0B_SIZE / STAGES; - static constexpr uint32_t L0C_PINGPONG_BUF_SIZE = ArchTag::L0C_SIZE / STAGES; - - // Check LayoutC - static_assert(tla::detail::isRowMajor::value || - ((std::is_same_v || std::is_same_v || - std::is_same_v) && - tla::detail::iszN::value), - "LayoutC only supports zN in half or bfloat16 or float, RowMajor in all dtype yet!"); - - // Check L1TileShape - static_assert(L1A_TILE_SIZE + L1B_TILE_SIZE * STAGES <= ArchTag::L1_SIZE, - "L1TileShape exceeding the L1 space!"); - - // Check L0TileShape - static_assert(L0A_TILE_SIZE * STAGES <= ArchTag::L0A_SIZE, "L0TileShape exceeding the L0A space!"); - static_assert(L0B_TILE_SIZE * STAGES <= ArchTag::L0B_SIZE, "L0TileShape exceeding the L0B space!"); - static_assert(L0C_TILE_SIZE <= ArchTag::L0C_SIZE, "L0TileShape exceeding the L0C space!"); - - static constexpr uint32_t BLOCK_SIZE = 16; - static constexpr uint32_t EMBED_SPLIT_SIZE = 128; - static constexpr uint32_t UNIT_BLOCK_STACK_NUM = 4; - static constexpr uint32_t KV_BASE_BLOCK = 512; - static constexpr uint32_t KV_SPLIT_SIZE = 128; - - CATLASS_DEVICE - BlockMmadTla(Arch::Resource &resource, uint32_t l1BufAddrStart = 0) - { - // Allocate L1 memory space - l1ATensor = resource.l1Buf.template GetBufferByByte(l1BufAddrStart); - for (uint32_t i = 0; i < STAGES; i++) { - l1BTensor[i] = - resource.l1Buf.template GetBufferByByte(l1BufAddrStart + L1A_TILE_SIZE + L1B_TILE_SIZE * i); - l0ATensor[i] = resource.l0ABuf.template GetBufferByByte(L0A_PINGPONG_BUF_SIZE * i); - l0BTensor[i] = resource.l0BBuf.template GetBufferByByte(L0B_PINGPONG_BUF_SIZE * i); - } - l0CTensor = resource.l0CBuf.template GetBufferByByte(0); - } - - /// Destructor - CATLASS_DEVICE - ~BlockMmadTla() {} - - CATLASS_DEVICE - void getBlockShape(GemmCoord &actualShape, uint32_t &nowNIdx, uint32_t &kIdx, - uint32_t &nLoop, uint32_t &kLoop, uint32_t &kvSeqlen, uint32_t &embed, bool firstBlock, uint32_t maskTailS = 0) - { - uint32_t nSplitSize = KV_SPLIT_SIZE; - uint32_t embedSplitSize = EMBED_SPLIT_SIZE; - if (nowNIdx == nLoop - 1) { - nSplitSize = kvSeqlen - nowNIdx * KV_SPLIT_SIZE; - } - if (firstBlock && maskTailS != 0) { - nSplitSize = nSplitSize - maskTailS; - } - if (kIdx == kLoop - 1) { - embedSplitSize = embed - kIdx * EMBED_SPLIT_SIZE; - } - actualShape[1] = nSplitSize; - actualShape[2] = embedSplitSize; - } - - CATLASS_DEVICE - void getKVOffset(AscendC::GlobalTensor &gBlockTable, uint32_t &kOffset, uint32_t &nowNIdx, uint32_t &kIdx, - uint32_t &nLoop, uint32_t &kLoop, uint32_t &strideKV, uint32_t &blockSize, uint32_t maskTailS = 0) - { - if (nowNIdx >= nLoop || kIdx >= kLoop) { - kOffset = 0; - } - if constexpr (PAGED_CACHE_FLAG_) { - uint32_t blockTableId = gBlockTable.GetValue(nowNIdx); - kOffset = blockTableId * blockSize * strideKV + maskTailS * strideKV + kIdx * EMBED_SPLIT_SIZE; - } else { - kOffset = nowNIdx * KV_SPLIT_SIZE * strideKV + kIdx * EMBED_SPLIT_SIZE; - } - } - - template - CATLASS_DEVICE void operator()(TensorA &tensorA, TensorB &tensorB, AscendC::GlobalTensor gC, - AscendC::GlobalTensor gBlockTable, GemmCoord actualOriShape, - uint32_t &nIdx, uint32_t &nLoop, uint32_t &blockSize, uint32_t kvSeqlen, uint32_t strideKV, uint32_t maskTailS, bool preloadFlag) - { - uint32_t rowNum = actualOriShape[0]; - uint32_t embed = actualOriShape[2]; - uint32_t kLoop = CeilDiv(embed); - uint32_t nkBlockLoop = nLoop * kLoop; - GemmCoord actualShape{rowNum, 0, 0}; - GemmCoord actualNextShape{rowNum, 0, 0}; - uint32_t nkBlockNextIdx = nIdx * kLoop + 1; - uint32_t gBOffset = 0; - uint32_t gBNextOffset = 0; - uint32_t stackTile = 0; - uint32_t preStackTile = 0; - for (uint32_t blockStackIdx = 0; (blockStackIdx < UNIT_BLOCK_STACK_NUM) && ((nIdx + blockStackIdx) < nLoop); - ++blockStackIdx) { - for (uint32_t kIdx = 0; kIdx < kLoop; kIdx++) { - uint32_t nowNIdx = nIdx + blockStackIdx; - uint32_t nLoopNextIdx = nkBlockNextIdx / kLoop; - uint32_t kLoopNextIdx = nkBlockNextIdx % kLoop; - uint32_t gCOffset = stackTile; - uint32_t startSeqOffset = nowNIdx == nIdx ? maskTailS : 0; - uint32_t startSeqNxtOffset = nLoopNextIdx == nIdx ? maskTailS : 0; - getBlockShape(actualShape, nowNIdx, kIdx, nLoop, kLoop, kvSeqlen, embed, nowNIdx == nIdx, maskTailS); - getBlockShape(actualNextShape, nLoopNextIdx, kLoopNextIdx, nLoop, kLoop, kvSeqlen, embed, nLoopNextIdx == nIdx, maskTailS); - getKVOffset(gBlockTable, gBOffset, nowNIdx, kIdx, nLoop, kLoop, strideKV, blockSize, startSeqOffset); - getKVOffset(gBlockTable, gBNextOffset, nLoopNextIdx, kLoopNextIdx, nLoop, kLoop, strideKV, blockSize, startSeqNxtOffset); - bool firstItr = kIdx == 0; - bool endItr = kIdx == kLoop - 1; - bool firstQtr = blockStackIdx == 0; - bool endQItr = ((nowNIdx == nLoop - 1) || (blockStackIdx == UNIT_BLOCK_STACK_NUM - 1)) && (kIdx == kLoop - 1); - bool initMmad = kIdx == 0; - preStackTile = stackTile; - stackTile += actualShape[1]; - auto layoutS = tla::MakeLayout(tla::MakeShape(rowNum, stackTile), tla::MakeStride(Int<512>{}, Int<1>{})); - auto tensorCOffset = tla::MakeTensor(gC[gCOffset], layoutS, Arch::PositionGM{}); - computeQK(tensorA, tensorB, tensorCOffset, gBOffset, gBNextOffset, - actualShape, actualNextShape, blockStackIdx, nkBlockNextIdx, nkBlockLoop, strideKV, firstItr, endItr, initMmad, firstQtr, endQItr, preloadFlag); - ++nkBlockNextIdx; - preloadFlag = false; - } - maskTailS = 0; - } - } - - template - CATLASS_DEVICE void computeQK( - TensorA &tensorA, TensorB &tensorB, TensorC &tensorC, uint32_t gBOffset, uint32_t gBNextOffset, - GemmCoord actualShape, GemmCoord actualNextShape, uint32_t nowIdx, uint32_t &nkblockIdx, - uint32_t &nkblockLoop, uint32_t strideKV, bool firstItr, bool endItr, bool initMmad, bool firstQItr, bool endQItr, bool preloadFlag) - { - using CopyGmToL1B = typename TileCopy_::template CopyGmToL1B; - using CopyL0CToGm = typename TileCopy_::template CopyL0CToGm; - CopyGmToL1B copyGmToL1B; - CopyL0CToGm copyL0CToGm; - - uint32_t mActual = actualShape.m(); - uint32_t kActual = actualShape.k(); - uint32_t nActual = actualShape.n(); - uint32_t mRound = RoundUp(mActual); - uint32_t kRound = RoundUp(kActual); - uint32_t nRound = RoundUp(nActual); - auto layoutAInL1 = tla::MakeLayout(mRound, kActual); - auto layoutAInL0 = tla::MakeLayout(mRound, kActual); - auto layoutBInL1 = tla::MakeLayout(kActual, nActual); - auto layoutBInL0 = tla::MakeLayout(kActual, nRound); - uint32_t locPingPongFlag = nkblockIdx % 2; - uint32_t l1KvPingPongFlag = nkblockIdx % 2; - uint32_t l0ABPingPongFlag = nkblockIdx % 2; - auto tensorL1A = tla::MakeTensor(l1ATensor, layoutAInL1, Arch::PositionL1{}); - auto tensorL0A = tla::MakeTensor(l0ATensor[l0ABPingPongFlag], layoutAInL0, Arch::PositionL0A{}); - auto tensorTileB = GetTile(tensorB, tla::MakeCoord(gBOffset % strideKV, gBOffset / strideKV), tla::MakeShape(kActual, nActual)); - auto tensorL1B = tla::MakeTensor(l1BTensor[l1KvPingPongFlag], layoutBInL1, Arch::PositionL1{}); - auto tensorL0B = tla::MakeTensor(l0BTensor[l0ABPingPongFlag], layoutBInL0, Arch::PositionL0B{}); - if (preloadFlag) { - AscendC::WaitFlag(l1KvPingPongFlag); - copyGmToL1B(tensorL1B, tensorTileB); - AscendC::SetFlag(l1KvPingPongFlag); - } - if (nkblockIdx != nkblockLoop) { - uint32_t nNextActual = actualNextShape.n(); - uint32_t kNextActual = actualNextShape.k(); - auto tensorTileBNext = GetTile(tensorB, tla::MakeCoord(gBNextOffset % strideKV, gBNextOffset / strideKV), tla::MakeShape(kNextActual, nNextActual)); - auto layoutBNextInL1 = tla::MakeLayout(kNextActual, nNextActual); - auto tensorL1BNext = tla::MakeTensor(l1BTensor[1 - l1KvPingPongFlag], layoutBNextInL1, Arch::PositionL1{}); - AscendC::WaitFlag(1 - l1KvPingPongFlag); - copyGmToL1B(tensorL1BNext, tensorTileBNext); - AscendC::SetFlag(1 - l1KvPingPongFlag); - } - AscendC::WaitFlag(l0ABPingPongFlag); - copyL1ToL0A(tensorL0A, tensorL1A); - - AscendC::WaitFlag(l1KvPingPongFlag); - AscendC::WaitFlag(l0ABPingPongFlag + 2); - copyL1ToL0B(tensorL0B, tensorL1B); - AscendC::SetFlag(l1KvPingPongFlag); - - AscendC::SetFlag(l0ABPingPongFlag); - AscendC::WaitFlag(l0ABPingPongFlag); - uint8_t unitFlag = 0b00; - if constexpr (!ENABLE_UNIT_FLAG_) { - if (firstItr) { - AscendC::WaitFlag(locPingPongFlag); - } - } else { - unitFlag = 0b11; - } - auto layoutCTileInL0 = tla::MakeLayoutL0C(mRound, nActual); - auto tensorL0CTile = tla::MakeTensor(l0CTensor[locPingPongFlag * mRound * 128], layoutCTileInL0, Arch::PositionL0C{}); - tileMmad( - tensorL0CTile, tensorL0A, tensorL0B, - mRound, nActual, kActual, initMmad, unitFlag); - AscendC::SetFlag(l0ABPingPongFlag); - AscendC::SetFlag(l0ABPingPongFlag + 2); - auto layoutCInL0 = tla::MakeLayoutL0C(mActual, nActual); - auto tensorL0C = tla::MakeTensor(l0CTensor[locPingPongFlag * mRound * 128], layoutCInL0, Arch::PositionL0C{}); - if (endItr) { - if constexpr (!ENABLE_UNIT_FLAG_) { - AscendC::SetFlag(locPingPongFlag); - AscendC::WaitFlag(locPingPongFlag); - copyL0CToGm(tensorC, tensorL0C); - AscendC::SetFlag(locPingPongFlag); - } else { - copyL0CToGm(tensorC, tensorL0C, unitFlag); - } - } - } - -protected: - /// Data members - AscendC::LocalTensor l1ATensor; - AscendC::LocalTensor l1BTensor[STAGES]; - AscendC::LocalTensor l0ATensor[STAGES]; - AscendC::LocalTensor l0BTensor[STAGES]; - AscendC::LocalTensor l0CTensor; - - TileMmad tileMmad; - CopyL1ToL0A copyL1ToL0A; - CopyL1ToL0B copyL1ToL0B; -}; - -//////////////////////////////////////////////////////////////////// - -} // namespace Catlass::Gemm::Block - -#endif // CATLASS_GEMM_BLOCK_MMAD_QK_TAIL_TLA_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_qk_normal.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_qk_normal.hpp deleted file mode 100644 index d43f63ba13..0000000000 --- a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_qk_normal.hpp +++ /dev/null @@ -1,337 +0,0 @@ -/** - * Copyright (c) 2025 Huawei Technologies Co., Ltd. - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of - * CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ - -#ifndef CATLASS_GEMM_BLOCK_MMAD_QK_HPP -#define CATLASS_GEMM_BLOCK_MMAD_QK_HPP - -#include "catlass/kda_catlass.hpp" -#include "catlass/arch/kda_resource.hpp" -#include "catlass/kda_coord.hpp" -#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" -#include "catlass/gemm/kda_gemm_helper.hpp" -#include "catlass/kda_gemm_coord.hpp" -#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" -#include "catlass/gemm/tile/kda_tile_mmad.hpp" - -//////////////////////////////////////////////////////////////////// - -namespace Catlass::Gemm::Block { -//////////////////////////////////////////////////////////////////// - -template < - bool PAGED_CACHE_FLAG_, - bool ENABLE_UNIT_FLAG_, - class L1TileShape_, - class L0TileShape_, - class AType_, - class BType_, - class CType_, - class BiasType_, - class TileCopy_, - class TileMmad_> -struct BlockMmad< - MmadAtlasA2FAIQK, - L1TileShape_, - L0TileShape_, - AType_, - BType_, - CType_, - BiasType_, - TileCopy_, - TileMmad_> { -public: - // Type Aliases - using DispatchPolicy = MmadAtlasA2FAIQK; - using ArchTag = typename DispatchPolicy::ArchTag; - using L1TileShape = L1TileShape_; - using L0TileShape = L0TileShape_; - using ElementA = typename AType_::Element; - using LayoutA = typename AType_::Layout; - using ElementB = typename BType_::Element; - using LayoutB = typename BType_::Layout; - using ElementC = typename CType_::Element; - using LayoutC = typename CType_::Layout; - using TileMmad = TileMmad_; - using CopyGmToL1A = typename TileCopy_::CopyGmToL1A; - using CopyGmToL1B = typename TileCopy_::CopyGmToL1B; - using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; - using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; - using CopyL0CToGm = typename TileCopy_::CopyL0CToGm; - using ElementAccumulator = - typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; - using LayoutAInL1 = typename CopyL1ToL0A::LayoutSrc; - using LayoutBInL1 = typename CopyL1ToL0B::LayoutSrc; - using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst; - using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst; - using LayoutCInL0 = layout::zN; - - using L1AAlignHelper = Gemm::helper::L1AlignHelper; - using L1BAlignHelper = Gemm::helper::L1AlignHelper; - - static constexpr uint32_t STAGES = DispatchPolicy::STAGES; - static constexpr uint32_t L1A_SIZE = L1TileShape::M * L1TileShape::K * sizeof(ElementA); - static constexpr uint32_t L1B_SIZE = L1TileShape::N * L1TileShape::K * sizeof(ElementB); - static constexpr uint32_t L0A_SIZE = ArchTag::L0A_SIZE; - static constexpr uint32_t L0B_SIZE = ArchTag::L0B_SIZE; - static constexpr uint32_t L0C_SIZE = ArchTag::L0C_SIZE; - static constexpr uint32_t L0A_PINGPONG_BUF_SIZE = L0A_SIZE / STAGES; - static constexpr uint32_t L0B_PINGPONG_BUF_SIZE = L0B_SIZE / STAGES; - static constexpr uint32_t L0C_PINGPONG_BUF_SIZE = L0C_SIZE / STAGES; - static constexpr uint32_t BLOCK_SIZE = 16; - static constexpr uint32_t EMBED_SPLIT_SIZE = 128; - static constexpr uint32_t UNIT_BLOCK_STACK_NUM = 4; - static constexpr uint32_t KV_BASE_BLOCK = 512; - static constexpr uint32_t KV_SPLIT_SIZE = 128; - - static_assert(std::is_same_v, "LayoutC only support RowMajor yet!"); - - CATLASS_DEVICE - BlockMmad(Arch::Resource &resource, uint32_t l1BufAddrStart = 0) - { - // Allocate L1 memory space - l1ATensor = resource.l1Buf.template GetBufferByByte(l1BufAddrStart); - for (uint32_t i = 0; i < STAGES; i++) { - l1BTensor[i] = - resource.l1Buf.template GetBufferByByte(l1BufAddrStart + L1A_SIZE + L1B_SIZE * i); - l0ATensor[i] = resource.l0ABuf.template GetBufferByByte(L0A_PINGPONG_BUF_SIZE * i); - l0BTensor[i] = resource.l0BBuf.template GetBufferByByte(L0B_PINGPONG_BUF_SIZE * i); - } - l0CTensor = resource.l0CBuf.template GetBufferByByte(0); - } - - CATLASS_DEVICE - ~BlockMmad() {} - - CATLASS_DEVICE - void loadQGM(AscendC::GlobalTensor gA, LayoutA layoutA, uint32_t rowNum, uint32_t &singleGroupHeads, uint32_t &qHeads) - { - uint32_t embed = layoutA.shape(1); - uint32_t rowNumRound = RoundUp(rowNum); - uint32_t tokenNumPerGroup = rowNum / singleGroupHeads; - auto layoutSingleANd = layoutA.GetTileLayout(MakeCoord(singleGroupHeads, embed)); - LayoutAInL1 layoutAInL1 = LayoutAInL1::template MakeLayout(rowNum, embed); - copyGmToL1A( - l1ATensor, gA, - layoutAInL1, layoutSingleANd, - tokenNumPerGroup, qHeads * embed, tokenNumPerGroup, BLOCK_SIZE, rowNumRound); - // AscendC::Nd2NzParams intriParams; - // intriParams.nValue = singleGroupHeads; - // intriParams.dValue = embed; - // intriParams.srcDValue = embed; - // intriParams.dstNzNStride = tokenNumPerGroup; - // intriParams.dstNzC0Stride = rowNumRound; - // intriParams.ndNum = tokenNumPerGroup; - // intriParams.dstNzMatrixStride = 16; - // AscendC::DataCopy(l1ATensor, gA, intriParams); - AscendC::SetFlag(EVENT_ID3); - AscendC::WaitFlag(EVENT_ID3); - } - - CATLASS_DEVICE - void getBlockShape(GemmCoord &actualShape, uint32_t &nowNIdx, uint32_t &kIdx, - uint32_t &nLoop, uint32_t &kLoop, uint32_t &kvSeqlen, uint32_t &embed, bool firstBlock, uint32_t maskTailS = 0) - { - uint32_t nSplitSize = KV_SPLIT_SIZE; - uint32_t embedSplitSize = EMBED_SPLIT_SIZE; - if (nowNIdx == nLoop - 1) { - nSplitSize = kvSeqlen - nowNIdx * KV_SPLIT_SIZE; - } - if (firstBlock && maskTailS != 0) { - nSplitSize = nSplitSize - maskTailS; - } - // } - if (kIdx == kLoop - 1) { - embedSplitSize = embed - kIdx * EMBED_SPLIT_SIZE; - } - actualShape[1] = nSplitSize; - actualShape[2] = embedSplitSize; - } - - CATLASS_DEVICE - void getKVOffset(AscendC::GlobalTensor &gBlockTable, uint32_t &kOffset, uint32_t &nowNIdx, uint32_t &kIdx, - uint32_t &nLoop, uint32_t &kLoop, uint32_t &strideKV, uint32_t &blockSize, uint32_t maskTailS = 0) - { - if (nowNIdx >= nLoop || kIdx >= kLoop) { - kOffset = 0; - } - if constexpr (PAGED_CACHE_FLAG_) { - uint32_t blockTableId = gBlockTable.GetValue(nowNIdx); - kOffset = blockTableId * blockSize * strideKV + maskTailS * strideKV + kIdx * EMBED_SPLIT_SIZE; - } else { - kOffset = nowNIdx * KV_SPLIT_SIZE * strideKV + kIdx * EMBED_SPLIT_SIZE; - } - } - - CATLASS_DEVICE - void operator()(AscendC::GlobalTensor gA, - AscendC::GlobalTensor gB, - AscendC::GlobalTensor gC, - AscendC::GlobalTensor gBlockTable, - LayoutA layoutA, LayoutB layoutB, GemmCoord actualOriShape, - uint32_t &nIdx, uint32_t &nLoop, uint32_t &blockSize, uint32_t kvSeqlen, uint32_t strideKV) - { - uint32_t rowNum = actualOriShape[0]; - uint32_t embed = actualOriShape[2]; - uint32_t kLoop = CeilDiv(embed); - uint32_t nkBlockLoop = nLoop * kLoop; - GemmCoord actualShape{rowNum, 0, 0}; - GemmCoord actualNextShape{rowNum, 0, 0}; - uint32_t nkBlockNextIdx = nIdx * kLoop + 1; - uint32_t gBOffset = 0; - uint32_t gBNextOffset = 0; - uint32_t stackTile = 0; - for (uint32_t blockStackIdx = 0; (blockStackIdx < UNIT_BLOCK_STACK_NUM) && ((nIdx + blockStackIdx) < nLoop); - ++blockStackIdx) { - for (uint32_t kIdx = 0; kIdx < kLoop; kIdx++) { - uint32_t nowNIdx = nIdx + blockStackIdx; - uint32_t nLoopNextIdx = nkBlockNextIdx / kLoop; - uint32_t kLoopNextIdx = nkBlockNextIdx % kLoop; - uint32_t gCOffset = blockStackIdx / 2 * 2 * KV_SPLIT_SIZE; - getBlockShape(actualShape, nowNIdx, kIdx, nLoop, kLoop, kvSeqlen, embed, nowNIdx == nIdx); - getBlockShape(actualNextShape, nLoopNextIdx, kLoopNextIdx, nLoop, kLoop, kvSeqlen, embed, nLoopNextIdx == nIdx); - getKVOffset(gBlockTable, gBOffset, nowNIdx, kIdx, nLoop, kLoop, strideKV, blockSize); - getKVOffset(gBlockTable, gBNextOffset, nLoopNextIdx, kLoopNextIdx, nLoop, kLoop, strideKV, blockSize); - bool firstItr = ((blockStackIdx % 2) == 0) && (kIdx == 0); - bool endItr = (((blockStackIdx % 2) == 1) || (nowNIdx == nLoop - 1)) && (kIdx == kLoop - 1); - bool firstQtr = blockStackIdx == 0; - bool endQItr = ((nowNIdx == nLoop - 1) || (blockStackIdx == UNIT_BLOCK_STACK_NUM - 1)) && (kIdx == kLoop - 1); - int cc = 1; - bool initMmad = kIdx == 0; - stackTile += actualShape[1]; - LayoutC layOutSTemp(rowNum, stackTile, 512); - // AscendC::printf("firstItr:%d\n", firstItr); - // AscendC::printf("endItr:%d\n", endItr); - // AscendC::printf("initMmad:%d\n", initMmad); - // AscendC::printf("stackTile:%d\n", stackTile); - // AscendC::printf("blockStackIdx:%d\n", blockStackIdx); - // AscendC::printf("gBOffset:%d\n", gBOffset); - // AscendC::printf("gBNextOffset:%d\n", gBNextOffset); - // AscendC::printf("actualShape[0]:%d\n", actualShape[0]); - // AscendC::printf("actualShape[1]:%d\n", actualShape[1]); - // AscendC::printf("actualShape[2]:%d\n", actualShape[2]); - // AscendC::printf("actualShape.m():%d\n", actualShape.m()); - // AscendC::printf("actualShape.n():%d\n", actualShape.n()); - // AscendC::printf("actualShape.k():%d\n", actualShape.k()); - // AscendC::printf("actualNextShape[0]:%d\n", actualNextShape[0]); - // AscendC::printf("actualNextShape[1]:%d\n", actualNextShape[1]); - // AscendC::printf("actualNextShape[2]:%d\n", actualNextShape[2]); - computeQK(gA, gB[gBOffset], gC[gCOffset], gB[gBNextOffset], layoutA, layoutB, layOutSTemp, - actualShape, actualNextShape, blockStackIdx, nkBlockNextIdx, nkBlockLoop, firstItr, endItr, initMmad, firstQtr, endQItr); - ++nkBlockNextIdx; - if (endItr) { - stackTile = 0; - } - } - } - } - - CATLASS_DEVICE void computeQK( - AscendC::GlobalTensor const &gA, - AscendC::GlobalTensor const &gB, - AscendC::GlobalTensor const &gC, - AscendC::GlobalTensor const &gmNextBlockB, - LayoutA layoutA, LayoutB layoutB, LayoutC layoutC, - GemmCoord actualShape, GemmCoord actualNextShape, uint32_t nowIdx, uint32_t &nkblockIdx, - uint32_t &nkblockLoop, bool firstItr, bool endItr, bool initMmad, bool firstQItr, bool endQItr) - { - uint32_t mActual = actualShape.m(); - uint32_t kActual = actualShape.k(); - uint32_t nActual = actualShape.n(); - uint32_t mRound = RoundUp(mActual); - uint32_t kRound = RoundUp(kActual); - uint32_t nRound = RoundUp(nActual); - LayoutAInL1 layoutAInL1 = LayoutAInL1::template MakeLayout(mRound, kActual); // embed - LayoutAInL0 layoutAInL0 = LayoutAInL0::template MakeLayout(mRound, kActual); - LayoutBInL1 layoutBInL1 = LayoutBInL1::template MakeLayout(kActual, nActual); - LayoutBInL0 layoutBInL0 = LayoutBInL0::template MakeLayout(kActual, nRound); - uint32_t locPingPongFlag = nowIdx % 2; - uint32_t l1KvPingPongFlag = nkblockIdx % 2; - uint32_t l0ABPingPongFlag = nkblockIdx % 2; - if (nkblockIdx == 1) { - auto layoutBTile = layoutB.GetTileLayout(MakeCoord(kActual, nActual)); - AscendC::WaitFlag(l1KvPingPongFlag); - copyGmToL1B(l1BTensor[l1KvPingPongFlag], gB, layoutBInL1, layoutBTile); - AscendC::SetFlag(l1KvPingPongFlag); - } - if (nkblockIdx != nkblockLoop) { - uint32_t nNextActual = actualNextShape.n(); - uint32_t kNextActual = actualNextShape.k(); - LayoutBInL1 layoutBNextInL1 = LayoutBInL1::template MakeLayout(kNextActual, nNextActual); - auto layoutNextBTile = layoutB.GetTileLayout(MakeCoord(kNextActual, nNextActual)); - AscendC::WaitFlag(1 - l1KvPingPongFlag); - copyGmToL1B(l1BTensor[1 - l1KvPingPongFlag], gmNextBlockB, layoutBNextInL1, layoutNextBTile); - AscendC::SetFlag(1 - l1KvPingPongFlag); - } - if (firstQItr) { - AscendC::WaitFlag(0); - AscendC::WaitFlag(1); - copyL1ToL0A(l0ATensor[0], l1ATensor, layoutAInL0, layoutAInL1); - } - - AscendC::WaitFlag(l1KvPingPongFlag); - AscendC::WaitFlag(l0ABPingPongFlag + 2); - copyL1ToL0B(l0BTensor[l0ABPingPongFlag], l1BTensor[l1KvPingPongFlag], layoutBInL0, layoutBInL1); - AscendC::SetFlag(l1KvPingPongFlag); - - AscendC::SetFlag(l0ABPingPongFlag); - AscendC::WaitFlag(l0ABPingPongFlag); - uint8_t unitFlag = 0b00; - if constexpr (!ENABLE_UNIT_FLAG_) { - if (firstItr) { - AscendC::WaitFlag(0); - AscendC::WaitFlag(1); - } - } else { - unitFlag = 0b11; - } - tileMmad( - l0CTensor[locPingPongFlag * mRound * 128], l0ATensor[0], l0BTensor[l0ABPingPongFlag], - mRound, nActual, kActual, initMmad, unitFlag); - // AscendC::PipeBarrier(); - if (endQItr) { - AscendC::SetFlag(0); - AscendC::SetFlag(1); - } - AscendC::SetFlag(l0ABPingPongFlag + 2); - auto layoutInL0C = LayoutCInL0::MakeLayoutInL0C(MakeCoord(mActual, (uint32_t)256)); - if (endItr) { - if constexpr (!ENABLE_UNIT_FLAG_) { - AscendC::SetFlag(l1KvPingPongFlag); - AscendC::WaitFlag(l1KvPingPongFlag); - copyL0CToGm(gC, l0CTensor, layoutC, layoutInL0C); - AscendC::SetFlag(0); - AscendC::SetFlag(1); - } else { - copyL0CToGm(gC, l0CTensor, layoutC, layoutInL0C, unitFlag); - } - } - } - -protected: - /// Data members - AscendC::LocalTensor l1ATensor; - AscendC::LocalTensor l1BTensor[STAGES]; - AscendC::LocalTensor l0ATensor[STAGES]; - AscendC::LocalTensor l0BTensor[STAGES]; - AscendC::LocalTensor l0CTensor; - - TileMmad tileMmad; - CopyGmToL1A copyGmToL1A; - CopyGmToL1B copyGmToL1B; - CopyL1ToL0A copyL1ToL0A; - CopyL1ToL0B copyL1ToL0B; - CopyL0CToGm copyL0CToGm; -}; - -//////////////////////////////////////////////////////////////////// - -} // namespace Catlass::Gemm::Block - -#endif // CATLASS_GEMM_BLOCK_MMAD_QK_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_qk_normal_tla.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_qk_normal_tla.hpp deleted file mode 100644 index 540cd1aedd..0000000000 --- a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_qk_normal_tla.hpp +++ /dev/null @@ -1,347 +0,0 @@ -/** - * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of - * CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ - -#ifndef CATLASS_GEMM_BLOCK_MMAD_QK_TLA_HPP -#define CATLASS_GEMM_BLOCK_MMAD_QK_TLA_HPP - -#include "catlass/kda_catlass.hpp" -#include "catlass/arch/kda_resource.hpp" -#include "catlass/kda_coord.hpp" -#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" -#include "catlass/gemm/kda_gemm_helper.hpp" -#include "catlass/kda_gemm_coord.hpp" -#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" -#include "catlass/gemm/tile/kda_tile_mmad.hpp" -#include "tla/kda_tla_layout.hpp" -#include "tla/kda_tensor.hpp" - -//////////////////////////////////////////////////////////////////// - -namespace Catlass::Gemm::Block { -//////////////////////////////////////////////////////////////////// - -template < - class ArchTag_, - bool PAGED_CACHE_FLAG_, - bool ENABLE_UNIT_FLAG_, - class L1TileShape_, - class L0TileShape_, - class ElementA_, - class ElementB_, - class ElementC_, - class ElementBias_, - class TileCopy_, - class TileMmad_> -struct BlockMmadTla< - MmadFAIQK, - L1TileShape_, - L0TileShape_, - ElementA_, - ElementB_, - ElementC_, - ElementBias_, - TileCopy_, - TileMmad_> { -public: - // Type Aliases - using DispatchPolicy = MmadFAIQK; - using ArchTag = typename DispatchPolicy::ArchTag; - using L1TileShape = L1TileShape_; - using L0TileShape = L0TileShape_; - using ElementA = ElementA_; - using LayoutA = typename TileCopy_::LayoutA; - using ElementB = ElementB_; - using LayoutB = typename TileCopy_::LayoutB; - using ElementC = ElementC_; - using LayoutC = typename TileCopy_::LayoutC; - - using TileMmad = TileMmad_; - - using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; - using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; - - using ElementAccumulator = typename TileCopy_::ElementAccumulator; - - using LayoutTagL1A = typename TileCopy_::LayoutTagL1A; - using LayoutTagL1B = typename TileCopy_::LayoutTagL1B; - using LayoutTagL0A = typename TileCopy_::LayoutTagL0A; - using LayoutTagL0B = typename TileCopy_::LayoutTagL0B; - - using L1AAlignHelper = typename TileCopy_::L1AAlignHelper; - using L1BAlignHelper = typename TileCopy_::L1BAlignHelper; - - static_assert(tla::is_tuple::value && tla::is_static::value, - "L1TileShape must be tla::tuple and static!"); - static_assert(tla::is_tuple::value && tla::is_static::value, - "L0TileShape must be tla::tuple and static!"); - - static constexpr bool ENABLE_UNIT_FLAG = DispatchPolicy::ENABLE_UNIT_FLAG; - static constexpr uint32_t STAGES = DispatchPolicy::STAGES; - static constexpr uint32_t L1_TILE_M = tla::get<0>(L1TileShape{}); - static constexpr uint32_t L1_TILE_N = tla::get<1>(L1TileShape{}); - static constexpr uint32_t L1_TILE_K = tla::get<2>(L1TileShape{}); - static constexpr uint32_t L0_TILE_M = tla::get<0>(L0TileShape{}); - static constexpr uint32_t L0_TILE_N = tla::get<1>(L0TileShape{}); - static constexpr uint32_t L0_TILE_K = tla::get<2>(L0TileShape{}); - - // L1 tile size - static constexpr uint32_t L1A_TILE_SIZE = L1_TILE_M * L1_TILE_K * sizeof(ElementA); - static constexpr uint32_t L1B_TILE_SIZE = L1_TILE_N * L1_TILE_K * sizeof(ElementB); - // L0 tile size - static constexpr uint32_t L0A_TILE_SIZE = L0_TILE_M * L0_TILE_K * sizeof(ElementA); - static constexpr uint32_t L0B_TILE_SIZE = L0_TILE_K * L0_TILE_N * sizeof(ElementB); - static constexpr uint32_t L0C_TILE_SIZE = L0_TILE_M * L0_TILE_N * sizeof(ElementAccumulator); - static constexpr uint32_t L0A_PINGPONG_BUF_SIZE = ArchTag::L0A_SIZE / STAGES; - static constexpr uint32_t L0B_PINGPONG_BUF_SIZE = ArchTag::L0B_SIZE / STAGES; - static constexpr uint32_t L0C_PINGPONG_BUF_SIZE = ArchTag::L0C_SIZE / STAGES; - - // Check LayoutC - static_assert(tla::detail::isRowMajor::value || - ((std::is_same_v || std::is_same_v || - std::is_same_v) && - tla::detail::iszN::value), - "LayoutC only supports zN in half or bfloat16 or float, RowMajor in all dtype yet!"); - - // Check L1TileShape - static_assert(L1A_TILE_SIZE + L1B_TILE_SIZE * STAGES <= ArchTag::L1_SIZE, - "L1TileShape exceeding the L1 space!"); - - // Check L0TileShape - static_assert(L0A_TILE_SIZE * STAGES <= ArchTag::L0A_SIZE, "L0TileShape exceeding the L0A space!"); - static_assert(L0B_TILE_SIZE * STAGES <= ArchTag::L0B_SIZE, "L0TileShape exceeding the L0B space!"); - static_assert(L0C_TILE_SIZE <= ArchTag::L0C_SIZE, "L0TileShape exceeding the L0C space!"); - - static constexpr uint32_t BLOCK_SIZE = 16; - static constexpr uint32_t EMBED_SPLIT_SIZE = 128; - static constexpr uint32_t UNIT_BLOCK_STACK_NUM = 4; - static constexpr uint32_t KV_BASE_BLOCK = 512; - static constexpr uint32_t KV_SPLIT_SIZE = 128; - - CATLASS_DEVICE - BlockMmadTla(Arch::Resource &resource, uint32_t l1BufAddrStart = 0) - { - // Allocate L1 memory space - l1ATensor = resource.l1Buf.template GetBufferByByte(l1BufAddrStart); - for (uint32_t i = 0; i < STAGES; i++) { - l1BTensor[i] = - resource.l1Buf.template GetBufferByByte(l1BufAddrStart + L1A_TILE_SIZE + L1B_TILE_SIZE * i); - l0ATensor[i] = resource.l0ABuf.template GetBufferByByte(L0A_PINGPONG_BUF_SIZE * i); - l0BTensor[i] = resource.l0BBuf.template GetBufferByByte(L0B_PINGPONG_BUF_SIZE * i); - } - l0CTensor = resource.l0CBuf.template GetBufferByByte(0); - } - - CATLASS_DEVICE - ~BlockMmadTla() {} - - template - CATLASS_DEVICE void loadQGM(TensorA &tensorA) - { - using CopyGmToL1A = Gemm::Tile::TileCopyFAQTla; - CopyGmToL1A copyGmToL1A; - - uint32_t rowNum = tla::get<0>(tensorA.shape()) * tla::get<1>(tensorA.shape()); - uint32_t embed = tla::get<2>(tensorA.shape()); - auto L1A_LAYOUT = tla::MakeLayout(rowNum, embed); - auto tensorL1A = tla::MakeTensor(l1ATensor, L1A_LAYOUT, Arch::PositionL1{}); - copyGmToL1A(tensorL1A, tensorA); - AscendC::SetFlag(EVENT_ID3); - AscendC::WaitFlag(EVENT_ID3); - } - - CATLASS_DEVICE - void getBlockShape(GemmCoord &actualShape, uint32_t &nowNIdx, uint32_t &kIdx, - uint32_t &nLoop, uint32_t &kLoop, uint32_t &kvSeqlen, uint32_t &embed, bool firstBlock, uint32_t maskTailS = 0) - { - uint32_t nSplitSize = KV_SPLIT_SIZE; - uint32_t embedSplitSize = EMBED_SPLIT_SIZE; - if (nowNIdx == nLoop - 1) { - nSplitSize = kvSeqlen - nowNIdx * KV_SPLIT_SIZE; - } - if (firstBlock && maskTailS != 0) { - nSplitSize = nSplitSize - maskTailS; - } - // } - if (kIdx == kLoop - 1) { - embedSplitSize = embed - kIdx * EMBED_SPLIT_SIZE; - } - actualShape[1] = nSplitSize; - actualShape[2] = embedSplitSize; - } - - CATLASS_DEVICE - void getKVOffset(AscendC::GlobalTensor &gBlockTable, uint32_t &kOffset, uint32_t &nowNIdx, uint32_t &kIdx, - uint32_t &nLoop, uint32_t &kLoop, uint32_t &strideKV, uint32_t &blockSize, uint32_t maskTailS = 0) - { - if (nowNIdx >= nLoop || kIdx >= kLoop) { - kOffset = 0; - } - if constexpr (PAGED_CACHE_FLAG_) { - uint32_t blockTableId = gBlockTable.GetValue(nowNIdx); - kOffset = blockTableId * blockSize * strideKV + maskTailS * strideKV + kIdx * EMBED_SPLIT_SIZE; - } else { - kOffset = nowNIdx * KV_SPLIT_SIZE * strideKV + kIdx * EMBED_SPLIT_SIZE; - } - } - - template - CATLASS_DEVICE void operator()(TensorA &tensorA, TensorB &tensorB, AscendC::GlobalTensor gC, - AscendC::GlobalTensor gBlockTable, GemmCoord actualOriShape, - uint32_t &nIdx, uint32_t &nLoop, uint32_t &blockSize, uint32_t kvSeqlen, uint32_t strideKV) - { - uint32_t rowNum = actualOriShape[0]; - uint32_t embed = actualOriShape[2]; - uint32_t kLoop = CeilDiv(embed); - uint32_t nkBlockLoop = nLoop * kLoop; - GemmCoord actualShape{rowNum, 0, 0}; - GemmCoord actualNextShape{rowNum, 0, 0}; - uint32_t nkBlockNextIdx = nIdx * kLoop + 1; - uint32_t gBOffset = 0; - uint32_t gBNextOffset = 0; - uint32_t stackTile = 0; - for (uint32_t blockStackIdx = 0; (blockStackIdx < UNIT_BLOCK_STACK_NUM) && ((nIdx + blockStackIdx) < nLoop); - ++blockStackIdx) { - for (uint32_t kIdx = 0; kIdx < kLoop; kIdx++) { - uint32_t nowNIdx = nIdx + blockStackIdx; - uint32_t nLoopNextIdx = nkBlockNextIdx / kLoop; - uint32_t kLoopNextIdx = nkBlockNextIdx % kLoop; - uint32_t gCOffset = blockStackIdx / 2 * 2 * KV_SPLIT_SIZE; - getBlockShape(actualShape, nowNIdx, kIdx, nLoop, kLoop, kvSeqlen, embed, nowNIdx == nIdx); - getBlockShape(actualNextShape, nLoopNextIdx, kLoopNextIdx, nLoop, kLoop, kvSeqlen, embed, nLoopNextIdx == nIdx); - getKVOffset(gBlockTable, gBOffset, nowNIdx, kIdx, nLoop, kLoop, strideKV, blockSize); - getKVOffset(gBlockTable, gBNextOffset, nLoopNextIdx, kLoopNextIdx, nLoop, kLoop, strideKV, blockSize); - bool firstItr = ((blockStackIdx % 2) == 0) && (kIdx == 0); - bool endItr = (((blockStackIdx % 2) == 1) || (nowNIdx == nLoop - 1)) && (kIdx == kLoop - 1); - bool firstQtr = blockStackIdx == 0; - bool endQItr = ((nowNIdx == nLoop - 1) || (blockStackIdx == UNIT_BLOCK_STACK_NUM - 1)) && (kIdx == kLoop - 1); - int cc = 1; - bool initMmad = kIdx == 0; - stackTile += actualShape[1]; - auto layoutS = tla::MakeLayout(tla::MakeShape(rowNum, stackTile), tla::MakeStride(Int<512>{}, Int<1>{})); - auto tensorCOffset = tla::MakeTensor(gC[gCOffset], layoutS, Arch::PositionGM{}); - computeQK(tensorA, tensorB, tensorCOffset, gBOffset, gBNextOffset, - actualShape, actualNextShape, blockStackIdx, nkBlockNextIdx, nkBlockLoop, strideKV, firstItr, endItr, initMmad, firstQtr, endQItr); - ++nkBlockNextIdx; - if (endItr) { - stackTile = 0; - } - } - } - } - - template - CATLASS_DEVICE void computeQK( - TensorA &tensorA, TensorB &tensorB, TensorC &tensorC, uint32_t gBOffset, uint32_t gBNextOffset, - GemmCoord actualShape, GemmCoord actualNextShape, uint32_t nowIdx, uint32_t &nkblockIdx, - uint32_t &nkblockLoop, uint32_t strideKV, bool firstItr, bool endItr, bool initMmad, bool firstQItr, bool endQItr) - { - using CopyGmToL1B = typename TileCopy_::template CopyGmToL1B; - using CopyL0CToGm = typename TileCopy_::template CopyL0CToGm; - CopyGmToL1B copyGmToL1B; - CopyL0CToGm copyL0CToGm; - - uint32_t mActual = actualShape.m(); - uint32_t kActual = actualShape.k(); - uint32_t nActual = actualShape.n(); - uint32_t mRound = RoundUp(mActual); - uint32_t kRound = RoundUp(kActual); - uint32_t nRound = RoundUp(nActual); - auto layoutAInL1 = tla::MakeLayout(mRound, kActual); - auto layoutAInL0 = tla::MakeLayout(mRound, kActual); - auto layoutBInL1 = tla::MakeLayout(kActual, nActual); - auto layoutBInL0 = tla::MakeLayout(kActual, nRound); - uint32_t locPingPongFlag = nowIdx % 2; - uint32_t l1KvPingPongFlag = nkblockIdx % 2; - uint32_t l0ABPingPongFlag = nkblockIdx % 2; - auto tensorL1A = tla::MakeTensor(l1ATensor, layoutAInL1, Arch::PositionL1{}); - auto tensorL0A = tla::MakeTensor(l0ATensor[0], layoutAInL0, Arch::PositionL0A{}); - auto tensorTileB = GetTile(tensorB, tla::MakeCoord(gBOffset % strideKV, gBOffset / strideKV), tla::MakeShape(kActual, nActual)); - auto tensorL1B = tla::MakeTensor(l1BTensor[l1KvPingPongFlag], layoutBInL1, Arch::PositionL1{}); - auto tensorL0B = tla::MakeTensor(l0BTensor[l0ABPingPongFlag], layoutBInL0, Arch::PositionL0B{}); - if (nkblockIdx == 1) { - AscendC::WaitFlag(l1KvPingPongFlag); - copyGmToL1B(tensorL1B, tensorTileB); - AscendC::SetFlag(l1KvPingPongFlag); - } - if (nkblockIdx != nkblockLoop) { - uint32_t nNextActual = actualNextShape.n(); - uint32_t kNextActual = actualNextShape.k(); - auto tensorTileBNext = GetTile(tensorB, tla::MakeCoord(gBNextOffset % strideKV, gBNextOffset / strideKV), tla::MakeShape(kNextActual, nNextActual)); - auto layoutBNextInL1 = tla::MakeLayout(kNextActual, nNextActual); - auto tensorL1BNext = tla::MakeTensor(l1BTensor[1 - l1KvPingPongFlag], layoutBNextInL1, Arch::PositionL1{}); - AscendC::WaitFlag(1 - l1KvPingPongFlag); - copyGmToL1B(tensorL1BNext, tensorTileBNext); - AscendC::SetFlag(1 - l1KvPingPongFlag); - } - if (firstQItr) { - AscendC::WaitFlag(0); - AscendC::WaitFlag(1); - copyL1ToL0A(tensorL0A, tensorL1A); - } - - AscendC::WaitFlag(l1KvPingPongFlag); - AscendC::WaitFlag(l0ABPingPongFlag + 2); - copyL1ToL0B(tensorL0B, tensorL1B); - AscendC::SetFlag(l1KvPingPongFlag); - - AscendC::SetFlag(l0ABPingPongFlag); - AscendC::WaitFlag(l0ABPingPongFlag); - uint8_t unitFlag = 0b00; - if constexpr (!ENABLE_UNIT_FLAG_) { - if (firstItr) { - AscendC::WaitFlag(0); - AscendC::WaitFlag(1); - } - } else { - unitFlag = 0b11; - } - auto layoutCTileInL0 = tla::MakeLayoutL0C(mRound, nActual); - auto tensorL0CTile = tla::MakeTensor(l0CTensor[locPingPongFlag * mRound * 128], layoutCTileInL0, Arch::PositionL0C{}); - tileMmad( - tensorL0CTile, tensorL0A, tensorL0B, - mRound, nActual, kActual, initMmad, unitFlag); - if (endQItr) { - AscendC::SetFlag(0); - AscendC::SetFlag(1); - } - AscendC::SetFlag(l0ABPingPongFlag + 2); - auto layoutCInL0 = tla::MakeLayoutL0C(mActual, (uint32_t)256); - auto tensorL0C = tla::MakeTensor(l0CTensor, layoutCInL0, Arch::PositionL0C{}); - if (endItr) { - if constexpr (!ENABLE_UNIT_FLAG_) { - AscendC::SetFlag(l1KvPingPongFlag); - AscendC::WaitFlag(l1KvPingPongFlag); - copyL0CToGm(tensorC, tensorL0C); - AscendC::SetFlag(0); - AscendC::SetFlag(1); - } else { - copyL0CToGm(tensorC, tensorL0C, unitFlag); - } - } - } - -protected: - /// Data members - AscendC::LocalTensor l1ATensor; - AscendC::LocalTensor l1BTensor[STAGES]; - AscendC::LocalTensor l0ATensor[STAGES]; - AscendC::LocalTensor l0BTensor[STAGES]; - AscendC::LocalTensor l0CTensor; - - TileMmad tileMmad; - CopyL1ToL0A copyL1ToL0A; - CopyL1ToL0B copyL1ToL0B; -}; - -//////////////////////////////////////////////////////////////////// - -} // namespace Catlass::Gemm::Block - -#endif // CATLASS_GEMM_BLOCK_MMAD_QK_TLA_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_qk_tla.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_qk_tla.hpp deleted file mode 100644 index 9903b37153..0000000000 --- a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_fai_qk_tla.hpp +++ /dev/null @@ -1,349 +0,0 @@ -/** - * Copyright (c) 2026 Huawei Technologies Co., Ltd. - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of - * CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ -#ifndef CATLASS_GEMM_BLOCK_BLOCK_MMAD_FAI_QK_TLA_HPP -#define CATLASS_GEMM_BLOCK_BLOCK_MMAD_FAI_QK_TLA_HPP - -#include "catlass/kda_catlass.hpp" -#include "catlass/arch/kda_resource.hpp" -#include "catlass/kda_coord.hpp" -#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" -#include "catlass/gemm/kda_gemm_helper.hpp" -#include "catlass/kda_gemm_coord.hpp" -#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" -#include "catlass/gemm/tile/kda_tile_mmad.hpp" -#include "tla/kda_tensor.hpp" -#include "tla/kda_tla_layout.hpp" - -//////////////////////////////////////////////////////////////////// - -namespace Catlass::Gemm::Block { -//////////////////////////////////////////////////////////////////// - -template < - bool PAGED_CACHE_FLAG_, - bool ENABLE_UNIT_FLAG_, - class L1TileShape_, - class L0TileShape_, - class ElementA_, - class ElementB_, - class ElementC_, - class ElementBias_, - class TileCopy_, - class TileMmad_> -struct BlockMmadTla< - MmadFAIQK, - L1TileShape_, - L0TileShape_, - ElementA_, - ElementB_, - ElementC_, - ElementBias_, - TileCopy_, - TileMmad_> { -public: - using DispatchPolicy = MmadFAIQK; - using ArchTag = typename DispatchPolicy::ArchTag; - using L1TileShape = L1TileShape_; - using L0TileShape = L0TileShape_; - using TileCopy = TileCopy_; - using TileMmad = TileMmad_; - - using ElementA = ElementA_; - using ElementB = ElementB_; - using ElementC = ElementC_; - using ElementAccumulator = typename TileCopy::ElementAccumulator; - - using CopyL1ToL0A = typename TileCopy::CopyL1ToL0A; - using CopyL1ToL0B = typename TileCopy::CopyL1ToL0B; - - using LayoutTagA = typename TileCopy::LayoutTagA; - using LayoutTagB = typename TileCopy::LayoutTagB; - using LayoutTagC = typename TileCopy::LayoutTagC; - - using LayoutTagL1A = typename TileCopy::LayoutTagL1A; - using LayoutTagL1B = typename TileCopy::LayoutTagL1B; - using LayoutTagL0A = typename TileCopy::LayoutTagL0A; - using LayoutTagL0B = typename TileCopy::LayoutTagL0B; - - static constexpr uint32_t STAGES = DispatchPolicy::STAGES; - // L1 tile shape - static constexpr uint32_t L1_TILE_M = tla::get<0>(L1TileShape{}); // s1 - static constexpr uint32_t L1_TILE_N = tla::get<1>(L1TileShape{}); // s2 - static constexpr uint32_t L1_TILE_K = tla::get<2>(L1TileShape{}); // d - // L0 tile shape - static constexpr uint32_t L0_TILE_M = tla::get<0>(L0TileShape{}); // s1 - static constexpr uint32_t L0_TILE_N = tla::get<1>(L0TileShape{}); // s2 - static constexpr uint32_t L0_TILE_K = tla::get<2>(L0TileShape{}); // d - // L1 tile size - static constexpr uint32_t L1A_TILE_SIZE = L1_TILE_M * L1_TILE_K * sizeof(ElementA); - static constexpr uint32_t L1B_TILE_SIZE = L1_TILE_N * L1_TILE_K * sizeof(ElementB); - // L0 tile size - static constexpr uint32_t L0A_TILE_SIZE = L0_TILE_M * L0_TILE_K * sizeof(ElementA); - static constexpr uint32_t L0B_TILE_SIZE = L0_TILE_K * L0_TILE_N * sizeof(ElementB); - static constexpr uint32_t L0C_TILE_SIZE = L0_TILE_M * L0_TILE_N * sizeof(ElementAccumulator); - - static constexpr uint32_t BLOCK_L1_SIZE = (L1A_TILE_SIZE + L1B_TILE_SIZE) * STAGES; - static constexpr uint32_t BLOCK_L0C_SIZE = L0C_TILE_SIZE * STAGES; - - // Check L1/L0TileShape - static_assert(L1_TILE_M == L0_TILE_M && L1_TILE_N == L0_TILE_N && L1_TILE_K == L0_TILE_K, - "The situation where the basic blocks of L1 and L0 differ on the m, n, k axes is not supported yet"); - static_assert(L0A_TILE_SIZE * STAGES <= ArchTag::L0A_SIZE, "L0TileShape exceeding the L0A space!"); - static_assert(L0B_TILE_SIZE * STAGES <= ArchTag::L0B_SIZE, "L0TileShape exceeding the L0B space!"); - -public: - /// Construct - CATLASS_DEVICE - BlockMmadTla(Arch::Resource &resource, uint32_t &l1BufAddrStart, uint32_t &l0CBufAddrStart) - { - uint32_t l1AOffset = l1BufAddrStart; - uint32_t l1BOffset = l1AOffset + L1A_TILE_SIZE * STAGES; - uint32_t l0COffset = l0CBufAddrStart; - - for (uint32_t i = 0; i < STAGES; i++) { - // Assign L1/L0A/L0B space for each stages - l1ATensorList_[i] = resource.l1Buf.template GetBufferByByte(l1AOffset + L1A_TILE_SIZE * i); - l1BTensorList_[i] = resource.l1Buf.template GetBufferByByte(l1BOffset + L1B_TILE_SIZE * i); - l0ATensorList_[i] = resource.l0ABuf.template GetBufferByByte(L0A_TILE_SIZE * i); - l0BTensorList_[i] = resource.l0BBuf.template GetBufferByByte(L0B_TILE_SIZE * i); - l0CTensorList_[i] = resource.l0CBuf.template GetBufferByByte(l0COffset + L0C_TILE_SIZE * i); - - // Assign event ID for each stages - l1AEventList_[i] = BLOCK_EVENT_ID + i; - l1BEventList_[i] = BLOCK_EVENT_ID + i + STAGES; - l0AEventList_[i] = i; - l0BEventList_[i] = i + STAGES; - l0CEventList_[i] = BLOCK_EVENT_ID + i; - - // The event id that needs to be set before the loop - AscendC::SetFlag(l1AEventList_[i]); - AscendC::SetFlag(l1BEventList_[i]); - AscendC::SetFlag(l0AEventList_[i]); - AscendC::SetFlag(l0BEventList_[i]); - AscendC::SetFlag(l0CEventList_[i]); - } - - l1BufAddrStart += BLOCK_L1_SIZE; - l0CBufAddrStart += BLOCK_L0C_SIZE; - } - - /// Destructor - CATLASS_DEVICE - ~BlockMmadTla() - { - for (uint32_t i = 0; i < STAGES; i++) { - AscendC::WaitFlag(l1AEventList_[i]); - AscendC::WaitFlag(l1BEventList_[i]); - AscendC::WaitFlag(l0AEventList_[i]); - AscendC::WaitFlag(l0BEventList_[i]); - AscendC::WaitFlag(l0CEventList_[i]); - } - } - -public: - template - CATLASS_DEVICE void operator()( - TensorA &tensorA, - TensorB &tensorB, - TensorC &tensorC, - AscendC::GlobalTensor &blockTable, - Shape &actualShape, - int64_t taskId, int32_t blockSize, - bool isFirstLoop = false, bool isLastLoop = false) - { - int32_t blockM = tla::get<0>(actualShape); - int32_t blockN = tla::get<1>(actualShape); - int32_t blockK = tla::get<2>(actualShape); - - AscendC::WaitFlag(l0CEventList_[l0CListId_]); - auto layoutInL0C = tla::MakeLayoutL0C(blockM, blockN); - auto tensorL0C = tla::MakeTensor(l0CTensorList_[l0CListId_], layoutInL0C, Arch::PositionL0C{}); - - if (unlikely(isFirstLoop)) { - AscendC::WaitFlag(l1AEventList_[l1AListId_]); - } - - int32_t kLoops = (blockK + L0_TILE_K - 1) / L0_TILE_K; - for (int32_t kIdx = 0; kIdx < kLoops; ++kIdx) { - int32_t tileK; - if (kIdx == kLoops - 1) { - int32_t tailSize = blockK % L0_TILE_K; - tileK = tailSize ? tailSize : L0_TILE_K; - } else { - tileK = L0_TILE_K; - } - - // load matrix A tile from GM to L1 - auto layoutAInL1 = tla::MakeLayout(blockM, blockK); - auto tensorL1A = tla::MakeTensor(l1ATensorList_[l1AListId_], layoutAInL1, Arch::PositionL1{}); - if (unlikely(isFirstLoop)) { - CopyInL1A(tensorL1A, tensorA, blockM, tileK, kIdx * L0_TILE_K); - AscendC::SetFlag(l1AEventList_[l1AListId_]); - AscendC::WaitFlag(l1AEventList_[l1AListId_]); - } - - // load matrix B tile from GM to L1 - auto layoutBInL1 = tla::MakeLayout(tileK, blockN); - auto tensorL1B = tla::MakeTensor(l1BTensorList_[l1BListId_], layoutBInL1, Arch::PositionL1{}); - AscendC::WaitFlag(l1BEventList_[l1BListId_]); - CopyInL1B(tensorL1B, tensorB, blockTable, tileK, blockN, kIdx * L0_TILE_K, blockSize); - AscendC::SetFlag(l1BEventList_[l1BListId_]); - AscendC::WaitFlag(l1BEventList_[l1BListId_]); - - // load matrix A tile from L1 to L0A - auto layoutAInL0 = tla::MakeLayout(blockM, tileK); - auto tensorL0A = tla::MakeTensor(l0ATensorList_[l0ListId_], layoutAInL0, Arch::PositionL0A{}); - auto tensorL1TileA = GetTile( - tensorL1A, - tla::MakeCoord(0, kIdx * L0_TILE_K), - tla::MakeShape(blockM, tileK)); - AscendC::WaitFlag(l0AEventList_[l0ListId_]); - copyL1ToL0A(tensorL0A, tensorL1TileA); - - // load matrix B tile from L1 to L0B - auto layoutBInL0 = tla::MakeLayout(tileK, blockN); - auto tensorL0B = tla::MakeTensor(l0BTensorList_[l0ListId_], layoutBInL0, Arch::PositionL0B{}); - auto tensorL1TileB = GetTile( - tensorL1B, - tla::MakeCoord(0, 0), - tla::MakeShape(tileK, blockN)); - AscendC::WaitFlag(l0BEventList_[l0ListId_]); - copyL1ToL0B(tensorL0B, tensorL1TileB); - - AscendC::SetFlag(l0CEventList_[l0CListId_]); - AscendC::WaitFlag(l0CEventList_[l0CListId_]); - - AscendC::SetFlag(l1BEventList_[l1BListId_]); - l1BListId_ = (l1BListId_ + 1 < STAGES) ? (l1BListId_ + 1) : 0; - - // Mmad - bool initC = kIdx == 0; - uint32_t tileM = (blockM == 1 ? M_ALIGN : blockM); - uint32_t tileN = blockN; - tileMmad(tensorL0C, tensorL0A, tensorL0B, tileM, tileN, tileK, initC); - - AscendC::SetFlag(l0AEventList_[l0ListId_]); - AscendC::SetFlag(l0BEventList_[l0ListId_]); - l0ListId_ = (l0ListId_ + 1 < STAGES) ? (l0ListId_ + 1) : 0; - } - - if (unlikely(isLastLoop)) { - AscendC::SetFlag(l1AEventList_[l1AListId_]); - l1AListId_ = (l1AListId_ + 1 < STAGES) ? (l1AListId_ + 1) : 0; - } - - AscendC::SetFlag(l0CEventList_[l0CListId_]); - AscendC::WaitFlag(l0CEventList_[l0CListId_]); - - AscendC::CrossCoreWaitFlag(MM1_RES_INTRA_EVENT[taskId]); - AscendC::CrossCoreWaitFlag(16 + MM1_RES_INTRA_EVENT[taskId]); - - // copy block out - using CopyL0CToDst = typename TileCopy::template CopyL0CToDst; - CopyL0CToDst copyL0CToDst; - copyL0CToDst(tensorC, tensorL0C); - - AscendC::SetFlag(l0CEventList_[l0CListId_]); - l0CListId_ = (l0CListId_ + 1 < STAGES) ? (l0CListId_ + 1) : 0; - } - -private: - template - CATLASS_DEVICE void CopyInL1A(TensorL1A &tensorL1A, TensorA &tensorA, int32_t tileM, int32_t tileK, int32_t kOffset) - { - using CopyGmToL1A = typename TileCopy::template CopyGmToL1A; - CopyGmToL1A copyGmToL1A; - - auto tensorTileA = GetTile( - tensorA, - tla::MakeCoord(0, kOffset), - tla::MakeShape(tileM, tileK)); - - auto tensorL1TileA = GetTile( - tensorL1A, - tla::MakeCoord(0, kOffset), - tla::MakeShape(tileM, tileK)); - - copyGmToL1A(tensorL1TileA, tensorTileA); - } - - template - CATLASS_DEVICE void CopyInL1B( - TensorL1B &tensorL1B, - TensorB &tensorB, - AscendC::GlobalTensor &blockTable, - int32_t tileK, int32_t tileN, - int32_t kOffset, int32_t blockSize) - { - using CopyGmToL1B = typename TileCopy::template CopyGmToL1B; - CopyGmToL1B copyGmToL1B; - - if constexpr (PAGED_CACHE_FLAG_) { - int32_t copyColCnt = 0; - int32_t blockLoops = (tileN + blockSize - 1) / blockSize; - - for (int32_t blockIdx = 0; blockIdx < blockLoops; ++blockIdx) { - int32_t curCopyCol = blockIdx == blockLoops - 1 ? tileN - copyColCnt : blockSize; - int32_t idInBlockTable = blockTable.GetValue(blockIdx); - - auto tensorL1TileB = GetTile( - tensorL1B, - tla::MakeCoord(0, copyColCnt), - tla::MakeShape(tileK, curCopyCol)); - - auto tensorTileB = GetTile( - tensorB, - tla::MakeCoord(kOffset, idInBlockTable * blockSize), - tla::MakeShape(tileK, curCopyCol)); - - copyGmToL1B(tensorL1TileB, tensorTileB); - copyColCnt += curCopyCol; - } - } else { - auto tensorTileB = GetTile( - tensorB, - tla::MakeCoord(kOffset, 0), - tla::MakeShape(tileK, tileN)); - - copyGmToL1B(tensorL1B, tensorTileB); - } - } - -private: - static constexpr uint16_t M_ALIGN = 16; - static constexpr uint32_t SYNC_MODE = 4; - static constexpr uint32_t MM1_RES_INTRA_EVENT[2] = {9, 10}; - static constexpr uint32_t BLOCK_EVENT_ID = 0; - - // Multi-stage tensors list - AscendC::LocalTensor l1ATensorList_[STAGES]; - AscendC::LocalTensor l1BTensorList_[STAGES]; - AscendC::LocalTensor l0ATensorList_[STAGES]; - AscendC::LocalTensor l0BTensorList_[STAGES]; - AscendC::LocalTensor l0CTensorList_[STAGES]; - - // Multi-stage event id list - int32_t l1AEventList_[STAGES]; - int32_t l1BEventList_[STAGES]; - int32_t l0AEventList_[STAGES]; - int32_t l0BEventList_[STAGES]; - int32_t l0CEventList_[STAGES]; - - uint32_t l1AListId_{0}; - uint32_t l1BListId_{0}; - uint32_t l0ListId_{0}; - uint32_t l0CListId_{0}; - - CopyL1ToL0A copyL1ToL0A; - CopyL1ToL0B copyL1ToL0B; - TileMmad tileMmad; -}; -} // namespace Catlass::Gemm::Block -#endif diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_pingpong_dequant_tla.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_pingpong_dequant_tla.hpp deleted file mode 100644 index d286e68a0d..0000000000 --- a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_pingpong_dequant_tla.hpp +++ /dev/null @@ -1,666 +0,0 @@ -/** - * Copyright (c) 2026 Huawei Technologies Co., Ltd. - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of - * CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ - -#ifndef CATLASS_GEMM_BLOCK_BLOCK_MMAD_PINGPONG_DEQUANT_TLA_HPP -#define CATLASS_GEMM_BLOCK_BLOCK_MMAD_PINGPONG_DEQUANT_TLA_HPP - -#include "catlass/kda_catlass.hpp" -#include "catlass/arch/kda_resource.hpp" -#include "catlass/kda_coord.hpp" -#include "catlass/kda_gemm_coord.hpp" -#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" -#include "catlass/gemm/kda_gemm_helper.hpp" -#include "tla/kda_tla_layout.hpp" -#include "tla/kda_tensor.hpp" - -namespace Catlass::Gemm::Block { - -template < - class ArchTag_, - bool ENABLE_UNIT_FLAG_, - bool USE_HF32_MODE_, - uint32_t L0C_STAGES_, - bool ENABLE_L1_RESIDENT_, - uint32_t L1A_STAGES_, - uint32_t L1B_STAGES_, - uint32_t L0A_STAGES_, - uint32_t L0B_STAGES_, - class L1TileShape_, - class L0TileShape_, - class ElementA_, - class ElementB_, - class ElementC_, - class ElementBias_, - class TileCopy_, - class TileMmad_> -struct BlockMmadTla< - MmadDequant, - L1TileShape_, - L0TileShape_, - ElementA_, - ElementB_, - ElementC_, - ElementBias_, - TileCopy_, - TileMmad_> { -public: - // Type Aliases - using DispatchPolicy = MmadDequant; - using ArchTag = typename DispatchPolicy::ArchTag; - using TileCopy = TileCopy_; - using L1TileShape = L1TileShape_; - using L0TileShape = L0TileShape_; - using ElementA = ElementA_; - using LayoutA = typename TileCopy::LayoutA; - using ElementB = ElementB_; - using LayoutB = typename TileCopy::LayoutB; - using ElementC = ElementC_; - using LayoutC = typename TileCopy::LayoutC; - using ElementBias = ElementBias_; - - using TileMmad = TileMmad_; - - using CopyL1ToL0A = typename TileCopy::CopyL1ToL0A; - using CopyL1ToL0B = typename TileCopy::CopyL1ToL0B; - using CopyL1ToBT = typename TileCopy::CopyL1ToBT; - - using ElementAccumulator = typename TileCopy::ElementAccumulator; - - static constexpr bool HAS_BIAS = TileCopy::HAS_BIAS; - static constexpr Tile::ScaleGranularity DEQUANT_MODE = TileCopy::DEQUANT_GRANULARITY; - static constexpr bool HAS_QUANT_TENSOR = TileCopy::HAS_QUANT_TENSOR; - - using LayoutTagL1A = typename TileCopy::LayoutTagL1A; - using LayoutTagL1B = typename TileCopy::LayoutTagL1B; - using LayoutTagL0A = typename TileCopy::LayoutTagL0A; - using LayoutTagL0B = typename TileCopy::LayoutTagL0B; - - static_assert(tla::is_tuple::value && tla::is_static::value, - "L1TileShape must be tla::tuple and static!"); - static_assert(tla::is_tuple::value && tla::is_static::value, - "L0TileShape must be tla::tuple and static!"); - - static constexpr bool ENABLE_UNIT_FLAG = DispatchPolicy::ENABLE_UNIT_FLAG; - static constexpr bool USE_HF32_MODE = DispatchPolicy::USE_HF32_MODE; - static constexpr bool ENABLE_L1_RESIDENT = DispatchPolicy::ENABLE_L1_RESIDENT; - static constexpr uint32_t L1A_STAGES = DispatchPolicy::L1A_STAGES; - static constexpr uint32_t L1B_STAGES = DispatchPolicy::L1B_STAGES; - static constexpr uint32_t L0A_STAGES = DispatchPolicy::L0A_STAGES; - static constexpr uint32_t L0B_STAGES = DispatchPolicy::L0B_STAGES; - static constexpr uint32_t L0C_STAGES = DispatchPolicy::L0C_STAGES; - static constexpr uint32_t L1_TILE_M = tla::get<0>(L1TileShape{}); - static constexpr uint32_t L1_TILE_N = tla::get<1>(L1TileShape{}); - static constexpr uint32_t L1_TILE_K = tla::get<2>(L1TileShape{}); - static constexpr uint32_t L0_TILE_M = tla::get<0>(L0TileShape{}); - static constexpr uint32_t L0_TILE_N = tla::get<1>(L0TileShape{}); - static constexpr uint32_t L0_TILE_K = tla::get<2>(L0TileShape{}); - - // L1 tile size - static constexpr uint32_t L1A_TILE_SIZE = L1_TILE_M * L1_TILE_K * sizeof(ElementA); - static constexpr uint32_t L1B_TILE_SIZE = L1_TILE_N * L1_TILE_K * sizeof(ElementB); - static constexpr uint32_t L1AB_TILE_SIZE = L1A_TILE_SIZE * L1A_STAGES + L1B_TILE_SIZE * L1B_STAGES; - // L0 tile size - static constexpr uint32_t L0A_TILE_SIZE = L0_TILE_M * L0_TILE_K * sizeof(ElementA); - static constexpr uint32_t L0B_TILE_SIZE = L0_TILE_K * L0_TILE_N * sizeof(ElementB); - static constexpr uint32_t L0C_TILE_SIZE = L1_TILE_M * L1_TILE_N * sizeof(ElementAccumulator); - static constexpr uint32_t BIAS_BUF_SIZE = L0_TILE_N * sizeof(ElementAccumulator); - - // Check HF32_MODE - static_assert( - !USE_HF32_MODE || (USE_HF32_MODE && std::is_same_v && std::is_same_v), - "HF32 MODE only supports in float!"); - - // Check L0C_STAGES - static_assert(!(ENABLE_UNIT_FLAG && L0C_STAGES != 1), "L0C_STAGES must be 1 when UnitFlag is true!"); - - // Check LayoutC - static_assert(tla::detail::isRowMajor::value || - ((std::is_same_v || std::is_same_v || - std::is_same_v) && - tla::detail::iszN::value), - "LayoutC only supports zN in half or bfloat16 or float, RowMajor in all dtype yet!"); - - // Check L1TileShape - static_assert(L1AB_TILE_SIZE <= ArchTag::L1_SIZE, "L1TileShape exceeding the L1 space!"); - - // Check L0TileShape - static_assert(L0A_TILE_SIZE * L0A_STAGES <= ArchTag::L0A_SIZE, "L0TileShape exceeding the L0A space!"); - static_assert(L0B_TILE_SIZE * L0B_STAGES <= ArchTag::L0B_SIZE, "L0TileShape exceeding the L0B space!"); - static_assert(L0C_TILE_SIZE * L0C_STAGES <= ArchTag::L0C_SIZE, "L0TileShape exceeding the L0C space!"); - static_assert(BIAS_BUF_SIZE <= ArchTag::BIAS_SIZE, "BIAS_BUF_SIZE exceeding the BT space! Reduce L0_TILE_N"); - - static_assert(L1_TILE_M == L0_TILE_M && L1_TILE_N == L0_TILE_N, - "The situation where the basic blocks of L1 and L0 differ on the m and n axes is not supported yet"); - static_assert(L0_TILE_K <= L1_TILE_K, "L0TileShape::K cannot exceed L1TileShape::K"); - - static_assert((!HAS_BIAS && !HAS_QUANT_TENSOR && (L1A_STAGES + L1B_STAGES) <= 8) || - (HAS_BIAS && !HAS_QUANT_TENSOR && (L1A_STAGES + L1B_STAGES) <= 7) || - (HAS_QUANT_TENSOR && (L1A_STAGES + L1B_STAGES) <= 6), - "L1 Buffer overflow: Exceeds the supported range of EVENT(0~7)"); - - static_assert((!HAS_BIAS && (L0A_STAGES + L0B_STAGES) <= 8) || (HAS_BIAS && (L0A_STAGES + L0B_STAGES) <= 7), - "L0 Buffer overflow: Exceeds the supported range of EVENT_ID(0~7)"); - - static constexpr auto L1A_LAYOUT = - tla::MakeLayout(tla::Int{}, tla::Int{}); - static constexpr auto L1B_LAYOUT = - tla::MakeLayout(tla::Int{}, tla::Int{}); - static constexpr auto L1BIAS_LAYOUT = tla::MakeLayout(tla::Int{}); - static constexpr auto L0BIAS_LAYOUT = tla::MakeLayout(tla::Int{}); - static constexpr auto L1QUANT_LAYOUT = tla::MakeLayout(tla::Int{}); - - // When enableing L1 resident mode, restore the pointer and coordinates that record the last state - // to the initial state. if two blockmmad instances need to be consecutively invoked at the kernel layer, - // RestoreStatus() must be inserted between them. - CATLASS_DEVICE - void RestoreStatus() - { - for (int i = 0; i < L1A_STAGES; ++i) { - lastAddrA[i] = nullptr; - lastCoordA[i] = MatrixCoord{0U, 0U}; - } - for (int i = 0; i < L1B_STAGES; ++i) { - lastAddrB[i] = nullptr; - lastCoordB[i] = MatrixCoord{0U, 0U}; - } - } - - /// Construct - CATLASS_DEVICE - BlockMmadTla(Arch::Resource &resource, uint32_t l1BufAddrStart = 0) - { - if ASCEND_IS_AIC { - // use HF32 when USE_HF32_MODE is true - if constexpr (USE_HF32_MODE) { - AscendC::SetHF32Mode(true); - } else { - AscendC::SetHF32Mode(false); - } - if constexpr (ENABLE_UNIT_FLAG && tla::detail::isRowMajor::value) { - AscendC::SetMMLayoutTransform(true); - } - uint32_t l1AOffset = l1BufAddrStart; - uint32_t l1BOffset = l1BufAddrStart + L1A_TILE_SIZE * L1A_STAGES; - uint32_t l1BiasOffset = l1BOffset + L1B_TILE_SIZE * L1B_STAGES; - uint32_t l1QuantOffset = l1BiasOffset; - if constexpr (HAS_BIAS) { - l1QuantOffset += L1_TILE_N * sizeof(ElementBias); - } - // Init buffers - for (uint32_t i = 0; i < L1A_STAGES; i++) { - // Assign L1/L0A/L0B space for each stages - l1ATensorList[i] = resource.l1Buf.template GetBufferByByte(l1AOffset + L1A_TILE_SIZE * i); - // Assign event ID for each stages - l1AEventList[i] = i; - // The event id that needs to be set before the loop - AscendC::SetFlag(l1AEventList[i]); - } - for (uint32_t i = 0; i < L1B_STAGES; i++) { - // Assign L1/L0A/L0B space for each stages - l1BTensorList[i] = resource.l1Buf.template GetBufferByByte(l1BOffset + L1B_TILE_SIZE * i); - // Assign event ID for each stages - l1BEventList[i] = i + L1A_STAGES; - // The event id that needs to be set before the loop - AscendC::SetFlag(l1BEventList[i]); - } - for (uint32_t i = 0; i < L0A_STAGES; i++) { - // Assign L1/L0A/L0B space for each stages - l0ATensorList[i] = resource.l0ABuf.template GetBufferByByte(L0A_TILE_SIZE * i); - // Assign event ID for each stages - l0AEventList[i] = i; - // The event id that needs to be set before the loop - AscendC::SetFlag(l0AEventList[i]); - } - for (uint32_t i = 0; i < L0B_STAGES; i++) { - // Assign L1/L0A/L0B space for each stages - l0BTensorList[i] = resource.l0BBuf.template GetBufferByByte(L0B_TILE_SIZE * i); - // Assign event ID for each stages - l0BEventList[i] = i + L0A_STAGES; - // The event id that needs to be set before the loop - AscendC::SetFlag(l0BEventList[i]); - } - if constexpr (!ENABLE_UNIT_FLAG) { - for (uint32_t i = 0; i < L0C_STAGES; i++) { - l0CTensorList[i] = resource.l0CBuf.template GetBufferByByte(L0C_TILE_SIZE * i); - l0CEventList[i] = i; - AscendC::SetFlag(l0CEventList[i]); - } - } else { - l0CTensorList[0] = resource.l0CBuf.template GetBufferByByte(0); - } - if constexpr (HAS_BIAS) { - l1BiasTensor = resource.l1Buf.template GetBufferByByte(l1BiasOffset); - l0BiasTensor = resource.btBuf.template GetBufferByByte(0); - AscendC::SetFlag(L1A_STAGES + L1B_STAGES); - AscendC::SetFlag(L0A_STAGES + L0B_STAGES); - } - if constexpr (HAS_QUANT_TENSOR) { - l1QuantTensor = resource.l1Buf.template GetBufferByByte(l1QuantOffset); - AscendC::SetFlag(QUANT_EVENT_ID); - } - - if constexpr (ENABLE_L1_RESIDENT) { - RestoreStatus(); - } - } - } - - /// Destructor - CATLASS_DEVICE - ~BlockMmadTla() - { - if ASCEND_IS_AIC { - if constexpr (USE_HF32_MODE) { - AscendC::SetHF32Mode(false); - } - if constexpr (ENABLE_UNIT_FLAG && tla::detail::isRowMajor::value) { - AscendC::SetMMLayoutTransform(false); - } - for (uint32_t i = 0; i < L1A_STAGES; i++) { - AscendC::WaitFlag(l1AEventList[i]); - } - for (uint32_t i = 0; i < L1B_STAGES; i++) { - AscendC::WaitFlag(l1BEventList[i]); - } - for (uint32_t i = 0; i < L0A_STAGES; i++) { - AscendC::WaitFlag(l0AEventList[i]); - } - for (uint32_t i = 0; i < L0B_STAGES; i++) { - AscendC::WaitFlag(l0BEventList[i]); - } - if constexpr (!ENABLE_UNIT_FLAG) { - for (uint32_t i = 0; i < L0C_STAGES; i++) { - AscendC::WaitFlag(l0CEventList[i]); - } - } - if constexpr (HAS_BIAS) { - AscendC::WaitFlag(L1A_STAGES + L1B_STAGES); - AscendC::WaitFlag(L0A_STAGES + L0B_STAGES); - } - if constexpr (HAS_QUANT_TENSOR) { - AscendC::WaitFlag(QUANT_EVENT_ID); - } - } - } - - /// Perform a block-scoped matrix multiply-accumulate - template - CATLASS_DEVICE void operator()(TensorA &tensorA, TensorB &tensorB, TensorC &tensorC, GemmCoord const &actualShape, - TensorBias const &tensorBias = {}, TensorQuant const &tensorQuant = {}, float scaleParam = 1.0) - { - // Check TileShape - if constexpr (HAS_BIAS) { - static constexpr uint32_t BIAS_BUF_SIZE = L0_TILE_N * sizeof(ElementAccumulator); - static constexpr uint32_t L1BIAS_SIZE = L1_TILE_N * sizeof(ElementBias); - static_assert(BIAS_BUF_SIZE <= ArchTag::BIAS_SIZE, - "BIAS_BUF_SIZE exceeding the BT space! Reduce L0_TILE_N"); - static_assert(L1AB_TILE_SIZE + L1BIAS_SIZE <= ArchTag::L1_SIZE, - "L1TileShape exceeding the L1 space!"); - } - if constexpr (HAS_QUANT_TENSOR) { - static constexpr uint32_t L1QUANT_SIZE = L1_TILE_N * sizeof(uint64_t); - if constexpr (HAS_BIAS) { - static constexpr uint32_t L1BIAS_SIZE = L1_TILE_N * sizeof(ElementBias); - static_assert(L1AB_TILE_SIZE + L1BIAS_SIZE + L1QUANT_SIZE <= ArchTag::L1_SIZE, - "L1TileShape exceeding the L1 space!"); - } else { - static_assert(L1AB_TILE_SIZE + L1QUANT_SIZE <= ArchTag::L1_SIZE, - "L1TileShape exceeding the L1 space!"); - } - } - - using CopyGmToL1A = typename TileCopy_::template CopyGmToL1A; - using CopyGmToL1B = typename TileCopy_::template CopyGmToL1B; - using CopyL0CToDst = typename TileCopy_::template CopyL0CToDst; - CopyGmToL1A copyGmToL1A; - CopyGmToL1B copyGmToL1B; - CopyL0CToDst copyL0CToDst; - - uint32_t mBlockActual = actualShape.m(); - uint32_t kBlockActual = actualShape.k(); - uint32_t nBlockActual = actualShape.n(); - - uint32_t mL1Actual = mBlockActual; - if constexpr (std::is_same_v) { - // Avoid using the gemv mode in mmad - if (mL1Actual == 1) { - mL1Actual = 16; - } - } - uint32_t nL1Actual = nBlockActual; - - auto layoutInL0C = tla::MakeLayoutL0C(mL1Actual, nL1Actual); - auto tensorL0C = tla::MakeTensor(l0CTensorList[l0CListId], layoutInL0C, Arch::PositionL0C{}); - auto tensorL0Bias = tla::MakeTensor(l0BiasTensor, L0BIAS_LAYOUT, Arch::PositionBias{}); - - uint32_t kL1Actual = min(kBlockActual, L1_TILE_K); - // load first matrix A tile from GM to L1 - AscendC::WaitFlag(l1AEventList[l1AListId]); - auto tensorL1A = tla::MakeTensor(l1ATensorList[l1AListId], L1A_LAYOUT, Arch::PositionL1{}); - auto tensorTileA = GetTileA(tensorA, 0, 0, mBlockActual, kL1Actual); - if constexpr (ENABLE_L1_RESIDENT) { - // If the currently loaded GM pointer and block coordinates are the same as the last loaded ones, - // skip this loading. - if (lastAddrA[l1AListId] != tensorTileA.data().GetPhyAddr() || tla::get<0>(tensorTileA.coord()) != lastCoordA[l1AListId].row() || tla::get<1>(tensorTileA.coord()) != lastCoordA[l1AListId].column()) { - copyGmToL1A(tensorL1A, tensorTileA); - lastCoordA[l1AListId] = MatrixCoord{tla::get<0>(tensorTileA.coord()), tla::get<1>(tensorTileA.coord())}; - lastAddrA[l1AListId] = const_cast<__gm__ typename AscendC::GlobalTensor::PrimType *>( - tensorTileA.data().GetPhyAddr()); - } - } else { - copyGmToL1A(tensorL1A, tensorTileA); - } - AscendC::SetFlag(l1AEventList[l1AListId]); - - // load first matrix B tile from GM to L1 - AscendC::WaitFlag(l1BEventList[l1BListId]); - auto tensorL1B = tla::MakeTensor(l1BTensorList[l1BListId], L1B_LAYOUT, Arch::PositionL1{}); - auto tensorTileB = GetTile(tensorB, tla::MakeCoord(0, 0), tla::MakeShape(kL1Actual, nBlockActual)); - if constexpr (ENABLE_L1_RESIDENT) { - if (lastAddrB[l1BListId] != tensorTileB.data().GetPhyAddr() || tla::get<0>(tensorTileB.coord()) != lastCoordB[l1BListId].row() || tla::get<1>(tensorTileB.coord()) != lastCoordB[l1BListId].column()) { - copyGmToL1B(tensorL1B, tensorTileB); - lastCoordB[l1BListId] = MatrixCoord{tla::get<0>(tensorTileB.coord()), tla::get<1>(tensorTileB.coord())}; - lastAddrB[l1BListId] = const_cast<__gm__ typename AscendC::GlobalTensor::PrimType *>( - tensorTileB.data().GetPhyAddr()); - } - } else { - copyGmToL1B(tensorL1B, tensorTileB); - } - AscendC::SetFlag(l1BEventList[l1BListId]); - - if constexpr (HAS_BIAS && !std::is_same_v) { - using CopyGmToL1Bias = typename TileCopy::template CopyGmToL1Bias; - CopyGmToL1Bias copyGmToL1Bias; - AscendC::WaitFlag(L1A_STAGES + L1B_STAGES); - auto l1Bias = l1BiasTensor.template ReinterpretCast(); - auto tensorL1Bias = tla::MakeTensor(l1Bias, L1BIAS_LAYOUT, Arch::PositionL1{}); - copyGmToL1Bias(tensorL1Bias, tensorBias); - AscendC::SetFlag(L1A_STAGES + L1B_STAGES); - } - - if constexpr (!ENABLE_UNIT_FLAG) { - AscendC::WaitFlag(l0CEventList[l0CListId]); - } - - uint32_t mL0Loop = CeilDiv(mL1Actual); - uint32_t nL0Loop = CeilDiv(nL1Actual); - - // main loop - uint32_t kL1Loop = CeilDiv(kBlockActual); - for (uint32_t kL1Idx = 0; kL1Idx < kL1Loop; kL1Idx++) { - uint32_t l1AListIdNext = (l1AListId + 1 < L1A_STAGES) ? (l1AListId + 1) : 0; - uint32_t l1BListIdNext = (l1BListId + 1 < L1B_STAGES) ? (l1BListId + 1) : 0; - uint32_t kL1ActualNext{0}; - // preload next tile from GM to L1 - if (kL1Idx < kL1Loop - 1) { - uint32_t kL1IdxNext = kL1Idx + 1; - kL1ActualNext = (kL1IdxNext < kL1Loop - 1) ? L1_TILE_K : (kBlockActual - kL1IdxNext * L1_TILE_K); - - // Get L1 tensor for next stage - auto l1ATensor = l1ATensorList[l1AListIdNext]; - auto l1BTensor = l1BTensorList[l1BListIdNext]; - auto tensorL1A = tla::MakeTensor(l1ATensor, L1A_LAYOUT, Arch::PositionL1{}); - auto tensorL1B = tla::MakeTensor(l1BTensor, L1B_LAYOUT, Arch::PositionL1{}); - // Get GM tile for next stage - auto tensorTileA = GetTileA(tensorA, 0, kL1IdxNext * L1_TILE_K, mBlockActual, kL1ActualNext); - auto tensorTileB = GetTile(tensorB, tla::MakeCoord(kL1IdxNext * L1_TILE_K, 0), - tla::MakeShape(kL1ActualNext, nBlockActual)); - - // load next matrix A tile from GM to L1 - AscendC::WaitFlag(l1AEventList[l1AListIdNext]); - if constexpr (ENABLE_L1_RESIDENT) { - if (lastAddrA[l1AListIdNext] != tensorTileA.data().GetPhyAddr() || tla::get<0>(tensorTileA.coord()) != lastCoordA[l1AListIdNext].row() || tla::get<1>(tensorTileA.coord()) != lastCoordA[l1AListIdNext].column()) { - copyGmToL1A(tensorL1A, tensorTileA); - lastCoordA[l1AListIdNext] = - MatrixCoord{tla::get<0>(tensorTileA.coord()), tla::get<1>(tensorTileA.coord())}; - lastAddrA[l1AListIdNext] = - const_cast<__gm__ typename AscendC::GlobalTensor::PrimType *>( - tensorTileA.data().GetPhyAddr()); - } - } else { - copyGmToL1A(tensorL1A, tensorTileA); - } - AscendC::SetFlag(l1AEventList[l1AListIdNext]); - - // load next matrix B tile from GM to L1 - AscendC::WaitFlag(l1BEventList[l1BListIdNext]); - if constexpr (ENABLE_L1_RESIDENT) { - if (lastAddrB[l1BListIdNext] != tensorTileB.data().GetPhyAddr() || tla::get<0>(tensorTileB.coord()) != lastCoordB[l1BListIdNext].row() || tla::get<1>(tensorTileB.coord()) != lastCoordB[l1BListIdNext].column()) { - copyGmToL1B(tensorL1B, tensorTileB); - lastCoordB[l1BListIdNext] = - MatrixCoord{tla::get<0>(tensorTileB.coord()), tla::get<1>(tensorTileB.coord())}; - lastAddrB[l1BListIdNext] = - const_cast<__gm__ typename AscendC::GlobalTensor::PrimType *>( - tensorTileB.data().GetPhyAddr()); - } - } else { - copyGmToL1B(tensorL1B, tensorTileB); - } - AscendC::SetFlag(l1BEventList[l1BListIdNext]); - } - - // Get L1 tensor for current stage - auto l1ATensor = l1ATensorList[l1AListId]; - auto l1BTensor = l1BTensorList[l1BListId]; - tensorL1A = tla::MakeTensor(l1ATensor, L1A_LAYOUT, Arch::PositionL1{}); - tensorL1B = tla::MakeTensor(l1BTensor, L1B_LAYOUT, Arch::PositionL1{}); - // Get the loop nums on L0 - uint32_t kL0Loop = CeilDiv(kL1Actual); - - for (int mL0Idx = 0; mL0Idx < mL0Loop; mL0Idx++) { - uint32_t mL0Actual = (mL0Idx < mL0Loop - 1) ? L0_TILE_M : (mL1Actual - mL0Idx * L0_TILE_M); - - for (int kL0Idx = 0; kL0Idx < kL0Loop; kL0Idx++) { - uint32_t kL0Actual = (kL0Idx < kL0Loop - 1) ? L0_TILE_K : (kL1Actual - kL0Idx * L0_TILE_K); - - // Locate the current tile on L0A - auto l0ATile = l0ATensorList[l0AListId]; - auto layoutAInL0 = tla::MakeLayout(mL0Actual, kL0Actual); - auto tensorL0A = tla::MakeTensor(l0ATile, layoutAInL0, Arch::PositionL0A{}); - // Locate the current tile of matrix A on L1 - auto tensorTileL1A = GetTileA(tensorL1A, mL0Idx * L0_TILE_M, kL0Idx * L0_TILE_K, mL0Actual, kL0Actual); - - AscendC::WaitFlag(l0AEventList[l0AListId]); - if ((mL0Idx == 0) && (kL0Idx == 0)) { - AscendC::WaitFlag(l1AEventList[l1AListId]); - } - - // Load current tile from L1 to L0A - copyL1ToL0A(tensorL0A, tensorTileL1A); - - if ((mL0Idx == mL0Loop - 1) && (kL0Idx == kL0Loop - 1)) { - AscendC::SetFlag(l1AEventList[l1AListId]); - } - - bool initC = ((kL1Idx == 0) && (kL0Idx == 0)); - for (int nL0Idx = 0; nL0Idx < nL0Loop; nL0Idx++) { - uint32_t nL0Actual = (nL0Idx < nL0Loop - 1) ? L0_TILE_N : (nL1Actual - nL0Idx * L0_TILE_N); - - // Locate the current tile on L0B - auto l0BTile = l0BTensorList[l0BListId]; - auto layoutBInL0 = tla::MakeLayout(kL0Actual, nL0Actual); - auto tensorL0B = tla::MakeTensor(l0BTile, layoutBInL0, Arch::PositionL0B{}); - // Locate the current tile of matrix B on L1 - auto tensorTileL1B = GetTile(tensorL1B, - tla::MakeCoord(kL0Idx * L0_TILE_K, nL0Idx * L0_TILE_N), - tla::MakeShape(kL0Actual, nL0Actual)); - - // Wait for mmad finished - AscendC::WaitFlag(l0BEventList[l0BListId]); - // If the current tile is the first one on the k&n axis, wait for loading matrix B from GM to L1 - if ((mL0Idx == 0) && (kL0Idx == 0) && (nL0Idx == 0)) { - AscendC::WaitFlag(l1BEventList[l1BListId]); - } - - // Load current tile from L1 to L0B - copyL1ToL0B(tensorL0B, tensorTileL1B); - - // If the current tile is the last one on the k&n axis, notify to load matrix B from GM to L1 - if ((mL0Idx == mL0Loop - 1) && (kL0Idx == kL0Loop - 1) && (nL0Idx == nL0Loop - 1)) { - AscendC::SetFlag(l1BEventList[l1BListId]); - } - - if constexpr (HAS_BIAS && !std::is_same_v) { - if (initC) { - if (nL0Idx == 0) { - AscendC::WaitFlag(L1A_STAGES + L1B_STAGES); - } - AscendC::WaitFlag(L0A_STAGES + L0B_STAGES); - auto l1Bias = l1BiasTensor.template ReinterpretCast(); - auto tensorL1Bias = tla::MakeTensor(l1Bias, L1BIAS_LAYOUT, Arch::PositionL1{}); - auto tensorTileL1Bias = GetTile(tensorL1Bias, - tla::MakeCoord(nL0Idx * L0_TILE_N), - tla::MakeShape(nL0Actual)); - // Load bias to l0 biastable - copyL1ToBT(tensorL0Bias, tensorTileL1Bias); - if (nL0Idx == nL0Loop - 1) { - AscendC::SetFlag(L1A_STAGES + L1B_STAGES); - } - } - } - - // Notify to do mmad - AscendC::SetFlag(l0CEventList[l0CListId]); - - // Locate the current tile on L0C - auto tensorTileL0C = GetTile(tensorL0C, - tla::MakeCoord(mL0Idx * L0_TILE_M, nL0Idx * L0_TILE_N), - tla::MakeShape(mL0Actual, nL0Actual)); - - // Compute the matrix multiplication on L0A and L0B and write the result to the accumulator - // Wait for loading L0B - AscendC::WaitFlag(l0CEventList[l0CListId]); - - // If the unit flag is enabled, the unit flag is set according to the calculation progress - uint8_t unitFlag = 0b00; - if constexpr (ENABLE_UNIT_FLAG) { - if ((kL1Idx == kL1Loop - 1) && (mL0Idx == mL0Loop - 1) && - (kL0Idx == kL0Loop - 1) && (nL0Idx == nL0Loop - 1)) { - unitFlag = 0b11; - } else { - unitFlag = 0b10; - } - } - - if constexpr (HAS_BIAS && !std::is_same_v) { - if (initC) { - tileMmad(tensorTileL0C, tensorL0A, tensorL0B, tensorL0Bias, - mL0Actual, nL0Actual, kL0Actual, initC, unitFlag); - AscendC::SetFlag(L0A_STAGES + L0B_STAGES); - } else { - tileMmad(tensorTileL0C, tensorL0A, tensorL0B, - mL0Actual, nL0Actual, kL0Actual, initC, unitFlag); - } - } else { - tileMmad(tensorTileL0C, tensorL0A, tensorL0B, - mL0Actual, nL0Actual, kL0Actual, initC, unitFlag); - } - - // Notify to move the next L0B tile - AscendC::SetFlag(l0BEventList[l0BListId]); - l0BListId = (l0BListId + 1 < L0B_STAGES) ? (l0BListId + 1) : 0; - } - AscendC::SetFlag(l0AEventList[l0AListId]); - l0AListId = (l0AListId + 1 < L0A_STAGES) ? (l0AListId + 1) : 0; - } - } - l1AListId = l1AListIdNext; - l1BListId = l1BListIdNext; - kL1Actual = kL1ActualNext; - } - - // copy block out - if constexpr (DEQUANT_MODE == Tile::ScaleGranularity::PER_TENSOR) { - copyL0CToDst.params.scale = scaleParam; - } - - if constexpr (HAS_QUANT_TENSOR && !std::is_same_v) { - using CopyGmToL1Scale = typename TileCopy::template CopyGmToL1Scale; - CopyGmToL1Scale copyGmToL1Scale; - auto tensorL1Quant = tla::MakeTensor(l1QuantTensor, L1QUANT_LAYOUT, Arch::PositionL1{}); - AscendC::WaitFlag(QUANT_EVENT_ID); - copyGmToL1Scale(tensorL1Quant, tensorQuant); - AscendC::SetFlag(QUANT_EVENT_ID); - AscendC::WaitFlag(QUANT_EVENT_ID); - if constexpr (!ENABLE_UNIT_FLAG) { - AscendC::SetFlag(l0CEventList[l0CListId]); - AscendC::WaitFlag(l0CEventList[l0CListId]); - copyL0CToDst(tensorC, tensorL0C, tensorL1Quant); - AscendC::SetFlag(l0CEventList[l0CListId]); - l0CListId = (l0CListId + 1 < L0C_STAGES) ? (l0CListId + 1) : 0; - } else { - copyL0CToDst(tensorC, tensorL0C, tensorL1Quant, 0b11); - } - AscendC::SetFlag(QUANT_EVENT_ID); - } else { - if constexpr (!ENABLE_UNIT_FLAG) { - AscendC::SetFlag(l0CEventList[l0CListId]); - AscendC::WaitFlag(l0CEventList[l0CListId]); - copyL0CToDst(tensorC, tensorL0C); - AscendC::SetFlag(l0CEventList[l0CListId]); - l0CListId = (l0CListId + 1 < L0C_STAGES) ? (l0CListId + 1) : 0; - } else { - copyL0CToDst(tensorC, tensorL0C, 0b11); - } - } - } - -protected: - template - CATLASS_DEVICE auto GetTileA(TensorA &tensorA, uint32_t mIndex, uint32_t kIndex, uint32_t mSize, uint32_t kSize) - { - if constexpr (tla::detail::isVector::value) { - return GetTile(tensorA, tla::MakeCoord(kIndex), tla::MakeShape(kSize)); - } else { - return GetTile(tensorA, tla::MakeCoord(mIndex, kIndex), tla::MakeShape(mSize, kSize)); - } - } - - // Multi-stage tensors list - AscendC::LocalTensor l1ATensorList[L1A_STAGES]; - AscendC::LocalTensor l1BTensorList[L1B_STAGES]; - AscendC::LocalTensor l0ATensorList[L0A_STAGES]; - AscendC::LocalTensor l0BTensorList[L0B_STAGES]; - AscendC::LocalTensor l0CTensorList[L0C_STAGES]; - AscendC::LocalTensor l1BiasTensor; - AscendC::LocalTensor l0BiasTensor; - AscendC::LocalTensor l1QuantTensor; - - // Multi-stage event id list - int32_t l1AEventList[L1A_STAGES]; - int32_t l1BEventList[L1B_STAGES]; - int32_t l0AEventList[L0A_STAGES]; - int32_t l0BEventList[L0B_STAGES]; - int32_t l0CEventList[L0C_STAGES]; - static constexpr int32_t QUANT_EVENT_ID = L1A_STAGES + L1B_STAGES + 1; - - __gm__ typename AscendC::GlobalTensor::PrimType *lastAddrA[L1A_STAGES]; - __gm__ typename AscendC::GlobalTensor::PrimType *lastAddrB[L1B_STAGES]; - MatrixCoord lastCoordA[L1A_STAGES]; - MatrixCoord lastCoordB[L1B_STAGES]; - - // The id of current stage - uint32_t l1AListId{0}; - uint32_t l1BListId{0}; - uint32_t l0AListId{0}; - uint32_t l0BListId{0}; - uint32_t l0CListId{0}; - - TileMmad tileMmad; - CopyL1ToL0A copyL1ToL0A; - CopyL1ToL0B copyL1ToL0B; - CopyL1ToBT copyL1ToBT; -}; - -} // namespace Catlass::Gemm::Block - -#endif // CATLASS_GEMM_BLOCK_BLOCK_MMAD_PINGPONG_DEQUANT_TLA_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_pingpong_full_loadA.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_pingpong_full_loadA.hpp deleted file mode 100644 index 802dc7a5bc..0000000000 --- a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_pingpong_full_loadA.hpp +++ /dev/null @@ -1,347 +0,0 @@ -/** - * Copyright (c) 2025 Huawei Technologies Co., Ltd. - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of - * CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ - -#ifndef CATLASS_GEMM_BLOCK_BLOCK_MMAD_PINGPONG_FULL_LOADA_HPP -#define CATLASS_GEMM_BLOCK_BLOCK_MMAD_PINGPONG_FULL_LOADA_HPP - -#include "catlass/kda_catlass.hpp" -#include "catlass/arch/kda_resource.hpp" -#include "catlass/kda_coord.hpp" -#include "catlass/kda_gemm_coord.hpp" -#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" -#include "catlass/gemm/kda_gemm_helper.hpp" -#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" -#include "catlass/gemm/tile/kda_tile_mmad.hpp" - -namespace Catlass::Gemm::Block { - -template < - bool ENABLE_UNIT_FLAG_, - class L1TileShape_, - class L0TileShape_, - class AType_, - class BType_, - class CType_, - class BiasType_, - class TileCopy_, - class TileMmad_> -struct BlockMmad< - MmadAtlasA2FullLoadA, - L1TileShape_, - L0TileShape_, - AType_, - BType_, - CType_, - BiasType_, - TileCopy_, - TileMmad_> { -public: - // Type Aliases - using DispatchPolicy = MmadAtlasA2FullLoadA; - using ArchTag = typename DispatchPolicy::ArchTag; - using L1TileShape = L1TileShape_; - using L0TileShape = L0TileShape_; - using ElementA = typename AType_::Element; - using LayoutA = typename AType_::Layout; - using ElementB = typename BType_::Element; - using LayoutB = typename BType_::Layout; - using ElementC = typename CType_::Element; - using LayoutC = typename CType_::Layout; - using TileMmad = TileMmad_; - using CopyGmToL1A = typename TileCopy_::CopyGmToL1A; - using CopyGmToL1B = typename TileCopy_::CopyGmToL1B; - using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; - using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; - using CopyL0CToGm = typename TileCopy_::CopyL0CToGm; - using ElementAccumulator = - typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; - using LayoutAInL1 = typename CopyL1ToL0A::LayoutSrc; - using LayoutBInL1 = typename CopyL1ToL0B::LayoutSrc; - using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst; - using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst; - using LayoutCInL0 = layout::zN; - - using L1AAlignHelper = Gemm::helper::L1AlignHelper; - using L1BAlignHelper = Gemm::helper::L1AlignHelper; - - static constexpr bool ENABLE_UNIT_FLAG = DispatchPolicy::ENABLE_UNIT_FLAG; - static constexpr uint32_t STAGES = DispatchPolicy::STAGES; - static constexpr uint32_t L1A_SIZE = ArchTag::L1_SIZE / 2; - static constexpr uint32_t L1B_SIZE = L1TileShape::N * L1TileShape::K * sizeof(ElementB); - static constexpr uint32_t L0A_SIZE = ArchTag::L0A_SIZE; - static constexpr uint32_t L0B_SIZE = ArchTag::L0B_SIZE; - static constexpr uint32_t L0C_SIZE = ArchTag::L0C_SIZE; - static constexpr uint32_t L0A_PINGPONG_BUF_SIZE = L0A_SIZE / STAGES; - static constexpr uint32_t L0B_PINGPONG_BUF_SIZE = L0B_SIZE / STAGES; - - // Check LayoutC - static_assert(std::is_same_v, "LayoutC only support RowMajor yet!"); - - // Check L1TileShape in example, because problemShape.k() cannot get here. - - // Check L0TileShape - static constexpr uint32_t L0A_TILE_SIZE = L0TileShape::M * L0TileShape::K * sizeof(ElementA); - static constexpr uint32_t L0B_TILE_SIZE = L0TileShape::K * L0TileShape::N * sizeof(ElementB); - static constexpr uint32_t L0C_TILE_SIZE = L0TileShape::M * L0TileShape::N * sizeof(ElementAccumulator); - static_assert((L0A_TILE_SIZE * STAGES) <= L0A_SIZE, "L0TileShape exceeding the L0A space!"); - static_assert((L0B_TILE_SIZE * STAGES) <= L0B_SIZE, "L0TileShape exceeding the L0B space!"); - static_assert(L0C_TILE_SIZE <= L0C_SIZE, "L0TileShape exceeding the L0C space!"); - - static_assert(L1TileShape::M == L0TileShape::M && L1TileShape::N == L0TileShape::N, - "The situation where the basic blocks of L1 and L0 differ on the m and n axes is not supported yet"); - // 32B (256b) aligned - static_assert(Gemm::helper::TileShapeAlignChecker::_ALIGN == 256, - "Tile shape must be 32B aligned."); - - /// Construct - CATLASS_DEVICE - BlockMmad(Arch::Resource &resource, uint32_t l1BufAddrStart = 0) - { - uint32_t l1AOffset = l1BufAddrStart; - uint32_t l1BOffset = l1BufAddrStart + L1A_SIZE; - // Init L1A related - l1ATensorList[0] = resource.l1Buf.template GetBufferByByte(l1AOffset); - l1AEventList[0] = 0; - - // Init buffers - for (uint32_t i = 0; i < STAGES; i++) { - // Assign L1/L0A/L0B space for each stages - l1BTensorList[i] = resource.l1Buf.template GetBufferByByte(l1BOffset + L1B_SIZE * i); - l0ATensorList[i] = resource.l0ABuf.template GetBufferByByte(L0A_PINGPONG_BUF_SIZE * i); - l0BTensorList[i] = resource.l0BBuf.template GetBufferByByte(L0B_PINGPONG_BUF_SIZE * i); - - // Assign event ID for each stages - l1BEventList[i] = i + STAGES; - l0AEventList[i] = i; - l0BEventList[i] = i + STAGES; - - // The event id that needs to be set before the loop - AscendC::SetFlag(l1BEventList[i]); - AscendC::SetFlag(l0AEventList[i]); - AscendC::SetFlag(l0BEventList[i]); - } - l0CTensor = resource.l0CBuf.template GetBufferByByte(0); - AscendC::SetFlag(EVENT_ID0); - } - - /// Destructor - CATLASS_DEVICE - ~BlockMmad() - { - for (uint32_t i = 0; i < STAGES; i++) { - AscendC::WaitFlag(l1BEventList[i]); - AscendC::WaitFlag(l0AEventList[i]); - AscendC::WaitFlag(l0BEventList[i]); - } - AscendC::WaitFlag(EVENT_ID0); - } - - /// Perform a block-scoped matrix multiply-accumulate - CATLASS_DEVICE - void operator()( - AscendC::GlobalTensor const &gmA, LayoutA const &layoutA, - AscendC::GlobalTensor const &gmB, LayoutB const &layoutB, - AscendC::GlobalTensor const &gmC, LayoutC const &layoutC, - GemmCoord const &actualShape, bool needLoadL1) - { - uint32_t mRound = RoundUp(actualShape.m()); - uint32_t nRound = RoundUp(actualShape.n()); - - auto layoutAInL1 = LayoutAInL1::template MakeLayout(L1TileShape::M, actualShape.k()); - auto layoutBInL1 = LayoutBInL1::template MakeLayout(L1TileShape::K, L1TileShape::N); - auto layoutInL0C = LayoutCInL0::MakeLayoutInL0C(MakeCoord(mRound, nRound)); - - uint32_t kActual = min(actualShape.k(), L1TileShape::K); - - if (needLoadL1) { - // load first matrix A tile from GM to L1 - auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShape.m(), actualShape.k())); - AscendC::SetFlag(l1AEventList[0]); - AscendC::WaitFlag(l1AEventList[0]); - copyGmToL1A(l1ATensorList[0], gmA, layoutAInL1, layoutTileA); - AscendC::SetFlag(l1AEventList[0]); - AscendC::WaitFlag(l1AEventList[0]); - } - - // load first matrix B tile from GM to L1 - AscendC::WaitFlag(l1BEventList[l1ListId]); - auto layoutTileB = layoutB.GetTileLayout(MakeCoord(kActual, actualShape.n())); - copyGmToL1B(l1BTensorList[l1ListId], gmB, layoutBInL1, layoutTileB); - AscendC::SetFlag(l1BEventList[l1ListId]); - - if constexpr (!ENABLE_UNIT_FLAG) { - AscendC::WaitFlag(EVENT_ID0); - } - - uint32_t mPartLoop = CeilDiv(mRound); - uint32_t nPartLoop = CeilDiv(nRound); - - // main loop - uint32_t kTileCount = CeilDiv(actualShape.k()); - for (uint32_t kLoopIdx = 0; kLoopIdx < kTileCount; kLoopIdx++) { - uint32_t l1ListIdNext = (l1ListId + 1 < STAGES) ? (l1ListId + 1) : 0; - uint32_t kActualNext{0}; - // preload next tile from GM to L1 - if (kLoopIdx < kTileCount - 1) { - uint32_t kLoopIdxNext = kLoopIdx + 1; - kActualNext = (kLoopIdxNext < kTileCount - 1) ? - L1TileShape::K : - (actualShape.k() - kLoopIdxNext * L1TileShape::K); - - // Get L1 tensor for next stage - auto l1BTensor = l1BTensorList[l1ListIdNext]; - // Get GM tile for next stage - MatrixCoord gmTileBOffset{kLoopIdxNext * L1TileShape::K, 0}; - auto gmTileB = gmB[layoutB.GetOffset(gmTileBOffset)]; - - // load next matrix B tile from GM to L1 - AscendC::WaitFlag(l1BEventList[l1ListIdNext]); - layoutTileB = layoutB.GetTileLayout(MakeCoord(kActualNext, actualShape.n())); - copyGmToL1B(l1BTensor, gmTileB, layoutBInL1, layoutTileB); - AscendC::SetFlag(l1BEventList[l1ListIdNext]); - } - - // Get L1 tensor for current stage - auto l1ATensor = l1ATensorList[0]; - auto l1BTensor = l1BTensorList[l1ListId]; - - // Get the loop nums on L0 - uint32_t kPartLoop = CeilDiv(kActual); - - for (int mPartIdx = 0; mPartIdx < mPartLoop; mPartIdx++) { - uint32_t mPartActual = (mPartIdx < mPartLoop - 1) ? - L0TileShape::M : - (mRound - mPartIdx * L0TileShape::M); - - for (int kPartIdx = 0; kPartIdx < kPartLoop; kPartIdx++) { - uint32_t kPartActual = (kPartIdx < kPartLoop - 1) ? - L0TileShape::K : - (kActual - kPartIdx * L0TileShape::K); - - // Locate the current tile on L0A - auto l0ATile = l0ATensorList[l0AListId]; - LayoutAInL0 layoutAInL0 = LayoutAInL0::template MakeLayout(mPartActual, kPartActual); - // Locate the current tile of matrix A on L1 - MatrixCoord l1AOffset{mPartIdx * L0TileShape::M, - kPartIdx * L0TileShape::K + kLoopIdx * L1TileShape::K}; - auto l1ATile = l1ATensor[layoutAInL1.GetOffset(l1AOffset)]; - - AscendC::WaitFlag(l0AEventList[l0AListId]); - - // Load current tile from L1 to L0A - copyL1ToL0A(l0ATile, l1ATile, layoutAInL0, layoutAInL1); - - for (int nPartIdx = 0; nPartIdx < nPartLoop; nPartIdx++) { - uint32_t nPartActual = (nPartIdx < nPartLoop - 1) ? - L0TileShape::N : - (nRound - nPartIdx * L0TileShape::N); - - // Locate the current tile on L0B - auto l0BTile = l0BTensorList[l0BListId]; - LayoutBInL0 layoutBInL0 = LayoutBInL0::template MakeLayout(kPartActual, nPartActual); - // Locate the current tile of matrix B on L1 - MatrixCoord l1BOffset{kPartIdx * L0TileShape::K, nPartIdx * L0TileShape::N}; - auto l1BTile = l1BTensor[layoutBInL1.GetOffset(l1BOffset)]; - - // Wait for mmad finished - AscendC::WaitFlag(l0BEventList[l0BListId]); - // If the current tile is the first one on the k&n axis, wait for loading matrix B from GM to L1 - if ((kPartIdx == 0) && (nPartIdx == 0)) { - AscendC::WaitFlag(l1BEventList[l1ListId]); - } - - // Load current tile from L1 to L0B - copyL1ToL0B(l0BTile, l1BTile, layoutBInL0, layoutBInL1); - - // If the current tile is the last one on the k&n axis, notify to load matrix B from GM to L1 - if ((kPartIdx == kPartLoop - 1) && (nPartIdx == nPartLoop - 1)) { - AscendC::SetFlag(l1BEventList[l1ListId]); - } - // Notify to do mmad - AscendC::SetFlag(EVENT_ID0); - - // Locate the current tile on L0C - MatrixCoord l0COffset{mPartIdx * L0TileShape::M, nPartIdx * L0TileShape::N}; - auto l0CTile = l0CTensor[layoutInL0C.GetOffset(l0COffset)]; - - // Compute the matrix multiplication on L0A and L0B and write the result to the accumulator - // Wait for loading L0B - AscendC::WaitFlag(EVENT_ID0); - - // If the current tile is the first tile on the k axis, the accumulator needs to be reset to 0 - bool initC = ((kLoopIdx == 0) && (kPartIdx == 0)); - // If the unit flag is enabled, the unit flag is set according to the calculation progress - uint8_t unitFlag = 0b00; - if constexpr (ENABLE_UNIT_FLAG) { - if ((kLoopIdx == kTileCount - 1) && (mPartIdx == mPartLoop - 1) && - (kPartIdx == kPartLoop - 1) && (nPartIdx == nPartLoop - 1)) { - unitFlag = 0b11; - } else { - unitFlag = 0b10; - } - } - // Perform calculation operations - tileMmad(l0CTile, l0ATile, l0BTile, mPartActual, nPartActual, kPartActual, initC, unitFlag); - - // Notify to move the next L0B tile - AscendC::SetFlag(l0BEventList[l0BListId]); - l0BListId = (l0BListId + 1 < STAGES) ? (l0BListId + 1) : 0; - } - AscendC::SetFlag(l0AEventList[l0AListId]); - l0AListId = (l0AListId + 1 < STAGES) ? (l0AListId + 1) : 0; - } - } - l1ListId = l1ListIdNext; - kActual = kActualNext; - } - - // copy block out - LayoutC layoutBlock = layoutC.GetTileLayout(actualShape.GetCoordMN()); - - if constexpr (!ENABLE_UNIT_FLAG) { - AscendC::SetFlag(EVENT_ID0); - AscendC::WaitFlag(EVENT_ID0); - copyL0CToGm(gmC, l0CTensor, layoutBlock, layoutInL0C); - AscendC::SetFlag(EVENT_ID0); - } else { - copyL0CToGm(gmC, l0CTensor, layoutBlock, layoutInL0C, 0b11); - } - } - -protected: - // Multi-stage tensors list - AscendC::LocalTensor l1ATensorList[STAGES]; - AscendC::LocalTensor l1BTensorList[STAGES]; - AscendC::LocalTensor l0ATensorList[STAGES]; - AscendC::LocalTensor l0BTensorList[STAGES]; - AscendC::LocalTensor l0CTensor; - - // Multi-stage event id list - int32_t l1AEventList[STAGES]; - int32_t l1BEventList[STAGES]; - int32_t l0AEventList[STAGES]; - int32_t l0BEventList[STAGES]; - - // The id of current stage - uint32_t l1ListId{0}; - uint32_t l0AListId{0}; - uint32_t l0BListId{0}; - - TileMmad tileMmad; - CopyGmToL1A copyGmToL1A; - CopyGmToL1B copyGmToL1B; - CopyL1ToL0A copyL1ToL0A; - CopyL1ToL0B copyL1ToL0B; - CopyL0CToGm copyL0CToGm; -}; - -} // namespace Catlass::Gemm::Block - -#endif // CATLASS_GEMM_BLOCK_BLOCK_MMAD_PINGPONG_FULL_LOADA_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_pingpong_full_loadA_tla.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_pingpong_full_loadA_tla.hpp deleted file mode 100644 index fd323844ad..0000000000 --- a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_pingpong_full_loadA_tla.hpp +++ /dev/null @@ -1,582 +0,0 @@ -/** - * Copyright (c) 2026 Huawei Technologies Co., Ltd. - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of - * CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ -#ifndef CATLASS_GEMM_BLOCK_BLOCK_MMAD_PINGPONG_FULL_LOADA_TLA_HPP -#define CATLASS_GEMM_BLOCK_BLOCK_MMAD_PINGPONG_FULL_LOADA_TLA_HPP - -#include "catlass/kda_catlass.hpp" -#include "catlass/arch/kda_resource.hpp" -#include "catlass/kda_coord.hpp" -#include "catlass/kda_gemm_coord.hpp" -#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" -#include "catlass/gemm/kda_gemm_helper.hpp" -#include "tla/kda_tla_layout.hpp" -#include "tla/kda_tensor.hpp" - -namespace Catlass::Gemm::Block { - -template < - class ArchTag_, - bool ENABLE_UNIT_FLAG_, - bool ENABLE_SHUFFLE_K_, - bool USE_HF32_MODE_, - uint32_t L0C_STAGES_, - bool ENABLE_L1_RESIDENT_, - uint32_t L1A_STAGES_, - uint32_t L1B_STAGES_, - uint32_t L0A_STAGES_, - uint32_t L0B_STAGES_, - class L1TileShape_, - class L0TileShape_, - class ElementA_, - class ElementB_, - class ElementC_, - class ElementBias_, - class TileCopy_, - class TileMmad_> -struct BlockMmadTla< - MmadFullLoadA, - L1TileShape_, - L0TileShape_, - ElementA_, - ElementB_, - ElementC_, - ElementBias_, - TileCopy_, - TileMmad_> { -public: - // Type Aliases - using DispatchPolicy = MmadFullLoadA; - using ArchTag = typename DispatchPolicy::ArchTag; - using TileCopy = TileCopy_; - using L1TileShape = L1TileShape_; - using L0TileShape = L0TileShape_; - using ElementA = ElementA_; - using LayoutA = typename TileCopy::LayoutA; - using ElementB = ElementB_; - using LayoutB = typename TileCopy::LayoutB; - using ElementC = ElementC_; - using LayoutC = typename TileCopy::LayoutC; - using ElementBias = ElementBias_; - - using TileMmad = TileMmad_; - - using CopyL1ToL0A = typename TileCopy::CopyL1ToL0A; - using CopyL1ToL0B = typename TileCopy::CopyL1ToL0B; - using CopyL1ToBT = typename TileCopy::CopyL1ToBT; - - using ElementAccumulator = typename TileCopy::ElementAccumulator; - - static constexpr bool HAS_BIAS = TileCopy::HAS_BIAS; - - using LayoutTagL1A = typename TileCopy::LayoutTagL1A; - using LayoutTagL1B = typename TileCopy::LayoutTagL1B; - using LayoutTagL0A = typename TileCopy::LayoutTagL0A; - using LayoutTagL0B = typename TileCopy::LayoutTagL0B; - - static_assert(tla::is_tuple::value && tla::is_static::value, - "L1TileShape must be tla::tuple and static!"); - static_assert(tla::is_tuple::value && tla::is_static::value, - "L0TileShape must be tla::tuple and static!"); - - static constexpr bool ENABLE_UNIT_FLAG = DispatchPolicy::ENABLE_UNIT_FLAG; - static constexpr bool ENABLE_SHUFFLE_K = DispatchPolicy::ENABLE_SHUFFLE_K; - static constexpr bool USE_HF32_MODE = DispatchPolicy::USE_HF32_MODE; - static constexpr bool ENABLE_L1_RESIDENT = DispatchPolicy::ENABLE_L1_RESIDENT; - static constexpr uint32_t L1A_STAGES = DispatchPolicy::L1A_STAGES; - static constexpr uint32_t L1B_STAGES = DispatchPolicy::L1B_STAGES; - static constexpr uint32_t L0A_STAGES = DispatchPolicy::L0A_STAGES; - static constexpr uint32_t L0B_STAGES = DispatchPolicy::L0B_STAGES; - static constexpr uint32_t L0C_STAGES = DispatchPolicy::L0C_STAGES; - static constexpr uint32_t L1_TILE_M = tla::get<0>(L1TileShape{}); - static constexpr uint32_t L1_TILE_N = tla::get<1>(L1TileShape{}); - static constexpr uint32_t L1_TILE_K = tla::get<2>(L1TileShape{}); - static constexpr uint32_t L0_TILE_M = tla::get<0>(L0TileShape{}); - static constexpr uint32_t L0_TILE_N = tla::get<1>(L0TileShape{}); - static constexpr uint32_t L0_TILE_K = tla::get<2>(L0TileShape{}); - - // L1 tile size - static constexpr uint32_t L1A_TILE_SIZE = ArchTag::L1_SIZE / 2; - static constexpr uint32_t L1B_TILE_SIZE = L1_TILE_N * L1_TILE_K * sizeof(ElementB); - // L0 tile size - static constexpr uint32_t L0A_TILE_SIZE = L0_TILE_M * L0_TILE_K * sizeof(ElementA); - static constexpr uint32_t L0B_TILE_SIZE = L0_TILE_K * L0_TILE_N * sizeof(ElementB); - static constexpr uint32_t L0C_TILE_SIZE = L1_TILE_M * L1_TILE_N * sizeof(ElementAccumulator); - - // Check HF32_MODE - static_assert( - !USE_HF32_MODE || (USE_HF32_MODE && std::is_same_v && std::is_same_v), - "HF32 MODE only supports in float!"); - - // Check L0C_STAGES - static_assert(!(ENABLE_UNIT_FLAG && L0C_STAGES != 1), "L0C_STAGES must be 1 when UnitFlag is true!"); - - // Check LayoutC - static_assert(tla::detail::isRowMajor::value || - ((std::is_same_v || std::is_same_v || - std::is_same_v) && - tla::detail::iszN::value), - "LayoutC only supports zN in half or bfloat16 or float, RowMajor in all dtype yet!"); - - // Check L1TileShape - static_assert(L1A_TILE_SIZE * L1A_STAGES + L1B_TILE_SIZE * L1B_STAGES <= ArchTag::L1_SIZE, - "L1TileShape exceeding the L1 space!"); - - // Check L0TileShape - static_assert(L0A_TILE_SIZE * L0A_STAGES <= ArchTag::L0A_SIZE, "L0TileShape exceeding the L0A space!"); - static_assert(L0B_TILE_SIZE * L0B_STAGES <= ArchTag::L0B_SIZE, "L0TileShape exceeding the L0B space!"); - static_assert(L0C_TILE_SIZE * L0C_STAGES <= ArchTag::L0C_SIZE, "L0TileShape exceeding the L0C space!"); - - static_assert(L1_TILE_M == L0_TILE_M && L1_TILE_N == L0_TILE_N, - "The situation where the basic blocks of L1 and L0 differ on the m and n axes is not supported yet"); - static_assert(L0_TILE_K <= L1_TILE_K, "L0TileShape::K cannot exceed L1TileShape::K"); - - static_assert((!HAS_BIAS && (L1A_STAGES + L1B_STAGES) <= 8) || (HAS_BIAS && (L1A_STAGES + L1B_STAGES) <= 7), - "L1 Buffer overflow: Exceeds the supported range of EVENT(0~7)"); - - static_assert((!HAS_BIAS && (L0A_STAGES + L0B_STAGES) <= 8) || (HAS_BIAS && (L0A_STAGES + L0B_STAGES) <= 7), - "L0 Buffer overflow: Exceeds the supported range of EVENT_ID(0~7)"); - - static constexpr auto L1B_LAYOUT = - tla::MakeLayout(tla::Int{}, tla::Int{}); - static constexpr auto L1BIAS_LAYOUT = tla::MakeLayout(tla::Int{}); - static constexpr auto L0BIAS_LAYOUT = tla::MakeLayout(tla::Int{}); - - // When enableing L1 resident mode, restore the pointer and coordinates that record the last state - // to the initial state. if tow blockmmad instances need to be consecutively invoked at the kernel layer, - // RestoreStatus() must be inserted between them. - CATLASS_DEVICE - void RestoreStatus() - { - for (int i = 0; i < L1A_STAGES; ++i) { - lastAddrA[i] = nullptr; - lastCoordA[i] = MatrixCoord{0U, 0U}; - } - for (int i = 0; i < L1B_STAGES; ++i) { - lastAddrB[i] = nullptr; - lastCoordB[i] = MatrixCoord{0U, 0U}; - } - } - - /// Construct - CATLASS_DEVICE - BlockMmadTla(Arch::Resource &resource, uint32_t l1BufAddrStart = 0) - { - if ASCEND_IS_AIC { - // use HF32 when USE_HF32_MODE is true - if constexpr (USE_HF32_MODE) { - AscendC::SetHF32Mode(true); - } else { - AscendC::SetHF32Mode(false); - } - if constexpr (ENABLE_UNIT_FLAG && tla::detail::isRowMajor::value) { - AscendC::SetMMLayoutTransform(true); - } - uint32_t l1AOffset = l1BufAddrStart; - uint32_t l1BOffset = l1BufAddrStart + L1A_TILE_SIZE * L1A_STAGES; - // Init buffers - for (uint32_t i = 0; i < L1A_STAGES; i++) { - // Assign L1/L0A/L0B space for each stages - l1ATensorList[i] = resource.l1Buf.template GetBufferByByte(l1AOffset + L1A_TILE_SIZE * i); - // Assign event ID for each stages - l1AEventList[i] = i; - } - for (uint32_t i = 0; i < L1B_STAGES; i++) { - // Assign L1/L0A/L0B space for each stages - l1BTensorList[i] = resource.l1Buf.template GetBufferByByte(l1BOffset + L1B_TILE_SIZE * i); - // Assign event ID for each stages - l1BEventList[i] = i + L1A_STAGES; - // The event id that needs to be set before the loop - AscendC::SetFlag(l1BEventList[i]); - } - for (uint32_t i = 0; i < L0A_STAGES; i++) { - // Assign L1/L0A/L0B space for each stages - l0ATensorList[i] = resource.l0ABuf.template GetBufferByByte(L0A_TILE_SIZE * i); - // Assign event ID for each stages - l0AEventList[i] = i; - // The event id that needs to be set before the loop - AscendC::SetFlag(l0AEventList[i]); - } - for (uint32_t i = 0; i < L0B_STAGES; i++) { - // Assign L1/L0A/L0B space for each stages - l0BTensorList[i] = resource.l0BBuf.template GetBufferByByte(L0B_TILE_SIZE * i); - // Assign event ID for each stages - l0BEventList[i] = i + L0A_STAGES; - // The event id that needs to be set before the loop - AscendC::SetFlag(l0BEventList[i]); - } - if constexpr (!ENABLE_UNIT_FLAG) { - for (uint32_t i = 0; i < L0C_STAGES; i++) { - l0CTensorList[i] = resource.l0CBuf.template GetBufferByByte(L0C_TILE_SIZE * i); - l0CEventList[i] = i; - AscendC::SetFlag(l0CEventList[i]); - } - } else { - l0CTensorList[0] = resource.l0CBuf.template GetBufferByByte(0); - } - if constexpr (HAS_BIAS) { - uint32_t l1BiasOffset = l1BOffset + L1B_TILE_SIZE * L1B_STAGES; - l1BiasTensor = resource.l1Buf.template GetBufferByByte(l1BiasOffset); - l0BiasTensor = resource.btBuf.template GetBufferByByte(0); - AscendC::SetFlag(L1A_STAGES + L1B_STAGES); - AscendC::SetFlag(L0A_STAGES + L0B_STAGES); - } - - if constexpr (ENABLE_L1_RESIDENT) { - RestoreStatus(); - } - } - } - - /// Destructor - CATLASS_DEVICE - ~BlockMmadTla() - { - if ASCEND_IS_AIC { - if constexpr (USE_HF32_MODE) { - AscendC::SetHF32Mode(false); - } - if constexpr (ENABLE_UNIT_FLAG && tla::detail::isRowMajor::value) { - AscendC::SetMMLayoutTransform(false); - } - for (uint32_t i = 0; i < L1B_STAGES; i++) { - AscendC::WaitFlag(l1BEventList[i]); - } - for (uint32_t i = 0; i < L0A_STAGES; i++) { - AscendC::WaitFlag(l0AEventList[i]); - } - for (uint32_t i = 0; i < L0B_STAGES; i++) { - AscendC::WaitFlag(l0BEventList[i]); - } - if constexpr (!ENABLE_UNIT_FLAG) { - for (uint32_t i = 0; i < L0C_STAGES; i++) { - AscendC::WaitFlag(l0CEventList[i]); - } - } - if constexpr (HAS_BIAS) { - AscendC::WaitFlag(L1A_STAGES + L1B_STAGES); - AscendC::WaitFlag(L0A_STAGES + L0B_STAGES); - } - } - } - - /// Perform a block-scoped matrix multiply-accumulate - template - CATLASS_DEVICE void operator()(TensorA &tensorA, TensorB &tensorB, TensorC &tensorC, GemmCoord const &actualShape, - bool needLoadL1, TensorBias const &tensorBias = {}) - { - // Check L1TileShape - if constexpr (HAS_BIAS) { - static constexpr uint32_t BIAS_BUF_SIZE = L0_TILE_N * sizeof(ElementAccumulator); - static constexpr uint32_t L1_BIAS_SIZE = L1_TILE_N * sizeof(ElementBias); - static_assert(BIAS_BUF_SIZE <= ArchTag::BIAS_SIZE, - "BIAS_BUF_SIZE exceeding the BT space! Reduce L0_TILE_N"); - static_assert(L1A_TILE_SIZE * L1A_STAGES + L1B_TILE_SIZE * L1B_STAGES + L1_BIAS_SIZE <= ArchTag::L1_SIZE, - "L1TileShape exceeding the L1 space!"); - } - - using CopyGmToL1A = typename TileCopy_::template CopyGmToL1A; - using CopyGmToL1B = typename TileCopy_::template CopyGmToL1B; - using CopyL0CToGm = typename TileCopy_::template CopyL0CToGm; - CopyGmToL1A copyGmToL1A; - CopyGmToL1B copyGmToL1B; - CopyL0CToGm copyL0CToGm; - - // load first matrix A tile from GM to L1 - auto L1A_LAYOUT = tla::MakeLayout(tla::Int{}, actualShape.k()); - if (needLoadL1) { - AscendC::SetFlag(l1AEventList[0]); - AscendC::WaitFlag(l1AEventList[0]); - auto tensorL1A = tla::MakeTensor(l1ATensorList[0], L1A_LAYOUT, Arch::PositionL1{}); - auto tensorTileA = GetTileA(tensorA, 0, 0, actualShape.m(), actualShape.k()); - copyGmToL1A(tensorL1A, tensorTileA); - AscendC::SetFlag(l1AEventList[0]); - AscendC::WaitFlag(l1AEventList[0]); - } - - uint32_t mBlockActual = actualShape.m(); - uint32_t kBlockActual = actualShape.k(); - uint32_t nBlockActual = actualShape.n(); - - uint32_t mL1Actual = mBlockActual; - if constexpr (std::is_same_v) { - // Avoid using the gemv mode in mmad - if (mL1Actual == 1) { - mL1Actual = 16; - } - } - uint32_t nL1Actual = nBlockActual; - - auto layoutInL0C = tla::MakeLayoutL0C(mL1Actual, nL1Actual); - auto tensorL0C = tla::MakeTensor(l0CTensorList[l0CListId], layoutInL0C, Arch::PositionL0C{}); - auto tensorL0Bias = tla::MakeTensor(l0BiasTensor, L0BIAS_LAYOUT, Arch::PositionBias{}); - - // uint32_t kL1Actual = min(kBlockActual, L1_TILE_K); - uint32_t kTileCount = CeilDiv(kBlockActual); - uint32_t startTileIdx = 0; - if constexpr (ENABLE_SHUFFLE_K) { - startTileIdx = AscendC::GetBlockIdx(); - } - uint32_t firstTileIdx = startTileIdx % kTileCount; - uint32_t lastTileIdx = (startTileIdx + kTileCount - 1) % kTileCount; - uint32_t kL1Actual = - (firstTileIdx < kTileCount - 1) ? L1_TILE_K : (kBlockActual - firstTileIdx * L1_TILE_K); - - // load first matrix B tile from GM to L1 - AscendC::WaitFlag(l1BEventList[l1BListId]); - auto tensorL1B = tla::MakeTensor(l1BTensorList[l1BListId], L1B_LAYOUT, Arch::PositionL1{}); - auto tensorTileB = GetTile(tensorB, tla::MakeCoord(firstTileIdx * L1_TILE_K, 0), tla::MakeShape(kL1Actual, nBlockActual)); - if constexpr (ENABLE_L1_RESIDENT) { - if (lastAddrB[l1BListId] != tensorTileB.data().GetPhyAddr() || tla::get<0>(tensorTileB.coord()) != lastCoordB[l1BListId].row() || tla::get<1>(tensorTileB.coord()) != lastCoordB[l1BListId].column()) { - copyGmToL1B(tensorL1B, tensorTileB); - lastCoordB[l1BListId] = MatrixCoord{tla::get<0>(tensorTileB.coord()), tla::get<1>(tensorTileB.coord())}; - lastAddrB[l1BListId] = const_cast<__gm__ typename AscendC::GlobalTensor::PrimType *>( - tensorTileB.data().GetPhyAddr()); - } - } else { - copyGmToL1B(tensorL1B, tensorTileB); - } - AscendC::SetFlag(l1BEventList[l1BListId]); - - if constexpr (HAS_BIAS && !std::is_same_v) { - using CopyGmToL1Bias = typename TileCopy::template CopyGmToL1Bias; - CopyGmToL1Bias copyGmToL1Bias; - AscendC::WaitFlag(L1A_STAGES + L1B_STAGES); - auto l1Bias = l1BiasTensor.template ReinterpretCast(); - auto tensorL1Bias = tla::MakeTensor(l1Bias, L1BIAS_LAYOUT, Arch::PositionL1{}); - copyGmToL1Bias(tensorL1Bias, tensorBias); - AscendC::SetFlag(L1A_STAGES + L1B_STAGES); - } - - if constexpr (!ENABLE_UNIT_FLAG) { - AscendC::WaitFlag(l0CEventList[l0CListId]); - } - - uint32_t mL0Loop = CeilDiv(mL1Actual); - uint32_t nL0Loop = CeilDiv(nL1Actual); - - // main loop - for (uint32_t kL1Idx = 0; kL1Idx < kTileCount; kL1Idx++) { - uint32_t shuffleKIdx = (startTileIdx + kL1Idx) % kTileCount; - uint32_t l1BListIdNext = (l1BListId + 1 < L1B_STAGES) ? (l1BListId + 1) : 0; - uint32_t kL1ActualNext{0}; - // preload next tile from GM to L1 - if (kL1Idx < kTileCount - 1) { - uint32_t kL1IdxNext = (startTileIdx + kL1Idx + 1) % kTileCount; - kL1ActualNext = (kL1IdxNext < kTileCount - 1) ? L1_TILE_K : (kBlockActual - kL1IdxNext * L1_TILE_K); - - // Get L1 tensor for next stage - auto l1BTensor = l1BTensorList[l1BListIdNext]; - auto tensorL1B = tla::MakeTensor(l1BTensor, L1B_LAYOUT, Arch::PositionL1{}); - // Get GM tile for next stage - auto tensorTileB = GetTile(tensorB, tla::MakeCoord(kL1IdxNext * L1_TILE_K, 0), - tla::MakeShape(kL1ActualNext, nBlockActual)); - - // load next matrix B tile from GM to L1 - AscendC::WaitFlag(l1BEventList[l1BListIdNext]); - if constexpr (ENABLE_L1_RESIDENT) { - if (lastAddrB[l1BListIdNext] != tensorTileB.data().GetPhyAddr() || tla::get<0>(tensorTileB.coord()) != lastCoordB[l1BListIdNext].row() || tla::get<1>(tensorTileB.coord()) != lastCoordB[l1BListIdNext].column()) { - copyGmToL1B(tensorL1B, tensorTileB); - lastCoordB[l1BListIdNext] = - MatrixCoord{tla::get<0>(tensorTileB.coord()), tla::get<1>(tensorTileB.coord())}; - lastAddrB[l1BListIdNext] = - const_cast<__gm__ typename AscendC::GlobalTensor::PrimType *>( - tensorTileB.data().GetPhyAddr()); - } - } else { - copyGmToL1B(tensorL1B, tensorTileB); - } - AscendC::SetFlag(l1BEventList[l1BListIdNext]); - } - - // Get L1 tensor for current stage - auto l1ATensor = l1ATensorList[0]; - auto l1BTensor = l1BTensorList[l1BListId]; - auto tensorL1A = tla::MakeTensor(l1ATensor, L1A_LAYOUT, Arch::PositionL1{}); - tensorL1B = tla::MakeTensor(l1BTensor, L1B_LAYOUT, Arch::PositionL1{}); - // Get the loop nums on L0 - uint32_t kL0Loop = CeilDiv(kL1Actual); - - for (int mL0Idx = 0; mL0Idx < mL0Loop; mL0Idx++) { - uint32_t mL0Actual = (mL0Idx < mL0Loop - 1) ? L0_TILE_M : (mL1Actual - mL0Idx * L0_TILE_M); - - for (int kL0Idx = 0; kL0Idx < kL0Loop; kL0Idx++) { - uint32_t kL0Actual = (kL0Idx < kL0Loop - 1) ? L0_TILE_K : (kL1Actual - kL0Idx * L0_TILE_K); - - // Locate the current tile on L0A - auto l0ATile = l0ATensorList[l0AListId]; - auto layoutAInL0 = tla::MakeLayout(mL0Actual, kL0Actual); - auto tensorL0A = tla::MakeTensor(l0ATile, layoutAInL0, Arch::PositionL0A{}); - // Locate the current tile of matrix A on L1 - auto tensorTileL1A = GetTileA(tensorL1A, mL0Idx * L0_TILE_M, kL0Idx * L0_TILE_K + shuffleKIdx * L1_TILE_K, - mL0Actual, kL0Actual); - - AscendC::WaitFlag(l0AEventList[l0AListId]); - - // Load current tile from L1 to L0A - copyL1ToL0A(tensorL0A, tensorTileL1A); - - bool initC = ((kL1Idx == 0) && (kL0Idx == 0)); - for (int nL0Idx = 0; nL0Idx < nL0Loop; nL0Idx++) { - uint32_t nL0Actual = (nL0Idx < nL0Loop - 1) ? L0_TILE_N : (nL1Actual - nL0Idx * L0_TILE_N); - - // Locate the current tile on L0B - auto l0BTile = l0BTensorList[l0BListId]; - auto layoutBInL0 = tla::MakeLayout(kL0Actual, nL0Actual); - auto tensorL0B = tla::MakeTensor(l0BTile, layoutBInL0, Arch::PositionL0B{}); - // Locate the current tile of matrix B on L1 - auto tensorTileL1B = GetTile(tensorL1B, - tla::MakeCoord(kL0Idx * L0_TILE_K, nL0Idx * L0_TILE_N), - tla::MakeShape(kL0Actual, nL0Actual)); - - // Wait for mmad finished - AscendC::WaitFlag(l0BEventList[l0BListId]); - // If the current tile is the first one on the k&n axis, wait for loading matrix B from GM to L1 - if ((mL0Idx == 0) && (kL0Idx == 0) && (nL0Idx == 0)) { - AscendC::WaitFlag(l1BEventList[l1BListId]); - } - - // Load current tile from L1 to L0B - copyL1ToL0B(tensorL0B, tensorTileL1B); - - // If the current tile is the last one on the k&n axis, notify to load matrix B from GM to L1 - if ((mL0Idx == mL0Loop - 1) && (kL0Idx == kL0Loop - 1) && (nL0Idx == nL0Loop - 1)) { - AscendC::SetFlag(l1BEventList[l1BListId]); - } - - if constexpr (HAS_BIAS && !std::is_same_v) { - if (initC) { - if (nL0Idx == 0) { - AscendC::WaitFlag(L1A_STAGES + L1B_STAGES); - } - AscendC::WaitFlag(L0A_STAGES + L0B_STAGES); - auto l1Bias = l1BiasTensor.template ReinterpretCast(); - auto tensorL1Bias = tla::MakeTensor(l1Bias, L1BIAS_LAYOUT, Arch::PositionL1{}); - auto tensorTileL1Bias = GetTile(tensorL1Bias, - tla::MakeCoord(nL0Idx * L0_TILE_N), - tla::MakeShape(nL0Actual)); - // Load bias to l0 biastable - copyL1ToBT(tensorL0Bias, tensorTileL1Bias); - if (nL0Idx == nL0Loop - 1) { - AscendC::SetFlag(L1A_STAGES + L1B_STAGES); - } - } - } - - // Notify to do mmad - AscendC::SetFlag(l0CEventList[l0CListId]); - - // Locate the current tile on L0C - auto tensorTileL0C = GetTile(tensorL0C, - tla::MakeCoord(mL0Idx * L0_TILE_M, nL0Idx * L0_TILE_N), - tla::MakeShape(mL0Actual, nL0Actual)); - - // Compute the matrix multiplication on L0A and L0B and write the result to the accumulator - // Wait for loading L0B - AscendC::WaitFlag(l0CEventList[l0CListId]); - - // If the unit flag is enabled, the unit flag is set according to the calculation progress - uint8_t unitFlag = 0b00; - if constexpr (ENABLE_UNIT_FLAG) { - if ((kL1Idx == kTileCount - 1) && (mL0Idx == mL0Loop - 1) && - (kL0Idx == kL0Loop - 1) && (nL0Idx == nL0Loop - 1)) { - unitFlag = 0b11; - } else { - unitFlag = 0b10; - } - } - - if constexpr (HAS_BIAS && !std::is_same_v) { - if (initC) { - tileMmad(tensorTileL0C, tensorL0A, tensorL0B, tensorL0Bias, - mL0Actual, nL0Actual, kL0Actual, initC, unitFlag); - AscendC::SetFlag(L0A_STAGES + L0B_STAGES); - } else { - tileMmad(tensorTileL0C, tensorL0A, tensorL0B, - mL0Actual, nL0Actual, kL0Actual, initC, unitFlag); - } - } else { - tileMmad(tensorTileL0C, tensorL0A, tensorL0B, - mL0Actual, nL0Actual, kL0Actual, initC, unitFlag); - } - - // Notify to move the next L0B tile - AscendC::SetFlag(l0BEventList[l0BListId]); - l0BListId = (l0BListId + 1 < L0B_STAGES) ? (l0BListId + 1) : 0; - } - AscendC::SetFlag(l0AEventList[l0AListId]); - l0AListId = (l0AListId + 1 < L0A_STAGES) ? (l0AListId + 1) : 0; - } - } - l1BListId = l1BListIdNext; - kL1Actual = kL1ActualNext; - } - - // copy block out - if constexpr (!ENABLE_UNIT_FLAG) { - AscendC::SetFlag(l0CEventList[l0CListId]); - AscendC::WaitFlag(l0CEventList[l0CListId]); - copyL0CToGm(tensorC, tensorL0C); - AscendC::SetFlag(l0CEventList[l0CListId]); - l0CListId = (l0CListId + 1 < L0C_STAGES) ? (l0CListId + 1) : 0; - } else { - copyL0CToGm(tensorC, tensorL0C, 0b11); - } - } - -protected: - template - CATLASS_DEVICE auto GetTileA(TensorA &tensorA, uint32_t mIndex, uint32_t kIndex, uint32_t mSize, uint32_t kSize) - { - if constexpr (tla::detail::isVector::value) { - return GetTile(tensorA, tla::MakeCoord(kIndex), tla::MakeShape(kSize)); - } else { - return GetTile(tensorA, tla::MakeCoord(mIndex, kIndex), tla::MakeShape(mSize, kSize)); - } - } - - // Multi-stage tensors list - AscendC::LocalTensor l1ATensorList[L1A_STAGES]; - AscendC::LocalTensor l1BTensorList[L1B_STAGES]; - AscendC::LocalTensor l0ATensorList[L0A_STAGES]; - AscendC::LocalTensor l0BTensorList[L0B_STAGES]; - AscendC::LocalTensor l0CTensorList[L0C_STAGES]; - AscendC::LocalTensor l1BiasTensor; - AscendC::LocalTensor l0BiasTensor; - - // Multi-stage event id list - int32_t l1AEventList[L1A_STAGES]; - int32_t l1BEventList[L1B_STAGES]; - int32_t l0AEventList[L0A_STAGES]; - int32_t l0BEventList[L0B_STAGES]; - int32_t l0CEventList[L0C_STAGES]; - - __gm__ typename AscendC::GlobalTensor::PrimType *lastAddrA[L1A_STAGES]; - __gm__ typename AscendC::GlobalTensor::PrimType *lastAddrB[L1B_STAGES]; - MatrixCoord lastCoordA[L1A_STAGES]; - MatrixCoord lastCoordB[L1B_STAGES]; - - // The id of current stage - uint32_t l1AListId{0}; - uint32_t l1BListId{0}; - uint32_t l0AListId{0}; - uint32_t l0BListId{0}; - uint32_t l0CListId{0}; - - TileMmad tileMmad; - CopyL1ToL0A copyL1ToL0A; - CopyL1ToL0B copyL1ToL0B; - CopyL1ToBT copyL1ToBT; -}; - -} // namespace Catlass::Gemm::Block - -#endif // CATLASS_GEMM_BLOCK_BLOCK_MMAD_PINGPONG_FULL_LOADA_TLA_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_pingpong_per_group_per_block_tla.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_pingpong_per_group_per_block_tla.hpp deleted file mode 100644 index 56b51d3c41..0000000000 --- a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_pingpong_per_group_per_block_tla.hpp +++ /dev/null @@ -1,320 +0,0 @@ -/** - * Copyright (c) 2026 Huawei Technologies Co., Ltd. - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of - * CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ - -#ifndef CATLASS_GEMM_BLOCK_BLOCK_MMAD_PINGPONG_PER_GROUP_PER_BLOCK_TLA_HPP -#define CATLASS_GEMM_BLOCK_BLOCK_MMAD_PINGPONG_PER_GROUP_PER_BLOCK_TLA_HPP - -#include "catlass/kda_catlass.hpp" -#include "catlass/arch/kda_resource.hpp" -#include "catlass/kda_coord.hpp" -#include "catlass/kda_gemm_coord.hpp" -#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" -#include "catlass/gemm/kda_gemm_helper.hpp" -#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" -#include "catlass/gemm/tile/kda_tile_mmad.hpp" -#include "tla/kda_tla_layout.hpp" -#include "tla/kda_tensor.hpp" - -namespace Catlass::Gemm::Block { - -template < - class ArchTag_, - bool ENABLE_UNIT_FLAG_, - class L1TileShape_, - class L0TileShape_, - class ElementA_, - class ElementB_, - class ElementC_, - class ElementBias_, - class TileCopy_, - class TileMmad_> -struct BlockMmadTla< - MmadPingpongPertile, - L1TileShape_, - L0TileShape_, - ElementA_, - ElementB_, - ElementC_, - ElementBias_, - TileCopy_, - TileMmad_> { -public: - struct Arguments { - GM_ADDR aGmAddr{nullptr}; - GM_ADDR bGmAddr{nullptr}; - GM_ADDR cGmAddr{nullptr}; - GM_ADDR workspaceGmAddr{nullptr}; - }; - - using Params = Arguments; - - using DispatchPolicy = MmadPingpongPertile; - using ArchTag = typename DispatchPolicy::ArchTag; - - using TileCopy = TileCopy_; - using L1TileShape = L1TileShape_; - using L0TileShape = L0TileShape_; - using ElementA = ElementA_; - using LayoutA = typename TileCopy::LayoutA; - using ElementB = ElementB_; - using LayoutB = typename TileCopy::LayoutB; - using ElementC = ElementC_; - using LayoutC = typename TileCopy::LayoutC; - using ElementBias = ElementBias_; - - using TileMmad = TileMmad_; - - using CopyL1ToL0A = typename TileCopy::CopyL1ToL0A; - using CopyL1ToL0B = typename TileCopy::CopyL1ToL0B; - using CopyL1ToBT = typename TileCopy::CopyL1ToBT; - - using ElementAccumulator = typename TileCopy::ElementAccumulator; - - static constexpr bool HAS_BIAS = TileCopy::HAS_BIAS; - - using LayoutTagA = typename TileCopy::LayoutTagA; - using LayoutTagB = typename TileCopy::LayoutTagB; - using LayoutTagC = typename TileCopy::LayoutTagC; - using LayoutTagL1A = typename TileCopy::LayoutTagL1A; - using LayoutTagL1B = typename TileCopy::LayoutTagL1B; - using LayoutTagL0A = typename TileCopy::LayoutTagL0A; - using LayoutTagL0B = typename TileCopy::LayoutTagL0B; - - uint64_t m_{1}; - uint64_t n_{1}; - uint64_t k_{1}; - - constexpr static uint64_t STAGES = DispatchPolicy::STAGES; - - constexpr static uint64_t ML1_ = tla::get<0>(L1TileShape{}); - constexpr static uint64_t NL1_ = tla::get<1>(L1TileShape{}); - constexpr static uint64_t KL1_ = tla::get<2>(L1TileShape{}); - - constexpr static uint64_t ML0_ = tla::get<0>(L0TileShape{}); - constexpr static uint64_t NL0_ = tla::get<1>(L0TileShape{}); - constexpr static uint64_t KL0_ = tla::get<2>(L0TileShape{}); - - constexpr static uint64_t L1A_TILE_SIZE = ML1_ * KL1_ * sizeof(ElementA); - constexpr static uint64_t L1B_TILE_SIZE = NL1_ * KL1_ * sizeof(ElementB); - - constexpr static uint64_t L0A_TILE_SIZE = ML0_ * KL0_ * sizeof(ElementA); - constexpr static uint64_t L0B_TILE_SIZE = NL0_ * KL0_ * sizeof(ElementB); - constexpr static uint64_t L0C_TILE_SIZE = ML0_ * NL0_ * sizeof(ElementAccumulator); - - CATLASS_DEVICE - BlockMmadTla(const GemmCoord &shape) - { - uint32_t l1AOffset = 0; - uint32_t l1BOffset = l1AOffset + L1A_TILE_SIZE * STAGES; - // Init buffers - for (uint32_t i = 0; i < STAGES; i++) { - l1ATensorList_[i] = l1Buf_[l1AOffset + L1A_TILE_SIZE * i].template ReinterpretCast(); - l1BTensorList_[i] = l1Buf_[l1BOffset + L1B_TILE_SIZE * i].template ReinterpretCast(); - l0ATensorList_[i] = l0ABuf_[L0A_TILE_SIZE * i].template ReinterpretCast(); - l0BTensorList_[i] = l0BBuf_[L0B_TILE_SIZE * i].template ReinterpretCast(); - l0CTensorList_[i] = l0CBuf_[L0C_TILE_SIZE * i].template ReinterpretCast(); - - l1AEventList_[i] = i; - l1BEventList_[i] = i + STAGES; - l0AEventList_[i] = i; - l0BEventList_[i] = i + STAGES; - l0CEventList_[i] = i; - - AscendC::SetFlag(l1AEventList_[i]); - AscendC::SetFlag(l1BEventList_[i]); - AscendC::SetFlag(l0AEventList_[i]); - AscendC::SetFlag(l0BEventList_[i]); - AscendC::SetFlag(l0CEventList_[i]); - } - m_ = shape.m(); - n_ = shape.n(); - k_ = shape.k(); - } - - CATLASS_DEVICE - ~BlockMmadTla() - { - for (uint32_t i = 0; i < STAGES; i++) { - AscendC::WaitFlag(l1AEventList_[i]); - AscendC::WaitFlag(l1BEventList_[i]); - AscendC::WaitFlag(l0AEventList_[i]); - AscendC::WaitFlag(l0BEventList_[i]); - AscendC::WaitFlag(l0CEventList_[i]); - } - } - - CATLASS_DEVICE void WaitForVector(uint32_t crossPingPongID) - { - AscendC::CrossCoreWaitFlag(AIV_SYNC_AIC_FLAG + (crossPingPongID_)); - AscendC::CrossCoreWaitFlag(AIV_SYNC_AIC_FLAG + FLAG_ID_MAX + (crossPingPongID_)); - } - - CATLASS_DEVICE void NotifyVector(uint32_t crossPingPongID) - { - AscendC::CrossCoreSetFlag(AIC_SYNC_AIV_FLAG + (crossPingPongID_)); - AscendC::CrossCoreSetFlag(AIC_SYNC_AIV_FLAG + FLAG_ID_MAX + (crossPingPongID_)); - } - -public: - template - CATLASS_DEVICE void operator()( - DstTensor &cTensor, - SrcATensor &aTensor, - SrcBTensor &bTensor, - GemmCoord tileShape) - { - using CopyL0CToOut = typename TileCopy::template CopyL0CToDst; - CopyL0CToOut copyL0CToOut; - - uint32_t curML0 = tileShape.m(); - uint32_t curNL0 = tileShape.n(); - auto l0CLayout = tla::MakeLayoutL0C(curML0, curNL0); - - uint32_t kL1Loops = CeilDiv(k_, KL1_); - - // GM -> L1 - uint32_t kL1OffsetLength = 0; - for (uint32_t kL1Idx = 0; kL1Idx < kL1Loops; ++kL1Idx) { - uint32_t curKL1 = (kL1Idx + 1 == kL1Loops) ? (k_ - kL1OffsetLength) : KL1_; - - AscendC::WaitFlag(l1AEventList_[l1ListId_]); - auto l1ALayout = tla::MakeLayout(curML0, curKL1); - auto l1ATensor = tla::MakeTensor(l1ATensorList_[l1ListId_], l1ALayout, Arch::PositionL1{}); - - auto aTileTensor = GetTile( - aTensor, - tla::MakeCoord(0, kL1OffsetLength), - tla::MakeShape(curML0, curKL1)); - - using CopyGmToL1A = typename TileCopy::template CopyGmToL1A; - CopyGmToL1A copyGmToL1A; - copyGmToL1A(l1ATensor, aTileTensor); - - AscendC::SetFlag(l1AEventList_[l1ListId_]); - AscendC::WaitFlag(l1AEventList_[l1ListId_]); - - AscendC::WaitFlag(l1BEventList_[l1ListId_]); - - auto l1BLayout = tla::MakeLayout(curKL1, curNL0); - auto l1BTensor = tla::MakeTensor(l1BTensorList_[l1ListId_], l1BLayout, Arch::PositionL1{}); - auto bTileTensor = GetTile( - bTensor, - tla::MakeCoord(kL1OffsetLength, 0), - tla::MakeShape(curKL1, curNL0)); - - using CopyGmToL1B = typename TileCopy::template CopyGmToL1B; - CopyGmToL1B copyGmToL1B; - copyGmToL1B(l1BTensor, bTileTensor); - - AscendC::SetFlag(l1BEventList_[l1ListId_]); - AscendC::WaitFlag(l1BEventList_[l1ListId_]); - - kL1OffsetLength += curKL1; - - // L1 -> L0 - uint32_t kL0Loops = CeilDiv(curKL1, KL0_); - for (uint32_t kL0Idx = 0; kL0Idx < kL0Loops; ++kL0Idx) { - uint32_t curKL0 = (kL0Idx + 1 == kL0Loops) ? (curKL1 - kL0Idx * KL0_) : KL0_; - - AscendC::WaitFlag(l0AEventList_[l0ListId_]); - - auto l0ALayout = tla::MakeLayout(curML0, curKL0); - auto l0ATensor = tla::MakeTensor(l0ATensorList_[l1ListId_], l0ALayout, Arch::PositionL0A{}); - auto l1ATileTensor = GetTile( - l1ATensor, - tla::MakeCoord(0, kL0Idx * KL0_), - tla::MakeShape(curML0, curKL0)); - - copyL1ToL0A(l0ATensor, l1ATileTensor); - - AscendC::WaitFlag(l0BEventList_[l0ListId_]); - - auto l0BLayout = tla::MakeLayout(curKL0, curNL0); - auto l0BTensor = tla::MakeTensor(l0BTensorList_[l0ListId_], l0BLayout, Arch::PositionL0B{}); - auto l1BTileTensor = GetTile( - l1BTensor, - tla::MakeCoord(kL0Idx * KL0_, 0), - tla::MakeShape(curKL0, curNL0)); - - copyL1ToL0B(l0BTensor, l1BTileTensor); - - AscendC::WaitFlag(l0CEventList_[l0CListId_]); - auto l0CTensor = tla::MakeTensor(l0CTensorList_[l0CListId_], l0CLayout, Arch::PositionL0C{}); - - AscendC::SetFlag(l0CEventList_[l0CListId_]); - AscendC::WaitFlag(l0CEventList_[l0CListId_]); - - // Mmad - bool initC = (kL1Idx == 0) && (kL0Idx == 0); - curML0 = (curML0 == 1 ? M_ALIGN : curML0); - tileMmad(l0CTensor, l0ATensor, l0BTensor, curML0, curNL0, curKL0, true); - - AscendC::SetFlag(l0AEventList_[l0ListId_]); - AscendC::SetFlag(l0BEventList_[l0ListId_]); - l0ListId_ = (l0ListId_ + 1 < STAGES) ? (l0ListId_ + 1) : 0; - - AscendC::SetFlag(l0CEventList_[l0CListId_]); - AscendC::WaitFlag(l0CEventList_[l0CListId_]); - - WaitForVector(crossPingPongID_); - auto outTensor = tla::MakeTensor(crossPingPongID_ == 0 ? cTensor.data() : cTensor.data()[UB_OFFSET], - cTensor.layout(), Arch::PositionUB{}); - copyL0CToOut(outTensor, l0CTensor); - NotifyVector(crossPingPongID_); - - AscendC::SetFlag(l0CEventList_[l0CListId_]); - l0CListId_ = (l0CListId_ + 1 < STAGES) ? (l0CListId_ + 1) : 0; - crossPingPongID_ = (crossPingPongID_ + 1 < STAGES) ? (crossPingPongID_ + 1) : 0; - } - AscendC::SetFlag(l1AEventList_[l1ListId_]); - AscendC::SetFlag(l1BEventList_[l1ListId_]); - l1ListId_ = (l1ListId_ + 1 < STAGES) ? (l1ListId_ + 1) : 0; - } - } - -private: - constexpr static uint16_t M_ALIGN = 16; - constexpr static uint16_t AIC_SYNC_AIV_MODE_4 = 4; - constexpr static uint16_t AIV_SYNC_AIC_FLAG = 8; - constexpr static uint16_t AIC_SYNC_AIV_FLAG = 6; - constexpr static uint16_t FLAG_ID_MAX = 16; - constexpr static uint32_t UB_TWO_BANK_ELEMS_B32 = 128U; - constexpr static int64_t PER_BLOCK_SIZE = 128LL; - constexpr static uint64_t UB_OFFSET = UB_TWO_BANK_ELEMS_B32 * PER_BLOCK_SIZE; - - CopyL1ToL0A copyL1ToL0A; - CopyL1ToL0B copyL1ToL0B; - TileMmad tileMmad; - - AscendC::LocalTensor l1Buf_{AscendC::LocalTensor(AscendC::TPosition::A1, 0, ArchTag::L1_SIZE)}; - AscendC::LocalTensor l0ABuf_{AscendC::LocalTensor(AscendC::TPosition::A2, 0, ArchTag::L0A_SIZE)}; - AscendC::LocalTensor l0BBuf_{AscendC::LocalTensor(AscendC::TPosition::B2, 0, ArchTag::L0B_SIZE)}; - AscendC::LocalTensor l0CBuf_{AscendC::LocalTensor(AscendC::TPosition::CO1, 0, ArchTag::L0C_SIZE)}; - - AscendC::LocalTensor l1ATensorList_[STAGES]; - AscendC::LocalTensor l1BTensorList_[STAGES]; - AscendC::LocalTensor l0ATensorList_[STAGES]; - AscendC::LocalTensor l0BTensorList_[STAGES]; - AscendC::LocalTensor l0CTensorList_[STAGES]; - - int32_t l1AEventList_[STAGES]; - int32_t l1BEventList_[STAGES]; - int32_t l0AEventList_[STAGES]; - int32_t l0BEventList_[STAGES]; - int32_t l0CEventList_[STAGES]; - - uint32_t l1ListId_{0}; - uint32_t l0ListId_{0}; - uint32_t l0CListId_{0}; - uint32_t crossPingPongID_{0}; -}; -} // namespace Catlass::Gemm::Block - -#endif // CATLASS_GEMM_BLOCK_BLOCK_MMAD_PINGPONG_PER_GROUP_PER_BLOCK_TLA_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_pingpong_slice_k_with_prologue.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_pingpong_slice_k_with_prologue.hpp deleted file mode 100644 index ca4f77289f..0000000000 --- a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_pingpong_slice_k_with_prologue.hpp +++ /dev/null @@ -1,806 +0,0 @@ -/** - * Copyright (c) 2025 Huawei Technologies Co., Ltd. - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of - * CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ - -#ifndef CATLASS_GEMM_BLOCK_BLOCK_MMAD_PINGPONG_SLICE_K_WITH_PROLOGUE_HPP -#define CATLASS_GEMM_BLOCK_BLOCK_MMAD_PINGPONG_SLICE_K_WITH_PROLOGUE_HPP - -#include "catlass/kda_catlass.hpp" -#include "catlass/arch/kda_resource.hpp" -#include "catlass/kda_coord.hpp" -#include "catlass/kda_gemm_coord.hpp" -#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" -#include "catlass/gemm/kda_gemm_helper.hpp" -#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" -#include "catlass/gemm/tile/kda_tile_traits.hpp" - -namespace Catlass::Gemm::Block { - -template < - bool ENABLE_UNIT_FLAG_, - class L1TileShape_, - class L0TileShape_, - class AType_, - class BType_, - class CType_, - class BiasType_, - class TileCopy_, - class TileMmad_> -struct BlockMmad< - MmadAtlasA2PingpongSliceKWithPrologue, - L1TileShape_, - L0TileShape_, - AType_, - BType_, - CType_, - BiasType_, - TileCopy_, - TileMmad_> { -public: - // Type Aliases - using DispatchPolicy = MmadAtlasA2PingpongSliceKWithPrologue; - using ArchTag = typename DispatchPolicy::ArchTag; - using L1TileShape = L1TileShape_; - using L0TileShape = L0TileShape_; - using ElementA = typename AType_::Element; - using LayoutA = typename AType_::Layout; - using ElementB = typename BType_::Element; - using LayoutB = typename BType_::Layout; - using ElementC = typename CType_::Element; - using LayoutC = typename CType_::Layout; - using TileMmad = TileMmad_; - using CopyGmToL1A = typename TileCopy_::CopyGmToL1A; - using CopyGmToL1B = typename TileCopy_::CopyGmToL1B; - using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; - using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; - using CopyL0CToGm = typename TileCopy_::CopyL0CToGm; - - using PrologueA = typename TileCopy_::PrologueA; - using PrologueB = typename TileCopy_::PrologueB; - using PrologueC = typename TileCopy_::PrologueB; // NOTE: SHOULD BE Epilogue - - using ElementAccumulator = - typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; - using LayoutAInL1 = typename CopyL1ToL0A::LayoutSrc; - using LayoutBInL1 = typename CopyL1ToL0B::LayoutSrc; - using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst; - using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst; - using LayoutCInL0 = layout::zN; - - using L1AAlignHelper = Gemm::helper::L1AlignHelper; - using L1BAlignHelper = Gemm::helper::L1AlignHelper; - - static constexpr bool ENABLE_UNIT_FLAG = DispatchPolicy::ENABLE_UNIT_FLAG; - static constexpr uint32_t STAGES = DispatchPolicy::STAGES; - static constexpr uint32_t L1A_SIZE = L1TileShape::M * L1TileShape::K * sizeof(ElementA); - static constexpr uint32_t L1B_SIZE = L1TileShape::N * L1TileShape::K * sizeof(ElementB); - static constexpr uint32_t L0A_SIZE = ArchTag::L0A_SIZE; - static constexpr uint32_t L0B_SIZE = ArchTag::L0B_SIZE; - static constexpr uint32_t L0A_PINGPONG_BUF_SIZE = L0A_SIZE / STAGES; - static constexpr uint32_t L0B_PINGPONG_BUF_SIZE = L0B_SIZE / STAGES; - - static constexpr bool HAS_PROLOGUE_A = !std::is_same_v; - static constexpr bool HAS_PROLOGUE_B = !std::is_same_v; - - // Check L1TileShape - static_assert((L1A_SIZE * STAGES + L1B_SIZE * STAGES) <= ArchTag::L1_SIZE, "L1TileShape exceeding the L1 space!"); - - // Check L0TileShape - static constexpr uint32_t L0A_TILE_SIZE = L0TileShape::M * L0TileShape::K * sizeof(ElementA); - static constexpr uint32_t L0B_TILE_SIZE = L0TileShape::K * L0TileShape::N * sizeof(ElementB); - static_assert((L0A_TILE_SIZE * STAGES) <= L0A_SIZE, "L0TileShape exceeding the L0A space!"); - static_assert((L0B_TILE_SIZE * STAGES) <= L0B_SIZE, "L0TileShape exceeding the L0B space!"); - - // 32B (256b) aligned - static_assert(Gemm::helper::TileShapeAlignChecker::_ALIGN == 256, - "Tile shape must be 32B aligned."); - - struct Params { - typename Tile::PrologueTraits::Params prologueAParams{}; - typename Tile::PrologueTraits::Params prologueBParams{}; - uint32_t mScalar; - uint32_t nScalar; - uint32_t splitkLength; - }; - - /// Construct - CATLASS_DEVICE - BlockMmad(Arch::Resource &resource, Params const ¶ms_ = {}, - uint32_t l1BufAddrStart = 0) - : params(params_) - { - // PREVIOUSLY: - // Initilize PrologueA & PrologueB instantly - // prologueA(resource, params_.prologueAParams), prologueB(resource, params_.prologueBParams) - uint32_t l1AOffset = l1BufAddrStart; - uint32_t l1BOffset = l1BufAddrStart + L1A_SIZE * STAGES; - // Init set/wait flags - flag0[0].id = flagID0; - flag0[1].id = flagID1; - flag1[0].id = flagID2; - flag1[1].id = flagID3; - - // AIC Type - if constexpr (g_coreType == AscendC::AIC) { - // Init buffers - for (uint32_t i = 0; i < STAGES; i++) { - // Assign L1/L0A/L0B space for each stages - l1ATensorList[i] = resource.l1Buf.template GetBufferByByte(l1AOffset + L1A_SIZE * i); - l1BTensorList[i] = resource.l1Buf.template GetBufferByByte(l1BOffset + L1B_SIZE * i); - l0ATensorList[i] = resource.l0ABuf.template GetBufferByByte(L0A_PINGPONG_BUF_SIZE * i); - l0BTensorList[i] = resource.l0BBuf.template GetBufferByByte(L0B_PINGPONG_BUF_SIZE * i); - - // Assign event ID for each stages - l1AEventList[i] = i; - l1BEventList[i] = i + STAGES; - l0AEventList[i] = i; - l0BEventList[i] = i + STAGES; - - // The event id that needs to be set before the loop - AscendC::SetFlag(l1AEventList[i]); - AscendC::SetFlag(l1BEventList[i]); - AscendC::SetFlag(l0AEventList[i]); - AscendC::SetFlag(l0BEventList[i]); - } - l0CTensor = resource.l0CBuf.template GetBufferByByte(0); - AscendC::SetFlag(EVENT_ID0); - - // CV 同步 - Catlass::Arch::CrossCoreSetFlag<0x2, PIPE_FIX>(flag0[0]); - Catlass::Arch::CrossCoreSetFlag<0x2, PIPE_FIX>(flag0[1]); - } - // AIV Coretype - if constexpr (g_coreType == AscendC::AIV) { - AscendC::SetFlag(EVENT_ID0); - AscendC::SetFlag(EVENT_ID1); - AscendC::SetFlag(EVENT_ID2); - AscendC::SetFlag(EVENT_ID3); - AscendC::SetFlag(EVENT_ID0); - AscendC::SetFlag(EVENT_ID1); - AscendC::SetFlag(EVENT_ID2); - AscendC::SetFlag(EVENT_ID3); - } - } - - /// Destructor - CATLASS_DEVICE - ~BlockMmad() - { - if constexpr (g_coreType == AscendC::AIC) { - for (uint32_t i = 0; i < STAGES; i++) { - AscendC::WaitFlag(l1AEventList[i]); - AscendC::WaitFlag(l1BEventList[i]); - AscendC::WaitFlag(l0AEventList[i]); - AscendC::WaitFlag(l0BEventList[i]); - } - AscendC::WaitFlag(EVENT_ID0); - } - - if constexpr (g_coreType == AscendC::AIV) { - AscendC::WaitFlag(EVENT_ID0); - AscendC::WaitFlag(EVENT_ID1); - AscendC::WaitFlag(EVENT_ID2); - AscendC::WaitFlag(EVENT_ID3); - AscendC::WaitFlag(EVENT_ID0); - AscendC::WaitFlag(EVENT_ID1); - AscendC::WaitFlag(EVENT_ID2); - AscendC::WaitFlag(EVENT_ID3); - - Catlass::Arch::CrossCoreWaitFlag(flag0[0]); - Catlass::Arch::CrossCoreWaitFlag(flag0[1]); - } - } - - /// Perform a block-scoped matrix multiply-accumulate - CATLASS_DEVICE - void operator()( - AscendC::GlobalTensor const &gmA, - AscendC::GlobalTensor const &gmB, - AscendC::GlobalTensor const &gmC, - GemmCoord actualBlockShape, - GemmCoord problemShape) - { - MNScalerMatmulImpl(gmA, gmB, gmC, {}, actualBlockShape, problemShape); - } - - /// Prologue - // template - CATLASS_DEVICE - void Prologue( - AscendC::GlobalTensor const &gmDstA, - AscendC::GlobalTensor const &gmSrcA, typename PrologueA::LayoutSrc const &layoutSrcA, - AscendC::GlobalTensor const &gmDstB, - AscendC::GlobalTensor const &gmSrcB, typename PrologueB::LayoutSrc const &layoutSrcB, - AscendC::GlobalTensor const &gmDstC, - AscendC::GlobalTensor const &gmSrcC, LayoutC const &layoutSrcC, - GemmCoord const &blockCoord, GemmCoord const &nextBlockCoord, - GemmCoord const &actualBlockShape, GemmCoord const &nextActualBlockShape, - GemmCoord const &problemShape, - bool isFirstBlock, bool hasNextBlock, - uint32_t &bufferIndex) - { - PrologueImpl( - gmDstA, gmSrcA, layoutSrcA, - gmDstB, gmSrcB, layoutSrcB, - gmDstC, gmSrcC, layoutSrcC, - blockCoord, nextBlockCoord, - actualBlockShape, nextActualBlockShape, problemShape, - isFirstBlock, hasNextBlock, bufferIndex); - } - -protected: - CATLASS_DEVICE - void PrologueImpl( - AscendC::GlobalTensor const &gmDstA, - AscendC::GlobalTensor const &gmSrcA, - typename PrologueA::LayoutSrc const &layoutSrcA, - AscendC::GlobalTensor const &gmDstB, - AscendC::GlobalTensor const &gmSrcB, - typename PrologueB::LayoutSrc const &layoutSrcB, - AscendC::GlobalTensor const &gmDstC, - AscendC::GlobalTensor const &gmSrcC, - LayoutC const &layoutSrcC, - GemmCoord blockCoord, - GemmCoord nextBlockCoord, - GemmCoord actualBlockShape, - GemmCoord nextActualBlockShape, - GemmCoord problemShape, - bool const isFirstBlock, - bool const hasNextBlock, - uint32_t &bufferIndex) - { - // 当前任务块 Current block - // 计算基础偏移量 - // layoutSrcA: params.layoutA, layoutSrcB: params.layoutB, layoutSrcC: params.layoutC - int64_t gmOffsetA, gmOffsetB, gmOffsetC, gmOffsetWA, gmOffsetWB, gmOffsetWC; - int64_t gmOffsetNextA, gmOffsetNextB, gmOffsetWADelta, gmOffsetWBDelta; - - gmOffsetA = layoutSrcA.GetOffset(MakeCoord(blockCoord.m() * (L1TileShape::M * params.mScalar), 0U)); - gmOffsetNextA = layoutSrcA.GetOffset(MakeCoord(nextBlockCoord.m() * (L1TileShape::M * params.mScalar), 0U)); - gmOffsetWA = (AscendC::GetBlockIdx() / AIVPERCORE) * - (params.mScalar * L1TileShape::M) * params.splitkLength * STAGES; - gmOffsetWADelta = (params.mScalar * L1TileShape::M) * params.splitkLength; - - gmOffsetB = layoutSrcB.GetOffset(MakeCoord(0U, blockCoord.n() * (L1TileShape::N * params.nScalar))); - gmOffsetNextB = layoutSrcB.GetOffset(MakeCoord(0U, nextBlockCoord.n() * (L1TileShape::N * params.nScalar))); - gmOffsetWB = (AscendC::GetBlockIdx() / AIVPERCORE) * - (params.nScalar * L1TileShape::N) * params.splitkLength * STAGES; - gmOffsetWBDelta = params.splitkLength * (params.nScalar * L1TileShape::N); - - MatrixCoord offsetC{blockCoord.m() * (L1TileShape::M * params.mScalar), - blockCoord.n() * (L1TileShape::N * params.nScalar)}; - gmOffsetC = layoutSrcC.GetOffset(offsetC); - gmOffsetWC = (AscendC::GetBlockIdx() / AIVPERCORE) * - (params.mScalar * L1TileShape::M) * (params.nScalar * L1TileShape::N); - - uint32_t srcAStride = layoutSrcA.stride( - std::is_same_v ? 0 : 1); - uint32_t srcBStride = layoutSrcB.stride( - std::is_same_v ? 0 : 1); - uint32_t mShape_ = actualBlockShape.m(); - uint32_t nShape_ = actualBlockShape.n(); - - // Begin resource - Arch::Resource resource; - - // SPLIT-K MainLoop - uint32_t kLoop = (problemShape.k() + params.splitkLength - 1) / params.splitkLength; - uint32_t kResidueLenght = problemShape.k() % params.splitkLength; - for (uint32_t ldk = 0; ldk < kLoop; ldk++) { - uint32_t kActual = (problemShape.k() < (ldk + 1) * params.splitkLength) ? kResidueLenght : params.splitkLength; - uint32_t kActual_ = kActual; - - if (ldk == 0 && isFirstBlock) { // 首块 - Catlass::Arch::CrossCoreWaitFlag(flag0[crossCoreBufferIndexAIV]); - - uint32_t kActualAligned_ = RoundUp<256, uint32_t>(kActual_); - LayoutA layoutWA(actualBlockShape.m(), kActual_, srcAStride); - LayoutA layoutDstA(actualBlockShape.m(), kActualAligned_); - LayoutB layoutWB(kActual_, actualBlockShape.n(), srcBStride); - LayoutB layoutDstB(kActualAligned_, actualBlockShape.n()); - - PrologueA prologueA(resource, params.prologueAParams); - PrologueB prologueB(resource, params.prologueBParams); - prologueA( - gmDstA[gmOffsetWA + crossCoreBufferIndexAIV * gmOffsetWADelta], - layoutDstA, - gmSrcA[gmOffsetA], - layoutWA, - bufferIndex); - prologueB( - gmDstB[gmOffsetWB + crossCoreBufferIndexAIV * gmOffsetWBDelta], - layoutDstB, - gmSrcB[gmOffsetB], - layoutWB, - bufferIndex); - } - if (ldk < kLoop - 1) { // 后续块 - kActual_ = (problemShape.k() < (ldk + 2) * params.splitkLength) ? kResidueLenght : params.splitkLength; - - Catlass::Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(flag1[crossCoreBufferIndexAIV]); - Catlass::Arch::CrossCoreWaitFlag(flag0[1 - crossCoreBufferIndexAIV]); - gmOffsetA += layoutSrcA.GetOffset(MakeCoord(0U, kActual)); - gmOffsetB += layoutSrcB.GetOffset(MakeCoord(kActual, 0U)); - - uint32_t kActualAligned_ = RoundUp<256, uint32_t>(kActual_); - LayoutA layoutWA(actualBlockShape.m(), kActual_, srcAStride); - LayoutA layoutDstA(actualBlockShape.m(), kActualAligned_); - LayoutB layoutWB(kActual_, actualBlockShape.n(), srcBStride); - LayoutB layoutDstB(kActualAligned_, actualBlockShape.n()); - - PrologueA prologueA(resource, params.prologueAParams); - PrologueB prologueB(resource, params.prologueBParams); - prologueA( - gmDstA[gmOffsetWA + (1 - crossCoreBufferIndexAIV) * gmOffsetWADelta], - layoutDstA, - gmSrcA[gmOffsetA], - layoutWA, - bufferIndex); - prologueB( - gmDstB[gmOffsetWB + (1 - crossCoreBufferIndexAIV) * gmOffsetWBDelta], - layoutDstB, - gmSrcB[gmOffsetB], - layoutWB, - bufferIndex); - } - if ((ldk == kLoop - 1) && hasNextBlock) { // SliceK尾部块 - kActual_ = (problemShape.k() < params.splitkLength) ? kResidueLenght : params.splitkLength; - - Catlass::Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(flag1[crossCoreBufferIndexAIV]); - Catlass::Arch::CrossCoreWaitFlag(flag0[1 - crossCoreBufferIndexAIV]); - - uint32_t kActualAligned_ = RoundUp<256, uint32_t>(kActual_); - LayoutA layoutWA(nextActualBlockShape.m(), kActual_, srcAStride); - LayoutA layoutDstA(nextActualBlockShape.m(), kActualAligned_); - LayoutB layoutWB(kActual_, nextActualBlockShape.n(), srcBStride); - LayoutB layoutDstB(kActualAligned_, nextActualBlockShape.n()); - - PrologueA prologueA(resource, params.prologueAParams); - PrologueB prologueB(resource, params.prologueBParams); - prologueA( - gmDstA[gmOffsetWA + (1 - crossCoreBufferIndexAIV) * gmOffsetWADelta], - layoutDstA, - gmSrcA[gmOffsetNextA], - layoutWA, - bufferIndex); - prologueB( - gmDstB[gmOffsetWB + (1 - crossCoreBufferIndexAIV) * gmOffsetWBDelta], - layoutDstB, - gmSrcB[gmOffsetNextB], - layoutWB, - bufferIndex); - } - - if (ldk == kLoop - 1) { // 尾部块 - if (!hasNextBlock) { - Catlass::Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(flag1[crossCoreBufferIndexAIV]); - } - Catlass::Arch::CrossCoreWaitFlag(flag4); - - Catlass::layout::RowMajor layoutBlockC( - actualBlockShape.m(), actualBlockShape.n(), params.nScalar * L1TileShape::N); - Catlass::layout::RowMajor layoutBlockDstC( - problemShape.m(), problemShape.n()); - - PrologueC epCast; - epCast.EpCastFp32ToFp16( - gmDstC[gmOffsetC], - layoutBlockDstC, - gmSrcC[gmOffsetWC], - layoutBlockC); - } - crossCoreBufferIndexAIV = 1 - crossCoreBufferIndexAIV; - } // end: for-loop (ldk) - } - - CATLASS_DEVICE - void MNScalerMatmulImpl( - AscendC::GlobalTensor const &gmWA, - AscendC::GlobalTensor const &gmWB, - AscendC::GlobalTensor const &gmWC, - [[maybe_unused]] GemmCoord blockCoord, - GemmCoord actualBlockShape, - GemmCoord problemShape) - { - constexpr bool isLayoutAColumnMajor = std::is_same_v; - constexpr bool isLayoutBColumnMajor = std::is_same_v; - - // uint32_t splitKLen = params.splitkLength; - uint32_t kLoop = CeilDiv(problemShape.k(), params.splitkLength); - uint32_t mLoop = CeilDiv(actualBlockShape.m()); - uint32_t nLoop = CeilDiv(actualBlockShape.n()); - - uint32_t nextMIdx, nextNIdx; - // SPLIT-K MainLoop - for (uint32_t ldk = 0; ldk < kLoop; ldk++) { // K-loop - bool isFirstKSlice = (ldk == 0); - uint32_t kActual = (problemShape.k() < (ldk + 1) * params.splitkLength) ? - problemShape.k() % params.splitkLength : - params.splitkLength; - uint32_t kActualAligned = RoundUp<256>(kActual); - - uint32_t srcStrideA = isLayoutAColumnMajor ? actualBlockShape.m() : kActualAligned; - uint32_t srcStrideB = isLayoutBColumnMajor ? kActualAligned : actualBlockShape.n(); - - // CV 同步 - Catlass::Arch::CrossCoreWaitFlag(flag1[crossCoreBufferIndexAIC]); - - // M-N-loop - for (uint32_t mIdx = 0; mIdx < mLoop; mIdx++) { // M-loop - for (uint32_t nIdx = 0; nIdx < nLoop; nIdx++) { // N-loop - bool isFirstBlock = (mIdx == 0 && nIdx == 0); - bool hasNextBlock = !(mIdx == mLoop - 1 && nIdx == nLoop - 1); - - // 当前偏移及实际尺寸计算(No padding) - int64_t gmOffsetWA_ = AscendC::GetBlockIdx() * (L1TileShape::M * params.mScalar) * - params.splitkLength * STAGES; - int64_t gmOffsetWB_ = AscendC::GetBlockIdx() * (L1TileShape::N * params.nScalar) * - params.splitkLength * STAGES; - gmOffsetWA_ += crossCoreBufferIndexAIC ? params.splitkLength * - (L1TileShape::M * params.mScalar) : - 0; - gmOffsetWB_ += crossCoreBufferIndexAIC ? params.splitkLength * - (L1TileShape::N * params.nScalar) : - 0; - int64_t gmOffsetWA = isLayoutAColumnMajor ? gmOffsetWA_ + mIdx * L1TileShape::M : - gmOffsetWA_ + mIdx * L1TileShape::M * kActualAligned; - int64_t gmOffsetWB = isLayoutBColumnMajor ? gmOffsetWB_ + nIdx * L1TileShape::N * - kActualAligned : - gmOffsetWB_ + nIdx * L1TileShape::N; - int64_t gmOffsetWC = AscendC::GetBlockIdx() * (L1TileShape::M * params.mScalar) * - (L1TileShape::N * params.nScalar) + - mIdx * L1TileShape::M * - (L1TileShape::N * params.nScalar) + - nIdx * L1TileShape::N; - - uint32_t mActual = (mIdx == mLoop - 1 && actualBlockShape.m() % L1TileShape::M != 0) ? - actualBlockShape.m() % L1TileShape::M : - L1TileShape::M; - uint32_t nActual = (nIdx == nLoop - 1 && actualBlockShape.n() % L1TileShape::N != 0) ? - actualBlockShape.n() % L1TileShape::N : - L1TileShape::N; - LayoutA layoutWA(mActual, kActual, srcStrideA); - LayoutB layoutWB(kActual, nActual, srcStrideB); - LayoutC layoutWC(mActual, nActual, params.nScalar * L1TileShape::N); - GemmCoord actualSmallBlockShape(mActual, nActual, kActual); - - // Next实际尺寸(No padding) - int64_t gmOffsetWANext, gmOffsetWBNext, gmOffsetWCNext; - LayoutA layoutWANext; - LayoutB layoutWBNext; - LayoutC layoutWCNext; - GemmCoord nextSmallBlockShape; - if (hasNextBlock) { - uint32_t nextNIdx = (nIdx + 1) % nLoop; - uint32_t nextMIdx = mIdx + static_cast(nextNIdx == 0); - - gmOffsetWANext = isLayoutAColumnMajor ? gmOffsetWA_ + nextMIdx * L1TileShape::M : - gmOffsetWA_ + nextMIdx * L1TileShape::M * kActualAligned; - gmOffsetWBNext = isLayoutBColumnMajor ? gmOffsetWB_ + nextNIdx * L1TileShape::N * - kActualAligned : - gmOffsetWB_ + nextNIdx * L1TileShape::N; - gmOffsetWCNext = AscendC::GetBlockIdx() * (L1TileShape::M * params.mScalar) * - (L1TileShape::N * params.nScalar) + - nextMIdx * L1TileShape::M * - (L1TileShape::N * params.nScalar) + - nextNIdx * L1TileShape::N; - - uint32_t mActualNext = (nextMIdx == mLoop - 1 && actualBlockShape.m() % L1TileShape::M != 0) ? - actualBlockShape.m() % L1TileShape::M : - L1TileShape::M; - uint32_t nActualNext = (nextNIdx == nLoop - 1 && actualBlockShape.n() % L1TileShape::N != 0) ? - actualBlockShape.n() % L1TileShape::N : - L1TileShape::N; - layoutWANext = LayoutA(mActualNext, kActual, srcStrideA); - layoutWBNext = LayoutB(kActual, nActualNext, srcStrideB); - layoutWCNext = LayoutC(mActualNext, nActualNext, params.nScalar * L1TileShape::N); - nextSmallBlockShape = GemmCoord(mActualNext, nActualNext, kActual); - } - - // 完成一个128 * 256的小结果矩阵基本块的运算 - smallBlockMmadImpl(gmWA[gmOffsetWA], - layoutWA, - gmWB[gmOffsetWB], - layoutWB, - gmWC[gmOffsetWC], - layoutWC, - gmWA[gmOffsetWANext], - layoutWANext, - gmWB[gmOffsetWBNext], - layoutWBNext, - actualSmallBlockShape, - nextSmallBlockShape, - isFirstKSlice, - isFirstBlock, - hasNextBlock); - } // End of 'N-loop' - } // End of 'M-loop' - - Catlass::Arch::CrossCoreSetFlag<0x2, PIPE_FIX>(flag0[crossCoreBufferIndexAIC]); - crossCoreBufferIndexAIC = 1 - crossCoreBufferIndexAIC; - if (ldk == kLoop - 1) { - // cast 256 * 512的fp32大结果基本块为fp16 - Catlass::Arch::CrossCoreSetFlag<0x2, PIPE_FIX>(flag4); - } - } // End of 'K-loop' - } - - CATLASS_DEVICE - void smallBlockMmadImpl( - AscendC::GlobalTensor const &gmWA, LayoutA const &layoutA, - AscendC::GlobalTensor const &gmWB, LayoutB const &layoutB, - AscendC::GlobalTensor const &gmC, LayoutC const &layoutC, - AscendC::GlobalTensor const &gmNextWA, LayoutA const &layoutNextA, - AscendC::GlobalTensor const &gmNextWB, LayoutB const &layoutNextB, - GemmCoord const &actualShape, GemmCoord const &nextActualShape, - bool isFirstKSlice, bool isFirstBlock, bool hasNextBlock) - { - uint32_t mRound = RoundUp(actualShape.m()); - uint32_t nRound = RoundUp(actualShape.n()); - - auto layoutAInL1 = LayoutAInL1::template MakeLayout(L1TileShape::M, L1TileShape::K); - auto layoutBInL1 = LayoutBInL1::template MakeLayout(L1TileShape::K, L1TileShape::N); - auto layoutInL0C = LayoutCInL0::MakeLayoutInL0C(MakeCoord(mRound, nRound)); - - uint32_t kActual = min(actualShape.k(), L1TileShape::K); - - // load first matrix A tile from GM to L1 - if (isFirstBlock) { - AscendC::WaitFlag(l1AEventList[l1ListId]); - auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShape.m(), kActual)); - copyGmToL1A(l1ATensorList[l1ListId], gmWA, layoutAInL1, layoutTileA); - AscendC::SetFlag(l1AEventList[l1ListId]); - - // load first matrix B tile from GM to L1 - AscendC::WaitFlag(l1BEventList[l1ListId]); - auto layoutTileB = layoutB.GetTileLayout(MakeCoord(kActual, actualShape.n())); - copyGmToL1B(l1BTensorList[l1ListId], gmWB, layoutBInL1, layoutTileB); - AscendC::SetFlag(l1BEventList[l1ListId]); - } - - if constexpr (!ENABLE_UNIT_FLAG) { - AscendC::WaitFlag(EVENT_ID0); - } - - uint32_t mPartLoop = CeilDiv(mRound); - uint32_t nPartLoop = CeilDiv(nRound); - - // main loop - uint32_t kTileCount = CeilDiv(actualShape.k()); - for (uint32_t kLoopIdx = 0; kLoopIdx < kTileCount; kLoopIdx++) { - uint32_t l1ListIdNext = (l1ListId + 1 < STAGES) ? (l1ListId + 1) : 0; - uint32_t kActualNext{0}; - // preload next tile from GM to L1 - if (kLoopIdx < kTileCount - 1) { - uint32_t kLoopIdxNext = kLoopIdx + 1; - kActualNext = (kLoopIdxNext < kTileCount - 1) ? - L1TileShape::K : - (actualShape.k() - kLoopIdxNext * L1TileShape::K); - - // Get L1 tensor for next stage - auto l1ATensor = l1ATensorList[l1ListIdNext]; - auto l1BTensor = l1BTensorList[l1ListIdNext]; - // Get GM tile for next stage - MatrixCoord gmTileAOffset{0, kLoopIdxNext * L1TileShape::K}; - MatrixCoord gmTileBOffset{kLoopIdxNext * L1TileShape::K, 0}; - auto gmTileA = gmWA[layoutA.GetOffset(gmTileAOffset)]; - auto gmTileB = gmWB[layoutB.GetOffset(gmTileBOffset)]; - - // load next matrix A tile from GM to L1 - AscendC::WaitFlag(l1AEventList[l1ListIdNext]); - auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShape.m(), kActualNext)); - copyGmToL1A(l1ATensor, gmTileA, layoutAInL1, layoutTileA); - AscendC::SetFlag(l1AEventList[l1ListIdNext]); - - // load next matrix B tile from GM to L1 - AscendC::WaitFlag(l1BEventList[l1ListIdNext]); - auto layoutTileB = layoutB.GetTileLayout(MakeCoord(kActualNext, actualShape.n())); - copyGmToL1B(l1BTensor, gmTileB, layoutBInL1, layoutTileB); - AscendC::SetFlag(l1BEventList[l1ListIdNext]); - } else if (hasNextBlock) { - uint32_t kLoopIdxNext = 0; - kActualNext = (kLoopIdxNext < kTileCount - 1) ? L1TileShape::K : (nextActualShape.k()); - - // Get L1 tensor for next stage - auto l1ATensor = l1ATensorList[l1ListIdNext]; - auto l1BTensor = l1BTensorList[l1ListIdNext]; - // Get GM tile for next stage - MatrixCoord gmTileAOffset{0, kLoopIdxNext * L1TileShape::K}; - MatrixCoord gmTileBOffset{kLoopIdxNext * L1TileShape::K, 0}; - auto gmTileA = gmNextWA[layoutNextA.GetOffset(gmTileAOffset)]; - auto gmTileB = gmNextWB[layoutNextB.GetOffset(gmTileBOffset)]; - - // load next matrix A tile from GM to L1 - AscendC::WaitFlag(l1AEventList[l1ListIdNext]); - auto layoutTileA = layoutNextA.GetTileLayout(MakeCoord(nextActualShape.m(), kActualNext)); - copyGmToL1A(l1ATensor, gmTileA, layoutAInL1, layoutTileA); - AscendC::SetFlag(l1AEventList[l1ListIdNext]); - - // load next matrix B tile from GM to L1 - AscendC::WaitFlag(l1BEventList[l1ListIdNext]); - auto layoutTileB = layoutNextB.GetTileLayout(MakeCoord(kActualNext, nextActualShape.n())); - copyGmToL1B(l1BTensor, gmTileB, layoutBInL1, layoutTileB); - AscendC::SetFlag(l1BEventList[l1ListIdNext]); - } - - // Get L1 tensor for current stage - auto l1ATensor = l1ATensorList[l1ListId]; - auto l1BTensor = l1BTensorList[l1ListId]; - - // Get the loop nums on L0 - uint32_t kPartLoop = CeilDiv(kActual); - - for (int mPartIdx = 0; mPartIdx < mPartLoop; mPartIdx++) { - uint32_t mPartActual = (mPartIdx < mPartLoop - 1) ? - L0TileShape::M : - (mRound - mPartIdx * L0TileShape::M); - - for (int kPartIdx = 0; kPartIdx < kPartLoop; kPartIdx++) { - uint32_t kPartActual = (kPartIdx < kPartLoop - 1) ? - L0TileShape::K : - (kActual - kPartIdx * L0TileShape::K); - - // Locate the current tile on L0A - auto l0ATile = l0ATensorList[l0AListId]; - LayoutAInL0 layoutAInL0 = LayoutAInL0::template MakeLayout(mPartActual, kPartActual); - // Locate the current tile of matrix A on L1 - MatrixCoord l1AOffset{mPartIdx * L0TileShape::M, kPartIdx * L0TileShape::K}; - auto l1ATile = l1ATensor[layoutAInL1.GetOffset(l1AOffset)]; - - AscendC::WaitFlag(l0AEventList[l0AListId]); - if ((mPartIdx == 0) && (kPartIdx == 0)) { - AscendC::WaitFlag(l1AEventList[l1ListId]); - } - - // Load current tile from L1 to L0A - copyL1ToL0A(l0ATile, l1ATile, layoutAInL0, layoutAInL1); - - if ((mPartIdx == mPartLoop - 1) && (kPartIdx == kPartLoop - 1)) { - AscendC::SetFlag(l1AEventList[l1ListId]); - } - - for (int nPartIdx = 0; nPartIdx < nPartLoop; nPartIdx++) { - uint32_t nPartActual = (nPartIdx < nPartLoop - 1) ? - L0TileShape::N : - (nRound - nPartIdx * L0TileShape::N); - - // Locate the current tile on L0B - auto l0BTile = l0BTensorList[l0BListId]; - LayoutBInL0 layoutBInL0 = LayoutBInL0::template MakeLayout(kPartActual, nPartActual); - // Locate the current tile of matrix B on L1 - MatrixCoord l1BOffset{kPartIdx * L0TileShape::K, nPartIdx * L0TileShape::N}; - auto l1BTile = l1BTensor[layoutBInL1.GetOffset(l1BOffset)]; - - // Wait for mmad finished - AscendC::WaitFlag(l0BEventList[l0BListId]); - // If the current tile is the first one on the k&n axis, wait for loading matrix B from GM to L1 - if ((kPartIdx == 0) && (nPartIdx == 0)) { - AscendC::WaitFlag(l1BEventList[l1ListId]); - } - - // Load current tile from L1 to L0B - copyL1ToL0B(l0BTile, l1BTile, layoutBInL0, layoutBInL1); - - // If the current tile is the last one on the k&n axis, notify to load matrix B from GM to L1 - if ((kPartIdx == kPartLoop - 1) && (nPartIdx == nPartLoop - 1)) { - AscendC::SetFlag(l1BEventList[l1ListId]); - } - // Notify to do mmad - AscendC::SetFlag(EVENT_ID0); - - // Locate the current tile on L0C - MatrixCoord l0COffset{mPartIdx * L0TileShape::M, nPartIdx * L0TileShape::N}; - auto l0CTile = l0CTensor[layoutInL0C.GetOffset(l0COffset)]; - - // Compute the matrix multiplication on L0A and L0B and write the result to the accumulator - // Wait for loading L0B - AscendC::WaitFlag(EVENT_ID0); - - // If the current tile is the first tile on the k axis, the accumulator needs to be reset to 0 - bool initC = ((kLoopIdx == 0) && (kPartIdx == 0)); - // If the unit flag is enabled, the unit flag is set according to the calculation progress - uint8_t unitFlag = 0b00; - if constexpr (ENABLE_UNIT_FLAG) { - if ((kLoopIdx == kTileCount - 1) && (mPartIdx == mPartLoop - 1) && - (kPartIdx == kPartLoop - 1) && (nPartIdx == nPartLoop - 1)) { - unitFlag = 0b11; - } else { - unitFlag = 0b10; - } - } - // Perform calculation operations - tileMmad(l0CTile, l0ATile, l0BTile, mPartActual, nPartActual, kPartActual, initC, unitFlag); - - // Notify to move the next L0B tile - AscendC::SetFlag(l0BEventList[l0BListId]); - l0BListId = (l0BListId + 1 < STAGES) ? (l0BListId + 1) : 0; - } - AscendC::SetFlag(l0AEventList[l0AListId]); - l0AListId = (l0AListId + 1 < STAGES) ? (l0AListId + 1) : 0; - } - } - l1ListId = l1ListIdNext; - kActual = kActualNext; - } - - // copy block out - LayoutC layoutBlock = layoutC.GetTileLayout(actualShape.GetCoordMN()); - - if (!isFirstKSlice) { // 切K后非第一块 开启原子加 - AscendC::SetAtomicAdd(); - } - if constexpr (!ENABLE_UNIT_FLAG) { - AscendC::SetFlag(EVENT_ID0); - AscendC::WaitFlag(EVENT_ID0); - copyL0CToGm(gmC, l0CTensor, layoutBlock, layoutInL0C); - AscendC::SetFlag(EVENT_ID0); - } else { - copyL0CToGm(gmC, l0CTensor, layoutBlock, layoutInL0C, 0b11); - } - if (!isFirstKSlice) { - AscendC::SetAtomicNone(); - } - } - -protected: - /// Data members - Params params; - - // static constexpr uint32_t STAGES = 2; - static constexpr uint32_t AIVPERCORE = 2; - - // Multi-stage tensors list - AscendC::LocalTensor l1ATensorList[STAGES]; - AscendC::LocalTensor l1BTensorList[STAGES]; - AscendC::LocalTensor l0ATensorList[STAGES]; - AscendC::LocalTensor l0BTensorList[STAGES]; - AscendC::LocalTensor l0CTensor; - - // Multi-stage event id list - int32_t l1AEventList[STAGES]; - int32_t l1BEventList[STAGES]; - int32_t l0AEventList[STAGES]; - int32_t l0BEventList[STAGES]; - - // The id of current stage - uint32_t l1ListId{0}; - uint32_t l0AListId{0}; - uint32_t l0BListId{0}; - - // CV Flag - static constexpr Arch::FlagID flagID0 = 0; - static constexpr Arch::FlagID flagID1 = 1; - static constexpr Arch::FlagID flagID2 = 2; - static constexpr Arch::FlagID flagID3 = 3; - static constexpr Arch::FlagID flagID4 = 4; - - Arch::CrossCoreFlag flag0[STAGES]; - Arch::CrossCoreFlag flag1[STAGES]; - Arch::CrossCoreFlag flag4{flagID4}; - - uint32_t crossCoreBufferIndexAIC{0}; - uint32_t crossCoreBufferIndexAIV{0}; - - // Tile action - TileMmad tileMmad; - CopyGmToL1A copyGmToL1A; - CopyGmToL1B copyGmToL1B; - CopyL1ToL0A copyL1ToL0A; - CopyL1ToL0B copyL1ToL0B; - CopyL0CToGm copyL0CToGm; - - // Buffer index (for prologue) - uint32_t bufferIndex{0}; - - Tile::PrologueTraits prologueA; - Tile::PrologueTraits prologueB; -}; - -} // namespace Catlass::Gemm::Block -#endif // CATLASS_GEMM_BLOCK_BLOCK_MMAD_PINGPONG_SLICE_K_WITH_PROLOGUE_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_pingpong_tla_v2.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_pingpong_tla_v2.hpp deleted file mode 100644 index de6198e50e..0000000000 --- a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_pingpong_tla_v2.hpp +++ /dev/null @@ -1,394 +0,0 @@ -/** - * Copyright (c) 2026 Huawei Technologies Co., Ltd. - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of - * CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ - -#ifndef CATLASS_GEMM_BLOCK_BLOCK_MMAD_PINGPONG_TLA_V2_HPP -#define CATLASS_GEMM_BLOCK_BLOCK_MMAD_PINGPONG_TLA_V2_HPP - -#include "catlass/kda_catlass.hpp" -#include "catlass/arch/kda_resource.hpp" -#include "catlass/kda_coord.hpp" -#include "catlass/kda_gemm_coord.hpp" -#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" -#include "catlass/gemm/kda_gemm_helper.hpp" -#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" -#include "catlass/gemm/tile/kda_tile_mmad.hpp" -#include "tla/kda_tla_layout.hpp" -#include "tla/kda_tensor.hpp" - -namespace Catlass::Gemm::Block { - -template < - class ArchTag_, - bool ENABLE_UNIT_FLAG_, - class L1TileShape_, - class L0TileShape_, - class ElementA_, - class ElementB_, - class ElementC_, - class ElementBias_, - class TileCopy_, - class TileMmad_> -struct BlockMmadTla< - MmadPingpongTlaV2, - L1TileShape_, - L0TileShape_, - ElementA_, - ElementB_, - ElementC_, - ElementBias_, - TileCopy_, - TileMmad_> { -public: - // Type Aliases - using DispatchPolicy = MmadPingpongTlaV2; - using ArchTag = typename DispatchPolicy::ArchTag; - using L1TileShape = L1TileShape_; - using L0TileShape = L0TileShape_; - using ElementA = ElementA_; - using LayoutA = typename TileCopy_::LayoutA; - using ElementB = ElementB_; - using LayoutB = typename TileCopy_::LayoutB; - using ElementC = ElementC_; - using LayoutC = typename TileCopy_::LayoutC; - - using TileMmad = TileMmad_; - - using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; - using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; - - using ElementAccumulator = typename TileCopy_::ElementAccumulator; - - using LayoutTagL1A = typename TileCopy_::LayoutTagL1A; - using LayoutTagL1B = typename TileCopy_::LayoutTagL1B; - using LayoutTagL0A = typename TileCopy_::LayoutTagL0A; - using LayoutTagL0B = typename TileCopy_::LayoutTagL0B; - using LayoutTagL0C = typename TileCopy_::LayoutTagL0C; - - using L1AAlignHelper = typename TileCopy_::L1AAlignHelper; - using L1BAlignHelper = typename TileCopy_::L1BAlignHelper; - - static_assert(tla::is_tuple::value && tla::is_static::value, - "L1TileShape must be tla::tuple and static!"); - static_assert(tla::is_tuple::value && tla::is_static::value, - "L0TileShape must be tla::tuple and static!"); - - static constexpr bool ENABLE_UNIT_FLAG = DispatchPolicy::ENABLE_UNIT_FLAG; - static constexpr uint32_t STAGES = DispatchPolicy::STAGES; - static constexpr uint32_t L1_TILE_M = tla::get<0>(L1TileShape{}); - static constexpr uint32_t L1_TILE_N = tla::get<1>(L1TileShape{}); - static constexpr uint32_t L1_TILE_K = tla::get<2>(L1TileShape{}); - static constexpr uint32_t L0_TILE_M = tla::get<0>(L0TileShape{}); - static constexpr uint32_t L0_TILE_N = tla::get<1>(L0TileShape{}); - static constexpr uint32_t L0_TILE_K = tla::get<2>(L0TileShape{}); - - // L1 tile size - static constexpr uint32_t L1A_TILE_SIZE = L1_TILE_M * L1_TILE_K * sizeof(ElementA); - static constexpr uint32_t L1B_TILE_SIZE = L1_TILE_N * L1_TILE_K * sizeof(ElementB); - // L0 tile size - static constexpr uint32_t L0A_TILE_SIZE = L0_TILE_M * L0_TILE_K * sizeof(ElementA); - static constexpr uint32_t L0B_TILE_SIZE = L0_TILE_K * L0_TILE_N * sizeof(ElementB); - static constexpr uint32_t L0C_TILE_SIZE = L1_TILE_M * L1_TILE_N * sizeof(ElementAccumulator); - - // Check LayoutC - static_assert(tla::detail::isRowMajor::value || - ((std::is_same_v || std::is_same_v || - std::is_same_v) && - tla::detail::iszN::value), - "LayoutC only supports zN in half or bfloat16 or float, RowMajor in all dtype yet!"); - - // Check L1TileShape - static_assert((L1A_TILE_SIZE + L1B_TILE_SIZE) * STAGES <= ArchTag::L1_SIZE, - "L1TileShape exceeding the L1 space!"); - - // Check L0TileShape - static_assert(L0A_TILE_SIZE * STAGES <= ArchTag::L0A_SIZE, "L0TileShape exceeding the L0A space!"); - static_assert(L0B_TILE_SIZE * STAGES <= ArchTag::L0B_SIZE, "L0TileShape exceeding the L0B space!"); - static_assert(L0C_TILE_SIZE <= ArchTag::L0C_SIZE, "L0TileShape exceeding the L0C space!"); - - static constexpr uint32_t _32B = 32 * 8; // in bits - static_assert(L1_TILE_M == L0_TILE_M && L1_TILE_N == L0_TILE_N, - "The situation where the basic blocks of L1 and L0 differ on the m and n axes is not supported yet"); - static_assert(L0_TILE_K <= L1_TILE_K, "L0TileShape::K cannot exceed L1TileShape::K"); - static_assert(L1_TILE_M * SizeOfBits::value % _32B == 0, "L1TileShape::M must be 32B aligned."); - static_assert(L1_TILE_K * SizeOfBits::value % _32B == 0, "L1TileShape::K must be 32B aligned."); - static_assert(L1_TILE_K * SizeOfBits::value % _32B == 0, "L1TileShape::K must be 32B aligned."); - static_assert(L1_TILE_N * SizeOfBits::value % _32B == 0, "L1TileShape::N must be 32B aligned."); - static_assert(L0_TILE_K * SizeOfBits::value % _32B == 0, "L0TileShape::K must be 32B aligned."); - - static constexpr auto L1A_LAYOUT = tla::MakeLayout(tla::Int{}, tla::Int{}); - static constexpr auto L1B_LAYOUT = tla::MakeLayout(tla::Int{}, tla::Int{}); - - /// Construct - CATLASS_DEVICE - BlockMmadTla(Arch::Resource &resource, uint32_t l1BufAddrStart = 0) - { - uint32_t l1AOffset = l1BufAddrStart; - uint32_t l1BOffset = l1BufAddrStart + L1A_TILE_SIZE * STAGES; - // Init buffers - for (uint32_t i = 0; i < STAGES; i++) { - // Assign L1/L0A/L0B space for each stages - l1ATensorList[i] = resource.l1Buf.template GetBufferByByte(l1AOffset + L1A_TILE_SIZE * i); - l1BTensorList[i] = resource.l1Buf.template GetBufferByByte(l1BOffset + L1B_TILE_SIZE * i); - l0ATensorList[i] = resource.l0ABuf.template GetBufferByByte(L0A_TILE_SIZE * i); - l0BTensorList[i] = resource.l0BBuf.template GetBufferByByte(L0B_TILE_SIZE * i); - - // Assign event ID for each stages - l1AEventList[i] = i; - l1BEventList[i] = i + STAGES; - l0AEventList[i] = i; - l0BEventList[i] = i + STAGES; - - // The event id that needs to be set before the loop - AscendC::SetFlag(l1AEventList[i]); - AscendC::SetFlag(l1BEventList[i]); - AscendC::SetFlag(l0AEventList[i]); - AscendC::SetFlag(l0BEventList[i]); - } - l0CTensor = resource.l0CBuf.template GetBufferByByte(0); - AscendC::SetFlag(EVENT_ID0); - } - - /// Destructor - CATLASS_DEVICE - ~BlockMmadTla() - { - for (uint32_t i = 0; i < STAGES; i++) { - AscendC::WaitFlag(l1AEventList[i]); - AscendC::WaitFlag(l1BEventList[i]); - AscendC::WaitFlag(l0AEventList[i]); - AscendC::WaitFlag(l0BEventList[i]); - } - AscendC::WaitFlag(EVENT_ID0); - } - - /// Perform a block-scoped matrix multiply-accumulate - template - CATLASS_DEVICE void operator()(TensorA &tensorA, TensorB &tensorB, TensorC &tensorC) - { - using CopyGmToL1A = typename TileCopy_::template CopyGmToL1A; - using CopyGmToL1B = typename TileCopy_::template CopyGmToL1B; - CopyGmToL1A copyGmToL1A; - CopyGmToL1B copyGmToL1B; -#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 2201) - using CopyL0CToGm = typename TileCopy_::template CopyL0CToGm; - CopyL0CToGm copyL0CToDst; -#endif -#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 3510) - using CopyL0CToDst = typename TileCopy_::template CopyL0CToDst; - CopyL0CToDst copyL0CToDst; -#endif - - // Create an accumulator tensor view on L0C buffer. - // - Logical size comes from tensorC.layout().originShape() (tail-aware) - // - Layout is constructed from LayoutTagL0C - // - coord is initialized to (0, 0) (new buffer view) - auto tensorL0C = tla::MakeTensorLike(l0CTensor, tensorC, Arch::PositionL0C{}); - - // load first matrix A tile from GM to L1 - AscendC::WaitFlag(l1AEventList[l1ListId]); - // TileView: tileCoord is in tile units (not element units). - // It internally converts tileCoord to elementOffset = tileCoord ⊙ tileShape and handles tail tiles via originShape. - auto tensorTileA = tla::TileView(tensorA, - tla::MakeCoord(0u, 0u), // (m_tile, k_tile) - tla::MakeShape(tla::Int{}, tla::Int{})); - auto tensorL1A = tla::MakeTensorLike(l1ATensorList[l1ListId], tensorTileA, Arch::PositionL1{}, L1A_LAYOUT); - copyGmToL1A(tensorL1A, tensorTileA); - AscendC::SetFlag(l1AEventList[l1ListId]); - - // load first matrix B tile from GM to L1 - AscendC::WaitFlag(l1BEventList[l1ListId]); - auto tensorTileB = tla::TileView(tensorB, - tla::MakeCoord(0u, 0u), // (k_tile, n_tile) - tla::MakeShape(tla::Int{}, tla::Int{})); - auto tensorL1B = tla::MakeTensorLike(l1BTensorList[l1ListId], tensorTileB, Arch::PositionL1{}, L1B_LAYOUT); - copyGmToL1B(tensorL1B, tensorTileB); - AscendC::SetFlag(l1BEventList[l1ListId]); - - if constexpr (!ENABLE_UNIT_FLAG) { - AscendC::WaitFlag(EVENT_ID0); - } - - uint32_t mPartLoop = CeilDiv(tla::get<0>(tensorL0C.originShape())); // dim 0 = M - uint32_t nPartLoop = CeilDiv(tla::get<1>(tensorL0C.originShape())); // dim 1 = N - - // main loop - uint32_t kTileCount = CeilDiv(tla::get<1>(tensorA.originShape())); // dim 1 = K - for (uint32_t kLoopIdx = 0; kLoopIdx < kTileCount; kLoopIdx++) { - uint32_t l1ListIdNext = (l1ListId + 1 < STAGES) ? (l1ListId + 1) : 0; - // preload next tile from GM to L1 - if (kLoopIdx < kTileCount - 1) { - uint32_t kLoopIdxNext = kLoopIdx + 1; - - // Get L1 tensor for next stage - auto l1ATensor = l1ATensorList[l1ListIdNext]; - auto l1BTensor = l1BTensorList[l1ListIdNext]; - // Get GM tile for next stage - auto tensorTileA = tla::TileView(tensorA, - tla::MakeCoord(0u, kLoopIdxNext), // (m_tile, k_tile) - tla::MakeShape(tla::Int{}, tla::Int{})); - auto tensorTileB = tla::TileView(tensorB, - tla::MakeCoord(kLoopIdxNext, 0u), // (k_tile, n_tile) - tla::MakeShape(tla::Int{}, tla::Int{})); - auto tensorL1A = tla::MakeTensorLike(l1ATensor, tensorTileA, Arch::PositionL1{}, L1A_LAYOUT); - auto tensorL1B = tla::MakeTensorLike(l1BTensor, tensorTileB, Arch::PositionL1{}, L1B_LAYOUT); - - // load next matrix A tile from GM to L1 - AscendC::WaitFlag(l1AEventList[l1ListIdNext]); - copyGmToL1A(tensorL1A, tensorTileA); - AscendC::SetFlag(l1AEventList[l1ListIdNext]); - - // load next matrix B tile from GM to L1 - AscendC::WaitFlag(l1BEventList[l1ListIdNext]); - copyGmToL1B(tensorL1B, tensorTileB); - AscendC::SetFlag(l1BEventList[l1ListIdNext]); - } - - // Get L1 tensor for current stage - auto l1ATensor = l1ATensorList[l1ListId]; - auto l1BTensor = l1BTensorList[l1ListId]; - // Create tile view for current K iteration - auto tensorTileA = tla::TileView(tensorA, - tla::MakeCoord(0u, kLoopIdx), // (m_tile, k_tile) - tla::MakeShape(tla::Int{}, tla::Int{})); - auto tensorTileB = tla::TileView(tensorB, - tla::MakeCoord(kLoopIdx, 0u), // (k_tile, n_tile) - tla::MakeShape(tla::Int{}, tla::Int{})); - auto tensorL1A = tla::MakeTensorLike(l1ATensor, tensorTileA, Arch::PositionL1{}, L1A_LAYOUT); - auto tensorL1B = tla::MakeTensorLike(l1BTensor, tensorTileB, Arch::PositionL1{}, L1B_LAYOUT); - // Get the loop nums on L0 based on current L1 tile's actual K size - uint32_t kPartLoop = CeilDiv(tla::get<1>(tensorL1A.originShape())); // dim 1 = K - - for (int mPartIdx = 0; mPartIdx < mPartLoop; mPartIdx++) { - for (int kPartIdx = 0; kPartIdx < kPartLoop; kPartIdx++) { - // Locate the current tile on L0A - auto l0ATile = l0ATensorList[l0AListId]; - // Locate the current tile of matrix A on L1 - // Take a (L0_TILE_M, L0_TILE_K) tile from the current L1A tile (tile coordinates within L1 tile). - auto tensorTileL1A = tla::TileView( - tensorL1A, - tla::MakeCoord(mPartIdx, kPartIdx), - tla::MakeShape(tla::Int{}, tla::Int{})); - auto tensorL0A = tla::MakeTensorLike(l0ATile, tensorTileL1A, Arch::PositionL0A{}); - - AscendC::WaitFlag(l0AEventList[l0AListId]); - if ((mPartIdx == 0) && (kPartIdx == 0)) { - AscendC::WaitFlag(l1AEventList[l1ListId]); - } - - // Load current tile from L1 to L0A - copyL1ToL0A(tensorL0A, tensorTileL1A); - - if ((mPartIdx == mPartLoop - 1) && (kPartIdx == kPartLoop - 1)) { - AscendC::SetFlag(l1AEventList[l1ListId]); - } - - for (int nPartIdx = 0; nPartIdx < nPartLoop; nPartIdx++) { - // Locate the current tile on L0B - auto l0BTile = l0BTensorList[l0BListId]; - // Locate the current tile of matrix B on L1 - // Take a (L0_TILE_K, L0_TILE_N) tile from the current L1B tile (tile coordinates within L1 tile). - auto tensorTileL1B = tla::TileView( - tensorL1B, - tla::MakeCoord(kPartIdx, nPartIdx), - tla::MakeShape(tla::Int{}, tla::Int{})); - auto tensorL0B = tla::MakeTensorLike(l0BTile, tensorTileL1B, Arch::PositionL0B{}); - - // Wait for mmad finished - AscendC::WaitFlag(l0BEventList[l0BListId]); - // If the current tile is the first one on the k&n axis, wait for loading matrix B from GM to L1 - if ((kPartIdx == 0) && (nPartIdx == 0)) { - AscendC::WaitFlag(l1BEventList[l1ListId]); - } - - // Load current tile from L1 to L0B - copyL1ToL0B(tensorL0B, tensorTileL1B); - - // If the current tile is the last one on the k&n axis, notify to load matrix B from GM to L1 - if ((kPartIdx == kPartLoop - 1) && (nPartIdx == nPartLoop - 1)) { - AscendC::SetFlag(l1BEventList[l1ListId]); - } - // Notify to do mmad - AscendC::SetFlag(EVENT_ID0); - - // Locate the current tile on L0C - // View into L0C accumulator tile (tile coordinates in (m_part, n_part)). - auto tensorTileL0C = tla::TileView( - tensorL0C, - tla::MakeCoord(mPartIdx, nPartIdx), - tla::MakeShape(tla::Int{}, tla::Int{})); - - // Compute the matrix multiplication on L0A and L0B and write the result to the accumulator - // Wait for loading L0B - AscendC::WaitFlag(EVENT_ID0); - - // If the current tile is the first tile on the k axis, the accumulator needs to be reset to 0 - bool initC = ((kLoopIdx == 0) && (kPartIdx == 0)); - // If the unit flag is enabled, the unit flag is set according to the calculation progress - uint8_t unitFlag = 0b00; - if constexpr (ENABLE_UNIT_FLAG) { - if ((kLoopIdx == kTileCount - 1) && (mPartIdx == mPartLoop - 1) && - (kPartIdx == kPartLoop - 1) && (nPartIdx == nPartLoop - 1)) { - unitFlag = 0b11; - } else { - unitFlag = 0b10; - } - } - // Perform calculation operations - tileMmad(tensorTileL0C, tensorL0A, tensorL0B, initC, unitFlag); - - // Notify to move the next L0B tile - AscendC::SetFlag(l0BEventList[l0BListId]); - l0BListId = (l0BListId + 1 < STAGES) ? (l0BListId + 1) : 0; - } - AscendC::SetFlag(l0AEventList[l0AListId]); - l0AListId = (l0AListId + 1 < STAGES) ? (l0AListId + 1) : 0; - } - } - l1ListId = l1ListIdNext; - } - - // copy block out - if constexpr (!ENABLE_UNIT_FLAG) { - AscendC::SetFlag(EVENT_ID0); - AscendC::WaitFlag(EVENT_ID0); - copyL0CToDst(tensorC, tensorL0C); - AscendC::SetFlag(EVENT_ID0); - } else { - copyL0CToDst(tensorC, tensorL0C, 0b11); - } - } - -protected: - // Multi-stage tensors list - AscendC::LocalTensor l1ATensorList[STAGES]; - AscendC::LocalTensor l1BTensorList[STAGES]; - AscendC::LocalTensor l0ATensorList[STAGES]; - AscendC::LocalTensor l0BTensorList[STAGES]; - AscendC::LocalTensor l0CTensor; - - // Multi-stage event id list - int32_t l1AEventList[STAGES]; - int32_t l1BEventList[STAGES]; - int32_t l0AEventList[STAGES]; - int32_t l0BEventList[STAGES]; - - // The id of current stage - uint32_t l1ListId{0}; - uint32_t l0AListId{0}; - uint32_t l0BListId{0}; - - TileMmad tileMmad; - CopyL1ToL0A copyL1ToL0A; - CopyL1ToL0B copyL1ToL0B; -}; - -} // namespace Catlass::Gemm::Block - -#endif // CATLASS_GEMM_BLOCK_BLOCK_MMAD_PINGPONG_TLA_V2_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_pingpong_with_prologue.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_pingpong_with_prologue.hpp deleted file mode 100644 index dc1c234b83..0000000000 --- a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_pingpong_with_prologue.hpp +++ /dev/null @@ -1,535 +0,0 @@ -/** - * Copyright (c) 2025 Huawei Technologies Co., Ltd. - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of - * CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ - -#ifndef CATLASS_GEMM_BLOCK_BLOCK_MMAD_PINGPONG_WITH_PROLOGUE -#define CATLASS_GEMM_BLOCK_BLOCK_MMAD_PINGPONG_WITH_PROLOGUE - -#include "catlass/kda_catlass.hpp" -#include "catlass/arch/kda_resource.hpp" -#include "catlass/kda_coord.hpp" -#include "catlass/kda_gemm_coord.hpp" -#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" -#include "catlass/gemm/kda_gemm_helper.hpp" -#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" -#include "catlass/gemm/tile/kda_tile_traits.hpp" - -namespace Catlass::Gemm::Block { - -template < - bool ENABLE_UNIT_FLAG_, - class L1TileShape_, - class L0TileShape_, - class AType_, - class BType_, - class CType_, - class BiasType_, - class TileCopy_, - class TileMmad_> -struct BlockMmad< - MmadAtlasA2PingPongWithPrologue, - L1TileShape_, - L0TileShape_, - AType_, - BType_, - CType_, - BiasType_, - TileCopy_, - TileMmad_> { -public: - // Type Aliases - using DispatchPolicy = MmadAtlasA2PingPongWithPrologue; - using ArchTag = typename DispatchPolicy::ArchTag; - using L1TileShape = L1TileShape_; - using L0TileShape = L0TileShape_; - using ElementA = typename AType_::Element; - using LayoutA = typename AType_::Layout; - using ElementB = typename BType_::Element; - using LayoutB = typename BType_::Layout; - using ElementC = typename CType_::Element; - using LayoutC = typename CType_::Layout; - using TileMmad = TileMmad_; - - using CopyGmToL1A = typename TileCopy_::CopyGmToL1A; - using CopyGmToL1B = typename TileCopy_::CopyGmToL1B; - using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; - using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; - - using PrologueA = typename TileCopy_::PrologueA; - using PrologueB = typename TileCopy_::PrologueB; - - using ElementAccumulator = - typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; - using CopyL0CToGm = typename TileCopy_::CopyL0CToGm; - using LayoutAInL1 = typename CopyL1ToL0A::LayoutSrc; - using LayoutBInL1 = typename CopyL1ToL0B::LayoutSrc; - using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst; - using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst; - using LayoutCInL0 = layout::zN; - - using L1AAlignHelper = Gemm::helper::L1AlignHelper; - using L1BAlignHelper = Gemm::helper::L1AlignHelper; - - static constexpr bool ENABLE_UNIT_FLAG = DispatchPolicy::ENABLE_UNIT_FLAG; - static constexpr uint32_t STAGES = DispatchPolicy::STAGES; - static constexpr uint32_t L1A_SIZE = L1TileShape::M * L1TileShape::K * sizeof(ElementA); - static constexpr uint32_t L1B_SIZE = L1TileShape::N * L1TileShape::K * sizeof(ElementB); - static constexpr uint32_t L0A_SIZE = ArchTag::L0A_SIZE; - static constexpr uint32_t L0B_SIZE = ArchTag::L0B_SIZE; - static constexpr uint32_t L0C_SIZE = ArchTag::L0C_SIZE; - static constexpr uint32_t L0A_PINGPONG_BUF_SIZE = L0A_SIZE / STAGES; - static constexpr uint32_t L0B_PINGPONG_BUF_SIZE = L0B_SIZE / STAGES; - - static constexpr bool HAS_PROLOGUE_A = !std::is_same_v; - static constexpr bool HAS_PROLOGUE_B = !std::is_same_v; - - // Check L1TileShape - static_assert((L1A_SIZE * STAGES + L1B_SIZE * STAGES) <= ArchTag::L1_SIZE, - "L1TileShape exceeding the L1 space!"); - - // Check L0TileShape - static constexpr uint32_t L0A_TILE_SIZE = L0TileShape::M * L0TileShape::K * sizeof(ElementA); - static constexpr uint32_t L0B_TILE_SIZE = L0TileShape::K * L0TileShape::N * sizeof(ElementB); - static constexpr uint32_t L0C_TILE_SIZE = L0TileShape::M * L0TileShape::N * sizeof(ElementAccumulator); - static_assert((L0A_TILE_SIZE * STAGES) <= L0A_SIZE, - "L0TileShape exceeding the L0A space!"); - static_assert((L0B_TILE_SIZE * STAGES) <= L0B_SIZE, - "L0TileShape exceeding the L0B space!"); - static_assert(L0C_TILE_SIZE <= L0C_SIZE, "L0TileShape exceeding the L0C space!"); - - static_assert(L1TileShape::M == L0TileShape::M && L1TileShape::N == L0TileShape::N, - "The situation where the basic blocks of L1 and L0 differ on the m and n axes is not supported yet"); - static_assert(L0TileShape::K <= L1TileShape::K, "L0TileShape::K cannot exceed L1TileShape::K"); - - // 32B (256b) aligned - static_assert(Gemm::helper::TileShapeAlignChecker::_ALIGN == 256, - "Tile shape must be 32B aligned."); - - struct Params { - typename Tile::PrologueTraits::Params prologueA{}; - typename Tile::PrologueTraits::Params prologueB{}; - typename CopyL0CToGm::Params copyL0CToGm{}; - }; - - /// Construct - CATLASS_DEVICE - BlockMmad(Arch::Resource const &resource, Params const ¶ms_ = {}) - : params(params_), - prologueA(resource, params_.prologueA), - prologueB(resource, params_.prologueB), - copyL0CToGm(params_.copyL0CToGm) - { - Arch::FlagID flagId = 0; - for (uint32_t i = 0; i < STAGES; ++i) { - if constexpr (HAS_PROLOGUE_A) { - flagCopyAFinish[i] = Arch::CrossCoreFlag(flagId++); - flagPrologueAFinish[i] = Arch::CrossCoreFlag(flagId++); - } - - if constexpr (HAS_PROLOGUE_B) { - flagCopyBFinish[i] = Arch::CrossCoreFlag(flagId++); - flagPrologueBFinish[i] = Arch::CrossCoreFlag(flagId++); - } - } - - if constexpr (g_coreType == AscendC::AIC) { - uint32_t l1AOffset = 0; - uint32_t l1BOffset = L1A_SIZE * STAGES; - // Init buffers - for (uint32_t i = 0; i < STAGES; i++) { - if constexpr (HAS_PROLOGUE_A) { - Catlass::Arch::CrossCoreSetFlag<0x02, PIPE_MTE2>(flagCopyAFinish[i]); - } - - if constexpr (HAS_PROLOGUE_B) { - Catlass::Arch::CrossCoreSetFlag<0x02, PIPE_MTE2>(flagCopyBFinish[i]); - } - - l1AEventList[i] = i; - l1BEventList[i] = i + STAGES; - l0AEventList[i] = i; - l0BEventList[i] = i + STAGES; - - AscendC::SetFlag(l1AEventList[i]); - AscendC::SetFlag(l1BEventList[i]); - AscendC::SetFlag(l0AEventList[i]); - AscendC::SetFlag(l0BEventList[i]); - - l1ATensorList[i] = resource.l1Buf.template GetBufferByByte(l1AOffset + L1A_SIZE * i); - l1BTensorList[i] = resource.l1Buf.template GetBufferByByte(l1BOffset + L1B_SIZE * i); - l0ATensorList[i] = resource.l0ABuf.template GetBufferByByte(L0A_PINGPONG_BUF_SIZE * i); - l0BTensorList[i] = resource.l0BBuf.template GetBufferByByte(L0B_PINGPONG_BUF_SIZE * i); - } - AscendC::SetFlag(EVENT_ID0); - l0CTensor = resource.l0CBuf.template GetBufferByByte(0); - } - } - - /// Destructor - CATLASS_DEVICE - ~BlockMmad() - { - if constexpr (g_coreType == AscendC::AIC) { - for (uint32_t i = 0; i < STAGES; i++) { - AscendC::WaitFlag(l1AEventList[i]); - AscendC::WaitFlag(l1BEventList[i]); - AscendC::WaitFlag(l0AEventList[i]); - AscendC::WaitFlag(l0BEventList[i]); - } - AscendC::WaitFlag(EVENT_ID0); - } else { - for (uint32_t i = 0; i < STAGES; i++) { - if constexpr (HAS_PROLOGUE_A) { - Catlass::Arch::CrossCoreWaitFlag(flagCopyAFinish[i]); - } - - if constexpr (HAS_PROLOGUE_B) { - Catlass::Arch::CrossCoreWaitFlag(flagCopyBFinish[i]); - } - } - } - } - - template - CATLASS_DEVICE - std::enable_if_t && !std::is_void_v, void> - Prologue( - AscendC::GlobalTensor const &gmSrcA, typename T::LayoutSrc const &layoutSrcA, - AscendC::GlobalTensor const &gmDstA, typename T::LayoutDst const &layoutDstA, - AscendC::GlobalTensor const &gmSrcB, typename U::LayoutSrc const &layoutSrcB, - AscendC::GlobalTensor const &gmDstB, typename U::LayoutDst const &layoutDstB, - GemmCoord const &actualBlockShape) - { - PrologueImpl( - gmSrcA, layoutSrcA, gmDstA, layoutDstA, - gmSrcB, layoutSrcB, gmDstB, layoutDstB, - actualBlockShape); - } - - template - CATLASS_DEVICE - std::enable_if_t && std::is_void_v, void> - Prologue( - AscendC::GlobalTensor const &gmSrcA, typename T::LayoutSrc const &layoutSrcA, - AscendC::GlobalTensor const &gmDstA, typename T::LayoutDst const &layoutDstA, - GemmCoord const &actualBlockShape) - { - PrologueImpl( - gmSrcA, layoutSrcA, gmDstA, layoutDstA, - {}, {}, {}, {}, - actualBlockShape); - } - - template - CATLASS_DEVICE - std::enable_if_t && !std::is_void_v, void> - Prologue( - AscendC::GlobalTensor const &gmSrcB, typename U::LayoutSrc const &layoutSrcB, - AscendC::GlobalTensor const &gmDstB, typename U::LayoutDst const &layoutDstB, - GemmCoord const &actualBlockShape) - { - PrologueImpl( - {}, {}, {}, {}, - gmSrcB, layoutSrcB, gmDstB, layoutDstB, - actualBlockShape); - } - - /// Perform a block-scoped matrix multiply-accumulate - CATLASS_DEVICE - void operator()( - AscendC::GlobalTensor const &gmA, LayoutA const &layoutA, - AscendC::GlobalTensor const &gmB, LayoutB const &layoutB, - AscendC::GlobalTensor const &gmC, LayoutC const &layoutC, - GemmCoord const &actualShape) - { - uint32_t mRound = RoundUp(actualShape.m()); - uint32_t nRound = RoundUp(actualShape.n()); - - auto layoutAInL1 = LayoutAInL1::template MakeLayout(L1TileShape::M, L1TileShape::K); - auto layoutBInL1 = LayoutBInL1::template MakeLayout(L1TileShape::K, L1TileShape::N); - auto layoutInL0C = LayoutCInL0::MakeLayoutInL0C(MakeCoord(mRound, nRound)); - - uint32_t kActual = min(actualShape.k(), L1TileShape::K); - - // load first matrix A tile from GM to L1 - AscendC::WaitFlag(l1AEventList[l1ListId]); - auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShape.m(), kActual)); - if constexpr (HAS_PROLOGUE_A) { - auto gmTileA = gmA[l1ListId * layoutA.Capacity()]; - Catlass::Arch::CrossCoreWaitFlag(flagPrologueAFinish[l1ListId]); - copyGmToL1A(l1ATensorList[l1ListId], gmTileA, layoutAInL1, layoutTileA); - Catlass::Arch::CrossCoreSetFlag<0x02, PIPE_MTE2>(flagCopyAFinish[l1ListId]); - } else { - copyGmToL1A(l1ATensorList[l1ListId], gmA, layoutAInL1, layoutTileA); - } - AscendC::SetFlag(l1AEventList[l1ListId]); - - // load first matrix B tile from GM to L1 - AscendC::WaitFlag(l1BEventList[l1ListId]); - auto layoutTileB = layoutB.GetTileLayout(MakeCoord(kActual, actualShape.n())); - if constexpr (HAS_PROLOGUE_B) { - auto gmTileB = gmB[l1ListId * layoutB.Capacity()]; - Catlass::Arch::CrossCoreWaitFlag(flagPrologueBFinish[l1ListId]); - copyGmToL1B(l1BTensorList[l1ListId], gmTileB, layoutBInL1, layoutTileB); - Catlass::Arch::CrossCoreSetFlag<0x02, PIPE_MTE2>(flagCopyBFinish[l1ListId]); - } else { - copyGmToL1B(l1BTensorList[l1ListId], gmB, layoutBInL1, layoutTileB); - } - AscendC::SetFlag(l1BEventList[l1ListId]); - - if constexpr (!ENABLE_UNIT_FLAG) { - AscendC::WaitFlag(EVENT_ID0); - } - - uint32_t mPartLoop = CeilDiv(mRound); - uint32_t nPartLoop = CeilDiv(nRound); - - // main loop - uint32_t kTileCount = CeilDiv(actualShape.k()); - for (uint32_t kLoopIdx = 0; kLoopIdx < kTileCount; kLoopIdx++) { - uint32_t l1ListIdNext = (l1ListId + 1 < STAGES) ? (l1ListId + 1) : 0; - uint32_t kActualNext{0}; - // preload next tile from GM to L1 - if (kLoopIdx < kTileCount - 1) { - uint32_t kLoopIdxNext = kLoopIdx + 1; - kActualNext = (kLoopIdxNext < kTileCount - 1) ? - L1TileShape::K : - (actualShape.k() - kLoopIdxNext * L1TileShape::K); - - // Get L1 tensor for next stage - auto l1ATensor = l1ATensorList[l1ListIdNext]; - auto l1BTensor = l1BTensorList[l1ListIdNext]; - - // load next matrix A tile from GM to L1 - AscendC::WaitFlag(l1AEventList[l1ListIdNext]); - layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShape.m(), kActualNext)); - if constexpr (HAS_PROLOGUE_A) { - auto gmTileA = gmA[l1ListIdNext * layoutA.Capacity()]; - Catlass::Arch::CrossCoreWaitFlag(flagPrologueAFinish[l1ListIdNext]); - copyGmToL1A(l1ATensor, gmTileA, layoutAInL1, layoutTileA); - Catlass::Arch::CrossCoreSetFlag<0x02, PIPE_MTE2>(flagCopyAFinish[l1ListIdNext]); - } else { - MatrixCoord gmTileAOffset{0, kLoopIdxNext * L1TileShape::K}; - auto gmTileA = gmA[layoutA.GetOffset(gmTileAOffset)]; - copyGmToL1A(l1ATensor, gmTileA, layoutAInL1, layoutTileA); - } - AscendC::SetFlag(l1AEventList[l1ListIdNext]); - - // load next matrix B tile from GM to L1 - AscendC::WaitFlag(l1BEventList[l1ListIdNext]); - layoutTileB = layoutB.GetTileLayout(MakeCoord(kActualNext, actualShape.n())); - if constexpr (HAS_PROLOGUE_B) { - auto gmTileB = gmB[l1ListIdNext * layoutB.Capacity()]; - Catlass::Arch::CrossCoreWaitFlag(flagPrologueBFinish[l1ListIdNext]); - copyGmToL1B(l1BTensor, gmTileB, layoutBInL1, layoutTileB); - Catlass::Arch::CrossCoreSetFlag<0x02, PIPE_MTE2>(flagCopyBFinish[l1ListIdNext]); - } else { - MatrixCoord gmTileBOffset{kLoopIdxNext * L1TileShape::K, 0}; - auto gmTileB = gmB[layoutB.GetOffset(gmTileBOffset)]; - copyGmToL1B(l1BTensor, gmTileB, layoutBInL1, layoutTileB); - } - AscendC::SetFlag(l1BEventList[l1ListIdNext]); - } - - // Get L1 tensor for current stage - auto l1ATensor = l1ATensorList[l1ListId]; - auto l1BTensor = l1BTensorList[l1ListId]; - - // Get the loop nums on L0 - uint32_t kPartLoop = CeilDiv(kActual); - - for (int mPartIdx = 0; mPartIdx < mPartLoop; mPartIdx++) { - uint32_t mPartActual = (mPartIdx < mPartLoop - 1) ? - L0TileShape::M : - (mRound - mPartIdx * L0TileShape::M); - - for (int kPartIdx = 0; kPartIdx < kPartLoop; kPartIdx++) { - uint32_t kPartActual = (kPartIdx < kPartLoop - 1) ? - L0TileShape::K : - (kActual - kPartIdx * L0TileShape::K); - - // Locate the current tile on L0A - auto l0ATile = l0ATensorList[l0AListId]; - LayoutAInL0 layoutAInL0 = LayoutAInL0::template MakeLayout(mPartActual, kPartActual); - // Locate the current tile of matrix A on L1 - MatrixCoord l1AOffset{mPartIdx * L0TileShape::M, kPartIdx * L0TileShape::K}; - auto l1ATile = l1ATensor[layoutAInL1.GetOffset(l1AOffset)]; - - AscendC::WaitFlag(l0AEventList[l0AListId]); - if ((mPartIdx == 0) && (kPartIdx == 0)) { - AscendC::WaitFlag(l1AEventList[l1ListId]); - } - - // Load current tile from L1 to L0A - copyL1ToL0A(l0ATile, l1ATile, layoutAInL0, layoutAInL1); - - if ((mPartIdx == mPartLoop - 1) && (kPartIdx == kPartLoop - 1)) { - AscendC::SetFlag(l1AEventList[l1ListId]); - } - - for (int nPartIdx = 0; nPartIdx < nPartLoop; nPartIdx++) { - uint32_t nPartActual = (nPartIdx < nPartLoop - 1) ? - L0TileShape::N : - (nRound - nPartIdx * L0TileShape::N); - - // Locate the current tile on L0B - auto l0BTile = l0BTensorList[l0BListId]; - LayoutBInL0 layoutBInL0 = LayoutBInL0::template MakeLayout(kPartActual, nPartActual); - // Locate the current tile of matrix B on L1 - MatrixCoord l1BOffset{kPartIdx * L0TileShape::K, nPartIdx * L0TileShape::N}; - auto l1BTile = l1BTensor[layoutBInL1.GetOffset(l1BOffset)]; - - // Wait for mmad finished - AscendC::WaitFlag(l0BEventList[l0BListId]); - // If the current tile is the first one on the k&n axis, wait for loading matrix B from GM to L1 - if ((kPartIdx == 0) && (nPartIdx == 0)) { - AscendC::WaitFlag(l1BEventList[l1ListId]); - } - - // Load current tile from L1 to L0B - copyL1ToL0B(l0BTile, l1BTile, layoutBInL0, layoutBInL1); - - // If the current tile is the last one on the k&n axis, notify to load matrix B from GM to L1 - if ((kPartIdx == kPartLoop - 1) && (nPartIdx == nPartLoop - 1)) { - AscendC::SetFlag(l1BEventList[l1ListId]); - } - // Notify to do mmad - AscendC::SetFlag(EVENT_ID0); - - // Locate the current tile on L0C - MatrixCoord l0COffset{mPartIdx * L0TileShape::M, nPartIdx * L0TileShape::N}; - auto l0CTile = l0CTensor[layoutInL0C.GetOffset(l0COffset)]; - - // Compute the matrix multiplication on L0A and L0B and write the result to the accumulator - // Wait for loading L0B - AscendC::WaitFlag(EVENT_ID0); - - // If the current tile is the first tile on the k axis, the accumulator needs to be reset to 0 - bool initC = ((kLoopIdx == 0) && (kPartIdx == 0)); - // If the unit flag is enabled, the unit flag is set according to the calculation progress - uint8_t unitFlag = 0b00; - if constexpr (ENABLE_UNIT_FLAG) { - if ((kLoopIdx == kTileCount - 1) && (mPartIdx == mPartLoop - 1) && - (kPartIdx == kPartLoop - 1) && (nPartIdx == nPartLoop - 1)) { - unitFlag = 0b11; - } else { - unitFlag = 0b10; - } - } - // Perform calculation operations - tileMmad(l0CTile, l0ATile, l0BTile, mPartActual, nPartActual, kPartActual, initC, unitFlag); - - // Notify to move the next L0B tile - AscendC::SetFlag(l0BEventList[l0BListId]); - l0BListId = (l0BListId + 1 < STAGES) ? (l0BListId + 1) : 0; - } - AscendC::SetFlag(l0AEventList[l0AListId]); - l0AListId = (l0AListId + 1 < STAGES) ? (l0AListId + 1) : 0; - } - } - l1ListId = l1ListIdNext; - kActual = kActualNext; - } - - // copy block out - LayoutC layoutBlock = layoutC.GetTileLayout(actualShape.GetCoordMN()); - - if constexpr (!ENABLE_UNIT_FLAG) { - AscendC::SetFlag(EVENT_ID0); - AscendC::WaitFlag(EVENT_ID0); - copyL0CToGm(gmC, l0CTensor, layoutBlock, layoutInL0C); - AscendC::SetFlag(EVENT_ID0); - } else { - copyL0CToGm(gmC, l0CTensor, layoutBlock, layoutInL0C, 0b11); - } - } - -protected: - CATLASS_DEVICE - void PrologueImpl( - typename Tile::PrologueTraits::TensorSrc const &gmSrcA, - typename Tile::PrologueTraits::LayoutSrc const &layoutSrcA, - typename Tile::PrologueTraits::TensorDst const &gmDstA, - typename Tile::PrologueTraits::LayoutDst const &layoutDstA, - typename Tile::PrologueTraits::TensorSrc const &gmSrcB, - typename Tile::PrologueTraits::LayoutSrc const &layoutSrcB, - typename Tile::PrologueTraits::TensorDst const &gmDstB, - typename Tile::PrologueTraits::LayoutDst const &layoutDstB, - GemmCoord const &actualShape) - { - uint32_t kTileCount = CeilDiv(actualShape.k()); - for (uint32_t kLoopIdx = 0; kLoopIdx < kTileCount; kLoopIdx++) { - uint32_t kOffset = kLoopIdx * L1TileShape::K; - uint32_t kActual = (kLoopIdx == kTileCount - 1) ? (actualShape.k() - kOffset) : L1TileShape::K; - - if constexpr (HAS_PROLOGUE_A) { - MatrixCoord offsetCoordA{0, kOffset}; - MatrixCoord actualTileShapeA{actualShape.m(), kActual}; - auto gmTileSrcA = gmSrcA[layoutSrcA.GetOffset(offsetCoordA)]; - auto layoutTileSrcA = layoutSrcA.GetTileLayout(actualTileShapeA); - auto gmTileDstA = gmDstA[l1ListId * layoutDstA.Capacity()]; - auto layoutTileDstA = layoutDstA.GetTileLayout(actualTileShapeA); - Catlass::Arch::CrossCoreWaitFlag(flagCopyAFinish[l1ListId]); - prologueA(gmTileDstA, layoutTileDstA, gmTileSrcA, layoutTileSrcA); - Catlass::Arch::CrossCoreSetFlag<0x02, PIPE_MTE3>(flagPrologueAFinish[l1ListId]); - } - - if constexpr (HAS_PROLOGUE_B) { - MatrixCoord offsetCoordB{kOffset, 0}; - MatrixCoord actualTileShapeB{kActual, actualShape.n()}; - auto gmTileSrcB = gmSrcB[layoutSrcB.GetOffset(offsetCoordB)]; - auto layoutTileSrcB = layoutSrcB.GetTileLayout(actualTileShapeB); - auto gmTileDstB = gmDstB[l1ListId * layoutDstB.Capacity()]; - auto layoutTileDstB = layoutDstB.GetTileLayout(actualTileShapeB); - Catlass::Arch::CrossCoreWaitFlag(flagCopyBFinish[l1ListId]); - prologueB(gmTileDstB, layoutTileDstB, gmTileSrcB, layoutTileSrcB); - Catlass::Arch::CrossCoreSetFlag<0x02, PIPE_MTE3>(flagPrologueBFinish[l1ListId]); - } - - l1ListId = (l1ListId + 1 == STAGES) ? 0 : (l1ListId + 1); - } - } - - /// Data members - Params params; - - AscendC::LocalTensor l1ATensorList[STAGES]; - AscendC::LocalTensor l1BTensorList[STAGES]; - AscendC::LocalTensor l0ATensorList[STAGES]; - AscendC::LocalTensor l0BTensorList[STAGES]; - AscendC::LocalTensor l0CTensor; - - int32_t l1AEventList[STAGES]; - int32_t l1BEventList[STAGES]; - int32_t l0AEventList[STAGES]; - int32_t l0BEventList[STAGES]; - - Arch::CrossCoreFlag flagPrologueAFinish[STAGES]; - Arch::CrossCoreFlag flagCopyAFinish[STAGES]; - Arch::CrossCoreFlag flagPrologueBFinish[STAGES]; - Arch::CrossCoreFlag flagCopyBFinish[STAGES]; - - uint32_t l1ListId{0}; - uint32_t l0AListId{0}; - uint32_t l0BListId{0}; - - TileMmad tileMmad; - CopyGmToL1A copyGmToL1A; - CopyGmToL1B copyGmToL1B; - CopyL1ToL0A copyL1ToL0A; - CopyL1ToL0B copyL1ToL0B; - CopyL0CToGm copyL0CToGm; - - Tile::PrologueTraits prologueA; - Tile::PrologueTraits prologueB; -}; - -} // namespace Catlass::Gemm::Block - -#endif // CATLASS_GEMM_BLOCK_BLOCK_MMAD_PINGPONG_WITH_PROLOGUE diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_preload_async_with_callback_tla.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_preload_async_with_callback_tla.hpp deleted file mode 100644 index b17a26df5c..0000000000 --- a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_preload_async_with_callback_tla.hpp +++ /dev/null @@ -1,650 +0,0 @@ -/** - * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of - * CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ -#ifndef CATLASS_GEMM_BLOCK_BLOCK_MMAD_PRELOAD_ASYNC_WITH_CALLBACK_TLA_HPP -#define CATLASS_GEMM_BLOCK_BLOCK_MMAD_PRELOAD_ASYNC_WITH_CALLBACK_TLA_HPP - -#include "catlass/arch/kda_resource.hpp" -#include "catlass/kda_catlass.hpp" -#include "catlass/kda_coord.hpp" -#include "catlass/detail/kda_callback.hpp" -#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" -#include "catlass/gemm/kda_gemm_helper.hpp" -#include "catlass/kda_gemm_coord.hpp" -#include "tla/kda_tla_layout.hpp" -#include "tla/kda_tensor.hpp" - -namespace Catlass::Gemm::Block { - -template < - class ArchTag_, - uint32_t PRELOAD_STAGES_, - uint32_t L1A_STAGES_, - uint32_t L1B_STAGES_, - uint32_t L0A_STAGES_, - uint32_t L0B_STAGES_, - uint32_t L0C_STAGES_, - bool ENABLE_UNIT_FLAG_, - bool ENABLE_SHUFFLE_K_, - bool ENABLE_L1_RESIDENT_, - bool USE_HF32_MODE_, - class L1TileShape_, - class L0TileShape_, - class ElementA_, - class ElementB_, - class ElementC_, - class ElementBias_, - class TileCopy_, - class TileMmad_> -struct BlockMmadTla< - MmadPreloadAsyncWithCallback< - ArchTag_, - PRELOAD_STAGES_, - L1A_STAGES_, - L1B_STAGES_, - L0A_STAGES_, - L0B_STAGES_, - L0C_STAGES_, - ENABLE_UNIT_FLAG_, - ENABLE_SHUFFLE_K_, - USE_HF32_MODE_, - ENABLE_L1_RESIDENT_>, - L1TileShape_, - L0TileShape_, - ElementA_, - ElementB_, - ElementC_, - ElementBias_, - TileCopy_, - TileMmad_> { -public: - // Type Aliases - using DispatchPolicy = MmadPreloadAsyncWithCallback< - ArchTag_, - PRELOAD_STAGES_, - L1A_STAGES_, - L1B_STAGES_, - L0A_STAGES_, - L0B_STAGES_, - L0C_STAGES_, - ENABLE_UNIT_FLAG_, - ENABLE_SHUFFLE_K_, - USE_HF32_MODE_, - ENABLE_L1_RESIDENT_>; - using ArchTag = typename DispatchPolicy::ArchTag; - using TileCopy = TileCopy_; - using L1TileShape = L1TileShape_; - using L0TileShape = L0TileShape_; - using ElementA = ElementA_; - using LayoutA = typename TileCopy::LayoutA; - using ElementB = ElementB_; - using LayoutB = typename TileCopy::LayoutB; - using ElementC = ElementC_; - using LayoutC = typename TileCopy::LayoutC; - using ElementBias = ElementBias_; - - using TileMmad = TileMmad_; - - using CopyL1ToL0A = typename TileCopy::CopyL1ToL0A; - using CopyL1ToL0B = typename TileCopy::CopyL1ToL0B; - using CopyL1ToBT = typename TileCopy::CopyL1ToBT; - - using ElementAccumulator = typename TileCopy::ElementAccumulator; - - static constexpr bool HAS_BIAS = TileCopy::HAS_BIAS; - - using LayoutTagL1A = typename TileCopy::LayoutTagL1A; - using LayoutTagL1B = typename TileCopy::LayoutTagL1B; - using LayoutTagL0A = typename TileCopy::LayoutTagL0A; - using LayoutTagL0B = typename TileCopy::LayoutTagL0B; - - static_assert( - tla::is_tuple::value && tla::is_static::value, - "L1TileShape must be tla::tuple and static!"); - static_assert( - tla::is_tuple::value && tla::is_static::value, - "L0TileShape must be tla::tuple and static!"); - - static constexpr uint32_t PRELOAD_STAGES = DispatchPolicy::PRELOAD_STAGES; - static constexpr uint32_t L1A_STAGES = DispatchPolicy::L1A_STAGES; - static constexpr uint32_t L1B_STAGES = DispatchPolicy::L1B_STAGES; - static constexpr uint32_t L0A_STAGES = DispatchPolicy::L0A_STAGES; - static constexpr uint32_t L0B_STAGES = DispatchPolicy::L0B_STAGES; - static constexpr uint32_t L0C_STAGES = DispatchPolicy::L0C_STAGES; - static constexpr uint32_t MIN_L1_STAGES = L1A_STAGES > L1B_STAGES ? L1B_STAGES : L1A_STAGES; - - static constexpr bool ENABLE_UNIT_FLAG = DispatchPolicy::ENABLE_UNIT_FLAG; - static constexpr bool ENABLE_SHUFFLE_K = DispatchPolicy::ENABLE_SHUFFLE_K; - static constexpr bool ENABLE_L1_RESIDENT = DispatchPolicy::ENABLE_L1_RESIDENT; - static constexpr bool USE_HF32_MODE = DispatchPolicy::USE_HF32_MODE; - - static constexpr uint32_t L1_TILE_M = tla::get<0>(L1TileShape{}); - static constexpr uint32_t L1_TILE_N = tla::get<1>(L1TileShape{}); - static constexpr uint32_t L1_TILE_K = tla::get<2>(L1TileShape{}); - static constexpr uint32_t L0_TILE_M = tla::get<0>(L0TileShape{}); - static constexpr uint32_t L0_TILE_N = tla::get<1>(L0TileShape{}); - static constexpr uint32_t L0_TILE_K = tla::get<2>(L0TileShape{}); - - // L1 tile size - static constexpr uint32_t L1A_TILE_SIZE = L1_TILE_M * L1_TILE_K * sizeof(ElementA); - static constexpr uint32_t L1B_TILE_SIZE = L1_TILE_N * L1_TILE_K * sizeof(ElementB); - // L0 tile size - static constexpr uint32_t L0A_TILE_SIZE = L0_TILE_M * L0_TILE_K * sizeof(ElementA); - static constexpr uint32_t L0B_TILE_SIZE = L0_TILE_K * L0_TILE_N * sizeof(ElementB); - static constexpr uint32_t L0C_TILE_SIZE = L1_TILE_M * L1_TILE_N * sizeof(ElementAccumulator); - - // Check HF32_MODE - static_assert( - !USE_HF32_MODE || (USE_HF32_MODE && std::is_same_v && std::is_same_v), - "HF32 MODE only supports in float!"); - - // Check LayoutC - static_assert( - tla::detail::isRowMajor::value || ((std::is_same_v || std::is_same_v || std::is_same_v) && tla::detail::iszN::value), - "LayoutC only supports zN in half or bfloat16 or float, RowMajor in all dtype yet!"); - - // Check L1TileShape - static_assert( - L1A_TILE_SIZE * L1A_STAGES + L1B_TILE_SIZE * L1B_STAGES <= ArchTag::L1_SIZE, - "L1TileShape exceeding the L1 space!"); - - // Check L0TileShape - static_assert(L0A_TILE_SIZE * L0A_STAGES <= ArchTag::L0A_SIZE, "L0TileShape exceeding the L0A space!"); - static_assert(L0B_TILE_SIZE * L0B_STAGES <= ArchTag::L0B_SIZE, "L0TileShape exceeding the L0B space!"); - static_assert(L0C_TILE_SIZE * L0C_STAGES <= ArchTag::L0C_SIZE, "L0TileShape exceeding the L0C space!"); - - static_assert( - L1_TILE_M == L0_TILE_M && L1_TILE_N == L0_TILE_N, - "The situation where the basic blocks of L1 and L0 differ on the m and n axes is not supported yet"); - static_assert(L0_TILE_K <= L1_TILE_K, "L0TileShape::K cannot exceed L1TileShape::K"); - static_assert(PRELOAD_STAGES <= MIN_L1_STAGES, - "PRELOAD_STAGES must not exceed the smaller of L1A_STAGES and L1B_STAGES!"); - - static_assert( - (!HAS_BIAS && (L1A_STAGES + L1B_STAGES) <= 8) || (HAS_BIAS && (L1A_STAGES + L1B_STAGES + MIN_L1_STAGES) <= 8), - "L1 Buffer overflow: Exceeds the supported range of EVENT(0~7)"); - static_assert( - (!HAS_BIAS && (L0A_STAGES + L0B_STAGES) <= 8) || (HAS_BIAS && (L0A_STAGES + L0B_STAGES) <= 7), - "L0 Buffer overflow: Exceeds the supported range of EVENT_ID(0~7)"); - - static constexpr auto L1A_LAYOUT = - tla::MakeLayout(tla::Int{}, tla::Int{}); - static constexpr auto L1B_LAYOUT = - tla::MakeLayout(tla::Int{}, tla::Int{}); - static constexpr auto L1BIAS_LAYOUT = tla::MakeLayout(tla::Int{}); - static constexpr auto L0BIAS_LAYOUT = tla::MakeLayout(tla::Int{}); - - CATLASS_DEVICE - BlockMmadTla(Arch::Resource &resource, uint32_t l1BufAddrStart = 0) - { - // use HF32 when USE_HF32_MODE is true - if constexpr (USE_HF32_MODE) { - AscendC::SetHF32Mode(true); - } else { - AscendC::SetHF32Mode(false); - } - if constexpr (ENABLE_UNIT_FLAG && tla::detail::isRowMajor::value) { - AscendC::SetMMLayoutTransform(true); - } - InitL1(resource, l1BufAddrStart); - InitL0A(resource); - InitL0B(resource); - InitL0C(resource); - if constexpr (HAS_BIAS) { - l0BiasTensor = resource.btBuf.template GetBufferByByte(0); - AscendC::SetFlag(L0A_STAGES + L0B_STAGES); - } - if constexpr (ENABLE_L1_RESIDENT) { - RestoreStatus(); - } - } - - CATLASS_DEVICE - ~BlockMmadTla() - { - if constexpr (USE_HF32_MODE) { - AscendC::SetHF32Mode(false); - } - if constexpr (ENABLE_UNIT_FLAG && tla::detail::isRowMajor::value) { - AscendC::SetMMLayoutTransform(false); - } - for (uint32_t i = 0; i < L1A_STAGES; ++i) { - AscendC::WaitFlag(l1AEventList[i]); - } - for (uint32_t i = 0; i < L1B_STAGES; ++i) { - AscendC::WaitFlag(l1BEventList[i]); - } - for (uint32_t i = 0; i < L0A_STAGES; ++i) { - AscendC::WaitFlag(l0AEventList[i]); - } - for (uint32_t i = 0; i < L0B_STAGES; ++i) { - AscendC::WaitFlag(l0BEventList[i]); - } - for (uint32_t i = 0; i < L0C_STAGES; ++i) { - AscendC::WaitFlag(l0CEventList[i]); - } - if constexpr (HAS_BIAS) { - for (uint32_t i = 0; i < MIN_L1_STAGES; ++i) { - AscendC::WaitFlag(l1BiasEventList[i]); - } - AscendC::WaitFlag(L0A_STAGES + L0B_STAGES); - } - } - - template - CATLASS_DEVICE void operator()( - TensorA &tensorA, - TensorB &tensorB, - TensorC &tensorC, - GemmCoord const &actualShape, - TensorBias const &tensorBias = {}, - Callback const &callbackBeforeFixpipe = {}, - Callback const &callbackAfterFixpipe = {}) - { - // Check L1TileShape - if constexpr (HAS_BIAS) { - static constexpr uint32_t BIAS_BUF_SIZE = L0_TILE_N * sizeof(ElementAccumulator); - static constexpr uint32_t L1BIAS_SIZE = L1_TILE_N * MIN_L1_STAGES * sizeof(ElementBias); - static_assert(BIAS_BUF_SIZE <= ArchTag::BIAS_SIZE, "BIAS_BUF_SIZE exceeding the BT space! Reduce L0_TILE_N"); - static_assert( - L1A_TILE_SIZE * L1A_STAGES + L1B_TILE_SIZE * L1B_STAGES + L1BIAS_SIZE <= ArchTag::L1_SIZE, - "L1TileShape exceeding the L1 space!"); - } - - using CopyGmToL1A = typename TileCopy::template CopyGmToL1A; - using CopyGmToL1B = typename TileCopy::template CopyGmToL1B; - CopyGmToL1A copyGmToL1A; - CopyGmToL1B copyGmToL1B; - - uint32_t mBlockActual = actualShape.m(); - uint32_t kBlockActual = actualShape.k(); - uint32_t nBlockActual = actualShape.n(); - - uint32_t kL1Loop = CeilDiv(kBlockActual); - - uint32_t mL1Actual = mBlockActual; - if constexpr (std::is_same_v) { - // Avoid using the gemv mode in mmad - if (mL1Actual == 1) { - mL1Actual = 16; - } - } - uint32_t nL1Actual = nBlockActual; - - uint32_t startTileIdx = 0; - if constexpr (ENABLE_SHUFFLE_K) { - startTileIdx = AscendC::GetBlockIdx() % kL1Loop; - } - - for (uint32_t kL1Idx = 0; kL1Idx < kL1Loop; ++kL1Idx) { - uint32_t kL1TileIdx = (startTileIdx + kL1Idx < kL1Loop) ? (startTileIdx + kL1Idx) : (startTileIdx + kL1Idx - kL1Loop); - - uint32_t kL1Actual = (kL1TileIdx < kL1Loop - 1) ? L1_TILE_K : (kBlockActual - kL1TileIdx * L1_TILE_K); - - // Load matrix A tile from GM to L1 - AscendC::WaitFlag(l1AEventList[l1AListId]); - auto tensorL1A = tla::MakeTensor(l1ATensorList[l1AListId], L1A_LAYOUT, Arch::PositionL1{}); - auto tensorTileA = - GetTile(tensorA, tla::MakeCoord(0, kL1TileIdx * L1_TILE_K), tla::MakeShape(mBlockActual, kL1Actual)); - if constexpr (ENABLE_L1_RESIDENT) { - // If the currently loaded GM pointer and block coordinates are the same as the last loaded ones, - // skip this loading. - if (lastAddrA[l1AListId] != tensorTileA.data().GetPhyAddr() || tla::get<0>(tensorTileA.coord()) != lastCoordA[l1AListId].row() || tla::get<1>(tensorTileA.coord()) != lastCoordA[l1AListId].column()) { - copyGmToL1A(tensorL1A, tensorTileA); - lastCoordA[l1AListId] = MatrixCoord{ - tla::get<0>(tensorTileA.coord()), tla::get<1>(tensorTileA.coord())}; - lastAddrA[l1AListId] = const_cast<__gm__ typename AscendC::GlobalTensor::PrimType *>( - tensorTileA.data().GetPhyAddr()); - } - } else { - copyGmToL1A(tensorL1A, tensorTileA); - } - AscendC::SetFlag(l1AEventList[l1AListId]); - // Load matrix B tile from GM to L1 - AscendC::WaitFlag(l1BEventList[l1BListId]); - auto tensorL1B = tla::MakeTensor(l1BTensorList[l1BListId], L1B_LAYOUT, Arch::PositionL1{}); - auto tensorTileB = - GetTile(tensorB, tla::MakeCoord(kL1TileIdx * L1_TILE_K, 0), tla::MakeShape(kL1Actual, nBlockActual)); - if constexpr (ENABLE_L1_RESIDENT) { - if (lastAddrB[l1BListId] != tensorTileB.data().GetPhyAddr() || tla::get<0>(tensorTileB.coord()) != lastCoordB[l1BListId].row() || tla::get<1>(tensorTileB.coord()) != lastCoordB[l1BListId].column()) { - copyGmToL1B(tensorL1B, tensorTileB); - lastCoordB[l1BListId] = MatrixCoord{ - tla::get<0>(tensorTileB.coord()), tla::get<1>(tensorTileB.coord())}; - lastAddrB[l1BListId] = const_cast<__gm__ typename AscendC::GlobalTensor::PrimType *>( - tensorTileB.data().GetPhyAddr()); - } - } else { - copyGmToL1B(tensorL1B, tensorTileB); - } - AscendC::SetFlag(l1BEventList[l1BListId]); - - if constexpr (HAS_BIAS && !std::is_same_v) { - if (kL1Idx == 0) { - using CopyGmToL1Bias = typename TileCopy::template CopyGmToL1Bias; - CopyGmToL1Bias copyGmToL1Bias; - AscendC::WaitFlag(l1BiasEventList[l1BiasListId]); - auto l1Bias = l1BiasTensor[l1BiasListId].template ReinterpretCast(); - auto tensorL1Bias = tla::MakeTensor(l1Bias, L1BIAS_LAYOUT, Arch::PositionL1{}); - copyGmToL1Bias(tensorL1Bias, tensorBias); - AscendC::SetFlag(l1BiasEventList[l1BiasListId]); - } - } - - // If the number of preload instructions reaches the upper limit, perform an mmad calculation on L1 tile - if (preloadCount == PRELOAD_STAGES) { - L1TileMmad(l1TileMmadParamsList[l1TileMmadParamsId]); - } - - // Store the current load status - uint32_t preloadL1TileMmadParamsId = (l1TileMmadParamsId + preloadCount < PRELOAD_STAGES) ? (l1TileMmadParamsId + preloadCount) : (l1TileMmadParamsId + preloadCount - PRELOAD_STAGES); - auto &l1TileMmadParams = l1TileMmadParamsList[preloadL1TileMmadParamsId]; - l1TileMmadParams.l1AListId = l1AListId; - l1TileMmadParams.l1BListId = l1BListId; - l1TileMmadParams.l1BiasListId = l1BiasListId; - l1TileMmadParams.mL1Actual = mL1Actual; - l1TileMmadParams.nL1Actual = nL1Actual; - l1TileMmadParams.kL1Actual = kL1Actual; - l1TileMmadParams.isKLoopFirst = (kL1Idx == 0); - l1TileMmadParams.isKLoopLast = (kL1Idx == kL1Loop - 1); - l1TileMmadParams.isNeedAddBias = !std::is_same_v; - if (kL1Idx == kL1Loop - 1) { - l1TileMmadParams.gmBlockC = tensorC.data(); - l1TileMmadParams.layoutCInGm = tensorC.layout(); - l1TileMmadParams.coord = tensorC.coord(); - l1TileMmadParams.callbackBeforeFixpipe = callbackBeforeFixpipe; - l1TileMmadParams.callbackAfterFixpipe = callbackAfterFixpipe; - } - - if (preloadCount < PRELOAD_STAGES) { - ++preloadCount; - } else { - l1TileMmadParamsId = (l1TileMmadParamsId + 1 < PRELOAD_STAGES) ? (l1TileMmadParamsId + 1) : 0; - } - l1AListId = (l1AListId + 1 < L1A_STAGES) ? (l1AListId + 1) : 0; - l1BListId = (l1BListId + 1 < L1B_STAGES) ? (l1BListId + 1) : 0; - if constexpr (HAS_BIAS && !std::is_same_v) { - if (kL1Idx == 0) { - l1BiasListId = (l1BiasListId + 1 < MIN_L1_STAGES) ? (l1BiasListId + 1) : 0; - } - } - } - } - - template - CATLASS_DEVICE void SynchronizeBlock() - { - while (preloadCount > 0) { - L1TileMmad(l1TileMmadParamsList[l1TileMmadParamsId]); - l1TileMmadParamsId = (l1TileMmadParamsId + 1 < PRELOAD_STAGES) ? (l1TileMmadParamsId + 1) : 0; - --preloadCount; - } - } - - // When enabling L1 resident mode, restore the pointer and coordinates that record the last state - // to the initial state. if two blockmmad instances need to be consecutively invoked at the kernel layer, - // RestoreStatus() must be inserted between them. - CATLASS_DEVICE - void RestoreStatus() - { - for (int i = 0; i < L1A_STAGES; ++i) { - lastAddrA[i] = nullptr; - lastCoordA[i] = MatrixCoord{0U, 0U}; - } - for (int i = 0; i < L1B_STAGES; ++i) { - lastAddrB[i] = nullptr; - lastCoordB[i] = MatrixCoord{0U, 0U}; - } - } - -private: - struct L1TileMmadParams { - uint32_t l1AListId; - uint32_t l1BListId; - uint32_t l1BiasListId; - uint32_t mL1Actual; - uint32_t nL1Actual; - uint32_t kL1Actual; - bool isKLoopFirst; - bool isKLoopLast; - bool isNeedAddBias; - AscendC::GlobalTensor gmBlockC; - LayoutC layoutCInGm; - tla::tuple coord; - Callback callbackBeforeFixpipe; - Callback callbackAfterFixpipe; - - CATLASS_DEVICE - L1TileMmadParams() = default; - }; - - CATLASS_DEVICE - void InitL1(Arch::Resource &resource, uint32_t l1BufAddrStart) - { - uint32_t l1AOffset = l1BufAddrStart; - uint32_t l1BOffset = l1BufAddrStart + L1A_TILE_SIZE * L1A_STAGES; - uint32_t l1BiasOffset = l1BOffset + L1B_TILE_SIZE * L1B_STAGES; - for (uint32_t i = 0; i < L1A_STAGES; ++i) { - l1ATensorList[i] = resource.l1Buf.template GetBufferByByte(l1AOffset + L1A_TILE_SIZE * i); - l1AEventList[i] = i; - AscendC::SetFlag(l1AEventList[i]); - } - for (uint32_t i = 0; i < L1B_STAGES; ++i) { - l1BTensorList[i] = resource.l1Buf.template GetBufferByByte(l1BOffset + L1B_TILE_SIZE * i); - l1BEventList[i] = i + L1A_STAGES; - AscendC::SetFlag(l1BEventList[i]); - } - for (uint32_t i = 0; i < MIN_L1_STAGES; ++i) { - if constexpr (HAS_BIAS) { - l1BiasTensor[i] = resource.l1Buf.template GetBufferByByte( - l1BiasOffset + L1_TILE_N * sizeof(ElementBias) * i); - l1BiasEventList[i] = i + L1A_STAGES + L1B_STAGES; - AscendC::SetFlag(l1BiasEventList[i]); - } - } - } - - CATLASS_DEVICE - void InitL0A(Arch::Resource &resource) - { - for (uint32_t i = 0; i < L0A_STAGES; ++i) { - l0ATensorList[i] = resource.l0ABuf.template GetBufferByByte(L0A_TILE_SIZE * i); - l0AEventList[i] = i; - AscendC::SetFlag(l0AEventList[i]); - } - } - - CATLASS_DEVICE - void InitL0B(Arch::Resource &resource) - { - for (uint32_t i = 0; i < L0B_STAGES; ++i) { - l0BTensorList[i] = resource.l0BBuf.template GetBufferByByte(L0B_TILE_SIZE * i); - l0BEventList[i] = i + L0A_STAGES; - AscendC::SetFlag(l0BEventList[i]); - } - } - - CATLASS_DEVICE - void InitL0C(Arch::Resource &resource) - { - for (uint32_t i = 0; i < L0C_STAGES; ++i) { - l0CTensorList[i] = resource.l0CBuf.template GetBufferByByte(L0C_TILE_SIZE * i); - l0CEventList[i] = i; - AscendC::SetFlag(l0CEventList[i]); - } - } - - template - CATLASS_DEVICE void L1TileMmad(L1TileMmadParams const ¶ms) - { - using CopyL0CToDst = typename TileCopy_::template CopyL0CToDst; - CopyL0CToDst copyL0CToDst; - - uint32_t kL0Loop = CeilDiv(params.kL1Actual); - auto &l1ATensor = l1ATensorList[params.l1AListId]; - auto &l1BTensor = l1BTensorList[params.l1BListId]; - auto tensorL1A = tla::MakeTensor(l1ATensor, L1A_LAYOUT, Arch::PositionL1{}); - auto tensorL1B = tla::MakeTensor(l1BTensor, L1B_LAYOUT, Arch::PositionL1{}); - - auto &l0CTensor = l0CTensorList[l0CListId]; - auto layoutInL0C = tla::MakeLayoutL0C(params.mL1Actual, params.nL1Actual); - auto tensorL0C = tla::MakeTensor(l0CTensor, layoutInL0C, Arch::PositionL0C{}); - auto tensorL0Bias = tla::MakeTensor(l0BiasTensor, L0BIAS_LAYOUT, Arch::PositionBias{}); - - if constexpr (!ENABLE_UNIT_FLAG) { - if (params.isKLoopFirst) { - AscendC::WaitFlag(l0CEventList[l0CListId]); - } - } - - for (uint32_t kL0Idx = 0; kL0Idx < kL0Loop; ++kL0Idx) { - uint32_t kL0Actual = (kL0Idx < kL0Loop - 1) ? L0_TILE_K : (params.kL1Actual - kL0Idx * L0_TILE_K); - - auto &l0ATile = l0ATensorList[l0AListId]; - auto layoutAInL0 = tla::MakeLayout(params.mL1Actual, kL0Actual); - auto tensorL0A = tla::MakeTensor(l0ATile, layoutAInL0, Arch::PositionL0A{}); - auto tensorTileL1A = - GetTile(tensorL1A, tla::MakeCoord(0, kL0Idx * L0_TILE_K), tla::MakeShape(params.mL1Actual, kL0Actual)); - - AscendC::WaitFlag(l0AEventList[l0AListId]); - if (kL0Idx == 0) { - AscendC::WaitFlag(l1AEventList[params.l1AListId]); - } - copyL1ToL0A(tensorL0A, tensorTileL1A); - if (kL0Idx == kL0Loop - 1) { - AscendC::SetFlag(l1AEventList[params.l1AListId]); - } - - auto &l0BTile = l0BTensorList[l0BListId]; - auto layoutBInL0 = tla::MakeLayout(kL0Actual, params.nL1Actual); - auto tensorL0B = tla::MakeTensor(l0BTile, layoutBInL0, Arch::PositionL0B{}); - // Locate the current tile of matrix B on L1 - auto tensorTileL1B = - GetTile(tensorL1B, tla::MakeCoord(kL0Idx * L0_TILE_K, 0), tla::MakeShape(kL0Actual, params.nL1Actual)); - - AscendC::WaitFlag(l0BEventList[l0BListId]); - if (kL0Idx == 0) { - AscendC::WaitFlag(l1BEventList[params.l1BListId]); - } - copyL1ToL0B(tensorL0B, tensorTileL1B); - if (kL0Idx == kL0Loop - 1) { - AscendC::SetFlag(l1BEventList[params.l1BListId]); - } - - bool initC = (params.isKLoopFirst && (kL0Idx == 0)); - if constexpr (HAS_BIAS) { - if (params.isNeedAddBias && initC) { - AscendC::WaitFlag(l1BiasEventList[params.l1BiasListId]); - AscendC::WaitFlag(L0A_STAGES + L0B_STAGES); - auto l1Bias = l1BiasTensor[params.l1BiasListId].template ReinterpretCast(); - auto tensorL1Bias = tla::MakeTensor(l1Bias, L1BIAS_LAYOUT, Arch::PositionL1{}); - auto tensorTileL1Bias = GetTile(tensorL1Bias, tla::MakeCoord(0), tla::MakeShape(params.nL1Actual)); - // Load bias to l0 biasTable - copyL1ToBT(tensorL0Bias, tensorTileL1Bias); - AscendC::SetFlag(l1BiasEventList[params.l1BiasListId]); - } - } - - AscendC::SetFlag(EVENT_ID0); - AscendC::WaitFlag(EVENT_ID0); - - // If the unit flag is enabled, the unit flag is set according to the calculation progress - uint8_t unitFlag = 0b00; - if constexpr (ENABLE_UNIT_FLAG) { - if (params.isKLoopLast && (kL0Idx == kL0Loop - 1)) { - unitFlag = 0b11; - } else { - unitFlag = 0b10; - } - } - - if constexpr (HAS_BIAS) { - if (params.isNeedAddBias && initC) { - tileMmad( - tensorL0C, tensorL0A, tensorL0B, tensorL0Bias, params.mL1Actual, params.nL1Actual, kL0Actual, - initC, unitFlag); - AscendC::SetFlag(L0A_STAGES + L0B_STAGES); - } else { - tileMmad( - tensorL0C, tensorL0A, tensorL0B, params.mL1Actual, params.nL1Actual, kL0Actual, initC, unitFlag); - } - } else { - tileMmad( - tensorL0C, tensorL0A, tensorL0B, params.mL1Actual, params.nL1Actual, kL0Actual, initC, unitFlag); - } - - AscendC::SetFlag(l0BEventList[l0BListId]); - l0BListId = (l0BListId + 1 < L0B_STAGES) ? (l0BListId + 1) : 0; - - AscendC::SetFlag(l0AEventList[l0AListId]); - l0AListId = (l0AListId + 1 < L0A_STAGES) ? (l0AListId + 1) : 0; - } - - if (params.isKLoopLast) { - auto layoutCInGm = params.layoutCInGm; - auto tensorTileC = - tla::MakeTensor(params.gmBlockC, layoutCInGm, params.coord, Arch::PositionType{}); - - params.callbackBeforeFixpipe(); - - if constexpr (!ENABLE_UNIT_FLAG) { - AscendC::SetFlag(l0CEventList[l0CListId]); - AscendC::WaitFlag(l0CEventList[l0CListId]); - copyL0CToDst(tensorTileC, tensorL0C); - AscendC::SetFlag(l0CEventList[l0CListId]); - } else { - copyL0CToDst(tensorTileC, tensorL0C, 0b11); - } - l0CListId = (l0CListId + 1 < L0C_STAGES) ? (l0CListId + 1) : 0; - - params.callbackAfterFixpipe(); - } - } - - AscendC::LocalTensor l1ATensorList[L1A_STAGES]; - AscendC::LocalTensor l1BTensorList[L1B_STAGES]; - int32_t l1AEventList[L1A_STAGES]; - int32_t l1BEventList[L1B_STAGES]; - uint32_t l1AListId{0}; - uint32_t l1BListId{0}; - uint32_t l1BiasListId{0}; - - AscendC::LocalTensor l0ATensorList[L0A_STAGES]; - int32_t l0AEventList[L0A_STAGES]; - uint32_t l0AListId{0}; - - AscendC::LocalTensor l0BTensorList[L0B_STAGES]; - int32_t l0BEventList[L0B_STAGES]; - uint32_t l0BListId{0}; - - AscendC::LocalTensor l0CTensorList[L0C_STAGES_]; - int32_t l0CEventList[L0C_STAGES_]; - uint32_t l0CListId{0}; - - AscendC::LocalTensor l1BiasTensor[MIN_L1_STAGES]; - int32_t l1BiasEventList[MIN_L1_STAGES]; - AscendC::LocalTensor l0BiasTensor; - - L1TileMmadParams l1TileMmadParamsList[PRELOAD_STAGES]; - uint32_t l1TileMmadParamsId{0}; - uint32_t preloadCount{0}; - - __gm__ typename AscendC::GlobalTensor::PrimType *lastAddrA[L1A_STAGES]; - __gm__ typename AscendC::GlobalTensor::PrimType *lastAddrB[L1B_STAGES]; - MatrixCoord lastCoordA[L1A_STAGES]; - MatrixCoord lastCoordB[L1B_STAGES]; - - TileMmad tileMmad; - CopyL1ToL0A copyL1ToL0A; - CopyL1ToL0B copyL1ToL0B; - CopyL1ToBT copyL1ToBT; -}; - -} // namespace Catlass::Gemm::Block - -#endif // CATLASS_GEMM_BLOCK_BLOCK_MMAD_PRELOAD_ASYNC_WITH_CALLBACK_TLA_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_preload_tla.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_preload_tla.hpp deleted file mode 100644 index 7bcdad82aa..0000000000 --- a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_preload_tla.hpp +++ /dev/null @@ -1,445 +0,0 @@ -/** - * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of - * CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ - -#ifndef CATLASS_GEMM_BLOCK_BLOCK_MMAD_PRELOAD_TLA_HPP -#define CATLASS_GEMM_BLOCK_BLOCK_MMAD_PRELOAD_TLA_HPP - -#include "catlass/kda_catlass.hpp" -#include "catlass/arch/kda_resource.hpp" -#include "catlass/kda_coord.hpp" -#include "catlass/detail/kda_callback.hpp" -#include "catlass/kda_gemm_coord.hpp" -#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" -#include "catlass/gemm/kda_gemm_helper.hpp" -#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" -#include "catlass/gemm/tile/kda_tile_mmad.hpp" -#include "tla/kda_tla_layout.hpp" -#include "tla/kda_tensor.hpp" - -namespace Catlass::Gemm::Block { - -template < - bool ENABLE_UNIT_FLAG_, - bool ENABLE_SHUFFLE_K_, - class L1TileShape_, - class L0TileShape_, - class TensorA_, - class TensorB_, - class TensorC_, - class TensorBias_, - class TileCopy_, - class TileMmad_> -struct BlockMmadTla< - MmadAtlasA2Preload, - L1TileShape_, - L0TileShape_, - TensorA_, - TensorB_, - TensorC_, - TensorBias_, - TileCopy_, - TileMmad_> { -public: - // Type Aliases - using DispatchPolicy = MmadAtlasA2Preload; - using ArchTag = typename DispatchPolicy::ArchTag; - using L1TileShape = L1TileShape_; - using L0TileShape = L0TileShape_; - using TensorA = TensorA_; - using TensorB = TensorB_; - using TensorC = TensorC_; - using ElementA = typename TensorA::Element; - using LayoutA = typename TensorA::Layout; - using ElementB = typename TensorB::Element; - using LayoutB = typename TensorB::Layout; - using ElementC = typename TensorC::Element; - using LayoutC = typename TensorC::Layout; - - using TileMmad = TileMmad_; - using CopyGmToL1A = typename TileCopy_::CopyGmToL1A; - using CopyGmToL1B = typename TileCopy_::CopyGmToL1B; - using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; - using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; - using CopyL0CToGm = typename TileCopy_::CopyL0CToGm; - using ElementAccumulator = typename CopyL0CToGm::ElementSrc; - - using LayoutTagL1A = typename TileCopy_::LayoutTagL1A; - using LayoutTagL1B = typename TileCopy_::LayoutTagL1B; - using LayoutTagL0A = typename TileCopy_::LayoutTagL0A; - using LayoutTagL0B = typename TileCopy_::LayoutTagL0B; - - using L1AAlignHelper = typename TileCopy_::L1AAlignHelper; - using L1BAlignHelper = typename TileCopy_::L1BAlignHelper; - - static_assert(tla::is_tuple::value && tla::is_static::value, - "L1TileShape must be tla::tuple and static!"); - static_assert(tla::is_tuple::value && tla::is_static::value, - "L0TileShape must be tla::tuple and static!"); - - static constexpr bool ENABLE_UNIT_FLAG = DispatchPolicy::ENABLE_UNIT_FLAG; - static constexpr bool ENABLE_SHUFFLE_K = DispatchPolicy::ENABLE_SHUFFLE_K; - static constexpr uint32_t STAGES = DispatchPolicy::STAGES; - static constexpr uint32_t L1_TILE_M = tla::get<0>(L1TileShape{}); - static constexpr uint32_t L1_TILE_N = tla::get<1>(L1TileShape{}); - static constexpr uint32_t L1_TILE_K = tla::get<2>(L1TileShape{}); - static constexpr uint32_t L0_TILE_M = tla::get<0>(L0TileShape{}); - static constexpr uint32_t L0_TILE_N = tla::get<1>(L0TileShape{}); - static constexpr uint32_t L0_TILE_K = tla::get<2>(L0TileShape{}); - - // L1 tile size - static constexpr uint32_t L1A_TILE_SIZE = L1_TILE_M * L1_TILE_K * sizeof(ElementA); - static constexpr uint32_t L1B_TILE_SIZE = L1_TILE_N * L1_TILE_K * sizeof(ElementB); - // L0 tile size - static constexpr uint32_t L0A_TILE_SIZE = L0_TILE_M * L0_TILE_K * sizeof(ElementA); - static constexpr uint32_t L0B_TILE_SIZE = L0_TILE_K * L0_TILE_N * sizeof(ElementB); - static constexpr uint32_t L0C_TILE_SIZE = L1_TILE_M * L1_TILE_N * sizeof(ElementAccumulator); - - // Check LayoutC - static_assert(tla::detail::isRowMajor::value, "LayoutC only support RowMajor yet!"); - - // Check L1TileShape - static_assert((L1A_TILE_SIZE + L1B_TILE_SIZE) * STAGES <= ArchTag::L1_SIZE, - "L1TileShape exceeding the L1 space!"); - - // Check L0TileShape - static_assert(L0A_TILE_SIZE * STAGES <= ArchTag::L0A_SIZE, "L0TileShape exceeding the L0A space!"); - static_assert(L0B_TILE_SIZE * STAGES <= ArchTag::L0B_SIZE, "L0TileShape exceeding the L0B space!"); - static_assert(L0C_TILE_SIZE <= ArchTag::L0C_SIZE, "L0TileShape exceeding the L0C space!"); - - static constexpr uint32_t _32B = 32 * 8; // in bits - static_assert(L1_TILE_M == L0_TILE_M && L1_TILE_N == L0_TILE_N, - "The situation where the basic blocks of L1 and L0 differ on the m and n axes is not supported yet"); - static_assert(L0_TILE_K <= L1_TILE_K, "L0TileShape::K cannot exceed L1TileShape::K"); - static_assert(L1_TILE_M * SizeOfBits::value % _32B == 0, "L1TileShape::M must be 32B aligned."); - static_assert(L1_TILE_K * SizeOfBits::value % _32B == 0, "L1TileShape::K must be 32B aligned."); - static_assert(L1_TILE_K * SizeOfBits::value % _32B == 0, "L1TileShape::K must be 32B aligned."); - static_assert(L1_TILE_N * SizeOfBits::value % _32B == 0, "L1TileShape::N must be 32B aligned."); - static_assert(L0_TILE_K * SizeOfBits::value % _32B == 0, "L0TileShape::K must be 32B aligned."); - - static constexpr auto L1A_LAYOUT = tla::MakeLayout(L1_TILE_M, L1_TILE_K); - static constexpr auto L1B_LAYOUT = tla::MakeLayout(L1_TILE_K, L1_TILE_N); - - /// Construct - CATLASS_DEVICE - BlockMmadTla(Arch::Resource &resource, uint32_t l1BufAddrStart = 0) - { - uint32_t l1AOffset = l1BufAddrStart; - uint32_t l1BOffset = l1BufAddrStart + L1A_TILE_SIZE * STAGES; - // Init buffers - for (uint32_t i = 0; i < STAGES; i++) { - // Assign L1/L0A/L0B space for each stages - l1ATensorList[i] = resource.l1Buf.template GetBufferByByte(l1AOffset + L1A_TILE_SIZE * i); - l1BTensorList[i] = resource.l1Buf.template GetBufferByByte(l1BOffset + L1B_TILE_SIZE * i); - l0ATensorList[i] = resource.l0ABuf.template GetBufferByByte(L0A_TILE_SIZE * i); - l0BTensorList[i] = resource.l0BBuf.template GetBufferByByte(L0B_TILE_SIZE * i); - - // Assign event ID for each stages - l1AEventList[i] = i; - l1BEventList[i] = i + STAGES; - l0AEventList[i] = i; - l0BEventList[i] = i + STAGES; - - // The event id that needs to be set before the loop - AscendC::SetFlag(l1AEventList[i]); - AscendC::SetFlag(l1BEventList[i]); - AscendC::SetFlag(l0AEventList[i]); - AscendC::SetFlag(l0BEventList[i]); - } - l0CTensor = resource.l0CBuf.template GetBufferByByte(0); - AscendC::SetFlag(EVENT_ID0); - } - - /// Destructor - CATLASS_DEVICE - ~BlockMmadTla() - { - for (uint32_t i = 0; i < STAGES; i++) { - AscendC::WaitFlag(l1AEventList[i]); - AscendC::WaitFlag(l1BEventList[i]); - AscendC::WaitFlag(l0AEventList[i]); - AscendC::WaitFlag(l0BEventList[i]); - } - AscendC::WaitFlag(EVENT_ID0); - } - - /// Perform a block-scoped matrix multiply-accumulate - template - CATLASS_DEVICE void operator()( - TensorA &tensorA, TensorB &tensorB, TensorC &tensorC, TensorA &tensorNextA, TensorB &tensorNextB, - GemmCoord const &actualShape, GemmCoord const &actualShapeNext, bool isFirstBlock, bool hasNextBlock, - Callback const &callbackBeforeFixpipe = {}, Callback const &callbackAfterFixpipe = {}) - { - uint32_t mBlockActual = actualShape.m(); - uint32_t kBlockActual = actualShape.k(); - uint32_t nBlockActual = actualShape.n(); - uint32_t mNextBlockActual = actualShapeNext.m(); - uint32_t kNextBlockActual = actualShapeNext.k(); - uint32_t nNextBlockActual = actualShapeNext.n(); - - uint32_t mRound = RoundUp(mBlockActual); - uint32_t nRound = RoundUp(nBlockActual); - - auto layoutInL0C = tla::MakeLayoutL0C(mRound, nRound); - auto tensorL0C = tla::MakeTensor(l0CTensor, layoutInL0C, Arch::PositionL0C{}); - - if constexpr (!ENABLE_UNIT_FLAG) { - AscendC::WaitFlag(EVENT_ID0); - } - uint32_t startTileIdx = 0; - if constexpr (ENABLE_SHUFFLE_K) { - startTileIdx = AscendC::GetBlockIdx(); - } - - uint32_t kTileCount = CeilDiv(kBlockActual); - uint32_t firstTileIdx = startTileIdx % kTileCount; - uint32_t lastTileIdx = (startTileIdx + kTileCount - 1) % kTileCount; - uint32_t kActual = - (firstTileIdx < kTileCount - 1) ? L1_TILE_K : (kBlockActual - firstTileIdx * L1_TILE_K); - uint32_t kTileCountNext = CeilDiv(kNextBlockActual); - uint32_t firstTileIdxNext = startTileIdx % kTileCountNext; - - if (isFirstBlock) { - // load first matrix A tile from GM to L1 - AscendC::WaitFlag(l1AEventList[l1ListId]); - auto tensorL1A = tla::MakeTensor(l1ATensorList[l1ListId], L1A_LAYOUT, Arch::PositionL1{}); - auto tensorTileA = GetTile( - tensorA, - tla::MakeCoord(0, firstTileIdx * L1_TILE_K), - tla::MakeShape(mBlockActual, kActual)); - copyGmToL1A(tensorL1A, tensorTileA, tla::MakeShape(mBlockActual, kActual)); - AscendC::SetFlag(l1AEventList[l1ListId]); - - // load first matrix B tile from GM to L1 - AscendC::WaitFlag(l1BEventList[l1ListId]); - auto tensorL1B = tla::MakeTensor(l1BTensorList[l1ListId], L1B_LAYOUT, Arch::PositionL1{}); - auto tensorTileB = GetTile( - tensorB, - tla::MakeCoord(firstTileIdx * L1_TILE_K, 0), - tla::MakeShape(kActual, nBlockActual)); - copyGmToL1B(tensorL1B, tensorTileB, tla::MakeShape(kActual, nBlockActual)); - AscendC::SetFlag(l1BEventList[l1ListId]); - } - - uint32_t mPartLoop = CeilDiv(mRound); - uint32_t nPartLoop = CeilDiv(nRound); - - // main loop - for (uint32_t kLoopIdx = 0; kLoopIdx < kTileCount; kLoopIdx++) { - uint32_t shuffleKIdx = (startTileIdx + kLoopIdx) % kTileCount; - uint32_t l1ListIdNext = (l1ListId + 1 < STAGES) ? (l1ListId + 1) : 0; - uint32_t kActualNext{0}; - - if (shuffleKIdx != lastTileIdx) { - uint32_t shuffleKIdxNext = (startTileIdx + kLoopIdx + 1) % kTileCount; - kActualNext = (shuffleKIdxNext < kTileCount - 1) ? - L1_TILE_K : - (kBlockActual - shuffleKIdxNext * L1_TILE_K); - - auto tensorL1A = tla::MakeTensor(l1ATensorList[l1ListIdNext], L1A_LAYOUT, Arch::PositionL1{}); - auto tensorL1B = tla::MakeTensor(l1BTensorList[l1ListIdNext], L1B_LAYOUT, Arch::PositionL1{}); - auto tensorTileA = GetTile( - tensorA, - tla::MakeCoord(0, shuffleKIdxNext * L1_TILE_K), - tla::MakeShape(mBlockActual, kActualNext)); - auto tensorTileB = GetTile( - tensorB, - tla::MakeCoord(shuffleKIdxNext * L1_TILE_K, 0), - tla::MakeShape(kActualNext, nBlockActual)); - - // load next matrix A tile from GM to L1 - AscendC::WaitFlag(l1AEventList[l1ListIdNext]); - copyGmToL1A(tensorL1A, tensorTileA, tla::MakeShape(mBlockActual, kActualNext)); - AscendC::SetFlag(l1AEventList[l1ListIdNext]); - - // load next matrix B tile from GM to L1 - AscendC::WaitFlag(l1BEventList[l1ListIdNext]); - copyGmToL1B(tensorL1B, tensorTileB, tla::MakeShape(kActualNext, nBlockActual)); - AscendC::SetFlag(l1BEventList[l1ListIdNext]); - } - - // preload next tile from GM to L1 - if (shuffleKIdx == lastTileIdx && hasNextBlock) { - kActualNext = (firstTileIdxNext < kTileCountNext - 1) ? - L1_TILE_K : - (kNextBlockActual - firstTileIdxNext * L1_TILE_K); - - // Get L1 tensor for next stage - auto tensorL1A = tla::MakeTensor(l1ATensorList[l1ListIdNext], L1A_LAYOUT, Arch::PositionL1{}); - auto tensorL1B = tla::MakeTensor(l1BTensorList[l1ListIdNext], L1B_LAYOUT, Arch::PositionL1{}); - // Get GM tile for next stage - auto tensorTileA = GetTile( - tensorNextA, - tla::MakeCoord(0, firstTileIdxNext * L1_TILE_K), - tla::MakeShape(mNextBlockActual, kActualNext)); - auto tensorTileB = GetTile( - tensorNextB, - tla::MakeCoord(firstTileIdxNext * L1_TILE_K, 0), - tla::MakeShape(kActualNext, nNextBlockActual)); - - // load next matrix A tile from GM to L1 - AscendC::WaitFlag(l1AEventList[l1ListIdNext]); - copyGmToL1A(tensorL1A, tensorTileA, tla::MakeShape(mNextBlockActual, kActualNext)); - AscendC::SetFlag(l1AEventList[l1ListIdNext]); - - // load next matrix B tile from GM to L1 - AscendC::WaitFlag(l1BEventList[l1ListIdNext]); - copyGmToL1B(tensorL1B, tensorTileB, tla::MakeShape(kActualNext, nNextBlockActual)); - AscendC::SetFlag(l1BEventList[l1ListIdNext]); - } - - // Get L1 tensor for current stage - auto l1ATensor = l1ATensorList[l1ListId]; - auto l1BTensor = l1BTensorList[l1ListId]; - auto tensorL1A = tla::MakeTensor(l1ATensor, L1A_LAYOUT, Arch::PositionL1{}); - auto tensorL1B = tla::MakeTensor(l1BTensor, L1B_LAYOUT, Arch::PositionL1{}); - // Get the loop nums on L0 - uint32_t kPartLoop = CeilDiv(kActual); - - for (int mPartIdx = 0; mPartIdx < mPartLoop; mPartIdx++) { - uint32_t mPartActual = (mPartIdx < mPartLoop - 1) ? - L0_TILE_M : - (mRound - mPartIdx * L0_TILE_M); - - for (int kPartIdx = 0; kPartIdx < kPartLoop; kPartIdx++) { - uint32_t kPartActual = (kPartIdx < kPartLoop - 1) ? - L0_TILE_K : - (kActual - kPartIdx * L0_TILE_K); - - // Locate the current tile on L0A - auto l0ATile = l0ATensorList[l0AListId]; - auto layoutAInL0 = tla::MakeLayout(mPartActual, kPartActual); - auto tensorL0A = tla::MakeTensor(l0ATile, layoutAInL0, Arch::PositionL0A{}); - // Locate the current tile of matrix A on L1 - auto tensorTileL1A = GetTile(tensorL1A, tla::MakeCoord(mPartIdx * L0_TILE_M, kPartIdx * L0_TILE_K), - tla::MakeShape(mPartActual, kPartActual)); - - AscendC::WaitFlag(l0AEventList[l0AListId]); - if ((mPartIdx == 0) && (kPartIdx == 0)) { - AscendC::WaitFlag(l1AEventList[l1ListId]); - } - - // Load current tile from L1 to L0A - copyL1ToL0A(tensorL0A, tensorTileL1A); - - if ((mPartIdx == mPartLoop - 1) && (kPartIdx == kPartLoop - 1)) { - AscendC::SetFlag(l1AEventList[l1ListId]); - } - - for (int nPartIdx = 0; nPartIdx < nPartLoop; nPartIdx++) { - uint32_t nPartActual = (nPartIdx < nPartLoop - 1) ? - L0_TILE_N : - (nRound - nPartIdx * L0_TILE_N); - - // Locate the current tile on L0B - auto l0BTile = l0BTensorList[l0BListId]; - auto layoutBInL0 = tla::MakeLayout(kPartActual, nPartActual); - auto tensorL0B = tla::MakeTensor(l0BTile, layoutBInL0, Arch::PositionL0B{}); - // Locate the current tile of matrix B on L1 - auto tensorTileL1B = GetTile(tensorL1B, - tla::MakeCoord(kPartIdx * L0_TILE_K, nPartIdx * L0_TILE_N), - tla::MakeShape(kPartActual, nPartActual)); - - // Wait for mmad finished - AscendC::WaitFlag(l0BEventList[l0BListId]); - // If the current tile is the first one on the k&n axis, wait for loading matrix B from GM to L1 - if ((kPartIdx == 0) && (nPartIdx == 0)) { - AscendC::WaitFlag(l1BEventList[l1ListId]); - } - - // Load current tile from L1 to L0B - copyL1ToL0B(tensorL0B, tensorTileL1B); - - // If the current tile is the last one on the k&n axis, notify to load matrix B from GM to L1 - if ((kPartIdx == kPartLoop - 1) && (nPartIdx == nPartLoop - 1)) { - AscendC::SetFlag(l1BEventList[l1ListId]); - } - // Notify to do mmad - AscendC::SetFlag(EVENT_ID0); - - // Locate the current tile on L0C - auto tensorTileL0C = GetTile(tensorL0C, - tla::MakeCoord(mPartIdx * L0_TILE_M, nPartIdx * L0_TILE_N), - tla::MakeShape(mPartActual, nPartActual)); - - // Compute the matrix multiplication on L0A and L0B and write the result to the accumulator - // Wait for loading L0B - AscendC::WaitFlag(EVENT_ID0); - - // If the current tile is the first tile on the k axis, the accumulator needs to be reset to 0 - bool initC = ((kLoopIdx == 0) && (kPartIdx == 0)); - // If the unit flag is enabled, the unit flag is set according to the calculation progress - uint8_t unitFlag = 0b00; - if constexpr (ENABLE_UNIT_FLAG) { - if ((kLoopIdx == kTileCount - 1) && (mPartIdx == mPartLoop - 1) && - (kPartIdx == kPartLoop - 1) && (nPartIdx == nPartLoop - 1)) { - unitFlag = 0b11; - } else { - unitFlag = 0b10; - } - } - // Perform calculation operations - tileMmad(tensorTileL0C, tensorL0A, tensorL0B, - mPartActual, nPartActual, kPartActual, initC, unitFlag); - - // Notify to move the next L0B tile - AscendC::SetFlag(l0BEventList[l0BListId]); - l0BListId = (l0BListId + 1 < STAGES) ? (l0BListId + 1) : 0; - } - AscendC::SetFlag(l0AEventList[l0AListId]); - l0AListId = (l0AListId + 1 < STAGES) ? (l0AListId + 1) : 0; - } - } - l1ListId = l1ListIdNext; - kActual = kActualNext; - } - - callbackBeforeFixpipe(); - // copy block out - if constexpr (!ENABLE_UNIT_FLAG) { - AscendC::SetFlag(EVENT_ID0); - AscendC::WaitFlag(EVENT_ID0); - copyL0CToGm(tensorC, tensorL0C); - AscendC::SetFlag(EVENT_ID0); - } else { - copyL0CToGm(tensorC, tensorL0C, 0b11); - } - callbackAfterFixpipe(); - } - -protected: - // Multi-stage tensors list - AscendC::LocalTensor l1ATensorList[STAGES]; - AscendC::LocalTensor l1BTensorList[STAGES]; - AscendC::LocalTensor l0ATensorList[STAGES]; - AscendC::LocalTensor l0BTensorList[STAGES]; - AscendC::LocalTensor l0CTensor; - - // Multi-stage event id list - int32_t l1AEventList[STAGES]; - int32_t l1BEventList[STAGES]; - int32_t l0AEventList[STAGES]; - int32_t l0BEventList[STAGES]; - - // The id of current stage - uint32_t l1ListId{0}; - uint32_t l0AListId{0}; - uint32_t l0BListId{0}; - - TileMmad tileMmad; - CopyGmToL1A copyGmToL1A; - CopyGmToL1B copyGmToL1B; - CopyL1ToL0A copyL1ToL0A; - CopyL1ToL0B copyL1ToL0B; - CopyL0CToGm copyL0CToGm; -}; - -} // namespace Catlass::Gemm::Block - -#endif // CATLASS_GEMM_BLOCK_BLOCK_MMAD_PRELOAD_TLA_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_single_core_splitk.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_single_core_splitk.hpp deleted file mode 100644 index 9101ffee2d..0000000000 --- a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_single_core_splitk.hpp +++ /dev/null @@ -1,391 +0,0 @@ -/** - * Copyright (c) 2025 Huawei Technologies Co., Ltd. - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of - * CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ - -#ifndef CATLASS_GEMM_BLOCK_BLOCK_MMAD_SINGLE_CORE_SPLITK_HPP -#define CATLASS_GEMM_BLOCK_BLOCK_MMAD_SINGLE_CORE_SPLITK_HPP - -#include "catlass/kda_catlass.hpp" -#include "catlass/arch/kda_resource.hpp" -#include "catlass/kda_coord.hpp" -#include "catlass/kda_gemm_coord.hpp" -#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" -#include "catlass/gemm/kda_gemm_helper.hpp" -#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" -#include "catlass/gemm/tile/kda_tile_mmad.hpp" - -namespace Catlass::Gemm::Block { - -template < - uint32_t L1A_STAGES_, - uint32_t L1B_STAGES_, - uint32_t L0C_STAGES_, - bool ENABLE_UNIT_FLAG_, - class L1TileShape_, - class L0TileShape_, - class AType_, - class BType_, - class CType_, - class BiasType_, - class TileCopy_, - class TileMmad_> -struct BlockMmad< - MmadAtlasA2SingleCoreSplitk, - L1TileShape_, - L0TileShape_, - AType_, - BType_, - CType_, - BiasType_, - TileCopy_, - TileMmad_> { -public: - // Type Aliases - using DispatchPolicy = MmadAtlasA2SingleCoreSplitk; - using ArchTag = typename DispatchPolicy::ArchTag; - using ElementA = typename AType_::Element; - using L1TileShape = L1TileShape_; - using L0TileShape = L0TileShape_; - using LayoutA = typename AType_::Layout; - using ElementB = typename BType_::Element; - using LayoutB = typename BType_::Layout; - using ElementC = typename CType_::Element; - using LayoutC = typename CType_::Layout; - using TileMmad = TileMmad_; - using CopyGmToL1A = typename TileCopy_::CopyGmToL1A; - using CopyGmToL1B = typename TileCopy_::CopyGmToL1B; - using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; - using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; - using CopyL0CToGm = typename TileCopy_::CopyL0CToGm; - using ElementAccumulator = - typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; - using LayoutAInL1 = typename CopyL1ToL0A::LayoutSrc; - using LayoutBInL1 = typename CopyL1ToL0B::LayoutSrc; - using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst; - using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst; - using LayoutCInL0 = layout::zN; - - using L1AAlignHelper = Gemm::helper::L1AlignHelper; - using L1BAlignHelper = Gemm::helper::L1AlignHelper; - - static constexpr bool ENABLE_UNIT_FLAG = DispatchPolicy::ENABLE_UNIT_FLAG; - static constexpr uint32_t L1A_STAGES = DispatchPolicy::L1A_STAGES; - static constexpr uint32_t L1B_STAGES = DispatchPolicy::L1B_STAGES; - static constexpr uint32_t L0AB_STAGES = DispatchPolicy::L0AB_STAGES; - static constexpr uint32_t L0C_STAGES = DispatchPolicy::L0C_STAGES; - - static constexpr uint32_t L1A_SIZE = L1TileShape::M * L1TileShape::K * sizeof(ElementA); - static constexpr uint32_t L1B_SIZE = L1TileShape::K * L1TileShape::N * sizeof(ElementB); - static constexpr uint32_t L0C_SIZE = L0TileShape::M * L0TileShape::N * sizeof(ElementC); - - // Check L1A & L1B - static_assert((L1A_SIZE * L1A_STAGES + L1B_SIZE * L1B_STAGES) <= ArchTag::L1_SIZE, - "L1TileShape exceeding the L1 space!"); - - static constexpr uint32_t L0A_TILE_SIZE = L0TileShape::M * L0TileShape::K * sizeof(ElementA); - static constexpr uint32_t L0B_TILE_SIZE = L0TileShape::K * L0TileShape::N * sizeof(ElementB); - static constexpr uint32_t L0C_TILE_SIZE = L0TileShape::M * L0TileShape::N * sizeof(ElementAccumulator); - - static_assert((L0C_TILE_SIZE * L0C_STAGES <= ArchTag::L0C_SIZE), "L0TileShape exceeding the L0C space!"); - static_assert(L0A_TILE_SIZE * L0AB_STAGES <= ArchTag::L0A_SIZE, "L0TileShape exceeding the L0A space!"); - static_assert(L0B_TILE_SIZE * L0AB_STAGES <= ArchTag::L0B_SIZE, "L0TileShape exceeding the L0B space!"); - - static_assert(L0TileShape::K <= L1TileShape::K, "L0TileShape::K cannot exceed L1TileShape::K"); - static_assert(L0TileShape::M <= L1TileShape::M, "L0TileShape::M cannot exceed L1TileShape::M"); - static_assert(L0TileShape::N <= L1TileShape::N, "L0TileShape::N cannot exceed L1TileShape::N"); - // 32B (256b) aligned - static_assert(Gemm::helper::TileShapeAlignChecker::_ALIGN == 256, - "Tile shape must be 32B aligned."); - - static constexpr uint32_t L0A_PINGPONG_BUF_SIZE = ArchTag::L0A_SIZE / L0AB_STAGES; - static constexpr uint32_t L0B_PINGPONG_BUF_SIZE = ArchTag::L0B_SIZE / L0AB_STAGES; - static constexpr uint32_t L0C_PINGPONG_BUF_SIZE = ArchTag::L0C_SIZE / L0C_STAGES; - - // Check LayoutC - static_assert(std::is_same_v, "LayoutC only support RowMajor yet!"); - - /// Construct - CATLASS_DEVICE - BlockMmad(Arch::Resource &resource, uint32_t l1BufAddrStart = 0) - { - uint32_t l1AOffset = l1BufAddrStart; - uint32_t l1BOffset = l1BufAddrStart + L1A_SIZE * L1A_STAGES; - // Init buffers - for (uint32_t i = 0; i < L1A_STAGES; i++) { - l1ATensorList[i] = resource.l1Buf.template GetBufferByByte(l1AOffset + L1A_SIZE * i); - l1AEventList[i] = i; - AscendC::SetFlag(l1AEventList[i]); - } - for (uint32_t i = 0; i < L1B_STAGES; i++) { - l1BTensorList[i] = resource.l1Buf.template GetBufferByByte(l1BOffset + L1B_SIZE * i); - l1BEventList[i] = i + L1A_STAGES; - AscendC::SetFlag(l1BEventList[i]); - } - - for (uint32_t i = 0; i < L0AB_STAGES; i++) { - l0ATensorList[i] = resource.l0ABuf.template GetBufferByByte(L0A_PINGPONG_BUF_SIZE * i); - l0BTensorList[i] = resource.l0BBuf.template GetBufferByByte(L0B_PINGPONG_BUF_SIZE * i); - l0AEventList[i] = i; - l0BEventList[i] = i + L0AB_STAGES; - AscendC::SetFlag(l0AEventList[i]); - AscendC::SetFlag(l0BEventList[i]); - } - for (uint32_t i = 0; i < L0C_STAGES; i++) { - l0CTensorList[i] = resource.l0CBuf.template GetBufferByByte(L0C_PINGPONG_BUF_SIZE * i); - l0CEventList[i] = i; - AscendC::SetFlag(l0CEventList[i]); - } - } - - /// Destructor - CATLASS_DEVICE - ~BlockMmad() - { - for (uint32_t i = 0; i < L1A_STAGES; i++) { - AscendC::WaitFlag(l1AEventList[i]); - } - for (uint32_t i = 0; i < L1B_STAGES; i++) { - AscendC::WaitFlag(l1BEventList[i]); - } - - for (uint32_t i = 0; i < L0AB_STAGES; i++) { - AscendC::WaitFlag(l0AEventList[i]); - AscendC::WaitFlag(l0BEventList[i]); - } - for (uint32_t i = 0; i < L0C_STAGES; i++) { - AscendC::WaitFlag(l0CEventList[i]); - } - } - - /// Perform a block-scoped matrix multiply-accumulate - CATLASS_DEVICE - void operator()(AscendC::GlobalTensor const &gmBlockA, LayoutA const &layoutA, - AscendC::GlobalTensor const &gmBlockB, LayoutB const &layoutB, - AscendC::GlobalTensor const &gmBlockC, LayoutC const &layoutC, - AscendC::GlobalTensor const &gmNextBlockA, AscendC::GlobalTensor const &gmNextBlockB, - GemmCoord const &actualShape, GemmCoord const &actualShapeNext, bool needLoadNextA, bool needLoadNextB, - bool atomicAdd) - { - auto layoutAInL1 = LayoutAInL1::template MakeLayout(L1TileShape::M, L1TileShape::K); - auto layoutBInL1 = LayoutBInL1::template MakeLayout(L1TileShape::K, L1TileShape::N); - - bool isFirstBlock = !preLoadedA && !preLoadedB; - bool isLastBlock = !needLoadNextA && !needLoadNextB; - uint32_t l1AListIdNext = (l1AListId + 1 < L1A_STAGES) ? (l1AListId + 1) : 0; - uint32_t l1BListIdNext = (l1BListId + 1 < L1B_STAGES) ? (l1BListId + 1) : 0; - - if (isFirstBlock) { - // Load matrix A tile from GM to L1 - AscendC::WaitFlag(l1AEventList[l1AListId]); - auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShape.m(), actualShape.k())); - copyGmToL1A(l1ATensorList[l1AListId], gmBlockA, layoutAInL1, layoutTileA); - AscendC::SetFlag(l1AEventList[l1AListId]); - - // Load matrix B tile from GM to L1 - AscendC::WaitFlag(l1BEventList[l1BListId]); - auto layoutTileB = layoutB.GetTileLayout(MakeCoord(actualShape.k(), actualShape.n())); - copyGmToL1B(l1BTensorList[l1BListId], gmBlockB, layoutBInL1, layoutTileB); - AscendC::SetFlag(l1BEventList[l1BListId]); - } - - if (needLoadNextA && (L1A_STAGES > 1)) { - AscendC::WaitFlag(l1AEventList[l1AListIdNext]); - auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShapeNext.m(), actualShapeNext.k())); - copyGmToL1A(l1ATensorList[l1AListIdNext], gmNextBlockA, layoutAInL1, layoutTileA); - AscendC::SetFlag(l1AEventList[l1AListIdNext]); - } - if (needLoadNextB && (L1B_STAGES > 1)) { - AscendC::WaitFlag(l1BEventList[l1BListIdNext]); - auto layoutTileB = layoutB.GetTileLayout(MakeCoord(actualShapeNext.k(), actualShapeNext.n())); - copyGmToL1B(l1BTensorList[l1BListIdNext], gmNextBlockB, layoutBInL1, layoutTileB); - AscendC::SetFlag(l1BEventList[l1BListIdNext]); - } - - // Get L1 tensor for current stage - auto l1ATensor = l1ATensorList[l1AListId]; - auto l1BTensor = l1BTensorList[l1BListId]; - - // Atomic set once in mmad level - if (atomicAdd) { - AscendC::SetAtomicAdd(); - } else { - AscendC::SetAtomicNone(); - } - - uint32_t l1mLoops = CeilDiv(actualShape.m(), L0TileShape::M); - uint32_t l1nLoops = CeilDiv(actualShape.n(), L0TileShape::N); - for (uint32_t l1mLoopIdx = 0; l1mLoopIdx < l1mLoops; l1mLoopIdx++) { - for (uint32_t l1nLoopIdx = 0; l1nLoopIdx < l1nLoops; l1nLoopIdx++) { - // layoutAInL0 - uint32_t mActual = (l1mLoopIdx < l1mLoops - 1) ? L0TileShape::M : (actualShape.m() - l1mLoopIdx * L0TileShape::M); - uint32_t nActual = (l1nLoopIdx < l1nLoops - 1) ? L0TileShape::N : (actualShape.n() - l1nLoopIdx * L0TileShape::N); - uint32_t mRound = RoundUp(mActual); - uint32_t nRound = RoundUp(nActual); - auto layoutInL0C = LayoutCInL0::MakeLayoutInL0C(MakeCoord(mRound, nRound)); - - uint32_t l0ABufId = 0; - uint32_t l0BBufId = 0; - - // wait/set conditions in block/l1mLoops/l1nLoops level - bool mnLoopStart = (l1mLoopIdx == 0 && l1nLoopIdx == 0); - bool mnLoopEnd = (l1mLoopIdx == l1mLoops - 1) && (l1nLoopIdx == l1nLoops - 1); - bool waitFlagA = (isFirstBlock || preLoadedA) && mnLoopStart; - bool setFlagA = (needLoadNextA || isLastBlock) && mnLoopEnd; - bool waitFlagB = (isFirstBlock || preLoadedB) && mnLoopStart; - bool setFlagB = (needLoadNextB || isLastBlock) && mnLoopEnd; - - if constexpr (!ENABLE_UNIT_FLAG) { - AscendC::WaitFlag(l0CEventList[l0CListId]); - } - - // Get the loop nums on L0 - uint32_t l1kLoops = CeilDiv(actualShape.k(), L0TileShape::K); - for (uint32_t l1kLoopIdx = 0; l1kLoopIdx < l1kLoops; l1kLoopIdx++) { - uint32_t kActual = - (l1kLoopIdx < l1kLoops - 1) ? L0TileShape::K : (actualShape.k() - l1kLoopIdx * L0TileShape::K); - - // Locate the current tile on L0A - auto l0ATile = l0ATensorList[l0ABufId]; - LayoutAInL0 layoutAInL0 = LayoutAInL0::template MakeLayout(mRound, kActual); - // Locate the current tile of matrix A on L1 - MatrixCoord l1ACoord{l1mLoopIdx * L0TileShape::M, l1kLoopIdx * L0TileShape::K}; - auto l1ATile = l1ATensor[layoutAInL1.GetOffset(l1ACoord)]; - - AscendC::WaitFlag(l0AEventList[l0ABufId]); - if (waitFlagA && (l1kLoopIdx == 0)) { - AscendC::WaitFlag(l1AEventList[l1AListId]); - } - - copyL1ToL0A(l0ATile, l1ATile, layoutAInL0, layoutAInL1); - - if (setFlagA && (l1kLoopIdx == l1kLoops - 1)) { - AscendC::SetFlag(l1AEventList[l1AListId]); - } - - // Locate the current tile on L0B - auto l0BTile = l0BTensorList[l0BBufId]; - LayoutBInL0 layoutBInL0 = LayoutBInL0::template MakeLayout(kActual, nRound); - // Locate the current tile of matrix B on L1 - MatrixCoord l1BCoord{l1kLoopIdx * L0TileShape::K, l1nLoopIdx * L0TileShape::N}; - auto l1BTile = l1BTensor[layoutBInL1.GetOffset(l1BCoord)]; - - // Wait for mmad finished - AscendC::WaitFlag(l0BEventList[l0BBufId]); - // If the current tile is the first one on the k&n axis, wait for loading matrix B from GM to L1 - if (waitFlagB && (l1kLoopIdx == 0)) { - AscendC::WaitFlag(l1BEventList[l1BListId]); - } - - // Load current tile from L1 to L0B - copyL1ToL0B(l0BTile, l1BTile, layoutBInL0, layoutBInL1); - - // If the current tile is the last one on the k&n axis, notify to load matrix B from GM to L1 - if (setFlagB && (l1kLoopIdx == l1kLoops - 1)) { - AscendC::SetFlag(l1BEventList[l1BListId]); - } - // Notify to do mmad - AscendC::SetFlag(l0CEventList[l0CListId]); - - // Locate the current tile on L0C - auto l0CTile = l0CTensorList[l0CListId]; - - // Compute the matrix multiplication on L0A and L0B and write the result to the accumulator - // Wait for loading L0B - AscendC::WaitFlag(l0CEventList[l0CListId]); - - // If the current tile is the first tile on the k axis, the accumulator needs to be reset to 0 - bool initC = (l1kLoopIdx == 0); - // If the unit flag is enabled, the unit flag is set according to the calculation progress - uint8_t unitFlag = 0b00; - if constexpr (ENABLE_UNIT_FLAG) { - if (l1kLoopIdx == l1kLoops - 1) { - unitFlag = 0b11; - } else { - unitFlag = 0b10; - } - } - // Perform calculation operations - tileMmad(l0CTile, l0ATile, l0BTile, mRound, nRound, kActual, initC, unitFlag); - - // Notify to move the next L0B tile - AscendC::SetFlag(l0BEventList[l0BBufId]); - l0BBufId = (l0BBufId + 1 < L0AB_STAGES) ? (l0BBufId + 1) : 0; - - AscendC::SetFlag(l0AEventList[l0ABufId]); - l0ABufId = (l0ABufId + 1 < L0AB_STAGES) ? (l0ABufId + 1) : 0; - } - // copy block out - LayoutC layoutBlock = layoutC.GetTileLayout(MakeCoord(mActual, nActual)); - MatrixCoord gmCoordC{l1mLoopIdx * L0TileShape::M, l1nLoopIdx * L0TileShape::N}; - auto gmTileC = gmBlockC[layoutC.GetOffset(gmCoordC)]; - if constexpr (!ENABLE_UNIT_FLAG) { - AscendC::SetFlag(l0CEventList[l0CListId]); - AscendC::WaitFlag(l0CEventList[l0CListId]); - copyL0CToGm(gmTileC, l0CTensorList[l0CListId], layoutBlock, layoutInL0C); - AscendC::SetFlag(l0CEventList[l0CListId]); - } else { - copyL0CToGm(gmTileC, l0CTensorList[l0CListId], layoutBlock, layoutInL0C, 0b11); - } - l0CListId = (l0CListId + 1 < L0C_STAGES) ? (l0CListId + 1) : 0; - } - } - - // if L1A_STAGES/L1B_STAGES == 1, load next block after compute - if (needLoadNextA && (L1A_STAGES == 1)) { - // Load first matrix A tile from GM to L1 - AscendC::WaitFlag(l1AEventList[l1AListIdNext]); - auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShapeNext.m(), actualShapeNext.k())); - copyGmToL1A(l1ATensorList[l1AListIdNext], gmNextBlockA, layoutAInL1, layoutTileA); - AscendC::SetFlag(l1AEventList[l1AListIdNext]); - } - if (needLoadNextB && (L1B_STAGES == 1)) { - // Load first matrix B tile from GM to L1 - AscendC::WaitFlag(l1BEventList[l1BListIdNext]); - auto layoutTileB = layoutB.GetTileLayout(MakeCoord(actualShapeNext.k(), actualShapeNext.n())); - copyGmToL1B(l1BTensorList[l1BListIdNext], gmNextBlockB, layoutBInL1, layoutTileB); - AscendC::SetFlag(l1BEventList[l1BListIdNext]); - } - - preLoadedA = needLoadNextA; - preLoadedB = needLoadNextB; - l1AListId = needLoadNextA ? l1AListIdNext : l1AListId; - l1BListId = needLoadNextB ? l1BListIdNext : l1BListId; - } - -protected: - /// Data members - AscendC::LocalTensor l1ATensorList[L1A_STAGES]; - AscendC::LocalTensor l1BTensorList[L1B_STAGES]; - AscendC::LocalTensor l0ATensorList[L0AB_STAGES]; - AscendC::LocalTensor l0BTensorList[L0AB_STAGES]; - AscendC::LocalTensor l0CTensorList[L0C_STAGES]; - - int32_t l1AEventList[L1A_STAGES]; - int32_t l1BEventList[L1B_STAGES]; - int32_t l0AEventList[L0AB_STAGES]; - int32_t l0BEventList[L0AB_STAGES]; - int32_t l0CEventList[L0C_STAGES]; - - bool preLoadedA{false}; - bool preLoadedB{false}; - uint32_t l1AListId{0}; - uint32_t l1BListId{0}; - uint32_t l0CListId{0}; - - TileMmad tileMmad; - CopyGmToL1A copyGmToL1A; - CopyGmToL1B copyGmToL1B; - CopyL1ToL0A copyL1ToL0A; - CopyL1ToL0B copyL1ToL0B; - CopyL0CToGm copyL0CToGm; -}; - -} // namespace Catlass::Gemm::Block - -#endif // CATLASS_GEMM_BLOCK_BLOCK_MMAD_SINGLE_CORE_SPLITK_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_small.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_small.hpp deleted file mode 100644 index d74e23afff..0000000000 --- a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_small.hpp +++ /dev/null @@ -1,319 +0,0 @@ -/** - * Copyright (c) 2025 Huawei Technologies Co., Ltd. - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of - * CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ - -#ifndef CATLASS_GEMM_BLOCK_BLOCK_MMAD_SMALL_HPP -#define CATLASS_GEMM_BLOCK_BLOCK_MMAD_SMALL_HPP - -#include "catlass/kda_catlass.hpp" -#include "catlass/arch/kda_resource.hpp" -#include "catlass/kda_coord.hpp" -#include "catlass/kda_gemm_coord.hpp" -#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" -#include "catlass/gemm/kda_gemm_helper.hpp" -#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" -#include "catlass/gemm/tile/kda_tile_mmad.hpp" - -namespace Catlass::Gemm::Block { - -template < - uint32_t STAGES_, - bool ENABLE_UNIT_FLAG_, - bool ENABLE_SHUFFLE_K_, - class L1TileShape_, - class L0TileShape_, - class AType_, - class BType_, - class CType_, - class BiasType_, - class TileCopy_, - class TileMmad_> -struct BlockMmad< - MmadAtlasA2Small, - L1TileShape_, - L0TileShape_, - AType_, - BType_, - CType_, - BiasType_, - TileCopy_, - TileMmad_> { -public: - // Type Aliases - using DispatchPolicy = MmadAtlasA2Small; - using ArchTag = typename DispatchPolicy::ArchTag; - using L1TileShape = L1TileShape_; - using L0TileShape = L0TileShape_; - using ElementA = typename AType_::Element; - using LayoutA = typename AType_::Layout; - using ElementB = typename BType_::Element; - using LayoutB = typename BType_::Layout; - using ElementC = typename CType_::Element; - using LayoutC = typename CType_::Layout; - using TileMmad = TileMmad_; - using CopyGmToL1A = typename TileCopy_::CopyGmToL1A; - using CopyGmToL1B = typename TileCopy_::CopyGmToL1B; - using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; - using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; - using CopyL0CToGm = typename TileCopy_::CopyL0CToGm; - using ElementAccumulator = - typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; - using LayoutAInL1 = typename CopyL1ToL0A::LayoutSrc; - using LayoutBInL1 = typename CopyL1ToL0B::LayoutSrc; - using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst; - using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst; - using LayoutCInL0 = layout::zN; - - using L1AAlignHelper = Gemm::helper::L1AlignHelper; - using L1BAlignHelper = Gemm::helper::L1AlignHelper; - - static constexpr bool ENABLE_UNIT_FLAG = DispatchPolicy::ENABLE_UNIT_FLAG; - static constexpr bool ENABLE_SHUFFLE_K = DispatchPolicy::ENABLE_SHUFFLE_K; - static constexpr uint32_t STAGES = DispatchPolicy::STAGES; - static constexpr uint32_t L1A_SIZE = L1TileShape::M * L1TileShape::K * sizeof(ElementA); - static constexpr uint32_t L1B_SIZE = L1TileShape::N * L1TileShape::K * sizeof(ElementB); - static constexpr uint32_t L0A_SIZE = ArchTag::L0A_SIZE; - static constexpr uint32_t L0B_SIZE = ArchTag::L0B_SIZE; - static constexpr uint32_t L0C_SIZE = ArchTag::L0C_SIZE; - static constexpr uint32_t L0A_PINGPONG_BUF_SIZE = L0A_SIZE / STAGES; - static constexpr uint32_t L0B_PINGPONG_BUF_SIZE = L0B_SIZE / STAGES; - - // Check LayoutC - static_assert(std::is_same_v, "LayoutC only support RowMajor yet!"); - - // Check L1TileShape - static_assert((L1A_SIZE * STAGES + L1B_SIZE * STAGES) <= ArchTag::L1_SIZE, "L1TileShape exceeding the L1 space!"); - - // Check L0TileShape - static constexpr uint32_t L0A_TILE_SIZE = L0TileShape::M * L0TileShape::K * sizeof(ElementA); - static constexpr uint32_t L0B_TILE_SIZE = L0TileShape::K * L0TileShape::N * sizeof(ElementB); - static constexpr uint32_t L0C_TILE_SIZE = L0TileShape::M * L0TileShape::N * sizeof(ElementAccumulator); - static_assert((L0A_TILE_SIZE * STAGES) <= L0A_SIZE, "L0TileShape exceeding the L0A space!"); - static_assert((L0B_TILE_SIZE * STAGES) <= L0B_SIZE, "L0TileShape exceeding the L0B space!"); - static_assert(L0C_TILE_SIZE <= L0C_SIZE, "L0TileShape exceeding the L0C space!"); - - static_assert(L1TileShape::M == L0TileShape::M && L1TileShape::N == L0TileShape::N, - "The situation where the basic blocks of L1 and L0 differ on the m and n axes is not supported yet"); - static_assert(L0TileShape::K <= L1TileShape::K, "L0TileShape::K cannot exceed L1TileShape::K"); - // 32B (256b) aligned - static_assert(Gemm::helper::TileShapeAlignChecker::_ALIGN == 256, - "Tile shape must be 32B aligned."); - - /// Construct - CATLASS_DEVICE - BlockMmad(Arch::Resource &resource, uint32_t l1BufAddrStart = 0) - { - uint32_t l1AOffset = l1BufAddrStart; - uint32_t l1BOffset = l1BufAddrStart + L1A_SIZE * STAGES; - // Init buffers - for (uint32_t i = 0; i < STAGES; i++) { - // Assign L1/L0A/L0B space for each stages - l1ATensorList[i] = resource.l1Buf.template GetBufferByByte(l1AOffset + L1A_SIZE * i); - l1BTensorList[i] = resource.l1Buf.template GetBufferByByte(l1BOffset + L1B_SIZE * i); - l0ATensorList[i] = resource.l0ABuf.template GetBufferByByte(L0A_PINGPONG_BUF_SIZE * i); - l0BTensorList[i] = resource.l0BBuf.template GetBufferByByte(L0B_PINGPONG_BUF_SIZE * i); - - // Assign event ID for each stages - l1AEventList[i] = i; - l1BEventList[i] = i + STAGES; - l0AEventList[i] = i; - l0BEventList[i] = i + STAGES; - - // The event id that needs to be set before the loop - AscendC::SetFlag(l1AEventList[i]); - AscendC::SetFlag(l1BEventList[i]); - AscendC::SetFlag(l0AEventList[i]); - AscendC::SetFlag(l0BEventList[i]); - } - l0CTensor = resource.l0CBuf.template GetBufferByByte(0); - if constexpr (!ENABLE_UNIT_FLAG) { - AscendC::SetFlag(EVENT_ID0); - } - } - - /// Destructor - CATLASS_DEVICE - ~BlockMmad() - { - for (uint32_t i = 0; i < STAGES; i++) { - AscendC::WaitFlag(l1AEventList[i]); - AscendC::WaitFlag(l1BEventList[i]); - AscendC::WaitFlag(l0AEventList[i]); - AscendC::WaitFlag(l0BEventList[i]); - } - if constexpr (!ENABLE_UNIT_FLAG) { - AscendC::WaitFlag(EVENT_ID0); - } - } - - /// Perform a block-scoped matrix multiply-accumulate - CATLASS_DEVICE - void operator()( - AscendC::GlobalTensor const &gmA, LayoutA const &layoutA, - AscendC::GlobalTensor const &gmB, LayoutB const &layoutB, - AscendC::GlobalTensor const &gmC, LayoutC const &layoutC, - GemmCoord const &actualShape) - { - uint32_t mRound = RoundUp(actualShape.m()); - uint32_t nRound = RoundUp(actualShape.n()); - - auto layoutAInL1 = LayoutAInL1::template MakeLayout(L1TileShape::M, L1TileShape::K); - auto layoutBInL1 = LayoutBInL1::template MakeLayout(L1TileShape::K, L1TileShape::N); - auto layoutInL0C = LayoutCInL0::MakeLayoutInL0C(MakeCoord(mRound, nRound)); - - if constexpr (!ENABLE_UNIT_FLAG) { - AscendC::WaitFlag(EVENT_ID0); - } - - uint32_t kActual = actualShape.k(); - - // k loop - { - // For small-shape matmul, the k-loop executes only once. - - // load first matrix A tile from GM to L1 - AscendC::WaitFlag(l1AEventList[l1ListId]); - auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShape.m(), kActual)); - copyGmToL1A(l1ATensorList[l1ListId], gmA, layoutAInL1, layoutTileA); - AscendC::SetFlag(l1AEventList[l1ListId]); - - // load first matrix B tile from GM to L1 - AscendC::WaitFlag(l1BEventList[l1ListId]); - auto layoutTileB = layoutB.GetTileLayout(MakeCoord(kActual, actualShape.n())); - copyGmToL1B(l1BTensorList[l1ListId], gmB, layoutBInL1, layoutTileB); - AscendC::SetFlag(l1BEventList[l1ListId]); - - // Get L1 tensor for current stage - auto l1ATensor = l1ATensorList[l1ListId]; - auto l1BTensor = l1BTensorList[l1ListId]; - - // Get the loop nums on L0 - uint32_t kPartLoop = CeilDiv(kActual); - - for (int kPartIdx = 0; kPartIdx < kPartLoop; kPartIdx++) { - uint32_t kPartActual = (kPartIdx < kPartLoop - 1) ? - L0TileShape::K : - (kActual - kPartIdx * L0TileShape::K); - - // Locate the current tile on L0A - auto l0ATile = l0ATensorList[l0AListId]; - LayoutAInL0 layoutAInL0 = LayoutAInL0::template MakeLayout(mRound, kPartActual); - // Locate the current tile of matrix A on L1 - MatrixCoord l1AOffset{0, kPartIdx * L0TileShape::K}; - auto l1ATile = l1ATensor[layoutAInL1.GetOffset(l1AOffset)]; - - AscendC::WaitFlag(l0AEventList[l0AListId]); - if (kPartIdx == 0) { - AscendC::WaitFlag(l1AEventList[l1ListId]); - } - - // Load current tile from L1 to L0A - copyL1ToL0A(l0ATile, l1ATile, layoutAInL0, layoutAInL1); - - if (kPartIdx == kPartLoop - 1) { - AscendC::SetFlag(l1AEventList[l1ListId]); - } - - // Locate the current tile on L0B - auto l0BTile = l0BTensorList[l0BListId]; - LayoutBInL0 layoutBInL0 = LayoutBInL0::template MakeLayout(kPartActual, nRound); - // Locate the current tile of matrix B on L1 - MatrixCoord l1BOffset{kPartIdx * L0TileShape::K, 0}; - auto l1BTile = l1BTensor[layoutBInL1.GetOffset(l1BOffset)]; - - // Wait for mmad finished - AscendC::WaitFlag(l0BEventList[l0BListId]); - // If the current tile is the first one on the k&n axis, wait for loading matrix B from GM to L1 - if (kPartIdx == 0) { - AscendC::WaitFlag(l1BEventList[l1ListId]); - } - - // Load current tile from L1 to L0B - copyL1ToL0B(l0BTile, l1BTile, layoutBInL0, layoutBInL1); - - // If the current tile is the last one on the k&n axis, notify to load matrix B from GM to L1 - if (kPartIdx == kPartLoop - 1) { - AscendC::SetFlag(l1BEventList[l1ListId]); - } - // Notify to do mmad - AscendC::SetFlag(EVENT_ID0); - - // Locate the current tile on L0C - auto l0CTile = l0CTensor; - - // Compute the matrix multiplication on L0A and L0B and write the result to the accumulator - // Wait for loading L0B - AscendC::WaitFlag(EVENT_ID0); - - // If the current tile is the first tile on the k axis, the accumulator needs to be reset to 0 - bool initC = (kPartIdx == 0); - // If the unit flag is enabled, the unit flag is set according to the calculation progress - uint8_t unitFlag = 0b00; - if constexpr (ENABLE_UNIT_FLAG) { - if (kPartIdx == kPartLoop - 1) { - unitFlag = 0b11; - } else { - unitFlag = 0b10; - } - } - // Perform calculation operations - tileMmad(l0CTile, l0ATile, l0BTile, mRound, nRound, kPartActual, initC, unitFlag); - - // Notify to move the next L0B tile - AscendC::SetFlag(l0BEventList[l0BListId]); - l0BListId = (l0BListId + 1 < STAGES) ? (l0BListId + 1) : 0; - - AscendC::SetFlag(l0AEventList[l0AListId]); - l0AListId = (l0AListId + 1 < STAGES) ? (l0AListId + 1) : 0; - } - } - - // copy block out - LayoutC layoutBlock = layoutC.GetTileLayout(actualShape.GetCoordMN()); - - if constexpr (!ENABLE_UNIT_FLAG) { - AscendC::SetFlag(EVENT_ID0); - AscendC::WaitFlag(EVENT_ID0); - copyL0CToGm(gmC, l0CTensor, layoutBlock, layoutInL0C); - AscendC::SetFlag(EVENT_ID0); - } else { - copyL0CToGm(gmC, l0CTensor, layoutBlock, layoutInL0C, 0b11); - } - } - -protected: - // Multi-stage tensors list - AscendC::LocalTensor l1ATensorList[STAGES]; - AscendC::LocalTensor l1BTensorList[STAGES]; - AscendC::LocalTensor l0ATensorList[STAGES]; - AscendC::LocalTensor l0BTensorList[STAGES]; - AscendC::LocalTensor l0CTensor; - - // Multi-stage event id list - int32_t l1AEventList[STAGES]; - int32_t l1BEventList[STAGES]; - int32_t l0AEventList[STAGES]; - int32_t l0BEventList[STAGES]; - - // The id of current stage - uint32_t l1ListId{0}; - uint32_t l0AListId{0}; - uint32_t l0BListId{0}; - - TileMmad tileMmad; - CopyGmToL1A copyGmToL1A; - CopyGmToL1B copyGmToL1B; - CopyL1ToL0A copyL1ToL0A; - CopyL1ToL0B copyL1ToL0B; - CopyL0CToGm copyL0CToGm; -}; - -} // namespace Catlass::Gemm::Block - -#endif // CATLASS_GEMM_BLOCK_BLOCK_MMAD_SMALL_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_sparse_tla.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_sparse_tla.hpp deleted file mode 100644 index 99433b669d..0000000000 --- a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_sparse_tla.hpp +++ /dev/null @@ -1,397 +0,0 @@ -/** - * Copyright (c) 2025 Huawei Technologies Co., Ltd. - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of - * CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ -#ifndef CATLASS_GEMM_BLOCK_BLOCK_MMAD_SPARSE_TLA_HPP -#define CATLASS_GEMM_BLOCK_BLOCK_MMAD_SPARSE_TLA_HPP - -#include "catlass/kda_catlass.hpp" -#include "catlass/arch/kda_resource.hpp" -#include "catlass/kda_coord.hpp" -#include "catlass/kda_gemm_coord.hpp" -#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" -#include "catlass/gemm/kda_gemm_helper.hpp" -#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" -#include "catlass/gemm/tile/kda_tile_mmad.hpp" -#include "tla/kda_tla_layout.hpp" -#include "tla/kda_tensor.hpp" - -namespace Catlass::Gemm::Block { - -template < - class ArchTag_, - bool ENABLE_UNIT_FLAG_, - class L1TileShape_, - class L0TileShape_, - class ElementA_, - class ElementB_, - class ElementC_, - class ElementBias_, - class TileCopy_> -struct BlockMmadSparseTla< - SparseMatmulMultiBlockOnKAxis, - L1TileShape_, - L0TileShape_, - ElementA_, - ElementB_, - ElementC_, - ElementBias_, - TileCopy_> { -public: - using DispatchPolicy = SparseMatmulMultiBlockOnKAxis; - using ArchTag = typename DispatchPolicy::ArchTag; - using L1Shape = L1TileShape_; - using L0Shape = L0TileShape_; - using ElementA = ElementA_; - using LayoutA = typename TileCopy_::LayoutA; - using ElementB = ElementB_; - using LayoutB = typename TileCopy_::LayoutB; - using ElementC = ElementC_; - using LayoutC = typename TileCopy_::LayoutC; - using ElementSparseIndex = uint8_t; - using ElementAccumulator = typename TileCopy_::ElementAccumulator; - - using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; - using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; - - using L1ALayout = typename TileCopy_::LayoutTagL1A; - using L1BLayout = typename TileCopy_::LayoutTagL1B; - using L1IdxLayout = typename TileCopy_::LayoutTagL1BIdx; // sparse index fractal 16*8 - using L0ALayout = detail::TagToLayout_t; - using L0BLayout = detail::TagToLayout_t; - using L0IdxLayout = detail::TagToLayout_t; - - using L1ATlaTensor = tla::Tensor, L1ALayout, - tla::Coord, AscendC::TPosition::GM>; - using L1BTlaTensor = tla::Tensor, L1BLayout, - tla::Coord, AscendC::TPosition::GM>; - - static constexpr uint16_t LIMIT_MNSIZE = 10; - static constexpr uint16_t ALIGN_NUM = 16; - static constexpr int32_t DENSE_MATRIX_B_OFFSET = 2; - static constexpr int32_t INDEX_MATRIX_OFFSET = 8; - static constexpr int32_t SHARE_LENS_COUNT = 3; - static constexpr auto L1_M = tla::get<0>(L1Shape{}); - static constexpr auto L1_N = tla::get<1>(L1Shape{}); - static constexpr auto L1_K = tla::get<2>(L1Shape{}); - static constexpr auto L0_M = tla::get<0>(L0Shape{}); - static constexpr auto L0_N = tla::get<1>(L0Shape{}); - static constexpr auto L0_K = tla::get<2>(L0Shape{}); - static constexpr uint32_t STAGES = DispatchPolicy::STAGES; - - // L1 tile size - static constexpr uint32_t L1A_TILE_SIZE = L1_M * L1_K * sizeof(ElementA); - static constexpr uint32_t L1B_TILE_SIZE = L1_N * L1_K * sizeof(ElementB) / DENSE_MATRIX_B_OFFSET; - static constexpr uint32_t L1B_IDX_TILE_SIZE = L1_N * L1_K * sizeof(ElementB) / INDEX_MATRIX_OFFSET; - // L0 tile size - static constexpr uint32_t L0A_TILE_SIZE = L0_M * L0_K * sizeof(ElementA); - static constexpr uint32_t L0B_TILE_SIZE = L0_K * L0_N * sizeof(ElementB); - static constexpr uint32_t L0C_TILE_SIZE = L0_M * L0_N * sizeof(ElementAccumulator); - - static constexpr uint32_t L0A_SIZE = ArchTag::L0A_SIZE; - static constexpr uint32_t L0B_SIZE = ArchTag::L0B_SIZE; - static constexpr uint32_t L0C_SIZE = ArchTag::L0C_SIZE; - static constexpr uint32_t L0A_PINGPONG_BUF_SIZE = L0A_SIZE / STAGES; - static constexpr uint32_t L0B_PINGPONG_BUF_SIZE = L0B_SIZE / STAGES; - - static constexpr bool TRANS_A = tla::detail::isColumnMajor::value; - static constexpr bool TRANS_B = tla::detail::isColumnMajor::value; - - // Check L1TileShape - static_assert((L1A_TILE_SIZE + L1B_TILE_SIZE + L1B_IDX_TILE_SIZE) * STAGES <= ArchTag::L1_SIZE, - "L1TileShape exceeding the L1 space!"); - - // Check L0TileShape - static_assert(L0A_TILE_SIZE <= ArchTag::L0A_SIZE, "L0TileShape exceeding the L0A space!"); - static_assert(L0B_TILE_SIZE <= ArchTag::L0B_SIZE, "L0TileShape exceeding the L0B space!"); - static_assert(L0C_TILE_SIZE <= ArchTag::L0C_SIZE, "L0TileShape exceeding the L0C space!"); - - static_assert(L1_M == L0_M && L1_N == L0_N, - "The situation where the basic blocks of L1 and L0 differ on the m and n axes is not supported yet"); - static_assert(L0_K <= L1_K, "L0TileShape::K cannot exceed L1TileShape::K"); - - static_assert((std::is_same_v && std::is_same_v && - std::is_same_v), - "Unsupported dtype"); - static_assert(tla::detail::isColumnMajor::value, "Only support MatrixB transpose"); - - CATLASS_DEVICE - BlockMmadSparseTla() - { - Init(); - } - - CATLASS_DEVICE - ~BlockMmadSparseTla() - { - End(); - } - - /** - * @brief Initialization function - * - * This function initializes the buffers and parameters required for matrix multiplication operations - */ - CATLASS_DEVICE - void Init() - { - constexpr static int32_t aL0MatrixByteSize = - L0_M * L0_K * sizeof(ElementA); - constexpr static int32_t bL0MatrixByteSize = - L0_N * L0_K * sizeof(ElementB); - constexpr static int32_t cMatrixByteSize = - L0_M * L0_N * sizeof(ElementAccumulator); - - constexpr static int dbL0AFlag = (aL0MatrixByteSize * STAGES > ArchTag::L0A_SIZE) ? 1 : STAGES; - constexpr static int dbL0BFlag = (bL0MatrixByteSize * STAGES > ArchTag::L0B_SIZE) ? 1 : STAGES; - - uint32_t l1AOffset = 0; - uint32_t l1BOffset = L1A_TILE_SIZE * STAGES; - uint32_t l1BIdxOffset = l1BOffset + L1B_TILE_SIZE * STAGES; - // Init buffers - for (uint32_t i = 0; i < STAGES; i++) { - // Assign L1/L0A/L0B space for each stages - l1ATensorList_[i] = l1Buf_[l1AOffset + L1A_TILE_SIZE * i].template ReinterpretCast(); - l1BTensorList_[i] = l1Buf_[l1BOffset + L1B_TILE_SIZE * i].template ReinterpretCast(); - l1BIdxTensorList_[i] = l1Buf_[l1BIdxOffset + L1B_IDX_TILE_SIZE * i].template ReinterpretCast(); - l0ATensorList_[i] = l0ABuf_[L0A_PINGPONG_BUF_SIZE * i].template ReinterpretCast(); - l0BTensorList_[i] = l0BBuf_[L0B_PINGPONG_BUF_SIZE * i].template ReinterpretCast(); - - // Assign event ID for each stages - l1AEventList_[i] = i; - l1BEventList_[i] = i + STAGES; - l1BIdxEventList_[i] = i + STAGES * 2; - l0AEventList_[i] = i; - l0BEventList_[i] = i + STAGES; - - // The event id that needs to be set before the loop - AscendC::SetFlag(l1AEventList_[i]); - AscendC::SetFlag(l1BEventList_[i]); - AscendC::SetFlag(l1BIdxEventList_[i]); - AscendC::SetFlag(l0AEventList_[i]); - AscendC::SetFlag(l0BEventList_[i]); - } - l0CTensor_ = l0CBuf_.template ReinterpretCast(); - AscendC::SetFlag(EVENT_ID0); - } - - /** - * @brief Perform the matrix multiplication operations - * @param [in] DstTensor: destination tensor type - * @param [in] SrcATensor: matrix A source tensor type - * @param [in] SrcBTensor: matrix B source tensor type - * @param [in] SparseIdxTensor: the type of the sparse index tensor - * @param [in] Shape: shape type - * @param [in] c: destination tensor - * @param [in] a: matrix A source tensor - * @param [in] b: matrix B source tensor - * @param [in] sparseIdx: the sparse index tensor - * @param [in] actualShape: actual shape - */ - template - CATLASS_DEVICE void operator()(DstTensor &c, const SrcATensor &a, const SrcBTensor &b, - const SparseIdxTensor &sparseIdx, const Shape &actualShape) - { - using CopyL0CToGm = typename TileCopy_::template CopyL0CToGm; - CopyL0CToGm copyL0CToGm; - - int32_t mIter = CeilDiv(tla::get<0>(actualShape), L0_M); - int32_t nIter = CeilDiv(tla::get<1>(actualShape), L0_N); - int32_t kOuterIter = CeilDiv(tla::get<2>(actualShape), L1_K); - int32_t tailBaseM = (tla::get<0>(actualShape) % L0_M) == 0 ? L0_M : (tla::get<0>(actualShape) % L0_M); - int32_t tailBaseN = (tla::get<1>(actualShape) % L0_N) == 0 ? L0_N : (tla::get<1>(actualShape) % L0_N); - int32_t tailBaseK = (tla::get<2>(actualShape) % L0_K) == 0 ? L0_K : (tla::get<2>(actualShape) % L0_K); - int32_t tailL1K = (tla::get<2>(actualShape) % L1_K) == 0 ? L1_K : (tla::get<2>(actualShape) % L1_K); - - for (auto mIndex = 0; mIndex < mIter; ++mIndex) { - for (auto nIndex = 0; nIndex < nIter; ++nIndex) { - int mPartActual = (mIndex + 1 == mIter) ? tailBaseM : L0_M; - int nPartActual = (nIndex + 1 == nIter) ? tailBaseN : L0_N; - auto l0CLayout = tla::MakeLayoutL0C(mPartActual, nPartActual); - AscendC::WaitFlag(EVENT_ID0); - auto l0CTlaTensor = tla::MakeTensor(l0CTensor_, l0CLayout, Arch::PositionL0C{}); - - for (auto kOuterIndex = 0; kOuterIndex < kOuterIter; ++kOuterIndex) { - uint32_t l1ListIdNext = (l1ListId + 1 < STAGES) ? (l1ListId + 1) : 0; - int l1K = (kOuterIndex + 1 == kOuterIter) ? tailL1K : L1_K; - // -----------------Step1: GM -> L1 ----------------- - auto aTileHeight = mPartActual; - auto aTileWidth = l1K; - auto aL1Row = mIndex; - auto aL1Col = kOuterIndex; - - auto l1ALayout = tla::MakeLayout(aTileHeight, aTileWidth); - AscendC::WaitFlag(l1AEventList_[l1ListIdNext]); - using CopyGmToL1A = typename TileCopy_::template CopyGmToL1A; - - auto l1ATensor = tla::MakeTensor(l1ATensorList_[l1ListIdNext], l1ALayout, Arch::PositionL1{}); - auto SrcATlaTensor = tla::MakeTensor(a.data(), a.layout(), - tla::MakeCoord(aL1Row * L0_M, - aL1Col * L1_K), - Arch::PositionGM{}); - CopyGmToL1A copyGmToL1A; - copyGmToL1A(l1ATensor, SrcATlaTensor); - - AscendC::SetFlag(l1AEventList_[l1ListIdNext]); - - auto bTileHeight = l1K; - auto bTileWidth = nPartActual; - auto bL1Row = kOuterIndex; - auto bL1Col = nIndex; - using CopyGmToL1B = typename TileCopy_::template CopyGmToL1B; - using CopyGmToL1BIdx = typename TileCopy_::template CopyGmToL1BIdx; - CopyGmToL1B copyGmToL1B; - CopyGmToL1BIdx copyGmToL1BIdx; - - auto l1BLayout = tla::MakeLayout(bTileHeight / DENSE_MATRIX_B_OFFSET, bTileWidth); - - AscendC::WaitFlag(l1BEventList_[l1ListIdNext]); - - auto l1BTensor = tla::MakeTensor(l1BTensorList_[l1ListIdNext], l1BLayout, Arch::PositionL1{}); - auto srcBTlaTensor = GetTile(b, tla::MakeCoord(bL1Row * L1_K / DENSE_MATRIX_B_OFFSET, bL1Col * L0_N), - tla::MakeShape(l1K / DENSE_MATRIX_B_OFFSET, nPartActual)); - - auto l1BIndexLayout = - tla::MakeLayout(bTileHeight / INDEX_MATRIX_OFFSET, bTileWidth); - AscendC::WaitFlag(l1BIdxEventList_[l1ListIdNext]); - - auto l1BIdxTensor = tla::MakeTensor(l1BIdxTensorList_[l1ListIdNext], l1BIndexLayout, Arch::PositionL1{}); - auto sparseIdxTlaTensor = tla::MakeTensor(sparseIdx.data(), sparseIdx.layout(), - tla::MakeCoord(bL1Row * L1_K / INDEX_MATRIX_OFFSET, bL1Col * L0_N), - Arch::PositionGM{}); - - copyGmToL1B(l1BTensor, srcBTlaTensor); - - copyGmToL1BIdx(l1BIdxTensor, sparseIdxTlaTensor); - - AscendC::SetFlag(l1BEventList_[l1ListIdNext]); - AscendC::SetFlag(l1BIdxEventList_[l1ListIdNext]); - AscendC::WaitFlag(l1AEventList_[l1ListIdNext]); - AscendC::WaitFlag(l1BEventList_[l1ListIdNext]); - AscendC::WaitFlag(l1BIdxEventList_[l1ListIdNext]); - - int32_t kInnerIter = - (kOuterIndex + 1 == kOuterIter) ? CeilDiv(tailL1K, L0_K) : CeilDiv(L1_K, L0_K); - for (auto kInnerIndex = 0; kInnerIndex < kInnerIter; ++kInnerIndex) { - int baseK = - (kOuterIndex + 1 == kOuterIter) && (kInnerIndex + 1 == kInnerIter) ? tailBaseK : L0_K; - - // -----------------Step2: L1 -> L0----------------- - AscendC::WaitFlag(l0AEventList_[l0AListId]); - - auto aBaseHeight = mPartActual; - auto aBaseWidth = baseK; - auto aL0Row = mIndex; - auto aL0Col = kInnerIndex; - auto l0ALayout = tla::MakeLayout(aBaseHeight, aBaseWidth); - auto l0ATensor = tla::MakeTensor(l0ATensorList_[l0AListId], l0ALayout, Arch::PositionL0A{}); - - auto l1ATileTensor = tla::MakeTensor(l1ATensor.data(), l1ATensor.layout(), - tla::MakeCoord(aL0Row * 0, - aL0Col * L0_K), - Arch::PositionL1{}); - copyL1ToL0A_(l0ATensor, l1ATileTensor); - - auto bBaseHeight = baseK; - auto bBaseWidth = nPartActual; - auto bL0Row = kInnerIndex; - auto bL0Col = nIndex; - auto l0BLayout = tla::MakeLayout( - bBaseHeight / DENSE_MATRIX_B_OFFSET, bBaseWidth); - - auto l0BTensor = tla::MakeTensor(l0BTensorList_[l0BListId], l0BLayout, Arch::PositionL0B{}); - auto l1BTileTensor = tla::MakeTensor(l1BTensor.data(), l1BTensor.layout(), - tla::MakeCoord(bL0Row * L0_K / DENSE_MATRIX_B_OFFSET, 0), - Arch::PositionL1{}); - - copyL1ToL0B_(l0BTensor, l1BTileTensor, l1BIdxTensor); - - AscendC::SetFlag(l0AEventList_[l0AListId]); - AscendC::WaitFlag(l0AEventList_[l0AListId]); - - // -----------------Step3: compute----------------- - AscendC::MmadParams mmadParams; - // GEMV is automatically enabled when setting M=1 in normal mode - mmadParams.m = (mPartActual == 1 ? ALIGN_NUM : mPartActual); - mmadParams.k = baseK; - mmadParams.n = nPartActual; - mmadParams.unitFlag = 0; - mmadParams.cmatrixInitVal = (kOuterIndex == 0) && (kInnerIndex == 0) ? true : false; - MmadWithSparse(l0CTensor_, l0ATensorList_[l0AListId], l0BTensorList_[l0BListId], mmadParams); - if ((mPartActual / ALIGN_NUM) * (nPartActual / ALIGN_NUM) < LIMIT_MNSIZE) { - AscendC::PipeBarrier(); - } - AscendC::SetFlag(l0AEventList_[l0AListId]); - l0AListId = (l0AListId + 1 < STAGES) ? (l0AListId + 1) : 0; - l0BListId = (l0BListId + 1 < STAGES) ? (l0BListId + 1) : 0; - } - AscendC::SetFlag(l1AEventList_[l1ListIdNext]); - AscendC::SetFlag(l1BEventList_[l1ListIdNext]); - AscendC::SetFlag(l1BIdxEventList_[l1ListIdNext]); - l1ListId = l1ListIdNext; - } - AscendC::SetFlag(EVENT_ID0); - AscendC::WaitFlag(EVENT_ID0); - // -----------------Step4: L0C -> GM----------------- - auto l0CSrcTensor = tla::MakeTensor(l0CTlaTensor.data(), l0CTlaTensor.layout(), - tla::MakeCoord(mIndex * L0_M, nIndex * L0_N), - Arch::PositionL0C{}); - copyL0CToGm(c, l0CSrcTensor); - - AscendC::SetFlag(EVENT_ID0); - } - } - } - -private: - /** - * @brief End function, release all events - */ - CATLASS_DEVICE - void End() - { - for (uint32_t i = 0; i < STAGES; i++) { - AscendC::WaitFlag(l1AEventList_[i]); - AscendC::WaitFlag(l1BEventList_[i]); - AscendC::WaitFlag(l1BIdxEventList_[i]); - AscendC::WaitFlag(l0AEventList_[i]); - AscendC::WaitFlag(l0BEventList_[i]); - } - AscendC::WaitFlag(EVENT_ID0); - } - -private: - CopyL1ToL0A copyL1ToL0A_; - CopyL1ToL0B copyL1ToL0B_; - - AscendC::LocalTensor l1Buf_{AscendC::LocalTensor(AscendC::TPosition::A1, 0, ArchTag::L1_SIZE)}; - AscendC::LocalTensor l0ABuf_{AscendC::LocalTensor(AscendC::TPosition::A2, 0, ArchTag::L0A_SIZE)}; - AscendC::LocalTensor l0BBuf_{AscendC::LocalTensor(AscendC::TPosition::B2, 0, ArchTag::L0B_SIZE)}; - AscendC::LocalTensor l0CBuf_{AscendC::LocalTensor(AscendC::TPosition::CO1, 0, ArchTag::L0C_SIZE)}; - AscendC::LocalTensor l1ATensorList_[STAGES]; - AscendC::LocalTensor l1BTensorList_[STAGES]; - AscendC::LocalTensor l1BIdxTensorList_[STAGES]; - AscendC::LocalTensor l0ATensorList_[STAGES]; - AscendC::LocalTensor l0BTensorList_[STAGES]; - AscendC::LocalTensor l0CTensor_; - - // Multi-stage event id list - int32_t l1AEventList_[STAGES]; - int32_t l1BEventList_[STAGES]; - int32_t l1BIdxEventList_[STAGES]; - int32_t l0AEventList_[STAGES]; - int32_t l0BEventList_[STAGES]; - - // The id of current stage - uint32_t l1ListId{0}; - uint32_t l0AListId{0}; - uint32_t l0BListId{0}; -}; -} // namespace Catlass::Gemm::Block - -#endif // CATLASS_GEMM_BLOCK_BLOCK_MMAD_SPARSE_TLA_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_streamk.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_streamk.hpp deleted file mode 100644 index 73f760d116..0000000000 --- a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_streamk.hpp +++ /dev/null @@ -1,406 +0,0 @@ -/** - * Copyright (c) 2025 Huawei Technologies Co., Ltd. - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of - * CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ - -#ifndef CATLASS_GEMM_BLOCK_BLOCK_MMAD_STREAMK_HPP -#define CATLASS_GEMM_BLOCK_BLOCK_MMAD_STREAMK_HPP - -#include "catlass/kda_catlass.hpp" -#include "catlass/arch/kda_resource.hpp" -#include "catlass/kda_coord.hpp" -#include "catlass/kda_gemm_coord.hpp" -#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" -#include "catlass/gemm/kda_gemm_helper.hpp" -#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" -#include "catlass/gemm/tile/kda_tile_mmad.hpp" - -namespace Catlass::Gemm::Block { - -template < - bool ENABLE_UNIT_FLAG_, - bool ENABLE_SHUFFLE_K_, - class L1TileShape_, - class L0TileShape_, - class AType_, - class BType_, - class CType_, - class BiasType_, - class TileCopy_, - class TileMmad_> -struct BlockMmad< - MmadAtlasA2Streamk, - L1TileShape_, - L0TileShape_, - AType_, - BType_, - CType_, - BiasType_, - TileCopy_, - TileMmad_> { -public: - // Type Aliases - using DispatchPolicy = MmadAtlasA2Streamk; - using ArchTag = typename DispatchPolicy::ArchTag; - using L1TileShape = L1TileShape_; - using L0TileShape = L0TileShape_; - using ElementA = typename AType_::Element; - using LayoutA = typename AType_::Layout; - using ElementB = typename BType_::Element; - using LayoutB = typename BType_::Layout; - using ElementC = typename CType_::Element; - using LayoutC = typename CType_::Layout; - using TileMmad = TileMmad_; - using CopyGmToL1A = typename TileCopy_::CopyGmToL1A; - using CopyGmToL1B = typename TileCopy_::CopyGmToL1B; - using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; - using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; - using CopyL0CToGmNormalBlock = typename TileCopy_::CopyL0CToGm; - using CopyL0CToGmStreamkBlock = Gemm::Tile::CopyL0CToGm>; - using ElementAccumulator = - typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; - using LayoutAInL1 = typename CopyL1ToL0A::LayoutSrc; - using LayoutBInL1 = typename CopyL1ToL0B::LayoutSrc; - using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst; - using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst; - using LayoutCInL0 = layout::zN; - - using L1AAlignHelper = Gemm::helper::L1AlignHelper; - using L1BAlignHelper = Gemm::helper::L1AlignHelper; - - static constexpr bool ENABLE_UNIT_FLAG = DispatchPolicy::ENABLE_UNIT_FLAG; - static constexpr bool ENABLE_SHUFFLE_K = DispatchPolicy::ENABLE_SHUFFLE_K; - static constexpr uint32_t STAGES = DispatchPolicy::STAGES; - static constexpr uint32_t L1A_SIZE = L1TileShape::M * L1TileShape::K * sizeof(ElementA); - static constexpr uint32_t L1B_SIZE = L1TileShape::N * L1TileShape::K * sizeof(ElementB); - static constexpr uint32_t L0A_SIZE = ArchTag::L0A_SIZE; - static constexpr uint32_t L0B_SIZE = ArchTag::L0B_SIZE; - static constexpr uint32_t L0C_SIZE = ArchTag::L0C_SIZE; - static constexpr uint32_t L0A_PINGPONG_BUF_SIZE = L0A_SIZE / STAGES; - static constexpr uint32_t L0B_PINGPONG_BUF_SIZE = L0B_SIZE / STAGES; - - // Check LayoutC - static_assert(std::is_same_v, "LayoutC only support RowMajor yet!"); - - // Check L1TileShape - static_assert((L1A_SIZE * STAGES + L1B_SIZE * STAGES) <= ArchTag::L1_SIZE, "L1TileShape exceeding the L1 space!"); - - // Check L0TileShape - static constexpr uint32_t L0A_TILE_SIZE = L0TileShape::M * L0TileShape::K * sizeof(ElementA); - static constexpr uint32_t L0B_TILE_SIZE = L0TileShape::K * L0TileShape::N * sizeof(ElementB); - static constexpr uint32_t L0C_TILE_SIZE = L0TileShape::M * L0TileShape::N * sizeof(ElementAccumulator); - static_assert((L0A_TILE_SIZE * STAGES) <= L0A_SIZE, "L0TileShape exceeding the L0A space!"); - static_assert((L0B_TILE_SIZE * STAGES) <= L0B_SIZE, "L0TileShape exceeding the L0B space!"); - static_assert(L0C_TILE_SIZE <= L0C_SIZE, "L0TileShape exceeding the L0C space!"); - - static_assert(L1TileShape::M == L0TileShape::M && L1TileShape::N == L0TileShape::N, - "The situation where the basic blocks of L1 and L0 differ on the m and n axes is not supported yet"); - static_assert(L0TileShape::K <= L1TileShape::K, "L0TileShape::K cannot exceed L1TileShape::K"); - // 32B (256b) aligned - static_assert(Gemm::helper::TileShapeAlignChecker::_ALIGN == 256, - "Tile shape must be 32B aligned."); - - /// Construct - CATLASS_DEVICE - BlockMmad(Arch::Resource &resource, uint32_t l1BufAddrStart = 0) - { - kPartLenMax = min(L0A_PINGPONG_BUF_SIZE / sizeof(ElementA) / L1TileShape::M / L1AAlignHelper::ELE_NUM_PER_C0 * - L1AAlignHelper::ELE_NUM_PER_C0, - L0B_PINGPONG_BUF_SIZE / sizeof(ElementB) / L1TileShape::N / L1BAlignHelper::ELE_NUM_PER_C0 * - L1BAlignHelper::ELE_NUM_PER_C0); - - if constexpr (std::is_same_v && std::is_same_v) { - kPartLenMax = RoundDown(kPartLenMax); - } - - uint32_t l1AOffset = l1BufAddrStart; - uint32_t l1BOffset = l1BufAddrStart + L1A_SIZE * STAGES; - // Init buffers - for (uint32_t i = 0; i < STAGES; i++) { - l1ATensorList[i] = resource.l1Buf.template GetBufferByByte(l1AOffset + L1A_SIZE * i); - l1BTensorList[i] = resource.l1Buf.template GetBufferByByte(l1BOffset + L1B_SIZE * i); - l0ATensorList[i] = resource.l0ABuf.template GetBufferByByte(L0A_PINGPONG_BUF_SIZE * i); - l0BTensorList[i] = resource.l0BBuf.template GetBufferByByte(L0B_PINGPONG_BUF_SIZE * i); - - l1AEventList[i] = i; - l1BEventList[i] = i + STAGES; - l0AEventList[i] = i; - l0BEventList[i] = i + STAGES; - AscendC::SetFlag(l1AEventList[i]); - AscendC::SetFlag(l1BEventList[i]); - AscendC::SetFlag(l0AEventList[i]); - AscendC::SetFlag(l0BEventList[i]); - } - l0CTensor = resource.l0CBuf.template GetBufferByByte(0); - if constexpr (!ENABLE_UNIT_FLAG) { - AscendC::SetFlag(EVENT_ID0); - } - } - - /// Destructor - CATLASS_DEVICE - ~BlockMmad() - { - for (uint32_t i = 0; i < STAGES; i++) { - AscendC::WaitFlag(l1AEventList[i]); - AscendC::WaitFlag(l1BEventList[i]); - AscendC::WaitFlag(l0AEventList[i]); - AscendC::WaitFlag(l0BEventList[i]); - } - if constexpr (!ENABLE_UNIT_FLAG) { - AscendC::WaitFlag(EVENT_ID0); - } - } - - /// Perform a block-scoped matrix multiply-accumulate - CATLASS_DEVICE - void operator()( - AscendC::GlobalTensor const &gmBlockA, LayoutA const &layoutA, - AscendC::GlobalTensor const &gmBlockB, LayoutB const &layoutB, - AscendC::GlobalTensor const &gmBlockC, LayoutC const &layoutC, - AscendC::GlobalTensor const &gmW, LayoutC const &layoutW, - AscendC::GlobalTensor const &gmNextBlockA, AscendC::GlobalTensor const &gmNextBlockB, - GemmCoord const &actualShape, GemmCoord const &actualShapeNext, - bool isFirstBlock, bool hasNextBlock, bool isSkBlock) - { - uint32_t mRound = RoundUp(actualShape.m()); - uint32_t nRound = RoundUp(actualShape.n()); - - auto layoutAInL1 = LayoutAInL1::template MakeLayout(L1TileShape::M, L1TileShape::K); - auto layoutBInL1 = LayoutBInL1::template MakeLayout(L1TileShape::K, L1TileShape::N); - auto layoutInL0C = LayoutCInL0::MakeLayoutInL0C(MakeCoord(mRound, nRound)); - - uint32_t kTileCount = CeilDiv(actualShape.k()); - - if constexpr (!ENABLE_UNIT_FLAG) { - AscendC::WaitFlag(EVENT_ID0); - } - uint32_t startTileIdx = 0; - if constexpr (ENABLE_SHUFFLE_K) { - startTileIdx = AscendC::GetBlockIdx(); - } - uint32_t firstTileIdx = startTileIdx % kTileCount; - uint32_t lastTileIdx = (startTileIdx + kTileCount - 1) % kTileCount; - uint32_t kActual = - (firstTileIdx < kTileCount - 1) ? L1TileShape::K : (actualShape.k() - firstTileIdx * L1TileShape::K); - - // k loop - for (uint32_t kLoopIdx = 0; kLoopIdx < kTileCount; kLoopIdx++) { - uint32_t shuffleKIdx = (startTileIdx + kLoopIdx) % kTileCount; - // Load first matrix A tile in total kernel loop from GM to L1 - if (shuffleKIdx == firstTileIdx && isFirstBlock) { - MatrixCoord gmTileACoord{0, shuffleKIdx * L1TileShape::K}; - MatrixCoord gmTileBCoord{shuffleKIdx * L1TileShape::K, 0}; - auto gmTileA = gmBlockA[layoutA.GetOffset(gmTileACoord)]; - auto gmTileB = gmBlockB[layoutB.GetOffset(gmTileBCoord)]; - // Load first matrix A tile from GM to L1 - AscendC::WaitFlag(l1AEventList[l1ListId]); - auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShape.m(), kActual)); - copyGmToL1A(l1ATensorList[l1ListId], gmTileA, layoutAInL1, layoutTileA); - AscendC::SetFlag(l1AEventList[l1ListId]); - - // Load first matrix B tile from GM to L1 - AscendC::WaitFlag(l1BEventList[l1ListId]); - auto layoutTileB = layoutB.GetTileLayout(MakeCoord(kActual, actualShape.n())); - copyGmToL1B(l1BTensorList[l1ListId], gmTileB, layoutBInL1, layoutTileB); - AscendC::SetFlag(l1BEventList[l1ListId]); - } - - uint32_t l1ListIdNext = (l1ListId + 1 < STAGES) ? (l1ListId + 1) : 0; - uint32_t kActualNext{0}; - - // preload next tile from GM to L1 - if (shuffleKIdx != lastTileIdx) { - uint32_t shuffleKIdxNext = (startTileIdx + kLoopIdx + 1) % kTileCount; - // Get L1 tensor for next stage - auto l1ATensor = l1ATensorList[l1ListIdNext]; - auto l1BTensor = l1BTensorList[l1ListIdNext]; - // Get GM tensor for next stage - kActualNext = (shuffleKIdxNext < kTileCount - 1) ? L1TileShape::K : (actualShape.k() - shuffleKIdxNext * L1TileShape::K); - MatrixCoord gmTileACoord{0, shuffleKIdxNext * L1TileShape::K}; - MatrixCoord gmTileBCoord{shuffleKIdxNext * L1TileShape::K, 0}; - auto gmTileA = gmBlockA[layoutA.GetOffset(gmTileACoord)]; - auto gmTileB = gmBlockB[layoutB.GetOffset(gmTileBCoord)]; - - // load next matrix A tile from GM to L1 - AscendC::WaitFlag(l1AEventList[l1ListIdNext]); - auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShape.m(), kActualNext)); - copyGmToL1A(l1ATensor, gmTileA, layoutAInL1, layoutTileA); - AscendC::SetFlag(l1AEventList[l1ListIdNext]); - - // load next matrix B tile from GM to L1 - AscendC::WaitFlag(l1BEventList[l1ListIdNext]); - auto layoutTileB = layoutB.GetTileLayout(MakeCoord(kActualNext, actualShape.n())); - copyGmToL1B(l1BTensor, gmTileB, layoutBInL1, layoutTileB); - AscendC::SetFlag(l1BEventList[l1ListIdNext]); - } - if (shuffleKIdx == lastTileIdx && hasNextBlock) { - uint32_t kTileCount = CeilDiv(actualShapeNext.k()); - uint32_t firstTileIdx = startTileIdx % kTileCount; - // Get L1 tensor for next stage - auto l1ATensor = l1ATensorList[l1ListIdNext]; - auto l1BTensor = l1BTensorList[l1ListIdNext]; - // Get GM tensor for next stage - kActualNext = (firstTileIdx < kTileCount - 1) ? L1TileShape::K : (actualShapeNext.k() - firstTileIdx * L1TileShape::K); - MatrixCoord gmTileACoord{0, firstTileIdx * L1TileShape::K}; - MatrixCoord gmTileBCoord{firstTileIdx * L1TileShape::K, 0}; - auto gmTileA = gmNextBlockA[layoutA.GetOffset(gmTileACoord)]; - auto gmTileB = gmNextBlockB[layoutB.GetOffset(gmTileBCoord)]; - // load next matrix A tile from GM to L1 - AscendC::WaitFlag(l1AEventList[l1ListIdNext]); - auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShapeNext.m(), kActualNext)); - copyGmToL1A(l1ATensor, gmTileA, layoutAInL1, layoutTileA); - AscendC::SetFlag(l1AEventList[l1ListIdNext]); - - // load next matrix B tile from GM to L1 - AscendC::WaitFlag(l1BEventList[l1ListIdNext]); - auto layoutTileB = layoutB.GetTileLayout(MakeCoord(kActualNext, actualShapeNext.n())); - copyGmToL1B(l1BTensor, gmTileB, layoutBInL1, layoutTileB); - AscendC::SetFlag(l1BEventList[l1ListIdNext]); - } - - // Get L1 tensor for current stage - auto l1ATensor = l1ATensorList[l1ListId]; - auto l1BTensor = l1BTensorList[l1ListId]; - - // Get the loop nums on L0 - uint32_t kPartLoop = CeilDiv(kActual, kPartLenMax); - - uint32_t l0ABufId = 0; - uint32_t l0BBufId = 0; - - for (int kPartIdx = 0; kPartIdx < kPartLoop; kPartIdx++) { - uint32_t kPartActual = - (kPartIdx < kPartLoop - 1) ? kPartLenMax : (kActual - kPartIdx * kPartLenMax); - - // Locate the current tile on L0A - auto l0ATile = l0ATensorList[l0ABufId]; - LayoutAInL0 layoutAInL0 = LayoutAInL0::template MakeLayout(mRound, kPartActual); - // Locate the current tile of matrix A on L1 - MatrixCoord l1ACoord{0, kPartIdx * kPartLenMax}; - auto l1ATile = l1ATensor[layoutAInL1.GetOffset(l1ACoord)]; - - AscendC::WaitFlag(l0AEventList[l0ABufId]); - if (kPartIdx == 0) { - AscendC::WaitFlag(l1AEventList[l1ListId]); - } - - copyL1ToL0A(l0ATile, l1ATile, layoutAInL0, layoutAInL1); - - if (kPartIdx == kPartLoop - 1) { - AscendC::SetFlag(l1AEventList[l1ListId]); - } - - // Locate the current tile on L0B - auto l0BTile = l0BTensorList[l0BBufId]; - LayoutBInL0 layoutBInL0 = LayoutBInL0::template MakeLayout(kPartActual, nRound); - // Locate the current tile of matrix B on L1 - MatrixCoord l1BCoord{kPartIdx * kPartLenMax, 0}; - auto l1BTile = l1BTensor[layoutBInL1.GetOffset(l1BCoord)]; - - // Wait for mmad finished - AscendC::WaitFlag(l0BEventList[l0BBufId]); - // If the current tile is the first one on the k&n axis, wait for loading matrix B from GM to L1 - if (kPartIdx == 0) { - AscendC::WaitFlag(l1BEventList[l1ListId]); - } - - // Load current tile from L1 to L0B - copyL1ToL0B(l0BTile, l1BTile, layoutBInL0, layoutBInL1); - - // If the current tile is the last one on the k&n axis, notify to load matrix B from GM to L1 - if (kPartIdx == kPartLoop - 1) { - AscendC::SetFlag(l1BEventList[l1ListId]); - } - // Notify to do mmad - AscendC::SetFlag(EVENT_ID0); - - // Locate the current tile on L0C - MatrixCoord l0CCoord{static_cast(0U), static_cast(0U)}; - auto l0CTile = l0CTensor[layoutInL0C.GetOffset(l0CCoord)]; - - // Compute the matrix multiplication on L0A and L0B and write the result to the accumulator - // Wait for loading L0B - AscendC::WaitFlag(EVENT_ID0); - - // If the current tile is the first tile on the k axis, the accumulator needs to be reset to 0 - bool initC = ((kLoopIdx == 0) && (kPartIdx == 0)); - // If the unit flag is enabled, the unit flag is set according to the calculation progress - uint8_t unitFlag = 0b00; - if constexpr (ENABLE_UNIT_FLAG) { - if ((kLoopIdx == kTileCount - 1) && (kPartIdx == kPartLoop - 1)) { - unitFlag = 0b11; - } else { - unitFlag = 0b10; - } - } - // Perform calculation operations - tileMmad(l0CTile, l0ATile, l0BTile, mRound, nRound, kPartActual, initC, unitFlag); - - // Notify to move the next L0B tile - AscendC::SetFlag(l0BEventList[l0BBufId]); - - l0BBufId = (l0BBufId + 1 < STAGES) ? (l0BBufId + 1) : 0; - - AscendC::SetFlag(l0AEventList[l0ABufId]); - l0ABufId = (l0ABufId + 1 < STAGES) ? (l0ABufId + 1) : 0; - } - - l1ListId = l1ListIdNext; - kActual = kActualNext; - } - - // copy block out - LayoutC layoutBlock = layoutC.GetTileLayout(actualShape.GetCoordMN()); - - if constexpr (!ENABLE_UNIT_FLAG) { - AscendC::SetFlag(EVENT_ID0); - AscendC::WaitFlag(EVENT_ID0); - if (isSkBlock) { - copyL0CToGmStreamkBlock(gmW, l0CTensor, layoutW, layoutInL0C); - } else { - copyL0CToGmNormalBlock(gmBlockC, l0CTensor, layoutBlock, layoutInL0C); - } - AscendC::SetFlag(EVENT_ID0); - } else { - if (isSkBlock) { - copyL0CToGmStreamkBlock(gmW, l0CTensor, layoutW, layoutInL0C, 0b11); - } else { - copyL0CToGmNormalBlock(gmBlockC, l0CTensor, layoutBlock, layoutInL0C, 0b11); - } - } - } - -protected: - /// Data members - AscendC::LocalTensor l1ATensorList[STAGES]; - AscendC::LocalTensor l1BTensorList[STAGES]; - AscendC::LocalTensor l0ATensorList[STAGES]; - AscendC::LocalTensor l0BTensorList[STAGES]; - AscendC::LocalTensor l0CTensor; - - int32_t l1AEventList[STAGES]; - int32_t l1BEventList[STAGES]; - int32_t l0AEventList[STAGES]; - int32_t l0BEventList[STAGES]; - - uint32_t l1ListId{0}; - uint32_t kPartLenMax; - - TileMmad tileMmad; - CopyGmToL1A copyGmToL1A; - CopyGmToL1B copyGmToL1B; - CopyL1ToL0A copyL1ToL0A; - CopyL1ToL0B copyL1ToL0B; - CopyL0CToGmNormalBlock copyL0CToGmNormalBlock; - CopyL0CToGmStreamkBlock copyL0CToGmStreamkBlock; -}; - -} // namespace Catlass::Gemm::Block - -#endif // CATLASS_GEMM_BLOCK_BLOCK_MMAD_STREAMK_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_w4a4_per_token_per_channel_dequant.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_w4a4_per_token_per_channel_dequant.hpp deleted file mode 100644 index fc30e82382..0000000000 --- a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_w4a4_per_token_per_channel_dequant.hpp +++ /dev/null @@ -1,494 +0,0 @@ -/** - * Copyright (c) 2025 Huawei Technologies Co., Ltd. - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of - * CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ -#ifndef CATLASS_GEMM_BLOCK_BLOCK_MMAD_W4A4_HPP -#define CATLASS_GEMM_BLOCK_BLOCK_MMAD_W4A4_HPP - -#include "catlass/arch/kda_resource.hpp" -#include "catlass/kda_catlass.hpp" -#include "catlass/kda_coord.hpp" -#include "catlass/detail/kda_callback.hpp" -#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" -#include "catlass/gemm/kda_gemm_helper.hpp" -#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" -#include "catlass/gemm/tile/kda_tile_mmad.hpp" -#include "catlass/kda_gemm_coord.hpp" - -namespace Catlass::Gemm::Block { - -template < - uint32_t PRELOAD_STAGES_, - uint32_t L1_STAGES_, - uint32_t L0A_STAGES_, - uint32_t L0B_STAGES_, - uint32_t L0C_STAGES_, - bool ENABLE_UNIT_FLAG_, - bool ENABLE_SHUFFLE_K_, - class L1TileShape_, - class L0TileShape_, - class AType_, - class BType_, - class CType_, - class BiasType_, - class TileCopy_, - class TileMmad_> -struct BlockMmad< - MmadAtlasA2W4A4MatmulPerTokenPerChannelDequant< - PRELOAD_STAGES_, - L1_STAGES_, - L0A_STAGES_, - L0B_STAGES_, - L0C_STAGES_, - ENABLE_UNIT_FLAG_, - ENABLE_SHUFFLE_K_>, - L1TileShape_, - L0TileShape_, - AType_, - BType_, - CType_, - BiasType_, - TileCopy_, - TileMmad_> { -public: - // Type Aliases - using DispatchPolicy = MmadAtlasA2W4A4MatmulPerTokenPerChannelDequant< - PRELOAD_STAGES_, - L1_STAGES_, - L0A_STAGES_, - L0B_STAGES_, - L0C_STAGES_, - ENABLE_UNIT_FLAG_, - ENABLE_SHUFFLE_K_>; - using ArchTag = typename DispatchPolicy::ArchTag; - using L1TileShape = L1TileShape_; - using L0TileShape = L0TileShape_; - using ElementA = typename AType_::Element; - using LayoutA = typename AType_::Layout; - using ElementB = typename BType_::Element; - using LayoutB = typename BType_::Layout; - using ElementC = typename CType_::Element; - using LayoutC = typename CType_::Layout; - using TileMmad = TileMmad_; - using CopyGmToL1A = typename TileCopy_::CopyGmToL1A; - using CopyGmToL1B = typename TileCopy_::CopyGmToL1B; - using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; - using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; - using CopyL0CToGm = typename TileCopy_::CopyL0CToGm; - using CopyGmToL1Scale = typename TileCopy_::CopyGmToL1Scale; - using CopyL1ToFP = typename TileCopy_::CopyL1ToFP; - using ElementAccumulator = typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; - using LayoutAInL1 = typename CopyL1ToL0A::LayoutSrc; - using LayoutBInL1 = typename CopyL1ToL0B::LayoutSrc; - using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst; - using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst; - using LayoutCInL0 = layout::zN; - using LayoutScale = layout::VectorLayout; - using LayoutL1Scale = layout::VectorLayout; - using LayoutFpBuf = layout::VectorLayout; - - using L1AAlignHelper = Gemm::helper::L1AlignHelper; - using L1BAlignHelper = Gemm::helper::L1AlignHelper; - - static constexpr uint32_t PRELOAD_STAGES = DispatchPolicy::PRELOAD_STAGES; - static constexpr uint32_t L1_STAGES = DispatchPolicy::L1_STAGES; - static constexpr uint32_t L0A_STAGES = DispatchPolicy::L0A_STAGES; - static constexpr uint32_t L0B_STAGES = DispatchPolicy::L0B_STAGES; - static constexpr uint32_t L0C_STAGES = DispatchPolicy::L0C_STAGES; - - static constexpr bool ENABLE_UNIT_FLAG = DispatchPolicy::ENABLE_UNIT_FLAG; - static constexpr bool ENABLE_SHUFFLE_K = DispatchPolicy::ENABLE_SHUFFLE_K; - - // L1 tile size - static constexpr uint32_t L1A_TILE_SIZE = L1TileShape::M * L1TileShape::K * sizeof(int8_t) / 2; - static constexpr uint32_t L1B_TILE_SIZE = L1TileShape::N * L1TileShape::K * sizeof(int8_t) / 2; - static constexpr uint32_t SCALE_KTILE_SIZE = L1TileShape::N * sizeof(uint64_t); - // L0 tile size - static constexpr uint32_t L0A_TILE_SIZE = L0TileShape::M * L0TileShape::K * sizeof(int8_t) / 2; - static constexpr uint32_t L0B_TILE_SIZE = L0TileShape::K * L0TileShape::N * sizeof(int8_t) / 2; - static constexpr uint32_t L0C_TILE_SIZE = L1TileShape::M * L1TileShape::N * sizeof(ElementAccumulator); - - // Check LayoutC - static_assert(std::is_same_v, "LayoutC only support RowMajor yet!"); - - // Check L1TileShape - static_assert( - (L1A_TILE_SIZE + L1B_TILE_SIZE + SCALE_KTILE_SIZE) * L1_STAGES <= ArchTag::L1_SIZE, - "L1TileShape exceeding the L1 space!"); - - // Check L0TileShape - static_assert(L0A_TILE_SIZE * L0A_STAGES <= ArchTag::L0A_SIZE, "L0TileShape exceeding the L0A space!"); - static_assert(L0B_TILE_SIZE * L0B_STAGES <= ArchTag::L0B_SIZE, "L0TileShape exceeding the L0B space!"); - static_assert(L0C_TILE_SIZE * L0C_STAGES <= ArchTag::L0C_SIZE, "L0TileShape exceeding the L0C space!"); - - static_assert( - L1TileShape::M == L0TileShape::M && L1TileShape::N == L0TileShape::N, - "The situation where the basic blocks of L1 and L0 differ on the m and n axes is not supported yet"); - static_assert(L0TileShape::K <= L1TileShape::K, "L0TileShape::K cannot exceed L1TileShape::K"); - - static constexpr auto L1A_LAYOUT = LayoutAInL1::template MakeLayout(L1TileShape::M, L1TileShape::K); - static constexpr auto L1B_LAYOUT = LayoutBInL1::template MakeLayout(L1TileShape::K, L1TileShape::N); - - CATLASS_DEVICE - BlockMmad(Arch::Resource &resource, uint32_t l1BufAddrStart = 0, uint32_t FpAddrStart = 0) - { - InitL1(resource, l1BufAddrStart); - InitFpBuf(resource, FpAddrStart); - InitL0A(resource); - InitL0B(resource); - InitL0C(resource); - } - - CATLASS_DEVICE - ~BlockMmad() - { - for (uint32_t i = 0; i < L1_STAGES; ++i) { - AscendC::WaitFlag(l1AEventList[i]); - AscendC::WaitFlag(l1BEventList[i]); - AscendC::WaitFlag(l1ScaleEventList[i]); - } - for (uint32_t i = 0; i < L0A_STAGES; ++i) { - AscendC::WaitFlag(l0AEventList[i]); - } - for (uint32_t i = 0; i < L0B_STAGES; ++i) { - AscendC::WaitFlag(l0BEventList[i]); - } - for (uint32_t i = 0; i < L0C_STAGES; ++i) { - AscendC::WaitFlag(l0CEventList[i]); - } - } - - CATLASS_DEVICE - void operator()( - AscendC::GlobalTensor const &gmScale, - LayoutScale const &layoutScale, - AscendC::GlobalTensor const &gmBlockA, - LayoutA const &layoutA, - AscendC::GlobalTensor const &gmBlockB, - LayoutB const &layoutB, - AscendC::GlobalTensor const &gmBlockC, - LayoutC const &layoutC, - GemmCoord const &actualShape, - Callback const &callbackBeforeFixpipe, - Callback const &callbackAfterFixpipe) - { - uint32_t kTileCount = CeilDiv(actualShape.k()); - - uint32_t mRound = RoundUp(actualShape.m()); - uint32_t nRound = RoundUp(actualShape.n()); - - uint32_t startTileIdx = 0; - if constexpr (ENABLE_SHUFFLE_K) { - startTileIdx = AscendC::GetBlockIdx() % kTileCount; - } - - // Copy scale from GM to L1 - auto layoutTileScale = layoutScale.GetTileLayout(MakeCoord(L1TileShape::N)); - LayoutL1Scale layoutL1KTileScale{L1TileShape::N}; - AscendC::WaitFlag(l1ScaleEventList[l1ScaleListId]); - copyGmToL1Scale(l1ScaleBufList[l1ScaleListId], gmScale, layoutL1KTileScale, layoutTileScale); - AscendC::SetFlag(l1ScaleEventList[l1ScaleListId]); - - for (uint32_t kLoopIdx = 0; kLoopIdx < kTileCount; ++kLoopIdx) { - uint32_t kTileIdx = (startTileIdx + kLoopIdx < kTileCount) ? (startTileIdx + kLoopIdx) : (startTileIdx + kLoopIdx - kTileCount); - - uint32_t kActual = (kTileIdx < kTileCount - 1) ? L1TileShape::K : (actualShape.k() - kTileIdx * L1TileShape::K); - - // Emission load instruction from GM to L1 - MatrixCoord gmTileAOffset{0, kTileIdx * L1TileShape::K}; - MatrixCoord gmTileBOffset{kTileIdx * L1TileShape::K, 0}; - auto gmTileA = gmBlockA[layoutA.GetOffset(gmTileAOffset)]; - auto gmTileB = gmBlockB[layoutB.GetOffset(gmTileBOffset)]; - // Load first matrix A tile from GM to L1 - AscendC::WaitFlag(l1AEventList[l1ListId]); - auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShape.m(), kActual)); - copyGmToL1A(l1ATensorList[l1ListId], gmTileA, L1A_LAYOUT, layoutTileA); - AscendC::SetFlag(l1AEventList[l1ListId]); - // Load first matrix B tile from GM to L1 - AscendC::WaitFlag(l1BEventList[l1ListId]); - auto layoutTileB = layoutB.GetTileLayout(MakeCoord(kActual, actualShape.n())); - copyGmToL1B(l1BTensorList[l1ListId], gmTileB, L1B_LAYOUT, layoutTileB); - AscendC::SetFlag(l1BEventList[l1ListId]); - - // If the number of preload instructions reaches the upper limit, perform an mmad calculation on L1 tile - if (preloadCount == PRELOAD_STAGES) { - L1TileMmad(l1TileMmadParamsList[l1TileMmadParamsId]); - } - - // Store the current load status - uint32_t preloadL1TileMmadParamsId = (l1TileMmadParamsId + preloadCount < PRELOAD_STAGES) ? (l1TileMmadParamsId + preloadCount) : (l1TileMmadParamsId + preloadCount - PRELOAD_STAGES); - auto &l1TileMmadParams = l1TileMmadParamsList[preloadL1TileMmadParamsId]; - l1TileMmadParams.l1ListId = l1ListId; - l1TileMmadParams.mRound = mRound; - l1TileMmadParams.nRound = nRound; - l1TileMmadParams.kActual = kActual; - l1TileMmadParams.isKLoopFirst = (kLoopIdx == 0); - l1TileMmadParams.isKLoopLast = (kLoopIdx == kTileCount - 1); - if (kLoopIdx == kTileCount - 1) { - l1TileMmadParams.l1ScaleListId = l1ScaleListId; - l1TileMmadParams.gmBlockC = gmBlockC; - l1TileMmadParams.layoutCInGm = layoutC.GetTileLayout(actualShape.GetCoordMN()); - l1TileMmadParams.callbackBeforeFixpipe = callbackBeforeFixpipe; - l1TileMmadParams.callbackAfterFixpipe = callbackAfterFixpipe; - } - - if (preloadCount < PRELOAD_STAGES) { - ++preloadCount; - } else { - l1TileMmadParamsId = (l1TileMmadParamsId + 1 < PRELOAD_STAGES) ? (l1TileMmadParamsId + 1) : 0; - } - l1ListId = (l1ListId + 1 < L1_STAGES) ? (l1ListId + 1) : 0; - } - l1ScaleListId = (l1ScaleListId + 1 < L1_STAGES) ? (l1ScaleListId + 1) : 0; - } - - CATLASS_DEVICE - void SynchronizeBlock() - { - while (preloadCount > 0) { - L1TileMmad(l1TileMmadParamsList[l1TileMmadParamsId]); - l1TileMmadParamsId = (l1TileMmadParamsId + 1 < PRELOAD_STAGES) ? (l1TileMmadParamsId + 1) : 0; - --preloadCount; - } - } - -private: - struct L1TileMmadParams { - uint32_t l1ListId; - uint32_t l1ScaleListId; - uint32_t mRound; - uint32_t nRound; - uint32_t kActual; - bool isKLoopFirst; - bool isKLoopLast; - AscendC::GlobalTensor gmBlockC; - LayoutC layoutCInGm; - Callback callbackBeforeFixpipe; - Callback callbackAfterFixpipe; - - CATLASS_DEVICE - L1TileMmadParams() = default; - }; - - CATLASS_DEVICE - void InitL1(Arch::Resource &resource, uint32_t l1BufAddrStart) - { - uint32_t l1AOffset = l1BufAddrStart; - uint32_t l1BOffset = l1BufAddrStart + L1A_TILE_SIZE * L1_STAGES; - uint32_t l1ScaleOffset = l1BOffset + L1B_TILE_SIZE * L1_STAGES; - for (uint32_t i = 0; i < L1_STAGES; ++i) { - l1ATensorList[i] = resource.l1Buf.template GetBufferByByte(l1AOffset + L1A_TILE_SIZE * i); - l1BTensorList[i] = resource.l1Buf.template GetBufferByByte(l1BOffset + L1B_TILE_SIZE * i); - l1ScaleBufList[i] = resource.l1Buf.template GetBufferByByte(l1ScaleOffset + SCALE_KTILE_SIZE * i); - l1AEventList[i] = i; - l1BEventList[i] = i + L1_STAGES; - l1ScaleEventList[i] = i; - AscendC::SetFlag(l1AEventList[i]); - AscendC::SetFlag(l1BEventList[i]); - AscendC::SetFlag(l1ScaleEventList[i]); - } - } - - CATLASS_DEVICE - void InitFpBuf(Arch::Resource &resource, uint32_t FpAddrStart) - { - uint32_t FpOffset = FpAddrStart; - fixpipeBuf = resource.fpBuf.template GetBufferByByte(FpOffset); - } - - CATLASS_DEVICE - void InitL0A(Arch::Resource &resource) - { - for (uint32_t i = 0; i < L0A_STAGES; ++i) { - l0ATensorList[i] = resource.l0ABuf.template GetBufferByByte(L0A_TILE_SIZE * i); - l0AEventList[i] = i; - AscendC::SetFlag(l0AEventList[i]); - } - } - - CATLASS_DEVICE - void InitL0B(Arch::Resource &resource) - { - for (uint32_t i = 0; i < L0B_STAGES; ++i) { - l0BTensorList[i] = resource.l0BBuf.template GetBufferByByte(L0B_TILE_SIZE * i); - l0BEventList[i] = i + L0A_STAGES; - AscendC::SetFlag(l0BEventList[i]); - } - } - - CATLASS_DEVICE - void InitL0C(Arch::Resource &resource) - { - for (uint32_t i = 0; i < L0C_STAGES; ++i) { - l0CTensorList[i] = resource.l0CBuf.template GetBufferByByte(L0C_TILE_SIZE * i); - l0CEventList[i] = i; - AscendC::SetFlag(l0CEventList[i]); - } - } - - CATLASS_DEVICE - void L1TileMmad(L1TileMmadParams const ¶ms) - { - uint32_t mPartLoop = CeilDiv(params.mRound); - uint32_t nPartLoop = CeilDiv(params.nRound); - uint32_t kPartLoop = CeilDiv(params.kActual); - auto &l1ATensor = l1ATensorList[params.l1ListId]; - auto &l1BTensor = l1BTensorList[params.l1ListId]; - - LayoutFpBuf layoutFpBuf{L1TileShape::N}; - LayoutL1Scale layoutL1Scale{L1TileShape::N}; - - auto &l0CTensor = l0CTensorList[l0CListId]; - LayoutCInL0 layoutCInL0 = LayoutCInL0::MakeLayoutInL0C(MakeCoord(params.mRound, params.nRound)); - - if constexpr (!ENABLE_UNIT_FLAG) { - if (params.isKLoopFirst) { - AscendC::WaitFlag(l0CEventList[l0CListId]); - } - } - - for (uint32_t mPartIdx = 0; mPartIdx < mPartLoop; ++mPartIdx) { - uint32_t mPartActual = (mPartIdx < mPartLoop - 1) ? L0TileShape::M : (params.mRound - mPartIdx * L0TileShape::M); - - for (uint32_t kPartIdx = 0; kPartIdx < kPartLoop; ++kPartIdx) { - uint32_t kPartActual = (kPartIdx < kPartLoop - 1) ? L0TileShape::K : (params.kActual - kPartIdx * L0TileShape::K); - - auto &l0ATile = l0ATensorList[l0AListId]; - auto layoutAInL0 = LayoutAInL0::template MakeLayout(mPartActual, kPartActual); - auto l1AOffset = MakeCoord(mPartIdx, kPartIdx) * L0TileShape::ToCoordMK(); - auto l1ATile = l1ATensor[L1A_LAYOUT.GetOffset(l1AOffset)]; - - AscendC::WaitFlag(l0AEventList[l0AListId]); - if ((mPartIdx == 0) && (kPartIdx == 0)) { - AscendC::WaitFlag(l1AEventList[params.l1ListId]); - } - copyL1ToL0A(l0ATile, l1ATile, layoutAInL0, L1A_LAYOUT); - if ((mPartIdx == mPartLoop - 1) && (kPartIdx == kPartLoop - 1)) { - AscendC::SetFlag(l1AEventList[params.l1ListId]); - } - - for (uint32_t nPartIdx = 0; nPartIdx < nPartLoop; ++nPartIdx) { - uint32_t nPartActual = (nPartIdx < nPartLoop - 1) ? L0TileShape::N : (params.nRound - nPartIdx * L0TileShape::N); - - auto &l0BTile = l0BTensorList[l0BListId]; - auto layoutBInL0 = LayoutBInL0::template MakeLayout(kPartActual, nPartActual); - auto l1BOffset = MakeCoord(kPartIdx, nPartIdx) * L0TileShape::ToCoordKN(); - auto l1BTile = l1BTensor[L1B_LAYOUT.GetOffset(l1BOffset)]; - - AscendC::WaitFlag(l0BEventList[l0BListId]); - if ((kPartIdx == 0) && (nPartIdx == 0)) { - AscendC::WaitFlag(l1BEventList[params.l1ListId]); - } - copyL1ToL0B(l0BTile, l1BTile, layoutBInL0, L1B_LAYOUT); - if ((kPartIdx == kPartLoop - 1) && (nPartIdx == nPartLoop - 1)) { - AscendC::SetFlag(l1BEventList[params.l1ListId]); - } - - AscendC::SetFlag(EVENT_ID0); - - auto l0COffset = MakeCoord(mPartIdx, nPartIdx) * L0TileShape::ToCoordMN(); - auto l0CTile = l0CTensor[layoutCInL0.GetOffset(l0COffset)]; - - AscendC::WaitFlag(EVENT_ID0); - // If the current tile is the first tile on the k axis, the accumulator needs to be reset to 0 - bool initC = (params.isKLoopFirst && (kPartIdx == 0)); - // If the unit flag is enabled, the unit flag is set according to the calculation progress - uint8_t unitFlag = 0b00; - if constexpr (ENABLE_UNIT_FLAG) { - if (params.isKLoopLast && (mPartIdx == mPartLoop - 1) && (kPartIdx == kPartLoop - 1) && (nPartIdx == nPartLoop - 1)) { - unitFlag = 0b11; - } else { - unitFlag = 0b10; - } - } - tileMmad(l0CTile, l0ATile, l0BTile, mPartActual, nPartActual, kPartActual, initC, unitFlag); - - AscendC::SetFlag(l0BEventList[l0BListId]); - l0BListId = (l0BListId + 1 < L0B_STAGES) ? (l0BListId + 1) : 0; - } - AscendC::SetFlag(l0AEventList[l0AListId]); - l0AListId = (l0AListId + 1 < L0A_STAGES) ? (l0AListId + 1) : 0; - } - } - - if (params.isKLoopLast) { - auto layoutCInGm = params.layoutCInGm; - params.callbackBeforeFixpipe(); - - if constexpr (!ENABLE_UNIT_FLAG) { - AscendC::SetFlag(l0CEventList[l0CListId]); - AscendC::WaitFlag(l0CEventList[l0CListId]); - - // Copy scale from L1 to FixpipeBuffer - auto &l1ScaleBuf = l1ScaleBufList[params.l1ScaleListId]; - AscendC::WaitFlag(params.l1ScaleListId); - copyL1ToFP(fixpipeBuf, l1ScaleBuf, layoutFpBuf, layoutL1Scale); - AscendC::SetFlag(params.l1ScaleListId); - - AscendC::PipeBarrier(); - - copyL0CToGm(params.gmBlockC, l0CTensor, fixpipeBuf, layoutCInGm, layoutCInL0); - AscendC::SetFlag(l0CEventList[l0CListId]); - } else { - // Copy scale from L1 to FixpipeBuffer - auto &l1ScaleBuf = l1ScaleBufList[params.l1ScaleListId]; - AscendC::WaitFlag(params.l1ScaleListId); - copyL1ToFP(fixpipeBuf, l1ScaleBuf, layoutFpBuf, layoutL1Scale); - AscendC::SetFlag(params.l1ScaleListId); - - AscendC::PipeBarrier(); - - copyL0CToGm(params.gmBlockC, l0CTensor, fixpipeBuf, layoutCInGm, layoutCInL0, 0b11); - } - - l0CListId = (l0CListId + 1 < L0C_STAGES) ? (l0CListId + 1) : 0; - - params.callbackAfterFixpipe(); - } - } - - AscendC::LocalTensor fixpipeBuf; - - AscendC::LocalTensor l1ATensorList[L1_STAGES]; - AscendC::LocalTensor l1BTensorList[L1_STAGES]; - AscendC::LocalTensor l1ScaleBufList[L1_STAGES]; - int32_t l1AEventList[L1_STAGES]; - int32_t l1BEventList[L1_STAGES]; - int32_t l1ScaleEventList[L1_STAGES]; - uint32_t l1ListId{0}; - uint32_t l1ScaleListId{0}; - - AscendC::LocalTensor l0ATensorList[L0A_STAGES]; - int32_t l0AEventList[L0A_STAGES]; - uint32_t l0AListId{0}; - - AscendC::LocalTensor l0BTensorList[L0B_STAGES]; - int32_t l0BEventList[L0B_STAGES]; - uint32_t l0BListId{0}; - - AscendC::LocalTensor l0CTensorList[L0C_STAGES_]; - int32_t l0CEventList[L0C_STAGES_]; - uint32_t l0CListId{0}; - - L1TileMmadParams l1TileMmadParamsList[PRELOAD_STAGES]; - uint32_t l1TileMmadParamsId{0}; - uint32_t preloadCount{0}; - - TileMmad tileMmad; - CopyGmToL1A copyGmToL1A; - CopyGmToL1B copyGmToL1B; - CopyL1ToL0A copyL1ToL0A; - CopyL1ToL0B copyL1ToL0B; - CopyL0CToGm copyL0CToGm; - CopyGmToL1Scale copyGmToL1Scale; - CopyL1ToFP copyL1ToFP; -}; - -} // namespace Catlass::Gemm::Block - -#endif // CATLASS_GEMM_BLOCK_BLOCK_MMAD_W4A4_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad.hpp index d832b08517..290d58344b 100644 --- a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad.hpp +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad.hpp @@ -17,65 +17,6 @@ namespace Catlass::Gemm::Block { -#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 2201) -template < - class DispatchPolicy, - class L1TileShape, - class L0TileShape, - class AType, - class BType, - class CType, - class BiasType = void, - class TileCopy = Gemm::Tile::TileCopy, - class TileMmad = Gemm::Tile::TileMmad > -struct BlockMmad { - static_assert(DEPENDENT_FALSE, "BlockMmad is not implemented for this DispatchPolicy"); -}; - -/// new add for the reason that i am using the dispatchpolicy which is same as the policy of the optimized_matmul -// so i add a new one class to avoid the conflict -template < - class DispatchPolicy, - class L1TileShape, - class L0TileShape, - class AType, - class BType, - class CType, - class BiasType = void, - class TileCopy = Gemm::Tile::TileCopyGemm, // change the name - class TileMmad = Gemm::Tile::TileMmad > -struct BlockGemm { - static_assert(DEPENDENT_FALSE, "BlockMmad is not implemented for this DispatchPolicy"); -}; - -template < - class DispatchPolicy, - class AType, - class BType, - class CType, - class BiasType, - class TileCopy, - class TileMmad> -struct BlockMmadAiv { - static_assert(DEPENDENT_FALSE, "BlockMmadAiv is not implemented for this DispatchPolicy"); -}; - -template < - class DispatchPolicy, - class L1TileShape, - class L0TileShape, - class ElementA, - class ElementB, - class ElementC, - class ElementBias = void, - class TileCopy = Gemm::Tile::SparseTileCopyTla > -struct BlockMmadSparseTla { - static_assert(DEPENDENT_FALSE, "BlockMmadSparseTla is not implemented for this DispatchPolicy"); -}; - -#endif - template < class DispatchPolicy, class L1TileShape, @@ -94,50 +35,9 @@ struct BlockMmadTla { } // namespace Catlass::Gemm::Block -#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 2201) -#include "catlass/gemm/block/kda_block_mmad_pingpong.hpp" -#include "catlass/gemm/block/kda_block_mmad_fa_qk.hpp" -#include "catlass/gemm/block/kda_block_mmad_fa_pv.hpp" -#include "catlass/gemm/block/kda_block_mmad_mla_qk.hpp" -#include "catlass/gemm/block/kda_block_mmad_mla_pv.hpp" -#include "catlass/gemm/block/kda_block_mmad_mla_qk_tp1_spec.hpp" -#include "catlass/gemm/block/kda_block_mmad_mla_pv_tp1_spec.hpp" -#include "catlass/gemm/block/kda_block_mmad_preload.hpp" -#include "catlass/gemm/block/kda_block_mmad_preload_async.hpp" -#include "catlass/gemm/block/kda_block_mmad_preload_async_with_callback.hpp" -#include "catlass/gemm/block/kda_block_mmad_gemm.hpp" -#include "catlass/gemm/block/kda_block_mmad_pingpong_bias.hpp" -#include "catlass/gemm/block/block_mmad_fai_qk_head_tail.hpp" -#include "catlass/gemm/block/block_mmad_fai_qk_normal.hpp" -#include "catlass/gemm/block/block_mmad_fai_pv_head_tail.hpp" -#include "catlass/gemm/block/block_mmad_fai_pv_normal.hpp" -#include "catlass/gemm/block/block_mmad_pingpong_full_loadA.hpp" -#include "catlass/gemm/block/block_mmad_pingpong_full_loadA_tla.hpp" -#include "catlass/gemm/block/block_mmad_pingpong_with_prologue.hpp" -#include "catlass/gemm/block/block_mmad_pingpong_slice_k_with_prologue.hpp" -#include "catlass/gemm/block/block_mmad_dynamic_common.hpp" -#include "catlass/gemm/block/block_mmad_dynamic_small.hpp" -#include "catlass/gemm/block/block_mmad_dynamic_streamk.hpp" -#include "catlass/gemm/block/block_mmad_dynamic_single_core_splitk.hpp" -#include "catlass/gemm/block/block_mmad_dynamic_preload_async_with_callback.hpp" -#include "catlass/gemm/block/block_mmad_small.hpp" -#include "catlass/gemm/block/block_mmad_single_core_splitk.hpp" -#include "catlass/gemm/block/block_mmad_dynamic_aiv.hpp" -#include "catlass/gemm/block/block_mmad_streamk.hpp" -#include "catlass/gemm/block/block_mmad_w4a4_per_token_per_channel_dequant.hpp" -#include "catlass/gemm/block/block_mmad_sparse_tla.hpp" -#include "catlass/gemm/block/block_mmad_fai_qk_head_tail_tla.hpp" -#include "catlass/gemm/block/block_mmad_fai_qk_normal_tla.hpp" -#include "catlass/gemm/block/block_mmad_fai_pv_head_tail_tla.hpp" -#include "catlass/gemm/block/block_mmad_fai_pv_normal_tla.hpp" -#endif - +// Only the pingpong-TLA implementation is instantiated by chunk_kda_fwd +// (Gemm::MmadPingpong). The MmadPingpongTlaMulti / MmadPingpongTlaPreloadAL1B +// specializations live under kernel_utils/block/ and are included by their users. #include "catlass/gemm/block/block_mmad_pingpong_tla.hpp" -#include "catlass/gemm/block/block_mmad_pingpong_dequant_tla.hpp" -#include "catlass/gemm/block/block_mmad_pingpong_tla_v2.hpp" -#include "catlass/gemm/block/block_mmad_preload_tla.hpp" -#include "catlass/gemm/block/block_mmad_preload_async_with_callback_tla.hpp" -#include "catlass/gemm/block/block_mmad_fai_pv_tla.hpp" -#include "catlass/gemm/block/block_mmad_fai_qk_tla.hpp" -#include "catlass/gemm/block/block_mmad_pingpong_per_group_per_block_tla.hpp" + #endif // CATLASS_KDA_BLOCK_MMAD_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_fa_pv.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_fa_pv.hpp deleted file mode 100644 index dfbe1ede13..0000000000 --- a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_fa_pv.hpp +++ /dev/null @@ -1,193 +0,0 @@ -/** - * Copyright (c) 2025 Huawei Technologies Co., Ltd. - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of - * CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ - -#ifndef CATLASS_KDA_BLOCK_MMAD_FA_PV_HPP -#define CATLASS_KDA_BLOCK_MMAD_FA_PV_HPP - -#include "catlass/kda_catlass.hpp" -#include "catlass/arch/kda_cross_core_sync.hpp" -#include "catlass/arch/kda_resource.hpp" -#include "catlass/kda_coord.hpp" -#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" -#include "catlass/gemm/kda_gemm_helper.hpp" -#include "catlass/kda_gemm_coord.hpp" -#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" -#include "catlass/gemm/tile/kda_tile_mmad.hpp" - -//////////////////////////////////////////////////////////////////// - -namespace Catlass::Gemm::Block { -//////////////////////////////////////////////////////////////////// - -template < - class L1TileShape_, - class L0TileShape_, - class AType_, - class BType_, - class CType_, - class BiasType_, - class TileCopy_, - class TileMmad_> -struct BlockMmad< - MmadAtlasA2FAPV, - L1TileShape_, - L0TileShape_, - AType_, - BType_, - CType_, - BiasType_, - TileCopy_, - TileMmad_> { -public: - // Type Aliases - using DispatchPolicy = MmadAtlasA2FAPV; - using ArchTag = typename DispatchPolicy::ArchTag; - using L1TileShape = L1TileShape_; - using L0TileShape = L0TileShape_; - using ElementA = typename AType_::Element; - using LayoutA = typename AType_::Layout; - using ElementB = typename BType_::Element; - using LayoutB = typename BType_::Layout; - using ElementC = typename CType_::Element; - using LayoutC = typename CType_::Layout; - using TileMmad = TileMmad_; - using CopyGmToL1A = typename TileCopy_::CopyGmToL1A; - using CopyGmToL1B = typename TileCopy_::CopyGmToL1B; - using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; - using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; - using CopyL0CToGm = typename TileCopy_::CopyL0CToGm; - using ElementAccumulator = - typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; - using LayoutAInL1 = typename CopyL1ToL0A::LayoutSrc; - using LayoutBInL1 = typename CopyL1ToL0B::LayoutSrc; - using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst; - using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst; - using LayoutCInL0 = layout::zN; - - using L1AAlignHelper = Gemm::helper::L1AlignHelper; - using L1BAlignHelper = Gemm::helper::L1AlignHelper; - - static constexpr uint32_t STAGES = DispatchPolicy::STAGES; - static constexpr uint32_t L1A_SIZE = L1TileShape::M * L1TileShape::K * sizeof(ElementA); - static constexpr uint32_t L1B_SIZE = L1TileShape::N * L1TileShape::K * sizeof(ElementB); - static constexpr uint32_t L0A_SIZE = ArchTag::L0A_SIZE; - static constexpr uint32_t L0B_SIZE = ArchTag::L0B_SIZE; - static constexpr uint32_t L0C_SIZE = ArchTag::L0C_SIZE; - static constexpr uint32_t L0A_PINGPONG_BUF_SIZE = L0A_SIZE / STAGES; - static constexpr uint32_t L0B_PINGPONG_BUF_SIZE = L0B_SIZE / STAGES; - static constexpr uint32_t L0C_PINGPONG_BUF_SIZE = L0C_SIZE / STAGES; - - // Check LayoutC - static_assert(std::is_same_v, "LayoutC only support RowMajor yet!"); - - /// Construct - CATLASS_DEVICE - BlockMmad(Arch::Resource &resource, uint32_t l1BufAddrStart = 0) - { - for (uint32_t i = 0; i < STAGES; i++) { - l1ATensor[i] = resource.l1Buf.template GetBufferByByte(l1BufAddrStart + L1A_SIZE * i); - l1BTensor[i] = - resource.l1Buf.template GetBufferByByte(l1BufAddrStart + L1A_SIZE * STAGES + L1B_SIZE * i); - l0ATensor[i] = resource.l0ABuf.template GetBufferByByte(L0A_PINGPONG_BUF_SIZE * i); - l0BTensor[i] = resource.l0BBuf.template GetBufferByByte(L0B_PINGPONG_BUF_SIZE * i); - l0CTensor[i] = resource.l0CBuf.template GetBufferByByte(L0C_PINGPONG_BUF_SIZE * i); - } - - AscendC::SetFlag(EVENT_ID2); - AscendC::SetFlag(EVENT_ID3); - } - - /// Destructor - CATLASS_DEVICE - ~BlockMmad() - { - AscendC::WaitFlag(EVENT_ID2); - AscendC::WaitFlag(EVENT_ID3); - } - - /// Perform a block-scoped matrix multiply-accumulate - CATLASS_DEVICE - void operator()( - AscendC::GlobalTensor gA, - AscendC::GlobalTensor gB, - AscendC::GlobalTensor gC, - LayoutA layoutA, LayoutB layoutB, LayoutC layoutC, - GemmCoord actualShape, uint32_t &pingpongFlag, - Arch::CrossCoreFlag softmaxReady) - { - LayoutAInL1 layoutAInL1 = LayoutAInL1::template MakeLayout(L1TileShape::M, L1TileShape::K); - LayoutBInL1 layoutBInL1 = LayoutBInL1::template MakeLayout(L1TileShape::K, L1TileShape::N); - - uint32_t mRound = RoundUp(actualShape.m()); - uint32_t nRound = RoundUp(actualShape.n()); - uint32_t kRound = RoundUp(actualShape.k()); - - constexpr uint32_t PINGPONG_FLAG_OFFSET = 2; - AscendC::WaitFlag(pingpongFlag + PINGPONG_FLAG_OFFSET); - auto layoutTileB = layoutB.GetTileLayout(MakeCoord(actualShape.k(), actualShape.n())); - copyGmToL1B(l1BTensor[pingpongFlag], gB, layoutBInL1, layoutTileB); - AscendC::SetFlag(EVENT_ID0); - AscendC::WaitFlag(EVENT_ID0); - - AscendC::WaitFlag(pingpongFlag); - LayoutBInL0 layoutBInL0 = LayoutBInL0::template MakeLayout(kRound, nRound); - copyL1ToL0B(l0BTensor[pingpongFlag], l1BTensor[pingpongFlag], layoutBInL0, layoutBInL1); - - Arch::CrossCoreWaitFlag(softmaxReady); - auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShape.m(), actualShape.k())); - copyGmToL1A(l1ATensor[pingpongFlag], gA, layoutAInL1, layoutTileA); - AscendC::SetFlag(EVENT_ID0); - AscendC::WaitFlag(EVENT_ID0); - LayoutAInL0 layoutAInL0 = LayoutAInL0::template MakeLayout(mRound, kRound); - copyL1ToL0A(l0ATensor[pingpongFlag], l1ATensor[pingpongFlag], layoutAInL0, layoutAInL1); - AscendC::SetFlag(pingpongFlag + PINGPONG_FLAG_OFFSET); - - AscendC::SetFlag(EVENT_ID0); - AscendC::WaitFlag(EVENT_ID0); - - AscendC::WaitFlag(pingpongFlag); - tileMmad( - l0CTensor[pingpongFlag], l0ATensor[pingpongFlag], l0BTensor[pingpongFlag], mRound, nRound, actualShape.k()); - AscendC::SetFlag(pingpongFlag); - - AscendC::SetFlag(EVENT_ID0); - AscendC::WaitFlag(EVENT_ID0); - - // copy block out - auto blockShape = MakeCoord(actualShape.m(), actualShape.n()); - auto layoutInL0C = LayoutCInL0::MakeLayoutInL0C(blockShape); - - copyL0CToGm(gC, l0CTensor[pingpongFlag], layoutC, layoutInL0C); - AscendC::SetFlag(pingpongFlag); - - pingpongFlag = 1 - pingpongFlag; - } - -protected: - /// Data members - AscendC::LocalTensor l1ATensor[STAGES]; - AscendC::LocalTensor l1BTensor[STAGES]; - AscendC::LocalTensor l0ATensor[STAGES]; - AscendC::LocalTensor l0BTensor[STAGES]; - AscendC::LocalTensor l0CTensor[STAGES]; - - TileMmad tileMmad; - CopyGmToL1A copyGmToL1A; - CopyGmToL1B copyGmToL1B; - CopyL1ToL0A copyL1ToL0A; - CopyL1ToL0B copyL1ToL0B; - CopyL0CToGm copyL0CToGm; -}; - -//////////////////////////////////////////////////////////////////// - -} // namespace Catlass::Gemm::Block - -#endif // CATLASS_KDA_BLOCK_MMAD_FA_PV_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_fa_qk.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_fa_qk.hpp deleted file mode 100644 index e58ceee9ff..0000000000 --- a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_fa_qk.hpp +++ /dev/null @@ -1,199 +0,0 @@ -/** - * Copyright (c) 2025 Huawei Technologies Co., Ltd. - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of - * CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ - -#ifndef CATLASS_KDA_BLOCK_MMAD_FA_QK_HPP -#define CATLASS_KDA_BLOCK_MMAD_FA_QK_HPP - -#include "catlass/kda_catlass.hpp" -#include "catlass/arch/kda_resource.hpp" -#include "catlass/kda_coord.hpp" -#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" -#include "catlass/gemm/kda_gemm_helper.hpp" -#include "catlass/kda_gemm_coord.hpp" -#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" -#include "catlass/gemm/tile/kda_tile_mmad.hpp" - -//////////////////////////////////////////////////////////////////// - -namespace Catlass::Gemm::Block { -//////////////////////////////////////////////////////////////////// - -template < - class L1TileShape_, - class L0TileShape_, - class AType_, - class BType_, - class CType_, - class BiasType_, - class TileCopy_, - class TileMmad_> -struct BlockMmad< - MmadAtlasA2FAQK, - L1TileShape_, - L0TileShape_, - AType_, - BType_, - CType_, - BiasType_, - TileCopy_, - TileMmad_> { -public: - // Type Aliases - using DispatchPolicy = MmadAtlasA2FAQK; - using ArchTag = typename DispatchPolicy::ArchTag; - using L1TileShape = L1TileShape_; - using L0TileShape = L0TileShape_; - using ElementA = typename AType_::Element; - using LayoutA = typename AType_::Layout; - using ElementB = typename BType_::Element; - using LayoutB = typename BType_::Layout; - using ElementC = typename CType_::Element; - using LayoutC = typename CType_::Layout; - using TileMmad = TileMmad_; - using CopyGmToL1A = typename TileCopy_::CopyGmToL1A; - using CopyGmToL1B = typename TileCopy_::CopyGmToL1B; - using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; - using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; - using CopyL0CToGm = typename TileCopy_::CopyL0CToGm; - using ElementAccumulator = - typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; - using LayoutAInL1 = typename CopyL1ToL0A::LayoutSrc; - using LayoutBInL1 = typename CopyL1ToL0B::LayoutSrc; - using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst; - using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst; - using LayoutCInL0 = layout::zN; - - using L1AAlignHelper = Gemm::helper::L1AlignHelper; - using L1BAlignHelper = Gemm::helper::L1AlignHelper; - - static constexpr uint32_t STAGES = DispatchPolicy::STAGES; - static constexpr uint32_t L1A_SIZE = L1TileShape::M * L1TileShape::K * sizeof(ElementA); - static constexpr uint32_t L1B_SIZE = L1TileShape::N * L1TileShape::K * sizeof(ElementB); - static constexpr uint32_t L0A_SIZE = ArchTag::L0A_SIZE; - static constexpr uint32_t L0B_SIZE = ArchTag::L0B_SIZE; - static constexpr uint32_t L0C_SIZE = ArchTag::L0C_SIZE; - static constexpr uint32_t L0A_PINGPONG_BUF_SIZE = L0A_SIZE / STAGES; - static constexpr uint32_t L0B_PINGPONG_BUF_SIZE = L0B_SIZE / STAGES; - static constexpr uint32_t L0C_PINGPONG_BUF_SIZE = L0C_SIZE / STAGES; - - // Check LayoutC - static_assert(std::is_same_v, "LayoutC only support RowMajor yet!"); - - /// Construct - CATLASS_DEVICE - BlockMmad(Arch::Resource &resource, uint32_t l1BufAddrStart = 0) - { - for (uint32_t i = 0; i < STAGES; i++) { - l1ATensor[i] = resource.l1Buf.template GetBufferByByte(l1BufAddrStart + L1A_SIZE * i); - l1BTensor[i] = - resource.l1Buf.template GetBufferByByte(l1BufAddrStart + L1A_SIZE * STAGES + L1B_SIZE * i); - l0ATensor[i] = resource.l0ABuf.template GetBufferByByte(L0A_PINGPONG_BUF_SIZE * i); - l0BTensor[i] = resource.l0BBuf.template GetBufferByByte(L0B_PINGPONG_BUF_SIZE * i); - l0CTensor[i] = resource.l0CBuf.template GetBufferByByte(L0C_PINGPONG_BUF_SIZE * i); - } - - AscendC::SetFlag(EVENT_ID0); - AscendC::SetFlag(EVENT_ID1); - AscendC::SetFlag(EVENT_ID0); - AscendC::SetFlag(EVENT_ID1); - AscendC::SetFlag(EVENT_ID0); - AscendC::SetFlag(EVENT_ID1); - } - - /// Destructor - CATLASS_DEVICE - ~BlockMmad() - { - AscendC::WaitFlag(EVENT_ID0); - AscendC::WaitFlag(EVENT_ID1); - AscendC::WaitFlag(EVENT_ID0); - AscendC::WaitFlag(EVENT_ID1); - AscendC::WaitFlag(EVENT_ID0); - AscendC::WaitFlag(EVENT_ID1); - } - - /// Perform a block-scoped matrix multiply-accumulate - CATLASS_DEVICE - void operator()( - AscendC::GlobalTensor gA, - AscendC::GlobalTensor gB, - AscendC::GlobalTensor gC, - LayoutA layoutA, LayoutB layoutB, LayoutC layoutC, - GemmCoord actualShape, uint32_t &pingpongFlag, bool isFirst) - { - LayoutAInL1 layoutAInL1 = LayoutAInL1::template MakeLayout(L1TileShape::M, L1TileShape::K); - LayoutBInL1 layoutBInL1 = LayoutBInL1::template MakeLayout(L1TileShape::K, L1TileShape::N); - - uint32_t mRound = RoundUp(actualShape.m()); - uint32_t nRound = RoundUp(actualShape.n()); - uint32_t kRound = RoundUp(actualShape.k()); - - if (isFirst) { - auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShape.m(), actualShape.k())); - copyGmToL1A(l1ATensor[0], gA, layoutAInL1, layoutTileA); - AscendC::SetFlag(EVENT_ID0); - AscendC::WaitFlag(EVENT_ID0); - } - - AscendC::WaitFlag(pingpongFlag); - LayoutAInL0 layoutAInL0 = LayoutAInL0::template MakeLayout(mRound, kRound); - copyL1ToL0A(l0ATensor[pingpongFlag], l1ATensor[0], layoutAInL0, layoutAInL1); - - AscendC::WaitFlag(pingpongFlag); - auto layoutTileB = layoutB.GetTileLayout(MakeCoord(actualShape.k(), actualShape.n())); - copyGmToL1B(l1BTensor[pingpongFlag], gB, layoutBInL1, layoutTileB); - AscendC::SetFlag(EVENT_ID0); - AscendC::WaitFlag(EVENT_ID0); - LayoutBInL0 layoutBInL0 = LayoutBInL0::template MakeLayout(kRound, nRound); - copyL1ToL0B(l0BTensor[pingpongFlag], l1BTensor[pingpongFlag], layoutBInL0, layoutBInL1); - AscendC::SetFlag(pingpongFlag); - - AscendC::SetFlag(EVENT_ID0); - AscendC::WaitFlag(EVENT_ID0); - - AscendC::WaitFlag(pingpongFlag); - tileMmad( - l0CTensor[pingpongFlag], l0ATensor[pingpongFlag], l0BTensor[pingpongFlag], mRound, nRound, actualShape.k()); - AscendC::SetFlag(pingpongFlag); - - AscendC::SetFlag(EVENT_ID0); - AscendC::WaitFlag(EVENT_ID0); - - // copy block out - auto blockShape = MakeCoord(actualShape.m(), actualShape.n()); - auto layoutInL0C = LayoutCInL0::MakeLayoutInL0C(blockShape); - - copyL0CToGm(gC, l0CTensor[pingpongFlag], layoutC, layoutInL0C); - AscendC::SetFlag(pingpongFlag); - - pingpongFlag = 1 - pingpongFlag; - } - -protected: - /// Data members - AscendC::LocalTensor l1ATensor[STAGES]; - AscendC::LocalTensor l1BTensor[STAGES]; - AscendC::LocalTensor l0ATensor[STAGES]; - AscendC::LocalTensor l0BTensor[STAGES]; - AscendC::LocalTensor l0CTensor[STAGES]; - - TileMmad tileMmad; - CopyGmToL1A copyGmToL1A; - CopyGmToL1B copyGmToL1B; - CopyL1ToL0A copyL1ToL0A; - CopyL1ToL0B copyL1ToL0B; - CopyL0CToGm copyL0CToGm; -}; - -//////////////////////////////////////////////////////////////////// - -} // namespace Catlass::Gemm::Block - -#endif // CATLASS_KDA_BLOCK_MMAD_FA_QK_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_gemm.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_gemm.hpp deleted file mode 100644 index 8f97aa18ea..0000000000 --- a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_gemm.hpp +++ /dev/null @@ -1,348 +0,0 @@ -/** - * Copyright (c) 2025 Huawei Technologies Co., Ltd. - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of - * CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ - -#ifndef CATLASS_KDA_BLOCK_MMAD_GEMM_HPP -#define CATLASS_KDA_BLOCK_MMAD_GEMM_HPP - -#include "catlass/kda_catlass.hpp" -#include "catlass/gemm/kda_gemm_helper.hpp" -#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" -#include "catlass/gemm/tile/kda_tile_mmad.hpp" -#include "catlass/kda_gemm_coord.hpp" -#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" -#include "catlass/arch/kda_resource.hpp" - -namespace Catlass::Gemm::Block { - -template < - bool ENABLE_UNIT_FLAG_, - bool ENABLE_SHUFFLE_K_, - bool ENABLE_ABBA_, - class L1TileShape_, - class L0TileShape_, - class AType_, - class BType_, - class CType_, - class BiasType_, - class TileCopy_, - class TileMmad_> -struct BlockGemm< - Gemm::GemmAtlasA2, - L1TileShape_, - L0TileShape_, - AType_, - BType_, - CType_, - BiasType_, - TileCopy_, - TileMmad_> { -public: - using DispatchPolicy = Gemm::GemmAtlasA2; - using ArchTag = typename DispatchPolicy::ArchTag; - using L1TileShape = L1TileShape_; - using L0TileShape = L0TileShape_; - using ElementA = typename AType_::Element; - using LayoutA = typename AType_::Layout; - using ElementB = typename BType_::Element; - using LayoutB = typename BType_::Layout; - using ElementC = typename CType_::Element; - using LayoutC = typename CType_::Layout; - - using TileMmad = TileMmad_; - using CopyGmToL1A = typename TileCopy_::CopyGmToL1A; - using CopyGmToL1B = typename TileCopy_::CopyGmToL1B; - using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; - using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; - using CopyL0CToGm = typename TileCopy_::CopyL0CToGm; - using ElementAccumulator = - typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; - using LayoutAInL1 = typename CopyGmToL1A::LayoutDst; - using LayoutBInL1 = typename CopyGmToL1B::LayoutDst; - using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst; - using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst; - using LayoutCInL0 = layout::zN; - - static constexpr uint32_t STAGES = DispatchPolicy::STAGES; - static constexpr bool ENABLE_UNIT_FLAG = DispatchPolicy::ENABLE_UNIT_FLAG; - static constexpr bool ENABLE_SHUFFLE_K = DispatchPolicy::ENABLE_SHUFFLE_K; - static constexpr bool ENABLE_ABBA = DispatchPolicy::ENABLE_ABBA; - const uint32_t L1Size = ArchTag::L1_SIZE; - const uint32_t L1ASize = L1TileShape::M * L1TileShape::K * sizeof(ElementA); - const uint32_t L1BSize = L1TileShape::K * L1TileShape::N * sizeof(ElementB); - const uint32_t cSize = L1TileShape::M * L1TileShape::N * sizeof(ElementAccumulator); - const uint32_t BlockCnt = L1TileShape::M * L1TileShape::N; - const uint32_t L0ASize = ArchTag::L0A_SIZE; - const uint32_t L0BSize = ArchTag::L0B_SIZE; - const uint32_t L0CSize = ArchTag::L0C_SIZE; - const uint32_t L0A_PINGPONG_BUF_LEN = (L0ASize / STAGES); - const uint32_t L0B_PINGPONG_BUF_LEN = (L0BSize / STAGES); - const uint32_t l0CBlockNum = ArchTag::L0C_SIZE / cSize; - - static_assert(std::is_same_v, "LayoutC only support RowMajor yet!"); - - // 32B (256b) aligned - static_assert(Gemm::helper::TileShapeAlignChecker::_ALIGN == 256, - "Tile shape must be 32B aligned."); - - CATLASS_DEVICE - BlockGemm(Arch::Resource &resource, uint32_t l1BufAddrStart = 0) - { - uint32_t l1AOffset = l1BufAddrStart; - uint32_t l1BOffset = l1BufAddrStart + L1ASize * STAGES; - for (uint32_t i = 0; i < STAGES; i++) { - l1ATensor[i] = resource.l1Buf.template GetBufferByByte(l1AOffset + L1ASize * i); - l1BTensor[i] = resource.l1Buf.template GetBufferByByte(l1BOffset + L1BSize * i); - l0ATensor[i] = resource.l0ABuf.template GetBufferByByte(L0A_PINGPONG_BUF_LEN * i); - l0BTensor[i] = resource.l0BBuf.template GetBufferByByte(L0B_PINGPONG_BUF_LEN * i); - // Assign event ID for each stages - l1AEventList[i] = i; - l1BEventList[i] = i + STAGES; - l0AEventList[i] = i; - l0BEventList[i] = i + STAGES; - // The event id that needs to be set before the loop - AscendC::SetFlag(l1AEventList[i]); - AscendC::SetFlag(l1BEventList[i]); - AscendC::SetFlag(l0AEventList[i]); - AscendC::SetFlag(l0BEventList[i]); - } - l0CTensor = resource.l0CBuf.template GetBufferByByte(0); - } - // destroy function - CATLASS_DEVICE - ~BlockGemm() - { - for (uint32_t i = 0; i < STAGES; i++) { - AscendC::WaitFlag(l1AEventList[i]); - AscendC::WaitFlag(l1BEventList[i]); - AscendC::WaitFlag(l0AEventList[i]); - AscendC::WaitFlag(l0BEventList[i]); - } - } - - CATLASS_DEVICE - void operator()( - AscendC::GlobalTensor const &gmA, LayoutA const &layoutA, - AscendC::GlobalTensor const &gmB, LayoutB const &layoutB, - AscendC::GlobalTensor const &gmC, LayoutC const &layoutC, - AscendC::GlobalTensor const &gmNextBlockA, - AscendC::GlobalTensor const &gmNextBlockB, - GemmCoord const &actualShape, GemmCoord const &actualShapeNext, - bool isFirstBlock, bool hasNextBlock, uint32_t singleIdx) - { - uint32_t K = actualShape.k(); - uint32_t maxKPerBlock = L1TileShape::K; - uint32_t kLoops = CeilDiv(K, maxKPerBlock); - uint32_t kLoopsNext = CeilDiv(actualShapeNext.k(), maxKPerBlock); - uint32_t startTileIdx{0}; - if (ENABLE_SHUFFLE_K) { - startTileIdx = AscendC::GetBlockIdx(); - } - uint32_t firstTileIdx = startTileIdx % kLoops; - uint32_t firstTileIdxNext = startTileIdx % kLoopsNext; - uint32_t lastTileIdx = (startTileIdx + kLoops - 1) % kLoops; - uint32_t kGmActual = (firstTileIdx == kLoops - 1) ? (K - firstTileIdx * maxKPerBlock) : maxKPerBlock; - auto layoutAInL1 = LayoutAInL1::template MakeLayout(L1TileShape::M, L1TileShape::K); - auto layoutBInL1 = LayoutBInL1::template MakeLayout(L1TileShape::K, L1TileShape::N); - for (uint32_t kIdx = 0; kIdx < kLoops; kIdx++) { - uint32_t shuffleKIdx = (startTileIdx + kIdx) % kLoops; - if (shuffleKIdx == firstTileIdx && isFirstBlock) { - auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShape.m(), kGmActual)); - auto layoutTileB = layoutB.GetTileLayout(MakeCoord(kGmActual, actualShape.n())); - MatrixCoord gmTileAOffset{0, shuffleKIdx * maxKPerBlock}; - auto gmTileA = gmA[layoutA.GetOffset(gmTileAOffset)]; - MatrixCoord gmTileBOffset{shuffleKIdx * maxKPerBlock, 0}; - auto gmTileB = gmB[layoutB.GetOffset(gmTileBOffset)]; - AscendC::WaitFlag(l1AEventList[l1ListId]); - copyGmToL1A(l1ATensor[l1ListId], gmTileA, layoutAInL1, layoutTileA); - AscendC::SetFlag(l1AEventList[l1ListId]); - AscendC::WaitFlag(l1BEventList[l1ListId]); - copyGmToL1B(l1BTensor[l1ListId], gmTileB, layoutBInL1, layoutTileB); - AscendC::SetFlag(l1BEventList[l1ListId]); - } - l1ListIdNext = 1 - l1ListId; - uint32_t kGmActualNext = 0; - if (shuffleKIdx != lastTileIdx) { - uint32_t shuffleKIdxNext = (startTileIdx + kIdx + 1) % kLoops; - kGmActualNext = (shuffleKIdxNext == kLoops - 1) ? (K - shuffleKIdxNext * maxKPerBlock) : maxKPerBlock; - auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShape.m(), kGmActualNext)); - auto layoutTileB = layoutB.GetTileLayout(MakeCoord(kGmActualNext, actualShape.n())); - MatrixCoord gmTileAOffset{0, shuffleKIdxNext * maxKPerBlock}; - auto gmTileA = gmA[layoutA.GetOffset(gmTileAOffset)]; - MatrixCoord gmTileBOffset{shuffleKIdxNext * maxKPerBlock, 0}; - auto gmTileB = gmB[layoutB.GetOffset(gmTileBOffset)]; - if (ENABLE_ABBA) { - if (shuffleKIdxNext % 2 == 1) { - AscendC::WaitFlag(l1BEventList[l1ListIdNext]); - copyGmToL1B(l1BTensor[l1ListIdNext], gmTileB, layoutBInL1, layoutTileB); - AscendC::SetFlag(l1BEventList[l1ListIdNext]); - AscendC::WaitFlag(l1AEventList[l1ListIdNext]); - copyGmToL1A(l1ATensor[l1ListIdNext], gmTileA, layoutAInL1, layoutTileA); - AscendC::SetFlag(l1AEventList[l1ListIdNext]); - } else { - AscendC::WaitFlag(l1AEventList[l1ListIdNext]); - copyGmToL1A(l1ATensor[l1ListIdNext], gmTileA, layoutAInL1, layoutTileA); - AscendC::SetFlag(l1AEventList[l1ListIdNext]); - AscendC::WaitFlag(l1BEventList[l1ListIdNext]); - copyGmToL1B(l1BTensor[l1ListIdNext], gmTileB, layoutBInL1, layoutTileB); - AscendC::SetFlag(l1BEventList[l1ListIdNext]); - } - } else { - AscendC::WaitFlag(l1AEventList[l1ListIdNext]); - copyGmToL1A(l1ATensor[l1ListIdNext], gmTileA, layoutAInL1, layoutTileA); - AscendC::SetFlag(l1AEventList[l1ListIdNext]); - AscendC::WaitFlag(l1BEventList[l1ListIdNext]); - copyGmToL1B(l1BTensor[l1ListIdNext], gmTileB, layoutBInL1, layoutTileB); - AscendC::SetFlag(l1BEventList[l1ListIdNext]); - } - } - if (shuffleKIdx == lastTileIdx && hasNextBlock) { - kGmActualNext = (firstTileIdxNext == kLoopsNext - 1) ? (actualShapeNext.k() - firstTileIdxNext * maxKPerBlock) : maxKPerBlock; - auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShapeNext.m(), kGmActualNext)); - auto layoutTileB = layoutB.GetTileLayout(MakeCoord(kGmActualNext, actualShapeNext.n())); - MatrixCoord gmTileAOffset{0, firstTileIdxNext * maxKPerBlock}; - auto gmNextTileA = gmNextBlockA[layoutA.GetOffset(gmTileAOffset)]; - MatrixCoord gmTileBOffset{firstTileIdxNext * maxKPerBlock, 0}; - auto gmNextTileB = gmNextBlockB[layoutB.GetOffset(gmTileBOffset)]; - if (ENABLE_ABBA) { - if (shuffleKIdx % 2 == 0) { - AscendC::WaitFlag(l1BEventList[l1ListIdNext]); - copyGmToL1B(l1BTensor[l1ListIdNext], gmNextTileB, layoutBInL1, layoutTileB); - AscendC::SetFlag(l1BEventList[l1ListIdNext]); - AscendC::WaitFlag(l1AEventList[l1ListIdNext]); - copyGmToL1A(l1ATensor[l1ListIdNext], gmNextTileA, layoutAInL1, layoutTileA); - AscendC::SetFlag(l1AEventList[l1ListIdNext]); - } else { - AscendC::WaitFlag(l1AEventList[l1ListIdNext]); - copyGmToL1A(l1ATensor[l1ListIdNext], gmNextTileA, layoutAInL1, layoutTileA); - AscendC::SetFlag(l1AEventList[l1ListIdNext]); - AscendC::WaitFlag(l1BEventList[l1ListIdNext]); - copyGmToL1B(l1BTensor[l1ListIdNext], gmNextTileB, layoutBInL1, layoutTileB); - AscendC::SetFlag(l1BEventList[l1ListIdNext]); - } - } else { - AscendC::WaitFlag(l1AEventList[l1ListIdNext]); - copyGmToL1A(l1ATensor[l1ListIdNext], gmNextTileA, layoutAInL1, layoutTileA); - AscendC::SetFlag(l1AEventList[l1ListIdNext]); - AscendC::WaitFlag(l1BEventList[l1ListIdNext]); - copyGmToL1B(l1BTensor[l1ListIdNext], gmNextTileB, layoutBInL1, layoutTileB); - AscendC::SetFlag(l1BEventList[l1ListIdNext]); - } - } - - uint32_t kL0TileSize = L0TileShape::K; - uint32_t kL0Loops = CeilDiv(kGmActual, kL0TileSize); - AscendC::WaitFlag(l1AEventList[l1ListId]); - AscendC::WaitFlag(l1BEventList[l1ListId]); - auto l1ATile = l1ATensor[l1ListId]; - auto l1BTile = l1BTensor[l1ListId]; - uint32_t mActual{0}; - uint32_t nActual{0}; - for (uint32_t kL0Idx = 0; kL0Idx < kL0Loops; kL0Idx++) { - uint32_t kL0Actual = (kL0Idx == kL0Loops - 1) ? (kGmActual - kL0Idx * kL0TileSize) : kL0TileSize; - LayoutAInL0 layoutAInL0 = LayoutAInL0::template MakeLayout(L1TileShape::M, kL0Actual); - LayoutBInL0 layoutBInL0 = LayoutBInL0::template MakeLayout(kL0Actual, L1TileShape::N); - uint32_t l1TileAOffset = layoutAInL1.GetOffset(MatrixCoord(0, kL0Idx * kL0TileSize)); - uint32_t l1TileBOffset = layoutBInL1.GetOffset(MatrixCoord(kL0Idx * kL0TileSize, 0)); - auto l1TileA = l1ATile[l1TileAOffset]; - auto l1TileB = l1BTile[l1TileBOffset]; - auto l0TileA = l0ATensor[l0ListId]; - auto l0TileB = l0BTensor[l0ListId]; - mActual = L1TileShape::M; - nActual = L1TileShape::N; - if (ENABLE_ABBA) { - if (shuffleKIdx % 2 == 0) { - if (kL0Idx % 2 == 0) { - AscendC::WaitFlag(l0BEventList[l0ListId]); - copyL1ToL0B(l0TileB, l1TileB, layoutBInL0, layoutBInL1); - AscendC::SetFlag(l0BEventList[l0ListId]); - AscendC::WaitFlag(l0AEventList[l0ListId]); - copyL1ToL0A(l0TileA, l1TileA, layoutAInL0, layoutAInL1); - AscendC::SetFlag(l0AEventList[l0ListId]); - } else { - AscendC::WaitFlag(l0AEventList[l0ListId]); - copyL1ToL0A(l0TileA, l1TileA, layoutAInL0, layoutAInL1); - AscendC::SetFlag(l0AEventList[l0ListId]); - AscendC::WaitFlag(l0BEventList[l0ListId]); - copyL1ToL0B(l0TileB, l1TileB, layoutBInL0, layoutBInL1); - AscendC::SetFlag(l0BEventList[l0ListId]); - } - } else { - if (kL0Idx % 2 == 0) { - AscendC::WaitFlag(l0AEventList[l0ListId]); - copyL1ToL0A(l0TileA, l1TileA, layoutAInL0, layoutAInL1); - AscendC::SetFlag(l0AEventList[l0ListId]); - AscendC::WaitFlag(l0BEventList[l0ListId]); - copyL1ToL0B(l0TileB, l1TileB, layoutBInL0, layoutBInL1); - AscendC::SetFlag(l0BEventList[l0ListId]); - } else { - AscendC::WaitFlag(l0BEventList[l0ListId]); - copyL1ToL0B(l0TileB, l1TileB, layoutBInL0, layoutBInL1); - AscendC::SetFlag(l0BEventList[l0ListId]); - AscendC::WaitFlag(l0AEventList[l0ListId]); - copyL1ToL0A(l0TileA, l1TileA, layoutAInL0, layoutAInL1); - AscendC::SetFlag(l0AEventList[l0ListId]); - } - } - } else { - AscendC::WaitFlag(l0AEventList[l0ListId]); - copyL1ToL0A(l0TileA, l1TileA, layoutAInL0, layoutAInL1); - AscendC::SetFlag(l0AEventList[l0ListId]); - AscendC::WaitFlag(l0BEventList[l0ListId]); - copyL1ToL0B(l0TileB, l1TileB, layoutBInL0, layoutBInL1); - AscendC::SetFlag(l0BEventList[l0ListId]); - } - if (kL0Idx == kL0Loops - 1) { - AscendC::SetFlag(l1AEventList[l1ListId]); - AscendC::SetFlag(l1BEventList[l1ListId]); - l1ListId = l1ListIdNext; - kGmActual = kGmActualNext; - } - AscendC::WaitFlag(l0BEventList[l0ListId]); - AscendC::WaitFlag(l0AEventList[l0ListId]); - tileMmad(l0CTensor[(singleIdx % l0CBlockNum) * BlockCnt], l0TileA, l0TileB, mActual, nActual, kL0Actual, (kIdx == 0) && (kL0Idx == 0)); - AscendC::SetFlag(l0AEventList[l0ListId]); - AscendC::SetFlag(l0BEventList[l0ListId]); - l0ListId = 1 - l0ListId; - } - } - AscendC::SetFlag((int32_t)(singleIdx % l0CBlockNum)); - AscendC::WaitFlag((int32_t)(singleIdx % l0CBlockNum)); - auto layoutInL0X = LayoutCInL0::MakeLayoutInL0C(MakeCoord(L1TileShape::M, L1TileShape::N)); - LayoutC layoutBlock = layoutC.GetTileLayout(MakeCoord(actualShape.m(), actualShape.n())); - copyL0CToGm(gmC, l0CTensor[(singleIdx % l0CBlockNum) * BlockCnt], layoutBlock, layoutInL0X); - } - -private: - AscendC::LocalTensor l1ATensor[STAGES]; - AscendC::LocalTensor l1BTensor[STAGES]; - AscendC::LocalTensor l0ATensor[STAGES]; - AscendC::LocalTensor l0BTensor[STAGES]; - AscendC::LocalTensor l0CTensor; - // Multi-stage event id list - int32_t l1AEventList[STAGES]; - int32_t l1BEventList[STAGES]; - int32_t l0AEventList[STAGES]; - int32_t l0BEventList[STAGES]; - - // The id of current stage - uint32_t l1ListId{0}; - uint32_t l0ListId{0}; - uint32_t l1ListIdNext{0}; - - TileMmad tileMmad; - CopyGmToL1A copyGmToL1A; - CopyGmToL1B copyGmToL1B; - CopyL1ToL0A copyL1ToL0A; - CopyL1ToL0B copyL1ToL0B; - CopyL0CToGm copyL0CToGm; -}; -} // namespace Catlass::Gemm::Block - -#endif // CATLASS_KDA_BLOCK_MMAD_GEMM_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_mla_pv.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_mla_pv.hpp deleted file mode 100644 index 3d37211bc2..0000000000 --- a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_mla_pv.hpp +++ /dev/null @@ -1,207 +0,0 @@ -/** - * Copyright (c) 2025 Huawei Technologies Co., Ltd. - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of - * CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ - -#ifndef CATLASS_KDA_BLOCK_MMAD_MLA_PV_HPP -#define CATLASS_KDA_BLOCK_MMAD_MLA_PV_HPP - -#include "catlass/kda_catlass.hpp" -#include "catlass/arch/kda_cross_core_sync.hpp" -#include "catlass/arch/kda_resource.hpp" -#include "catlass/kda_coord.hpp" -#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" -#include "catlass/gemm/kda_gemm_helper.hpp" -#include "catlass/kda_gemm_coord.hpp" -#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" -#include "catlass/gemm/tile/kda_tile_mmad.hpp" - -//////////////////////////////////////////////////////////////////// - -namespace Catlass::Gemm::Block { -//////////////////////////////////////////////////////////////////// - -template < - class L1TileShape_, - class L0TileShape_, - class AType_, - class BType_, - class CType_, - class BiasType_, - class TileCopy_, - class TileMmad_> -struct BlockMmad< - MmadAtlasA2MLAPV, - L1TileShape_, - L0TileShape_, - AType_, - BType_, - CType_, - BiasType_, - TileCopy_, - TileMmad_> { -public: - // Type Aliases - using DispatchPolicy = MmadAtlasA2MLAPV; - using ArchTag = typename DispatchPolicy::ArchTag; - using L1TileShape = L1TileShape_; - using L0TileShape = L0TileShape_; - using ElementA = typename AType_::Element; - using LayoutA = typename AType_::Layout; - using ElementB = typename BType_::Element; - using LayoutB = typename BType_::Layout; - using ElementC = typename CType_::Element; - using LayoutC = typename CType_::Layout; - using TileMmad = TileMmad_; - using CopyGmToL1A = typename TileCopy_::CopyGmToL1A; - using CopyGmToL1B = typename TileCopy_::CopyGmToL1B; - using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; - using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; - using CopyL0CToGm = typename TileCopy_::CopyL0CToGm; - using ElementAccumulator = - typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; - using LayoutAInL1 = typename CopyL1ToL0A::LayoutSrc; - using LayoutBInL1 = typename CopyL1ToL0B::LayoutSrc; - using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst; - using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst; - using LayoutCInL0 = layout::zN; - - using L1AAlignHelper = Gemm::helper::L1AlignHelper; - using L1BAlignHelper = Gemm::helper::L1AlignHelper; - - static constexpr uint32_t STAGES = DispatchPolicy::STAGES; - static constexpr uint32_t L1A_SIZE = L1TileShape::M * L1TileShape::K * sizeof(ElementA); - static constexpr uint32_t L1B_SIZE = L1TileShape::N * L1TileShape::K * sizeof(ElementB); - static constexpr uint32_t L0A_SIZE = ArchTag::L0A_SIZE; - static constexpr uint32_t L0B_SIZE = ArchTag::L0B_SIZE; - static constexpr uint32_t L0C_SIZE = ArchTag::L0C_SIZE; - static constexpr uint32_t L0A_PINGPONG_BUF_SIZE = L0A_SIZE / STAGES; - static constexpr uint32_t L0B_PINGPONG_BUF_SIZE = L0B_SIZE / STAGES; - static constexpr uint32_t L0C_PINGPONG_BUF_SIZE = L0C_SIZE / STAGES; - static constexpr uint32_t BLOCK_SIZE = 16; - static constexpr uint32_t EMBED_SPLIT_SIZE = 128; - static constexpr uint32_t EMBED_SPLIT_LOOP = 4; - - // Check LayoutC - static_assert(std::is_same_v, "LayoutC only support RowMajor yet!"); - - /// Construct - CATLASS_DEVICE - BlockMmad(Arch::Resource &resource, uint32_t l1BufAddrStart = 0) - { - // Allocate L1 memory space - l1ATensor = - resource.l1Buf.template GetBufferByByte(l1BufAddrStart + L1A_SIZE + L1B_SIZE * STAGES); - for (uint32_t i = 0; i < STAGES; i++) { - l1BTensor[i] = - resource.l1Buf.template GetBufferByByte(l1BufAddrStart + L1A_SIZE + L1B_SIZE * i); - l0ATensor[i] = resource.l0ABuf.template GetBufferByByte(L0A_PINGPONG_BUF_SIZE * i); - l0BTensor[i] = resource.l0BBuf.template GetBufferByByte(L0B_PINGPONG_BUF_SIZE * i); - l0CTensor[i] = resource.l0CBuf.template GetBufferByByte(L0C_PINGPONG_BUF_SIZE * i); - } - } - - /// Destructor - CATLASS_DEVICE - ~BlockMmad() {} - - /// Perform a block-scoped matrix multiply-accumulate - CATLASS_DEVICE - void operator()( - AscendC::GlobalTensor gA, - AscendC::GlobalTensor gC, - LayoutA layoutA, LayoutB layoutB, LayoutC layoutC, - GemmCoord actualShape, uint32_t &nIdx, Arch::CrossCoreFlag softmaxReady) - { - uint32_t rowNum = actualShape.m(); - uint32_t vSeqTile = actualShape.k(); - uint32_t embed = actualShape.n(); - uint32_t embedSplitSize = EMBED_SPLIT_SIZE; - uint32_t embedSplitLoopV = EMBED_SPLIT_LOOP; - uint32_t rowNumRound = RoundUp(rowNum); - uint32_t embedSplitSizeRound = RoundUp(embedSplitSize); - uint32_t vSeqTileRound = RoundUp(vSeqTile); - uint32_t L1BPingPongFlag = (nIdx - 1) % 2; - uint32_t L0APingPongFlag = (nIdx - 1) % 2; - - for (uint32_t embedSplitIdx = 0; embedSplitIdx < embedSplitLoopV; embedSplitIdx++) { - uint32_t L0CPingPongFlag = (nIdx + embedSplitIdx) % 2; - uint32_t L0BPingPongFlag = (embedSplitIdx + 1) % 2; - AscendC::WaitFlag(L0BPingPongFlag + 2); - LayoutBInL1 layoutBInL1 = LayoutBInL1::template MakeLayout(vSeqTile, embed); - LayoutBInL0 layoutBInL0 = LayoutBInL0::template MakeLayout(vSeqTile, embedSplitSize); - // copy V from L1 to L0B - copyL1ToL0B( - l0BTensor[L0BPingPongFlag], - l1BTensor[L1BPingPongFlag][embedSplitIdx * vSeqTileRound * EMBED_SPLIT_SIZE], - layoutBInL0, layoutBInL1); - if (embedSplitIdx == embedSplitLoopV - 1) { - AscendC::SetFlag(L1BPingPongFlag); - } - - if (embedSplitIdx == 0) { - LayoutAInL1 layoutAInL1 = LayoutAInL1::template MakeLayout(rowNum, vSeqTile); - Arch::CrossCoreWaitFlag(softmaxReady); - AscendC::WaitFlag(EVENT_ID7); - // copy P to L1 - copyGmToL1A(l1ATensor, gA, layoutAInL1, layoutA); - AscendC::SetFlag(EVENT_ID7); - AscendC::WaitFlag(EVENT_ID7); - // move p from l1 to l0a - AscendC::WaitFlag(L0APingPongFlag); - LayoutAInL0 layoutAInL0 = LayoutAInL0::template MakeLayout(rowNum, vSeqTile); - copyL1ToL0A(l0ATensor[L0APingPongFlag], l1ATensor, layoutAInL0, layoutAInL1); - AscendC::SetFlag(EVENT_ID7); - } - - AscendC::SetFlag(L0BPingPongFlag); - AscendC::WaitFlag(L0BPingPongFlag); - AscendC::WaitFlag(L0CPingPongFlag); - // mmad - tileMmad( - l0CTensor[L0CPingPongFlag], l0ATensor[L0APingPongFlag], l0BTensor[L0BPingPongFlag], - rowNumRound, embedSplitSize, vSeqTile); - AscendC::SetFlag(L0BPingPongFlag + 2); - if (embedSplitIdx == embedSplitLoopV - 1) { - AscendC::SetFlag(L0APingPongFlag); - } - AscendC::SetFlag(L0CPingPongFlag); - - AscendC::WaitFlag(L0CPingPongFlag); - auto blockShape = MakeCoord(rowNum, embedSplitSize); - auto layoutInL0C = LayoutCInL0::MakeLayoutInL0C(blockShape); - auto layoutCSplitK = layoutC.GetTileLayout(MakeCoord(rowNumRound, embedSplitSizeRound)); - // copy Otmp to gm - copyL0CToGm( - gC[embedSplitIdx * embedSplitSizeRound], l0CTensor[L0CPingPongFlag], - layoutCSplitK, layoutInL0C); - AscendC::SetFlag(L0CPingPongFlag); - } - } - -protected: - /// Data members - AscendC::LocalTensor l1ATensor; - AscendC::LocalTensor l1BTensor[STAGES]; - AscendC::LocalTensor l0ATensor[STAGES]; - AscendC::LocalTensor l0BTensor[STAGES]; - AscendC::LocalTensor l0CTensor[STAGES]; - - TileMmad tileMmad; - CopyGmToL1A copyGmToL1A; - CopyGmToL1B copyGmToL1B; - CopyL1ToL0A copyL1ToL0A; - CopyL1ToL0B copyL1ToL0B; - CopyL0CToGm copyL0CToGm; -}; - -//////////////////////////////////////////////////////////////////// - -} // namespace Catlass::Gemm::Block - -#endif // CATLASS_KDA_BLOCK_MMAD_MLA_PV_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_mla_pv_tp1_spec.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_mla_pv_tp1_spec.hpp deleted file mode 100644 index 36a8ba0188..0000000000 --- a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_mla_pv_tp1_spec.hpp +++ /dev/null @@ -1,221 +0,0 @@ -/** - * Copyright (c) 2025 Huawei Technologies Co., Ltd. - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of - * CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ - -#ifndef CATLASS_KDA_BLOCK_MMAD_MLA_PV_TP1_SPEC_HPP -#define CATLASS_KDA_BLOCK_MMAD_MLA_PV_TP1_SPEC_HPP - -#include "catlass/kda_catlass.hpp" -#include "catlass/arch/kda_cross_core_sync.hpp" -#include "catlass/arch/kda_resource.hpp" -#include "catlass/kda_coord.hpp" -#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" -#include "catlass/gemm/kda_gemm_helper.hpp" -#include "catlass/kda_gemm_coord.hpp" -#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" -#include "catlass/gemm/tile/kda_tile_mmad.hpp" - -//////////////////////////////////////////////////////////////////// - -namespace Catlass::Gemm::Block { -//////////////////////////////////////////////////////////////////// - -template < - class L1TileShape_, - class L0TileShape_, - class AType_, - class BType_, - class CType_, - class BiasType_, - class TileCopy_, - class TileMmad_> -struct BlockMmad< - MmadAtlasA2MLAPVTp1Spec, - L1TileShape_, - L0TileShape_, - AType_, - BType_, - CType_, - BiasType_, - TileCopy_, - TileMmad_> { -public: - // Type Aliases - using DispatchPolicy = MmadAtlasA2MLAPV; - using ArchTag = typename DispatchPolicy::ArchTag; - using L1TileShape = L1TileShape_; - using L0TileShape = L0TileShape_; - using ElementA = typename AType_::Element; - using LayoutA = typename AType_::Layout; - using ElementB = typename BType_::Element; - using LayoutB = typename BType_::Layout; - using ElementC = typename CType_::Element; - using LayoutC = typename CType_::Layout; - using TileMmad = TileMmad_; - using CopyGmToL1A = typename TileCopy_::CopyGmToL1A; - using CopyGmToL1B = typename TileCopy_::CopyGmToL1B; - using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; - using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; - using CopyL0CToGm = typename TileCopy_::CopyL0CToGm; - using ElementAccumulator = - typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; - using LayoutAInL1 = typename CopyL1ToL0A::LayoutSrc; - using LayoutBInL1 = typename CopyL1ToL0B::LayoutSrc; - using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst; - using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst; - using LayoutCInL0 = layout::zN; - - using L1AAlignHelper = Gemm::helper::L1AlignHelper; - using L1BAlignHelper = Gemm::helper::L1AlignHelper; - - static constexpr uint32_t STAGES = DispatchPolicy::STAGES; - static constexpr uint32_t L0A_SIZE = ArchTag::L0A_SIZE; - static constexpr uint32_t L0B_SIZE = ArchTag::L0B_SIZE; - static constexpr uint32_t L0C_SIZE = ArchTag::L0C_SIZE; - static constexpr uint32_t L0A_PINGPONG_BUF_SIZE = L0A_SIZE / STAGES; - static constexpr uint32_t L0B_PINGPONG_BUF_SIZE = L0B_SIZE / STAGES; - static constexpr uint32_t L0C_PINGPONG_BUF_SIZE = L0C_SIZE / STAGES; - static constexpr uint32_t UNIT_BLOCK_STACK_NUM = 4; - static constexpr uint32_t EMBED_SPLIT_LOOP = 4; - static constexpr uint32_t BLOCK_SIZE = 16; - static constexpr uint32_t EMBED_SPLIT_SIZE = 128; - static constexpr uint32_t L1_PV_ADDR_START = 311296; // reserved for Q(no db) and K(db) - static constexpr uint32_t L1A_SIZE = L1TileShape::M * L1TileShape::N * sizeof(ElementA); - static constexpr uint32_t L1B_SIZE = L1TileShape::N * EMBED_SPLIT_SIZE * sizeof(ElementB); - - /// Construct - CATLASS_DEVICE - BlockMmad(Arch::Resource &resource, uint32_t l1BufAddrStart = L1_PV_ADDR_START) - { - // Allocate L1 memory space - for (uint32_t i = 0; i < STAGES; i++) { - l1ATensor[i] = - resource.l1Buf.template GetBufferByByte(l1BufAddrStart + L1B_SIZE * STAGES + L1A_SIZE * i); - l1BTensor[i] = resource.l1Buf.template GetBufferByByte(l1BufAddrStart + L1B_SIZE * i); - l0ATensor[i] = resource.l0ABuf.template GetBufferByByte(L0A_PINGPONG_BUF_SIZE * i); - l0BTensor[i] = resource.l0BBuf.template GetBufferByByte(L0B_PINGPONG_BUF_SIZE * i); - l0CTensor[i] = resource.l0CBuf.template GetBufferByByte(L0C_PINGPONG_BUF_SIZE * i); - } - } - - /// Destructor - CATLASS_DEVICE - ~BlockMmad() {} - - /// Perform a block-scoped matrix multiply-accumulate - CATLASS_DEVICE - void operator()(AscendC::GlobalTensor gA, AscendC::GlobalTensor gB, - AscendC::GlobalTensor gblockTable, AscendC::GlobalTensor gC, LayoutA layoutA, - LayoutB layoutB, LayoutC layoutC, GemmCoord actualShape, uint32_t &nIdx, uint32_t &nLoop, - uint32_t &blockSize, uint32_t kvSeqlen, Arch::CrossCoreFlag softmaxReady) - { - uint32_t rowNum = actualShape.m(); - uint32_t stackSeqTile = actualShape.k(); - uint32_t seqTile = blockSize; - uint32_t embed = actualShape.n(); - uint32_t embedSplitSize = EMBED_SPLIT_SIZE; - uint32_t embedSplitLoopV = EMBED_SPLIT_LOOP; - uint32_t rowNumRound = RoundUp(rowNum); - uint32_t embedSplitSizeRound = RoundUp(embedSplitSize); - uint32_t stackSeqTileRound = layoutA.stride(0); - uint32_t seqTileRound = RoundUp(seqTile); - - for (uint32_t embedSplitIdx = 0; embedSplitIdx < embedSplitLoopV; embedSplitIdx++) { - uint32_t L0CPingPongFlag = embedSplitIdx % 2; - for (uint32_t blockStackIdx = 0; - (blockStackIdx < UNIT_BLOCK_STACK_NUM) && ((nIdx + blockStackIdx) < (nLoop + UNIT_BLOCK_STACK_NUM)); - blockStackIdx++) { - uint32_t nIdxActual = nIdx + blockStackIdx; - if (nIdxActual == (nLoop + UNIT_BLOCK_STACK_NUM - 1)) { - seqTile = (kvSeqlen - (nIdxActual - UNIT_BLOCK_STACK_NUM) * blockSize); - } else { - seqTile = blockSize; - } - seqTileRound = RoundUp(seqTile); - uint32_t L1ABPingPongFlag = nIdxActual % 2; - uint32_t L0ABPingPongFlag = nIdxActual % 2; - uint32_t blockTableId = gblockTable.GetValue(nIdxActual - UNIT_BLOCK_STACK_NUM); - uint64_t kvOffset = (uint64_t)blockTableId * blockSize * embed + embedSplitIdx * embedSplitSize; - - AscendC::WaitFlag(L1ABPingPongFlag + 4); - auto layoutUnitBSplitN = layoutB.GetTileLayout(MakeCoord(seqTile, embedSplitSize)); - // copy V to L1 - LayoutBInL1 layoutUnitBSplitNInL1 = LayoutBInL1::template MakeLayout(seqTile, embedSplitSize); - copyGmToL1B(l1BTensor[L1ABPingPongFlag], gB[kvOffset], layoutUnitBSplitNInL1, layoutUnitBSplitN); - AscendC::SetFlag(L1ABPingPongFlag); - - AscendC::WaitFlag(L1ABPingPongFlag); - AscendC::WaitFlag(L0ABPingPongFlag + 2); - LayoutBInL0 layoutUnitBSplitNInL0 = LayoutBInL0::template MakeLayout(seqTile, embedSplitSize); - // copy V from L1 to L0B - copyL1ToL0B(l0BTensor[L0ABPingPongFlag], l1BTensor[L1ABPingPongFlag], layoutUnitBSplitNInL0, - layoutUnitBSplitNInL1); - AscendC::SetFlag(L1ABPingPongFlag + 4); - - if (embedSplitIdx == 0 && blockStackIdx == 0) { - Arch::CrossCoreWaitFlag(softmaxReady); - } - AscendC::WaitFlag(L1ABPingPongFlag + 6); - auto layoutASplitK = layoutA.GetTileLayout(MakeCoord(rowNum, seqTile)); - LayoutAInL1 layoutASplitKInL1 = LayoutAInL1::template MakeLayout(rowNum, seqTile); - // copy P to L1 - copyGmToL1A(l1ATensor[L1ABPingPongFlag], gA[blockStackIdx * blockSize], layoutASplitKInL1, - layoutASplitK); - AscendC::SetFlag(EVENT_ID7); - // copy P to l0a - AscendC::WaitFlag(EVENT_ID7); - AscendC::WaitFlag(L0ABPingPongFlag); - LayoutAInL0 layoutASplitKInL0 = LayoutAInL0::template MakeLayout(rowNum, seqTile); - copyL1ToL0A(l0ATensor[L0ABPingPongFlag], l1ATensor[L1ABPingPongFlag], layoutASplitKInL0, - layoutASplitKInL1); - AscendC::SetFlag(L1ABPingPongFlag + 6); - AscendC::SetFlag(L0ABPingPongFlag); - // mmad - AscendC::WaitFlag(L0ABPingPongFlag); - if (blockStackIdx == 0) { - AscendC::WaitFlag(L0CPingPongFlag); - } - tileMmad(l0CTensor[L0CPingPongFlag], l0ATensor[L0ABPingPongFlag], l0BTensor[L0ABPingPongFlag], - rowNumRound, embedSplitSize, seqTile, blockStackIdx == 0); - AscendC::SetFlag(L0ABPingPongFlag + 2); - AscendC::SetFlag(L0ABPingPongFlag); - } - // copy Otmp to gm - AscendC::SetFlag(L0CPingPongFlag); - AscendC::WaitFlag(L0CPingPongFlag); - auto blockShape = MakeCoord(rowNum, embedSplitSize); - auto layoutInL0C = LayoutCInL0::MakeLayoutInL0C(blockShape); - auto layoutCSplitN = layoutC.GetTileLayout(MakeCoord(rowNumRound, embedSplitSizeRound)); - copyL0CToGm(gC[embedSplitIdx * embedSplitSizeRound], l0CTensor[L0CPingPongFlag], layoutCSplitN, - layoutInL0C); - AscendC::SetFlag(L0CPingPongFlag); - } - } - -protected: - /// Data members - AscendC::LocalTensor l1ATensor[STAGES]; - AscendC::LocalTensor l1BTensor[STAGES]; - AscendC::LocalTensor l0ATensor[STAGES]; - AscendC::LocalTensor l0BTensor[STAGES]; - AscendC::LocalTensor l0CTensor[STAGES]; - - TileMmad tileMmad; - CopyGmToL1A copyGmToL1A; - CopyGmToL1B copyGmToL1B; - CopyL1ToL0A copyL1ToL0A; - CopyL1ToL0B copyL1ToL0B; - CopyL0CToGm copyL0CToGm; -}; - -//////////////////////////////////////////////////////////////////// - -} // namespace Catlass::Gemm::Block - -#endif // CATLASS_KDA_BLOCK_MMAD_MLA_PV_TP1_SPEC_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_mla_qk.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_mla_qk.hpp deleted file mode 100644 index 3b1fedb199..0000000000 --- a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_mla_qk.hpp +++ /dev/null @@ -1,241 +0,0 @@ -/** - * Copyright (c) 2025 Huawei Technologies Co., Ltd. - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of - * CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ - -#ifndef CATLASS_KDA_BLOCK_MMAD_MLA_QK_HPP -#define CATLASS_KDA_BLOCK_MMAD_MLA_QK_HPP - -#include "catlass/kda_catlass.hpp" -#include "catlass/arch/kda_resource.hpp" -#include "catlass/kda_coord.hpp" -#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" -#include "catlass/gemm/kda_gemm_helper.hpp" -#include "catlass/kda_gemm_coord.hpp" -#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" -#include "catlass/gemm/tile/kda_tile_mmad.hpp" - -//////////////////////////////////////////////////////////////////// - -namespace Catlass::Gemm::Block { -//////////////////////////////////////////////////////////////////// - -template < - class L1TileShape_, - class L0TileShape_, - class AType_, - class BType_, - class CType_, - class BiasType_, - class TileCopy_, - class TileMmad_> -struct BlockMmad< - MmadAtlasA2MLAQK, - L1TileShape_, - L0TileShape_, - AType_, - BType_, - CType_, - BiasType_, - TileCopy_, - TileMmad_> { -public: - // Type Aliases - using DispatchPolicy = MmadAtlasA2MLAQK; - using ArchTag = typename DispatchPolicy::ArchTag; - using L1TileShape = L1TileShape_; - using L0TileShape = L0TileShape_; - using ElementA = typename AType_::Element; - using LayoutA = typename AType_::Layout; - using ElementB = typename BType_::Element; - using LayoutB = typename BType_::Layout; - using ElementC = typename CType_::Element; - using LayoutC = typename CType_::Layout; - using TileMmad = TileMmad_; - using CopyGmToL1A = typename TileCopy_::CopyGmToL1A; - using CopyGmToL1B = typename TileCopy_::CopyGmToL1B; - using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; - using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; - using CopyL0CToGm = typename TileCopy_::CopyL0CToGm; - using ElementAccumulator = - typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; - using LayoutAInL1 = typename CopyL1ToL0A::LayoutSrc; - using LayoutBInL1 = typename CopyL1ToL0B::LayoutSrc; - using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst; - using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst; - using LayoutCInL0 = layout::zN; - - using L1AAlignHelper = Gemm::helper::L1AlignHelper; - using L1BAlignHelper = Gemm::helper::L1AlignHelper; - - static constexpr uint32_t STAGES = DispatchPolicy::STAGES; - static constexpr uint32_t L1A_SIZE = L1TileShape::M * L1TileShape::K * sizeof(ElementA); - static constexpr uint32_t L1B_SIZE = L1TileShape::N * L1TileShape::K * sizeof(ElementB); - static constexpr uint32_t L0A_SIZE = ArchTag::L0A_SIZE; - static constexpr uint32_t L0B_SIZE = ArchTag::L0B_SIZE; - static constexpr uint32_t L0C_SIZE = ArchTag::L0C_SIZE; - static constexpr uint32_t L0A_PINGPONG_BUF_SIZE = L0A_SIZE / STAGES; - static constexpr uint32_t L0B_PINGPONG_BUF_SIZE = L0B_SIZE / STAGES; - static constexpr uint32_t L0C_PINGPONG_BUF_SIZE = L0C_SIZE / STAGES; - static constexpr uint32_t BLOCK_SIZE = 16; - static constexpr uint32_t EMBED_SPLIT_SIZE = 128; - static constexpr uint32_t EMBED_SPLIT_LOOP = 5; - - // Check LayoutC - static_assert(std::is_same_v, "LayoutC only support RowMajor yet!"); - - /// Construct - CATLASS_DEVICE - BlockMmad(Arch::Resource &resource, uint32_t l1BufAddrStart = 0) - { - // Allocate L1 memory space - l1ATensor = resource.l1Buf.template GetBufferByByte(l1BufAddrStart); - for (uint32_t i = 0; i < STAGES; i++) { - l1BTensor[i] = - resource.l1Buf.template GetBufferByByte(l1BufAddrStart + L1A_SIZE + L1B_SIZE * i); - l0ATensor[i] = resource.l0ABuf.template GetBufferByByte(L0A_PINGPONG_BUF_SIZE * i); - l0BTensor[i] = resource.l0BBuf.template GetBufferByByte(L0B_PINGPONG_BUF_SIZE * i); - l0CTensor[i] = resource.l0CBuf.template GetBufferByByte(L0C_PINGPONG_BUF_SIZE * i); - } - } - - /// Destructor - CATLASS_DEVICE - ~BlockMmad() {} - - /// Perform a block-scoped matrix multiply-accumulate - CATLASS_DEVICE - void operator()( - AscendC::GlobalTensor gA, - AscendC::GlobalTensor gARope, - AscendC::GlobalTensor gB, - AscendC::GlobalTensor gBRope, - AscendC::GlobalTensor gC, - LayoutA layoutA, LayoutA layoutARope, LayoutB layoutB, LayoutB layoutBRope, LayoutC layoutC, - GemmCoord actualShape, MatrixCoord qShapeSingleNd, - uint32_t &qHeads, uint32_t &nIdx) - { - uint32_t rowNum = actualShape.m(); - uint32_t kSeqTile = actualShape.n(); - uint32_t embed = layoutB.shape(0); - uint32_t embedRope = layoutBRope.shape(0); - uint32_t embedCat = actualShape.k(); - uint32_t embedSplitSize = EMBED_SPLIT_SIZE; - uint32_t embedSplitLoopK = EMBED_SPLIT_LOOP; - uint32_t curHeadNum = qShapeSingleNd.row(); - uint32_t tokenNumPerHead = rowNum / curHeadNum; - uint32_t kSeqTileRound = RoundUp(kSeqTile); - uint32_t rowNumRound = RoundUp(rowNum); - uint32_t l1KvPingPongFlag = nIdx % 2; - - if (nIdx == 0) { - // copy Q to L1 - auto layoutASingleNd = layoutA.GetTileLayout(MakeCoord(curHeadNum, embed)); - LayoutAInL1 layoutAInL1 = LayoutAInL1::template MakeLayout(rowNum, embed); - copyGmToL1A( - l1ATensor, gA, - layoutAInL1, layoutASingleNd, - tokenNumPerHead, qHeads * embed, tokenNumPerHead, BLOCK_SIZE, rowNumRound); - - // copy QRope to L1 - auto layoutARopeSingleNd = layoutARope.GetTileLayout(MakeCoord(curHeadNum, embedRope)); - LayoutAInL1 layoutARopeInL1 = LayoutAInL1::template MakeLayout(rowNum, embedRope); - copyGmToL1A( - l1ATensor[rowNumRound * embed], gARope, - layoutARopeInL1, layoutARopeSingleNd, - tokenNumPerHead, qHeads * embedRope, tokenNumPerHead, BLOCK_SIZE, rowNumRound); - AscendC::SetFlag(EVENT_ID0); - AscendC::WaitFlag(EVENT_ID0); - } - - AscendC::WaitFlag(l1KvPingPongFlag); - // copy K to L1 - LayoutBInL1 layoutBInL1 = LayoutBInL1::template MakeLayout(embed, kSeqTile); - copyGmToL1B(l1BTensor[l1KvPingPongFlag], gB, layoutBInL1, layoutB); - AscendC::SetFlag(l1KvPingPongFlag); - - AscendC::WaitFlag(l1KvPingPongFlag + 2); - // copy KRope to L1 - LayoutBInL1 layoutBRopeInL1 = LayoutBInL1::template MakeLayout(embedRope, kSeqTile); - copyGmToL1B(l1BTensor[l1KvPingPongFlag][kSeqTileRound * embed], gBRope, layoutBRopeInL1, layoutBRope); - AscendC::SetFlag(l1KvPingPongFlag + 2); - - for (uint32_t embedSplitIdx = 0; embedSplitIdx < embedSplitLoopK; embedSplitIdx++) { - uint32_t l0ABPingPongFlag = embedSplitIdx % 2; - if (embedSplitIdx == embedSplitLoopK - 1) { - embedSplitSize = embedRope; - } - // copy Q from L1 to l0a - AscendC::WaitFlag(l0ABPingPongFlag); - LayoutAInL1 layoutACatInL1 = LayoutAInL1::template MakeLayout(rowNum, embedCat); - LayoutAInL0 layoutACatInL0 = LayoutAInL0::template MakeLayout(rowNum, embedSplitSize); - copyL1ToL0A( - l0ATensor[l0ABPingPongFlag], l1ATensor[embedSplitIdx * rowNumRound * EMBED_SPLIT_SIZE], - layoutACatInL0, layoutACatInL1); - AscendC::SetFlag(l0ABPingPongFlag); - - if (embedSplitIdx == 0) { - AscendC::WaitFlag(l1KvPingPongFlag); - } else if (embedSplitIdx == embedSplitLoopK - 1) { - AscendC::WaitFlag(l1KvPingPongFlag + 2); - } - AscendC::WaitFlag(l0ABPingPongFlag + 2); - // copy K from L1 to l0b - LayoutBInL1 layoutBCatInL1 = LayoutBInL1::template MakeLayout(embedCat, kSeqTile); - LayoutBInL0 layoutBCatInL0 = LayoutBInL0::template MakeLayout(embedSplitSize, kSeqTile); - copyL1ToL0B(l0BTensor[l0ABPingPongFlag], - l1BTensor[l1KvPingPongFlag][embedSplitIdx * kSeqTileRound * EMBED_SPLIT_SIZE], - layoutBCatInL0, layoutBCatInL1); - if (embedSplitIdx == embedSplitLoopK - 1) { - AscendC::SetFlag(l1KvPingPongFlag + 2); - } - AscendC::SetFlag(l0ABPingPongFlag + 2); - // mmad - AscendC::WaitFlag(l0ABPingPongFlag); - AscendC::WaitFlag(l0ABPingPongFlag + 2); - if (embedSplitIdx == 0) { - AscendC::WaitFlag(l1KvPingPongFlag); - } - // mmad - tileMmad( - l0CTensor[l1KvPingPongFlag], l0ATensor[l0ABPingPongFlag], l0BTensor[l0ABPingPongFlag], - rowNumRound, kSeqTile, embedSplitSize, embedSplitIdx == 0); - AscendC::SetFlag(l0ABPingPongFlag); - AscendC::SetFlag(l0ABPingPongFlag + 2); - } - // copy S to gm - AscendC::SetFlag(l1KvPingPongFlag); - AscendC::WaitFlag(l1KvPingPongFlag); - auto blockShape = MakeCoord(rowNum, kSeqTile); - auto layoutInL0C = LayoutCInL0::MakeLayoutInL0C(blockShape); - // copy L0C to gm - copyL0CToGm(gC, l0CTensor[l1KvPingPongFlag], layoutC, layoutInL0C); - AscendC::SetFlag(l1KvPingPongFlag); - } - -protected: - /// Data members - AscendC::LocalTensor l1ATensor; - AscendC::LocalTensor l1BTensor[STAGES]; - AscendC::LocalTensor l0ATensor[STAGES]; - AscendC::LocalTensor l0BTensor[STAGES]; - AscendC::LocalTensor l0CTensor[STAGES]; - - TileMmad tileMmad; - CopyGmToL1A copyGmToL1A; - CopyGmToL1B copyGmToL1B; - CopyL1ToL0A copyL1ToL0A; - CopyL1ToL0B copyL1ToL0B; - CopyL0CToGm copyL0CToGm; -}; - -//////////////////////////////////////////////////////////////////// - -} // namespace Catlass::Gemm::Block - -#endif // CATLASS_KDA_BLOCK_MMAD_MLA_QK_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_mla_qk_tp1_spec.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_mla_qk_tp1_spec.hpp deleted file mode 100644 index 999cfa0bf0..0000000000 --- a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_mla_qk_tp1_spec.hpp +++ /dev/null @@ -1,259 +0,0 @@ -/** - * Copyright (c) 2025 Huawei Technologies Co., Ltd. - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of - * CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ - -#ifndef CATLASS_KDA_BLOCK_MMAD_MLA_QK_TP1_SPEC_HPP -#define CATLASS_KDA_BLOCK_MMAD_MLA_QK_TP1_SPEC_HPP - -#include "catlass/kda_catlass.hpp" -#include "catlass/arch/kda_resource.hpp" -#include "catlass/kda_coord.hpp" -#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" -#include "catlass/gemm/kda_gemm_helper.hpp" -#include "catlass/kda_gemm_coord.hpp" -#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" -#include "catlass/gemm/tile/kda_tile_mmad.hpp" - -//////////////////////////////////////////////////////////////////// - -namespace Catlass::Gemm::Block { -//////////////////////////////////////////////////////////////////// - -template < - class L1TileShape_, - class L0TileShape_, - class AType_, - class BType_, - class CType_, - class BiasType_, - class TileCopy_, - class TileMmad_> -struct BlockMmad< - MmadAtlasA2MLAQKTp1Spec, - L1TileShape_, - L0TileShape_, - AType_, - BType_, - CType_, - BiasType_, - TileCopy_, - TileMmad_> { -public: - // Type Aliases - using DispatchPolicy = MmadAtlasA2MLAQKTp1Spec; - using ArchTag = typename DispatchPolicy::ArchTag; - using L1TileShape = L1TileShape_; - using L0TileShape = L0TileShape_; - using ElementA = typename AType_::Element; - using LayoutA = typename AType_::Layout; - using ElementB = typename BType_::Element; - using LayoutB = typename BType_::Layout; - using ElementC = typename CType_::Element; - using LayoutC = typename CType_::Layout; - using TileMmad = TileMmad_; - using CopyGmToL1A = typename TileCopy_::CopyGmToL1A; - using CopyGmToL1B = typename TileCopy_::CopyGmToL1B; - using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; - using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; - using CopyL0CToGm = typename TileCopy_::CopyL0CToGm; - using ElementAccumulator = - typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; - using LayoutAInL1 = typename CopyL1ToL0A::LayoutSrc; - using LayoutBInL1 = typename CopyL1ToL0B::LayoutSrc; - using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst; - using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst; - using LayoutCInL0 = layout::zN; - - using L1AAlignHelper = Gemm::helper::L1AlignHelper; - using L1BAlignHelper = Gemm::helper::L1AlignHelper; - - static constexpr uint32_t STAGES = DispatchPolicy::STAGES; - static constexpr uint32_t L1A_SIZE = L1TileShape::M * L1TileShape::K * sizeof(ElementA); - static constexpr uint32_t L0A_SIZE = ArchTag::L0A_SIZE; - static constexpr uint32_t L0B_SIZE = ArchTag::L0B_SIZE; - static constexpr uint32_t L0C_SIZE = ArchTag::L0C_SIZE; - static constexpr uint32_t L0A_PINGPONG_BUF_SIZE = L0A_SIZE / STAGES; - static constexpr uint32_t L0B_PINGPONG_BUF_SIZE = L0B_SIZE / STAGES; - static constexpr uint32_t L0C_PINGPONG_BUF_SIZE = L0C_SIZE / STAGES; - static constexpr uint32_t UNIT_BLOCK_STACK_NUM = 4; - static constexpr uint32_t BLOCK_SIZE = 16; - static constexpr uint32_t EMBED_SPLIT_SIZE = 128; - static constexpr uint32_t EMBED_ROPE = 64; - static constexpr uint32_t GM_L1_EMBED_SPLIT_SIZE = 256; - static constexpr uint32_t EMBED_SPLIT_LOOP = 5; - static constexpr uint32_t L1B_SIZE = L1TileShape::N * GM_L1_EMBED_SPLIT_SIZE * sizeof(ElementB); - static constexpr uint32_t L1BROPE_SIZE = L1TileShape::N * EMBED_ROPE * sizeof(ElementB); - static constexpr uint32_t L1B_ROPE_START = L1TileShape::N * GM_L1_EMBED_SPLIT_SIZE; - - /// Construct - CATLASS_DEVICE - BlockMmad(Arch::Resource &resource, uint32_t l1BufAddrStart = 0) - { - // Allocate L1 memory space - l1ATensor = resource.l1Buf.template GetBufferByByte(l1BufAddrStart); - for (uint32_t i = 0; i < STAGES; i++) { - l1BTensor[i] = resource.l1Buf.template GetBufferByByte(l1BufAddrStart + L1A_SIZE + L1B_SIZE * i); - l1BRopeTensor[i] = resource.l1Buf.template GetBufferByByte(l1BufAddrStart + L1A_SIZE + - L1B_SIZE * STAGES + L1BROPE_SIZE * i); - l0ATensor[i] = resource.l0ABuf.template GetBufferByByte(L0A_PINGPONG_BUF_SIZE * i); - l0BTensor[i] = resource.l0BBuf.template GetBufferByByte(L0B_PINGPONG_BUF_SIZE * i); - l0CTensor[i] = resource.l0CBuf.template GetBufferByByte(L0C_PINGPONG_BUF_SIZE * i); - } - } - - /// Destructor - CATLASS_DEVICE - ~BlockMmad() {} - - /// Perform a block-scoped matrix multiply-accumulate - CATLASS_DEVICE - void operator()(AscendC::GlobalTensor gA, AscendC::GlobalTensor gARope, - AscendC::GlobalTensor gB, AscendC::GlobalTensor gBRope, - AscendC::GlobalTensor gblockTable, AscendC::GlobalTensor gC, LayoutA layoutA, - LayoutA layoutARope, LayoutB layoutB, LayoutB layoutBRope, LayoutC layoutC, GemmCoord actualShape, - uint32_t &nIdx, uint32_t &nLoop, uint32_t &blockSize, uint32_t kvSeqlen) - { - uint32_t rowNum = actualShape.m(); - uint32_t stackSeqTile = actualShape.n(); - uint32_t seqTile = blockSize; - uint32_t embed = layoutA.shape(1); - uint32_t embedRope = layoutARope.shape(1); - uint32_t embedCat = actualShape.k(); - uint32_t embedSplitGm2L1 = GM_L1_EMBED_SPLIT_SIZE; - uint32_t stackSeqTileRound = layoutC.shape(1); - uint32_t rowNumRound = layoutC.shape(0); - uint32_t seqTileRound = RoundUp(seqTile); - - if (nIdx == 0) { - // copy Q to L1 - LayoutAInL1 layoutAInL1 = LayoutAInL1::template MakeLayout(rowNum, embed); - copyGmToL1A(l1ATensor, gA, layoutAInL1, layoutA); - - // copy QRope to L1 - LayoutAInL1 layoutARopeInL1 = LayoutAInL1::template MakeLayout(rowNum, embedRope); - copyGmToL1A(l1ATensor[rowNumRound * embed], gARope, layoutARopeInL1, layoutARope); - AscendC::SetFlag(EVENT_ID0); - AscendC::WaitFlag(EVENT_ID0); - } - - for (uint32_t blockStackIdx = 0; (blockStackIdx < UNIT_BLOCK_STACK_NUM) && ((nIdx + blockStackIdx) < nLoop); - blockStackIdx++) { - uint32_t nIdxActual = nIdx + blockStackIdx; - uint32_t L0CPingPongFlag = nIdxActual % 2; - uint32_t L1BRopePingPongFlag = nIdxActual % 2; - if (nIdxActual == (nLoop - 1)) { - seqTile = (kvSeqlen - nIdxActual * blockSize); - seqTileRound = RoundUp(seqTile); - } - uint32_t blockTableId = gblockTable.GetValue(nIdxActual); - uint64_t kvOffset = (uint64_t)blockTableId * blockSize * embed; - uint64_t kvOffsetRope = (uint64_t)blockTableId * blockSize * embedRope; - uint64_t l1bSplitOffset = 0; - uint32_t embedSplitSize = EMBED_SPLIT_SIZE; - uint32_t embedSplitLoopK = EMBED_SPLIT_LOOP; - for (uint32_t embedSplitIdx = 0; embedSplitIdx < embedSplitLoopK; embedSplitIdx++) { - uint32_t L0ABPingPongFlag = (blockStackIdx + embedSplitIdx) % 2; - uint32_t innerSplitIdxL12L0 = embedSplitIdx % 2; - uint32_t L1BPingPongFlag = embedSplitIdx / 2; - if (embedSplitIdx == 4) { - embedSplitSize = embedRope; - } - // copy Q to l0a - AscendC::WaitFlag(L0ABPingPongFlag); - LayoutAInL1 layoutACatSplitKInL1 = LayoutAInL1::template MakeLayout(rowNum, embedSplitSize); - LayoutAInL0 layoutACatSplitKInL0 = LayoutAInL0::template MakeLayout(rowNum, embedSplitSize); - copyL1ToL0A(l0ATensor[L0ABPingPongFlag], l1ATensor[embedSplitIdx * rowNumRound * EMBED_SPLIT_SIZE], - layoutACatSplitKInL0, layoutACatSplitKInL1); - AscendC::SetFlag(L0ABPingPongFlag); - - if (embedSplitIdx == 0 || embedSplitIdx == 2) { - // copy K to l1b - AscendC::WaitFlag(L1BPingPongFlag); - auto layoutUnitBSplitK = layoutB.GetTileLayout(MakeCoord(embedSplitGm2L1, seqTile)); - LayoutBInL1 layoutUnitBSplitKInL1 = - LayoutBInL1::template MakeLayout(embedSplitGm2L1, seqTile); - copyGmToL1B(l1BTensor[L1BPingPongFlag], gB[kvOffset + embedSplitIdx * embedSplitSize], - layoutUnitBSplitKInL1, layoutUnitBSplitK); - AscendC::SetFlag(L1BPingPongFlag); - AscendC::WaitFlag(L1BPingPongFlag); - } else if (embedSplitIdx == 4) { - // copy KRope to L1 - AscendC::WaitFlag(L1BRopePingPongFlag + 2); - auto layoutUnitBRope = layoutBRope.GetTileLayout(MakeCoord(embedRope, seqTile)); - LayoutBInL1 layoutBRopeInL1 = LayoutBInL1::template MakeLayout(embedRope, seqTile); - copyGmToL1B(l1BRopeTensor[L1BRopePingPongFlag], gBRope[kvOffsetRope], layoutBRopeInL1, - layoutUnitBRope); - AscendC::SetFlag(L1BRopePingPongFlag); - AscendC::WaitFlag(L1BRopePingPongFlag); - } - // copy K to l0b - AscendC::WaitFlag(L0ABPingPongFlag + 2); - LayoutBInL1 layoutBCatSplitKInL1 = LayoutBInL1::template MakeLayout(embedSplitSize, seqTile); - LayoutBInL0 layoutBCatSplitKInL0 = LayoutBInL0::template MakeLayout(embedSplitSize, seqTile); - if (embedSplitIdx != 4) { - copyL1ToL0B(l0BTensor[L0ABPingPongFlag], - l1BTensor[L1BPingPongFlag][innerSplitIdxL12L0 * EMBED_SPLIT_SIZE * seqTileRound], - layoutBCatSplitKInL0, layoutBCatSplitKInL1); - } else { - copyL1ToL0B(l0BTensor[L0ABPingPongFlag], l1BRopeTensor[L1BRopePingPongFlag], layoutBCatSplitKInL0, - layoutBCatSplitKInL1); - } - - if (embedSplitIdx == 1 || embedSplitIdx == 3) { - AscendC::SetFlag(L1BPingPongFlag); - } else if (embedSplitIdx == 4) { - AscendC::SetFlag(L1BRopePingPongFlag + 2); - } - AscendC::SetFlag(L0ABPingPongFlag + 2); - - AscendC::WaitFlag(L0ABPingPongFlag); - AscendC::WaitFlag(L0ABPingPongFlag + 2); - if (embedSplitIdx == 0) { - AscendC::WaitFlag(L0CPingPongFlag); - } - // mmad - tileMmad(l0CTensor[L0CPingPongFlag], l0ATensor[L0ABPingPongFlag], l0BTensor[L0ABPingPongFlag], - rowNumRound, seqTile, embedSplitSize, embedSplitIdx == 0); - AscendC::SetFlag(L0ABPingPongFlag); - AscendC::SetFlag(L0ABPingPongFlag + 2); - } - // copy S to gm - AscendC::SetFlag(L0CPingPongFlag); - AscendC::WaitFlag(L0CPingPongFlag); - auto blockShape = MakeCoord(rowNum, seqTile); - auto layoutInL0C = LayoutCInL0::MakeLayoutInL0C(blockShape); - auto layoutCSplitN = layoutC.GetTileLayout(MakeCoord(rowNumRound, seqTileRound)); - // copy L0C to gm - copyL0CToGm(gC[blockStackIdx * blockSize], l0CTensor[L0CPingPongFlag], layoutCSplitN, layoutInL0C); - AscendC::SetFlag(L0CPingPongFlag); - } - } - -protected: - /// Data members - AscendC::LocalTensor l1ATensor; - AscendC::LocalTensor l1BTensor[STAGES]; - AscendC::LocalTensor l1BRopeTensor[STAGES]; - AscendC::LocalTensor l0ATensor[STAGES]; - AscendC::LocalTensor l0BTensor[STAGES]; - AscendC::LocalTensor l0CTensor[STAGES]; - - TileMmad tileMmad; - CopyGmToL1A copyGmToL1A; - CopyGmToL1B copyGmToL1B; - CopyL1ToL0A copyL1ToL0A; - CopyL1ToL0B copyL1ToL0B; - CopyL0CToGm copyL0CToGm; -}; - -//////////////////////////////////////////////////////////////////// - -} // namespace Catlass::Gemm::Block - -#endif // CATLASS_KDA_BLOCK_MMAD_MLA_QK_TP1_SPEC_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_pingpong.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_pingpong.hpp deleted file mode 100644 index 685821c0e4..0000000000 --- a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_pingpong.hpp +++ /dev/null @@ -1,362 +0,0 @@ -/** - * Copyright (c) 2025 Huawei Technologies Co., Ltd. - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of - * CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ - -#ifndef CATLASS_KDA_BLOCK_MMAD_PINGPONG_HPP -#define CATLASS_KDA_BLOCK_MMAD_PINGPONG_HPP - -#include "catlass/kda_catlass.hpp" -#include "catlass/arch/kda_resource.hpp" -#include "catlass/kda_coord.hpp" -#include "catlass/kda_gemm_coord.hpp" -#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" -#include "catlass/gemm/kda_gemm_helper.hpp" -#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" -#include "catlass/gemm/tile/kda_tile_mmad.hpp" - -namespace Catlass::Gemm::Block { - -template < - bool ENABLE_UNIT_FLAG_, - class L1TileShape_, - class L0TileShape_, - class AType_, - class BType_, - class CType_, - class BiasType_, - class TileCopy_, - class TileMmad_> -struct BlockMmad< - MmadAtlasA2Pingpong, - L1TileShape_, - L0TileShape_, - AType_, - BType_, - CType_, - BiasType_, - TileCopy_, - TileMmad_> { -public: - // Type Aliases - using DispatchPolicy = MmadAtlasA2Pingpong; - using ArchTag = typename DispatchPolicy::ArchTag; - using L1TileShape = L1TileShape_; - using L0TileShape = L0TileShape_; - using ElementA = typename AType_::Element; - using LayoutA = typename AType_::Layout; - using ElementB = typename BType_::Element; - using LayoutB = typename BType_::Layout; - using ElementC = typename CType_::Element; - using LayoutC = typename CType_::Layout; - using TileMmad = TileMmad_; - using CopyGmToL1A = typename TileCopy_::CopyGmToL1A; - using CopyGmToL1B = typename TileCopy_::CopyGmToL1B; - using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; - using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; - using CopyL0CToGm = typename TileCopy_::CopyL0CToGm; - using ElementAccumulator = - typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; - using LayoutAInL1 = typename CopyL1ToL0A::LayoutSrc; - using LayoutBInL1 = typename CopyL1ToL0B::LayoutSrc; - using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst; - using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst; - using LayoutCInL0 = layout::zN; - - using L1AAlignHelper = Gemm::helper::L1AlignHelper; - using L1BAlignHelper = Gemm::helper::L1AlignHelper; - - static constexpr bool ENABLE_UNIT_FLAG = DispatchPolicy::ENABLE_UNIT_FLAG; - static constexpr uint32_t STAGES = DispatchPolicy::STAGES; - static constexpr uint32_t L1A_SIZE = L1TileShape::M * L1TileShape::K * sizeof(ElementA); - static constexpr uint32_t L1B_SIZE = L1TileShape::N * L1TileShape::K * sizeof(ElementB); - static constexpr uint32_t L0A_SIZE = ArchTag::L0A_SIZE; - static constexpr uint32_t L0B_SIZE = ArchTag::L0B_SIZE; - static constexpr uint32_t L0C_SIZE = ArchTag::L0C_SIZE; - static constexpr uint32_t L0A_PINGPONG_BUF_SIZE = L0A_SIZE / STAGES; - static constexpr uint32_t L0B_PINGPONG_BUF_SIZE = L0B_SIZE / STAGES; - - // Check LayoutC - static_assert(std::is_same_v, "LayoutC only support RowMajor yet!"); - - // Check L1TileShape - static_assert((L1A_SIZE * STAGES + L1B_SIZE * STAGES) <= ArchTag::L1_SIZE, "L1TileShape exceeding the L1 space!"); - - // Check L0TileShape - static constexpr uint32_t L0A_TILE_SIZE = L0TileShape::M * L0TileShape::K * sizeof(ElementA); - static constexpr uint32_t L0B_TILE_SIZE = L0TileShape::K * L0TileShape::N * sizeof(ElementB); - static constexpr uint32_t L0C_TILE_SIZE = L0TileShape::M * L0TileShape::N * sizeof(ElementAccumulator); - static_assert((L0A_TILE_SIZE * STAGES) <= L0A_SIZE, "L0TileShape exceeding the L0A space!"); - static_assert((L0B_TILE_SIZE * STAGES) <= L0B_SIZE, "L0TileShape exceeding the L0B space!"); - static_assert(L0C_TILE_SIZE <= L0C_SIZE, "L0TileShape exceeding the L0C space!"); - - // Check tileshape - static_assert(L1TileShape::M == L0TileShape::M && L1TileShape::N == L0TileShape::N, - "The situation where the basic blocks of L1 and L0 differ on the m and n axes is not supported yet"); - static_assert(L0TileShape::K <= L1TileShape::K, "L0TileShape::K cannot exceed L1TileShape::K"); - - // 32B (256b) aligned - static_assert(Gemm::helper::TileShapeAlignChecker::_ALIGN == 256, - "Tile shape must be 32B aligned."); - - /// Construct - CATLASS_DEVICE - BlockMmad(Arch::Resource &resource, uint32_t l1BufAddrStart = 0) - { - uint32_t l1AOffset = l1BufAddrStart; - uint32_t l1BOffset = l1BufAddrStart + L1A_SIZE * STAGES; - // Init buffers - for (uint32_t i = 0; i < STAGES; i++) { - // Assign L1/L0A/L0B space for each stages - l1ATensorList[i] = resource.l1Buf.template GetBufferByByte(l1AOffset + L1A_SIZE * i); - l1BTensorList[i] = resource.l1Buf.template GetBufferByByte(l1BOffset + L1B_SIZE * i); - l0ATensorList[i] = resource.l0ABuf.template GetBufferByByte(L0A_PINGPONG_BUF_SIZE * i); - l0BTensorList[i] = resource.l0BBuf.template GetBufferByByte(L0B_PINGPONG_BUF_SIZE * i); - - // Assign event ID for each stages - l1AEventList[i] = i; - l1BEventList[i] = i + STAGES; - l0AEventList[i] = i; - l0BEventList[i] = i + STAGES; - - // The event id that needs to be set before the loop - AscendC::SetFlag(l1AEventList[i]); - AscendC::SetFlag(l1BEventList[i]); - AscendC::SetFlag(l0AEventList[i]); - AscendC::SetFlag(l0BEventList[i]); - } - l0CTensor = resource.l0CBuf.template GetBufferByByte(0); - AscendC::SetFlag(EVENT_ID0); - } - - /// Destructor - CATLASS_DEVICE - ~BlockMmad() - { - for (uint32_t i = 0; i < STAGES; i++) { - AscendC::WaitFlag(l1AEventList[i]); - AscendC::WaitFlag(l1BEventList[i]); - AscendC::WaitFlag(l0AEventList[i]); - AscendC::WaitFlag(l0BEventList[i]); - } - AscendC::WaitFlag(EVENT_ID0); - } - - /// Perform a block-scoped matrix multiply-accumulate - CATLASS_DEVICE - void operator()( - AscendC::GlobalTensor const &gmA, LayoutA const &layoutA, - AscendC::GlobalTensor const &gmB, LayoutB const &layoutB, - AscendC::GlobalTensor const &gmC, LayoutC const &layoutC, - GemmCoord const &actualShape) - { - uint32_t mRound = RoundUp(actualShape.m()); - uint32_t nRound = RoundUp(actualShape.n()); - - auto layoutAInL1 = LayoutAInL1::template MakeLayout(L1TileShape::M, L1TileShape::K); - auto layoutBInL1 = LayoutBInL1::template MakeLayout(L1TileShape::K, L1TileShape::N); - auto layoutInL0C = LayoutCInL0::MakeLayoutInL0C(MakeCoord(mRound, nRound)); - - uint32_t kActual = min(actualShape.k(), L1TileShape::K); - - // load first matrix A tile from GM to L1 - AscendC::WaitFlag(l1AEventList[l1ListId]); - auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShape.m(), kActual)); - copyGmToL1A(l1ATensorList[l1ListId], gmA, layoutAInL1, layoutTileA); - AscendC::SetFlag(l1AEventList[l1ListId]); - - // load first matrix B tile from GM to L1 - AscendC::WaitFlag(l1BEventList[l1ListId]); - auto layoutTileB = layoutB.GetTileLayout(MakeCoord(kActual, actualShape.n())); - copyGmToL1B(l1BTensorList[l1ListId], gmB, layoutBInL1, layoutTileB); - AscendC::SetFlag(l1BEventList[l1ListId]); - - if constexpr (!ENABLE_UNIT_FLAG) { - AscendC::WaitFlag(EVENT_ID0); - } - - uint32_t mPartLoop = CeilDiv(mRound); - uint32_t nPartLoop = CeilDiv(nRound); - - // main loop - uint32_t kTileCount = CeilDiv(actualShape.k()); - for (uint32_t kLoopIdx = 0; kLoopIdx < kTileCount; kLoopIdx++) { - uint32_t l1ListIdNext = (l1ListId + 1 < STAGES) ? (l1ListId + 1) : 0; - uint32_t kActualNext{0}; - // preload next tile from GM to L1 - if (kLoopIdx < kTileCount - 1) { - uint32_t kLoopIdxNext = kLoopIdx + 1; - kActualNext = (kLoopIdxNext < kTileCount - 1) ? - L1TileShape::K : - (actualShape.k() - kLoopIdxNext * L1TileShape::K); - - // Get L1 tensor for next stage - auto l1ATensor = l1ATensorList[l1ListIdNext]; - auto l1BTensor = l1BTensorList[l1ListIdNext]; - // Get GM tile for next stage - MatrixCoord gmTileAOffset{0, kLoopIdxNext * L1TileShape::K}; - MatrixCoord gmTileBOffset{kLoopIdxNext * L1TileShape::K, 0}; - auto gmTileA = gmA[layoutA.GetOffset(gmTileAOffset)]; - auto gmTileB = gmB[layoutB.GetOffset(gmTileBOffset)]; - - // load next matrix A tile from GM to L1 - AscendC::WaitFlag(l1AEventList[l1ListIdNext]); - layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShape.m(), kActualNext)); - copyGmToL1A(l1ATensor, gmTileA, layoutAInL1, layoutTileA); - AscendC::SetFlag(l1AEventList[l1ListIdNext]); - - // load next matrix B tile from GM to L1 - AscendC::WaitFlag(l1BEventList[l1ListIdNext]); - layoutTileB = layoutB.GetTileLayout(MakeCoord(kActualNext, actualShape.n())); - copyGmToL1B(l1BTensor, gmTileB, layoutBInL1, layoutTileB); - AscendC::SetFlag(l1BEventList[l1ListIdNext]); - } - - // Get L1 tensor for current stage - auto l1ATensor = l1ATensorList[l1ListId]; - auto l1BTensor = l1BTensorList[l1ListId]; - - // Get the loop nums on L0 - uint32_t kPartLoop = CeilDiv(kActual); - - for (int mPartIdx = 0; mPartIdx < mPartLoop; mPartIdx++) { - uint32_t mPartActual = (mPartIdx < mPartLoop - 1) ? - L0TileShape::M : - (mRound - mPartIdx * L0TileShape::M); - - for (int kPartIdx = 0; kPartIdx < kPartLoop; kPartIdx++) { - uint32_t kPartActual = (kPartIdx < kPartLoop - 1) ? - L0TileShape::K : - (kActual - kPartIdx * L0TileShape::K); - - // Locate the current tile on L0A - auto l0ATile = l0ATensorList[l0AListId]; - LayoutAInL0 layoutAInL0 = LayoutAInL0::template MakeLayout(mPartActual, kPartActual); - // Locate the current tile of matrix A on L1 - MatrixCoord l1AOffset{mPartIdx * L0TileShape::M, kPartIdx * L0TileShape::K}; - auto l1ATile = l1ATensor[layoutAInL1.GetOffset(l1AOffset)]; - - AscendC::WaitFlag(l0AEventList[l0AListId]); - if ((mPartIdx == 0) && (kPartIdx == 0)) { - AscendC::WaitFlag(l1AEventList[l1ListId]); - } - - // Load current tile from L1 to L0A - copyL1ToL0A(l0ATile, l1ATile, layoutAInL0, layoutAInL1); - - if ((mPartIdx == mPartLoop - 1) && (kPartIdx == kPartLoop - 1)) { - AscendC::SetFlag(l1AEventList[l1ListId]); - } - - for (int nPartIdx = 0; nPartIdx < nPartLoop; nPartIdx++) { - uint32_t nPartActual = (nPartIdx < nPartLoop - 1) ? - L0TileShape::N : - (nRound - nPartIdx * L0TileShape::N); - - // Locate the current tile on L0B - auto l0BTile = l0BTensorList[l0BListId]; - LayoutBInL0 layoutBInL0 = LayoutBInL0::template MakeLayout(kPartActual, nPartActual); - // Locate the current tile of matrix B on L1 - MatrixCoord l1BOffset{kPartIdx * L0TileShape::K, nPartIdx * L0TileShape::N}; - auto l1BTile = l1BTensor[layoutBInL1.GetOffset(l1BOffset)]; - - // Wait for mmad finished - AscendC::WaitFlag(l0BEventList[l0BListId]); - // If the current tile is the first one on the k&n axis, wait for loading matrix B from GM to L1 - if ((kPartIdx == 0) && (nPartIdx == 0)) { - AscendC::WaitFlag(l1BEventList[l1ListId]); - } - - // Load current tile from L1 to L0B - copyL1ToL0B(l0BTile, l1BTile, layoutBInL0, layoutBInL1); - - // If the current tile is the last one on the k&n axis, notify to load matrix B from GM to L1 - if ((kPartIdx == kPartLoop - 1) && (nPartIdx == nPartLoop - 1)) { - AscendC::SetFlag(l1BEventList[l1ListId]); - } - // Notify to do mmad - AscendC::SetFlag(EVENT_ID0); - - // Locate the current tile on L0C - MatrixCoord l0COffset{mPartIdx * L0TileShape::M, nPartIdx * L0TileShape::N}; - auto l0CTile = l0CTensor[layoutInL0C.GetOffset(l0COffset)]; - - // Compute the matrix multiplication on L0A and L0B and write the result to the accumulator - // Wait for loading L0B - AscendC::WaitFlag(EVENT_ID0); - - // If the current tile is the first tile on the k axis, the accumulator needs to be reset to 0 - bool initC = ((kLoopIdx == 0) && (kPartIdx == 0)); - // If the unit flag is enabled, the unit flag is set according to the calculation progress - uint8_t unitFlag = 0b00; - if constexpr (ENABLE_UNIT_FLAG) { - if ((kLoopIdx == kTileCount - 1) && (mPartIdx == mPartLoop - 1) && - (kPartIdx == kPartLoop - 1) && (nPartIdx == nPartLoop - 1)) { - unitFlag = 0b11; - } else { - unitFlag = 0b10; - } - } - // Perform calculation operations - tileMmad(l0CTile, l0ATile, l0BTile, mPartActual, nPartActual, kPartActual, initC, unitFlag); - - // Notify to move the next L0B tile - AscendC::SetFlag(l0BEventList[l0BListId]); - l0BListId = (l0BListId + 1 < STAGES) ? (l0BListId + 1) : 0; - } - AscendC::SetFlag(l0AEventList[l0AListId]); - l0AListId = (l0AListId + 1 < STAGES) ? (l0AListId + 1) : 0; - } - } - l1ListId = l1ListIdNext; - kActual = kActualNext; - } - - // copy block out - LayoutC layoutBlock = layoutC.GetTileLayout(actualShape.GetCoordMN()); - - if constexpr (!ENABLE_UNIT_FLAG) { - AscendC::SetFlag(EVENT_ID0); - AscendC::WaitFlag(EVENT_ID0); - copyL0CToGm(gmC, l0CTensor, layoutBlock, layoutInL0C); - AscendC::SetFlag(EVENT_ID0); - } else { - copyL0CToGm(gmC, l0CTensor, layoutBlock, layoutInL0C, 0b11); - } - } - -protected: - // Multi-stage tensors list - AscendC::LocalTensor l1ATensorList[STAGES]; - AscendC::LocalTensor l1BTensorList[STAGES]; - AscendC::LocalTensor l0ATensorList[STAGES]; - AscendC::LocalTensor l0BTensorList[STAGES]; - AscendC::LocalTensor l0CTensor; - - // Multi-stage event id list - int32_t l1AEventList[STAGES]; - int32_t l1BEventList[STAGES]; - int32_t l0AEventList[STAGES]; - int32_t l0BEventList[STAGES]; - - // The id of current stage - uint32_t l1ListId{0}; - uint32_t l0AListId{0}; - uint32_t l0BListId{0}; - - TileMmad tileMmad; - CopyGmToL1A copyGmToL1A; - CopyGmToL1B copyGmToL1B; - CopyL1ToL0A copyL1ToL0A; - CopyL1ToL0B copyL1ToL0B; - CopyL0CToGm copyL0CToGm; -}; - -} // namespace Catlass::Gemm::Block - -#endif // CATLASS_KDA_BLOCK_MMAD_PINGPONG_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_pingpong_bias.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_pingpong_bias.hpp deleted file mode 100644 index f03d3a75d8..0000000000 --- a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_pingpong_bias.hpp +++ /dev/null @@ -1,421 +0,0 @@ -/** - * Copyright (c) 2025 Huawei Technologies Co., Ltd. - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of - * CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ - -#ifndef CATLASS_KDA_BLOCK_MMAD_PINGPONG_BIAS_HPP -#define CATLASS_KDA_BLOCK_MMAD_PINGPONG_BIAS_HPP - -#include "catlass/kda_catlass.hpp" -#include "catlass/arch/kda_resource.hpp" -#include "catlass/kda_coord.hpp" -#include "catlass/kda_gemm_coord.hpp" -#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" -#include "catlass/gemm/kda_gemm_helper.hpp" -#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" -#include "catlass/gemm/tile/kda_tile_mmad.hpp" - -namespace Catlass::Gemm::Block { - -template < - bool ENABLE_UNIT_FLAG_, - class L1TileShape_, - class L0TileShape_, - class AType_, - class BType_, - class CType_, - class BiasType_, - class TileCopy_, - class TileMmad_> -struct BlockMmad< - MmadAtlasA2PingpongBias, - L1TileShape_, - L0TileShape_, - AType_, - BType_, - CType_, - BiasType_, - TileCopy_, - TileMmad_> { -public: - // Type Aliases - using DispatchPolicy = MmadAtlasA2PingpongBias; - using ArchTag = typename DispatchPolicy::ArchTag; - using L1TileShape = L1TileShape_; - using L0TileShape = L0TileShape_; - using ElementA = typename AType_::Element; - using LayoutA = typename AType_::Layout; - using ElementB = typename BType_::Element; - using LayoutB = typename BType_::Layout; - using ElementC = typename CType_::Element; - using LayoutC = typename CType_::Layout; - using ElementBias = typename BiasType_::Element; - using LayoutBias = typename BiasType_::Layout; - using TileMmad = TileMmad_; - using CopyGmToL1A = typename TileCopy_::CopyGmToL1A; - using CopyGmToL1B = typename TileCopy_::CopyGmToL1B; - using CopyGmToL1Bias = typename TileCopy_::CopyGmToL1Bias; - using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; - using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; - using CopyL0CToGm = typename TileCopy_::CopyL0CToGm; - using ElementAccumulator = - typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; - using CopyL1ToBT = typename TileCopy_::CopyL1ToBT; - using LayoutAInL1 = typename CopyL1ToL0A::LayoutSrc; - using LayoutBInL1 = typename CopyL1ToL0B::LayoutSrc; - using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst; - using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst; - using LayoutCInL0 = layout::zN; - - using L1AAlignHelper = Gemm::helper::L1AlignHelper; - using L1BAlignHelper = Gemm::helper::L1AlignHelper; - - static constexpr bool ENABLE_UNIT_FLAG = DispatchPolicy::ENABLE_UNIT_FLAG; - static constexpr uint32_t STAGES = DispatchPolicy::STAGES; - static constexpr uint32_t L1A_SIZE = L1TileShape::M * L1TileShape::K * sizeof(ElementA); - static constexpr uint32_t L1B_SIZE = L1TileShape::N * L1TileShape::K * sizeof(ElementB); - static constexpr uint32_t L1BIAS_SIZE = L1TileShape::N * sizeof(ElementBias); - static constexpr uint32_t L0A_SIZE = ArchTag::L0A_SIZE; - static constexpr uint32_t L0B_SIZE = ArchTag::L0B_SIZE; - static constexpr uint32_t L0C_SIZE = ArchTag::L0C_SIZE; - static constexpr uint32_t BT_SIZE = ArchTag::BIAS_SIZE; - static constexpr uint32_t L0A_PINGPONG_BUF_SIZE = L0A_SIZE / STAGES; - static constexpr uint32_t L0B_PINGPONG_BUF_SIZE = L0B_SIZE / STAGES; - static constexpr uint32_t BIAS_BUF_SIZE = L0TileShape::N * sizeof(ElementAccumulator); - - // Check LayoutC - static_assert(std::is_same_v, "LayoutC only support RowMajor yet!"); - - // Check L1TileShape - static_assert((L1A_SIZE * STAGES + L1B_SIZE * STAGES + L1BIAS_SIZE) <= ArchTag::L1_SIZE, - "L1TileShape exceeding the L1 space!"); - - // Check L0TileShape - static constexpr uint32_t L0A_TILE_SIZE = L0TileShape::M * L0TileShape::K * sizeof(ElementA); - static constexpr uint32_t L0B_TILE_SIZE = L0TileShape::K * L0TileShape::N * sizeof(ElementB); - static constexpr uint32_t L0C_TILE_SIZE = L0TileShape::M * L0TileShape::N * sizeof(ElementAccumulator); - static_assert((L0A_TILE_SIZE * STAGES) <= L0A_SIZE, "L0TileShape exceeding the L0A space!"); - static_assert((L0B_TILE_SIZE * STAGES) <= L0B_SIZE, "L0TileShape exceeding the L0B space!"); - static_assert(L0C_TILE_SIZE <= L0C_SIZE, "L0TileShape exceeding the L0C space!"); - static_assert(BIAS_BUF_SIZE <= BT_SIZE, "BIAS_BUF_SIZE exceeding the BT space! Reduce L0TileShape::N"); - - static_assert(L1TileShape::M == L0TileShape::M && L1TileShape::N == L0TileShape::N, - "The situation where the basic blocks of L1 and L0 differ on the m and n axes is not supported yet"); - static_assert(L0TileShape::K <= L1TileShape::K, "L0TileShape::K cannot exceed L1TileShape::K"); - static_assert(!std::is_same_v, "Bias does not support bfloat16 yet"); - static_assert(Gemm::helper::TileShapeAlignChecker::_ALIGN == 256, - "Tile shape must be 32B aligned."); - - /// Construct - CATLASS_DEVICE - BlockMmad(Arch::Resource &resource, uint32_t l1BufAddrStart = 0) - { - uint32_t l1AOffset = l1BufAddrStart; - uint32_t l1BOffset = l1BufAddrStart + L1A_SIZE * STAGES; - uint32_t l1BiasOffset = l1BufAddrStart + L1A_SIZE * STAGES + L1B_SIZE * STAGES; - // Init buffers - for (uint32_t i = 0; i < STAGES; i++) { - // Assign L1/L0A/L0B space for each stages - l1ATensorList[i] = resource.l1Buf.template GetBufferByByte(l1AOffset + L1A_SIZE * i); - l1BTensorList[i] = resource.l1Buf.template GetBufferByByte(l1BOffset + L1B_SIZE * i); - l0ATensorList[i] = resource.l0ABuf.template GetBufferByByte(L0A_PINGPONG_BUF_SIZE * i); - l0BTensorList[i] = resource.l0BBuf.template GetBufferByByte(L0B_PINGPONG_BUF_SIZE * i); - - // Assign event ID for each stages - l1AEventList[i] = i; - l1BEventList[i] = i + STAGES; - l0AEventList[i] = i; - l0BEventList[i] = i + STAGES; - - // The event id that needs to be set before the loop - AscendC::SetFlag(l1AEventList[i]); - AscendC::SetFlag(l1BEventList[i]); - AscendC::SetFlag(l0AEventList[i]); - AscendC::SetFlag(l0BEventList[i]); - } - EVENT_ID_BIAS = 2 * STAGES; - l0CTensor = resource.l0CBuf.template GetBufferByByte(0); - l1BiasTensor = resource.l1Buf.template GetBufferByByte(l1BiasOffset); - l0BiasTensor = resource.btBuf.template GetBufferByByte(0); - AscendC::SetFlag(EVENT_ID0); - // Initialize bias independent event - AscendC::SetFlag(EVENT_ID_BIAS); - AscendC::SetFlag(EVENT_ID_BIAS); - } - - /// Destructor - CATLASS_DEVICE - ~BlockMmad() - { - for (uint32_t i = 0; i < STAGES; i++) { - AscendC::WaitFlag(l1AEventList[i]); - AscendC::WaitFlag(l1BEventList[i]); - AscendC::WaitFlag(l0AEventList[i]); - AscendC::WaitFlag(l0BEventList[i]); - } - AscendC::WaitFlag(EVENT_ID0); - // Wait for bias independent event - AscendC::WaitFlag(EVENT_ID_BIAS); - AscendC::WaitFlag(EVENT_ID_BIAS); - } - - /// Perform a block-scoped matrix multiply-accumulate - CATLASS_DEVICE - void operator()( - AscendC::GlobalTensor const &gmA, LayoutA const &layoutA, - AscendC::GlobalTensor const &gmB, LayoutB const &layoutB, - AscendC::GlobalTensor const &gmC, LayoutC const &layoutC, - AscendC::GlobalTensor const &gmBias, - GemmCoord const &actualShape) - { - uint32_t mRound = RoundUp(actualShape.m()); - uint32_t nRound = RoundUp(actualShape.n()); - - auto layoutAInL1 = LayoutAInL1::template MakeLayout(L1TileShape::M, L1TileShape::K); - auto layoutBInL1 = LayoutBInL1::template MakeLayout(L1TileShape::K, L1TileShape::N); - auto layoutBiasInL1 = layout::VectorLayout(L1TileShape::N); - auto layoutBiasInL0 = layout::VectorLayout(L0TileShape::N); - auto layoutInL0C = LayoutCInL0::MakeLayoutInL0C(MakeCoord(mRound, nRound)); - - uint32_t kActual = min(actualShape.k(), L1TileShape::K); - - // load first matrix A tile from GM to L1 - AscendC::WaitFlag(l1AEventList[l1ListId]); - auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShape.m(), kActual)); - copyGmToL1A(l1ATensorList[l1ListId], gmA, layoutAInL1, layoutTileA); - AscendC::SetFlag(l1AEventList[l1ListId]); - - // load first matrix B tile from GM to L1 - AscendC::WaitFlag(l1BEventList[l1ListId]); - auto layoutTileB = layoutB.GetTileLayout(MakeCoord(kActual, actualShape.n())); - copyGmToL1B(l1BTensorList[l1ListId], gmB, layoutBInL1, layoutTileB); - AscendC::SetFlag(l1BEventList[l1ListId]); - - // load bias from GM to L1 independently - AscendC::WaitFlag(EVENT_ID_BIAS); - auto layoutTileBias = layout::VectorLayout(actualShape.n()); - copyGmToL1Bias(l1BiasTensor, gmBias, layoutBiasInL1, layoutTileBias); - AscendC::SetFlag(EVENT_ID_BIAS); - - if constexpr (!ENABLE_UNIT_FLAG) { - AscendC::WaitFlag(EVENT_ID0); - } - - uint32_t mPartLoop = CeilDiv(mRound); - uint32_t nPartLoop = CeilDiv(nRound); - - // main loop - uint32_t kTileCount = CeilDiv(actualShape.k()); - for (uint32_t kLoopIdx = 0; kLoopIdx < kTileCount; kLoopIdx++) { - uint32_t l1ListIdNext = (l1ListId + 1 < STAGES) ? (l1ListId + 1) : 0; - uint32_t kActualNext{0}; - // preload next tile from GM to L1 - if (kLoopIdx < kTileCount - 1) { - uint32_t kLoopIdxNext = kLoopIdx + 1; - kActualNext = (kLoopIdxNext < kTileCount - 1) ? - L1TileShape::K : - (actualShape.k() - kLoopIdxNext * L1TileShape::K); - - // Get L1 tensor for next stage - auto l1ATensor = l1ATensorList[l1ListIdNext]; - auto l1BTensor = l1BTensorList[l1ListIdNext]; - // Get GM tile for next stage - MatrixCoord gmTileAOffset{0, kLoopIdxNext * L1TileShape::K}; - MatrixCoord gmTileBOffset{kLoopIdxNext * L1TileShape::K, 0}; - auto gmTileA = gmA[layoutA.GetOffset(gmTileAOffset)]; - auto gmTileB = gmB[layoutB.GetOffset(gmTileBOffset)]; - - // load next matrix A tile from GM to L1 - AscendC::WaitFlag(l1AEventList[l1ListIdNext]); - layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShape.m(), kActualNext)); - copyGmToL1A(l1ATensor, gmTileA, layoutAInL1, layoutTileA); - AscendC::SetFlag(l1AEventList[l1ListIdNext]); - - // load next matrix B tile from GM to L1 - AscendC::WaitFlag(l1BEventList[l1ListIdNext]); - layoutTileB = layoutB.GetTileLayout(MakeCoord(kActualNext, actualShape.n())); - copyGmToL1B(l1BTensor, gmTileB, layoutBInL1, layoutTileB); - AscendC::SetFlag(l1BEventList[l1ListIdNext]); - } - - // Get L1 tensor for current stage - auto l1ATensor = l1ATensorList[l1ListId]; - auto l1BTensor = l1BTensorList[l1ListId]; - - // Get the loop nums on L0 - uint32_t kPartLoop = CeilDiv(kActual); - - for (int mPartIdx = 0; mPartIdx < mPartLoop; mPartIdx++) { - uint32_t mPartActual = (mPartIdx < mPartLoop - 1) ? - L0TileShape::M : - (mRound - mPartIdx * L0TileShape::M); - - for (int kPartIdx = 0; kPartIdx < kPartLoop; kPartIdx++) { - uint32_t kPartActual = (kPartIdx < kPartLoop - 1) ? - L0TileShape::K : - (kActual - kPartIdx * L0TileShape::K); - - // Locate the current tile on L0A - auto l0ATile = l0ATensorList[l0AListId]; - LayoutAInL0 layoutAInL0 = LayoutAInL0::template MakeLayout(mPartActual, kPartActual); - // Locate the current tile of matrix A on L1 - MatrixCoord l1AOffset{mPartIdx * L0TileShape::M, kPartIdx * L0TileShape::K}; - auto l1ATile = l1ATensor[layoutAInL1.GetOffset(l1AOffset)]; - - AscendC::WaitFlag(l0AEventList[l0AListId]); - if ((mPartIdx == 0) && (kPartIdx == 0)) { - AscendC::WaitFlag(l1AEventList[l1ListId]); - } - - // Load current tile from L1 to L0A - copyL1ToL0A(l0ATile, l1ATile, layoutAInL0, layoutAInL1); - - if ((mPartIdx == mPartLoop - 1) && (kPartIdx == kPartLoop - 1)) { - AscendC::SetFlag(l1AEventList[l1ListId]); - } - - for (int nPartIdx = 0; nPartIdx < nPartLoop; nPartIdx++) { - uint32_t nPartActual = (nPartIdx < nPartLoop - 1) ? - L0TileShape::N : - (nRound - nPartIdx * L0TileShape::N); - - // Locate the current tile on L0B - auto l0BTile = l0BTensorList[l0BListId]; - LayoutBInL0 layoutBInL0 = LayoutBInL0::template MakeLayout(kPartActual, nPartActual); - // Locate the current tile of matrix B on L1 - MatrixCoord l1BOffset{kPartIdx * L0TileShape::K, nPartIdx * L0TileShape::N}; - auto l1BTile = l1BTensor[layoutBInL1.GetOffset(l1BOffset)]; - - // Wait for mmad finished - AscendC::WaitFlag(l0BEventList[l0BListId]); - // If the current tile is the first one on the k&n axis, wait for loading matrix B from GM to L1 - if ((kPartIdx == 0) && (nPartIdx == 0)) { - AscendC::WaitFlag(l1BEventList[l1ListId]); - } - - // Load current tile from L1 to L0B - copyL1ToL0B(l0BTile, l1BTile, layoutBInL0, layoutBInL1); - - // If the current tile is the last one on the k&n axis, notify to load matrix B from GM to L1 - if ((kPartIdx == kPartLoop - 1) && (nPartIdx == nPartLoop - 1)) { - AscendC::SetFlag(l1BEventList[l1ListId]); - } - - // If the current tile is the first tile on the k axis, the accumulator needs to be reset to 0 - bool initC = ((kLoopIdx == 0) && (kPartIdx == 0)); - - // Load bias to l0 biastable only when initC is true (kLoopIdx==0 && kPartIdx==0) - // Load L0_TILE_N length data in nPartIdx loop, starting from nPartIdx * L0_TILE_N - if (initC) { - if (nPartIdx == 0) { - AscendC::WaitFlag(EVENT_ID_BIAS); - } - AscendC::WaitFlag(EVENT_ID_BIAS); - // Create layout for actual nPartActual size - auto layoutBiasInL0Actual = layout::VectorLayout(nPartActual); - // Get L1 bias tile starting from nPartIdx * L0_TILE_N - // For VectorLayout, directly use index offset - auto l1BiasTile = l1BiasTensor[nPartIdx * L0TileShape::N]; - // Create layout for L1 source with offset - auto layoutBiasInL1Offset = layout::VectorLayout(L1TileShape::N); - copyL1ToBT(l0BiasTensor, l1BiasTile, layoutBiasInL0Actual, layoutBiasInL1Offset); - if (nPartIdx == nPartLoop - 1) { - AscendC::SetFlag(EVENT_ID_BIAS); - } - } - // Notify to do mmad - AscendC::SetFlag(EVENT_ID0); - - // Locate the current tile on L0C - MatrixCoord l0COffset{mPartIdx * L0TileShape::M, nPartIdx * L0TileShape::N}; - auto l0CTile = l0CTensor[layoutInL0C.GetOffset(l0COffset)]; - - // Compute the matrix multiplication on L0A and L0B and write the result to the accumulator - // Wait for loading L0B - AscendC::WaitFlag(EVENT_ID0); - // If the unit flag is enabled, the unit flag is set according to the calculation progress - uint8_t unitFlag = 0b00; - if constexpr (ENABLE_UNIT_FLAG) { - if ((kLoopIdx == kTileCount - 1) && (mPartIdx == mPartLoop - 1) && - (kPartIdx == kPartLoop - 1) && (nPartIdx == nPartLoop - 1)) { - unitFlag = 0b11; - } else { - unitFlag = 0b10; - } - } - // Perform calculation operations - if (initC) { - tileMmad(l0CTile, l0ATile, l0BTile, l0BiasTensor, mPartActual, nPartActual, kPartActual, - initC, unitFlag); - AscendC::SetFlag(EVENT_ID_BIAS); - } else { - tileMmad(l0CTile, l0ATile, l0BTile, mPartActual, nPartActual, kPartActual, initC, unitFlag); - } - - // Notify to move the next L0B tile - AscendC::SetFlag(l0BEventList[l0BListId]); - l0BListId = (l0BListId + 1 < STAGES) ? (l0BListId + 1) : 0; - } - AscendC::SetFlag(l0AEventList[l0AListId]); - l0AListId = (l0AListId + 1 < STAGES) ? (l0AListId + 1) : 0; - } - } - l1ListId = l1ListIdNext; - kActual = kActualNext; - } - - // copy block out - LayoutC layoutBlock = layoutC.GetTileLayout(actualShape.GetCoordMN()); - - if constexpr (!ENABLE_UNIT_FLAG) { - AscendC::SetFlag(EVENT_ID0); - AscendC::WaitFlag(EVENT_ID0); - copyL0CToGm(gmC, l0CTensor, layoutBlock, layoutInL0C); - AscendC::SetFlag(EVENT_ID0); - } else { - copyL0CToGm(gmC, l0CTensor, layoutBlock, layoutInL0C, 0b11); - } - } - -protected: - // Multi-stage tensors list - AscendC::LocalTensor l1ATensorList[STAGES]; - AscendC::LocalTensor l1BTensorList[STAGES]; - AscendC::LocalTensor l1BiasTensor; - AscendC::LocalTensor l0ATensorList[STAGES]; - AscendC::LocalTensor l0BTensorList[STAGES]; - AscendC::LocalTensor l0CTensor; - AscendC::LocalTensor l0BiasTensor; - - // Multi-stage event id list - int32_t l1AEventList[STAGES]; - int32_t l1BEventList[STAGES]; - int32_t l0AEventList[STAGES]; - int32_t l0BEventList[STAGES]; - - // Event ID for Bias - int32_t EVENT_ID_BIAS; - - // The id of current stage - uint32_t l1ListId{0}; - uint32_t l0AListId{0}; - uint32_t l0BListId{0}; - - TileMmad tileMmad; - CopyGmToL1A copyGmToL1A; - CopyGmToL1B copyGmToL1B; - CopyGmToL1Bias copyGmToL1Bias; - CopyL1ToL0A copyL1ToL0A; - CopyL1ToL0B copyL1ToL0B; - CopyL1ToBT copyL1ToBT; - CopyL0CToGm copyL0CToGm; -}; - -} // namespace Catlass::Gemm::Block - -#endif // CATLASS_KDA_BLOCK_MMAD_PINGPONG_BIAS_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_preload.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_preload.hpp deleted file mode 100644 index b720d161a6..0000000000 --- a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_preload.hpp +++ /dev/null @@ -1,402 +0,0 @@ -/** - * Copyright (c) 2025 Huawei Technologies Co., Ltd. - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of - * CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ - -#ifndef CATLASS_KDA_BLOCK_MMAD_PRELOAD_HPP -#define CATLASS_KDA_BLOCK_MMAD_PRELOAD_HPP - -#include "catlass/kda_catlass.hpp" -#include "catlass/arch/kda_resource.hpp" -#include "catlass/kda_coord.hpp" -#include "catlass/kda_gemm_coord.hpp" -#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" -#include "catlass/gemm/kda_gemm_helper.hpp" -#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" -#include "catlass/gemm/tile/kda_tile_mmad.hpp" - -namespace Catlass::Gemm::Block { - -template < - bool ENABLE_UNIT_FLAG_, - bool ENABLE_SHUFFLE_K_, - class L1TileShape_, - class L0TileShape_, - class AType_, - class BType_, - class CType_, - class BiasType_, - class TileCopy_, - class TileMmad_> -struct BlockMmad< - MmadAtlasA2Preload, - L1TileShape_, - L0TileShape_, - AType_, - BType_, - CType_, - BiasType_, - TileCopy_, - TileMmad_> { -public: - // Type Aliases - using DispatchPolicy = MmadAtlasA2Preload; - using ArchTag = typename DispatchPolicy::ArchTag; - using L1TileShape = L1TileShape_; - using L0TileShape = L0TileShape_; - using ElementA = typename AType_::Element; - using LayoutA = typename AType_::Layout; - using ElementB = typename BType_::Element; - using LayoutB = typename BType_::Layout; - using ElementC = typename CType_::Element; - using LayoutC = typename CType_::Layout; - using TileMmad = TileMmad_; - using CopyGmToL1A = typename TileCopy_::CopyGmToL1A; - using CopyGmToL1B = typename TileCopy_::CopyGmToL1B; - using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; - using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; - using CopyL0CToGm = typename TileCopy_::CopyL0CToGm; - using ElementAccumulator = - typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; - using LayoutAInL1 = typename CopyL1ToL0A::LayoutSrc; - using LayoutBInL1 = typename CopyL1ToL0B::LayoutSrc; - using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst; - using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst; - using LayoutCInL0 = layout::zN; - - using L1AAlignHelper = Gemm::helper::L1AlignHelper; - using L1BAlignHelper = Gemm::helper::L1AlignHelper; - - static constexpr bool ENABLE_UNIT_FLAG = DispatchPolicy::ENABLE_UNIT_FLAG; - static constexpr bool ENABLE_SHUFFLE_K = DispatchPolicy::ENABLE_SHUFFLE_K; - static constexpr uint32_t STAGES = DispatchPolicy::STAGES; - static constexpr uint32_t L1A_SIZE = L1TileShape::M * L1TileShape::K * sizeof(ElementA); - static constexpr uint32_t L1B_SIZE = L1TileShape::N * L1TileShape::K * sizeof(ElementB); - static constexpr uint32_t L0A_SIZE = ArchTag::L0A_SIZE; - static constexpr uint32_t L0B_SIZE = ArchTag::L0B_SIZE; - static constexpr uint32_t L0C_SIZE = ArchTag::L0C_SIZE; - static constexpr uint32_t L0A_PINGPONG_BUF_SIZE = L0A_SIZE / STAGES; - static constexpr uint32_t L0B_PINGPONG_BUF_SIZE = L0B_SIZE / STAGES; - - // Check LayoutC - static_assert(std::is_same_v, "LayoutC only support RowMajor yet!"); - - // Check L1TileShape - static_assert((L1A_SIZE * STAGES + L1B_SIZE * STAGES) <= ArchTag::L1_SIZE, "L1TileShape exceeding the L1 space!"); - - // Check L0TileShape - static constexpr uint32_t L0A_TILE_SIZE = L0TileShape::M * L0TileShape::K * sizeof(ElementA); - static constexpr uint32_t L0B_TILE_SIZE = L0TileShape::K * L0TileShape::N * sizeof(ElementB); - static constexpr uint32_t L0C_TILE_SIZE = L0TileShape::M * L0TileShape::N * sizeof(ElementAccumulator); - static_assert((L0A_TILE_SIZE * STAGES) <= L0A_SIZE, "L0TileShape exceeding the L0A space!"); - static_assert((L0B_TILE_SIZE * STAGES) <= L0B_SIZE, "L0TileShape exceeding the L0B space!"); - static_assert(L0C_TILE_SIZE <= L0C_SIZE, "L0TileShape exceeding the L0C space!"); - - static_assert(L1TileShape::M == L0TileShape::M && L1TileShape::N == L0TileShape::N, - "The situation where the basic blocks of L1 and L0 differ on the m and n axes is not supported yet"); - static_assert(L0TileShape::K <= L1TileShape::K, "L0TileShape::K cannot exceed L1TileShape::K"); - // 32B (256b) aligned - static_assert(Gemm::helper::TileShapeAlignChecker::_ALIGN == 256, - "Tile shape must be 32B aligned."); - - /// Construct - CATLASS_DEVICE - BlockMmad(Arch::Resource &resource, uint32_t l1BufAddrStart = 0) - { - uint32_t l1AOffset = l1BufAddrStart; - uint32_t l1BOffset = l1BufAddrStart + L1A_SIZE * STAGES; - // Init buffers - for (uint32_t i = 0; i < STAGES; i++) { - l1ATensorList[i] = resource.l1Buf.template GetBufferByByte(l1AOffset + L1A_SIZE * i); - l1BTensorList[i] = resource.l1Buf.template GetBufferByByte(l1BOffset + L1B_SIZE * i); - l0ATensorList[i] = resource.l0ABuf.template GetBufferByByte(L0A_PINGPONG_BUF_SIZE * i); - l0BTensorList[i] = resource.l0BBuf.template GetBufferByByte(L0B_PINGPONG_BUF_SIZE * i); - - l1AEventList[i] = i; - l1BEventList[i] = i + STAGES; - l0AEventList[i] = i; - l0BEventList[i] = i + STAGES; - AscendC::SetFlag(l1AEventList[i]); - AscendC::SetFlag(l1BEventList[i]); - AscendC::SetFlag(l0AEventList[i]); - AscendC::SetFlag(l0BEventList[i]); - } - l0CTensor = resource.l0CBuf.template GetBufferByByte(0); - AscendC::SetFlag(EVENT_ID0); - } - - /// Destructor - CATLASS_DEVICE - ~BlockMmad() - { - for (uint32_t i = 0; i < STAGES; i++) { - AscendC::WaitFlag(l1AEventList[i]); - AscendC::WaitFlag(l1BEventList[i]); - AscendC::WaitFlag(l0AEventList[i]); - AscendC::WaitFlag(l0BEventList[i]); - } - AscendC::WaitFlag(EVENT_ID0); - } - - /// Perform a block-scoped matrix multiply-accumulate - CATLASS_DEVICE - void operator()( - AscendC::GlobalTensor const &gmBlockA, LayoutA const &layoutA, - AscendC::GlobalTensor const &gmBlockB, LayoutB const &layoutB, - AscendC::GlobalTensor const &gmBlockC, LayoutC const &layoutC, - AscendC::GlobalTensor const &gmNextBlockA, - AscendC::GlobalTensor const &gmNextBlockB, - GemmCoord const &actualShape, GemmCoord const &actualShapeNext, - bool isFirstBlock, bool hasNextBlock) - { - uint32_t mRound = RoundUp(actualShape.m()); - uint32_t nRound = RoundUp(actualShape.n()); - - auto layoutAInL1 = LayoutAInL1::template MakeLayout(L1TileShape::M, L1TileShape::K); - auto layoutBInL1 = LayoutBInL1::template MakeLayout(L1TileShape::K, L1TileShape::N); - auto layoutInL0C = LayoutCInL0::MakeLayoutInL0C(MakeCoord(mRound, nRound)); - - uint32_t kTileCount = CeilDiv(actualShape.k()); - uint32_t kTileCountNext = CeilDiv(actualShapeNext.k()); - - if constexpr (!ENABLE_UNIT_FLAG) { - AscendC::WaitFlag(EVENT_ID0); - } - uint32_t startTileIdx = 0; - if constexpr (ENABLE_SHUFFLE_K) { - startTileIdx = AscendC::GetBlockIdx(); - } - uint32_t firstTileIdx = startTileIdx % kTileCount; - uint32_t lastTileIdx = (startTileIdx + kTileCount - 1) % kTileCount; - uint32_t kActual = - (firstTileIdx < kTileCount - 1) ? L1TileShape::K : (actualShape.k() - firstTileIdx * L1TileShape::K); - uint32_t firstTileIdxNext = startTileIdx % kTileCountNext; - - uint32_t mPartLoop = CeilDiv(mRound); - uint32_t nPartLoop = CeilDiv(nRound); - - // k loop - for (uint32_t kLoopIdx = 0; kLoopIdx < kTileCount; kLoopIdx++) { - uint32_t shuffleKIdx = (startTileIdx + kLoopIdx) % kTileCount; - // Load first matrix A tile in total kernel loop from GM to L1 - if (shuffleKIdx == firstTileIdx && isFirstBlock) { - MatrixCoord gmTileAOffset{0, shuffleKIdx * L1TileShape::K}; - MatrixCoord gmTileBOffset{shuffleKIdx * L1TileShape::K, 0}; - auto gmTileA = gmBlockA[layoutA.GetOffset(gmTileAOffset)]; - auto gmTileB = gmBlockB[layoutB.GetOffset(gmTileBOffset)]; - // Load first matrix A tile from GM to L1 - AscendC::WaitFlag(l1AEventList[l1ListId]); - auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShape.m(), kActual)); - copyGmToL1A(l1ATensorList[l1ListId], gmTileA, layoutAInL1, layoutTileA); - AscendC::SetFlag(l1AEventList[l1ListId]); - - // Load first matrix B tile from GM to L1 - AscendC::WaitFlag(l1BEventList[l1ListId]); - auto layoutTileB = layoutB.GetTileLayout(MakeCoord(kActual, actualShape.n())); - copyGmToL1B(l1BTensorList[l1ListId], gmTileB, layoutBInL1, layoutTileB); - AscendC::SetFlag(l1BEventList[l1ListId]); - } - - uint32_t l1ListIdNext = (l1ListId + 1 < STAGES) ? (l1ListId + 1) : 0; - uint32_t kActualNext{0}; - - // preload next tile from GM to L1 - if (shuffleKIdx != lastTileIdx) { - uint32_t shuffleKIdxNext = (startTileIdx + kLoopIdx + 1) % kTileCount; - // Get L1 tensor for next stage - auto l1ATensor = l1ATensorList[l1ListIdNext]; - auto l1BTensor = l1BTensorList[l1ListIdNext]; - // Get GM tensor for next stage - kActualNext = (shuffleKIdxNext < kTileCount - 1) ? - L1TileShape::K : - (actualShape.k() - shuffleKIdxNext * L1TileShape::K); - MatrixCoord gmTileAOffset{0, shuffleKIdxNext * L1TileShape::K}; - MatrixCoord gmTileBOffset{shuffleKIdxNext * L1TileShape::K, 0}; - auto gmTileA = gmBlockA[layoutA.GetOffset(gmTileAOffset)]; - auto gmTileB = gmBlockB[layoutB.GetOffset(gmTileBOffset)]; - - // load next matrix A tile from GM to L1 - AscendC::WaitFlag(l1AEventList[l1ListIdNext]); - auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShape.m(), kActualNext)); - copyGmToL1A(l1ATensor, gmTileA, layoutAInL1, layoutTileA); - AscendC::SetFlag(l1AEventList[l1ListIdNext]); - - // load next matrix B tile from GM to L1 - AscendC::WaitFlag(l1BEventList[l1ListIdNext]); - auto layoutTileB = layoutB.GetTileLayout(MakeCoord(kActualNext, actualShape.n())); - copyGmToL1B(l1BTensor, gmTileB, layoutBInL1, layoutTileB); - AscendC::SetFlag(l1BEventList[l1ListIdNext]); - } - if (shuffleKIdx == lastTileIdx && hasNextBlock) { - // Get L1 tensor for next stage - auto l1ATensor = l1ATensorList[l1ListIdNext]; - auto l1BTensor = l1BTensorList[l1ListIdNext]; - // Get GM tensor for next stage - kActualNext = (firstTileIdxNext < kTileCountNext - 1) ? - L1TileShape::K : - (actualShapeNext.k() - firstTileIdxNext * L1TileShape::K); - MatrixCoord gmTileAOffset{0, firstTileIdxNext * L1TileShape::K}; - MatrixCoord gmTileBOffset{firstTileIdxNext * L1TileShape::K, 0}; - auto gmTileA = gmNextBlockA[layoutA.GetOffset(gmTileAOffset)]; - auto gmTileB = gmNextBlockB[layoutB.GetOffset(gmTileBOffset)]; - // load next matrix A tile from GM to L1 - AscendC::WaitFlag(l1AEventList[l1ListIdNext]); - auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShapeNext.m(), kActualNext)); - copyGmToL1A(l1ATensor, gmTileA, layoutAInL1, layoutTileA); - AscendC::SetFlag(l1AEventList[l1ListIdNext]); - - // load next matrix B tile from GM to L1 - AscendC::WaitFlag(l1BEventList[l1ListIdNext]); - auto layoutTileB = layoutB.GetTileLayout(MakeCoord(kActualNext, actualShapeNext.n())); - copyGmToL1B(l1BTensor, gmTileB, layoutBInL1, layoutTileB); - AscendC::SetFlag(l1BEventList[l1ListIdNext]); - } - - // Get L1 tensor for current stage - auto l1ATensor = l1ATensorList[l1ListId]; - auto l1BTensor = l1BTensorList[l1ListId]; - - // Get the loop nums on L0 - uint32_t kPartLoop = CeilDiv(kActual); - - uint32_t l0ABufId = 0; - uint32_t l0BBufId = 0; - - for (int mPartIdx = 0; mPartIdx < mPartLoop; mPartIdx++) { - uint32_t mPartActual = (mPartIdx < mPartLoop - 1) ? - L0TileShape::M : - (mRound - mPartIdx * L0TileShape::M); - - for (int kPartIdx = 0; kPartIdx < kPartLoop; kPartIdx++) { - uint32_t kPartActual = (kPartIdx < kPartLoop - 1) ? - L0TileShape::K : - (kActual - kPartIdx * L0TileShape::K); - - // Locate the current tile on L0A - auto l0ATile = l0ATensorList[l0ABufId]; - LayoutAInL0 layoutAInL0 = LayoutAInL0::template MakeLayout(mPartActual, kPartActual); - // Locate the current tile of matrix A on L1 - MatrixCoord l1AOffset{mPartIdx * L0TileShape::M, kPartIdx * L0TileShape::K}; - auto l1ATile = l1ATensor[layoutAInL1.GetOffset(l1AOffset)]; - - AscendC::WaitFlag(l0AEventList[l0ABufId]); - if ((mPartIdx == 0) && (kPartIdx == 0)) { - AscendC::WaitFlag(l1AEventList[l1ListId]); - } - - copyL1ToL0A(l0ATile, l1ATile, layoutAInL0, layoutAInL1); - - if ((mPartIdx == mPartLoop - 1) && (kPartIdx == kPartLoop - 1)) { - AscendC::SetFlag(l1AEventList[l1ListId]); - } - - for (int nPartIdx = 0; nPartIdx < nPartLoop; nPartIdx++) { - uint32_t nPartActual = (nPartIdx < nPartLoop - 1) ? - L0TileShape::N : - (nRound - nPartIdx * L0TileShape::N); - - // Locate the current tile on L0B - auto l0BTile = l0BTensorList[l0BBufId]; - LayoutBInL0 layoutBInL0 = LayoutBInL0::template MakeLayout(kPartActual, nPartActual); - // Locate the current tile of matrix B on L1 - MatrixCoord l1BOffset{kPartIdx * L0TileShape::K, nPartIdx * L0TileShape::N}; - auto l1BTile = l1BTensor[layoutBInL1.GetOffset(l1BOffset)]; - - // Wait for mmad finished - AscendC::WaitFlag(l0BEventList[l0BBufId]); - // If the current tile is the first one on the k&n axis, wait for loading matrix B from GM to L1 - if ((kPartIdx == 0) && (nPartIdx == 0)) { - AscendC::WaitFlag(l1BEventList[l1ListId]); - } - - // Load current tile from L1 to L0B - copyL1ToL0B(l0BTile, l1BTile, layoutBInL0, layoutBInL1); - - // If the current tile is the last one on the k&n axis, notify to load matrix B from GM to L1 - if ((kPartIdx == kPartLoop - 1) && (nPartIdx == nPartLoop - 1)) { - AscendC::SetFlag(l1BEventList[l1ListId]); - } - // Notify to do mmad - AscendC::SetFlag(EVENT_ID0); - - // Locate the current tile on L0C - MatrixCoord l0COffset{mPartIdx * L0TileShape::M, nPartIdx * L0TileShape::N}; - auto l0CTile = l0CTensor[layoutInL0C.GetOffset(l0COffset)]; - - // Compute the matrix multiplication on L0A and L0B and write the result to the accumulator - // Wait for loading L0B - AscendC::WaitFlag(EVENT_ID0); - - // If the current tile is the first tile on the k axis, the accumulator needs to be reset to 0 - bool initC = ((kLoopIdx == 0) && (kPartIdx == 0)); - // If the unit flag is enabled, the unit flag is set according to the calculation progress - uint8_t unitFlag = 0b00; - if constexpr (ENABLE_UNIT_FLAG) { - if ((kLoopIdx == kTileCount - 1) && (mPartIdx == mPartLoop - 1) && - (kPartIdx == kPartLoop - 1) && (nPartIdx == nPartLoop - 1)) { - unitFlag = 0b11; - } else { - unitFlag = 0b10; - } - } - // Perform calculation operations - tileMmad(l0CTile, l0ATile, l0BTile, mPartActual, nPartActual, kPartActual, initC, unitFlag); - - // Notify to move the next L0B tile - AscendC::SetFlag(l0BEventList[l0BBufId]); - - l0BBufId = (l0BBufId + 1 < STAGES) ? (l0BBufId + 1) : 0; - } - AscendC::SetFlag(l0AEventList[l0ABufId]); - l0ABufId = (l0ABufId + 1 < STAGES) ? (l0ABufId + 1) : 0; - } - } - - l1ListId = l1ListIdNext; - kActual = kActualNext; - } - - // copy block out - LayoutC layoutBlock = layoutC.GetTileLayout(actualShape.GetCoordMN()); - - if constexpr (!ENABLE_UNIT_FLAG) { - AscendC::SetFlag(EVENT_ID0); - AscendC::WaitFlag(EVENT_ID0); - copyL0CToGm(gmBlockC, l0CTensor, layoutBlock, layoutInL0C); - AscendC::SetFlag(EVENT_ID0); - } else { - copyL0CToGm(gmBlockC, l0CTensor, layoutBlock, layoutInL0C, 0b11); - } - } - -protected: - /// Data members - AscendC::LocalTensor l1ATensorList[STAGES]; - AscendC::LocalTensor l1BTensorList[STAGES]; - AscendC::LocalTensor l0ATensorList[STAGES]; - AscendC::LocalTensor l0BTensorList[STAGES]; - AscendC::LocalTensor l0CTensor; - - int32_t l1AEventList[STAGES]; - int32_t l1BEventList[STAGES]; - int32_t l0AEventList[STAGES]; - int32_t l0BEventList[STAGES]; - - uint32_t l1ListId{0}; - - TileMmad tileMmad; - CopyGmToL1A copyGmToL1A; - CopyGmToL1B copyGmToL1B; - CopyL1ToL0A copyL1ToL0A; - CopyL1ToL0B copyL1ToL0B; - CopyL0CToGm copyL0CToGm; -}; - -} // namespace Catlass::Gemm::Block - -#endif // CATLASS_KDA_BLOCK_MMAD_PRELOAD_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_preload_async.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_preload_async.hpp deleted file mode 100644 index 490cb3b11b..0000000000 --- a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_preload_async.hpp +++ /dev/null @@ -1,448 +0,0 @@ -/** - * Copyright (c) 2025 Huawei Technologies Co., Ltd. - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of - * CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ - -#ifndef CATLASS_KDA_BLOCK_MMAD_PRELOAD_ASYNC_HPP -#define CATLASS_KDA_BLOCK_MMAD_PRELOAD_ASYNC_HPP - -#include "catlass/kda_catlass.hpp" -#include "catlass/arch/kda_resource.hpp" -#include "catlass/kda_coord.hpp" -#include "catlass/detail/kda_callback.hpp" -#include "catlass/kda_gemm_coord.hpp" -#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" -#include "catlass/gemm/kda_gemm_helper.hpp" -#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" -#include "catlass/gemm/tile/kda_tile_mmad.hpp" - -namespace Catlass::Gemm::Block { - -template < - uint32_t PRELOAD_STAGES_, - uint32_t L1_STAGES_, - uint32_t L0A_STAGES_, - uint32_t L0B_STAGES_, - uint32_t L0C_STAGES_, - bool ENABLE_UNIT_FLAG_, - bool ENABLE_SHUFFLE_K_, - class L1TileShape_, - class L0TileShape_, - class AType_, - class BType_, - class CType_, - class BiasType_, - class TileCopy_, - class TileMmad_> -struct BlockMmad< - MmadAtlasA2PreloadAsync< - PRELOAD_STAGES_, - L1_STAGES_, - L0A_STAGES_, - L0B_STAGES_, - L0C_STAGES_, - ENABLE_UNIT_FLAG_, - ENABLE_SHUFFLE_K_>, - L1TileShape_, - L0TileShape_, - AType_, - BType_, - CType_, - BiasType_, - TileCopy_, - TileMmad_> { -public: - // Type Aliases - using DispatchPolicy = MmadAtlasA2PreloadAsync< - PRELOAD_STAGES_, - L1_STAGES_, - L0A_STAGES_, - L0B_STAGES_, - L0C_STAGES_, - ENABLE_UNIT_FLAG_, - ENABLE_SHUFFLE_K_>; - using ArchTag = typename DispatchPolicy::ArchTag; - using L1TileShape = L1TileShape_; - using L0TileShape = L0TileShape_; - using ElementA = typename AType_::Element; - using LayoutA = typename AType_::Layout; - using ElementB = typename BType_::Element; - using LayoutB = typename BType_::Layout; - using ElementC = typename CType_::Element; - using LayoutC = typename CType_::Layout; - using TileMmad = TileMmad_; - using CopyGmToL1A = typename TileCopy_::CopyGmToL1A; - using CopyGmToL1B = typename TileCopy_::CopyGmToL1B; - using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; - using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; - using CopyL0CToGm = typename TileCopy_::CopyL0CToGm; - using ElementAccumulator = - typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; - using LayoutAInL1 = typename CopyL1ToL0A::LayoutSrc; - using LayoutBInL1 = typename CopyL1ToL0B::LayoutSrc; - using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst; - using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst; - using LayoutCInL0 = layout::zN; - - using L1AAlignHelper = Gemm::helper::L1AlignHelper; - using L1BAlignHelper = Gemm::helper::L1AlignHelper; - - static constexpr uint32_t PRELOAD_STAGES = DispatchPolicy::PRELOAD_STAGES; - static constexpr uint32_t L1_STAGES = DispatchPolicy::L1_STAGES; - static constexpr uint32_t L0A_STAGES = DispatchPolicy::L0A_STAGES; - static constexpr uint32_t L0B_STAGES = DispatchPolicy::L0B_STAGES; - static constexpr uint32_t L0C_STAGES = DispatchPolicy::L0C_STAGES; - - static constexpr bool ENABLE_UNIT_FLAG = DispatchPolicy::ENABLE_UNIT_FLAG; - static constexpr bool ENABLE_SHUFFLE_K = DispatchPolicy::ENABLE_SHUFFLE_K; - - // L1 tile size - static constexpr uint32_t L1A_TILE_SIZE = L1TileShape::M * L1TileShape::K * sizeof(ElementA); - static constexpr uint32_t L1B_TILE_SIZE = L1TileShape::N * L1TileShape::K * sizeof(ElementB); - // L0 tile size - static constexpr uint32_t L0A_TILE_SIZE = L0TileShape::M * L0TileShape::K * sizeof(ElementA); - static constexpr uint32_t L0B_TILE_SIZE = L0TileShape::K * L0TileShape::N * sizeof(ElementB); - static constexpr uint32_t L0C_TILE_SIZE = L1TileShape::M * L1TileShape::N * sizeof(ElementAccumulator); - - // Check LayoutC - static_assert(std::is_same_v, "LayoutC only support RowMajor yet!"); - - // Check L1TileShape - static_assert((L1A_TILE_SIZE + L1B_TILE_SIZE) * L1_STAGES <= ArchTag::L1_SIZE, - "L1TileShape exceeding the L1 space!"); - - // Check L0TileShape - static_assert(L0A_TILE_SIZE * L0A_STAGES <= ArchTag::L0A_SIZE, "L0TileShape exceeding the L0A space!"); - static_assert(L0B_TILE_SIZE * L0B_STAGES <= ArchTag::L0B_SIZE, "L0TileShape exceeding the L0B space!"); - static_assert(L0C_TILE_SIZE * L0C_STAGES <= ArchTag::L0C_SIZE, "L0TileShape exceeding the L0C space!"); - - static_assert(L1TileShape::M == L0TileShape::M && L1TileShape::N == L0TileShape::N, - "The situation where the basic blocks of L1 and L0 differ on the m and n axes is not supported yet"); - static_assert(L0TileShape::K <= L1TileShape::K, "L0TileShape::K cannot exceed L1TileShape::K"); - // 32B (256b) aligned - static_assert(Gemm::helper::TileShapeAlignChecker::_ALIGN == 256, - "Tile shape must be 32B aligned."); - - static constexpr auto L1A_LAYOUT = LayoutAInL1::template MakeLayout( - L1TileShape::M, L1TileShape::K); - static constexpr auto L1B_LAYOUT = LayoutBInL1::template MakeLayout( - L1TileShape::K, L1TileShape::N); - - CATLASS_DEVICE - BlockMmad(Arch::Resource &resource, uint32_t l1BufAddrStart = 0) - { - InitL1(resource, l1BufAddrStart); - InitL0A(resource); - InitL0B(resource); - InitL0C(resource); - } - - CATLASS_DEVICE - ~BlockMmad() - { - SynchronizeBlock(); - for (uint32_t i = 0; i < L1_STAGES; ++i) { - AscendC::WaitFlag(l1AEventList[i]); - AscendC::WaitFlag(l1BEventList[i]); - } - for (uint32_t i = 0; i < L0A_STAGES; ++i) { - AscendC::WaitFlag(l0AEventList[i]); - } - for (uint32_t i = 0; i < L0B_STAGES; ++i) { - AscendC::WaitFlag(l0BEventList[i]); - } - for (uint32_t i = 0; i < L0C_STAGES; ++i) { - AscendC::WaitFlag(l0CEventList[i]); - } - } - - CATLASS_DEVICE - void operator()( - AscendC::GlobalTensor const &gmBlockA, LayoutA const &layoutA, - AscendC::GlobalTensor const &gmBlockB, LayoutB const &layoutB, - AscendC::GlobalTensor const &gmBlockC, LayoutC const &layoutC, - GemmCoord const &actualShape, Callback &&callback = Callback{}) - { - uint32_t kTileCount = CeilDiv(actualShape.k()); - - uint32_t mRound = RoundUp(actualShape.m()); - uint32_t nRound = RoundUp(actualShape.n()); - - uint32_t startTileIdx = 0; - if constexpr (ENABLE_SHUFFLE_K) { - startTileIdx = AscendC::GetBlockIdx() % kTileCount; - } - - for (uint32_t kLoopIdx = 0; kLoopIdx < kTileCount; ++kLoopIdx) { - uint32_t kTileIdx = (startTileIdx + kLoopIdx < kTileCount) ? - (startTileIdx + kLoopIdx) : - (startTileIdx + kLoopIdx - kTileCount); - - uint32_t kActual = (kTileIdx < kTileCount - 1) ? - L1TileShape::K : - (actualShape.k() - kTileIdx * L1TileShape::K); - - // Emission load instruction from GM to L1 - MatrixCoord gmTileAOffset{0, kTileIdx * L1TileShape::K}; - MatrixCoord gmTileBOffset{kTileIdx * L1TileShape::K, 0}; - auto gmTileA = gmBlockA[layoutA.GetOffset(gmTileAOffset)]; - auto gmTileB = gmBlockB[layoutB.GetOffset(gmTileBOffset)]; - // Load first matrix A tile from GM to L1 - AscendC::WaitFlag(l1AEventList[l1ListId]); - auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShape.m(), kActual)); - copyGmToL1A(l1ATensorList[l1ListId], gmTileA, L1A_LAYOUT, layoutTileA); - AscendC::SetFlag(l1AEventList[l1ListId]); - // Load first matrix B tile from GM to L1 - AscendC::WaitFlag(l1BEventList[l1ListId]); - auto layoutTileB = layoutB.GetTileLayout(MakeCoord(kActual, actualShape.n())); - copyGmToL1B(l1BTensorList[l1ListId], gmTileB, L1B_LAYOUT, layoutTileB); - AscendC::SetFlag(l1BEventList[l1ListId]); - - // If the number of preload instructions reaches the upper limit, perform an mmad calculation on L1 tile - if (preloadCount == PRELOAD_STAGES) { - L1TileMmad(l1TileMmadParamsList[l1TileMmadParamsId]); - } - - // Store the current load status - uint32_t preloadL1TileMmadParamsId = (l1TileMmadParamsId + preloadCount < PRELOAD_STAGES) ? - (l1TileMmadParamsId + preloadCount) : - (l1TileMmadParamsId + preloadCount - PRELOAD_STAGES); - auto &l1TileMmadParams = l1TileMmadParamsList[preloadL1TileMmadParamsId]; - l1TileMmadParams.l1ListId = l1ListId; - l1TileMmadParams.mRound = mRound; - l1TileMmadParams.nRound = nRound; - l1TileMmadParams.kActual = kActual; - l1TileMmadParams.isKLoopFirst = (kLoopIdx == 0); - l1TileMmadParams.isKLoopLast = (kLoopIdx == kTileCount - 1); - if (kLoopIdx == kTileCount - 1) { - l1TileMmadParams.gmBlockC = gmBlockC; - l1TileMmadParams.layoutCInGm = layoutC.GetTileLayout(actualShape.GetCoordMN()); - l1TileMmadParams.callback = callback; - } - - if (preloadCount < PRELOAD_STAGES) { - ++preloadCount; - } else { - l1TileMmadParamsId = (l1TileMmadParamsId + 1 < PRELOAD_STAGES) ? (l1TileMmadParamsId + 1) : 0; - } - l1ListId = (l1ListId + 1 < L1_STAGES) ? (l1ListId + 1) : 0; - } - } - - CATLASS_DEVICE - void SynchronizeBlock() - { - while (preloadCount > 0) { - L1TileMmad(l1TileMmadParamsList[l1TileMmadParamsId]); - l1TileMmadParamsId = (l1TileMmadParamsId + 1 < PRELOAD_STAGES) ? (l1TileMmadParamsId + 1) : 0; - --preloadCount; - } - } - -private: - struct L1TileMmadParams { - uint32_t l1ListId; - uint32_t mRound; - uint32_t nRound; - uint32_t kActual; - bool isKLoopFirst; - bool isKLoopLast; - AscendC::GlobalTensor gmBlockC; - LayoutC layoutCInGm; - Callback callback; - - CATLASS_DEVICE - L1TileMmadParams() = default; - }; - - CATLASS_DEVICE - void InitL1(Arch::Resource &resource, uint32_t l1BufAddrStart) - { - uint32_t l1AOffset = l1BufAddrStart; - uint32_t l1BOffset = l1BufAddrStart + L1A_TILE_SIZE * L1_STAGES; - for (uint32_t i = 0; i < L1_STAGES; ++i) { - l1ATensorList[i] = resource.l1Buf.template GetBufferByByte(l1AOffset + L1A_TILE_SIZE * i); - l1BTensorList[i] = resource.l1Buf.template GetBufferByByte(l1BOffset + L1B_TILE_SIZE * i); - l1AEventList[i] = i; - l1BEventList[i] = i + L1_STAGES; - AscendC::SetFlag(l1AEventList[i]); - AscendC::SetFlag(l1BEventList[i]); - } - } - - CATLASS_DEVICE - void InitL0A(Arch::Resource &resource) - { - for (uint32_t i = 0; i < L0A_STAGES; ++i) { - l0ATensorList[i] = resource.l0ABuf.template GetBufferByByte(L0A_TILE_SIZE * i); - l0AEventList[i] = i; - AscendC::SetFlag(l0AEventList[i]); - } - } - - CATLASS_DEVICE - void InitL0B(Arch::Resource &resource) - { - for (uint32_t i = 0; i < L0B_STAGES; ++i) { - l0BTensorList[i] = resource.l0BBuf.template GetBufferByByte(L0B_TILE_SIZE * i); - l0BEventList[i] = i + L0A_STAGES; - AscendC::SetFlag(l0BEventList[i]); - } - } - - CATLASS_DEVICE - void InitL0C(Arch::Resource &resource) - { - for (uint32_t i = 0; i < L0C_STAGES; ++i) { - l0CTensorList[i] = resource.l0CBuf.template GetBufferByByte(L0C_TILE_SIZE * i); - l0CEventList[i] = i; - AscendC::SetFlag(l0CEventList[i]); - } - } - - CATLASS_DEVICE - void L1TileMmad(L1TileMmadParams const ¶ms) - { - uint32_t mPartLoop = CeilDiv(params.mRound); - uint32_t nPartLoop = CeilDiv(params.nRound); - uint32_t kPartLoop = CeilDiv(params.kActual); - auto &l1ATensor = l1ATensorList[params.l1ListId]; - auto &l1BTensor = l1BTensorList[params.l1ListId]; - - auto &l0CTensor = l0CTensorList[l0CListId]; - LayoutCInL0 layoutCInL0 = LayoutCInL0::MakeLayoutInL0C(MakeCoord(params.mRound, params.nRound)); - - if constexpr (!ENABLE_UNIT_FLAG) { - if (params.isKLoopFirst) { - AscendC::WaitFlag(l0CEventList[l0CListId]); - } - } - - for (uint32_t mPartIdx = 0; mPartIdx < mPartLoop; ++mPartIdx) { - uint32_t mPartActual = (mPartIdx < mPartLoop - 1) ? - L0TileShape::M : - (params.mRound - mPartIdx * L0TileShape::M); - - for (uint32_t kPartIdx = 0; kPartIdx < kPartLoop; ++kPartIdx) { - uint32_t kPartActual = (kPartIdx < kPartLoop - 1) ? - L0TileShape::K : - (params.kActual - kPartIdx * L0TileShape::K); - - auto &l0ATile = l0ATensorList[l0AListId]; - auto layoutAInL0 = LayoutAInL0::template MakeLayout(mPartActual, kPartActual); - auto l1AOffset = MakeCoord(mPartIdx, kPartIdx) * L0TileShape::ToCoordMK(); - auto l1ATile = l1ATensor[L1A_LAYOUT.GetOffset(l1AOffset)]; - - AscendC::WaitFlag(l0AEventList[l0AListId]); - if ((mPartIdx == 0) && (kPartIdx == 0)) { - AscendC::WaitFlag(l1AEventList[params.l1ListId]); - } - copyL1ToL0A(l0ATile, l1ATile, layoutAInL0, L1A_LAYOUT); - if ((mPartIdx == mPartLoop - 1) && (kPartIdx == kPartLoop - 1)) { - AscendC::SetFlag(l1AEventList[params.l1ListId]); - } - - for (uint32_t nPartIdx = 0; nPartIdx < nPartLoop; ++nPartIdx) { - uint32_t nPartActual = (nPartIdx < nPartLoop - 1) ? - L0TileShape::N : - (params.nRound - nPartIdx * L0TileShape::N); - - auto &l0BTile = l0BTensorList[l0BListId]; - auto layoutBInL0 = LayoutBInL0::template MakeLayout(kPartActual, nPartActual); - auto l1BOffset = MakeCoord(kPartIdx, nPartIdx) * L0TileShape::ToCoordKN(); - auto l1BTile = l1BTensor[L1B_LAYOUT.GetOffset(l1BOffset)]; - - AscendC::WaitFlag(l0BEventList[l0BListId]); - if ((kPartIdx == 0) && (nPartIdx == 0)) { - AscendC::WaitFlag(l1BEventList[params.l1ListId]); - } - copyL1ToL0B(l0BTile, l1BTile, layoutBInL0, L1B_LAYOUT); - if ((kPartIdx == kPartLoop - 1) && (nPartIdx == nPartLoop - 1)) { - AscendC::SetFlag(l1BEventList[params.l1ListId]); - } - - AscendC::SetFlag(EVENT_ID0); - - auto l0COffset = MakeCoord(mPartIdx, nPartIdx) * L0TileShape::ToCoordMN(); - auto l0CTile = l0CTensor[layoutCInL0.GetOffset(l0COffset)]; - - AscendC::WaitFlag(EVENT_ID0); - // If the current tile is the first tile on the k axis, the accumulator needs to be reset to 0 - bool initC = (params.isKLoopFirst && (kPartIdx == 0)); - // If the unit flag is enabled, the unit flag is set according to the calculation progress - uint8_t unitFlag = 0b00; - if constexpr (ENABLE_UNIT_FLAG) { - if (params.isKLoopLast && - (mPartIdx == mPartLoop - 1) && (kPartIdx == kPartLoop - 1) && (nPartIdx == nPartLoop - 1)) { - unitFlag = 0b11; - } else { - unitFlag = 0b10; - } - } - tileMmad(l0CTile, l0ATile, l0BTile, mPartActual, nPartActual, kPartActual, initC, unitFlag); - - AscendC::SetFlag(l0BEventList[l0BListId]); - l0BListId = (l0BListId + 1 < L0B_STAGES) ? (l0BListId + 1) : 0; - } - AscendC::SetFlag(l0AEventList[l0AListId]); - l0AListId = (l0AListId + 1 < L0A_STAGES) ? (l0AListId + 1) : 0; - } - } - - if (params.isKLoopLast) { - auto layoutCInGm = params.layoutCInGm; - - if constexpr (!ENABLE_UNIT_FLAG) { - AscendC::SetFlag(l0CEventList[l0CListId]); - AscendC::WaitFlag(l0CEventList[l0CListId]); - copyL0CToGm(params.gmBlockC, l0CTensor, layoutCInGm, layoutCInL0); - AscendC::SetFlag(l0CEventList[l0CListId]); - } else { - copyL0CToGm(params.gmBlockC, l0CTensor, layoutCInGm, layoutCInL0, 0b11); - } - l0CListId = (l0CListId + 1 < L0C_STAGES) ? (l0CListId + 1) : 0; - - if (params.callback) { - params.callback(); - } - } - } - - AscendC::LocalTensor l1ATensorList[L1_STAGES]; - AscendC::LocalTensor l1BTensorList[L1_STAGES]; - int32_t l1AEventList[L1_STAGES]; - int32_t l1BEventList[L1_STAGES]; - uint32_t l1ListId{0}; - - AscendC::LocalTensor l0ATensorList[L0A_STAGES]; - int32_t l0AEventList[L0A_STAGES]; - uint32_t l0AListId{0}; - - AscendC::LocalTensor l0BTensorList[L0B_STAGES]; - int32_t l0BEventList[L0B_STAGES]; - uint32_t l0BListId{0}; - - AscendC::LocalTensor l0CTensorList[L0C_STAGES_]; - int32_t l0CEventList[L0C_STAGES_]; - uint32_t l0CListId{0}; - - L1TileMmadParams l1TileMmadParamsList[PRELOAD_STAGES]; - uint32_t l1TileMmadParamsId{0}; - uint32_t preloadCount{0}; - - TileMmad tileMmad; - CopyGmToL1A copyGmToL1A; - CopyGmToL1B copyGmToL1B; - CopyL1ToL0A copyL1ToL0A; - CopyL1ToL0B copyL1ToL0B; - CopyL0CToGm copyL0CToGm; -}; - -} // namespace Catlass::Gemm::Block - -#endif // CATLASS_KDA_BLOCK_MMAD_PRELOAD_ASYNC_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_preload_async_with_callback.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_preload_async_with_callback.hpp deleted file mode 100644 index df26fa7548..0000000000 --- a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad_preload_async_with_callback.hpp +++ /dev/null @@ -1,449 +0,0 @@ -/** - * Copyright (c) 2025 Huawei Technologies Co., Ltd. - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of - * CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ - -#ifndef CATLASS_KDA_BLOCK_MMAD_PRELOAD_ASYNC_WITH_CALLBACK_HPP -#define CATLASS_KDA_BLOCK_MMAD_PRELOAD_ASYNC_WITH_CALLBACK_HPP - -#include "catlass/kda_catlass.hpp" -#include "catlass/arch/kda_resource.hpp" -#include "catlass/kda_coord.hpp" -#include "catlass/detail/kda_callback.hpp" -#include "catlass/kda_gemm_coord.hpp" -#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" -#include "catlass/gemm/kda_gemm_helper.hpp" -#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" -#include "catlass/gemm/tile/kda_tile_mmad.hpp" - -namespace Catlass::Gemm::Block { - -template < - uint32_t PRELOAD_STAGES_, - uint32_t L1_STAGES_, - uint32_t L0A_STAGES_, - uint32_t L0B_STAGES_, - uint32_t L0C_STAGES_, - bool ENABLE_UNIT_FLAG_, - bool ENABLE_SHUFFLE_K_, - class L1TileShape_, - class L0TileShape_, - class AType_, - class BType_, - class CType_, - class BiasType_, - class TileCopy_, - class TileMmad_> -struct BlockMmad< - MmadAtlasA2PreloadAsyncWithCallback< - PRELOAD_STAGES_, - L1_STAGES_, - L0A_STAGES_, - L0B_STAGES_, - L0C_STAGES_, - ENABLE_UNIT_FLAG_, - ENABLE_SHUFFLE_K_>, - L1TileShape_, - L0TileShape_, - AType_, - BType_, - CType_, - BiasType_, - TileCopy_, - TileMmad_> { -public: - // Type Aliases - using DispatchPolicy = MmadAtlasA2PreloadAsyncWithCallback< - PRELOAD_STAGES_, - L1_STAGES_, - L0A_STAGES_, - L0B_STAGES_, - L0C_STAGES_, - ENABLE_UNIT_FLAG_, - ENABLE_SHUFFLE_K_>; - using ArchTag = typename DispatchPolicy::ArchTag; - using L1TileShape = L1TileShape_; - using L0TileShape = L0TileShape_; - using ElementA = typename AType_::Element; - using LayoutA = typename AType_::Layout; - using ElementB = typename BType_::Element; - using LayoutB = typename BType_::Layout; - using ElementC = typename CType_::Element; - using LayoutC = typename CType_::Layout; - using TileMmad = TileMmad_; - using CopyGmToL1A = typename TileCopy_::CopyGmToL1A; - using CopyGmToL1B = typename TileCopy_::CopyGmToL1B; - using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; - using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; - using CopyL0CToGm = typename TileCopy_::CopyL0CToGm; - using ElementAccumulator = - typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; - using LayoutAInL1 = typename CopyL1ToL0A::LayoutSrc; - using LayoutBInL1 = typename CopyL1ToL0B::LayoutSrc; - using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst; - using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst; - using LayoutCInL0 = layout::zN; - - using L1AAlignHelper = Gemm::helper::L1AlignHelper; - using L1BAlignHelper = Gemm::helper::L1AlignHelper; - - static constexpr uint32_t PRELOAD_STAGES = DispatchPolicy::PRELOAD_STAGES; - static constexpr uint32_t L1_STAGES = DispatchPolicy::L1_STAGES; - static constexpr uint32_t L0A_STAGES = DispatchPolicy::L0A_STAGES; - static constexpr uint32_t L0B_STAGES = DispatchPolicy::L0B_STAGES; - static constexpr uint32_t L0C_STAGES = DispatchPolicy::L0C_STAGES; - - static constexpr bool ENABLE_UNIT_FLAG = DispatchPolicy::ENABLE_UNIT_FLAG; - static constexpr bool ENABLE_SHUFFLE_K = DispatchPolicy::ENABLE_SHUFFLE_K; - - // L1 tile size - static constexpr uint32_t L1A_TILE_SIZE = L1TileShape::M * L1TileShape::K * sizeof(ElementA); - static constexpr uint32_t L1B_TILE_SIZE = L1TileShape::N * L1TileShape::K * sizeof(ElementB); - // L0 tile size - static constexpr uint32_t L0A_TILE_SIZE = L0TileShape::M * L0TileShape::K * sizeof(ElementA); - static constexpr uint32_t L0B_TILE_SIZE = L0TileShape::K * L0TileShape::N * sizeof(ElementB); - static constexpr uint32_t L0C_TILE_SIZE = L1TileShape::M * L1TileShape::N * sizeof(ElementAccumulator); - - // Check LayoutC - static_assert(std::is_same_v, "LayoutC only support RowMajor yet!"); - - // Check L1TileShape - static_assert((L1A_TILE_SIZE + L1B_TILE_SIZE) * L1_STAGES <= ArchTag::L1_SIZE, - "L1TileShape exceeding the L1 space!"); - - // Check L0TileShape - static_assert(L0A_TILE_SIZE * L0A_STAGES <= ArchTag::L0A_SIZE, "L0TileShape exceeding the L0A space!"); - static_assert(L0B_TILE_SIZE * L0B_STAGES <= ArchTag::L0B_SIZE, "L0TileShape exceeding the L0B space!"); - static_assert(L0C_TILE_SIZE * L0C_STAGES <= ArchTag::L0C_SIZE, "L0TileShape exceeding the L0C space!"); - - static_assert(L1TileShape::M == L0TileShape::M && L1TileShape::N == L0TileShape::N, - "The situation where the basic blocks of L1 and L0 differ on the m and n axes is not supported yet"); - static_assert(L0TileShape::K <= L1TileShape::K, "L0TileShape::K cannot exceed L1TileShape::K"); - // 32B (256b) aligned - static_assert(Gemm::helper::TileShapeAlignChecker::_ALIGN == 256, - "Tile shape must be 32B aligned."); - - static constexpr auto L1A_LAYOUT = LayoutAInL1::template MakeLayout(L1TileShape::M, L1TileShape::K); - static constexpr auto L1B_LAYOUT = LayoutBInL1::template MakeLayout(L1TileShape::K, L1TileShape::N); - - CATLASS_DEVICE - BlockMmad(Arch::Resource &resource, uint32_t l1BufAddrStart = 0) - { - InitL1(resource, l1BufAddrStart); - InitL0A(resource); - InitL0B(resource); - InitL0C(resource); - } - - CATLASS_DEVICE - ~BlockMmad() - { - SynchronizeBlock(); - for (uint32_t i = 0; i < L1_STAGES; ++i) { - AscendC::WaitFlag(l1AEventList[i]); - AscendC::WaitFlag(l1BEventList[i]); - } - for (uint32_t i = 0; i < L0A_STAGES; ++i) { - AscendC::WaitFlag(l0AEventList[i]); - } - for (uint32_t i = 0; i < L0B_STAGES; ++i) { - AscendC::WaitFlag(l0BEventList[i]); - } - for (uint32_t i = 0; i < L0C_STAGES; ++i) { - AscendC::WaitFlag(l0CEventList[i]); - } - } - - CATLASS_DEVICE - void operator()( - AscendC::GlobalTensor const &gmBlockA, LayoutA const &layoutA, - AscendC::GlobalTensor const &gmBlockB, LayoutB const &layoutB, - AscendC::GlobalTensor const &gmBlockC, LayoutC const &layoutC, - GemmCoord const &actualShape, - Callback const &callbackBeforeFixpipe, Callback const &callbackAfterFixpipe) - { - uint32_t kTileCount = CeilDiv(actualShape.k()); - - uint32_t mRound = RoundUp(actualShape.m()); - uint32_t nRound = RoundUp(actualShape.n()); - - uint32_t startTileIdx = 0; - if constexpr (ENABLE_SHUFFLE_K) { - startTileIdx = AscendC::GetBlockIdx() % kTileCount; - } - - for (uint32_t kLoopIdx = 0; kLoopIdx < kTileCount; ++kLoopIdx) { - uint32_t kTileIdx = (startTileIdx + kLoopIdx < kTileCount) ? - (startTileIdx + kLoopIdx) : - (startTileIdx + kLoopIdx - kTileCount); - - uint32_t kActual = (kTileIdx < kTileCount - 1) ? - L1TileShape::K : - (actualShape.k() - kTileIdx * L1TileShape::K); - - // Emission load instruction from GM to L1 - MatrixCoord gmTileAOffset{0, kTileIdx * L1TileShape::K}; - MatrixCoord gmTileBOffset{kTileIdx * L1TileShape::K, 0}; - auto gmTileA = gmBlockA[layoutA.GetOffset(gmTileAOffset)]; - auto gmTileB = gmBlockB[layoutB.GetOffset(gmTileBOffset)]; - // Load first matrix A tile from GM to L1 - AscendC::WaitFlag(l1AEventList[l1ListId]); - auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShape.m(), kActual)); - copyGmToL1A(l1ATensorList[l1ListId], gmTileA, L1A_LAYOUT, layoutTileA); - AscendC::SetFlag(l1AEventList[l1ListId]); - // Load first matrix B tile from GM to L1 - AscendC::WaitFlag(l1BEventList[l1ListId]); - auto layoutTileB = layoutB.GetTileLayout(MakeCoord(kActual, actualShape.n())); - copyGmToL1B(l1BTensorList[l1ListId], gmTileB, L1B_LAYOUT, layoutTileB); - AscendC::SetFlag(l1BEventList[l1ListId]); - - // If the number of preload instructions reaches the upper limit, perform an mmad calculation on L1 tile - if (preloadCount == PRELOAD_STAGES) { - L1TileMmad(l1TileMmadParamsList[l1TileMmadParamsId]); - } - - // Store the current load status - uint32_t preloadL1TileMmadParamsId = (l1TileMmadParamsId + preloadCount < PRELOAD_STAGES) ? - (l1TileMmadParamsId + preloadCount) : - (l1TileMmadParamsId + preloadCount - PRELOAD_STAGES); - auto &l1TileMmadParams = l1TileMmadParamsList[preloadL1TileMmadParamsId]; - l1TileMmadParams.l1ListId = l1ListId; - l1TileMmadParams.mRound = mRound; - l1TileMmadParams.nRound = nRound; - l1TileMmadParams.kActual = kActual; - l1TileMmadParams.isKLoopFirst = (kLoopIdx == 0); - l1TileMmadParams.isKLoopLast = (kLoopIdx == kTileCount - 1); - if (kLoopIdx == kTileCount - 1) { - l1TileMmadParams.gmBlockC = gmBlockC; - l1TileMmadParams.layoutCInGm = layoutC.GetTileLayout(actualShape.GetCoordMN()); - l1TileMmadParams.callbackBeforeFixpipe = callbackBeforeFixpipe; - l1TileMmadParams.callbackAfterFixpipe = callbackAfterFixpipe; - } - - if (preloadCount < PRELOAD_STAGES) { - ++preloadCount; - } else { - l1TileMmadParamsId = (l1TileMmadParamsId + 1 < PRELOAD_STAGES) ? (l1TileMmadParamsId + 1) : 0; - } - l1ListId = (l1ListId + 1 < L1_STAGES) ? (l1ListId + 1) : 0; - } - } - - CATLASS_DEVICE - void SynchronizeBlock() - { - while (preloadCount > 0) { - L1TileMmad(l1TileMmadParamsList[l1TileMmadParamsId]); - l1TileMmadParamsId = (l1TileMmadParamsId + 1 < PRELOAD_STAGES) ? (l1TileMmadParamsId + 1) : 0; - --preloadCount; - } - } - -private: - struct L1TileMmadParams { - uint32_t l1ListId; - uint32_t mRound; - uint32_t nRound; - uint32_t kActual; - bool isKLoopFirst; - bool isKLoopLast; - AscendC::GlobalTensor gmBlockC; - LayoutC layoutCInGm; - Callback callbackBeforeFixpipe; - Callback callbackAfterFixpipe; - - CATLASS_DEVICE - L1TileMmadParams() = default; - }; - - CATLASS_DEVICE - void InitL1(Arch::Resource &resource, uint32_t l1BufAddrStart) - { - uint32_t l1AOffset = l1BufAddrStart; - uint32_t l1BOffset = l1BufAddrStart + L1A_TILE_SIZE * L1_STAGES; - for (uint32_t i = 0; i < L1_STAGES; ++i) { - l1ATensorList[i] = resource.l1Buf.template GetBufferByByte(l1AOffset + L1A_TILE_SIZE * i); - l1BTensorList[i] = resource.l1Buf.template GetBufferByByte(l1BOffset + L1B_TILE_SIZE * i); - l1AEventList[i] = i; - l1BEventList[i] = i + L1_STAGES; - AscendC::SetFlag(l1AEventList[i]); - AscendC::SetFlag(l1BEventList[i]); - } - } - - CATLASS_DEVICE - void InitL0A(Arch::Resource &resource) - { - for (uint32_t i = 0; i < L0A_STAGES; ++i) { - l0ATensorList[i] = resource.l0ABuf.template GetBufferByByte(L0A_TILE_SIZE * i); - l0AEventList[i] = i; - AscendC::SetFlag(l0AEventList[i]); - } - } - - CATLASS_DEVICE - void InitL0B(Arch::Resource &resource) - { - for (uint32_t i = 0; i < L0B_STAGES; ++i) { - l0BTensorList[i] = resource.l0BBuf.template GetBufferByByte(L0B_TILE_SIZE * i); - l0BEventList[i] = i + L0A_STAGES; - AscendC::SetFlag(l0BEventList[i]); - } - } - - CATLASS_DEVICE - void InitL0C(Arch::Resource &resource) - { - for (uint32_t i = 0; i < L0C_STAGES; ++i) { - l0CTensorList[i] = resource.l0CBuf.template GetBufferByByte(L0C_TILE_SIZE * i); - l0CEventList[i] = i; - AscendC::SetFlag(l0CEventList[i]); - } - } - - CATLASS_DEVICE - void L1TileMmad(L1TileMmadParams const ¶ms) - { - uint32_t mPartLoop = CeilDiv(params.mRound); - uint32_t nPartLoop = CeilDiv(params.nRound); - uint32_t kPartLoop = CeilDiv(params.kActual); - auto &l1ATensor = l1ATensorList[params.l1ListId]; - auto &l1BTensor = l1BTensorList[params.l1ListId]; - - auto &l0CTensor = l0CTensorList[l0CListId]; - LayoutCInL0 layoutCInL0 = LayoutCInL0::MakeLayoutInL0C(MakeCoord(params.mRound, params.nRound)); - - if constexpr (!ENABLE_UNIT_FLAG) { - if (params.isKLoopFirst) { - AscendC::WaitFlag(l0CEventList[l0CListId]); - } - } - - for (uint32_t mPartIdx = 0; mPartIdx < mPartLoop; ++mPartIdx) { - uint32_t mPartActual = (mPartIdx < mPartLoop - 1) ? - L0TileShape::M : - (params.mRound - mPartIdx * L0TileShape::M); - - for (uint32_t kPartIdx = 0; kPartIdx < kPartLoop; ++kPartIdx) { - uint32_t kPartActual = (kPartIdx < kPartLoop - 1) ? - L0TileShape::K : - (params.kActual - kPartIdx * L0TileShape::K); - - auto &l0ATile = l0ATensorList[l0AListId]; - auto layoutAInL0 = LayoutAInL0::template MakeLayout(mPartActual, kPartActual); - auto l1AOffset = MakeCoord(mPartIdx, kPartIdx) * L0TileShape::ToCoordMK(); - auto l1ATile = l1ATensor[L1A_LAYOUT.GetOffset(l1AOffset)]; - - AscendC::WaitFlag(l0AEventList[l0AListId]); - if ((mPartIdx == 0) && (kPartIdx == 0)) { - AscendC::WaitFlag(l1AEventList[params.l1ListId]); - } - copyL1ToL0A(l0ATile, l1ATile, layoutAInL0, L1A_LAYOUT); - if ((mPartIdx == mPartLoop - 1) && (kPartIdx == kPartLoop - 1)) { - AscendC::SetFlag(l1AEventList[params.l1ListId]); - } - - for (uint32_t nPartIdx = 0; nPartIdx < nPartLoop; ++nPartIdx) { - uint32_t nPartActual = (nPartIdx < nPartLoop - 1) ? - L0TileShape::N : - (params.nRound - nPartIdx * L0TileShape::N); - - auto &l0BTile = l0BTensorList[l0BListId]; - auto layoutBInL0 = LayoutBInL0::template MakeLayout(kPartActual, nPartActual); - auto l1BOffset = MakeCoord(kPartIdx, nPartIdx) * L0TileShape::ToCoordKN(); - auto l1BTile = l1BTensor[L1B_LAYOUT.GetOffset(l1BOffset)]; - - AscendC::WaitFlag(l0BEventList[l0BListId]); - if ((kPartIdx == 0) && (nPartIdx == 0)) { - AscendC::WaitFlag(l1BEventList[params.l1ListId]); - } - copyL1ToL0B(l0BTile, l1BTile, layoutBInL0, L1B_LAYOUT); - if ((kPartIdx == kPartLoop - 1) && (nPartIdx == nPartLoop - 1)) { - AscendC::SetFlag(l1BEventList[params.l1ListId]); - } - - AscendC::SetFlag(EVENT_ID0); - - auto l0COffset = MakeCoord(mPartIdx, nPartIdx) * L0TileShape::ToCoordMN(); - auto l0CTile = l0CTensor[layoutCInL0.GetOffset(l0COffset)]; - - AscendC::WaitFlag(EVENT_ID0); - // If the current tile is the first tile on the k axis, the accumulator needs to be reset to 0 - bool initC = (params.isKLoopFirst && (kPartIdx == 0)); - // If the unit flag is enabled, the unit flag is set according to the calculation progress - uint8_t unitFlag = 0b00; - if constexpr (ENABLE_UNIT_FLAG) { - if (params.isKLoopLast && - (mPartIdx == mPartLoop - 1) && (kPartIdx == kPartLoop - 1) && (nPartIdx == nPartLoop - 1)) { - unitFlag = 0b11; - } else { - unitFlag = 0b10; - } - } - tileMmad(l0CTile, l0ATile, l0BTile, mPartActual, nPartActual, kPartActual, initC, unitFlag); - - AscendC::SetFlag(l0BEventList[l0BListId]); - l0BListId = (l0BListId + 1 < L0B_STAGES) ? (l0BListId + 1) : 0; - } - AscendC::SetFlag(l0AEventList[l0AListId]); - l0AListId = (l0AListId + 1 < L0A_STAGES) ? (l0AListId + 1) : 0; - } - } - - if (params.isKLoopLast) { - auto layoutCInGm = params.layoutCInGm; - - params.callbackBeforeFixpipe(); - - if constexpr (!ENABLE_UNIT_FLAG) { - AscendC::SetFlag(l0CEventList[l0CListId]); - AscendC::WaitFlag(l0CEventList[l0CListId]); - copyL0CToGm(params.gmBlockC, l0CTensor, layoutCInGm, layoutCInL0); - AscendC::SetFlag(l0CEventList[l0CListId]); - } else { - copyL0CToGm(params.gmBlockC, l0CTensor, layoutCInGm, layoutCInL0, 0b11); - } - l0CListId = (l0CListId + 1 < L0C_STAGES) ? (l0CListId + 1) : 0; - - params.callbackAfterFixpipe(); - } - } - - AscendC::LocalTensor l1ATensorList[L1_STAGES]; - AscendC::LocalTensor l1BTensorList[L1_STAGES]; - int32_t l1AEventList[L1_STAGES]; - int32_t l1BEventList[L1_STAGES]; - uint32_t l1ListId{0}; - - AscendC::LocalTensor l0ATensorList[L0A_STAGES]; - int32_t l0AEventList[L0A_STAGES]; - uint32_t l0AListId{0}; - - AscendC::LocalTensor l0BTensorList[L0B_STAGES]; - int32_t l0BEventList[L0B_STAGES]; - uint32_t l0BListId{0}; - - AscendC::LocalTensor l0CTensorList[L0C_STAGES_]; - int32_t l0CEventList[L0C_STAGES_]; - uint32_t l0CListId{0}; - - L1TileMmadParams l1TileMmadParamsList[PRELOAD_STAGES]; - uint32_t l1TileMmadParamsId{0}; - uint32_t preloadCount{0}; - - TileMmad tileMmad; - CopyGmToL1A copyGmToL1A; - CopyGmToL1B copyGmToL1B; - CopyL1ToL0A copyL1ToL0A; - CopyL1ToL0B copyL1ToL0B; - CopyL0CToGm copyL0CToGm; -}; - -} // namespace Catlass::Gemm::Block - -#endif // CATLASS_KDA_BLOCK_MMAD_PRELOAD_ASYNC_WITH_CALLBACK_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_swizzle.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_swizzle.hpp deleted file mode 100644 index db6bb00841..0000000000 --- a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_swizzle.hpp +++ /dev/null @@ -1,1018 +0,0 @@ -/** - * Copyright (c) 2025 Huawei Technologies Co., Ltd. - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of - * CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ - -#ifndef CATLASS_KDA_BLOCK_SWIZZLE_HPP -#define CATLASS_KDA_BLOCK_SWIZZLE_HPP - -#include "catlass/kda_catlass.hpp" -#include "catlass/detail/kda_alignment.hpp" -#include "catlass/kda_gemm_coord.hpp" -#include "catlass/kda_matrix_coord.hpp" - -namespace Catlass::Gemm::Block { - -///////////////////////////////////////////////////////////////////////////////////////////////// - -/// Block swizzling function for Gemms -template -struct GemmIdentityBlockSwizzle { - /// Data members - - GemmCoord problemShape; - MatrixCoord tileMN; - MatrixCoord loopsMN; - - /// Methods - - CATLASS_DEVICE - GemmIdentityBlockSwizzle() {} - - CATLASS_DEVICE - GemmIdentityBlockSwizzle(GemmCoord const &problemShape_, MatrixCoord const &tileMN_) - : problemShape(problemShape_), - tileMN(tileMN_) - { - loopsMN = CeilDiv(MatrixCoord(problemShape.GetCoordMN()), tileMN); - } - - CATLASS_DEVICE - GemmIdentityBlockSwizzle(GemmCoord const &problemShape_, MatrixCoord const &tileMN_, - MatrixCoord const &loopsMN_) - : problemShape(problemShape_), - tileMN(tileMN_), - loopsMN(loopsMN_) - {} - - CATLASS_DEVICE - void Update(GemmCoord const &problemShape_, MatrixCoord const &tileMN_) - { - problemShape = problemShape_; - tileMN = tileMN_; - - loopsMN = CeilDiv(MatrixCoord(problemShape.GetCoordMN()), tileMN); - } - - CATLASS_DEVICE - void Update(GemmCoord const &problemShape_, MatrixCoord const &tileMN_, MatrixCoord const &loopsMN_) - { - problemShape = problemShape_; - tileMN = tileMN_; - loopsMN = loopsMN_; - } - - CATLASS_DEVICE - uint32_t GetCoreLoops() const - { - return loopsMN.row() * loopsMN.column(); - } - - CATLASS_DEVICE - uint32_t GetBatchIdx(uint32_t taskIdx) - { - return taskIdx / (GetCoreLoops()); - } - - CATLASS_DEVICE - GemmCoord GetBlockCoord(uint32_t taskIdx) - { - uint32_t innerIdx = taskIdx % GetCoreLoops(); - if constexpr (SwizzleDirection == 0) { // Zn - uint32_t tileBlockLoop = CeilDiv(loopsMN.row(), SwizzleOffset); - uint32_t tileBlockIdx = innerIdx / (SwizzleOffset * loopsMN.column()); - uint32_t inTileBlockIdx = innerIdx % (SwizzleOffset * loopsMN.column()); - - uint32_t nRow = SwizzleOffset; - if (tileBlockIdx == tileBlockLoop - 1) { - nRow = loopsMN.row() - SwizzleOffset * tileBlockIdx; - } - uint32_t mIdx = tileBlockIdx * SwizzleOffset + inTileBlockIdx % nRow; - uint32_t nIdx = inTileBlockIdx / nRow; - if (tileBlockIdx % 2 == 1) { - nIdx = loopsMN.column() - nIdx - 1; - } - return GemmCoord{mIdx, nIdx, 0}; - } else if constexpr (SwizzleDirection == 1) { // Nz - uint32_t tileBlockLoop = CeilDiv(loopsMN.column(), SwizzleOffset); - uint32_t tileBlockIdx = innerIdx / (SwizzleOffset * loopsMN.row()); - uint32_t inTileBlockIdx = innerIdx % (SwizzleOffset * loopsMN.row()); - - uint32_t nCol = SwizzleOffset; - if (tileBlockIdx == tileBlockLoop - 1) { - nCol = loopsMN.column() - SwizzleOffset * tileBlockIdx; - } - uint32_t mIdx = inTileBlockIdx / nCol; - uint32_t nIdx = tileBlockIdx * SwizzleOffset + inTileBlockIdx % nCol; - if (tileBlockIdx % 2 == 1) { - mIdx = loopsMN.row() - mIdx - 1; - } - return GemmCoord{mIdx, nIdx, 0}; - } - } - - CATLASS_DEVICE - GemmCoord GetActualBlockShape(GemmCoord blockCoord) - { - uint32_t mActual = (blockCoord.m() == (loopsMN.row() - 1)) ? - (problemShape.m() - blockCoord.m() * tileMN.row()) : - tileMN.row(); - uint32_t nActual = (blockCoord.n() == (loopsMN.column() - 1)) ? - (problemShape.n() - blockCoord.n() * tileMN.column()) : - tileMN.column(); - uint32_t kActual = problemShape.k(); - return GemmCoord{mActual, nActual, kActual}; - } -}; - -struct DynamicGemmIdentityBlockSwizzle : public GemmIdentityBlockSwizzle<> { - uint32_t swizzleOffset{1}; - uint32_t swizzleDirection{0}; - - CATLASS_DEVICE - DynamicGemmIdentityBlockSwizzle(GemmCoord const &problemShape_, MatrixCoord const &tileMN_, - uint32_t swizzleOffset_, uint32_t swizzleDirection_) - : swizzleOffset(swizzleOffset_), - swizzleDirection(swizzleDirection_), - GemmIdentityBlockSwizzle<>(problemShape_, tileMN_) - {} - - CATLASS_DEVICE - DynamicGemmIdentityBlockSwizzle(GemmCoord const &problemShape_, MatrixCoord const &tileMN_) - : GemmIdentityBlockSwizzle<>(problemShape_, tileMN_) - {} - - CATLASS_DEVICE - DynamicGemmIdentityBlockSwizzle() {} - - CATLASS_DEVICE - void SetSwizzleParams(uint32_t swizzleOffset_, uint32_t swizzleDirection_) - { - swizzleOffset = swizzleOffset_; - swizzleDirection = swizzleDirection_; - } - - CATLASS_DEVICE - GemmCoord GetBlockCoord(uint32_t taskIdx) - { - uint32_t innerIdx = taskIdx % GetCoreLoops(); - if (swizzleDirection == 0) { // Zn - uint32_t tileBlockLoop = CeilDiv(loopsMN.row(), swizzleOffset); - uint32_t tileBlockIdx = innerIdx / (swizzleOffset * loopsMN.column()); - uint32_t inTileBlockIdx = innerIdx % (swizzleOffset * loopsMN.column()); - - uint32_t nRow = swizzleOffset; - if (tileBlockIdx == tileBlockLoop - 1) { - nRow = loopsMN.row() - swizzleOffset * tileBlockIdx; - } - uint32_t mIdx = tileBlockIdx * swizzleOffset + inTileBlockIdx % nRow; - uint32_t nIdx = inTileBlockIdx / nRow; - if (tileBlockIdx % 2 == 1) { - nIdx = loopsMN.column() - nIdx - 1; - } - return GemmCoord{mIdx, nIdx, 0}; - } else { // Nz - uint32_t tileBlockLoop = CeilDiv(loopsMN.column(), swizzleOffset); - uint32_t tileBlockIdx = innerIdx / (swizzleOffset * loopsMN.row()); - uint32_t inTileBlockIdx = innerIdx % (swizzleOffset * loopsMN.row()); - - uint32_t nCol = swizzleOffset; - if (tileBlockIdx == tileBlockLoop - 1) { - nCol = loopsMN.column() - swizzleOffset * tileBlockIdx; - } - uint32_t mIdx = inTileBlockIdx / nCol; - uint32_t nIdx = tileBlockIdx * swizzleOffset + inTileBlockIdx % nCol; - if (tileBlockIdx % 2 == 1) { - mIdx = loopsMN.row() - mIdx - 1; - } - return GemmCoord{mIdx, nIdx, 0}; - } - } -}; - -/// Block swizzling function for Splitk Gemms -template -struct SplitkGemmIdentityBlockSwizzle { - /// Data members - - GemmCoord problemShape; - GemmCoord tileShape; - GemmCoord loopsMNK; - uint32_t splitkFactor = 1; // splite k dim into virtual cores - - /// Methods - - CATLASS_DEVICE - SplitkGemmIdentityBlockSwizzle() {} - - CATLASS_DEVICE - SplitkGemmIdentityBlockSwizzle( - GemmCoord const &problemShape_, GemmCoord const &tileShape_, uint32_t splitkFactor_ = 1) - : problemShape(problemShape_), - tileShape(tileShape_), - splitkFactor(splitkFactor_) - { - loopsMNK = CeilDiv(problemShape, tileShape); - } - - CATLASS_DEVICE - uint32_t GetKIdxBySplitkSliceIdx(uint32_t splitkSliceIdx) const - { - if (splitkSliceIdx < loopsMNK.k() % splitkFactor) { - return (loopsMNK.k() / splitkFactor + 1) * splitkSliceIdx; - } else { - return splitkSliceIdx * (loopsMNK.k() / splitkFactor) + loopsMNK.k() % splitkFactor; - } - } - - CATLASS_DEVICE - uint32_t GetSplitkSliceIdx(uint32_t taskIdx) const - { - uint32_t mnLoops = loopsMNK.m() * loopsMNK.n(); - return taskIdx % GetCoreLoops() / mnLoops; - } - - CATLASS_DEVICE - uint32_t GetCoreLoops() const - { - return loopsMNK.m() * loopsMNK.n() * splitkFactor; - } - - CATLASS_DEVICE - uint32_t GetBatchIdx(uint32_t taskIdx) - { - return taskIdx / GetCoreLoops(); - } - - CATLASS_DEVICE - GemmCoord GetBlockCoord(uint32_t taskIdx) - { - uint32_t splitkSliceIdx = GetSplitkSliceIdx(taskIdx); - uint32_t kIdx = GetKIdxBySplitkSliceIdx(splitkSliceIdx); - - uint32_t innerIdx = taskIdx % (loopsMNK.m() * loopsMNK.n()); - if constexpr (SwizzleDirection == 0) { // Zn - uint32_t tileBlockLoop = CeilDiv(loopsMNK.m(), SwizzleOffset); - uint32_t tileBlockIdx = innerIdx / (SwizzleOffset * loopsMNK.n()); - uint32_t inTileBlockIdx = innerIdx % (SwizzleOffset * loopsMNK.n()); - - uint32_t nRow = SwizzleOffset; - if (tileBlockIdx == tileBlockLoop - 1) { - nRow = loopsMNK.m() - SwizzleOffset * tileBlockIdx; - } - uint32_t mIdx = tileBlockIdx * SwizzleOffset + inTileBlockIdx % nRow; - uint32_t nIdx = inTileBlockIdx / nRow; - if (tileBlockIdx % 2 == 1) { - nIdx = loopsMNK.n() - nIdx - 1; - } - return GemmCoord{mIdx, nIdx, kIdx}; - } else if constexpr (SwizzleDirection == 1) { // Nz - uint32_t tileBlockLoop = CeilDiv(loopsMNK.n(), SwizzleOffset); - uint32_t tileBlockIdx = innerIdx / (SwizzleOffset * loopsMNK.m()); - uint32_t inTileBlockIdx = innerIdx % (SwizzleOffset * loopsMNK.m()); - - uint32_t nCol = SwizzleOffset; - if (tileBlockIdx == tileBlockLoop - 1) { - nCol = loopsMNK.n() - SwizzleOffset * tileBlockIdx; - } - uint32_t mIdx = inTileBlockIdx / nCol; - uint32_t nIdx = tileBlockIdx * SwizzleOffset + inTileBlockIdx % nCol; - if (tileBlockIdx % 2 == 1) { - mIdx = loopsMNK.m() - mIdx - 1; - } - return GemmCoord{mIdx, nIdx, kIdx}; - } - } - - CATLASS_DEVICE - GemmCoord GetActualBlockShape(GemmCoord blockCoord, uint32_t splitkSliceIdx) - { - uint32_t splitkSliceLen; - if (splitkSliceIdx < loopsMNK.k() % splitkFactor) { - splitkSliceLen = (loopsMNK.k() / splitkFactor + 1) * tileShape.k(); - } else { - splitkSliceLen = (loopsMNK.k() / splitkFactor) * tileShape.k(); - } - uint32_t mActual = (blockCoord.m() == (loopsMNK.m() - 1)) ? - (problemShape.m() - blockCoord.m() * tileShape.m()) : - tileShape.m(); - uint32_t nActual = (blockCoord.n() == (loopsMNK.n() - 1)) ? - (problemShape.n() - blockCoord.n() * tileShape.n()) : - tileShape.n(); - uint32_t kActual = (splitkSliceIdx == (splitkFactor - 1)) ? - (problemShape.k() - blockCoord.k() * tileShape.k()) : - splitkSliceLen; - return GemmCoord{mActual, nActual, kActual}; - } -}; - -/// Block swizzleing function for single core splitk matmul -template -struct SingleCoreSplitkGemmIdentityBlockSwizzle { - /// Data members - static_assert(SwizzleDirection == 0 || SwizzleDirection == 1, "SwizzleDirection only support 0/1"); - - GemmCoord problemShape; - GemmCoord tileMNK; - GemmCoord loopsMNK; - uint32_t mnLoopRangeStart; - uint32_t mnLoopRangeLen; - - /// Methods - CATLASS_DEVICE - SingleCoreSplitkGemmIdentityBlockSwizzle() {} - - /// Methods - CATLASS_DEVICE - SingleCoreSplitkGemmIdentityBlockSwizzle(GemmCoord const &problemShape_, GemmCoord const &l1TileShape) - : problemShape(problemShape_), - tileMNK(l1TileShape) - { - loopsMNK = CeilDiv(problemShape, tileMNK); - GetMNLoopRange(AscendC::GetBlockNum(), AscendC::GetBlockIdx()); - } - - CATLASS_DEVICE - uint32_t GetSingleCoreLoops() const - { - return loopsMNK.k() * mnLoopRangeLen; - } - - CATLASS_DEVICE - void GetMNLoopRange(uint32_t coreNum, uint32_t coreIdx) - { - uint32_t mnLoops = loopsMNK.m() * loopsMNK.n(); - uint32_t bPerCore = CeilDiv(mnLoops, coreNum); - // the first tmpNum cores process bPerCore blocks - // and the last (coreNum - tmpNum) cores process (bPerCore - 1) blocks - // total mnLoops = bPerCore * tmpNum + (bPerCore - 1) * (coreNum - tmpNum) - uint32_t tmpNum = mnLoops - coreNum * (bPerCore - 1); - - mnLoopRangeStart = (coreIdx < tmpNum) ? (coreIdx * bPerCore) : (tmpNum * bPerCore + (coreIdx - tmpNum) * (bPerCore - 1)); - mnLoopRangeLen = (coreIdx < tmpNum) ? bPerCore : (bPerCore - 1); - } - - CATLASS_DEVICE - bool IsAtomicAdd(uint32_t loopIdx) - { - uint32_t kIdx = (loopIdx / mnLoopRangeLen); - return kIdx != 0; - } - - CATLASS_DEVICE - GemmCoord GetBlockCoord(uint32_t loopIdx) - { - uint32_t kIdx = (loopIdx / mnLoopRangeLen); - // mnLoopIdx - uint32_t innerIdx = mnLoopRangeStart + loopIdx % mnLoopRangeLen; - if constexpr (SwizzleDirection == 0) { // Zn - uint32_t tileBlockLoop = CeilDiv(loopsMNK.m(), SwizzleOffset); - uint32_t tileBlockIdx = innerIdx / (SwizzleOffset * loopsMNK.n()); - uint32_t inTileBlockIdx = innerIdx % (SwizzleOffset * loopsMNK.n()); - uint32_t tailOffset = loopsMNK.m() - (tileBlockLoop - 1) * SwizzleOffset; - uint32_t nRow = SwizzleOffset; - - if (tileBlockLoop >= 2 && tileBlockIdx >= tileBlockLoop - 2 && tailOffset <= SwizzleOffset / 2) { - tileBlockLoop = (loopsMNK.m() - tailOffset) / SwizzleOffset; - tileBlockIdx = tileBlockLoop - 1; - nRow = SwizzleOffset + tailOffset; - inTileBlockIdx = innerIdx - SwizzleOffset * loopsMNK.n() * tileBlockIdx; - } else if (tileBlockIdx == tileBlockLoop - 1) { - nRow = loopsMNK.m() - SwizzleOffset * tileBlockIdx; - } - - uint32_t nIdx = inTileBlockIdx / nRow; - if (tileBlockIdx % 2 == 1) { - nIdx = loopsMNK.n() - nIdx - 1; - } - - uint32_t mIdx = tileBlockIdx * SwizzleOffset + inTileBlockIdx % nRow; - if (nIdx % 2 == 1) { - mIdx = tileBlockIdx * SwizzleOffset + (nRow - inTileBlockIdx % nRow - 1); - } - return GemmCoord{mIdx, nIdx, kIdx}; - } else if constexpr (SwizzleDirection == 1) { // Nz - uint32_t tileBlockLoop = CeilDiv(loopsMNK.n(), SwizzleOffset); - uint32_t tileBlockIdx = innerIdx / (SwizzleOffset * loopsMNK.m()); - uint32_t inTileBlockIdx = innerIdx % (SwizzleOffset * loopsMNK.m()); - uint32_t tailOffset = loopsMNK.n() - (tileBlockLoop - 1) * SwizzleOffset; - uint32_t nCol = SwizzleOffset; - - if (tileBlockLoop >= 2 && tileBlockIdx >= tileBlockLoop - 2 && tailOffset < SwizzleOffset / 2) { - tileBlockLoop = (loopsMNK.n() - tailOffset) / SwizzleOffset; - tileBlockIdx = tileBlockLoop - 1; - nCol = SwizzleOffset + tailOffset; - inTileBlockIdx = innerIdx - SwizzleOffset * loopsMNK.m() * tileBlockIdx; - } else if (tileBlockIdx == tileBlockLoop - 1) { - nCol = loopsMNK.n() - SwizzleOffset * tileBlockIdx; - } - - uint32_t mIdx = inTileBlockIdx / nCol; - if (tileBlockIdx % 2 == 1) { - mIdx = loopsMNK.m() - mIdx - 1; - } - - uint32_t nIdx = tileBlockIdx * SwizzleOffset + inTileBlockIdx % nCol; - if (mIdx % 2 == 1) { - nIdx = tileBlockIdx * SwizzleOffset + (nCol - inTileBlockIdx % nCol - 1); - } - return GemmCoord{mIdx, nIdx, kIdx}; - } - } - - CATLASS_DEVICE - GemmCoord GetActualBlockShape(GemmCoord &blockCoord) - { - uint32_t mActual = (blockCoord.m() == (loopsMNK.m() - 1)) ? - (problemShape.m() - blockCoord.m() * tileMNK.m()) : - tileMNK.m(); - uint32_t nActual = (blockCoord.n() == (loopsMNK.n() - 1)) ? - (problemShape.n() - blockCoord.n() * tileMNK.n()) : - tileMNK.n(); - uint32_t kActual = (blockCoord.k() == (loopsMNK.k() - 1)) ? - (problemShape.k() - (loopsMNK.k() - 1) * tileMNK.k()) : - tileMNK.k(); - return GemmCoord{mActual, nActual, kActual}; - } -}; - -struct DynamicSingleCoreSplitkGemmIdentityBlockSwizzle : public SingleCoreSplitkGemmIdentityBlockSwizzle<> { - uint32_t swizzleOffset; - uint32_t swizzleDirection; - - CATLASS_DEVICE - DynamicSingleCoreSplitkGemmIdentityBlockSwizzle(GemmCoord const &problemShape_, GemmCoord const &l1TileShape, - uint32_t swizzleOffset_, uint32_t swizzleDirection_) - : swizzleOffset(swizzleOffset_), - swizzleDirection(swizzleDirection_), - SingleCoreSplitkGemmIdentityBlockSwizzle<>(problemShape_, l1TileShape) - {} - - CATLASS_DEVICE - GemmCoord GetBlockCoord(uint32_t loopIdx) - { - uint32_t kIdx = (loopIdx / mnLoopRangeLen); - // mnLoopIdx - uint32_t innerIdx = mnLoopRangeStart + loopIdx % mnLoopRangeLen; - if (swizzleDirection == 0) { // Zn - uint32_t tileBlockLoop = CeilDiv(loopsMNK.m(), swizzleOffset); - uint32_t tileBlockIdx = innerIdx / (swizzleOffset * loopsMNK.n()); - uint32_t inTileBlockIdx = innerIdx % (swizzleOffset * loopsMNK.n()); - uint32_t tailOffset = loopsMNK.m() - (tileBlockLoop - 1) * swizzleOffset; - uint32_t nRow = swizzleOffset; - - if (tileBlockLoop >= 2 && tileBlockIdx >= tileBlockLoop - 2 && tailOffset <= swizzleOffset / 2) { - tileBlockLoop = (loopsMNK.m() - tailOffset) / swizzleOffset; - tileBlockIdx = tileBlockLoop - 1; - nRow = swizzleOffset + tailOffset; - inTileBlockIdx = innerIdx - swizzleOffset * loopsMNK.n() * tileBlockIdx; - } else if (tileBlockIdx == tileBlockLoop - 1) { - nRow = loopsMNK.m() - swizzleOffset * tileBlockIdx; - } - - uint32_t nIdx = inTileBlockIdx / nRow; - if (tileBlockIdx % 2 == 1) { - nIdx = loopsMNK.n() - nIdx - 1; - } - - uint32_t mIdx = tileBlockIdx * swizzleOffset + inTileBlockIdx % nRow; - if (nIdx % 2 == 1) { - mIdx = tileBlockIdx * swizzleOffset + (nRow - inTileBlockIdx % nRow - 1); - } - return GemmCoord{mIdx, nIdx, kIdx}; - } else { // Nz - uint32_t tileBlockLoop = CeilDiv(loopsMNK.n(), swizzleOffset); - uint32_t tileBlockIdx = innerIdx / (swizzleOffset * loopsMNK.m()); - uint32_t inTileBlockIdx = innerIdx % (swizzleOffset * loopsMNK.m()); - uint32_t tailOffset = loopsMNK.n() - (tileBlockLoop - 1) * swizzleOffset; - uint32_t nCol = swizzleOffset; - - if (tileBlockLoop >= 2 && tileBlockIdx >= tileBlockLoop - 2 && tailOffset < swizzleOffset / 2) { - tileBlockLoop = (loopsMNK.n() - tailOffset) / swizzleOffset; - tileBlockIdx = tileBlockLoop - 1; - nCol = swizzleOffset + tailOffset; - inTileBlockIdx = innerIdx - swizzleOffset * loopsMNK.m() * tileBlockIdx; - } else if (tileBlockIdx == tileBlockLoop - 1) { - nCol = loopsMNK.n() - swizzleOffset * tileBlockIdx; - } - - uint32_t mIdx = inTileBlockIdx / nCol; - if (tileBlockIdx % 2 == 1) { - mIdx = loopsMNK.m() - mIdx - 1; - } - - uint32_t nIdx = tileBlockIdx * swizzleOffset + inTileBlockIdx % nCol; - if (mIdx % 2 == 1) { - nIdx = tileBlockIdx * swizzleOffset + (nCol - inTileBlockIdx % nCol - 1); - } - return GemmCoord{mIdx, nIdx, kIdx}; - } - } -}; - -/// Block swizzling function for Gemms -template -struct GemmIdentityBlockSwizzleL1FullLoad { - /// Data members - - GemmCoord problemShape; - MatrixCoord tileMN; - MatrixCoord loopsMN; - - uint32_t loopsPerCore; - uint32_t loopsTail; - uint32_t aicCoreNum; - - /// Methods - - CATLASS_DEVICE - GemmIdentityBlockSwizzleL1FullLoad() {} - - CATLASS_DEVICE - GemmIdentityBlockSwizzleL1FullLoad(GemmCoord const &problemShape_, MatrixCoord const &tileMN_) - : problemShape(problemShape_), - tileMN(tileMN_) - { - loopsMN = CeilDiv(MatrixCoord(problemShape.GetCoordMN()), tileMN); - uint32_t loopsTotalNum = GetCoreLoops(); - aicCoreNum = AscendC::GetBlockNum(); - loopsPerCore = loopsTotalNum / aicCoreNum; - loopsTail = loopsTotalNum % aicCoreNum; - } - - CATLASS_DEVICE - GemmIdentityBlockSwizzleL1FullLoad(GemmCoord const &problemShape_, MatrixCoord const &tileMN_, - MatrixCoord const &loopsMN_) - : problemShape(problemShape_), - tileMN(tileMN_), - loopsMN(loopsMN_) - { - uint32_t loopsTotalNum = GetCoreLoops(); - aicCoreNum = AscendC::GetBlockNum(); - loopsPerCore = loopsTotalNum / aicCoreNum; - loopsTail = loopsTotalNum % aicCoreNum; - } - - CATLASS_DEVICE - void Update(GemmCoord const &problemShape_, MatrixCoord const &tileMN_) - { - problemShape = problemShape_; - tileMN = tileMN_; - loopsMN = CeilDiv(MatrixCoord(problemShape.GetCoordMN()), tileMN); - - uint32_t loopsTotalNum = GetCoreLoops(); - aicCoreNum = AscendC::GetBlockNum(); - loopsPerCore = loopsTotalNum / aicCoreNum; - loopsTail = loopsTotalNum % aicCoreNum; - } - - CATLASS_DEVICE - void Update(GemmCoord const &problemShape_, MatrixCoord const &tileMN_, MatrixCoord const &loopsMN_) - { - problemShape = problemShape_; - tileMN = tileMN_; - loopsMN = loopsMN_; - - uint32_t loopsTotalNum = GetCoreLoops(); - aicCoreNum = AscendC::GetBlockNum(); - loopsPerCore = loopsTotalNum / aicCoreNum; - loopsTail = loopsTotalNum % aicCoreNum; - } - - CATLASS_DEVICE - uint32_t GetCoreLoops() const - { - return loopsMN.row() * loopsMN.column(); - } - - ////// WARNING: current strategy not support GetBatchIdx() - - CATLASS_DEVICE - GemmCoord GetBlockCoord(uint32_t taskIdx) - { - // calculate innerIdx from taskIdx - uint32_t CoreIdx = taskIdx % aicCoreNum; - uint32_t innerCoreIdx = taskIdx / aicCoreNum; - uint32_t innerIdx = CoreIdx * loopsPerCore + innerCoreIdx; - if (CoreIdx < loopsTail) { - innerIdx += CoreIdx; - } else { - innerIdx += loopsTail; - } - // calculate block location in swizzle, using innerIdx - if constexpr (SwizzleDirection == 0) { // Zn - uint32_t tileBlockLoop = CeilDiv(loopsMN.row(), SwizzleOffset); - uint32_t tileBlockIdx = innerIdx / (SwizzleOffset * loopsMN.column()); - uint32_t inTileBlockIdx = innerIdx % (SwizzleOffset * loopsMN.column()); - - uint32_t nRow = SwizzleOffset; - if (tileBlockIdx == tileBlockLoop - 1) { - nRow = loopsMN.row() - SwizzleOffset * tileBlockIdx; - } - uint32_t mIdx = tileBlockIdx * SwizzleOffset + inTileBlockIdx % nRow; - uint32_t nIdx = inTileBlockIdx / nRow; - if (tileBlockIdx % 2 == 1) { - nIdx = loopsMN.column() - nIdx - 1; - } - return GemmCoord{mIdx, nIdx, 0}; - } else if constexpr (SwizzleDirection == 1) { // Nz - uint32_t tileBlockLoop = CeilDiv(loopsMN.column(), SwizzleOffset); - uint32_t tileBlockIdx = innerIdx / (SwizzleOffset * loopsMN.row()); - uint32_t inTileBlockIdx = innerIdx % (SwizzleOffset * loopsMN.row()); - - uint32_t nCol = SwizzleOffset; - if (tileBlockIdx == tileBlockLoop - 1) { - nCol = loopsMN.column() - SwizzleOffset * tileBlockIdx; - } - uint32_t mIdx = inTileBlockIdx / nCol; - uint32_t nIdx = tileBlockIdx * SwizzleOffset + inTileBlockIdx % nCol; - if (tileBlockIdx % 2 == 1) { - mIdx = loopsMN.row() - mIdx - 1; - } - return GemmCoord{mIdx, nIdx, 0}; - } - } - - CATLASS_DEVICE - GemmCoord GetActualBlockShape(GemmCoord blockCoord) - { - uint32_t mActual = (blockCoord.m() == (loopsMN.row() - 1)) ? - (problemShape.m() - blockCoord.m() * tileMN.row()) : - tileMN.row(); - uint32_t nActual = (blockCoord.n() == (loopsMN.column() - 1)) ? - (problemShape.n() - blockCoord.n() * tileMN.column()) : - tileMN.column(); - uint32_t kActual = problemShape.k(); - return GemmCoord{mActual, nActual, kActual}; - } -}; - -struct DynamicSplitkGemmIdentityBlockSwizzle : public SplitkGemmIdentityBlockSwizzle<> { - uint32_t swizzleOffset{1}; - uint32_t swizzleDirection{0}; - - CATLASS_DEVICE - DynamicSplitkGemmIdentityBlockSwizzle(GemmCoord const &problemShape_, GemmCoord const &tileMNK_, - uint32_t splitkFactor_, uint32_t swizzleOffset_, uint32_t swizzleDirection_) - : swizzleOffset(swizzleOffset_), - swizzleDirection(swizzleDirection_), - SplitkGemmIdentityBlockSwizzle<>(problemShape_, tileMNK_, splitkFactor_) - {} - - CATLASS_DEVICE - DynamicSplitkGemmIdentityBlockSwizzle(GemmCoord const &problemShape_, GemmCoord const &tileMNK_, - uint32_t splitkFactor_) - : SplitkGemmIdentityBlockSwizzle<>(problemShape_, tileMNK_, splitkFactor_) - {} - - CATLASS_DEVICE - DynamicSplitkGemmIdentityBlockSwizzle() {} - - CATLASS_DEVICE - void SetSwizzleParams(uint32_t swizzleOffset_, uint32_t swizzleDirection_) - { - swizzleOffset = swizzleOffset_; - swizzleDirection = swizzleDirection_; - } - - CATLASS_DEVICE - GemmCoord GetBlockCoord(uint32_t taskIdx) - { - uint32_t splitkSliceIdx = GetSplitkSliceIdx(taskIdx); - uint32_t kIdx = GetKIdxBySplitkSliceIdx(splitkSliceIdx); - - uint32_t innerIdx = taskIdx % (loopsMNK.m() * loopsMNK.n()); - if (swizzleDirection == 0) { // Zn - uint32_t tileBlockLoop = CeilDiv(loopsMNK.m(), swizzleOffset); - uint32_t tileBlockIdx = innerIdx / (swizzleOffset * loopsMNK.n()); - uint32_t inTileBlockIdx = innerIdx % (swizzleOffset * loopsMNK.n()); - - uint32_t nRow = swizzleOffset; - if (tileBlockIdx == tileBlockLoop - 1) { - nRow = loopsMNK.m() - swizzleOffset * tileBlockIdx; - } - uint32_t mIdx = tileBlockIdx * swizzleOffset + inTileBlockIdx % nRow; - uint32_t nIdx = inTileBlockIdx / nRow; - if (tileBlockIdx % 2 == 1) { - nIdx = loopsMNK.n() - nIdx - 1; - } - return GemmCoord{mIdx, nIdx, kIdx}; - } else { // Nz - uint32_t tileBlockLoop = CeilDiv(loopsMNK.n(), swizzleOffset); - uint32_t tileBlockIdx = innerIdx / (swizzleOffset * loopsMNK.m()); - uint32_t inTileBlockIdx = innerIdx % (swizzleOffset * loopsMNK.m()); - - uint32_t nCol = swizzleOffset; - if (tileBlockIdx == tileBlockLoop - 1) { - nCol = loopsMNK.n() - swizzleOffset * tileBlockIdx; - } - uint32_t mIdx = inTileBlockIdx / nCol; - uint32_t nIdx = tileBlockIdx * swizzleOffset + inTileBlockIdx % nCol; - if (tileBlockIdx % 2 == 1) { - mIdx = loopsMNK.m() - mIdx - 1; - } - return GemmCoord{mIdx, nIdx, kIdx}; - } - } -}; - -/// Block swizzling function for streamk Gemms -template -struct StreamkGemmIdentityBlockSwizzle { - /// Data members - - GemmCoord problemShape; - GemmCoord tileMNK; - GemmCoord loopsMNK; - uint32_t blockDim; - // Represents the number of ordinary blocks. - uint32_t normalBlocks; - // Represents the number of blocks that need to be split according to the Streamk algorithm. - uint32_t streamkBlocks; - // If computing streamk blocks, how many kTile should each core process - uint32_t kTileNumPerCore; - // Since the number of kTiles may not be evenly divisible by the number of cores, - // some cores need to compute one additoinal kTile compared to others - uint32_t kTileRemain; - - struct StreamkBlockDec { - // If the current block is a ordinary block, set it to false. - // If it is a streamk block, set it to true. - bool isStreamkBlock{false}; - // If the current block is a cross block-meaning it consists partly of the tail k portion of the - // current basic block and partly of the head k portion of the next basic block, set it to true. - bool isCrossBlock{false}; - - GemmCoord blockCoord; - GemmCoord actualBlockShape; - GemmCoord streamkBlockCoord; - GemmCoord streamkActualBlockShape; - }; - - /// Methods - - CATLASS_DEVICE - StreamkGemmIdentityBlockSwizzle() {} - - CATLASS_DEVICE - StreamkGemmIdentityBlockSwizzle( - GemmCoord const &problemShape_, GemmCoord const &tileMNK_, uint32_t blockDim_) - : problemShape(problemShape_), - tileMNK(tileMNK_), - blockDim(blockDim_) - { - loopsMNK = CeilDiv(problemShape, tileMNK); - streamkBlocks = loopsMNK.m() * loopsMNK.n() % blockDim; - normalBlocks = loopsMNK.m() * loopsMNK.n() - streamkBlocks; - kTileNumPerCore = streamkBlocks * loopsMNK.k() / blockDim; - kTileRemain = streamkBlocks * loopsMNK.k() % blockDim; - } - - CATLASS_DEVICE - uint32_t GetCoreLoops() const - { - uint32_t blocksNum = loopsMNK.m() * loopsMNK.n(); - uint32_t coreLoops = blocksNum / blockDim * blockDim + min(streamkBlocks * loopsMNK.k(), blockDim); - return coreLoops; - } - - CATLASS_DEVICE - GemmCoord GetBlockCoord(uint32_t taskIdx) - { - uint32_t innerIdx = taskIdx % GetCoreLoops(); - if constexpr (SwizzleDirection == 0) { // Zn - uint32_t tileBlockLoop = CeilDiv(loopsMNK.m(), SwizzleOffset); - uint32_t tileBlockIdx = innerIdx / (SwizzleOffset * loopsMNK.n()); - uint32_t inTileBlockIdx = innerIdx % (SwizzleOffset * loopsMNK.n()); - - uint32_t nRow = SwizzleOffset; - if (tileBlockIdx == tileBlockLoop - 1) { - nRow = loopsMNK.m() - SwizzleOffset * tileBlockIdx; - } - uint32_t mIdx = tileBlockIdx * SwizzleOffset + inTileBlockIdx % nRow; - uint32_t nIdx = inTileBlockIdx / nRow; - if (tileBlockIdx % 2 == 1) { - nIdx = loopsMNK.n() - nIdx - 1; - } - return GemmCoord{mIdx, nIdx, 0}; - } else if constexpr (SwizzleDirection == 1) { // Nz - uint32_t tileBlockLoop = CeilDiv(loopsMNK.n(), SwizzleOffset); - uint32_t tileBlockIdx = innerIdx / (SwizzleOffset * loopsMNK.m()); - uint32_t inTileBlockIdx = innerIdx % (SwizzleOffset * loopsMNK.m()); - - uint32_t nCol = SwizzleOffset; - if (tileBlockIdx == tileBlockLoop - 1) { - nCol = loopsMNK.n() - SwizzleOffset * tileBlockIdx; - } - uint32_t mIdx = inTileBlockIdx / nCol; - uint32_t nIdx = tileBlockIdx * SwizzleOffset + inTileBlockIdx % nCol; - if (tileBlockIdx % 2 == 1) { - mIdx = loopsMNK.m() - mIdx - 1; - } - return GemmCoord{mIdx, nIdx, 0}; - } - } - - CATLASS_DEVICE - GemmCoord GetActualBlockShape(GemmCoord blockCoord) - { - uint32_t mActual = - (blockCoord.m() == (loopsMNK.m() - 1)) ? (problemShape.m() - blockCoord.m() * tileMNK.m()) : tileMNK.m(); - uint32_t nActual = - (blockCoord.n() == (loopsMNK.n() - 1)) ? (problemShape.n() - blockCoord.n() * tileMNK.n()) : tileMNK.n(); - uint32_t kActual = problemShape.k(); - return GemmCoord{mActual, nActual, kActual}; - } - - CATLASS_DEVICE - uint32_t GetStreamkBlockNum() { return streamkBlocks; } - - CATLASS_DEVICE - uint32_t GetNormalBlockNum() { return normalBlocks; } - - CATLASS_DEVICE - void GetStreamkBlockDec(uint32_t loopIdx, StreamkBlockDec &streamkBlockDec) - { - if (loopIdx < normalBlocks) { - streamkBlockDec.isStreamkBlock = false; - streamkBlockDec.isCrossBlock = false; - streamkBlockDec.blockCoord = GetBlockCoord(loopIdx); - streamkBlockDec.actualBlockShape = GetActualBlockShape(streamkBlockDec.blockCoord); - } else { - streamkBlockDec.isStreamkBlock = true; - uint32_t curKTileNum = kTileNumPerCore; - uint32_t kTileIdx = ((loopIdx - normalBlocks) * kTileNumPerCore + kTileRemain); - if (loopIdx - normalBlocks < kTileRemain) { - curKTileNum += 1; - kTileIdx = ((loopIdx - normalBlocks) * curKTileNum); - } - uint32_t streamkBlockIdx = kTileIdx / loopsMNK.k(); - GemmCoord blockCoordTmp = GetBlockCoord(normalBlocks + streamkBlockIdx); - streamkBlockDec.blockCoord = GemmCoord{blockCoordTmp.m(), blockCoordTmp.n(), kTileIdx % loopsMNK.k()}; - - uint32_t nextStreamkBlockIdx = (kTileIdx + curKTileNum) / loopsMNK.k(); - streamkBlockDec.isCrossBlock = false; - GemmCoord actualBlockShapeTmp = GetActualBlockShape(streamkBlockDec.blockCoord); - if ((kTileIdx % loopsMNK.k() + curKTileNum) * tileMNK.k() > problemShape.k()) { - streamkBlockDec.actualBlockShape = GemmCoord{actualBlockShapeTmp.m(), actualBlockShapeTmp.n(), - problemShape.k() - (kTileIdx % loopsMNK.k()) * tileMNK.k()}; - } else { - streamkBlockDec.actualBlockShape = - GemmCoord{actualBlockShapeTmp.m(), actualBlockShapeTmp.n(), curKTileNum * tileMNK.k()}; - } - if (kTileIdx % loopsMNK.k() + curKTileNum > loopsMNK.k()) { - streamkBlockDec.isCrossBlock = true; - streamkBlockDec.streamkBlockCoord = GetBlockCoord(normalBlocks + nextStreamkBlockIdx); - GemmCoord streamkActualBlockShapeTmp = GetActualBlockShape(streamkBlockDec.streamkBlockCoord); - streamkBlockDec.streamkActualBlockShape = GemmCoord{streamkActualBlockShapeTmp.m(), - streamkActualBlockShapeTmp.n(), ((kTileIdx + curKTileNum) % loopsMNK.k()) * tileMNK.k()}; - } - } - } - - CATLASS_DEVICE - uint32_t GetCoreIdx(uint32_t streamkBlockIdx) - { - if (streamkBlockIdx * loopsMNK.k() > kTileRemain * (kTileNumPerCore + 1)) { - return kTileRemain + - (streamkBlockIdx * loopsMNK.k() - kTileRemain * (kTileNumPerCore + 1)) / kTileNumPerCore; - } else { - return streamkBlockIdx * loopsMNK.k() / (kTileNumPerCore + 1); - } - } - - CATLASS_DEVICE - bool IsCross(uint32_t streamkBlockIdx) - { - // If the head part of the current block and the tail part of the previous block are computed by - // the same core, return true. - if (streamkBlockIdx * loopsMNK.k() > kTileRemain * (kTileNumPerCore + 1)) { - if ((streamkBlockIdx * loopsMNK.k() - kTileRemain * (kTileNumPerCore + 1)) % kTileNumPerCore == 0) { - return false; - } else { - return true; - } - } else { - if ((streamkBlockIdx * loopsMNK.k()) % (kTileNumPerCore + 1) == 0) { - return false; - } else { - return true; - } - } - } -}; - -struct DynamicStreamkGemmIdentityBlockSwizzle : public StreamkGemmIdentityBlockSwizzle<> { - uint32_t swizzleOffset{1}; - uint32_t swizzleDirection{0}; - - CATLASS_DEVICE - DynamicStreamkGemmIdentityBlockSwizzle(GemmCoord const &problemShape_, GemmCoord const &tileShape_, - uint32_t coreNum_, uint32_t swizzleOffset_, uint32_t swizzleDirection_) - : swizzleOffset(swizzleOffset_), - swizzleDirection(swizzleDirection_), - StreamkGemmIdentityBlockSwizzle<>(problemShape_, tileShape_, coreNum_) - {} - - CATLASS_DEVICE - DynamicStreamkGemmIdentityBlockSwizzle(GemmCoord const &problemShape_, GemmCoord const &tileShape_, - uint32_t coreNum_) - : StreamkGemmIdentityBlockSwizzle<>(problemShape_, tileShape_, coreNum_) - {} - - CATLASS_DEVICE - DynamicStreamkGemmIdentityBlockSwizzle() {} - - CATLASS_DEVICE - void SetSwizzleParams(uint32_t swizzleOffset_, uint32_t swizzleDirection_) - { - swizzleOffset = swizzleOffset_; - swizzleDirection = swizzleDirection_; - } - - CATLASS_DEVICE - void GetStreamkBlockDec(uint32_t loopIdx, StreamkBlockDec &streamkBlockDec) - { - if (loopIdx < normalBlocks) { - streamkBlockDec.isStreamkBlock = false; - streamkBlockDec.isCrossBlock = false; - streamkBlockDec.blockCoord = GetBlockCoord(loopIdx); - streamkBlockDec.actualBlockShape = GetActualBlockShape(streamkBlockDec.blockCoord); - } else { - streamkBlockDec.isStreamkBlock = true; - uint32_t curKTileNum = kTileNumPerCore; - uint32_t kTileIdx = ((loopIdx - normalBlocks) * kTileNumPerCore + kTileRemain); - if (loopIdx - normalBlocks < kTileRemain) { - curKTileNum += 1; - kTileIdx = ((loopIdx - normalBlocks) * curKTileNum); - } - uint32_t streamkBlockIdx = kTileIdx / loopsMNK.k(); - GemmCoord blockCoordTmp = GetBlockCoord(normalBlocks + streamkBlockIdx); - streamkBlockDec.blockCoord = GemmCoord{blockCoordTmp.m(), blockCoordTmp.n(), kTileIdx % loopsMNK.k()}; - - uint32_t nextStreamkBlockIdx = (kTileIdx + curKTileNum) / loopsMNK.k(); - streamkBlockDec.isCrossBlock = false; - GemmCoord actualBlockShapeTmp = GetActualBlockShape(streamkBlockDec.blockCoord); - if ((kTileIdx % loopsMNK.k() + curKTileNum) * tileMNK.k() > problemShape.k()) { - streamkBlockDec.actualBlockShape = GemmCoord{actualBlockShapeTmp.m(), actualBlockShapeTmp.n(), - problemShape.k() - (kTileIdx % loopsMNK.k()) * tileMNK.k()}; - } else { - streamkBlockDec.actualBlockShape = - GemmCoord{actualBlockShapeTmp.m(), actualBlockShapeTmp.n(), curKTileNum * tileMNK.k()}; - } - if (kTileIdx % loopsMNK.k() + curKTileNum > loopsMNK.k()) { - streamkBlockDec.isCrossBlock = true; - streamkBlockDec.streamkBlockCoord = GetBlockCoord(normalBlocks + nextStreamkBlockIdx); - GemmCoord streamkActualBlockShapeTmp = GetActualBlockShape(streamkBlockDec.streamkBlockCoord); - streamkBlockDec.streamkActualBlockShape = GemmCoord{streamkActualBlockShapeTmp.m(), - streamkActualBlockShapeTmp.n(), ((kTileIdx + curKTileNum) % loopsMNK.k()) * tileMNK.k()}; - } - } - } - - CATLASS_DEVICE - GemmCoord GetBlockCoord(uint32_t taskIdx) - { - uint32_t innerIdx = taskIdx % GetCoreLoops(); - if (swizzleDirection == 0) { // Zn - uint32_t tileBlockLoop = CeilDiv(loopsMNK.m(), swizzleOffset); - uint32_t tileBlockIdx = innerIdx / (swizzleOffset * loopsMNK.n()); - uint32_t inTileBlockIdx = innerIdx % (swizzleOffset * loopsMNK.n()); - - uint32_t nRow = swizzleOffset; - if (tileBlockIdx == tileBlockLoop - 1) { - nRow = loopsMNK.m() - swizzleOffset * tileBlockIdx; - } - uint32_t mIdx = tileBlockIdx * swizzleOffset + inTileBlockIdx % nRow; - uint32_t nIdx = inTileBlockIdx / nRow; - if (tileBlockIdx % 2 == 1) { - nIdx = loopsMNK.n() - nIdx - 1; - } - return GemmCoord{mIdx, nIdx, 0}; - } else { // Nz - uint32_t tileBlockLoop = CeilDiv(loopsMNK.n(), swizzleOffset); - uint32_t tileBlockIdx = innerIdx / (swizzleOffset * loopsMNK.m()); - uint32_t inTileBlockIdx = innerIdx % (swizzleOffset * loopsMNK.m()); - - uint32_t nCol = swizzleOffset; - if (tileBlockIdx == tileBlockLoop - 1) { - nCol = loopsMNK.n() - swizzleOffset * tileBlockIdx; - } - uint32_t mIdx = inTileBlockIdx / nCol; - uint32_t nIdx = tileBlockIdx * swizzleOffset + inTileBlockIdx % nCol; - if (tileBlockIdx % 2 == 1) { - mIdx = loopsMNK.m() - mIdx - 1; - } - return GemmCoord{mIdx, nIdx, 0}; - } - } -}; - -} // namespace Catlass::Gemm::Block - -#endif // CATLASS_KDA_BLOCK_SWIZZLE_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_tile_muls.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_tile_muls.hpp deleted file mode 100644 index bf4c313aa7..0000000000 --- a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_tile_muls.hpp +++ /dev/null @@ -1,49 +0,0 @@ -/** - * Copyright (c) 2025 Huawei Technologies Co., Ltd. - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of - * CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ -#ifndef CATLASS_KDA_TILE_MULS_HPP -#define CATLASS_KDA_TILE_MULS_HPP - -#include "catlass/kda_catlass.hpp" -#include "catlass/gemm/kda_gemm_helper.hpp" -#include "tla/kda_tensor.hpp" -namespace Catlass::Gemm::Tile { - -/////////////////////////////////////////////////////// -template -struct TileMuls { - using ArchTag = ArchTag_; - using Element = typename ComputeType_::Element; - - static constexpr uint32_t COMPUTE_LENGTH = COMPUTE_LENGTH_; - - CATLASS_DEVICE - TileMuls() - { - } - - CATLASS_DEVICE - void operator()( - AscendC::LocalTensor dstTensor, - AscendC::LocalTensor srcTensor, - Element scalar, - uint32_t len) - { - AscendC::SetMaskCount(); - AscendC::SetVectorMask(len); - AscendC::Muls( - dstTensor, srcTensor, scalar, AscendC::MASK_PLACEHOLDER, 1, AscendC::UnaryRepeatParams{}); - AscendC::SetMaskNorm(); - AscendC::ResetMask(); - } -}; -/////////////////////////////////////////////////////// -} // namespace Catlass::Gemm::Tile - -#endif // CATLASS_KDA_TILE_MULS_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_tile_traits.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_tile_traits.hpp deleted file mode 100644 index 8a1c3017ad..0000000000 --- a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_tile_traits.hpp +++ /dev/null @@ -1,48 +0,0 @@ -/** - * Copyright (c) 2025 Huawei Technologies Co., Ltd. - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of - * CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ - -#ifndef CATLASS_KDA_TILE_TRAITS_HPP -#define CATLASS_KDA_TILE_TRAITS_HPP - -#include "catlass/kda_catlass.hpp" - -namespace Catlass::Gemm::Tile { - -struct EmptyType {}; - -template -struct PrologueTraits : public Prologue { - using Prologue::Prologue; - - using TensorSrc = AscendC::GlobalTensor; - using TensorDst = AscendC::GlobalTensor; -}; - -template <> -struct PrologueTraits { - using ElementSrc = EmptyType; - using LayoutSrc = EmptyType; - using ElementDst = EmptyType; - using LayoutDst = EmptyType; - - using TensorSrc = EmptyType; - using TensorDst = EmptyType; - - using Params = EmptyType; - - template - CATLASS_DEVICE - PrologueTraits(Args...) - {} -}; - -} // namespace Catlass::Gemm::Tile - -#endif diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemv/kda_gemv_helper.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemv/kda_gemv_helper.hpp deleted file mode 100644 index aef1fb9e18..0000000000 --- a/csrc/chunk_kda_fwd/op_kernel/catlass/gemv/kda_gemv_helper.hpp +++ /dev/null @@ -1,122 +0,0 @@ -/** - * Copyright (c) 2025 Huawei Technologies Co., Ltd. - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of - * CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ - -#ifndef CATLASS_KDA_GEMV_HELPER_HPP -#define CATLASS_KDA_GEMV_HELPER_HPP - -#include "catlass/kda_catlass.hpp" -#include "catlass/layout/kda_layout.hpp" -#include "catlass/gemm/kda_gemm_type.hpp" - -namespace Catlass::Gemv::helper { - -template -struct UBAlignHelper { - static constexpr uint32_t ALIGN = BYTE_PER_BLK / sizeof(Element); -}; - -template -struct AtomicAddSelector { - static_assert(DEPENDENT_FALSE, - "Unsupported layout selector, can not find the specialization."); -}; - -template -struct AtomicAddSelector> { - static constexpr bool value = false; -}; - -template -struct AtomicAddSelector> { - static constexpr bool value = true; -}; - -template -struct L1AlignHelper { - static_assert(DEPENDENT_FALSE, "Unsupported align helper, can not find the specialization."); -}; - -template -struct L1AlignHelper { - static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); - static constexpr uint32_t M_ALIGNED = C0_NUM_PER_FRACTAL; - static constexpr uint32_t N_ALIGNED = ELE_NUM_PER_C0; -}; - -template -struct L1AlignHelper { - static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); - - static constexpr uint32_t getNAligned() - { - if constexpr (std::is_same::value) { - return ELE_NUM_PER_C0 / sizeof(Element); - } else { - return C0_NUM_PER_FRACTAL; - } - } - - static constexpr uint32_t getMAligned() - { - if constexpr (std::is_same::value) { - return ELE_NUM_PER_C0 / sizeof(Element); - } else { - return C0_NUM_PER_FRACTAL; - } - } - - static constexpr uint32_t N_ALIGNED = getNAligned(); - static constexpr uint32_t M_ALIGNED = getMAligned(); -}; - -template -struct L1AlignHelper { - static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); - static constexpr uint32_t M_ALIGNED = C0_NUM_PER_FRACTAL; - static constexpr uint32_t N_ALIGNED = ELE_NUM_PER_C0; -}; - -//////////////////////////////// -// new add gemvaic selector -template -struct L1AndL0TypeSelectorGemv { - static_assert(DEPENDENT_FALSE, - "Unsupported layout selector, can not find the specialization."); - static_assert(DEPENDENT_FALSE, - "Unsupported layout selector, can not find the specialization."); -}; - -template -struct L1AndL0TypeSelectorGemv, Gemm::GemmType> { - using L1AType = Gemm::GemmType; - using L1BType = Gemm::GemmType; - using L0AType = Gemm::GemmType; - using L0BType = Gemm::GemmType; -}; - -template -struct L1AndL0TypeSelectorGemv, Gemm::GemmType> { - using L1AType = Gemm::GemmType; - using L1BType = Gemm::GemmType; - using L0AType = Gemm::GemmType; - using L0BType = Gemm::GemmType; -}; - -template <> -struct L1AndL0TypeSelectorGemv, Gemm::GemmType> { - using L1AType = Gemm::GemmType; - using L1BType = Gemm::GemmType; - using L0AType = Gemm::GemmType; - using L0BType = Gemm::GemmType; -}; - -} // namespace Catlass::Gemv::helper - -#endif // CATLASS_KDA_GEMV_HELPER_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/kda_gemv_coord.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/kda_gemv_coord.hpp deleted file mode 100644 index f209676924..0000000000 --- a/csrc/chunk_kda_fwd/op_kernel/catlass/kda_gemv_coord.hpp +++ /dev/null @@ -1,108 +0,0 @@ -/** - * Copyright (c) 2025 Huawei Technologies Co., Ltd. - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of - * CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ - -#ifndef CATLASS_KDA_GEMV_COORD_HPP -#define CATLASS_KDA_GEMV_COORD_HPP - -#include "catlass/kda_coord.hpp" - -namespace Catlass { - -/// Shape of a matrix multiply-add operation -template < - /// Rows of matrix product - uint32_t M_ = 1, - /// Columns of the matrix (number of elements in the input vector) - uint32_t N_ = 1> -struct GemvShape { - static constexpr uint32_t M = M_; - static constexpr uint32_t N = N_; - - static constexpr int64_t MN = M * N; - - static constexpr int64_t COUNT = MN; - - /// Returns a Coord object - CATLASS_HOST_DEVICE - static Coord<2> ToCoord() - { - return MakeCoord(M, N); - } -}; - -/// GemvCoord is a structure derived from Coord<2> that specifies a location within the -/// coordinate space of a GEMV problem. -struct GemvCoord : public Coord<2, uint32_t> { - /// Integer-valued index - using Index = uint32_t; - - /// Base type is a Coord of rank=2 - using Base = Coord<2, Index>; - - /// GEMV M dimension - rows of the output vector (y) - static constexpr int M_INDEX = 0; - - /// GEMV N dimension - columns of the matrix (length of the input vector x) - static constexpr int N_INDEX = 1; - - /// Default ctor - CATLASS_HOST_DEVICE - GemvCoord() {} - - /// Constructs from Coord<2> and a batch - CATLASS_HOST_DEVICE - GemvCoord(Coord<2, Index> const &coord) - : Base(coord) - {} - - /// Helper to construct from M, N coordinates - CATLASS_HOST_DEVICE - GemvCoord(Index m, Index n) - : Base(MakeCoord(m, n)) - {} - - /// Returns the GEMV M coordinate (row of the result y) - CATLASS_HOST_DEVICE - Index const &m() const - { - return this->At(M_INDEX); - } - - /// Returns reference to the GEMV M coordinate - CATLASS_HOST_DEVICE - Index &m() - { - return this->At(M_INDEX); - } - - /// Returns the GEMV N coordinate (column of the matrix A or the input vector x) - CATLASS_HOST_DEVICE - Index const &n() const - { - return this->At(N_INDEX); - } - - /// Returns reference to the GEMV N coordinate - CATLASS_HOST_DEVICE - Index &n() - { - return this->At(N_INDEX); - } - - CATLASS_HOST_DEVICE - auto GetCoordMN() const - { - return this->GetCoordByAxis(); - } -}; - -} // namespace Catlass - -#endif // CATLASS_KDA_GEMV_COORD_HPP From 772d5c8eda8712a258da3a82efe5363eb371def1 Mon Sep 17 00:00:00 2001 From: chenjz Date: Fri, 4 Sep 2026 09:24:31 +0800 Subject: [PATCH 07/11] refactor(chunk_kda_fwd): drop the dead BETA_T parameter on PostWu/Finalize ChunkKdaFwdPostWuKernel and ChunkKdaFwdFinalizeKernel are templated on the beta element type, but their beta_ tensor is only ever bound in Init() and is never read -- in either the arch35 or the generic implementation: $ grep -n "beta_" arch35/chunk_kda_fwd_finalize_arch35.h 255: beta_.SetGlobalBuffer((__gm__ BETA_T *)beta); 1430: GlobalTensor beta_; So the parameter bought nothing and cost a full second instantiation of each kernel. Removing it also frees RunGenericBackEnd, whose only use of BETA_T was forwarding it to Finalize; that one matters twice over because the function has two overloads (with and without a TPipe) and each inlines the whole FwdH + Finalize body. Instantiations in the A5 translation unit: arch35 PostWu (1959 lines) 2 -> 1 arch35 Finalize (1475 lines) 2 -> 1 RunGenericBackEnd 4 -> 2 The beta GM_ADDR arguments stay in the signatures so no call site changes, and RunFrontEnd keeps BETA_T for now because Prepare still needs it. No behavioural change: the removed state was write-only. Co-Authored-By: Claude Opus 5 --- .../arch35/chunk_kda_fwd_finalize_arch35.h | 10 ++++------ .../op_kernel/arch35/chunk_kda_fwd_impl.h | 4 ++-- .../arch35/chunk_kda_fwd_post_wu_arch35.h | 10 ++++------ .../arch35/chunk_kda_fwd_prepare_arch35.h | 2 +- .../op_kernel/chunk_kda_fwd_common.h | 16 ++++++++-------- .../op_kernel/chunk_kda_fwd_finalize.h | 10 ++++------ .../op_kernel/chunk_kda_fwd_post_wu.h | 10 ++++------ 7 files changed, 27 insertions(+), 35 deletions(-) diff --git a/csrc/chunk_kda_fwd/op_kernel/arch35/chunk_kda_fwd_finalize_arch35.h b/csrc/chunk_kda_fwd/op_kernel/arch35/chunk_kda_fwd_finalize_arch35.h index 39d021f0da..7f0fa9cb2e 100644 --- a/csrc/chunk_kda_fwd/op_kernel/arch35/chunk_kda_fwd_finalize_arch35.h +++ b/csrc/chunk_kda_fwd/op_kernel/arch35/chunk_kda_fwd_finalize_arch35.h @@ -235,7 +235,7 @@ __aicore__ inline T FloatToType(float value) return static_cast(value); } -template +template class ChunkKdaFwdFinalizeKernel { public: using OUT_T = float; @@ -252,7 +252,6 @@ class ChunkKdaFwdFinalizeKernel { k_.SetGlobalBuffer((__gm__ T *)k); v_.SetGlobalBuffer((__gm__ T *)v); gk_.SetGlobalBuffer((__gm__ GK_T *)gk); - beta_.SetGlobalBuffer((__gm__ BETA_T *)beta); if (initialState != nullptr) { initialState_.SetGlobalBuffer((__gm__ float *)initialState); } @@ -1427,7 +1426,6 @@ class ChunkKdaFwdFinalizeKernel { GlobalTensor k_; GlobalTensor v_; GlobalTensor gk_; - GlobalTensor beta_; GlobalTensor initialState_; GlobalTensor o_; GlobalTensor finalState_; @@ -1488,7 +1486,7 @@ class ChunkKdaFwdFinalizeKernel { }; } // namespace -template +template __aicore__ inline void RunChunkKdaOutput(GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR gk, GM_ADDR beta, GM_ADDR initialState, GM_ADDR cuSeqlens, GM_ADDR chunkIndices, GM_ADDR qgScaled, GM_ADDR aqk, GM_ADDR propagatedVNew, GM_ADDR propagatedH, @@ -1500,14 +1498,14 @@ __aicore__ inline void RunChunkKdaOutput(GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADD GM_ADDR stateScratch = outputScratch; GM_ADDR localScratch = outputScratch + outputElements * sizeof(float); if ASCEND_IS_AIC { - ChunkKdaFwdFinalizeKernel op; + ChunkKdaFwdFinalizeKernel op; op.Init(q, k, v, gk, beta, initialState, cuSeqlens, chunkIndices, qgScaled, aqk, propagatedVNew, propagatedH, stateScratch, userWorkspace, aqk, userWorkspace, userWorkspace, localScratch, userWorkspace, userWorkspace, o, propagatedH, outputScratch, tiling, &pipe, false); op.ProcessAic(); } if ASCEND_IS_AIV { - ChunkKdaFwdFinalizeKernel op; + ChunkKdaFwdFinalizeKernel op; op.Init(q, k, v, gk, beta, initialState, cuSeqlens, chunkIndices, qgScaled, aqk, propagatedVNew, propagatedH, stateScratch, userWorkspace, aqk, userWorkspace, userWorkspace, localScratch, userWorkspace, userWorkspace, o, propagatedH, outputScratch, tiling, &pipe); diff --git a/csrc/chunk_kda_fwd/op_kernel/arch35/chunk_kda_fwd_impl.h b/csrc/chunk_kda_fwd/op_kernel/arch35/chunk_kda_fwd_impl.h index dab78aadb6..31cdab9bc6 100644 --- a/csrc/chunk_kda_fwd/op_kernel/arch35/chunk_kda_fwd_impl.h +++ b/csrc/chunk_kda_fwd/op_kernel/arch35/chunk_kda_fwd_impl.h @@ -48,11 +48,11 @@ __aicore__ inline void Run( fwdHTaskCount > tiling.prepareUsedCoreNum; if (isolateGenericBackEnd) { pipe.Destroy(); - RunGenericBackEnd( + RunGenericBackEnd( q, k, v, beta, initialState, cuSeqlens, chunkIndices, aqk, attnOut, addresses, userWorkspace, tiling); } else { - RunGenericBackEnd( + RunGenericBackEnd( q, k, v, beta, initialState, cuSeqlens, chunkIndices, aqk, attnOut, addresses, userWorkspace, tiling, pipe); } diff --git a/csrc/chunk_kda_fwd/op_kernel/arch35/chunk_kda_fwd_post_wu_arch35.h b/csrc/chunk_kda_fwd/op_kernel/arch35/chunk_kda_fwd_post_wu_arch35.h index a27984d2be..5e9eb01125 100644 --- a/csrc/chunk_kda_fwd/op_kernel/arch35/chunk_kda_fwd_post_wu_arch35.h +++ b/csrc/chunk_kda_fwd/op_kernel/arch35/chunk_kda_fwd_post_wu_arch35.h @@ -321,7 +321,7 @@ __aicore__ inline T FloatToType(float value) return static_cast(value); } -template +template class ChunkKdaFwdPostWuKernel { public: using OUT_T = T; @@ -338,7 +338,6 @@ class ChunkKdaFwdPostWuKernel { k_.SetGlobalBuffer((__gm__ T *)k); v_.SetGlobalBuffer((__gm__ T *)v); gk_.SetGlobalBuffer((__gm__ GK_T *)gk); - beta_.SetGlobalBuffer((__gm__ BETA_T *)beta); if (initialState != nullptr) { initialState_.SetGlobalBuffer((__gm__ float *)initialState); } @@ -1950,7 +1949,6 @@ class ChunkKdaFwdPostWuKernel { GlobalTensor k_; GlobalTensor v_; GlobalTensor gk_; - GlobalTensor beta_; GlobalTensor initialState_; GlobalTensor o_; GlobalTensor finalState_; @@ -2010,7 +2008,7 @@ class ChunkKdaFwdPostWuKernel { }; } // namespace -template +template __aicore__ inline void RunChunkKdaPostWu(GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR gk, GM_ADDR beta, GM_ADDR initialState, GM_ADDR cuSeqlens, GM_ADDR chunkIndices, GM_ADDR wSeed, GM_ADDR akk, GM_ADDR uSeed, GM_ADDR w, GM_ADDR u, GM_ADDR kg, GM_ADDR vNew, @@ -2018,14 +2016,14 @@ __aicore__ inline void RunChunkKdaPostWu(GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADD { GM_ADDR postScratch = userWorkspace + tiling.postWuScratchOffset; if ASCEND_IS_AIC { - ChunkKdaFwdPostWuKernel op; + ChunkKdaFwdPostWuKernel op; op.Init(q, k, v, gk, beta, initialState, cuSeqlens, chunkIndices, wSeed, akk, uSeed, nullptr, userWorkspace, userWorkspace, userWorkspace, akk, w, u, userWorkspace, kg, vNew, postScratch, postScratch, tiling, &pipe, false); op.ProcessAic(); } if ASCEND_IS_AIV { - ChunkKdaFwdPostWuKernel op; + ChunkKdaFwdPostWuKernel op; op.Init(q, k, v, gk, beta, initialState, cuSeqlens, chunkIndices, wSeed, akk, uSeed, nullptr, userWorkspace, userWorkspace, userWorkspace, akk, w, u, userWorkspace, kg, vNew, postScratch, postScratch, tiling, &pipe); diff --git a/csrc/chunk_kda_fwd/op_kernel/arch35/chunk_kda_fwd_prepare_arch35.h b/csrc/chunk_kda_fwd/op_kernel/arch35/chunk_kda_fwd_prepare_arch35.h index 852bad8d21..de2d5ce4a9 100644 --- a/csrc/chunk_kda_fwd/op_kernel/arch35/chunk_kda_fwd_prepare_arch35.h +++ b/csrc/chunk_kda_fwd/op_kernel/arch35/chunk_kda_fwd_prepare_arch35.h @@ -4756,7 +4756,7 @@ __aicore__ inline void RunChunkKdaPrepare(GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_AD prepareScratch, tiling, &pipe, false, storeQG); if (tiling.fusePostWu) { #if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 - KdaPostWu::ChunkKdaFwdPostWuKernel postWu; + KdaPostWu::ChunkKdaFwdPostWuKernel postWu; postWu.Init(nullptr, k, nullptr, gk, beta, initialState, cuSeqlens, chunkIndices, wSeed, akk, uSeed, nullptr, userWorkspace, userWorkspace, userWorkspace, akk, wSeed, uSeed, userWorkspace, finalKg, userWorkspace, prepareScratch, prepareScratch, tiling, &pipe, false); diff --git a/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_common.h b/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_common.h index 2a570ab3cb..a6c99f782e 100644 --- a/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_common.h +++ b/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_common.h @@ -220,7 +220,7 @@ __aicore__ inline void RunFrontEnd( pipe.Reset(); if (!tiling.fusePostWu && !tiling.fusePostWuIntoFwdH) { - KdaPostWu::RunChunkKdaPostWu( + KdaPostWu::RunChunkKdaPostWu( q, k, v, addresses.gk, beta, initialState, cuSeqlens, chunkIndices, wSeed, akk, uSeed, addresses.w, addresses.u, addresses.kg, addresses.vNew, @@ -248,7 +248,7 @@ __aicore__ inline void RunFwdH( stateOp.Process(); } -template +template __aicore__ inline void RunGenericBackEnd( GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR beta, GM_ADDR initialState, GM_ADDR cuSeqlens, GM_ADDR chunkIndices, GM_ADDR aqk, GM_ADDR attnOut, @@ -266,13 +266,13 @@ __aicore__ inline void RunGenericBackEnd( } SyncAll(); TPipe pipe; - KdaFinalize::RunChunkKdaOutput( + KdaFinalize::RunChunkKdaOutput( q, k, v, addresses.gk, beta, initialState, cuSeqlens, chunkIndices, addresses.qgScaled, aqk, addresses.vNew, addresses.h, attnOut, userWorkspace, tiling, pipe); } -template +template __aicore__ inline void RunGenericBackEnd( GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR beta, GM_ADDR initialState, GM_ADDR cuSeqlens, GM_ADDR chunkIndices, GM_ADDR aqk, GM_ADDR attnOut, @@ -289,7 +289,7 @@ __aicore__ inline void RunGenericBackEnd( userWorkspace, tiling); } SyncAll(); - KdaFinalize::RunChunkKdaOutput( + KdaFinalize::RunChunkKdaOutput( q, k, v, addresses.gk, beta, initialState, cuSeqlens, chunkIndices, addresses.qgScaled, aqk, addresses.vNew, addresses.h, attnOut, userWorkspace, tiling, pipe); @@ -315,11 +315,11 @@ __aicore__ inline void RunGeneric( chunkIndices, aqk, akk, addresses, userWorkspace, tiling, pipe); if (!tiling.isVarLen && tiling.seqlen % tiling.chunkSize == 0) { pipe.Destroy(); - RunGenericBackEnd( + RunGenericBackEnd( q, k, v, beta, initialState, cuSeqlens, chunkIndices, aqk, attnOut, addresses, userWorkspace, tiling); } else { - RunGenericBackEnd( + RunGenericBackEnd( q, k, v, beta, initialState, cuSeqlens, chunkIndices, aqk, attnOut, addresses, userWorkspace, tiling, pipe); } @@ -331,7 +331,7 @@ __aicore__ inline void RunGeneric( q, k, v, g, beta, aLog, dtBias, initialState, cuSeqlens, chunkIndices, aqk, akk, addresses, userWorkspace, tiling, pipe); } - RunGenericBackEnd( + RunGenericBackEnd( q, k, v, beta, initialState, cuSeqlens, chunkIndices, aqk, attnOut, addresses, userWorkspace, tiling); #endif diff --git a/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_finalize.h b/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_finalize.h index bc99f96057..6654d36203 100644 --- a/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_finalize.h +++ b/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_finalize.h @@ -146,7 +146,7 @@ __aicore__ inline T FloatToType(float value) return static_cast(value); } -template +template class ChunkKdaFwdFinalizeKernel { public: using OUT_T = float; @@ -163,7 +163,6 @@ class ChunkKdaFwdFinalizeKernel { k_.SetGlobalBuffer((__gm__ T *)k); v_.SetGlobalBuffer((__gm__ T *)v); gk_.SetGlobalBuffer((__gm__ GK_T *)gk); - beta_.SetGlobalBuffer((__gm__ BETA_T *)beta); if (initialState != nullptr) { initialState_.SetGlobalBuffer((__gm__ float *)initialState); } @@ -721,7 +720,6 @@ class ChunkKdaFwdFinalizeKernel { GlobalTensor k_; GlobalTensor v_; GlobalTensor gk_; - GlobalTensor beta_; GlobalTensor initialState_; GlobalTensor o_; GlobalTensor finalState_; @@ -780,7 +778,7 @@ class ChunkKdaFwdFinalizeKernel { }; } // namespace -template +template __aicore__ inline void RunChunkKdaOutput(GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR gk, GM_ADDR beta, GM_ADDR initialState, GM_ADDR cuSeqlens, GM_ADDR chunkIndices, GM_ADDR qgScaled, GM_ADDR aqk, GM_ADDR propagatedVNew, GM_ADDR propagatedH, @@ -792,14 +790,14 @@ __aicore__ inline void RunChunkKdaOutput(GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADD GM_ADDR stateScratch = outputScratch; GM_ADDR localScratch = outputScratch + outputElements * sizeof(float); if ASCEND_IS_AIC { - ChunkKdaFwdFinalizeKernel op; + ChunkKdaFwdFinalizeKernel op; op.Init(q, k, v, gk, beta, initialState, cuSeqlens, chunkIndices, qgScaled, aqk, propagatedVNew, propagatedH, stateScratch, userWorkspace, aqk, userWorkspace, userWorkspace, localScratch, userWorkspace, userWorkspace, o, propagatedH, outputScratch, tiling, &pipe, false); op.ProcessAic(); } if ASCEND_IS_AIV { - ChunkKdaFwdFinalizeKernel op; + ChunkKdaFwdFinalizeKernel op; op.Init(q, k, v, gk, beta, initialState, cuSeqlens, chunkIndices, qgScaled, aqk, propagatedVNew, propagatedH, stateScratch, userWorkspace, aqk, userWorkspace, userWorkspace, localScratch, userWorkspace, userWorkspace, o, propagatedH, outputScratch, tiling, &pipe); diff --git a/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_post_wu.h b/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_post_wu.h index b3942a08d4..1146108de2 100644 --- a/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_post_wu.h +++ b/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_post_wu.h @@ -150,7 +150,7 @@ __aicore__ inline T FloatToType(float value) return static_cast(value); } -template +template class ChunkKdaFwdPostWuKernel { public: using OUT_T = T; @@ -167,7 +167,6 @@ class ChunkKdaFwdPostWuKernel { k_.SetGlobalBuffer((__gm__ T *)k); v_.SetGlobalBuffer((__gm__ T *)v); gk_.SetGlobalBuffer((__gm__ GK_T *)gk); - beta_.SetGlobalBuffer((__gm__ BETA_T *)beta); if (initialState != nullptr) { initialState_.SetGlobalBuffer((__gm__ float *)initialState); } @@ -840,7 +839,6 @@ class ChunkKdaFwdPostWuKernel { GlobalTensor k_; GlobalTensor v_; GlobalTensor gk_; - GlobalTensor beta_; GlobalTensor initialState_; GlobalTensor o_; GlobalTensor finalState_; @@ -900,7 +898,7 @@ class ChunkKdaFwdPostWuKernel { }; } // namespace -template +template __aicore__ inline void RunChunkKdaPostWu(GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR gk, GM_ADDR beta, GM_ADDR initialState, GM_ADDR cuSeqlens, GM_ADDR chunkIndices, GM_ADDR wSeed, GM_ADDR akk, GM_ADDR uSeed, GM_ADDR w, GM_ADDR u, GM_ADDR kg, GM_ADDR vNew, @@ -908,14 +906,14 @@ __aicore__ inline void RunChunkKdaPostWu(GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADD { GM_ADDR postScratch = userWorkspace + tiling.postWuScratchOffset; if ASCEND_IS_AIC { - ChunkKdaFwdPostWuKernel op; + ChunkKdaFwdPostWuKernel op; op.Init(q, k, v, gk, beta, initialState, cuSeqlens, chunkIndices, wSeed, akk, uSeed, nullptr, userWorkspace, userWorkspace, userWorkspace, akk, w, u, userWorkspace, kg, vNew, postScratch, postScratch, tiling, &pipe, false); op.ProcessAic(); } if ASCEND_IS_AIV { - ChunkKdaFwdPostWuKernel op; + ChunkKdaFwdPostWuKernel op; op.Init(q, k, v, gk, beta, initialState, cuSeqlens, chunkIndices, wSeed, akk, uSeed, nullptr, userWorkspace, userWorkspace, userWorkspace, akk, w, u, userWorkspace, kg, vNew, postScratch, postScratch, tiling, &pipe); From 2f0684003b317e609157373d6bd38c0cd6508686 Mon Sep 17 00:00:00 2001 From: chenjz Date: Fri, 4 Sep 2026 09:25:02 +0800 Subject: [PATCH 08/11] perf(chunk_kda_fwd): resolve the beta dtype at runtime instead of by template Prepare was templated on BETA_T, so the 4336-line arch35 kernel was instantiated twice over a parameter that only selects one vector cast: if constexpr (IsSameType::value) { Adds(betaFp32, betaTyped, 0.0f, tileRows); } else { Cast(betaFp32, betaTyped, RoundMode::CAST_NONE, tileRows); } Bind beta under both element types and pick between them with a runtime flag taken from tiling.betaDataType -- the same field the kernel entry was already dispatching on. Three call shapes are affected: the five (arch35) / four (generic) LoadAsFloatRow sites become LoadBetaAsFloatRow, the two staged conversions become ConvertStagedBeta, and the fused-path staging copy branches on the flag. Both branches now live in the binary at those sites, which costs roughly 50 lines against a whole 4336-line instantiation. With Prepare independent of the beta dtype, BETA_T drops out of RunFrontEnd, RunGeneric and arch35::Run too, so the entry point's three `if (betaDataType == 0) ... else ...` blocks collapse to single calls. Instantiations of arch35 Prepare in the A5 translation unit: 8 -> 4 (SAFE_GATE x COMPILE_BT/K/V remain). Combined with the previous commit, total weighted instantiated body lines drop from ~41.5k to ~20.8k. The UB staging slot is now always sized for fp32. Checked that this changes nothing observable: with K=128 both the old bf16 (3076 B/row) and fp32 (3080 B/row) layouts already saturate KDA_GATE_TILE_ROWS=16, so the tile row count is unchanged; the bf16-beta path grows by 64 bytes of UB (49216 -> 49280) into the layout the fp32-beta path already used, leaving 14.7 KB of headroom under the 192 KB budget. The surrounding PipeBarrier / SetFlag / WaitFlag sequences are untouched. Note for review: the staging path (ConvertStagedBeta and the fused copy) sits inside `if constexpr (SAFE_GATE && COMPILE_BT == 64 && COMPILE_K == 128 && COMPILE_V == 128)`, which safe_gate=False callers never reach. Accuracy runs covering that combination with bf16 beta are the ones that matter here. Co-Authored-By: Claude Opus 5 --- .../op_kernel/arch35/chunk_kda_fwd_impl.h | 4 +- .../arch35/chunk_kda_fwd_prepare_arch35.h | 84 ++++++++++++------- .../chunk_kda_fwd/op_kernel/chunk_kda_fwd.cpp | 69 +++++---------- .../op_kernel/chunk_kda_fwd_common.h | 10 +-- .../op_kernel/chunk_kda_fwd_prepare.h | 42 +++++++--- 5 files changed, 112 insertions(+), 97 deletions(-) diff --git a/csrc/chunk_kda_fwd/op_kernel/arch35/chunk_kda_fwd_impl.h b/csrc/chunk_kda_fwd/op_kernel/arch35/chunk_kda_fwd_impl.h index 31cdab9bc6..beee398641 100644 --- a/csrc/chunk_kda_fwd/op_kernel/arch35/chunk_kda_fwd_impl.h +++ b/csrc/chunk_kda_fwd/op_kernel/arch35/chunk_kda_fwd_impl.h @@ -15,7 +15,7 @@ namespace KdaForward::arch35 { -template __aicore__ inline void Run( GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR g, GM_ADDR beta, @@ -27,7 +27,7 @@ __aicore__ inline void Run( { const auto addresses = ResolveAddresses( finalState, gk, w, u, qg, kg, vNew, h, userWorkspace, tiling); - RunFrontEnd( q, k, v, g, beta, aLog, dtBias, initialState, cuSeqlens, chunkIndices, aqk, akk, addresses, userWorkspace, tiling, pipe); diff --git a/csrc/chunk_kda_fwd/op_kernel/arch35/chunk_kda_fwd_prepare_arch35.h b/csrc/chunk_kda_fwd/op_kernel/arch35/chunk_kda_fwd_prepare_arch35.h index de2d5ce4a9..bcb92bd14a 100644 --- a/csrc/chunk_kda_fwd/op_kernel/arch35/chunk_kda_fwd_prepare_arch35.h +++ b/csrc/chunk_kda_fwd/op_kernel/arch35/chunk_kda_fwd_prepare_arch35.h @@ -121,6 +121,9 @@ constexpr uint32_t KDA_SCORE_SCRATCH_KG = 2; constexpr uint64_t KDA_WORKSPACE_ALIGN = 512; constexpr uint32_t KDA_GATE_TILE_ROWS = 16; constexpr uint32_t KDA_GATE_PIPELINE_DEPTH = 3; +// beta is staged in UB as fp32 or bf16; the slot is always sized for the wider one +// so the UB layout does not depend on the beta dtype. +constexpr uint64_t KDA_BETA_STAGE_BYTES = sizeof(float); constexpr uint32_t KDA_AIV_UB_BUDGET_BYTES = 192 * 1024; constexpr uint32_t KDA_LOCAL_GK_FLOAT_OFFSET = 10 * 1024; constexpr uint32_t KDA_SCALED_QG_FLOAT_OFFSET = 18 * 1024; @@ -777,7 +780,7 @@ __aicore__ inline T FloatToType(float value) return static_cast(value); } -template class ChunkKdaFwdPrepareKernel { public: @@ -800,7 +803,8 @@ class ChunkKdaFwdPrepareKernel { rawG_.SetGlobalBuffer((__gm__ float *)rawG); aLog_.SetGlobalBuffer((__gm__ float *)aLog); dtBias_.SetGlobalBuffer((__gm__ float *)dtBias); - beta_.SetGlobalBuffer((__gm__ BETA_T *)beta); + betaFp32_.SetGlobalBuffer((__gm__ float *)beta); + betaBf16_.SetGlobalBuffer((__gm__ bfloat16_t *)beta); if (initialState != nullptr) { initialState_.SetGlobalBuffer((__gm__ float *)initialState); } @@ -851,6 +855,7 @@ class ChunkKdaFwdPrepareKernel { lowerBound_ = tiling.lowerBound; storeQG_ = storeQG; usedCoreNum_ = tiling.prepareUsedCoreNum; + betaIsFp32_ = tiling.betaDataType == 0; #if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 if constexpr (SAFE_GATE && !IsSameType::value) { headPairMode_ = KDA_ARCH35_ENABLE_HEAD_PAIR && HV_ % KDA_SCORE_LANES == 0; @@ -1143,7 +1148,7 @@ class ChunkKdaFwdPrepareKernel { uint64_t bytesPerRow = K_ * KDA_GATE_PIPELINE_DEPTH * (3 * sizeof(T) + sizeof(GK_T)); #if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 if constexpr (SAFE_GATE && COMPILE_BT == 64 && COMPILE_K == 128 && COMPILE_V == 128) { - bytesPerRow = GateBufferDepth() * (K_ * (4 * sizeof(T) + sizeof(GK_T)) + sizeof(BETA_T)); + bytesPerRow = GateBufferDepth() * (K_ * (4 * sizeof(T) + sizeof(GK_T)) + KDA_BETA_STAGE_BYTES); } #endif uint64_t rows = bytesPerRow == 0 ? 0 : availableBytes / bytesPerRow; @@ -1164,7 +1169,7 @@ class ChunkKdaFwdPrepareKernel { { #if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 if constexpr (SAFE_GATE && COMPILE_BT == 64 && COMPILE_K == 128 && COMPILE_V == 128) { - return GateStageElems() * (3 * sizeof(T) + sizeof(GK_T)) + GatePipelineRows() * sizeof(BETA_T); + return GateStageElems() * (3 * sizeof(T) + sizeof(GK_T)) + GatePipelineRows() * KDA_BETA_STAGE_BYTES; } #endif return GateStageElems() * (2 * sizeof(T) + sizeof(GK_T)); @@ -1194,10 +1199,11 @@ class ChunkKdaFwdPrepareKernel { return gateWritebackBuf_.Get()[byteOffset / sizeof(T)]; } - __aicore__ inline LocalTensor GateBetaTyped(uint64_t slot) + template + __aicore__ inline LocalTensor GateBetaStaging(uint64_t slot) { uint64_t byteOffset = slot * GateInputSlotBytes() + GateStageElems() * (3 * sizeof(T) + sizeof(GK_T)); - return gateWritebackBuf_.Get()[byteOffset / sizeof(BETA_T)]; + return gateWritebackBuf_.Get()[byteOffset / sizeof(BetaT)]; } __aicore__ inline LocalTensor GateKgTyped(uint64_t slot) @@ -1250,9 +1256,14 @@ class ChunkKdaFwdPrepareKernel { #if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 if constexpr (SAFE_GATE && COMPILE_BT == 64 && COMPILE_K == 128 && COMPILE_V == 128) { LocalTensor vTyped = GateVTyped(slot); - LocalTensor betaTyped = GateBetaTyped(slot); CopyRowsIn(vTyped, v_, VInputOffset(b, hv, token, 0), rows, V_, inputSequenceMajor_ ? HV_ * V_ : V_); - CopyVectorIn(betaTyped, beta_, BetaOffset(b, hv, token), rows); + if (betaIsFp32_) { + LocalTensor betaTyped = GateBetaStaging(slot); + CopyVectorIn(betaTyped, betaFp32_, BetaOffset(b, hv, token), rows); + } else { + LocalTensor betaTyped = GateBetaStaging(slot); + CopyVectorIn(betaTyped, betaBf16_, BetaOffset(b, hv, token), rows); + } } #endif SetFlag(mte2ToVEvent_); @@ -1442,6 +1453,29 @@ class ChunkKdaFwdPrepareKernel { PipeBarrier(); } + // beta arrives as fp32 or bf16 (tiling.betaDataType). Both are resolved at + // runtime so the whole prepare kernel stays independent of the beta dtype -- + // templating on it doubled every Prepare instantiation for one vector cast. + __aicore__ inline void LoadBetaAsFloatRow(uint64_t srcOffset, LocalTensor &dst, uint64_t count) + { + if (betaIsFp32_) { + LoadAsFloatRow(betaFp32_, srcOffset, dst, count); + } else { + LoadAsFloatRow(betaBf16_, srcOffset, dst, count); + } + } + + __aicore__ inline void ConvertStagedBeta(uint64_t slot, LocalTensor &dst, uint64_t count) + { + if (betaIsFp32_) { + LocalTensor staged = GateBetaStaging(slot); + Adds(dst, staged, 0.0f, static_cast(count)); + } else { + LocalTensor staged = GateBetaStaging(slot); + Cast(dst, staged, RoundMode::CAST_NONE, static_cast(count)); + } + } + template __aicore__ inline void LoadAsFloatVector(GlobalTensor &src, uint64_t srcOffset, LocalTensor &dst, LocalTensor &typedScratch, uint64_t count) @@ -1541,7 +1575,6 @@ class ChunkKdaFwdPrepareKernel { #if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 LocalTensor kgScore = GateKgTyped(qwSlot).template ReinterpretCast(); LocalTensor vTyped = GateVTyped(qwSlot); - LocalTensor betaTyped = GateBetaTyped(qwSlot); LocalTensor directQ = GateDirectQ(qwSlot); LocalTensor directW = GateDirectW(qwSlot); LocalTensor directV = GateDirectV(qwSlot); @@ -1567,11 +1600,7 @@ class ChunkKdaFwdPrepareKernel { } #endif #if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 - if constexpr (IsSameType::value) { - Adds(betaFp32, betaTyped, 0.0f, static_cast(tileRows)); - } else { - Cast(betaFp32, betaTyped, RoundMode::CAST_NONE, static_cast(tileRows)); - } + ConvertStagedBeta(qwSlot, betaFp32, tileRows); PipeBarrier(); #endif if (qwOutputPending) { @@ -1858,7 +1887,6 @@ class ChunkKdaFwdPrepareKernel { LocalTensor outFp32 = arena[4 * elems]; #else LocalTensor vTyped = GateVTyped(gateSlot); - LocalTensor betaTyped = GateBetaTyped(gateSlot); LocalTensor directQ = GateDirectQ(gateSlot); LocalTensor directW = GateDirectW(gateSlot); LocalTensor directV = GateDirectV(gateSlot); @@ -1877,11 +1905,7 @@ class ChunkKdaFwdPrepareKernel { } #endif #if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 - if constexpr (IsSameType::value) { - Adds(betaFp32, betaTyped, 0.0f, static_cast(tileRows)); - } else { - Cast(betaFp32, betaTyped, RoundMode::CAST_NONE, static_cast(tileRows)); - } + ConvertStagedBeta(gateSlot, betaFp32, tileRows); PipeBarrier(); #endif uint64_t nextTileRow = tileRow + maxRows; @@ -2896,7 +2920,7 @@ class ChunkKdaFwdPrepareKernel { LocalTensor maskLocal = arena[3 * KDA_SOLVE_MATRIX_ELEMENTS + KDA_SOLVE_BT + 512]; LocalTensor oneHotLocal = arena[3 * KDA_SOLVE_MATRIX_ELEMENTS + KDA_SOLVE_BT + 512 + KDA_SOLVE_BT]; - LoadAsFloatRow(beta_, BetaOffset(b, hv, start), betaLocal, KDA_SOLVE_BT); + LoadBetaAsFloatRow(BetaOffset(b, hv, start), betaLocal, KDA_SOLVE_BT); Brcb(betaBrcb, betaLocal, 8, {1, 8}); PipeBarrier(); @@ -2949,7 +2973,7 @@ class ChunkKdaFwdPrepareKernel { FillLocalFloat(betaLocal, 0.0f, KDA_SOLVE_BT); SetFlag(vToMte2Event_); WaitFlag(vToMte2Event_); - LoadAsFloatRow(beta_, BetaOffset(b, hv, start), betaLocal, curT); + LoadBetaAsFloatRow(BetaOffset(b, hv, start), betaLocal, curT); Brcb(betaBrcb, betaLocal, 8, {1, 8}); PipeBarrier(); @@ -3099,7 +3123,7 @@ class ChunkKdaFwdPrepareKernel { WaitFlag(vToMte2Event_); if (validRowCount > 0) { if (!scoreIncludesBeta) { - LoadAsFloatRow(beta_, BetaOffset(b, hv, token), betaLocal, validRowCount); + LoadBetaAsFloatRow(BetaOffset(b, hv, token), betaLocal, validRowCount); } if (curT < BT_) { DataCopy( @@ -3838,7 +3862,7 @@ class ChunkKdaFwdPrepareKernel { if (tileCount > tileRows) { tileCount = tileRows; } - LoadAsFloatRow(beta_, BetaOffset(b, hv, start + rowBegin + tileRow), betaLocal, tileCount); + LoadBetaAsFloatRow(BetaOffset(b, hv, start + rowBegin + tileRow), betaLocal, tileCount); ScaleRowsByBeta(w_, w_, b, hv, start, rowBegin + tileRow, tileCount, K_, betaLocal, betaBrcb, matrixLocal); ScaleRowsByBeta(v_, vNew_, b, hv, start, rowBegin + tileRow, tileCount, V_, betaLocal, betaBrcb, @@ -3847,7 +3871,7 @@ class ChunkKdaFwdPrepareKernel { return; } #endif - LoadAsFloatRow(beta_, BetaOffset(b, hv, start + rowBegin), betaLocal, rowCount); + LoadBetaAsFloatRow(BetaOffset(b, hv, start + rowBegin), betaLocal, rowCount); #if !defined(__CCE_AICORE__) || __CCE_AICORE__ != 310 Brcb(betaBrcb, betaLocal, static_cast((rowCount + 7) / 8), {1, 8}); PipeBarrier(); @@ -4662,7 +4686,8 @@ class ChunkKdaFwdPrepareKernel { GlobalTensor rawG_; GlobalTensor aLog_; GlobalTensor dtBias_; - GlobalTensor beta_; + GlobalTensor betaFp32_; + GlobalTensor betaBf16_; GlobalTensor initialState_; GlobalTensor o_; GlobalTensor finalState_; @@ -4727,6 +4752,7 @@ class ChunkKdaFwdPrepareKernel { bool hasALog_ = false; bool hasDtBias_ = false; bool storeQG_ = true; + bool betaIsFp32_ = true; float lowerBound_ = -5.0f; uint64_t usedCoreNum_ = 1; uint64_t solveCoreIdx_ = 0; @@ -4737,7 +4763,7 @@ class ChunkKdaFwdPrepareKernel { }; } // namespace -template __aicore__ inline void RunChunkKdaPrepare(GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR gk, GM_ADDR rawG, GM_ADDR aLog, GM_ADDR dtBias, GM_ADDR beta, GM_ADDR initialState, GM_ADDR cuSeqlens, @@ -4750,7 +4776,7 @@ __aicore__ inline void RunChunkKdaPrepare(GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_AD GM_ADDR prepareScratch = userWorkspace + tiling.prepareScratchOffset; if ASCEND_IS_AIC { - ChunkKdaFwdPrepareKernel op; + ChunkKdaFwdPrepareKernel op; op.Init(q, k, v, gk, rawG, aLog, dtBias, beta, initialState, cuSeqlens, chunkIndices, nullptr, nullptr, nullptr, nullptr, aqk, userWorkspace, aqkFp32, akkFp32, wSeed, akk, qg, qgScaled, uSeed, userWorkspace, finalKg, prepareScratch, tiling, &pipe, false, storeQG); @@ -4769,7 +4795,7 @@ __aicore__ inline void RunChunkKdaPrepare(GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_AD } } if ASCEND_IS_AIV { - ChunkKdaFwdPrepareKernel op; + ChunkKdaFwdPrepareKernel op; op.Init(q, k, v, gk, rawG, aLog, dtBias, beta, initialState, cuSeqlens, chunkIndices, nullptr, nullptr, nullptr, nullptr, aqk, userWorkspace, aqkFp32, akkFp32, wSeed, akk, qg, qgScaled, uSeed, userWorkspace, finalKg, prepareScratch, tiling, &pipe, true, storeQG); diff --git a/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd.cpp b/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd.cpp index e03c8aff33..1b4c07461b 100644 --- a/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd.cpp +++ b/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd.cpp @@ -22,7 +22,7 @@ namespace KdaForward { -template __aicore__ inline void DispatchGeneric( GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR g, GM_ADDR beta, @@ -32,7 +32,7 @@ __aicore__ inline void DispatchGeneric( GM_ADDR w, GM_ADDR u, GM_ADDR qg, GM_ADDR kg, GM_ADDR vNew, GM_ADDR h, GM_ADDR userWorkspace, const TilingData &tiling) { - RunGeneric( q, k, v, g, beta, aLog, dtBias, initialState, cuSeqlens, chunkIndices, attnOut, finalState, gk, aqk, akk, w, u, qg, kg, @@ -40,7 +40,7 @@ __aicore__ inline void DispatchGeneric( } #if KDA_COMPILE_ARCH35_FAST_PATH -template __aicore__ inline void DispatchArch35SafeGate( GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR g, GM_ADDR beta, @@ -52,13 +52,13 @@ __aicore__ inline void DispatchArch35SafeGate( { AscendC::TPipe pipe; if (tiling.safeGate) { - arch35::Run( q, k, v, g, beta, aLog, dtBias, initialState, cuSeqlens, chunkIndices, attnOut, finalState, gk, aqk, akk, w, u, qg, kg, vNew, h, userWorkspace, tiling, pipe); } else { - arch35::Run( q, k, v, g, beta, aLog, dtBias, initialState, cuSeqlens, chunkIndices, attnOut, finalState, gk, aqk, akk, w, u, qg, @@ -66,7 +66,7 @@ __aicore__ inline void DispatchArch35SafeGate( } } #elif defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 -template __aicore__ inline void DispatchArch35SafeGate( GM_ADDR, GM_ADDR, GM_ADDR, GM_ADDR, GM_ADDR, @@ -80,7 +80,7 @@ __aicore__ inline void DispatchArch35SafeGate( } #endif -template __aicore__ inline void DispatchGenericSafeGate( GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR g, GM_ADDR beta, @@ -91,13 +91,13 @@ __aicore__ inline void DispatchGenericSafeGate( GM_ADDR userWorkspace, const TilingData &tiling) { if (tiling.safeGate) { - DispatchGeneric( q, k, v, g, beta, aLog, dtBias, initialState, cuSeqlens, chunkIndices, attnOut, finalState, gk, aqk, akk, w, u, qg, kg, vNew, h, userWorkspace, tiling); } else { - DispatchGeneric( q, k, v, g, beta, aLog, dtBias, initialState, cuSeqlens, chunkIndices, attnOut, finalState, gk, aqk, akk, w, u, qg, @@ -156,48 +156,21 @@ extern "C" KDA_FWD_SCHEDULE_MODE __global__ __aicore__ void chunk_kda_fwd( tilingData->vHeadDim == 128; if (useFastPath) { #if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 - if (tilingData->betaDataType == 0) { - KdaForward::DispatchArch35SafeGate( - q, k, v, g, beta, a_log, dt_bias, initial_state, cu_seqlens, - chunk_indices, attn_out, final_state, gk, aqk, akk, w, u, qg, - kg, v_new, h, userWorkspace, *tilingData); - } else { - KdaForward::DispatchArch35SafeGate( - q, k, v, g, beta, a_log, dt_bias, initial_state, cu_seqlens, - chunk_indices, attn_out, final_state, gk, aqk, akk, w, u, qg, - kg, v_new, h, userWorkspace, *tilingData); - } + KdaForward::DispatchArch35SafeGate( + q, k, v, g, beta, a_log, dt_bias, initial_state, cu_seqlens, + chunk_indices, attn_out, final_state, gk, aqk, akk, w, u, qg, + kg, v_new, h, userWorkspace, *tilingData); #else - if (tilingData->betaDataType == 0) { - KdaForward::DispatchGenericSafeGate( - q, k, v, g, beta, a_log, dt_bias, initial_state, cu_seqlens, - chunk_indices, attn_out, final_state, gk, aqk, akk, w, u, qg, - kg, v_new, h, userWorkspace, *tilingData); - } else { - KdaForward::DispatchGenericSafeGate( - q, k, v, g, beta, a_log, dt_bias, initial_state, cu_seqlens, - chunk_indices, attn_out, final_state, gk, aqk, akk, w, u, qg, - kg, v_new, h, userWorkspace, *tilingData); - } + KdaForward::DispatchGenericSafeGate( + q, k, v, g, beta, a_log, dt_bias, initial_state, cu_seqlens, + chunk_indices, attn_out, final_state, gk, aqk, akk, w, u, qg, + kg, v_new, h, userWorkspace, *tilingData); #endif } else { - if (tilingData->betaDataType == 0) { - KdaForward::DispatchGenericSafeGate( - q, k, v, g, beta, a_log, dt_bias, initial_state, cu_seqlens, - chunk_indices, attn_out, final_state, gk, aqk, akk, w, u, qg, - kg, v_new, h, userWorkspace, *tilingData); - } else { - KdaForward::DispatchGenericSafeGate( - q, k, v, g, beta, a_log, dt_bias, initial_state, cu_seqlens, - chunk_indices, attn_out, final_state, gk, aqk, akk, w, u, qg, - kg, v_new, h, userWorkspace, *tilingData); - } + KdaForward::DispatchGenericSafeGate( + q, k, v, g, beta, a_log, dt_bias, initial_state, cu_seqlens, + chunk_indices, attn_out, final_state, gk, aqk, akk, w, u, qg, + kg, v_new, h, userWorkspace, *tilingData); } } diff --git a/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_common.h b/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_common.h index a6c99f782e..ec0cfcf6b0 100644 --- a/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_common.h +++ b/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_common.h @@ -185,7 +185,7 @@ __aicore__ inline void RunGateCumsum( } } -template __aicore__ inline void RunFrontEnd( @@ -210,7 +210,7 @@ __aicore__ inline void RunFrontEnd( } #endif - KdaPrepare::RunChunkKdaPrepare( q, k, v, addresses.gk, g, aLog, dtBias, beta, initialState, cuSeqlens, chunkIndices, aqk, akk, addresses.qg, @@ -295,7 +295,7 @@ __aicore__ inline void RunGenericBackEnd( addresses.vNew, addresses.h, attnOut, userWorkspace, tiling, pipe); } -template __aicore__ inline void RunGeneric( GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR g, GM_ADDR beta, @@ -309,7 +309,7 @@ __aicore__ inline void RunGeneric( finalState, gk, w, u, qg, kg, vNew, h, userWorkspace, tiling); #if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 TPipe pipe; - RunFrontEnd( q, k, v, g, beta, aLog, dtBias, initialState, cuSeqlens, chunkIndices, aqk, akk, addresses, userWorkspace, tiling, pipe); @@ -326,7 +326,7 @@ __aicore__ inline void RunGeneric( #else { TPipe pipe; - RunFrontEnd( q, k, v, g, beta, aLog, dtBias, initialState, cuSeqlens, chunkIndices, aqk, akk, addresses, userWorkspace, tiling, pipe); diff --git a/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_prepare.h b/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_prepare.h index 0bb747f8b6..ee95145a5a 100644 --- a/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_prepare.h +++ b/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_prepare.h @@ -159,7 +159,7 @@ __aicore__ inline T FloatToType(float value) return static_cast(value); } -template +template class ChunkKdaFwdPrepareKernel { public: using OUT_T = T; @@ -177,7 +177,8 @@ class ChunkKdaFwdPrepareKernel { k_.SetGlobalBuffer((__gm__ T *)k); v_.SetGlobalBuffer((__gm__ T *)v); gk_.SetGlobalBuffer((__gm__ GK_T *)gk); - beta_.SetGlobalBuffer((__gm__ BETA_T *)beta); + betaFp32_.SetGlobalBuffer((__gm__ float *)beta); + betaBf16_.SetGlobalBuffer((__gm__ bfloat16_t *)beta); if (initialState != nullptr) { initialState_.SetGlobalBuffer((__gm__ float *)initialState); } @@ -221,6 +222,7 @@ class ChunkKdaFwdPrepareKernel { isVarLen_ = tiling.isVarLen; inputSequenceMajor_ = tiling.inputSequenceMajor; usedCoreNum_ = tiling.prepareUsedCoreNum; + betaIsFp32_ = tiling.betaDataType == 0; constexpr uint64_t solvePipelineDepth = SAFE_GATE ? KDA_SOLVE_PIPELINE_DEPTH : 1; const uint64_t solveBytes = usedCoreNum_ * solvePipelineDepth * KDA_SOLVE_SCRATCH_SLOTS * BT_ * BT_ * sizeof(float); @@ -591,6 +593,18 @@ class ChunkKdaFwdPrepareKernel { PipeBarrier(); } + // beta arrives as fp32 or bf16 (tiling.betaDataType). Resolving that at + // runtime keeps the whole prepare kernel independent of the beta dtype; + // templating on it doubled every Prepare instantiation for one vector cast. + __aicore__ inline void LoadBetaAsFloatRow(uint64_t srcOffset, LocalTensor &dst, uint64_t count) + { + if (betaIsFp32_) { + LoadAsFloatRow(betaFp32_, srcOffset, dst, count); + } else { + LoadAsFloatRow(betaBf16_, srcOffset, dst, count); + } + } + template __aicore__ inline void LoadAsFloatVector(GlobalTensor &src, uint64_t srcOffset, LocalTensor &dst, LocalTensor &typedScratch, uint64_t count) @@ -1271,7 +1285,7 @@ class ChunkKdaFwdPrepareKernel { LocalTensor maskLocal = arena[3 * KDA_SOLVE_MATRIX_ELEMENTS + KDA_SOLVE_BT + 512]; LocalTensor oneHotLocal = arena[3 * KDA_SOLVE_MATRIX_ELEMENTS + KDA_SOLVE_BT + 512 + KDA_SOLVE_BT]; - LoadAsFloatRow(beta_, BetaOffset(b, hv, start), betaLocal, KDA_SOLVE_BT); + LoadBetaAsFloatRow(BetaOffset(b, hv, start), betaLocal, KDA_SOLVE_BT); Brcb(betaBrcb, betaLocal, 8, {1, 8}); PipeBarrier(); @@ -1324,7 +1338,7 @@ class ChunkKdaFwdPrepareKernel { FillLocalFloat(betaLocal, 0.0f, KDA_SOLVE_BT); SetFlag(vToMte2Event_); WaitFlag(vToMte2Event_); - LoadAsFloatRow(beta_, BetaOffset(b, hv, start), betaLocal, curT); + LoadBetaAsFloatRow(BetaOffset(b, hv, start), betaLocal, curT); Brcb(betaBrcb, betaLocal, 8, {1, 8}); PipeBarrier(); @@ -1414,7 +1428,7 @@ class ChunkKdaFwdPrepareKernel { SetFlag(vToMte2Event_); WaitFlag(vToMte2Event_); if (validRowCount > 0) { - LoadAsFloatRow(beta_, BetaOffset(b, hv, token), betaLocal, validRowCount); + LoadBetaAsFloatRow(BetaOffset(b, hv, token), betaLocal, validRowCount); DataCopy(aqkMat, aqk_[AOffset(b, hv, token, 0)], static_cast(validElemCount)); DataCopy(akkMat, akk_[AOffset(b, hv, token, 0)], static_cast(validElemCount)); SetFlag(mte2ToVEvent_); @@ -1851,7 +1865,7 @@ class ChunkKdaFwdPrepareKernel { LocalTensor betaLocal = arena; LocalTensor betaBrcb = arena[KDA_SOLVE_BT]; LocalTensor matrixLocal = arena[KDA_SOLVE_BT + 512]; - LoadAsFloatRow(beta_, BetaOffset(b, hv, start + rowBegin), betaLocal, rowCount); + LoadBetaAsFloatRow(BetaOffset(b, hv, start + rowBegin), betaLocal, rowCount); Brcb(betaBrcb, betaLocal, static_cast((rowCount + 7) / 8), {1, 8}); PipeBarrier(); ScaleRowsByBeta(w_, w_, b, hv, start, rowBegin, rowCount, K_, betaLocal, betaBrcb, matrixLocal); @@ -2415,7 +2429,8 @@ class ChunkKdaFwdPrepareKernel { GlobalTensor k_; GlobalTensor v_; GlobalTensor gk_; - GlobalTensor beta_; + GlobalTensor betaFp32_; + GlobalTensor betaBf16_; GlobalTensor initialState_; GlobalTensor o_; GlobalTensor finalState_; @@ -2469,6 +2484,7 @@ class ChunkKdaFwdPrepareKernel { bool hasInitial_ = false; bool isVarLen_ = false; bool inputSequenceMajor_ = false; + bool betaIsFp32_ = true; bool isAivOnly_ = false; uint64_t usedCoreNum_ = 1; uint64_t solveCoreIdx_ = 0; @@ -2478,7 +2494,7 @@ class ChunkKdaFwdPrepareKernel { }; } // namespace -template +template __aicore__ inline void RunChunkKdaPrepare(GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR gk, GM_ADDR beta, GM_ADDR initialState, GM_ADDR cuSeqlens, GM_ADDR chunkIndices, GM_ADDR aqk, GM_ADDR akk, GM_ADDR qg, GM_ADDR qgScaled, GM_ADDR wSeed, GM_ADDR uSeed, @@ -2489,14 +2505,14 @@ __aicore__ inline void RunChunkKdaPrepare(GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_AD GM_ADDR prepareScratch = userWorkspace + tiling.prepareScratchOffset; if ASCEND_IS_AIC { - ChunkKdaFwdPrepareKernel op; + ChunkKdaFwdPrepareKernel op; op.Init(q, k, v, gk, beta, initialState, cuSeqlens, chunkIndices, nullptr, nullptr, nullptr, nullptr, aqk, userWorkspace, aqkFp32, akkFp32, wSeed, akk, qg, qgScaled, uSeed, userWorkspace, prepareScratch, tiling, &pipe, false); op.ProcessAic(); } if ASCEND_IS_AIV { - ChunkKdaFwdPrepareKernel op; + ChunkKdaFwdPrepareKernel op; op.Init(q, k, v, gk, beta, initialState, cuSeqlens, chunkIndices, nullptr, nullptr, nullptr, nullptr, aqk, userWorkspace, aqkFp32, akkFp32, wSeed, akk, qg, qgScaled, uSeed, userWorkspace, prepareScratch, tiling, &pipe); @@ -2504,7 +2520,7 @@ __aicore__ inline void RunChunkKdaPrepare(GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_AD } } -template __aicore__ inline void RunChunkKdaPrepare(GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR gk, GM_ADDR, GM_ADDR, GM_ADDR, GM_ADDR beta, GM_ADDR initialState, GM_ADDR cuSeqlens, GM_ADDR chunkIndices, @@ -2512,8 +2528,8 @@ __aicore__ inline void RunChunkKdaPrepare(GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_AD GM_ADDR uSeed, GM_ADDR, GM_ADDR userWorkspace, const TilingData &tiling, TPipe &pipe, bool = true) { - RunChunkKdaPrepare(q, k, v, gk, beta, initialState, cuSeqlens, chunkIndices, aqk, akk, - qg, qgScaled, wSeed, uSeed, userWorkspace, tiling, pipe); + RunChunkKdaPrepare(q, k, v, gk, beta, initialState, cuSeqlens, chunkIndices, aqk, akk, qg, + qgScaled, wSeed, uSeed, userWorkspace, tiling, pipe); } } // namespace KdaPrepare From b66eea4624b74ef3f1bc7f3b76675a21996ec072 Mon Sep 17 00:00:00 2001 From: chenjz Date: Fri, 4 Sep 2026 09:25:12 +0800 Subject: [PATCH 09/11] build: add an opt-in incremental mode and scale -j to the host build_cmake_modules() unconditionally wipes the build directory, so every invocation is a cold build -- including one where nothing changed. That is the dominant cost while iterating, because the AscendC kernel targets (chunk_kda_fwd above all) do not benefit from the ccache launcher configured for the host compiler. KEEP_BUILD_DIR=1 reuses the existing CMake cache and object files. Default behaviour is unchanged, so nothing about CI or a release build moves; drop the variable whenever the CMake configuration itself changes. Also replace the hardcoded -j 16 with nproc (overridable via BUILD_JOBS). bash build.sh -a kernels Ascend950PR_9599 # as before KEEP_BUILD_DIR=1 bash build.sh -a kernels Ascend950PR_9599 # incremental Co-Authored-By: Claude Opus 5 --- build.sh | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/build.sh b/build.sh index 691568b88f..b14c27d9ef 100755 --- a/build.sh +++ b/build.sh @@ -399,13 +399,21 @@ function build_cmake_modules() ) fi - rm -rf "$BUILD_DIR" + # KEEP_BUILD_DIR=1 reuses the existing CMake cache and object files. The + # AscendC kernel targets (chunk_kda_fwd in particular) dominate a cold build, + # so keeping them lets an unrelated edit rebuild in seconds. Default stays a + # clean build; drop the variable whenever the CMake configuration changes. + if [[ "${KEEP_BUILD_DIR:-0}" == "1" && -f "$BUILD_DIR/CMakeCache.txt" ]]; then + echo "KEEP_BUILD_DIR=1: reusing $BUILD_DIR (incremental build)" + else + rm -rf "$BUILD_DIR" + fi mkdir -p "$BUILD_DIR" cmake "${cmake_args[@]}" \ -B "$BUILD_DIR" \ -S "$PROJECT_ROOT" - cmake --build "$BUILD_DIR" --target install -j 16 + cmake --build "$BUILD_DIR" --target install -j "${BUILD_JOBS:-$(nproc)}" } function build_deepep_kernels() From 9ff9cfcf6c5eaf5f9965f6f36062234a991b7b84 Mon Sep 17 00:00:00 2001 From: chenjz Date: Fri, 4 Sep 2026 10:14:11 +0800 Subject: [PATCH 10/11] fix(build): restore the -j 16 default instead of nproc Scaling -j to nproc was wrong: build parallelism here is bounded by RAM, not by cores. The CCE compiler needs roughly 1-2 GB per job and the AscendC kernel translation units are the hungry ones, so on a 28-core host with 15 GiB of RAM (9 GiB actually available) -j 28 wants ~40 GB and drives the machine into swap thrashing or the OOM killer. It hung a build at the no_workspace_kernel target. The change also bought very little, since the long pole is a single serial translation unit that extra jobs cannot help. Restore the historical default and keep BUILD_JOBS for hosts that can take more (or, as here, need less): BUILD_JOBS=8 bash build.sh -a kernels Ascend950PR_9599 Co-Authored-By: Claude Opus 5 --- build.sh | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/build.sh b/build.sh index b14c27d9ef..f5e7dd9be4 100755 --- a/build.sh +++ b/build.sh @@ -413,7 +413,10 @@ function build_cmake_modules() cmake "${cmake_args[@]}" \ -B "$BUILD_DIR" \ -S "$PROJECT_ROOT" - cmake --build "$BUILD_DIR" --target install -j "${BUILD_JOBS:-$(nproc)}" + # Parallelism is bounded by RAM, not cores: the CCE compiler needs ~1-2 GB per + # job and the AscendC kernel TUs are the memory-hungry ones. Keep the historical + # default and let callers tune it with BUILD_JOBS. + cmake --build "$BUILD_DIR" --target install -j "${BUILD_JOBS:-16}" } function build_deepep_kernels() From f2c84fae9342cc133c06ab72c9d80dd5d3542ada Mon Sep 17 00:00:00 2001 From: chenjz Date: Mon, 7 Sep 2026 17:41:41 +0800 Subject: [PATCH 11/11] delete repeated casual conv1d --- csrc/pytorch_extensions.cpp | 6 ------ 1 file changed, 6 deletions(-) diff --git a/csrc/pytorch_extensions.cpp b/csrc/pytorch_extensions.cpp index e5df0aedfc..b368b52e66 100644 --- a/csrc/pytorch_extensions.cpp +++ b/csrc/pytorch_extensions.cpp @@ -102,12 +102,6 @@ TORCH_LIBRARY_FRAGMENT(npu, m) "causal_conv1d_update(Tensor x, Tensor weight, Tensor(a!) conv_state, " "Tensor conv_state_indices, Tensor? bias=None, Tensor? num_accepted_tokens=None, " "Tensor? query_start_loc=None, bool activation_mode=False, int pad_slot_id=-1) -> Tensor"); - - m.def( - "causal_conv1d(Tensor x, Tensor weight, Tensor(a!) conv_states, Tensor? bias=None, " - "Tensor? query_start_loc=None, Tensor? cache_indices=None, Tensor? has_initial_state=None, " - "Tensor? num_accepted_tokens=None, int activation_mode=0, int pad_slot_id=-1, " - "int run_mode=0) -> Tensor"); m.def( "causal_conv1d(Tensor x, Tensor weight, Tensor(a!) conv_states, Tensor? bias=None, "