diff --git a/example/ck_tile/01_unified_attention/example_unified_attention.cpp b/example/ck_tile/01_unified_attention/example_unified_attention.cpp index b178d8ef83..36bbd31f63 100644 --- a/example/ck_tile/01_unified_attention/example_unified_attention.cpp +++ b/example/ck_tile/01_unified_attention/example_unified_attention.cpp @@ -159,9 +159,6 @@ struct Problem scale_s = 1.0f / ck_tile::sqrt(static_cast(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 query_lens; std::vector kv_lens; 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 36adf7ccd5..37193947f6 100644 --- a/example/ck_tile/01_unified_attention/unified_attention_impl.hpp +++ b/example/ck_tile/01_unified_attention/unified_attention_impl.hpp @@ -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, diff --git a/include/ck_tile/ops/gemm/warp/warp_gemm.hpp b/include/ck_tile/ops/gemm/warp/warp_gemm.hpp index a2c320f3e6..fb59fa6fa1 100644 --- a/include/ck_tile/ops/gemm/warp/warp_gemm.hpp +++ b/include/ck_tile/ops/gemm/warp/warp_gemm.hpp @@ -374,9 +374,11 @@ using WarpGemmMfma_f32_32x32x64_bf8_bf8 = WarpGemmImpl< WarpGemmAttributeMfma, AttrNumAccess>>; +template using WarpGemmMfma_f32_32x32x16_fp8_fp8_CTransposed = WarpGemmImpl>>; + WarpGemmAttributeMfmaImpl_f32_32x32x16_fp8_fp8, + AttrNumAccess>>; using WarpGemmMfma_f32_32x32x16_fp8_bf8_CTransposed = WarpGemmImpl struct Dispatcher { using Ty template<> struct Dispatcher { using Type = WarpGemmMfma_f32_32x32x32_fp8_fp8; }; template<> struct Dispatcher { using Type = WarpGemmMfma_f32_16x16x32_fp8_fp8; }; template<> struct Dispatcher { using Type = WarpGemmMfma_f32_16x16x64_fp8_fp8; }; -template<> struct Dispatcher { using Type = WarpGemmMfma_f32_32x32x16_fp8_fp8_CTransposed; }; +template<> struct Dispatcher { using Type = WarpGemmMfma_f32_32x32x16_fp8_fp8_CTransposed<>; }; +template<> struct Dispatcher { using Type = WarpGemmMfma_f32_32x32x16_fp8_fp8_CTransposed; }; template<> struct Dispatcher { using Type = WarpGemmMfma_f32_16x16x32_fp8_fp8_CTransposed; }; template<> struct Dispatcher { using Type = WarpGemmMfma_f32_32x32x16_fp8_bf8; }; template<> struct Dispatcher { using Type = WarpGemmMfma_f32_32x32x16_fp8_bf8_CTransposed; }; 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 0e1305809e..f14039483d 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 @@ -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 // 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> + struct UnifiedAttentionVarlenKargs + : UnifiedAttentionCommonKargs, + std::conditional_t> { const int32_t* block_tables_ptr; ck_tile::index_t block_table_stride; @@ -115,7 +114,7 @@ struct UnifiedAttentionKernel using Kargs = UnifiedAttentionVarlenKargs; - template + template CK_TILE_HOST static constexpr std::enable_if_t MakeKargs(const void* q_ptr, const void* k_ptr, @@ -175,7 +174,7 @@ struct UnifiedAttentionKernel return kargs; } - template + template CK_TILE_HOST static constexpr std::enable_if_t 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) { diff --git a/include/ck_tile/ops/unified_attention/pipeline/unified_attention_pipeline.hpp b/include/ck_tile/ops/unified_attention/pipeline/unified_attention_pipeline.hpp index ebc3888155..0eb05d1d1a 100644 --- a/include/ck_tile/ops/unified_attention/pipeline/unified_attention_pipeline.hpp +++ b/include/ck_tile/ops/unified_attention/pipeline/unified_attention_pipeline.hpp @@ -869,7 +869,7 @@ struct UnifiedAttentionPipeline get_slice_tile(q_tile, sequence<0, (k0_loops - 1) * HEAD_SIZE_PADDED>{}, sequence{}), - 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{}), - 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{}), - 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{}), - dequantized_v_tile); + dequantized_v_tile * scale_v); fmha_alu0(number<1>{} - sp_reg_idx); } else 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 609444e3c0..c41d430fa7 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 @@ -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) - // { - // 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; + WGAttrNumAccessEnum::Double>; using BlockGemmPolicy = BlockGemmARegBRegCRegV2CustomPolicy< typename Problem::PDataType, typename Problem::VDataType,