diff --git a/example/ck_tile/42_unified_attention/instances/unified_attention_d64_bf16_mask_gqa8_decode_t.cpp b/example/ck_tile/42_unified_attention/instances/unified_attention_d64_bf16_mask_gqa8_decode_t.cpp new file mode 100644 index 0000000000..31de799776 --- /dev/null +++ b/example/ck_tile/42_unified_attention/instances/unified_attention_d64_bf16_mask_gqa8_decode_t.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_tiny_kernel_traits; + +INST_UNIFIED_ATTENTION_DISPATCH_DECODE(kernel_traits) + +} // namespace ck_tile diff --git a/example/ck_tile/42_unified_attention/instances/unified_attention_d64_bf16_nmask_gqa8_decode_t.cpp b/example/ck_tile/42_unified_attention/instances/unified_attention_d64_bf16_nmask_gqa8_decode_t.cpp new file mode 100644 index 0000000000..d4d4306638 --- /dev/null +++ b/example/ck_tile/42_unified_attention/instances/unified_attention_d64_bf16_nmask_gqa8_decode_t.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_tiny_kernel_traits; + +INST_UNIFIED_ATTENTION_DISPATCH_DECODE(kernel_traits) + +} // namespace ck_tile diff --git a/example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp16_mask_gqa8_decode_t.cpp b/example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp16_mask_gqa8_decode_t.cpp new file mode 100644 index 0000000000..c3a4c1f850 --- /dev/null +++ b/example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp16_mask_gqa8_decode_t.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_tiny_kernel_traits; + +INST_UNIFIED_ATTENTION_DISPATCH_DECODE(kernel_traits) + +} // namespace ck_tile diff --git a/example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp16_nmask_gqa8_decode_t.cpp b/example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp16_nmask_gqa8_decode_t.cpp new file mode 100644 index 0000000000..7c127cc130 --- /dev/null +++ b/example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp16_nmask_gqa8_decode_t.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_tiny_kernel_traits; + +INST_UNIFIED_ATTENTION_DISPATCH_DECODE(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 e642f3a60f..6cd32bd5da 100644 --- a/example/ck_tile/42_unified_attention/unified_attention.cpp +++ b/example/ck_tile/42_unified_attention/unified_attention.cpp @@ -39,15 +39,25 @@ std::ostream& operator<<(std::ostream& stream, return unified_attention_kernel_dispatch_decode(args, config); \ } -enum class tile_tier { large, medium, small }; +#define DISPATCH_UNIFIED_ATTENTION_DECODE_TINY(DType, IsMask, HSize, BM, NQPKV) \ + { \ + using kernel_traits = unified_attention_decode_tiny_kernel_traits; \ + return unified_attention_kernel_dispatch_decode(args, config); \ + } + +enum class tile_tier { large, medium, small, tiny }; static tile_tier select_tile_tier(const unified_attention_args& args) { const index_t avg_q = args.num_seqs > 0 ? args.num_tokens / args.num_seqs : args.num_tokens; - const index_t kBlockQ_small = 64 / args.num_queries_per_kv; // kBlockQ for 2-warp kernel + const index_t kBlockQ_tiny = 16 / args.num_queries_per_kv; // kBlockQ for 1-warp 16x16 kernel + if(avg_q <= kBlockQ_tiny) + return tile_tier::tiny; // pure decode: 1 warp, 16x16 MFMA, kBlockM=16 + + const index_t kBlockQ_small = 64 / args.num_queries_per_kv; // kBlockQ for 2-warp kernel if(avg_q <= kBlockQ_small) - return tile_tier::small; // pure decode: 2 warps, kBlockM=64 + return tile_tier::small; // decode: 2 warps, kBlockM=64 const index_t kBlockQ_medium = 128 / args.num_queries_per_kv; // kBlockQ for 4-warp kernel if(avg_q <= kBlockQ_medium * 8) @@ -80,7 +90,21 @@ 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(tier == tile_tier::small) + if(tier == tile_tier::tiny) + { + // Tiny decode: 1 warp, 16x16 MFMA, kBlockM=16 (kBlockQ=2) + if(args.data_type == unified_attention_args::data_type_enum::fp16) + { + if(!is_mask) DISPATCH_UNIFIED_ATTENTION_DECODE_TINY(unified_attention_args::data_type_enum::fp16, false, 64, 16, 8) + else DISPATCH_UNIFIED_ATTENTION_DECODE_TINY(unified_attention_args::data_type_enum::fp16, true, 64, 16, 8) + } + else if(args.data_type == unified_attention_args::data_type_enum::bf16) + { + if(!is_mask) DISPATCH_UNIFIED_ATTENTION_DECODE_TINY(unified_attention_args::data_type_enum::bf16, false, 64, 16, 8) + else DISPATCH_UNIFIED_ATTENTION_DECODE_TINY(unified_attention_args::data_type_enum::bf16, true, 64, 16, 8) + } + } + else if(tier == tile_tier::small) { // Small decode: 2 warps, kBlockM=64 (kBlockQ=8) if(args.data_type == unified_attention_args::data_type_enum::fp16) @@ -130,6 +154,7 @@ std::pair unified_attention(const unified_attention_args& args, return std::make_pair(false, -1.f); } +#undef DISPATCH_UNIFIED_ATTENTION_DECODE_TINY #undef DISPATCH_UNIFIED_ATTENTION_DECODE_SMALL #undef DISPATCH_UNIFIED_ATTENTION_DECODE_MEDIUM #undef DISPATCH_UNIFIED_ATTENTION 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 7172a42cfb..fd2e2d825a 100644 --- a/example/ck_tile/42_unified_attention/unified_attention_impl.hpp +++ b/example/ck_tile/42_unified_attention/unified_attention_impl.hpp @@ -252,6 +252,68 @@ struct unified_attention_decode_small_kernel_traits using kernel = UnifiedAttentionKernel; }; +// Tiny decode traits: 1 warp, 16x16 MFMA, kBlockM=16, kBlockQ=2 for GQA-8. +// Matches Triton's BLOCK_M=16 / BLOCK_Q=2 decode configuration. +// Uses block_tile_reduce_sync instead of permlane32_swap for 16x16 MFMA. +template +struct unified_attention_decode_tiny_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<16, 16, 32>; + // 1 warp: kBlockM=1*16=16, kBlockSize=64, NumWarpGroups=1 + using unified_attention_block_warps = sequence<1, 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; +}; + template float unified_attention_kernel_launch(const unified_attention_args& args, const stream_config& config) diff --git a/include/ck_tile/ops/unified_attention/pipeline/unified_attention_pipeline.hpp b/include/ck_tile/ops/unified_attention/pipeline/unified_attention_pipeline.hpp index d91a122634..394532b49d 100644 --- a/include/ck_tile/ops/unified_attention/pipeline/unified_attention_pipeline.hpp +++ b/include/ck_tile/ops/unified_attention/pipeline/unified_attention_pipeline.hpp @@ -60,6 +60,9 @@ struct UnifiedAttentionPipeline static constexpr ck_tile::index_t kBlockM = UnifiedAttentionShape::kBlockM; static constexpr ck_tile::index_t kBlockQ = UnifiedAttentionShape::kBlockQ; + static constexpr ck_tile::index_t kWarpGemmM = + UnifiedAttentionShape::Gemm0WarpTile::at(ck_tile::number<0>{}); + static constexpr ck_tile::index_t kPageBlockSize = UnifiedAttentionShape::kPageBlockSize; static constexpr ck_tile::index_t kHeadDim = UnifiedAttentionShape::kHeadDim; static constexpr ck_tile::index_t kHeadDimPadded = UnifiedAttentionShape::kHeadDimPadded; @@ -513,15 +516,20 @@ struct UnifiedAttentionPipeline auto m_latest = block_tile_reduce( sp(sp_reg_idx).sp_compute, sequence<1>{}, f_max, m.thread_buf_[0]); #if defined(__gfx950__) - // assuming that we are using 32x32 mfma - int32x2_t swapped_regs = - __builtin_amdgcn_permlane32_swap(bit_cast(m_latest.thread_buf_[0]), - bit_cast(m_latest.thread_buf_[0]), - false, - false); - /// TODO: eliminate 2 redudant v_max_f32 instructions generated by the compiler - m_latest.thread_buf_[0] = f_max(bit_cast(swapped_regs.x), - bit_cast(swapped_regs.y)); + if constexpr(kWarpGemmM == 32) + { + int32x2_t swapped_regs = + __builtin_amdgcn_permlane32_swap(bit_cast(m_latest.thread_buf_[0]), + bit_cast(m_latest.thread_buf_[0]), + false, + false); + m_latest.thread_buf_[0] = f_max(bit_cast(swapped_regs.x), + bit_cast(swapped_regs.y)); + } + else + { + block_tile_reduce_sync(m_latest, f_max, bool_constant{}); + } #else block_tile_reduce_sync(m_latest, f_max, bool_constant{}); #endif @@ -558,14 +566,20 @@ struct UnifiedAttentionPipeline static_assert(rowsum_p.thread_buf_.size() == 1, "assuming that each thread holds 1 rowsum value"); #if defined(__gfx950__) - // assuming that we are using 32x32 mfma - int32x2_t swapped_regs = - __builtin_amdgcn_permlane32_swap(bit_cast(rowsum_p.thread_buf_[0]), - bit_cast(rowsum_p.thread_buf_[0]), - false, - false); - rowsum_p.thread_buf_[0] = f_sum(bit_cast(swapped_regs.x), - bit_cast(swapped_regs.y)); + if constexpr(kWarpGemmM == 32) + { + int32x2_t swapped_regs = + __builtin_amdgcn_permlane32_swap(bit_cast(rowsum_p.thread_buf_[0]), + bit_cast(rowsum_p.thread_buf_[0]), + false, + false); + rowsum_p.thread_buf_[0] = f_sum(bit_cast(swapped_regs.x), + bit_cast(swapped_regs.y)); + } + else + { + block_tile_reduce_sync(rowsum_p, f_sum, bool_constant{}); + } #else block_tile_reduce_sync(rowsum_p, f_sum, bool_constant{}); #endif diff --git a/include/ck_tile/ops/unified_attention/pipeline/unified_attention_pipeline_default_policy.hpp b/include/ck_tile/ops/unified_attention/pipeline/unified_attention_pipeline_default_policy.hpp index adc6bb3271..c2773f9076 100644 --- a/include/ck_tile/ops/unified_attention/pipeline/unified_attention_pipeline_default_policy.hpp +++ b/include/ck_tile/ops/unified_attention/pipeline/unified_attention_pipeline_default_policy.hpp @@ -603,4 +603,11 @@ struct UnifiedAttentionPipelineDecodePolicy : UnifiedAttentionPipelineDefaultPol NumWarpPerGroup * ck_tile::get_warp_size(); }; +struct UnifiedAttentionPipelineTinyDecodePolicy : UnifiedAttentionPipelineDefaultPolicy +{ + static constexpr ck_tile::index_t NumWarpPerGroup = 1; + static constexpr ck_tile::index_t NumThreadPerWarpGroup = + NumWarpPerGroup * ck_tile::get_warp_size(); +}; + } // namespace ck_tile