Add decode-tuned unified attention kernel (4 warps, kBlockM=128)

Generalize the unified attention pipeline to support NumWarpGroups=1
(single warp group) with a serial K/V loop, in addition to the existing
NumWarpGroups=2 interleaved pipeline.

New decode kernel traits use 4 warps (sequence<4,1,1>) with kBlockM=128
and kBlockQ=16 for GQA-8, reducing Q tile padding waste from 31/32 to
15/16 for decode workloads (max_seqlen_q=1).

Host-side dispatch (is_decode_shape) routes low-token workloads to the
decode kernel automatically.

Benchmark results on d64 GQA-8 (via aiter):
- 64-seq decode:  2.2x slower -> 1.27x slower (1.73x speedup)
- 512-seq decode: 3.5x slower -> 1.6x slower  (2.2x speedup)
- 1-seq decode:   0.83x (CK wins) -> 0.81x    (no regression)
- Prefill:        unchanged (uses original 8-warp kernel)

Made-with: Cursor

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<unified_attention_args::data_type_enum::bf16, true, 64, 128, 8>;
+
+INST_UNIFIED_ATTENTION_DISPATCH(kernel_traits)
+
+} // namespace ck_tile

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<unified_attention_args::data_type_enum::bf16, false, 64, 128, 8>;
+
+INST_UNIFIED_ATTENTION_DISPATCH(kernel_traits)
+
+} // namespace ck_tile

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<unified_attention_args::data_type_enum::fp16, true, 64, 128, 8>;
+
+INST_UNIFIED_ATTENTION_DISPATCH(kernel_traits)
+
+} // namespace ck_tile

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<unified_attention_args::data_type_enum::fp16, false, 64, 128, 8>;
+
+INST_UNIFIED_ATTENTION_DISPATCH(kernel_traits)
+
+} // namespace ck_tile

example/ck_tile/42_unified_attention/unified_attention.cpp

@@ -26,14 +26,24 @@ std::ostream& operator<<(std::ostream& stream,
         return unified_attention_kernel_dispatch<kernel_traits>(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<DType, IsMask, HSize, BM, NQPKV>; \
+        return unified_attention_kernel_dispatch<kernel_traits>(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<bool, float> unified_attention(const unified_attention_args& args,
                                          const stream_config& config)
 {
     const bool is_mask = (args.mask_type != static_cast<int>(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<bool, float> 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<bool, float> 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

example/ck_tile/42_unified_attention/unified_attention_impl.hpp

@@ -122,13 +122,12 @@ struct unified_attention_kernel_traits
     using kernel = UnifiedAttentionKernel<unified_attention_pipeline, epilogue>;
 };
 
-// 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 <unified_attention_args::data_type_enum DataType,
           bool IsMasking,
           index_t HeadSize_  = 128,
-          index_t BlockM_    = 64,
+          index_t BlockM_    = 128,
           index_t NumQPerKV_ = 1>
 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<kBlockM, kBlockQ, BLOCK_SIZE, HEAD_SIZE>;
+    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<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>;
+
+    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>;
+};
+
+// Aggressive decode traits: 4 warps (2x2 layout), kBlockM=64 for maximum decode throughput.
+template <unified_attention_args::data_type_enum DataType,
+          bool IsMasking,
+          index_t HeadSize_  = 64,
+          index_t BlockM_    = 64,
+          index_t NumQPerKV_ = 8>
+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<kBlockM, kBlockQ, BLOCK_SIZE, HEAD_SIZE>;
     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<unified_attention_block_tile,
                                                               unified_attention_block_warps,

include/ck_tile/ops/unified_attention/pipeline/unified_attention_pipeline.hpp

@@ -381,7 +381,7 @@ struct UnifiedAttentionPipeline
         // static_assert(kPageBlockSize == kHeadDimPadded);
 
         constexpr index_t NumWarpGroups = Problem::kBlockSize / Policy::NumThreadPerWarpGroup;
-        static_assert(NumWarpGroups == 2);
+        static_assert(NumWarpGroups == 1 || NumWarpGroups == 2);
 
         [[maybe_unused]] auto print_dist_tensor = [&](const auto& dist_tensor, const char* name) {
             printf("[POYENC] %s (size=%d): %5.2f",
@@ -960,22 +960,109 @@ struct UnifiedAttentionPipeline
 
         if(1 < num_total_loop)
         {
-            if(warp_group_id == 0)
+            if constexpr(NumWarpGroups == 1)
             {
-                V_mem_load(number<1>{}); // 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: