ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-13 10:37:42 +00:00
Code Issues Packages Projects Releases Wiki Activity

group mode for block scale

Browse Source
This commit is contained in:
ltqin
2025-11-21 13:53:45 +00:00
parent a7048cf4d4
commit b8c99da8f6
3 changed files with 41 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
example/ck_tile/01_fmha/fmha_fwd.hpp
View File

@@ -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,

10
example/ck_tile/01_fmha/fmha_fwd_runner.hpp
View File

@@ -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<const float*>(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<DataTypeConfig>(init_method);
bool cur_pass = ck_tile::check_err(o_host_result,
o_host_ref,

29
include/ck_tile/ops/fmha/kernel/fmha_fwd_kernel.hpp
View File

@@ -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<QScaleEnum == BlockAttentionQuantScaleEnum::BLOCKSCALE,
FmhaFwdCommonBlockScaleKargs,
FmhaFwdGroupBlockScaleKargs,
FmhaFwdEmptyKargs<3>>>,
std::conditional_t<kHasDropout, FmhaFwdCommonDropoutKargs, FmhaFwdEmptyKargs<4>>,
std::conditional_t<kHasLogitsSoftCap, FmhaFwdLogitsSoftCapKargs, FmhaFwdEmptyKargs<5>>,
@@ -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<const int32_t*>(bseqstart_q_ptr);
kargs.bseqstart_k_ptr = reinterpret_cast<const int32_t*>(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;