From b8c99da8f678cb41fc5e01a53bafb1feedc9f442 Mon Sep 17 00:00:00 2001 From: ltqin Date: Fri, 21 Nov 2025 13:53:45 +0000 Subject: [PATCH] group mode for block scale --- example/ck_tile/01_fmha/fmha_fwd.hpp | 10 +++++-- example/ck_tile/01_fmha/fmha_fwd_runner.hpp | 10 ++++++- .../ops/fmha/kernel/fmha_fwd_kernel.hpp | 29 ++++++++++++++++--- 3 files changed, 41 insertions(+), 8 deletions(-) diff --git a/example/ck_tile/01_fmha/fmha_fwd.hpp b/example/ck_tile/01_fmha/fmha_fwd.hpp index 16cb1a0566..d7634d0b5f 100644 --- a/example/ck_tile/01_fmha/fmha_fwd.hpp +++ b/example/ck_tile/01_fmha/fmha_fwd.hpp @@ -230,6 +230,8 @@ struct fmha_fwd_args // array [batch + 1]. (Used with padding) const void* cu_seqlen_k_ptr = nullptr; // Cumulative logical (excluding padding) sequence length // array [batch + 1]. (Used with padding) + const void* bseqstart_q_ptr; + const void* bseqstart_k_ptr; ck_tile::index_t seqlen_q; ck_tile::index_t seqlen_k; @@ -601,6 +603,8 @@ auto fmha_fwd_create_kargs_and_grids(fmha_fwd_args args) args.seqstart_k_ptr, args.seqlen_q_ptr, args.seqlen_k_ptr, + args.bseqstart_q_ptr, + args.bseqstart_k_ptr, args.hdim_q, args.hdim_v, args.nhead_q, @@ -613,9 +617,6 @@ auto fmha_fwd_create_kargs_and_grids(fmha_fwd_args args) args.stride_bias, args.stride_randval, args.stride_o, - args.nhead_stride_q_descale, - args.nhead_stride_k_descale, - args.nhead_stride_v_descale, args.nhead_stride_q, args.nhead_stride_k, args.nhead_stride_v, @@ -623,6 +624,9 @@ auto fmha_fwd_create_kargs_and_grids(fmha_fwd_args args) args.nhead_stride_randval, args.nhead_stride_lse, args.nhead_stride_o, + args.nhead_stride_q_descale, + args.nhead_stride_k_descale, + args.nhead_stride_v_descale, args.window_size_left, args.window_size_right, args.mask_type, diff --git a/example/ck_tile/01_fmha/fmha_fwd_runner.hpp b/example/ck_tile/01_fmha/fmha_fwd_runner.hpp index cbeb8c2a88..28966132c0 100644 --- a/example/ck_tile/01_fmha/fmha_fwd_runner.hpp +++ b/example/ck_tile/01_fmha/fmha_fwd_runner.hpp @@ -741,6 +741,8 @@ fwd_result fmha_fwd_run(mode_enum mode, ck_tile::DeviceMem q_descale_buf(q_descale_host.get_element_space_size_in_bytes()); ck_tile::DeviceMem k_descale_buf(k_descale_host.get_element_space_size_in_bytes()); ck_tile::DeviceMem v_descale_buf(v_descale_host.get_element_space_size_in_bytes()); + ck_tile::DeviceMem bseqstart_q_buf(bseqstart_q_host.size() * sizeof(int32_t)); + ck_tile::DeviceMem bseqstart_k_buf(bseqstart_k_host.size() * sizeof(int32_t)); ck_tile::DeviceMem lse_acc_buf(lse_acc_host.get_element_space_size_in_bytes()); ck_tile::DeviceMem o_acc_buf(o_acc_host.get_element_space_size_in_bytes()); ck_tile::DeviceMem lse_buf(lse_host.get_element_space_size_in_bytes()); @@ -785,6 +787,8 @@ fwd_result fmha_fwd_run(mode_enum mode, q_descale_buf.ToDevice(q_descale_host.data()); k_descale_buf.ToDevice(k_descale_host.data()); v_descale_buf.ToDevice(v_descale_host.data()); + bseqstart_q_buf.ToDevice(bseqstart_q_host.data()); + bseqstart_k_buf.ToDevice(bseqstart_k_host.data()); seqstart_q.ToDevice(seqstart_q_host.data()); // Keep logical starts in seqstart_k; pass padded K via separate pointer seqstart_k.ToDevice(seqstart_k_host.data()); @@ -1097,6 +1101,11 @@ fwd_result fmha_fwd_run(mode_enum mode, args.v_descale_ptr = reinterpret_cast(v_descale_buf.GetDeviceBuffer()); + args.bseqstart_q_ptr = + (mode == mode_enum::group ? bseqstart_q_buf.GetDeviceBuffer() : nullptr); + args.bseqstart_k_ptr = + (mode == mode_enum::group ? bseqstart_k_buf.GetDeviceBuffer() : nullptr); + args.nhead_stride_q_descale = nhead_stride_q_descale; args.nhead_stride_k_descale = nhead_stride_k_descale; args.nhead_stride_v_descale = nhead_stride_v_descale; @@ -1844,7 +1853,6 @@ fwd_result fmha_fwd_run(mode_enum mode, if(o_perm) o_host_result.ForEach([&](auto& self, auto idx) { self(idx) = o_host(b_idx, idx[0], idx[1] + query_offset, idx[2]); }); else o_host_result.ForEach([&](auto& self, auto idx) { self(idx) = o_host(b_idx, idx[1] + query_offset, idx[0], idx[2]); }); // clang-format on - auto [rtol, atol] = get_elimit(init_method); bool cur_pass = ck_tile::check_err(o_host_result, o_host_ref, diff --git a/include/ck_tile/ops/fmha/kernel/fmha_fwd_kernel.hpp b/include/ck_tile/ops/fmha/kernel/fmha_fwd_kernel.hpp index 2c5121e6e6..f348e4425f 100644 --- a/include/ck_tile/ops/fmha/kernel/fmha_fwd_kernel.hpp +++ b/include/ck_tile/ops/fmha/kernel/fmha_fwd_kernel.hpp @@ -231,6 +231,12 @@ struct FmhaFwdKernel ck_tile::index_t batch_stride_v_descale; }; + struct FmhaFwdGroupBlockScaleKargs : public FmhaFwdCommonBlockScaleKargs + { + const int32_t* bseqstart_q_ptr; + const int32_t* bseqstart_k_ptr; + }; + struct FmhaFwdCommonLSEKargs { void* lse_ptr = nullptr; @@ -339,7 +345,7 @@ struct FmhaFwdKernel QScaleEnum == BlockAttentionQuantScaleEnum::PERTENSOR, FmhaFwdCommonQScaleKargs, std::conditional_t>>, std::conditional_t>, std::conditional_t>, @@ -780,6 +786,8 @@ struct FmhaFwdKernel const void* seqstart_k_ptr, const void* seqlen_q_ptr, const void* seqlen_k_ptr, + const void* bseqstart_q_ptr, + const void* bseqstart_k_ptr, ck_tile::index_t hdim_q, ck_tile::index_t hdim_v, ck_tile::index_t num_head_q, @@ -890,6 +898,9 @@ struct FmhaFwdKernel kargs.block_scale_m = block_scale_m; kargs.block_scale_n = block_scale_n; + + kargs.bseqstart_q_ptr = reinterpret_cast(bseqstart_q_ptr); + kargs.bseqstart_k_ptr = reinterpret_cast(bseqstart_k_ptr); } if constexpr(kHasDropout) { @@ -942,6 +953,8 @@ struct FmhaFwdKernel const void* seqstart_k_ptr, const void* seqlen_q_ptr, const void* seqlen_k_ptr, + const void* bseqstart_q_ptr, + const void* bseqstart_k_ptr, ck_tile::index_t hdim_q, ck_tile::index_t hdim_v, ck_tile::index_t num_head_q, @@ -991,6 +1004,8 @@ struct FmhaFwdKernel seqstart_k_ptr, seqlen_q_ptr, seqlen_k_ptr, + bseqstart_q_ptr, + bseqstart_k_ptr, hdim_q, hdim_v, num_head_q, @@ -1043,6 +1058,8 @@ struct FmhaFwdKernel const void* seqstart_k_ptr, const void* seqlen_q_ptr, const void* seqlen_k_ptr, + const void* bseqstart_q_ptr, + const void* bseqstart_k_ptr, ck_tile::index_t hdim_q, ck_tile::index_t hdim_v, ck_tile::index_t num_head_q, @@ -1092,6 +1109,8 @@ struct FmhaFwdKernel seqstart_k_ptr, seqlen_q_ptr, seqlen_k_ptr, + bseqstart_q_ptr, + bseqstart_k_ptr, hdim_q, hdim_v, num_head_q, @@ -1296,9 +1315,11 @@ struct FmhaFwdKernel } if constexpr(QScaleEnum == BlockAttentionQuantScaleEnum::BLOCKSCALE) { - batch_offset_q_descale = query_start/128; - batch_offset_k_descale = key_start/128; - batch_offset_v_descale = key_start/128; + const long_index_t bquery_start = kargs.bseqstart_q_ptr[i_batch]; + const long_index_t bkey_start = kargs.bseqstart_k_ptr[i_batch]; + batch_offset_q_descale = bquery_start; + batch_offset_k_descale = bkey_start; + batch_offset_v_descale = bkey_start; } batch_offset_o = query_start * kargs.stride_o;