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

fixing bugs

Browse Source
This commit is contained in:
Juuso Korhonen
2025-10-23 09:47:30 +00:00
parent e144872308
commit 3c0e6d37bf
5 changed files with 36 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
example/ck_tile/01_unified_attention/unified_attention_impl.hpp
View File

@@ -110,7 +110,9 @@ template <typename Kernel>
float unified_attention_kernel_launch(const unified_attention_args& args, const stream_config& config)
{
index_t total_num_q_blocks = args.num_tokens / Kernel::BLOCK_Q + args.num_seqs;
index_t BLOCK_Q = Kernel::BLOCK_M / args.num_queries_per_kv;
index_t total_num_q_blocks = args.num_tokens / BLOCK_Q + args.num_seqs;
auto kargs = Kernel::MakeKargs(args.q_ptr,
args.k_ptr,

10
include/ck_tile/ops/unified_attention/kernel/unified_attention_kernel.hpp
View File

@@ -43,8 +43,8 @@ struct UnifiedAttentionKernel
// BLOCK_Q = BLOCK_M // num_queries_per_kv
// BLOCK_Q is the block size for q seqlen
static constexpr index_t BLOCK_Q = UnifiedAttentionPipeline::BLOCK_Q;
// static constexpr index_t BLOCK_M = UnifiedAttentionPipeline::BLOCK_M;
/// static constexpr index_t BLOCK_Q = UnifiedAttentionPipeline::BLOCK_Q;
static constexpr index_t BLOCK_M = UnifiedAttentionPipeline::BLOCK_M;
// BLOCK size for K seqlen
static constexpr index_t BLOCK_SIZE = UnifiedAttentionPipeline::BLOCK_SIZE;
@@ -276,10 +276,12 @@ struct UnifiedAttentionKernel
ck_tile::index_t pid = blockIdx.x;
const index_t BLOCK_M = BLOCK_Q * kargs.num_queries_per_kv;
const index_t num_queries_per_kv = kargs.num_queries_per_kv;
const index_t BLOCK_Q = BLOCK_M / num_queries_per_kv;
// for simplicity, batch stride we just modify the pointer
// const index_t num_head_q = kargs.num_head_q;
const index_t num_queries_per_kv = kargs.num_queries_per_kv;
// const index_t num_head_k = num_head_q / num_queries_per_kv;
pid = RemapTileIndices(pid, kargs);

2
include/ck_tile/ops/unified_attention/pipeline/tile_unified_attention_shape.hpp
View File

@@ -47,7 +47,7 @@ struct TileUnifiedAttentionShape
static constexpr index_t NumWarps = max(NumGemm0Warps, NumGemm1Warps);
static constexpr index_t BLOCK_M = BlockTile::at(number<0>{}); // tile size along q seqlen
static constexpr index_t BLOCK_M = BlockTile::at(number<0>{}); // tile size along the flattened batch dimension (: num_queries_per_kv * BS)
// static constexpr index_t BLOCK_M = BlockTile::at(number<1>{}); // tile size along q seqlen * num_queries_per_kv (q_head//kv_head)
static constexpr index_t BLOCK_SIZE = BlockTile::at(number<1>{}); // BLOCK size for K seqlen
static constexpr index_t HEAD_SIZE = BlockTile::at(number<2>{}); // BLOCK size for K seqlen

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

@@ -395,7 +395,7 @@ struct UnifiedAttentionPipeline
void* smem_ptr) const
{
using namespace ck_tile;
index_t BLOCK_Q = BLOCK_M / num_queries_per_kv;
constexpr index_t BLOCK_Q = BLOCK_M / num_queries_per_kv;
static_assert(
std::is_same_v<QDataType, remove_cvref_t<typename QDramBlockWindowTmp::DataType>> &&
@@ -410,7 +410,7 @@ struct UnifiedAttentionPipeline
BLOCK_SIZE == VDramBlockWindowTmp{}.get_window_lengths()[number<1>{}],
"wrong!");
static_assert(sizeof(SaccDataType) * BLOCK_Q * BLOCK_SIZE <= GetSmemSize());
static_assert(sizeof(SaccDataType) * BLOCK_SIZE <= GetSmemSize());
auto s_lds = make_tensor_view<address_space_enum::lds>(
reinterpret_cast<SaccDataType*>(static_cast<char*>(smem_ptr)),
MakeSimpleLdsDesc<BLOCK_Q, BLOCK_SIZE>());

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

@@ -92,8 +92,8 @@ struct UnifiedAttentionPipelineDefaultPolicy
{
using namespace ck_tile;
constexpr index_t kNPerBlock = Problem::UnifiedAttentionShape::kN0;
constexpr index_t kKPerBlock = Problem::UnifiedAttentionShape::kK0;
constexpr index_t kNPerBlock = Problem::UnifiedAttentionShape::BLOCK_SIZE;
constexpr index_t kKPerBlock = Problem::UnifiedAttentionShape::HEAD_SIZE;
constexpr index_t kBlockSize = Problem::kBlockSize;
constexpr index_t NumWarps = Problem::UnifiedAttentionShape::NumWarps;
constexpr index_t WarpSize = ck_tile::get_warp_size();
@@ -126,8 +126,8 @@ struct UnifiedAttentionPipelineDefaultPolicy
{
using namespace ck_tile;
constexpr index_t kNPerBlock = Problem::UnifiedAttentionShape::kK1;
constexpr index_t kKPerBlock = Problem::UnifiedAttentionShape::kN1;
constexpr index_t kNPerBlock = Problem::UnifiedAttentionShape::BLOCK_SIZE;
constexpr index_t kKPerBlock = Problem::UnifiedAttentionShape::HEAD_SIZE;
constexpr index_t kBlockSize = Problem::kBlockSize;
constexpr index_t NumWarps = Problem::UnifiedAttentionShape::NumWarps;
constexpr index_t WarpSize = ck_tile::get_warp_size();
@@ -197,8 +197,8 @@ struct UnifiedAttentionPipelineDefaultPolicy
constexpr index_t MWarp = Problem::UnifiedAttentionShape::Gemm1BlockWarps::at(number<0>{});
constexpr index_t NWarp = Problem::UnifiedAttentionShape::Gemm1BlockWarps::at(number<1>{});
constexpr index_t kNPerBlock = Problem::UnifiedAttentionShape::kN1;
constexpr index_t kKPerBlock = Problem::UnifiedAttentionShape::kK1;
constexpr index_t kNPerBlock = Problem::UnifiedAttentionShape::HEAD_SIZE;
constexpr index_t kKPerBlock = Problem::UnifiedAttentionShape::BLOCK_SIZE;
constexpr index_t NIterPerWarp = kNPerBlock / (NWarp * WarpGemm::kN);
constexpr index_t KIterPerWarp = kKPerBlock / WarpGemm::kK;
@@ -233,9 +233,9 @@ struct UnifiedAttentionPipelineDefaultPolicy
typename Problem::KDataType,
typename Problem::SaccDataType,
Problem::kBlockSize,
TileGemmShape<sequence<Problem::UnifiedAttentionShape::kM0,
Problem::UnifiedAttentionShape::kN0,
Problem::UnifiedAttentionShape::kK0>,
TileGemmShape<sequence<Problem::UnifiedAttentionShape::BLOCK_M,
Problem::UnifiedAttentionShape::BLOCK_SIZE,
Problem::UnifiedAttentionShape::HEAD_SIZE>,
typename Problem::UnifiedAttentionShape::Gemm0BlockWarps,
typename Problem::UnifiedAttentionShape::Gemm0WarpTile>>;
@@ -279,9 +279,9 @@ struct UnifiedAttentionPipelineDefaultPolicy
typename Problem::VDataType,
typename Problem::OaccDataType,
Problem::kBlockSize,
TileGemmShape<sequence<Problem::UnifiedAttentionShape::kM0,
Problem::UnifiedAttentionShape::kN1,
Problem::UnifiedAttentionShape::kK1>,
TileGemmShape<sequence<Problem::UnifiedAttentionShape::BLOCK_M,
Problem::UnifiedAttentionShape::HEAD_SIZE,
Problem::UnifiedAttentionShape::BLOCK_SIZE>,
typename Problem::UnifiedAttentionShape::Gemm1BlockWarps,
typename Problem::UnifiedAttentionShape::Gemm1WarpTile>>;
/// NOTICE: in order to use load_tile_transpose() later for V tiles, we have to pass
@@ -317,8 +317,8 @@ struct UnifiedAttentionPipelineDefaultPolicy
using namespace ck_tile;
// K is always k-major, we use async-copy to load into LDS
constexpr index_t kNPerBlock = Problem::UnifiedAttentionShape::kN0;
constexpr index_t kKPerBlock = Problem::UnifiedAttentionShape::kK0;
constexpr index_t kNPerBlock = Problem::UnifiedAttentionShape::BLOCK_SIZE;
constexpr index_t kKPerBlock = Problem::UnifiedAttentionShape::HEAD_SIZE;
constexpr index_t kBlockSize = Problem::kBlockSize;
constexpr index_t NumWarps = Problem::UnifiedAttentionShape::NumWarps;
constexpr index_t WarpSize = ck_tile::get_warp_size();
@@ -374,8 +374,8 @@ struct UnifiedAttentionPipelineDefaultPolicy
using namespace ck_tile;
// K is always k-major, we use async-copy to load into LDS
constexpr index_t kNPerBlock = Problem::UnifiedAttentionShape::kN0;
constexpr index_t kKPerBlock = Problem::UnifiedAttentionShape::kK0;
constexpr index_t kNPerBlock = Problem::UnifiedAttentionShape::BLOCK_SIZE;
constexpr index_t kKPerBlock = Problem::UnifiedAttentionShape::HEAD_SIZE;
constexpr index_t kBlockSize = Problem::kBlockSize;
constexpr index_t NumWarps = Problem::UnifiedAttentionShape::NumWarps;
constexpr index_t WarpSize = ck_tile::get_warp_size();
@@ -423,8 +423,8 @@ struct UnifiedAttentionPipelineDefaultPolicy
{
// this function assume K/V can share smem
constexpr index_t SingleKSize = [&]() {
constexpr index_t kNPerBlock = Problem::UnifiedAttentionShape::kN0;
constexpr index_t kKPerBlock = Problem::UnifiedAttentionShape::kK1;
constexpr index_t kNPerBlock = Problem::UnifiedAttentionShape::BLOCK_SIZE;
constexpr index_t kKPerBlock = Problem::UnifiedAttentionShape::BLOCK_SIZE;
constexpr index_t NumWarps = Problem::UnifiedAttentionShape::NumWarps;
constexpr index_t WarpSize = ck_tile::get_warp_size();
@@ -447,8 +447,8 @@ struct UnifiedAttentionPipelineDefaultPolicy
constexpr index_t kKPack = GetSmemKPackK<Problem>();
static_assert(PixelsPerRow % kKPack == 0);
constexpr index_t NPerRow = PixelsPerRow / kKPack;
constexpr index_t kNPerBlock = Problem::UnifiedAttentionShape::kN1;
constexpr index_t kKPerBlock = Problem::UnifiedAttentionShape::kK1;
constexpr index_t kNPerBlock = Problem::UnifiedAttentionShape::HEAD_SIZE;
constexpr index_t kKPerBlock = Problem::UnifiedAttentionShape::BLOCK_SIZE;
static_assert(kNPerBlock % NPerRow == 0);
static_assert(kKPerBlock % kKPack == 0);
@@ -465,8 +465,8 @@ struct UnifiedAttentionPipelineDefaultPolicy
using namespace ck_tile;
/// FIXME: rename the kNPerBlock & kKPerBlock since the kN1 is congtigous dimension
constexpr index_t kNPerBlock = Problem::UnifiedAttentionShape::kK1;
constexpr index_t kKPerBlock = Problem::UnifiedAttentionShape::kN1;
constexpr index_t kNPerBlock = Problem::UnifiedAttentionShape::BLOCK_SIZE;
constexpr index_t kKPerBlock = Problem::UnifiedAttentionShape::HEAD_SIZE;
constexpr index_t kBlockSize = Problem::kBlockSize;
constexpr index_t NumWarps = Problem::UnifiedAttentionShape::NumWarps;
constexpr index_t WarpSize = ck_tile::get_warp_size();
@@ -522,8 +522,8 @@ struct UnifiedAttentionPipelineDefaultPolicy
using namespace ck_tile;
/// FIXME: rename the kNPerBlock & kKPerBlock since the kN1 is congtigous dimension
constexpr index_t kNPerBlock = Problem::UnifiedAttentionShape::kK1;
constexpr index_t kKPerBlock = Problem::UnifiedAttentionShape::kN1;
constexpr index_t kNPerBlock = Problem::UnifiedAttentionShape::BLOCK_SIZE;
constexpr index_t kKPerBlock = Problem::UnifiedAttentionShape::HEAD_SIZE;
constexpr index_t kBlockSize = Problem::kBlockSize;
constexpr index_t NumWarps = Problem::UnifiedAttentionShape::NumWarps;
constexpr index_t WarpSize = ck_tile::get_warp_size();