ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-11 17:51:40 +00:00
Code Issues Packages Projects Releases Wiki Activity

Debug

Browse Source
This commit is contained in:
Tianxing Wu
2025-12-08 09:11:52 +00:00
parent 4cec09547f
commit 00b1d9794f
15 changed files with 237 additions and 106 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
example/ck_tile/01_unified_attention/instances/unified_attention_d128_bf16_fp8_mask.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_kernel_traits<unified_attention_args::data_type_enum::bf16, 1, true>;
INST_UNIFIED_ATTENTION_DISPATCH(kernel_traits)
} // namespace ck_tile

14
example/ck_tile/01_unified_attention/instances/unified_attention_d128_bf16_fp8_nmask.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_kernel_traits<unified_attention_args::data_type_enum::bf16, 1, false>;
INST_UNIFIED_ATTENTION_DISPATCH(kernel_traits)
} // namespace ck_tile

2
example/ck_tile/01_unified_attention/instances/unified_attention_d128_bf16_mask.cpp
View File

@@ -7,7 +7,7 @@
namespace ck_tile {
using kernel_traits =
unified_attention_kernel_traits<unified_attention_args::data_type_enum::bf16, 1, true>;
unified_attention_kernel_traits<unified_attention_args::data_type_enum::bf16, 0, true>;
INST_UNIFIED_ATTENTION_DISPATCH(kernel_traits)

2
example/ck_tile/01_unified_attention/instances/unified_attention_d128_bf16_nmask.cpp
View File

@@ -7,7 +7,7 @@
namespace ck_tile {
using kernel_traits =
unified_attention_kernel_traits<unified_attention_args::data_type_enum::bf16, 1, false>;
unified_attention_kernel_traits<unified_attention_args::data_type_enum::bf16, 0, false>;
INST_UNIFIED_ATTENTION_DISPATCH(kernel_traits)

14
example/ck_tile/01_unified_attention/instances/unified_attention_d128_fp16_fp8_mask.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_kernel_traits<unified_attention_args::data_type_enum::fp16, 1, true>;
INST_UNIFIED_ATTENTION_DISPATCH(kernel_traits)
} // namespace ck_tile

14
example/ck_tile/01_unified_attention/instances/unified_attention_d128_fp16_fp8_nmask.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_kernel_traits<unified_attention_args::data_type_enum::fp16, 1, false>;
INST_UNIFIED_ATTENTION_DISPATCH(kernel_traits)
} // namespace ck_tile

2
example/ck_tile/01_unified_attention/instances/unified_attention_d128_fp16_mask.cpp
View File

@@ -7,7 +7,7 @@
namespace ck_tile {
using kernel_traits =
unified_attention_kernel_traits<unified_attention_args::data_type_enum::fp16, 1, true>;
unified_attention_kernel_traits<unified_attention_args::data_type_enum::fp16, 0, true>;
INST_UNIFIED_ATTENTION_DISPATCH(kernel_traits)

2
example/ck_tile/01_unified_attention/instances/unified_attention_d128_fp16_nmask.cpp
View File

@@ -7,7 +7,7 @@
namespace ck_tile {
using kernel_traits =
unified_attention_kernel_traits<unified_attention_args::data_type_enum::fp16, 1, false>;
unified_attention_kernel_traits<unified_attention_args::data_type_enum::fp16, 0, false>;
INST_UNIFIED_ATTENTION_DISPATCH(kernel_traits)

14
example/ck_tile/01_unified_attention/instances/unified_attention_d128_q_bf16_fp8_mask.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_kernel_traits<unified_attention_args::data_type_enum::bf16, 2, true>;
INST_UNIFIED_ATTENTION_DISPATCH(kernel_traits)
} // namespace ck_tile

14
example/ck_tile/01_unified_attention/instances/unified_attention_d128_q_bf16_fp8_nmask.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_kernel_traits<unified_attention_args::data_type_enum::bf16, 2, false>;
INST_UNIFIED_ATTENTION_DISPATCH(kernel_traits)
} // namespace ck_tile

14
example/ck_tile/01_unified_attention/instances/unified_attention_d128_q_fp16_fp8_mask.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_kernel_traits<unified_attention_args::data_type_enum::fp16, 2, true>;
INST_UNIFIED_ATTENTION_DISPATCH(kernel_traits)
} // namespace ck_tile

14
example/ck_tile/01_unified_attention/instances/unified_attention_d128_q_fp16_fp8_nmask.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_kernel_traits<unified_attention_args::data_type_enum::fp16, 2, false>;
INST_UNIFIED_ATTENTION_DISPATCH(kernel_traits)
} // namespace ck_tile

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

@@ -177,74 +177,73 @@ float unified_attention_kernel_launch(const unified_attention_args& args,
assert(BLOCK_Q == BLOCK_M / args.num_queries_per_kv &&
"BLOCK_Q must equal BLOCK_M / num_queries_per_kv");
index_t total_num_q_blocks = args.num_tokens / BLOCK_Q + args.num_seqs;
auto kargs_lambda =
[&]() {
if constexpr(Kernel::kIsQuantized)
{
return Kernel::MakeKargs(args.q_ptr,
args.k_ptr,
args.v_ptr,
args.o_ptr,
args.num_blks,
args.num_head_q,
args.num_queries_per_kv,
args.scale_s,
args.scale_q,
args.scale_k,
args.scale_v,
args.scale_out,
total_num_q_blocks,
args.query_stride_0,
args.query_stride_1,
args.stride_k_cache_0,
args.stride_k_cache_1,
args.stride_k_cache_2,
args.stride_k_cache_3,
args.stride_v_cache_0,
args.stride_v_cache_1,
args.stride_v_cache_2,
args.stride_v_cache_3,
args.output_stride_0,
args.output_stride_1,
args.block_tables_ptr,
args.block_table_stride,
args.seq_lens_ptr,
args.query_start_len_ptr,
args.num_seqs);
}
else
{
return Kernel::MakeKargs(args.q_ptr,
args.k_ptr,
args.v_ptr,
args.o_ptr,
args.num_blks,
args.num_head_q,
args.num_queries_per_kv,
args.scale_s,
total_num_q_blocks,
args.query_stride_0,
args.query_stride_1,
args.stride_k_cache_0,
args.stride_k_cache_1,
args.stride_k_cache_2,
args.stride_k_cache_3,
args.stride_v_cache_0,
args.stride_v_cache_1,
args.stride_v_cache_2,
args.stride_v_cache_3,
args.output_stride_0,
args.output_stride_1,
args.block_tables_ptr,
args.block_table_stride,
args.seq_lens_ptr,
args.query_start_len_ptr,
args.num_seqs);
}
};
auto kargs = kargs_lambda();
auto kargs_lambda = [&]() {
if constexpr(Kernel::kIsQuantized)
{
return Kernel::MakeKargs(args.q_ptr,
args.k_ptr,
args.v_ptr,
args.o_ptr,
args.num_blks,
args.num_head_q,
args.num_queries_per_kv,
args.scale_s,
args.scale_q,
args.scale_k,
args.scale_v,
args.scale_out,
total_num_q_blocks,
args.query_stride_0,
args.query_stride_1,
args.stride_k_cache_0,
args.stride_k_cache_1,
args.stride_k_cache_2,
args.stride_k_cache_3,
args.stride_v_cache_0,
args.stride_v_cache_1,
args.stride_v_cache_2,
args.stride_v_cache_3,
args.output_stride_0,
args.output_stride_1,
args.block_tables_ptr,
args.block_table_stride,
args.seq_lens_ptr,
args.query_start_len_ptr,
args.num_seqs);
}
else
{
return Kernel::MakeKargs(args.q_ptr,
args.k_ptr,
args.v_ptr,
args.o_ptr,
args.num_blks,
args.num_head_q,
args.num_queries_per_kv,
args.scale_s,
total_num_q_blocks,
args.query_stride_0,
args.query_stride_1,
args.stride_k_cache_0,
args.stride_k_cache_1,
args.stride_k_cache_2,
args.stride_k_cache_3,
args.stride_v_cache_0,
args.stride_v_cache_1,
args.stride_v_cache_2,
args.stride_v_cache_3,
args.output_stride_0,
args.output_stride_1,
args.block_tables_ptr,
args.block_table_stride,
args.seq_lens_ptr,
args.query_start_len_ptr,
args.num_seqs);
}
};
auto kargs = kargs_lambda();
dim3 grids = Kernel::GridSize2D(args.num_head_q / args.num_queries_per_kv, total_num_q_blocks);
constexpr dim3 blocks = Kernel::BlockSize();
constexpr index_t kBlockPerCu = Kernel::kBlockPerCu;

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

@@ -25,7 +25,6 @@ struct UnifiedAttentionKernel
static constexpr ck_tile::index_t kBlockPerCu = UnifiedAttentionPipeline::kBlockPerCu;
static_assert(kBlockPerCu > 0);
using QDataType = ck_tile::remove_cvref_t<typename UnifiedAttentionPipeline::QDataType>;
using KDataType = ck_tile::remove_cvref_t<typename UnifiedAttentionPipeline::KDataType>;
using VDataType = ck_tile::remove_cvref_t<typename UnifiedAttentionPipeline::VDataType>;
@@ -39,7 +38,6 @@ struct UnifiedAttentionKernel
static constexpr bool kPadHeadDimQ = UnifiedAttentionPipeline::kPadHeadDimQ;
static constexpr bool kPadHeadDimV = UnifiedAttentionPipeline::kPadHeadDimV;
static constexpr auto QuantEnum = UnifiedAttentionPipeline::Problem::QuantEnum;
// TODO add yjese
static constexpr index_t HEAD_SIZE = UnifiedAttentionPipeline::HEAD_SIZE;
@@ -52,7 +50,7 @@ struct UnifiedAttentionKernel
static constexpr index_t BLOCK_Q = UnifiedAttentionPipeline::BLOCK_Q;
// BLOCK size for K seqlen
static constexpr index_t BLOCK_SIZE = UnifiedAttentionPipeline::BLOCK_SIZE;
static constexpr bool kIsQuantized = (QuantEnum != UnifiedAttentionQuantScaleEnum::NO_SCALE);
static constexpr bool kIsQuantized = (QuantEnum != UnifiedAttentionQuantScaleEnum::NO_SCALE);
template <ck_tile::index_t I> // to avoid duplicated base class problem, introduce an template
// arg
struct UnifiedAttentionEmptyKargs
@@ -102,10 +100,10 @@ struct UnifiedAttentionKernel
ck_tile::index_t output_stride_1;
};
struct UnifiedAttentionVarlenKargs
: UnifiedAttentionCommonKargs,
std::conditional_t<kIsQuantized, UnifiedAttentionQuantKargs, UnifiedAttentionEmptyKargs<0>>
struct UnifiedAttentionVarlenKargs : UnifiedAttentionCommonKargs,
std::conditional_t<kIsQuantized,
UnifiedAttentionQuantKargs,
UnifiedAttentionEmptyKargs<0>>
{
const int32_t* block_tables_ptr;
ck_tile::index_t block_table_stride;
@@ -299,9 +297,6 @@ struct UnifiedAttentionKernel
{
using namespace ck_tile;
// TODO fast exp
float scale_s = kargs.scale_s;
// allocate LDS
__shared__ char smem_ptr[GetSmemSize()];
@@ -485,8 +480,9 @@ struct UnifiedAttentionKernel
return FmhaMask{cur_batch_query_len, seq_len};
}();
// both quantized
float scale_s = kargs.scale_s;
// both quantized
auto o_acc_tile = [&]() {
if constexpr(kIsQuantized)
{

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

@@ -4,15 +4,6 @@
#pragma once
#include "ck_tile/ops/gemm/warp/warp_gemm.hpp"
#include "ck_tile/ops/gemm/warp/warp_gemm_dispatcher.hpp"
#include "ck_tile/ops/gemm/block/block_gemm_asmem_bsmem_creg_v1_custom_policy.hpp"
#include "ck_tile/ops/gemm/block/block_gemm_asmem_bsmem_creg_v1.hpp"
#include "ck_tile/ops/gemm/block/block_gemm_areg_bsmem_creg_v1_custom_policy.hpp"
#include "ck_tile/ops/gemm/block/block_gemm_areg_bsmem_creg_v2_custom_policy.hpp"
#include "ck_tile/ops/gemm/block/block_gemm_areg_bsmem_creg_v2.hpp"
#include "ck_tile/ops/gemm/block/block_gemm_areg_bsmem_creg_one_warp_v1.hpp"
#include "ck_tile/core.hpp"
#include "ck_tile/ops/gemm/block/block_gemm_areg_breg_creg_v2.hpp"
#include "ck_tile/ops/gemm/block/block_gemm_areg_breg_creg_v2_custom_policy.hpp"
@@ -272,7 +263,7 @@ struct UnifiedAttentionPipelineDefaultPolicy
{
/// NOTICE: in order to use load_tile_transpose() later for V tile, we cannot use
/// WarpGemmMfmaBf16Bf16F32M32N32K16SwizzleBTransposedCDistribution here
return WarpGemmMfmaFp8Fp8F32M32N32K32SwizzleBTransposedCDistribution<>{};
return WarpGemmMfma_f32_32x32x16_fp8_fp8_CTransposed{};
// return WarpGemmMfmaFp8Fp8F32M32N32K16SwizzleBTransposedCDistribution<>{};
}
}();
@@ -296,7 +287,7 @@ struct UnifiedAttentionPipelineDefaultPolicy
// Follow the V type for this gemm
using GemmProblem =
BlockGemmProblem<typename Problem::VDataType,
typename Problem::VDataType,
typename Problem::PDataType,
typename Problem::OaccDataType,
Problem::kBlockSize,
TileGemmShape<sequence<Problem::UnifiedAttentionShape::BLOCK_M,
@@ -306,21 +297,40 @@ struct UnifiedAttentionPipelineDefaultPolicy
typename Problem::UnifiedAttentionShape::Gemm1WarpTile>>;
/// NOTICE: in order to use load_tile_transpose() later for V tiles, we have to pass
/// WGAttrNumAccessEnum::Double instead of WGAttrNumAccessEnum::Single
using WarpGemm =
WarpGemmDispatcher<typename Problem::VDataType,
typename Problem::VDataType,
typename Problem::OaccDataType,
Problem::UnifiedAttentionShape::Gemm1WarpTile::at(number<0>{}),
Problem::UnifiedAttentionShape::Gemm1WarpTile::at(number<1>{}),
Problem::UnifiedAttentionShape::Gemm1WarpTile::at(number<2>{}),
true,
false,
false,
WGAttrNumAccessEnum::Double>;
using WarpGemm = decltype([]() {
if constexpr(std::is_same_v<typename Problem::VDataType, fp8_t>)
{
return WarpGemmDispatcher<
typename Problem::VDataType,
typename Problem::PDataType,
typename Problem::OaccDataType,
Problem::UnifiedAttentionShape::Gemm1WarpTile::at(number<0>{}),
Problem::UnifiedAttentionShape::Gemm1WarpTile::at(number<1>{}),
Problem::UnifiedAttentionShape::Gemm1WarpTile::at(number<2>{}),
true,
false,
false,
WGAttrNumAccessEnum::Single>{};
}
else
{
return WarpGemmDispatcher<
typename Problem::VDataType,
typename Problem::PDataType,
typename Problem::OaccDataType,
Problem::UnifiedAttentionShape::Gemm1WarpTile::at(number<0>{}),
Problem::UnifiedAttentionShape::Gemm1WarpTile::at(number<1>{}),
Problem::UnifiedAttentionShape::Gemm1WarpTile::at(number<2>{}),
true,
false,
false,
WGAttrNumAccessEnum::Double>{};
}
}());
using BlockGemmPolicy = BlockGemmARegBRegCRegV2CustomPolicy<
typename Problem::VDataType,
typename Problem::VDataType,
typename Problem::PDataType,
typename Problem::OaccDataType,
typename Problem::UnifiedAttentionShape::Gemm1BlockWarps,
WarpGemm,
@@ -427,6 +437,20 @@ struct UnifiedAttentionPipelineDefaultPolicy
number<KPack>{},
number<1>{});
// Debug: verify constexpr components
// constexpr index_t NumIssues = kNPerBlock / (LaneGroups * NumWarps);
static_assert(kNPerBlock > 0, "kNPerBlock must be constexpr");
static_assert(LaneGroups > 0, "LaneGroups must be constexpr");
static_assert(NumWarps > 0, "NumWarps must be constexpr");
// static_assert(NumIssues > 0, "NumIssues must be constexpr");
static_assert(LaneGroups > 0, "LaneGroups must be constexpr");
static_assert(NumWarps > 0, "NumWarps must be constexpr");
static_assert(kKPerBlock > 0, "kKPerBlock must be constexpr");
static_assert(KPack > 0 && kKPerBlock % KPack == 0,
"KPack must be constexpr and divide kKPerBlock");
constexpr auto k_lds_block_desc = transform_tensor_descriptor(
k_lds_block_desc_0,
make_tuple(