diff --git a/example/ck_tile/42_unified_attention/instances/unified_attention_d64_bf16_mask_gqa8_decode.cpp b/example/ck_tile/42_unified_attention/instances/unified_attention_d64_bf16_mask_gqa8_decode.cpp new file mode 100644 index 0000000000..a01ff6a23f --- /dev/null +++ b/example/ck_tile/42_unified_attention/instances/unified_attention_d64_bf16_mask_gqa8_decode.cpp @@ -0,0 +1,14 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +#include "unified_attention.hpp" +#include "unified_attention_impl.hpp" + +namespace ck_tile { + +using kernel_traits = + unified_attention_decode_kernel_traits; + +INST_UNIFIED_ATTENTION_DISPATCH(kernel_traits) + +} // namespace ck_tile diff --git a/example/ck_tile/42_unified_attention/instances/unified_attention_d64_bf16_nmask_gqa8_decode.cpp b/example/ck_tile/42_unified_attention/instances/unified_attention_d64_bf16_nmask_gqa8_decode.cpp new file mode 100644 index 0000000000..15ad7a8565 --- /dev/null +++ b/example/ck_tile/42_unified_attention/instances/unified_attention_d64_bf16_nmask_gqa8_decode.cpp @@ -0,0 +1,14 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +#include "unified_attention.hpp" +#include "unified_attention_impl.hpp" + +namespace ck_tile { + +using kernel_traits = + unified_attention_decode_kernel_traits; + +INST_UNIFIED_ATTENTION_DISPATCH(kernel_traits) + +} // namespace ck_tile diff --git a/example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp16_mask_gqa8_decode.cpp b/example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp16_mask_gqa8_decode.cpp new file mode 100644 index 0000000000..f9087e4147 --- /dev/null +++ b/example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp16_mask_gqa8_decode.cpp @@ -0,0 +1,14 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +#include "unified_attention.hpp" +#include "unified_attention_impl.hpp" + +namespace ck_tile { + +using kernel_traits = + unified_attention_decode_kernel_traits; + +INST_UNIFIED_ATTENTION_DISPATCH(kernel_traits) + +} // namespace ck_tile diff --git a/example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp16_nmask_gqa8_decode.cpp b/example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp16_nmask_gqa8_decode.cpp new file mode 100644 index 0000000000..06ae388ba7 --- /dev/null +++ b/example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp16_nmask_gqa8_decode.cpp @@ -0,0 +1,14 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +#include "unified_attention.hpp" +#include "unified_attention_impl.hpp" + +namespace ck_tile { + +using kernel_traits = + unified_attention_decode_kernel_traits; + +INST_UNIFIED_ATTENTION_DISPATCH(kernel_traits) + +} // namespace ck_tile diff --git a/example/ck_tile/42_unified_attention/unified_attention.cpp b/example/ck_tile/42_unified_attention/unified_attention.cpp index 75a1629933..a3e470b546 100644 --- a/example/ck_tile/42_unified_attention/unified_attention.cpp +++ b/example/ck_tile/42_unified_attention/unified_attention.cpp @@ -26,14 +26,24 @@ std::ostream& operator<<(std::ostream& stream, return unified_attention_kernel_dispatch(args, config); \ } +// Helper macro for decode-tuned dispatch (4 warps, kBlockM=128). +#define DISPATCH_UNIFIED_ATTENTION_DECODE(DType, IsMask, HSize, BM, NQPKV) \ + { \ + using kernel_traits = unified_attention_decode_kernel_traits; \ + return unified_attention_kernel_dispatch(args, config); \ + } + +static bool is_decode_shape(const unified_attention_args& args) +{ + const index_t kBlockQ_prefill = 256 / args.num_queries_per_kv; + return args.num_tokens <= args.num_seqs * kBlockQ_prefill; +} + std::pair unified_attention(const unified_attention_args& args, const stream_config& config) { const bool is_mask = (args.mask_type != static_cast(mask_enum::no_mask)); - - // Route based on (data_type, mask, hdim, num_queries_per_kv). - // Decode-tuned instances require pipeline changes (NumWarpGroups must == 2, - // which means exactly 8 warps; fewer warps are not supported). + const bool use_decode = is_decode_shape(args); // d128, MHA (num_queries_per_kv == 1) if(args.hdim == 128 && args.num_queries_per_kv == 1) @@ -53,15 +63,33 @@ std::pair unified_attention(const unified_attention_args& args, // d64, GQA-8 (num_queries_per_kv == 8) if(args.hdim == 64 && args.num_queries_per_kv == 8) { - if(args.data_type == unified_attention_args::data_type_enum::fp16) + if(use_decode) { - if(!is_mask) DISPATCH_UNIFIED_ATTENTION(unified_attention_args::data_type_enum::fp16, false, 64, 256, 8) - else DISPATCH_UNIFIED_ATTENTION(unified_attention_args::data_type_enum::fp16, true, 64, 256, 8) + // Decode-tuned: 4 warps, kBlockM=128 (kBlockQ=16) + if(args.data_type == unified_attention_args::data_type_enum::fp16) + { + if(!is_mask) DISPATCH_UNIFIED_ATTENTION_DECODE(unified_attention_args::data_type_enum::fp16, false, 64, 128, 8) + else DISPATCH_UNIFIED_ATTENTION_DECODE(unified_attention_args::data_type_enum::fp16, true, 64, 128, 8) + } + else if(args.data_type == unified_attention_args::data_type_enum::bf16) + { + if(!is_mask) DISPATCH_UNIFIED_ATTENTION_DECODE(unified_attention_args::data_type_enum::bf16, false, 64, 128, 8) + else DISPATCH_UNIFIED_ATTENTION_DECODE(unified_attention_args::data_type_enum::bf16, true, 64, 128, 8) + } } - else if(args.data_type == unified_attention_args::data_type_enum::bf16) + else { - if(!is_mask) DISPATCH_UNIFIED_ATTENTION(unified_attention_args::data_type_enum::bf16, false, 64, 256, 8) - else DISPATCH_UNIFIED_ATTENTION(unified_attention_args::data_type_enum::bf16, true, 64, 256, 8) + // Prefill: 8 warps, kBlockM=256 (kBlockQ=32) + if(args.data_type == unified_attention_args::data_type_enum::fp16) + { + if(!is_mask) DISPATCH_UNIFIED_ATTENTION(unified_attention_args::data_type_enum::fp16, false, 64, 256, 8) + else DISPATCH_UNIFIED_ATTENTION(unified_attention_args::data_type_enum::fp16, true, 64, 256, 8) + } + else if(args.data_type == unified_attention_args::data_type_enum::bf16) + { + if(!is_mask) DISPATCH_UNIFIED_ATTENTION(unified_attention_args::data_type_enum::bf16, false, 64, 256, 8) + else DISPATCH_UNIFIED_ATTENTION(unified_attention_args::data_type_enum::bf16, true, 64, 256, 8) + } } } @@ -71,6 +99,8 @@ std::pair unified_attention(const unified_attention_args& args, return std::make_pair(false, -1.f); } +#undef DISPATCH_UNIFIED_ATTENTION_DECODE + #undef DISPATCH_UNIFIED_ATTENTION } // namespace ck_tile diff --git a/example/ck_tile/42_unified_attention/unified_attention_impl.hpp b/example/ck_tile/42_unified_attention/unified_attention_impl.hpp index 77c97ce109..12881e3d86 100644 --- a/example/ck_tile/42_unified_attention/unified_attention_impl.hpp +++ b/example/ck_tile/42_unified_attention/unified_attention_impl.hpp @@ -122,13 +122,12 @@ struct unified_attention_kernel_traits using kernel = UnifiedAttentionKernel; }; -// Decode-tuned traits: fewer warps, smaller kBlockM for low-token workloads. -// NOTE: Currently cannot compile due to pipeline constraint (NumWarpGroups must == 2). -// Kept for future pipeline work. +// Decode-tuned traits: 4 warps (1 warp group), kBlockM=128, serial pipeline. +// Uses the single-warp-group path in UnifiedAttentionPipeline. template struct unified_attention_decode_kernel_traits { @@ -136,15 +135,75 @@ struct unified_attention_decode_kernel_traits static constexpr bool is_masking = IsMasking; static constexpr index_t kBlockM = BlockM_; - static constexpr index_t BLOCK_SIZE = 32; static constexpr index_t HEAD_SIZE = HeadSize_; + static constexpr index_t BLOCK_SIZE = (HEAD_SIZE <= 64) ? 64 : 32; + + static constexpr index_t num_queries_per_kv = NumQPerKV_; + static constexpr index_t kBlockQ = kBlockM / num_queries_per_kv; + + // kBlockM kBlockQ BLOCK_SIZE HEAD_SIZE + using unified_attention_block_tile = sequence; + using unified_attention_warp_gemm_shape = sequence<32, 32, 16>; + // 4 warps -> kBlockSize = 256 threads -> NumWarpGroups = 1 + using unified_attention_block_warps = sequence<4, 1, 1>; + + using unified_attention_shape = TileUnifiedAttentionShape; + + using unified_attention_traits = TileUnifiedAttentionTraits; + using unified_attention_mask = GenericAttentionMask; + + using unified_attention_pipeline_problem = UnifiedAttentionPipelineProblem< + typename unified_attention_problem_traits::qkvp_dtype, + typename unified_attention_problem_traits::qkvp_dtype, + typename unified_attention_problem_traits::qkvp_dtype, + typename unified_attention_problem_traits::acc_dtype, + typename unified_attention_problem_traits::acc_dtype, + typename unified_attention_problem_traits::acc_dtype, + typename unified_attention_problem_traits::lse_dtype, + typename unified_attention_problem_traits::qkvp_dtype, + typename unified_attention_problem_traits::acc_dtype, + typename unified_attention_problem_traits::o_dtype, + unified_attention_shape, + unified_attention_mask, + unified_attention_traits>; + + using unified_attention_pipeline = UnifiedAttentionPipeline; + + using epilogue = Default2DEpilogue< + Default2DEpilogueProblem::acc_dtype, + typename unified_attention_problem_traits::o_dtype, + true, true, true>>; + + using kernel = UnifiedAttentionKernel; +}; + +// Aggressive decode traits: 4 warps (2x2 layout), kBlockM=64 for maximum decode throughput. +template +struct unified_attention_decode_small_kernel_traits +{ + static constexpr auto date_type = DataType; + static constexpr bool is_masking = IsMasking; + + static constexpr index_t kBlockM = BlockM_; + static constexpr index_t HEAD_SIZE = HeadSize_; + static constexpr index_t BLOCK_SIZE = (HEAD_SIZE <= 64) ? 64 : 32; static constexpr index_t num_queries_per_kv = NumQPerKV_; static constexpr index_t kBlockQ = kBlockM / num_queries_per_kv; using unified_attention_block_tile = sequence; using unified_attention_warp_gemm_shape = sequence<32, 32, 16>; - using unified_attention_block_warps = sequence<2, 1, 1>; + // 2x2 warp layout: 4 warps total, kBlockM=2*32=64, N split=2*32=64 + using unified_attention_block_warps = sequence<2, 2, 1>; using unified_attention_shape = TileUnifiedAttentionShape{}); // V1 - K_lds_load(number<1>{}); // K1 + // --- Single warp group: serial pipeline --- + // After pre-stage: + // sp(0) has QK for block 0 (alu0 + alu_D_upd done, alu1 NOT done) + // V0 loading to LDS (V buf 0) + // K1 in LDS (K buf 1) if num_total_loop >= 2 + // K2 loading to LDS (K buf 0) if num_total_loop >= 3 - __builtin_amdgcn_s_setprio(0); + // Step 1: consume V0, K1 -> produce PV(0), QK(1) + s_waitcnt_vmcnt<0>(); __builtin_amdgcn_s_barrier(); - while(core_loop(number<0>{})) - ; + + V_lds_load(number<0>{}); // V0 from LDS + s_waitcnt_lgkmcnt<0>(); + fmha_alu1(number<0>{}); // finalize sp(0) -> P(0) + gemm(number<0>{}, /*gemm_idx=*/number<1>{}); // PV: P(0)*V0 + + __builtin_amdgcn_s_barrier(); + K_lds_load(number<1>{}); // K1 from LDS + s_waitcnt_lgkmcnt<0>(); + + V_mem_load(number<1>{}); // start V1 -> LDS buf 1 + + gemm(number<1>{}, /*gemm_idx=*/number<0>{}); // QK: Q*K1 -> sp(1) + fmha_mask(number<1>{}); + fmha_alu0(number<1>{}); + fmha_alu_D_upd(); + i_total_loops++; + + while(i_total_loops < num_total_loop) + { + // Even step: V from buf 1, K from buf 0, QK -> sp(0) + // kv_tile is a union: must finish PV GEMM (v_tile) before K load + s_waitcnt_vmcnt<0>(); + __builtin_amdgcn_s_barrier(); + + V_lds_load(number<1>{}); // V from buf 1 -> kv_tile.v_tile + s_waitcnt_lgkmcnt<0>(); + fmha_alu1(number<1>{}); // finalize sp(1) -> P(1) + gemm(number<1>{}, /*gemm_idx=*/number<1>{}); // PV: P(1)*V + + __builtin_amdgcn_s_barrier(); + K_lds_load(number<0>{}); // K from buf 0 -> kv_tile.k_tile + s_waitcnt_lgkmcnt<0>(); + + if(i_total_loops + 1 < num_total_loop) + K_mem_load(number<1>{}); // next K -> buf 1 + V_mem_load(number<0>{}); // next V -> buf 0 + + gemm(number<0>{}, /*gemm_idx=*/number<0>{}); // QK -> sp(0) + fmha_mask(number<0>{}); + fmha_alu0(number<0>{}); + fmha_alu_D_upd(); + i_total_loops++; + + if(i_total_loops >= num_total_loop) + break; + + // Odd step: V from buf 0, K from buf 1, QK -> sp(1) + s_waitcnt_vmcnt<0>(); + __builtin_amdgcn_s_barrier(); + + V_lds_load(number<0>{}); // V from buf 0 -> kv_tile.v_tile + s_waitcnt_lgkmcnt<0>(); + fmha_alu1(number<0>{}); // finalize sp(0) -> P(0) + gemm(number<0>{}, /*gemm_idx=*/number<1>{}); // PV: P(0)*V + + __builtin_amdgcn_s_barrier(); + K_lds_load(number<1>{}); // K from buf 1 -> kv_tile.k_tile + s_waitcnt_lgkmcnt<0>(); + + if(i_total_loops + 1 < num_total_loop) + K_mem_load(number<0>{}); // next K -> buf 0 + V_mem_load(number<1>{}); // next V -> buf 1 + + gemm(number<1>{}, /*gemm_idx=*/number<0>{}); // QK -> sp(1) + fmha_mask(number<1>{}); + fmha_alu0(number<1>{}); + fmha_alu_D_upd(); + i_total_loops++; + } } - if(warp_group_id != 0) + else { - __builtin_amdgcn_s_setprio(1); - __builtin_amdgcn_s_barrier(); - while(core_loop(number<1>{})) - ; + // --- Two warp groups: interleaved pipeline --- + if(warp_group_id == 0) + { + V_mem_load(number<1>{}); // V1 + K_lds_load(number<1>{}); // K1 + + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + while(core_loop(number<0>{})) + ; + } + if(warp_group_id != 0) + { + __builtin_amdgcn_s_setprio(1); + __builtin_amdgcn_s_barrier(); + while(core_loop(number<1>{})) + ; + } } } label_main_loops_exit: