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

New gemms

Browse Source
This commit is contained in:
Tianxing Wu
2025-12-10 13:40:42 +00:00
parent 5f9c363ba8
commit 12a2420994
7 changed files with 19 additions and 53 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
example/ck_tile/01_unified_attention/example_unified_attention.cpp
View File

@@ -159,9 +159,6 @@ struct Problem
scale_s = 1.0f / ck_tile::sqrt(static_cast<float>(hdim));
// Initialize other scales
scale = args.get_float("scale");
scale_k = args.get_float("scale_k");
scale_v = args.get_float("scale_v");
num_tokens = 0;
for(const auto& len : query_lens)
{
@@ -192,9 +189,6 @@ struct Problem
ck_tile::index_t num_tokens;
ck_tile::index_t fp8_mode;
float scale_s;
float scale;
float scale_k;
float scale_v;
mask_info mask;
std::vector<int> query_lens;
std::vector<int> kv_lens;

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

@@ -178,7 +178,7 @@ float unified_attention_kernel_launch(const unified_attention_args& args,
"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)
if constexpr(Kernel::IsQuantized)
{
return Kernel::MakeKargs(args.q_ptr,
args.k_ptr,

4
include/ck_tile/ops/gemm/warp/warp_gemm.hpp
View File

@@ -374,9 +374,11 @@ using WarpGemmMfma_f32_32x32x64_bf8_bf8 = WarpGemmImpl<
WarpGemmAttributeMfma<WarpGemmAttributeMfmaImpl_f32_32x32x64_bf8_bf8<WGAttrCtlEnum::Default_>,
AttrNumAccess>>;
template <WGAttrNumAccessEnum AttrNumAccess = WGAttrNumAccessEnum::Single>
using WarpGemmMfma_f32_32x32x16_fp8_fp8_CTransposed =
WarpGemmImpl<WarpGemmAttributeMfmaTransposedCDistribution<
WarpGemmAttributeMfmaImpl_f32_32x32x16_fp8_fp8<WGAttrCtlEnum::Default_>>>;
WarpGemmAttributeMfmaImpl_f32_32x32x16_fp8_fp8<WGAttrCtlEnum::Default_>,
AttrNumAccess>>;
using WarpGemmMfma_f32_32x32x16_fp8_bf8_CTransposed =
WarpGemmImpl<WarpGemmAttributeMfmaTransposedCDistribution<

3
include/ck_tile/ops/gemm/warp/warp_gemm_dispatcher.hpp
View File

@@ -100,7 +100,8 @@ template<> struct Dispatcher<fp8_t, fp8_t, float, 32, 32, 16, false> { using Ty
template<> struct Dispatcher<fp8_t, fp8_t, float, 32, 32, 32, false> { using Type = WarpGemmMfma_f32_32x32x32_fp8_fp8; };
template<> struct Dispatcher<fp8_t, fp8_t, float, 16, 16, 32, false> { using Type = WarpGemmMfma_f32_16x16x32_fp8_fp8; };
template<> struct Dispatcher<fp8_t, fp8_t, float, 16, 16, 64, false> { using Type = WarpGemmMfma_f32_16x16x64_fp8_fp8; };
template<> struct Dispatcher<fp8_t, fp8_t, float, 32, 32, 16, true> { using Type = WarpGemmMfma_f32_32x32x16_fp8_fp8_CTransposed; };
template<> struct Dispatcher<fp8_t, fp8_t, float, 32, 32, 16, true> { using Type = WarpGemmMfma_f32_32x32x16_fp8_fp8_CTransposed<>; };
template<> struct Dispatcher<fp8_t, fp8_t, float, 32, 32, 16, true, false, false, EDouble> { using Type = WarpGemmMfma_f32_32x32x16_fp8_fp8_CTransposed<EDouble>; };
template<> struct Dispatcher<fp8_t, fp8_t, float, 16, 16, 32, true> { using Type = WarpGemmMfma_f32_16x16x32_fp8_fp8_CTransposed; };
template<> struct Dispatcher<fp8_t, bf8_t, float, 32, 32, 16, false> { using Type = WarpGemmMfma_f32_32x32x16_fp8_bf8; };
template<> struct Dispatcher<fp8_t, bf8_t, float, 32, 32, 16, true> { using Type = WarpGemmMfma_f32_32x32x16_fp8_bf8_CTransposed; };

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

@@ -50,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 IsQuantized = (QuantEnum != UnifiedAttentionQuantScaleEnum::NO_SCALE);
template <ck_tile::index_t I> // to avoid duplicated base class problem, introduce an template
// arg
struct UnifiedAttentionEmptyKargs
@@ -100,10 +100,9 @@ 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<IsQuantized, UnifiedAttentionQuantKargs, UnifiedAttentionEmptyKargs<0>>
{
const int32_t* block_tables_ptr;
ck_tile::index_t block_table_stride;
@@ -115,7 +114,7 @@ struct UnifiedAttentionKernel
using Kargs = UnifiedAttentionVarlenKargs;
template <bool Cond = !kIsQuantized>
template <bool Cond = !IsQuantized>
CK_TILE_HOST static constexpr std::enable_if_t<Cond, Kargs>
MakeKargs(const void* q_ptr,
const void* k_ptr,
@@ -175,7 +174,7 @@ struct UnifiedAttentionKernel
return kargs;
}
template <bool Cond = kIsQuantized>
template <bool Cond = IsQuantized>
CK_TILE_HOST static constexpr std::enable_if_t<Cond, Kargs>
MakeKargs(const void* q_ptr,
const void* k_ptr,
@@ -484,7 +483,7 @@ struct UnifiedAttentionKernel
// both quantized
auto o_acc_tile = [&]() {
if constexpr(kIsQuantized)
if constexpr(IsQuantized)
{
if(std::is_same_v<QDataType, KDataType>)
{

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

@@ -869,7 +869,7 @@ struct UnifiedAttentionPipeline
get_slice_tile(q_tile,
sequence<0, (k0_loops - 1) * HEAD_SIZE_PADDED>{},
sequence<BLOCK_M, k0_loops * HEAD_SIZE_PADDED>{}),
dequantized_k_tile);
dequantized_k_tile * scale_k);
}
else
{
@@ -897,7 +897,7 @@ struct UnifiedAttentionPipeline
get_slice_tile(sp(sp_reg_idx).p,
sequence<0, (k1_loops - 1) * BLOCK_SIZE>{},
sequence<BLOCK_M, k1_loops * BLOCK_SIZE>{}),
dequantized_v_tile);
dequantized_v_tile * scale_v);
}
else
{
@@ -929,7 +929,7 @@ struct UnifiedAttentionPipeline
get_slice_tile(q_tile,
sequence<0, (k0_loops - 1) * HEAD_SIZE_PADDED>{},
sequence<BLOCK_M, k0_loops * HEAD_SIZE_PADDED>{}),
dequantized_k_tile);
dequantized_k_tile * scale_k);
}
else
{
@@ -957,7 +957,7 @@ struct UnifiedAttentionPipeline
get_slice_tile(sp(sp_reg_idx).p,
sequence<0, (k1_loops - 1) * BLOCK_SIZE>{},
sequence<BLOCK_M, k1_loops * BLOCK_SIZE>{}),
dequantized_v_tile);
dequantized_v_tile * scale_v);
fmha_alu0(number<1>{} - sp_reg_idx);
}
else

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

@@ -264,7 +264,7 @@ struct UnifiedAttentionPipelineDefaultPolicy
{
/// NOTICE: in order to use load_tile_transpose() later for V tile, we cannot use
/// WarpGemmMfmaBf16Bf16F32M32N32K16SwizzleBTransposedCDistribution here
return WarpGemmMfma_f32_32x32x16_fp8_fp8_CTransposed{};
return WarpGemmMfma_f32_32x32x16_fp8_fp8_CTransposed<>{};
// return WarpGemmMfmaFp8Fp8F32M32N32K16SwizzleBTransposedCDistribution<>{};
}
}();
@@ -298,36 +298,6 @@ 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 = decltype([]() {
// if constexpr(std::is_same_v<typename Problem::VDataType, fp8_t>)
// {
// return WarpGemmDispatcher<
// typename Problem::PDataType,
// 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::Single>{};
// }
// else
// {
// return WarpGemmDispatcher<
// typename Problem::PDataType,
// 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 =
WarpGemmDispatcher<typename Problem::PDataType,
typename Problem::VDataType,
@@ -338,7 +308,7 @@ struct UnifiedAttentionPipelineDefaultPolicy
true,
false,
false,
WGAttrNumAccessEnum::Single>;
WGAttrNumAccessEnum::Double>;
using BlockGemmPolicy = BlockGemmARegBRegCRegV2CustomPolicy<
typename Problem::PDataType,
typename Problem::VDataType,