diff --git a/example/ck_tile/01_unified_attention/instances/unified_attention_d128_bf16_fp8_mask.cpp b/example/ck_tile/01_unified_attention/instances/unified_attention_d128_bf16_fp8_mask.cpp new file mode 100644 index 0000000000..ead08b4678 --- /dev/null +++ b/example/ck_tile/01_unified_attention/instances/unified_attention_d128_bf16_fp8_mask.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_kernel_traits; + +INST_UNIFIED_ATTENTION_DISPATCH(kernel_traits) + +} // namespace ck_tile diff --git a/example/ck_tile/01_unified_attention/instances/unified_attention_d128_bf16_fp8_nmask.cpp b/example/ck_tile/01_unified_attention/instances/unified_attention_d128_bf16_fp8_nmask.cpp new file mode 100644 index 0000000000..739f090277 --- /dev/null +++ b/example/ck_tile/01_unified_attention/instances/unified_attention_d128_bf16_fp8_nmask.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_kernel_traits; + +INST_UNIFIED_ATTENTION_DISPATCH(kernel_traits) + +} // namespace ck_tile diff --git a/example/ck_tile/01_unified_attention/instances/unified_attention_d128_bf16_mask.cpp b/example/ck_tile/01_unified_attention/instances/unified_attention_d128_bf16_mask.cpp index ead08b4678..07d929998f 100644 --- a/example/ck_tile/01_unified_attention/instances/unified_attention_d128_bf16_mask.cpp +++ b/example/ck_tile/01_unified_attention/instances/unified_attention_d128_bf16_mask.cpp @@ -7,7 +7,7 @@ namespace ck_tile { using kernel_traits = - unified_attention_kernel_traits; + unified_attention_kernel_traits; INST_UNIFIED_ATTENTION_DISPATCH(kernel_traits) diff --git a/example/ck_tile/01_unified_attention/instances/unified_attention_d128_bf16_nmask.cpp b/example/ck_tile/01_unified_attention/instances/unified_attention_d128_bf16_nmask.cpp index 739f090277..9e0a0623ed 100644 --- a/example/ck_tile/01_unified_attention/instances/unified_attention_d128_bf16_nmask.cpp +++ b/example/ck_tile/01_unified_attention/instances/unified_attention_d128_bf16_nmask.cpp @@ -7,7 +7,7 @@ namespace ck_tile { using kernel_traits = - unified_attention_kernel_traits; + unified_attention_kernel_traits; INST_UNIFIED_ATTENTION_DISPATCH(kernel_traits) diff --git a/example/ck_tile/01_unified_attention/instances/unified_attention_d128_fp16_fp8_mask.cpp b/example/ck_tile/01_unified_attention/instances/unified_attention_d128_fp16_fp8_mask.cpp new file mode 100644 index 0000000000..6399d650cf --- /dev/null +++ b/example/ck_tile/01_unified_attention/instances/unified_attention_d128_fp16_fp8_mask.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_kernel_traits; + +INST_UNIFIED_ATTENTION_DISPATCH(kernel_traits) + +} // namespace ck_tile diff --git a/example/ck_tile/01_unified_attention/instances/unified_attention_d128_fp16_fp8_nmask.cpp b/example/ck_tile/01_unified_attention/instances/unified_attention_d128_fp16_fp8_nmask.cpp new file mode 100644 index 0000000000..e925f68d19 --- /dev/null +++ b/example/ck_tile/01_unified_attention/instances/unified_attention_d128_fp16_fp8_nmask.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_kernel_traits; + +INST_UNIFIED_ATTENTION_DISPATCH(kernel_traits) + +} // namespace ck_tile diff --git a/example/ck_tile/01_unified_attention/instances/unified_attention_d128_fp16_mask.cpp b/example/ck_tile/01_unified_attention/instances/unified_attention_d128_fp16_mask.cpp index 6399d650cf..e0f1fe2c25 100644 --- a/example/ck_tile/01_unified_attention/instances/unified_attention_d128_fp16_mask.cpp +++ b/example/ck_tile/01_unified_attention/instances/unified_attention_d128_fp16_mask.cpp @@ -7,7 +7,7 @@ namespace ck_tile { using kernel_traits = - unified_attention_kernel_traits; + unified_attention_kernel_traits; INST_UNIFIED_ATTENTION_DISPATCH(kernel_traits) diff --git a/example/ck_tile/01_unified_attention/instances/unified_attention_d128_fp16_nmask.cpp b/example/ck_tile/01_unified_attention/instances/unified_attention_d128_fp16_nmask.cpp index e925f68d19..509a403849 100644 --- a/example/ck_tile/01_unified_attention/instances/unified_attention_d128_fp16_nmask.cpp +++ b/example/ck_tile/01_unified_attention/instances/unified_attention_d128_fp16_nmask.cpp @@ -7,7 +7,7 @@ namespace ck_tile { using kernel_traits = - unified_attention_kernel_traits; + unified_attention_kernel_traits; INST_UNIFIED_ATTENTION_DISPATCH(kernel_traits) diff --git a/example/ck_tile/01_unified_attention/instances/unified_attention_d128_q_bf16_fp8_mask.cpp b/example/ck_tile/01_unified_attention/instances/unified_attention_d128_q_bf16_fp8_mask.cpp new file mode 100644 index 0000000000..0f74e0a5fe --- /dev/null +++ b/example/ck_tile/01_unified_attention/instances/unified_attention_d128_q_bf16_fp8_mask.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_kernel_traits; + +INST_UNIFIED_ATTENTION_DISPATCH(kernel_traits) + +} // namespace ck_tile diff --git a/example/ck_tile/01_unified_attention/instances/unified_attention_d128_q_bf16_fp8_nmask.cpp b/example/ck_tile/01_unified_attention/instances/unified_attention_d128_q_bf16_fp8_nmask.cpp new file mode 100644 index 0000000000..97ca9d9dcd --- /dev/null +++ b/example/ck_tile/01_unified_attention/instances/unified_attention_d128_q_bf16_fp8_nmask.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_kernel_traits; + +INST_UNIFIED_ATTENTION_DISPATCH(kernel_traits) + +} // namespace ck_tile diff --git a/example/ck_tile/01_unified_attention/instances/unified_attention_d128_q_fp16_fp8_mask.cpp b/example/ck_tile/01_unified_attention/instances/unified_attention_d128_q_fp16_fp8_mask.cpp new file mode 100644 index 0000000000..dd6a3fdcde --- /dev/null +++ b/example/ck_tile/01_unified_attention/instances/unified_attention_d128_q_fp16_fp8_mask.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_kernel_traits; + +INST_UNIFIED_ATTENTION_DISPATCH(kernel_traits) + +} // namespace ck_tile diff --git a/example/ck_tile/01_unified_attention/instances/unified_attention_d128_q_fp16_fp8_nmask.cpp b/example/ck_tile/01_unified_attention/instances/unified_attention_d128_q_fp16_fp8_nmask.cpp new file mode 100644 index 0000000000..a0562bc265 --- /dev/null +++ b/example/ck_tile/01_unified_attention/instances/unified_attention_d128_q_fp16_fp8_nmask.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_kernel_traits; + +INST_UNIFIED_ATTENTION_DISPATCH(kernel_traits) + +} // namespace ck_tile diff --git a/example/ck_tile/01_unified_attention/unified_attention_impl.hpp b/example/ck_tile/01_unified_attention/unified_attention_impl.hpp index 94c2fe0a0d..947ecb136d 100644 --- a/example/ck_tile/01_unified_attention/unified_attention_impl.hpp +++ b/example/ck_tile/01_unified_attention/unified_attention_impl.hpp @@ -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; diff --git a/include/ck_tile/ops/unified_attention/kernel/unified_attention_kernel.hpp b/include/ck_tile/ops/unified_attention/kernel/unified_attention_kernel.hpp index 8fa430ff5e..0e1305809e 100644 --- a/include/ck_tile/ops/unified_attention/kernel/unified_attention_kernel.hpp +++ b/include/ck_tile/ops/unified_attention/kernel/unified_attention_kernel.hpp @@ -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; using KDataType = ck_tile::remove_cvref_t; using VDataType = ck_tile::remove_cvref_t; @@ -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 // 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> + struct UnifiedAttentionVarlenKargs : UnifiedAttentionCommonKargs, + std::conditional_t> { 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) { diff --git a/include/ck_tile/ops/unified_attention/pipeline/unified_attention_pipeline_default_policy.hpp b/include/ck_tile/ops/unified_attention/pipeline/unified_attention_pipeline_default_policy.hpp index d7eb369d0f..12b691680f 100644 --- a/include/ck_tile/ops/unified_attention/pipeline/unified_attention_pipeline_default_policy.hpp +++ b/include/ck_tile/ops/unified_attention/pipeline/unified_attention_pipeline_default_policy.hpp @@ -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>; /// 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{}), - 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) + { + 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{}, 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(