From 86f7ebcf2768827f38fcd8200c20e9be5e374668 Mon Sep 17 00:00:00 2001 From: Amir Ghamarian Date: Mon, 30 Mar 2026 17:04:38 +0000 Subject: [PATCH] Add bs32 narrow tier: 2-warp 16x16 MFMA kBlockM=32 for decode New unified_attention_decode_bs32_kernel_traits: 2 warps, 16x16 MFMA, kBlockM=32, kBlockQ=4 for GQA-8 with block_size=32. Uses DecodePolicy (NumWarpPerGroup=2) for deterministic LDS access. Replaces the small tier (kBlockQ=8, 87.5% waste) for bs32 tiny decode with 2x less query waste (kBlockQ=4, 75% waste). Made-with: Cursor --- ...tention_d64_bf16_mask_gqa8_bs32_narrow.cpp | 14 +++++ ...ention_d64_bf16_nmask_gqa8_bs32_narrow.cpp | 14 +++++ ...tention_d64_fp16_mask_gqa8_bs32_narrow.cpp | 14 +++++ ...ention_d64_fp16_nmask_gqa8_bs32_narrow.cpp | 14 +++++ .../unified_attention.cpp | 47 +++++++++----- .../unified_attention_impl.hpp | 62 +++++++++++++++++++ 6 files changed, 148 insertions(+), 17 deletions(-) create mode 100644 example/ck_tile/42_unified_attention/instances/unified_attention_d64_bf16_mask_gqa8_bs32_narrow.cpp create mode 100644 example/ck_tile/42_unified_attention/instances/unified_attention_d64_bf16_nmask_gqa8_bs32_narrow.cpp create mode 100644 example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp16_mask_gqa8_bs32_narrow.cpp create mode 100644 example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp16_nmask_gqa8_bs32_narrow.cpp diff --git a/example/ck_tile/42_unified_attention/instances/unified_attention_d64_bf16_mask_gqa8_bs32_narrow.cpp b/example/ck_tile/42_unified_attention/instances/unified_attention_d64_bf16_mask_gqa8_bs32_narrow.cpp new file mode 100644 index 0000000000..204319568f --- /dev/null +++ b/example/ck_tile/42_unified_attention/instances/unified_attention_d64_bf16_mask_gqa8_bs32_narrow.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_bs32_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_bs32_narrow.cpp b/example/ck_tile/42_unified_attention/instances/unified_attention_d64_bf16_nmask_gqa8_bs32_narrow.cpp new file mode 100644 index 0000000000..27ccae7b06 --- /dev/null +++ b/example/ck_tile/42_unified_attention/instances/unified_attention_d64_bf16_nmask_gqa8_bs32_narrow.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_bs32_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_bs32_narrow.cpp b/example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp16_mask_gqa8_bs32_narrow.cpp new file mode 100644 index 0000000000..22d1f71e5b --- /dev/null +++ b/example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp16_mask_gqa8_bs32_narrow.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_bs32_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_bs32_narrow.cpp b/example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp16_nmask_gqa8_bs32_narrow.cpp new file mode 100644 index 0000000000..47a8dd7939 --- /dev/null +++ b/example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp16_nmask_gqa8_bs32_narrow.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_bs32_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 0ae38a480a..aa54396740 100644 --- a/example/ck_tile/42_unified_attention/unified_attention.cpp +++ b/example/ck_tile/42_unified_attention/unified_attention.cpp @@ -58,9 +58,9 @@ std::ostream& operator<<(std::ostream& stream, return unified_attention_kernel_dispatch_decode(args, config); \ } -#define DISPATCH_UNIFIED_ATTENTION_DECODE_TINY_BS32(DType, IsMask, HSize, BM, NQPKV) \ +#define DISPATCH_UNIFIED_ATTENTION_DECODE_BS32_NARROW(DType, IsMask, HSize, BM, NQPKV) \ { \ - using kernel_traits = unified_attention_decode_tiny_kernel_traits; \ + using kernel_traits = unified_attention_decode_bs32_kernel_traits; \ return unified_attention_kernel_dispatch_decode(args, config); \ } @@ -109,23 +109,36 @@ std::pair unified_attention(const unified_attention_args& args, { const bool use_bs32 = (args.page_blk_size < 64); - if(tier == tile_tier::tiny && !use_bs32) + if(tier == tile_tier::tiny) { - // Tiny tier (1-warp, 16x16 MFMA) only for bs64+. - // bs32 tiny has a race condition in block_tile_reduce_sync - // when the last KV block has 9+ valid tokens; promote to small. - 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) + if(use_bs32) { + // bs32 narrow: 2 warps, 16x16 MFMA, kBlockM=32, kBlockQ=4. + // Avoids 1-warp race condition; 2x less waste than small tier. + if(args.data_type == unified_attention_args::data_type_enum::fp16) + { + if(!is_mask) DISPATCH_UNIFIED_ATTENTION_DECODE_BS32_NARROW(unified_attention_args::data_type_enum::fp16, false, 64, 32, 8) + else DISPATCH_UNIFIED_ATTENTION_DECODE_BS32_NARROW(unified_attention_args::data_type_enum::fp16, true, 64, 32, 8) + } + else if(args.data_type == unified_attention_args::data_type_enum::bf16) + { + if(!is_mask) DISPATCH_UNIFIED_ATTENTION_DECODE_BS32_NARROW(unified_attention_args::data_type_enum::bf16, false, 64, 32, 8) + else DISPATCH_UNIFIED_ATTENTION_DECODE_BS32_NARROW(unified_attention_args::data_type_enum::bf16, true, 64, 32, 8) + } + } else { + // bs64 tiny: 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 || (tier == tile_tier::tiny && use_bs32)) + else if(tier == tile_tier::small) { if(args.data_type == unified_attention_args::data_type_enum::fp16) { @@ -194,7 +207,7 @@ std::pair unified_attention(const unified_attention_args& args, return std::make_pair(false, -1.f); } -#undef DISPATCH_UNIFIED_ATTENTION_DECODE_TINY_BS32 +#undef DISPATCH_UNIFIED_ATTENTION_DECODE_BS32_NARROW #undef DISPATCH_UNIFIED_ATTENTION_DECODE_SMALL_BS32 #undef DISPATCH_UNIFIED_ATTENTION_DECODE_MEDIUM_BS32 #undef DISPATCH_UNIFIED_ATTENTION_DECODE_TINY 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 52deb491b7..31e5c4c6ad 100644 --- a/example/ck_tile/42_unified_attention/unified_attention_impl.hpp +++ b/example/ck_tile/42_unified_attention/unified_attention_impl.hpp @@ -316,6 +316,68 @@ struct unified_attention_decode_tiny_kernel_traits using kernel = UnifiedAttentionKernel; }; +// bs32 decode traits: 2 warps, 16x16 MFMA, kBlockM=32, kBlockQ=4 for GQA-8. +// Used for block_size=32 decode: avoids the 1-warp pipeline race condition +// and reduces query waste from 87.5% (small tier kBlockQ=8) to 75% (kBlockQ=4). +template +struct unified_attention_decode_bs32_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 = BlockSize_; + + 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>; + using unified_attention_block_warps = sequence<2, 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)