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 38ba926b4e..ecaca2683a 100644 --- a/example/ck_tile/01_fmha/codegen/ops/fmha_fwd.py +++ b/example/ck_tile/01_fmha/codegen/ops/fmha_fwd.py @@ -355,18 +355,12 @@ class FmhaFwdApiTrait: @property def dcheck(self) -> str: - 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.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_trload", "qr_async_trload_v3"]: + elif self.pipeline_tag in ["qr", "qs", "qr_async", "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) @@ -377,18 +371,12 @@ class FmhaFwdApiTrait: @property def dvcheck(self) -> str: - 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.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_trload", "qr_async_trload_v3"]: + elif self.pipeline_tag in ["qr", "qs", "qr_async", "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) @@ -1046,7 +1034,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", "t", "f", "t", "t", logits, bias, lse, dropout, qscale, mask, skip, "f", sink)) # fmt: skip + 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", "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 @@ -1060,10 +1048,10 @@ class KernelComponentFactoryGfx9(CompatibilityRuleFactoryGfx9): ["f", "t"], ): if hdim == 64: - 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", "f", "f", "f", "f", 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", "t", "f", "t", "t", logits, bias, "f", "f", qscale, mask, "f", "f", sink)) # fmt: skip + 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", "t", "t", "t", logits, bias, "f", "f", qscale, mask, "f", "f", sink)) # fmt: skip return pipelines @@ -1169,6 +1157,9 @@ 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 f2b3a48dcb..12dc305ab5 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 +$EXE $base_group_args -s_qpad=2048,1536,1024,512 -s_kpad=2048,1536,1024,512 \ No newline at end of file diff --git a/include/ck_tile/core/tensor/tile_window.hpp b/include/ck_tile/core/tensor/tile_window.hpp index 7dcf498695..bb91e552d5 100644 --- a/include/ck_tile/core/tensor/tile_window.hpp +++ b/include/ck_tile/core/tensor/tile_window.hpp @@ -506,9 +506,8 @@ 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::PackedSize * Traits::ScalarPerVector) == 0); - using vectorized_tbuf = - array; + static_assert(YElementSize % Traits::ScalarPerVector == 0); + using vectorized_tbuf = array; constexpr auto tile_dstr = typename Base::TileDstr{}; @@ -534,10 +533,11 @@ 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(d % Traits::ScalarPerVector == 0); + static_assert(Traits::ScalarPerVector % Traits::PackedSize == 0); + static_assert(d % (Traits::ScalarPerVector / Traits::PackedSize) == 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 5155794b16..b621e9ceab 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,6 +7,7 @@ #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" @@ -29,6 +30,10 @@ 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; @@ -50,21 +55,20 @@ struct BlockFmhaPipelineQRKSVSAsync static_assert(kSubQKHeaddim <= 256, "hdim bigger than 256 is not suitable for this pipeline!"); - 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 kIsGroupMode = Problem::kIsGroupMode; + static constexpr bool kPadSeqLenQ = Problem::kPadSeqLenQ; static constexpr bool kPadSeqLenK = Problem::kPadSeqLenK; - 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 kPadHeadDimQ = Problem::kPadHeadDimQ; + static constexpr bool kPadHeadDimV = Problem::kPadHeadDimV; 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; @@ -82,7 +86,8 @@ struct BlockFmhaPipelineQRKSVSAsync if constexpr(std::is_same_v) return Policy::template GetAlignmentV(); else - return kPadSeqLenK ? 1 : Policy::template GetAlignmentV(); + return kPadSeqLenK ? numeric_traits::PackedSize + : Policy::template GetAlignmentV(); }(); static constexpr index_t kAlignmentO = Policy::template GetAlignmentO(); static constexpr index_t kAlignmentBias = @@ -201,9 +206,12 @@ struct BlockFmhaPipelineQRKSVSAsync const float* k_descale_ptr, const float* v_descale_ptr, const index_t block_scale_size_kv, - const QScaleDramBlockWindowTmp&, // M0*(K0/kQKScaleGranularity) tile - const KScaleDramBlockWindowTmp&, // N0*(K0/kQKScaleGranularity) tile - const VScaleDramBlockWindowTmp&, // N1*(K1/kVScaleGranularity) tile + 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 float sink_v) const { static_assert( @@ -213,6 +221,8 @@ 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>{}] && @@ -220,8 +230,28 @@ 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(QScaleEnum != BlockAttentionQuantScaleEnum::MX); + 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); + } constexpr auto LdsSeq = Policy::template GetLdsBufferSequence(); @@ -400,12 +430,52 @@ 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); @@ -426,6 +496,56 @@ 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) @@ -442,12 +562,8 @@ struct BlockFmhaPipelineQRKSVSAsync async_load_fence(k_dram_window.get_num_of_access()); __builtin_amdgcn_s_barrier(); __builtin_amdgcn_sched_barrier(0); - 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>{})); + run_gemm_0(number{}); + k_scale_block_tile = load_k_scale_block_tile(); }); } @@ -463,13 +579,7 @@ struct BlockFmhaPipelineQRKSVSAsync auto v_buf = load_tile(v_dram_window, number<-1>{}, bool_constant{}); __builtin_amdgcn_sched_barrier(0); { // tail - 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>{})); + run_gemm_0(number{}); } __builtin_amdgcn_sched_barrier(1); // dequant @@ -638,11 +748,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 } @@ -828,21 +938,19 @@ struct BlockFmhaPipelineQRKSVSAsync randval_ptr, seq_offset, p_compute, randval_dram_window); } - 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)); + 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; + } else - return cast_tile( - tile_elementwise_in(p_compute_element_func, p_compute)); -#endif - }(); + { + return make_null_tile_window(make_tuple()); + } + }; + auto v_scale_block_tile = load_v_scale_block_tile(); float v_descale = 1.0f; if constexpr(QScaleEnum == BlockAttentionQuantScaleEnum::BLOCKSCALE) @@ -851,6 +959,48 @@ 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); @@ -865,22 +1015,42 @@ 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 && i_k1 < k1_loops - 1) + if constexpr(i_k1 != 0) { v_buf = load_tile( v_dram_window, number<-1>{}, bool_constant{}); // load next v_buf } block_sync_lds(); - 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>{})); + run_gemm_1(number{}); if constexpr(std::is_same_v) { @@ -904,8 +1074,8 @@ struct BlockFmhaPipelineQRKSVSAsync store_tile(v_lds_window_tmp, tile_elementwise_in(v_element_func, v_buf)); // store next v_buf } - if constexpr(i_k1 < k1_loops - 1) - move_tile_window(v_dram_window, {0, kK1}); + move_tile_window(v_dram_window, {0, kK1}); + v_scale_block_tile = load_v_scale_block_tile(); }); } i_total_loops++; @@ -913,6 +1083,7 @@ 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}); @@ -920,6 +1091,10 @@ 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()); @@ -936,13 +1111,7 @@ struct BlockFmhaPipelineQRKSVSAsync // tail { block_sync_lds(); - 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>{})); + run_gemm_1(number{}); } 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 7f47220683..93d81d497e 100644 --- a/test/ck_tile/fmha/test_fmha_fwd.cpp +++ b/test/ck_tile/fmha/test_fmha_fwd.cpp @@ -49,6 +49,10 @@ 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 <> @@ -67,6 +71,10 @@ 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 <> @@ -82,6 +90,10 @@ 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 <> @@ -100,6 +112,10 @@ 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 <> @@ -123,6 +139,7 @@ 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); @@ -135,6 +152,7 @@ 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) @@ -253,6 +271,70 @@ 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 // --------------------------------------------------------------- @@ -430,8 +512,8 @@ TEST_P(HDimPadding, DataTypeConfig) nhead_k, {adjust_seqlen(seqlen_q)}, {adjust_seqlen(seqlen_k)}, - hdim_q, - hdim_v, + adjust_hdim(hdim_q), + adjust_hdim(hdim_v), 0, // seqlen_knew {-1}, // seqlen_qpads {seqlen_kpad}, // seqlen_kpads