diff --git a/example/ck_tile/01_fmha/codegen/ops/fmha_fwd.py b/example/ck_tile/01_fmha/codegen/ops/fmha_fwd.py index ecaca2683a..38ba926b4e 100644 --- a/example/ck_tile/01_fmha/codegen/ops/fmha_fwd.py +++ b/example/ck_tile/01_fmha/codegen/ops/fmha_fwd.py @@ -355,12 +355,18 @@ class FmhaFwdApiTrait: @property def dcheck(self) -> str: - if self.pipeline_tag == "qr_hpad": + if self.pipeline_tag == "qr_async": + vec = int((32 * 4) / DTYPE_BITS[self.dtype]) + if self.dpad == "t": + return f"a.hdim_q % {vec} == 0" + else: + assert False + elif self.pipeline_tag == "qr_hpad": if self.dpad == "t": return "a.hdim_q % 8 == 0" else: assert False - elif self.pipeline_tag in ["qr", "qs", "qr_async", "qr_async_trload", "qr_async_trload_v3"]: + elif self.pipeline_tag in ["qr", "qs", "qr_async_trload", "qr_async_trload_v3"]: bk0submax = K0_MAX_SUBMAX_MAP[self.bk0max] if self.dpad == "t": return f"true /*a.hdim_q % {bk0submax} != 0*/" # TODO: order of get_pipelines() matters! (ugly) @@ -371,12 +377,18 @@ class FmhaFwdApiTrait: @property def dvcheck(self) -> str: - if self.pipeline_tag == "qr_hpad": + if self.pipeline_tag == "qr_async": + vec = int((32 * 4) / DTYPE_BITS[self.dtype]) + if self.dvpad == "t": + return f"a.hdim_v % {vec} == 0" + else: + assert False + elif self.pipeline_tag == "qr_hpad": if self.dvpad == "t": return "a.hdim_v % 8 == 0" else: assert False - elif self.pipeline_tag in ["qr", "qs", "qr_async", "qr_async_trload", "qr_async_trload_v3"]: + elif self.pipeline_tag in ["qr", "qs", "qr_async_trload", "qr_async_trload_v3"]: bk0submax = K0_MAX_SUBMAX_MAP[self.bk0max] if self.dvpad == "t": return f"true /*a.hdim_v % {bk0submax} != 0*/" # TODO: order of get_pipelines() matters! (ugly) @@ -1034,7 +1046,7 @@ class KernelComponentFactoryGfx9(CompatibilityRuleFactoryGfx9): pipelines.append(FmhaFwdPipeline("qr", "row", "f", "f", "f", "f", logits, bias, lse, dropout, qscale, mask, skip, "f", sink)) # fmt: skip pipelines.append(FmhaFwdPipeline("qr", "row", "t", "t", "t", "t", logits, bias, lse, dropout, qscale, mask, skip, "f", sink)) # fmt: skip else: - pipelines.append(FmhaFwdPipeline("qr_async", "row", "f", "f", "f", "f", logits, bias, lse, dropout, qscale, mask, skip, "f", sink)) # fmt: skip + pipelines.append(FmhaFwdPipeline("qr_async", "row", "t", "f", "t", "t", logits, bias, lse, dropout, qscale, mask, skip, "f", sink)) # fmt: skip pipelines.append(FmhaFwdPipeline("qr_async", "row", "t", "t", "t", "t", logits, bias, lse, dropout, qscale, mask, skip, "f", sink)) # fmt: skip if receipt == 1 and bias != "bias": pipelines.append(FmhaFwdPipeline("qr", "row", "t", "t", "t", "t", logits, bias, lse, dropout, qscale, mask, skip, "f", sink)) # fmt: skip # TODO: cover arbitraty hdim# fmt: skip @@ -1048,10 +1060,10 @@ class KernelComponentFactoryGfx9(CompatibilityRuleFactoryGfx9): ["f", "t"], ): if hdim == 64: - pipelines.append(FmhaFwdPipeline("qr", "row", "f", "f", "f", "f", logits, bias, "f", "f", qscale, mask, "f", "f", sink)) # fmt: skip + pipelines.append(FmhaFwdPipeline("qr", "row", "t", "f", "t", "t", logits, bias, "f", "f", qscale, mask, "f", "f", sink)) # fmt: skip pipelines.append(FmhaFwdPipeline("qr", "row", "t", "t", "t", "t", logits, bias, "f", "f", qscale, mask, "f", "f", sink)) # fmt: skip else: - pipelines.append(FmhaFwdPipeline("qr_async", "row", "f", "f", "f", "f", logits, bias, "f", "f", qscale, mask, "f", "f", sink)) # fmt: skip + pipelines.append(FmhaFwdPipeline("qr_async", "row", "t", "f", "t", "t", logits, bias, "f", "f", qscale, mask, "f", "f", sink)) # fmt: skip pipelines.append(FmhaFwdPipeline("qr_async", "row", "t", "t", "t", "t", logits, bias, "f", "f", qscale, mask, "f", "f", sink)) # fmt: skip return pipelines @@ -1157,9 +1169,6 @@ class KernelComponentFactoryGfx950( ): pipelines.append(FmhaFwdPipeline("qr", "col", "f", "f", "f", "f", logits, bias, lse, dropout, qscale, mask, "f", "f", sink)) # fmt: skip pipelines.append(FmhaFwdPipeline("qr", "col", "t", "t", "t", "t", logits, bias, lse, dropout, qscale, mask, "f", "f", sink)) # fmt: skip - if hdim > 64 and dtype in cls._DT_MXFP8: - pipelines.append(FmhaFwdPipeline("qr_async", "col", "f", "f", "f", "f", logits, bias, lse, dropout, qscale, mask, "f", "f", sink)) # fmt: skip - pipelines.append(FmhaFwdPipeline("qr_async", "col", "t", "t", "t", "t", logits, bias, lse, dropout, qscale, mask, "f", "f", sink)) # fmt: skip return pipelines diff --git a/example/ck_tile/01_fmha/script/benchmark_fwd.sh b/example/ck_tile/01_fmha/script/benchmark_fwd.sh index 12dc305ab5..f2b3a48dcb 100755 --- a/example/ck_tile/01_fmha/script/benchmark_fwd.sh +++ b/example/ck_tile/01_fmha/script/benchmark_fwd.sh @@ -53,4 +53,4 @@ $EXE $base_group_args -s_qpad=1152,896,576,320 -s_kpad=1152,896,576,320 $EXE $base_group_args -s_qpad=1536,1152,768,384 -s_kpad=1536,1152,768,384 # high physical pad -$EXE $base_group_args -s_qpad=2048,1536,1024,512 -s_kpad=2048,1536,1024,512 \ No newline at end of file +$EXE $base_group_args -s_qpad=2048,1536,1024,512 -s_kpad=2048,1536,1024,512 diff --git a/include/ck_tile/core/tensor/tile_window.hpp b/include/ck_tile/core/tensor/tile_window.hpp index bb91e552d5..7dcf498695 100644 --- a/include/ck_tile/core/tensor/tile_window.hpp +++ b/include/ck_tile/core/tensor/tile_window.hpp @@ -506,8 +506,9 @@ struct tile_window_with_static_distribution using SFC_Ys = typename Traits::SFC_Ys; static constexpr index_t YElementSize = typename Base::TileDstr{}.get_ys_to_d_descriptor().get_element_space_size(); - static_assert(YElementSize % Traits::ScalarPerVector == 0); - using vectorized_tbuf = array; + static_assert(YElementSize % (Traits::PackedSize * Traits::ScalarPerVector) == 0); + using vectorized_tbuf = + array; constexpr auto tile_dstr = typename Base::TileDstr{}; @@ -533,11 +534,10 @@ struct tile_window_with_static_distribution constexpr index_t d = tile_dstr.get_ys_to_d_descriptor().calculate_offset(idx_ys_start) / Traits::PackedSize; - static_assert(Traits::ScalarPerVector % Traits::PackedSize == 0); - static_assert(d % (Traits::ScalarPerVector / Traits::PackedSize) == 0); + static_assert(d % Traits::ScalarPerVector == 0); this->get_bottom_tensor_view().template get_vectorized_elements_raw( - dst_vec_tbuf.template at(), + dst_vec_tbuf.template at(), bottom_tensor_thread_coord, 0 /**/, bool_constant{}, diff --git a/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_async.hpp b/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_async.hpp index b621e9ceab..5155794b16 100644 --- a/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_async.hpp +++ b/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_async.hpp @@ -7,7 +7,6 @@ #include "ck_tile/ops/common/tensor_layout.hpp" #include "ck_tile/ops/fmha/block/block_attention_bias_enum.hpp" #include "ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_async_default_policy.hpp" -#include "ck_tile/ops/fmha/block/cast_tile_mx.hpp" #include "ck_tile/ops/fmha/block/block_dropout.hpp" #include "ck_tile/ops/reduce/block/block_reduce.hpp" @@ -30,10 +29,6 @@ struct BlockFmhaPipelineQRKSVSAsync using PDataType = remove_cvref_t; using OaccDataType = remove_cvref_t; using ODataType = remove_cvref_t; - using QScaleDataType = remove_cvref_t; - using KScaleDataType = remove_cvref_t; - using VScaleDataType = remove_cvref_t; - using PScaleDataType = remove_cvref_t; using AttentionVariant = remove_cvref_t; using FmhaMask = remove_cvref_t; @@ -55,20 +50,21 @@ struct BlockFmhaPipelineQRKSVSAsync static_assert(kSubQKHeaddim <= 256, "hdim bigger than 256 is not suitable for this pipeline!"); - static constexpr bool kIsGroupMode = Problem::kIsGroupMode; - static constexpr bool kPadSeqLenQ = Problem::kPadSeqLenQ; + static constexpr bool kIsGroupMode = Problem::kIsGroupMode; + // TODO: seq_q always support padding, hdim_q/v support multiple of vector(like 8x) + // only need special care about seq_k padding (oob need set -INF of p instead of zero) + static_assert(Problem::kPadSeqLenQ == true && Problem::kPadHeadDimQ == true && + Problem::kPadHeadDimV == true); + static constexpr bool kPadSeqLenQ = true; static constexpr bool kPadSeqLenK = Problem::kPadSeqLenK; - static constexpr bool kPadHeadDimQ = Problem::kPadHeadDimQ; - static constexpr bool kPadHeadDimV = Problem::kPadHeadDimV; + static constexpr bool kPadHeadDimQ = true; // support multiple of vector(like 8x) + static constexpr bool kPadHeadDimV = true; // support multiple of vector(like 8x) static constexpr bool kHasLogitsSoftCap = Problem::kHasLogitsSoftCap; static constexpr auto BiasEnum = Problem::BiasEnum; static constexpr bool kStoreLSE = Problem::kStoreLSE; static constexpr bool kHasDropout = Problem::kHasDropout; static constexpr bool kHasSink = Problem::kHasSink; - static constexpr ck_tile::index_t kQKScaleGranularity = Problem::kQKScaleGranularity; - static constexpr ck_tile::index_t kVScaleGranularity = Problem::kVScaleGranularity; - // For BLOCKSCALE: shift value for exp2(x + shift) to scale P to [0, 2^shift] static constexpr float OCP_FP8_SHIFT = 8.0f; static constexpr float FNUZ_FP8_SHIFT = 7.0f; @@ -86,8 +82,7 @@ struct BlockFmhaPipelineQRKSVSAsync if constexpr(std::is_same_v) return Policy::template GetAlignmentV(); else - return kPadSeqLenK ? numeric_traits::PackedSize - : Policy::template GetAlignmentV(); + return kPadSeqLenK ? 1 : Policy::template GetAlignmentV(); }(); static constexpr index_t kAlignmentO = Policy::template GetAlignmentO(); static constexpr index_t kAlignmentBias = @@ -206,12 +201,9 @@ struct BlockFmhaPipelineQRKSVSAsync const float* k_descale_ptr, const float* v_descale_ptr, const index_t block_scale_size_kv, - const QScaleDramBlockWindowTmp& - q_scale_dram_block_window_tmp, // M0*(K0/kQKScaleGranularity) tile - const KScaleDramBlockWindowTmp& - k_scale_dram_block_window_tmp, // N0*(K0/kQKScaleGranularity) tile - const VScaleDramBlockWindowTmp& - v_scale_dram_block_window_tmp, // N1*(K1/kVScaleGranularity) tile + const QScaleDramBlockWindowTmp&, // M0*(K0/kQKScaleGranularity) tile + const KScaleDramBlockWindowTmp&, // N0*(K0/kQKScaleGranularity) tile + const VScaleDramBlockWindowTmp&, // N1*(K1/kVScaleGranularity) tile const float sink_v) const { static_assert( @@ -221,8 +213,6 @@ struct BlockFmhaPipelineQRKSVSAsync "wrong!"); static_assert(kM0 == QDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] && - kSubQKHeaddim == - QDramBlockWindowTmp{}.get_window_lengths()[number<1>{}] && kN0 == KDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] && kK0 == KDramBlockWindowTmp{}.get_window_lengths()[number<1>{}] && kN1 == VDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] && @@ -230,28 +220,8 @@ struct BlockFmhaPipelineQRKSVSAsync kM0 == BiasDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] && kN0 == BiasDramBlockWindowTmp{}.get_window_lengths()[number<1>{}], "wrong!"); - if constexpr(QScaleEnum == BlockAttentionQuantScaleEnum::MX) - { - static_assert(std::is_same_v); - static_assert( - std::is_same_v> && - std::is_same_v> && - std::is_same_v>); - static_assert(kM0 == QScaleDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] && - kSubQKHeaddim == - QScaleDramBlockWindowTmp{}.get_window_lengths()[number<1>{}] * - kQKScaleGranularity && - kN0 == KScaleDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] && - kK0 == KScaleDramBlockWindowTmp{}.get_window_lengths()[number<1>{}] * - kQKScaleGranularity && - kN1 == VScaleDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] && - kK1 == VScaleDramBlockWindowTmp{}.get_window_lengths()[number<1>{}] * - kVScaleGranularity); - } + static_assert(QScaleEnum != BlockAttentionQuantScaleEnum::MX); constexpr auto LdsSeq = Policy::template GetLdsBufferSequence(); @@ -430,52 +400,12 @@ struct BlockFmhaPipelineQRKSVSAsync {0, kv_load_start}, // TODO: hdim split? Policy::template MakeVDramTileDistribution()); - auto q_scale = [&] { - if constexpr(QScaleEnum == BlockAttentionQuantScaleEnum::MX) - { - auto q_scale_dram_window = - make_tile_window(q_scale_dram_block_window_tmp.get_bottom_tensor_view(), - q_scale_dram_block_window_tmp.get_window_lengths(), - q_scale_dram_block_window_tmp.get_window_origin(), - Policy::template MakeQScaleRegTileDistribution()); - return load_tile(q_scale_dram_window); - } - else - { - return null_tensor{}; - } - }(); - auto k_scale_dram_block_window = [&] { - if constexpr(QScaleEnum == BlockAttentionQuantScaleEnum::MX) - { - return make_tile_window(k_scale_dram_block_window_tmp.get_bottom_tensor_view(), - k_scale_dram_block_window_tmp.get_window_lengths(), - {seqlen_k_start, 0}); - } - else - { - return make_null_tile_window(make_tuple()); - } - }(); - auto v_scale_dram_window = [&] { - if constexpr(QScaleEnum == BlockAttentionQuantScaleEnum::MX) - { - return make_tile_window(v_scale_dram_block_window_tmp.get_bottom_tensor_view(), - v_scale_dram_block_window_tmp.get_window_lengths(), - {0, seqlen_k_start / kVScaleGranularity}, - Policy::template MakeVScaleRegTileDistribution()); - } - else - { - return make_null_tile_window(make_tuple()); - } - }(); - // prefetch K tile async_load_tile_raw( k_lds_store(LdsSeq.at(number<0>{})), k_dram_window, number<-1>{}, k_oob_ck, k_pre_np); move_tile_window(k_dram_window, {0, kK0}); __builtin_amdgcn_sched_barrier(0); + buffer_load_fence(k_dram_window.get_num_of_access(), q.get_thread_buffer()); (void)q_element_func; // ??? rocm-6.x if use q element func will have scratch on hdim=64/32 // auto q_tile = q; // tile_elementwise_in(q_element_func, q); @@ -496,56 +426,6 @@ struct BlockFmhaPipelineQRKSVSAsync const index_t kv_idx = (kv_load_start + i_total_loops * kN0) / block_scale_size_kv; k_descale = k_descale_ptr[kv_idx]; } - - auto k_scale_dram_window = [&] { - if constexpr(QScaleEnum == BlockAttentionQuantScaleEnum::MX) - { - return make_tile_window( - k_scale_dram_block_window.get_bottom_tensor_view(), - k_scale_dram_block_window.get_window_lengths(), - k_scale_dram_block_window.get_window_origin(), - Policy::template MakeKScaleRegTileDistribution()); - } - else - { - return make_null_tile_window(make_tuple()); - } - }(); - - auto load_k_scale_block_tile = [&] { - if constexpr(QScaleEnum == BlockAttentionQuantScaleEnum::MX) - { - auto t = load_tile(k_scale_dram_window); - move_tile_window(k_scale_dram_window, {0, kK0 / kQKScaleGranularity}); - return t; - } - else - { - return make_null_tile_window(make_tuple()); - } - }; - auto k_scale_block_tile = load_k_scale_block_tile(); - auto run_gemm_0 = [&](auto i_k0) { - auto q_slice = - get_slice_tile(q, sequence<0, i_k0 * kK0>{}, sequence{}); - auto k_slice = - get_slice_tile(k_lds_load, - sequence<(LdsSeq.at(number{})) * kN0, 0>{}, - sequence<(LdsSeq.at(number{}) + 1) * kN0, kK0>{}); - if constexpr(QScaleEnum == BlockAttentionQuantScaleEnum::MX) - { - auto q_scale_slice = - get_slice_tile(q_scale, - sequence<0, i_k0*(kK0 / kQKScaleGranularity)>{}, - sequence{}); - gemm_0(s_acc, q_slice, q_scale_slice, k_slice, k_scale_block_tile); - } - else - { - gemm_0(s_acc, q_slice, k_slice); - } - }; - // STAGE 1, QK gemm clear_tile(s_acc); // initialize C if constexpr(k0_loops > 1) @@ -562,8 +442,12 @@ struct BlockFmhaPipelineQRKSVSAsync async_load_fence(k_dram_window.get_num_of_access()); __builtin_amdgcn_s_barrier(); __builtin_amdgcn_sched_barrier(0); - run_gemm_0(number{}); - k_scale_block_tile = load_k_scale_block_tile(); + gemm_0(s_acc, + get_slice_tile( + q, sequence<0, i_k0 * kK0>{}, sequence{}), + get_slice_tile(k_lds_load, + sequence<(LdsSeq.at(number{})) * kN0, 0>{}, + sequence<(LdsSeq.at(number{}) + 1) * kN0, kK0>{})); }); } @@ -579,7 +463,13 @@ struct BlockFmhaPipelineQRKSVSAsync auto v_buf = load_tile(v_dram_window, number<-1>{}, bool_constant{}); __builtin_amdgcn_sched_barrier(0); { // tail - run_gemm_0(number{}); + gemm_0( + s_acc, + get_slice_tile( + q, sequence<0, (k0_loops - 1) * kK0>{}, sequence{}), + get_slice_tile(k_lds_load, + sequence<(LdsSeq.at(number{})) * kN0, 0>{}, + sequence<(LdsSeq.at(number{}) + 1) * kN0, kK0>{})); } __builtin_amdgcn_sched_barrier(1); // dequant @@ -748,11 +638,11 @@ struct BlockFmhaPipelineQRKSVSAsync tile_elementwise_in(v_element_func, v_buf)); // store the prefetch } - move_tile_window( - v_dram_window, - {0, kK1}); // will have scratch if move this right after load_tile(v_dram)... if constexpr(k1_loops > 1) { + move_tile_window( + v_dram_window, + {0, kK1}); // will have scratch if move this right after load_tile(v_dram)... v_buf = load_tile( v_dram_window, number<-1>{}, bool_constant{}); // load next v_buf } @@ -938,19 +828,21 @@ struct BlockFmhaPipelineQRKSVSAsync randval_ptr, seq_offset, p_compute, randval_dram_window); } - auto load_v_scale_block_tile = [&] { - if constexpr(QScaleEnum == BlockAttentionQuantScaleEnum::MX) - { - auto t = load_tile(v_scale_dram_window); - move_tile_window(v_scale_dram_window, {0, kK1 / kVScaleGranularity}); - return t; - } + const auto p = [&]() { +#if CK_TILE_FMHA_FLOAT_TO_FLOAT16_RTN + // For fp32 to fp16, + // impl::cast_tile_pkrtz_fp16_fp32 would cause precision issue, + // since it uses __builtin_amdgcn_cvt_pkrtz, which is round to zero. + return cast_tile(tile_elementwise_in(p_compute_element_func, p_compute)); +#else + if constexpr(std::is_same_v) + return impl::cast_tile_pkrtz_fp16_fp32( + tile_elementwise_in(p_compute_element_func, p_compute)); else - { - return make_null_tile_window(make_tuple()); - } - }; - auto v_scale_block_tile = load_v_scale_block_tile(); + return cast_tile( + tile_elementwise_in(p_compute_element_func, p_compute)); +#endif + }(); float v_descale = 1.0f; if constexpr(QScaleEnum == BlockAttentionQuantScaleEnum::BLOCKSCALE) @@ -959,48 +851,6 @@ struct BlockFmhaPipelineQRKSVSAsync const index_t kv_idx = (kv_load_start + i_total_loops * kN0) / block_scale_size_kv; v_descale = v_descale_ptr[kv_idx]; } - - const auto p_p_scale = [&] { - if constexpr(QScaleEnum == BlockAttentionQuantScaleEnum::MX) - { - auto p_result = make_static_distributed_tensor( - p_compute.get_tile_distribution()); - auto p_scale_result = make_static_distributed_tensor( - Policy::template MakePScaleRegTileDistribution()); - - constexpr auto config = - decltype(gemm_1)::Policy::template GetWarpGemmMWarpNWarp(); - using WG = remove_cvref_t())>; - - cast_tile_mx( - p_result, p_scale_result, p_compute); - - return make_tuple(p_result, p_scale_result); - } - else - { - const auto p_result = [&]() { -#if CK_TILE_FMHA_FLOAT_TO_FLOAT16_RTN - // For fp32 to fp16, - // impl::cast_tile_pkrtz_fp16_fp32 would cause precision issue, - // since it uses __builtin_amdgcn_cvt_pkrtz, which is round to zero. - return cast_tile( - tile_elementwise_in(p_compute_element_func, p_compute)); -#else - if constexpr(std::is_same_v) - return impl::cast_tile_pkrtz_fp16_fp32( - tile_elementwise_in(p_compute_element_func, p_compute)); - else - return cast_tile( - tile_elementwise_in(p_compute_element_func, p_compute)); -#endif - }(); - return make_tuple(p_result, null_tensor{}); - } - }(); - const auto p = p_p_scale[number<0>{}]; - const auto p_scale = p_p_scale[number<1>{}]; - // STAGE 3, KV gemm auto o_acc0 = decltype(o_acc){}; clear_tile(o_acc0); @@ -1015,42 +865,22 @@ struct BlockFmhaPipelineQRKSVSAsync return o_acc; } }(); - - auto run_gemm_1 = [&](auto i_k1) { - auto p_slice = - get_slice_tile(p, sequence<0, i_k1 * kK1>{}, sequence{}); - auto v_slice = get_slice_tile( - v_lds_window, - sequence<(LdsSeq.at(number{})) * kN1, 0>{}, - sequence<(LdsSeq.at(number{}) + 1) * kN1, kK1>{}); - if constexpr(QScaleEnum == BlockAttentionQuantScaleEnum::MX) - { - auto p_scale_slice = - get_slice_tile(p_scale, - sequence<0, i_k1*(kK1 / kVScaleGranularity)>{}, - sequence{}); - gemm_1(o_acc_, p_slice, p_scale_slice, v_slice, v_scale_block_tile); - } - else if constexpr(QScaleEnum == BlockAttentionQuantScaleEnum::BLOCKSCALE) - { - gemm_1(o_acc0, p_slice, v_slice); - } - else - { - gemm_1(o_acc_, p_slice, v_slice); - } - }; - if constexpr(k1_loops > 1) { static_for<0, k1_loops - 1, 1>{}([&](auto i_k1) { - if constexpr(i_k1 != 0) + if constexpr(i_k1 != 0 && i_k1 < k1_loops - 1) { v_buf = load_tile( v_dram_window, number<-1>{}, bool_constant{}); // load next v_buf } block_sync_lds(); - run_gemm_1(number{}); + gemm_1(o_acc_, + get_slice_tile( + p, sequence<0, i_k1 * kK1>{}, sequence{}), + get_slice_tile( + v_lds_window, + sequence<(LdsSeq.at(number{})) * kN1, 0>{}, + sequence<(LdsSeq.at(number{}) + 1) * kN1, kK1>{})); if constexpr(std::is_same_v) { @@ -1074,8 +904,8 @@ struct BlockFmhaPipelineQRKSVSAsync store_tile(v_lds_window_tmp, tile_elementwise_in(v_element_func, v_buf)); // store next v_buf } - move_tile_window(v_dram_window, {0, kK1}); - v_scale_block_tile = load_v_scale_block_tile(); + if constexpr(i_k1 < k1_loops - 1) + move_tile_window(v_dram_window, {0, kK1}); }); } i_total_loops++; @@ -1083,7 +913,6 @@ struct BlockFmhaPipelineQRKSVSAsync { if constexpr(kHasSink) { - // TODO: this never happens because of i_total_loops++ if(i_total_loops == 0) { move_tile_window(k_dram_block_window, {seqlen_k_start - sink_seq_end, 0}); @@ -1091,10 +920,6 @@ struct BlockFmhaPipelineQRKSVSAsync } } move_tile_window(k_dram_block_window, {kN0, 0}); - if constexpr(QScaleEnum == BlockAttentionQuantScaleEnum::MX) - { - move_tile_window(k_scale_dram_block_window, {kN0, 0}); - } k_dram_window.set_window_origin(k_dram_block_window.get_window_origin()); @@ -1111,7 +936,13 @@ struct BlockFmhaPipelineQRKSVSAsync // tail { block_sync_lds(); - run_gemm_1(number{}); + gemm_1( + o_acc_, + get_slice_tile(p, sequence<0, (k1_loops - 1) * kK1>{}, sequence{}), + get_slice_tile( + v_lds_window, + sequence<(LdsSeq.at(number{})) * kN1, 0>{}, + sequence<(LdsSeq.at(number{}) + 1) * kN1, kK1>{})); } if constexpr(QScaleEnum == BlockAttentionQuantScaleEnum::BLOCKSCALE) diff --git a/test/ck_tile/fmha/test_fmha_fwd.cpp b/test/ck_tile/fmha/test_fmha_fwd.cpp index 93d81d497e..7f47220683 100644 --- a/test/ck_tile/fmha/test_fmha_fwd.cpp +++ b/test/ck_tile/fmha/test_fmha_fwd.cpp @@ -49,10 +49,6 @@ struct TestConfigs static constexpr bool def_is_v_rowmajor = true; static constexpr auto init_method = "uf"; static int adjust_seqlen(int seqlen) { return seqlen; } - static int adjust_hdim(int hdim) - { - return hdim < 0 ? hdim : ck_tile::integer_least_multiple(hdim, 8); - } }; template <> @@ -71,10 +67,6 @@ struct TestConfigs // When there are no fp8 instances with padding, pad seqlen to avoid skipping most of the tests: // return ck_tile::integer_least_multiple(seqlen, 128); static int adjust_seqlen(int seqlen) { return seqlen; } - static int adjust_hdim(int hdim) - { - return hdim < 0 ? hdim : ck_tile::integer_least_multiple(hdim, 16); - } }; template <> @@ -90,10 +82,6 @@ struct TestConfigs static constexpr bool def_is_v_rowmajor = false; static constexpr auto init_method = "3"; static int adjust_seqlen(int seqlen) { return seqlen; } - static int adjust_hdim(int hdim) - { - return hdim < 0 ? hdim : ck_tile::integer_least_multiple(hdim, 16); - } }; template <> @@ -112,10 +100,6 @@ struct TestConfigs { return seqlen < 0 ? seqlen : ck_tile::integer_least_multiple(seqlen, 2); } - static int adjust_hdim(int hdim) - { - return hdim < 0 ? hdim : ck_tile::integer_least_multiple(hdim, 32); - } }; template <> @@ -139,7 +123,6 @@ struct TestConfigs static constexpr bool def_is_v_rowmajor = true; static constexpr auto init_method = "uf"; static int adjust_seqlen(int seqlen) { return seqlen; } - static int adjust_hdim(int hdim) { return hdim; } }; static auto HDimValues = ValuesIn(TestConfigs::HDimValues); @@ -152,7 +135,6 @@ constexpr bool def_lse = TestConfigs::def_lse; constexpr bool def_is_v_rowmajor = TestConfigs::def_is_v_rowmajor; constexpr auto init_method = TestConfigs::init_method; int adjust_seqlen(int seqlen) { return TestConfigs::adjust_seqlen(seqlen); } -int adjust_hdim(int hdim) { return TestConfigs::adjust_hdim(hdim); } // Random seed used for initializing input tensors. 0 for non-deterministic seed CK_TILE_DECLARE_ENV_VAR(CK_TILE_TEST_SEED, uint64_t, 123456) @@ -271,70 +253,6 @@ TEST_P(AllLong, DataTypeConfig) CHECK_RESULT(result); } -class General - : public TestWithParam, - bool, - bool, - mode_enum, - std::tuple>> -{ -}; - -INSTANTIATE_TEST_SUITE_P(TestCkTileFmhaFwd, - General, - Combine(HDimValues, - Bool(), - IsVRowmajorValues, - ModeValues, - Values(std::tuple{2, 2, 1, 55, 256, -1, "0"}, - std::tuple{1, 3, -1, 100, 51, -1, "0"}, - std::tuple{2, 1, -1, 99, 256, -1, "1"}, - std::tuple{1, 2, 1, 1024, 256, -1, "2"}, - std::tuple{2, 1, -1, 3, 99, -1, "2"}, - std::tuple{1, 2, 1, 33, 0, -1, "2"}, - std::tuple{1, 2, 1, 1, 10, 32, "2"}))); - -TEST_P(General, DataTypeConfig) -{ - auto [hdims, perm, is_v_rowmajor, mode, dims_mask] = GetParam(); - auto [hdim_q, hdim_v] = hdims; - auto [batch, nhead, nhead_k, seqlen_q, seqlen_k, seqlen_kpad, mask_str] = dims_mask; - - auto result = fmha_fwd_run(mode, - batch, - nhead, - nhead_k, - {adjust_seqlen(seqlen_q)}, - {adjust_seqlen(seqlen_k)}, - hdim_q, - hdim_v, - 0, // seqlen_knew - {-1}, // seqlen_qpads - {seqlen_kpad}, // seqlen_kpads - {}, // q_eff_lens_per_batch - {}, // kv_eff_lens_per_batch - 0, // rotary_dim - perm, // i_perm - perm, // o_perm - 0, // scale_s - 0, // logits_soft_cap - is_v_rowmajor, // is_v_rowmajor - def_lse, // lse - 0, // page_block_size - false, // use_cache_batch_idx - "n", // bias_str - 0.0f, // p_drop - 0, // drop_seed - 0, // drop_offset - false, // drop_prefs - mask_str, // mask_str - qscale_str, - true, // is_rotary_interleaved - 1, // num_splits - COMMON_ARGS); - CHECK_RESULT(result); -} - // --------------------------------------------------------------- // Negative tests: padding not supported with appendkv/splitkv/pagedkv // --------------------------------------------------------------- @@ -512,8 +430,8 @@ TEST_P(HDimPadding, DataTypeConfig) nhead_k, {adjust_seqlen(seqlen_q)}, {adjust_seqlen(seqlen_k)}, - adjust_hdim(hdim_q), - adjust_hdim(hdim_v), + hdim_q, + hdim_v, 0, // seqlen_knew {-1}, // seqlen_qpads {seqlen_kpad}, // seqlen_kpads