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
2026-07-18 17:48:06 +00:00 · 2026-03-30 17:04:38 +00:00
parent e21b915381
commit 86f7ebcf27
6 changed files with 148 additions and 17 deletions
--- 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
@@ -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<unified_attention_args::data_type_enum::bf16, true, 64, 32, 8, 32>;
+
+INST_UNIFIED_ATTENTION_DISPATCH_DECODE(kernel_traits)
+
+} // namespace ck_tile
--- 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
@@ -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<unified_attention_args::data_type_enum::bf16, false, 64, 32, 8, 32>;
+
+INST_UNIFIED_ATTENTION_DISPATCH_DECODE(kernel_traits)
+
+} // namespace ck_tile
--- 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
@@ -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<unified_attention_args::data_type_enum::fp16, true, 64, 32, 8, 32>;
+
+INST_UNIFIED_ATTENTION_DISPATCH_DECODE(kernel_traits)
+
+} // namespace ck_tile
--- 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
@@ -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<unified_attention_args::data_type_enum::fp16, false, 64, 32, 8, 32>;
+
+INST_UNIFIED_ATTENTION_DISPATCH_DECODE(kernel_traits)
+
+} // namespace ck_tile
--- 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<kernel_traits>(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<DType, IsMask, HSize, BM, NQPKV, 32>; \
+        using kernel_traits = unified_attention_decode_bs32_kernel_traits<DType, IsMask, HSize, BM, NQPKV, 32>; \
        return unified_attention_kernel_dispatch_decode<kernel_traits>(args, config); \
    }

@@ -109,23 +109,36 @@ std::pair<bool, float> 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<bool, float> 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
--- 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<unified_attention_pipeline, epilogue>;
 };

+// 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 <unified_attention_args::data_type_enum DataType,
+          bool IsMasking,
+          index_t HeadSize_  = 64,
+          index_t BlockM_    = 32,
+          index_t NumQPerKV_ = 8,
+          index_t BlockSize_ = 32>
+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<kBlockM, kBlockQ, BLOCK_SIZE, HEAD_SIZE>;
+    using unified_attention_warp_gemm_shape = sequence<16, 16, 32>;
+    using unified_attention_block_warps     = sequence<2, 1, 1>;
+
+    using unified_attention_shape = TileUnifiedAttentionShape<unified_attention_block_tile,
+                                                              unified_attention_block_warps,
+                                                              unified_attention_warp_gemm_shape,
+                                                              unified_attention_block_warps,
+                                                              unified_attention_warp_gemm_shape,
+                                                              true>;
+
+    using unified_attention_traits = TileUnifiedAttentionTraits<true, false, -1>;
+    using unified_attention_mask   = GenericAttentionMask<IsMasking, false>;
+
+    using unified_attention_pipeline_problem = UnifiedAttentionPipelineProblem<
+        typename unified_attention_problem_traits<date_type>::qkvp_dtype,
+        typename unified_attention_problem_traits<date_type>::qkvp_dtype,
+        typename unified_attention_problem_traits<date_type>::qkvp_dtype,
+        typename unified_attention_problem_traits<date_type>::acc_dtype,
+        typename unified_attention_problem_traits<date_type>::acc_dtype,
+        typename unified_attention_problem_traits<date_type>::acc_dtype,
+        typename unified_attention_problem_traits<date_type>::lse_dtype,
+        typename unified_attention_problem_traits<date_type>::qkvp_dtype,
+        typename unified_attention_problem_traits<date_type>::acc_dtype,
+        typename unified_attention_problem_traits<date_type>::o_dtype,
+        unified_attention_shape,
+        unified_attention_mask,
+        unified_attention_traits>;
+
+    using unified_attention_pipeline =
+        UnifiedAttentionPipeline<unified_attention_pipeline_problem,
+                                 UnifiedAttentionPipelineDecodePolicy>;
+
+    using epilogue = Default2DEpilogue<
+        Default2DEpilogueProblem<typename unified_attention_problem_traits<date_type>::acc_dtype,
+                                 typename unified_attention_problem_traits<date_type>::o_dtype,
+                                 true, true, true>>;
+
+    using kernel = UnifiedAttentionKernel<unified_attention_pipeline, epilogue>;
+};
+
 template <typename Kernel, bool UseDecodeGrid = false>
 float unified_attention_kernel_launch(const unified_attention_args& args,
                                      const stream_config& config)