ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-07 23:57:14 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add 16x16 MFMA tiny decode kernel (1 warp, kBlockM=16, kBlockQ=2)

Browse Source 
Enable 16x16 MFMA for decode by making the softmax cross-warp reduction
conditional on the warp gemm M dimension: use permlane32_swap for 32x32
MFMA (2 lanes per row), fall back to block_tile_reduce_sync for 16x16
MFMA (4 lanes per row).

New tiny decode traits: 1 warp, sequence<1,1,1>, warp_gemm 16x16x32,
kBlockM=16, kBlockQ=2 for GQA-8. This matches Triton's BLOCK_M=16 /
BLOCK_Q=2 decode configuration exactly.

Also adds 4-tier dispatch: tiny (avg_q<=2) -> small (avg_q<=8) ->
medium (avg_q<=128) -> large (prefill).

Benchmark results (d64 GQA-8 via aiter, 363 shapes):
  Before: CK faster 135 (37.2%), Triton faster 228 (62.8%)
  After:  CK faster 247 (68.0%), Triton faster 116 (32.0%)

Key shapes:
  1-seq decode:   0.021ms (CK 0.75x, wins 25%)
  64-seq decode:  0.025ms vs Triton 0.029ms (CK wins 14%)
  512-seq decode: 0.018ms vs Triton 0.021ms (CK wins)
  Weighted end-to-end: CK/Triton = 0.999x (tied)

Verified correct on 10 shapes: bf16+fp16, d64 GQA-8 + d128 MHA,
batch 1-64, all 4 dispatch tiers.

Made-with: Cursor
This commit is contained in:
Amir Ghamarian
2026-03-28 12:19:34 +00:00
parent 5f9b03746d
commit 33b2015939
8 changed files with 185 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
example/ck_tile/42_unified_attention/instances/unified_attention_d64_bf16_mask_gqa8_decode_t.cpp Normal file
View File

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

14
example/ck_tile/42_unified_attention/instances/unified_attention_d64_bf16_nmask_gqa8_decode_t.cpp Normal file
View File

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

14
example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp16_mask_gqa8_decode_t.cpp Normal file
View File

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

14
example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp16_nmask_gqa8_decode_t.cpp Normal file
View File

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

33
example/ck_tile/42_unified_attention/unified_attention.cpp
View File

@@ -39,15 +39,25 @@ std::ostream& operator<<(std::ostream& stream,
return unified_attention_kernel_dispatch_decode<kernel_traits>(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<DType, IsMask, HSize, BM, NQPKV>; \
return unified_attention_kernel_dispatch_decode<kernel_traits>(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<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(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<bool, float> 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

62
example/ck_tile/42_unified_attention/unified_attention_impl.hpp
View File

@@ -252,6 +252,68 @@ struct unified_attention_decode_small_kernel_traits
using kernel = UnifiedAttentionKernel<unified_attention_pipeline, epilogue>;
};
// 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 <unified_attention_args::data_type_enum DataType,
bool IsMasking,
index_t HeadSize_ = 64,
index_t BlockM_ = 16,
index_t NumQPerKV_ = 8>
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<kBlockM, kBlockQ, BLOCK_SIZE, HEAD_SIZE>;
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<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,
UnifiedAttentionPipelineTinyDecodePolicy>;
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)

48
include/ck_tile/ops/unified_attention/pipeline/unified_attention_pipeline.hpp
View File

@@ -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<SMPLComputeDataType>(
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<int32_t>(m_latest.thread_buf_[0]),
bit_cast<int32_t>(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<SMPLComputeDataType>(swapped_regs.x),
bit_cast<SMPLComputeDataType>(swapped_regs.y));
if constexpr(kWarpGemmM == 32)
{
int32x2_t swapped_regs =
__builtin_amdgcn_permlane32_swap(bit_cast<int32_t>(m_latest.thread_buf_[0]),
bit_cast<int32_t>(m_latest.thread_buf_[0]),
false,
false);
m_latest.thread_buf_[0] = f_max(bit_cast<SMPLComputeDataType>(swapped_regs.x),
bit_cast<SMPLComputeDataType>(swapped_regs.y));
}
else
{
block_tile_reduce_sync(m_latest, f_max, bool_constant<false>{});
}
#else
block_tile_reduce_sync(m_latest, f_max, bool_constant<false>{});
#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<int32_t>(rowsum_p.thread_buf_[0]),
bit_cast<int32_t>(rowsum_p.thread_buf_[0]),
false,
false);
rowsum_p.thread_buf_[0] = f_sum(bit_cast<SMPLComputeDataType>(swapped_regs.x),
bit_cast<SMPLComputeDataType>(swapped_regs.y));
if constexpr(kWarpGemmM == 32)
{
int32x2_t swapped_regs =
__builtin_amdgcn_permlane32_swap(bit_cast<int32_t>(rowsum_p.thread_buf_[0]),
bit_cast<int32_t>(rowsum_p.thread_buf_[0]),
false,
false);
rowsum_p.thread_buf_[0] = f_sum(bit_cast<SMPLComputeDataType>(swapped_regs.x),
bit_cast<SMPLComputeDataType>(swapped_regs.y));
}
else
{
block_tile_reduce_sync(rowsum_p, f_sum, bool_constant<false>{});
}
#else
block_tile_reduce_sync(rowsum_p, f_sum, bool_constant<false>{});
#endif

7
include/ck_tile/ops/unified_attention/pipeline/unified_attention_pipeline_default_policy.hpp
View File

@@ -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