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

Replace template kUseSoftmax/kStoreLSE by boolean parameters in reference fwd codes to save compiling time

Browse Source
This commit is contained in:
Qianfeng Zhang
2026-06-04 15:12:20 +00:00
parent 8f091efc4c
commit 0fdca0e940
2 changed files with 72 additions and 80 deletions
Download Patch File Download Diff File Expand all files Collapse all files

113
example/ck_tile/18_hstu_attention/example_hstu_attention_fwd.cpp
View File

@@ -587,35 +587,32 @@ bool run_no_group_hstu(const ck_tile::ArgParser& arg_parser, bool is_jagged)
{
using GemmAccDataType = typename HstuAttentionFwdTypeConfig<InOutDataType>::GemmAccDataType;
BOOL_SWITCH_3(is_jagged, kIsJagged, use_softmax, kUseSoftmax, use_causal, kUseCausal, [&] {
BOOL_SWITCH(is_training, kIsTraining, [&] {
constexpr bool kStoreLSE = (kIsTraining && kUseSoftmax);
ck_tile::reference_no_group_hstu_attention_fwd<
InOutDataType,
GemmAccDataType,
CompDataType,
kIsJagged,
kUseSoftmax,
kStoreLSE,
kUseCausal>::Run(is_cross_attention,
q_host,
k_host,
v_host,
o_host_ref,
lse_host_ref,
mask_host,
num_batch,
scale_s,
attn_scale,
max_seqlen_q,
max_seqlen_kv,
seq_offsets_q,
seq_offsets_kv,
num_targets,
contextual_seqlen,
window_size,
min_full_attn_seqlen);
});
BOOL_SWITCH_2(is_jagged, kIsJagged, use_causal, kUseCausal, [&] {
bool store_lse = (is_training && use_softmax);
ck_tile::reference_no_group_hstu_attention_fwd<InOutDataType,
GemmAccDataType,
CompDataType,
kIsJagged,
kUseCausal>::Run(is_cross_attention,
use_softmax,
store_lse,
q_host,
k_host,
v_host,
o_host_ref,
lse_host_ref,
mask_host,
num_batch,
scale_s,
attn_scale,
max_seqlen_q,
max_seqlen_kv,
seq_offsets_q,
seq_offsets_kv,
num_targets,
contextual_seqlen,
window_size,
min_full_attn_seqlen);
});
ck_tile::HostTensor<InOutDataType> o_host(
@@ -1042,36 +1039,34 @@ bool run_group_hstu(const ck_tile::ArgParser& arg_parser, int num_group)
{
using GemmAccDataType = typename HstuAttentionFwdTypeConfig<InOutDataType>::GemmAccDataType;
BOOL_SWITCH_2(use_softmax, kUseSoftmax, use_causal, kUseCausal, [&] {
BOOL_SWITCH(is_training, kIsTraining, [&] {
constexpr bool kStoreLSE = (kIsTraining && kUseSoftmax);
ck_tile::reference_group_hstu_attention_fwd<
InOutDataType,
GemmAccDataType,
CompDataType,
kUseSoftmax,
kStoreLSE,
kUseCausal>::Run(is_cross_attention,
q_host,
k_host,
v_host,
o_host_ref,
lse_host_ref,
mask_host,
num_batch,
num_batch / num_group,
scale_s,
max_max_seqlen_q,
max_max_seqlen_kv,
seq_offsets_q,
seq_offsets_kv,
num_targets,
group_max_seqlens_q,
group_contextual_seqlens,
group_window_sizes,
group_min_full_attn_seqlens,
group_attn_scales);
});
BOOL_SWITCH(use_causal, kUseCausal, [&] {
bool store_lse = (is_training && use_softmax);
ck_tile::reference_group_hstu_attention_fwd<
InOutDataType,
GemmAccDataType,
CompDataType,
kUseCausal>::Run(is_cross_attention,
use_softmax,
store_lse,
q_host,
k_host,
v_host,
o_host_ref,
lse_host_ref,
mask_host,
num_batch,
num_batch / num_group,
scale_s,
max_max_seqlen_q,
max_max_seqlen_kv,
seq_offsets_q,
seq_offsets_kv,
num_targets,
group_max_seqlens_q,
group_contextual_seqlens,
group_window_sizes,
group_min_full_attn_seqlens,
group_attn_scales);
});
ck_tile::HostTensor<InOutDataType> o_host(

39
example/ck_tile/18_hstu_attention/reference_hstu_attention_fwd.hpp
View File

@@ -29,12 +29,12 @@ template <typename InOutDataType,
typename GemmAccDataType,
typename CompDataType,
bool kIsJagged,
bool kUseSoftmax,
bool kStoreLSE,
bool kUseCausal>
struct reference_no_group_hstu_attention_fwd
{
static void Run(bool is_cross_attention,
bool use_softmax,
bool store_lse,
const HostTensor<InOutDataType>& q_batch_seq_nhead_hdim,
const HostTensor<InOutDataType>& k_batch_seq_nhead_hdim,
const HostTensor<InOutDataType>& v_batch_seq_nhead_hdim,
@@ -66,7 +66,7 @@ struct reference_no_group_hstu_attention_fwd
assert(v_batch_seq_nhead_hdim.get_lengths()[0] == 1);
assert(o_batch_seq_nhead_hdim.get_lengths()[0] == 1);
if constexpr(kStoreLSE)
if(store_lse)
{
assert(lse_batch_seq_nhead.get_lengths()[0] == 1);
}
@@ -80,7 +80,7 @@ struct reference_no_group_hstu_attention_fwd
assert(v_batch_seq_nhead_hdim.get_lengths()[0] == num_batch);
assert(o_batch_seq_nhead_hdim.get_lengths()[0] == num_batch);
if constexpr(kStoreLSE)
if(store_lse)
{
assert(lse_batch_seq_nhead.get_lengths()[0] == num_batch);
}
@@ -90,7 +90,7 @@ struct reference_no_group_hstu_attention_fwd
assert(q_batch_seq_nhead_hdim.get_lengths()[1] == k_batch_seq_nhead_hdim.get_lengths()[1]);
assert(q_batch_seq_nhead_hdim.get_lengths()[1] == v_batch_seq_nhead_hdim.get_lengths()[1]);
assert(q_batch_seq_nhead_hdim.get_lengths()[1] == o_batch_seq_nhead_hdim.get_lengths()[1]);
if constexpr(kStoreLSE)
if(store_lse)
{
assert(q_batch_seq_nhead_hdim.get_lengths()[1] == lse_batch_seq_nhead.get_lengths()[1]);
}
@@ -100,7 +100,7 @@ struct reference_no_group_hstu_attention_fwd
assert(num_head == k_batch_seq_nhead_hdim.get_lengths()[2]);
assert(num_head == v_batch_seq_nhead_hdim.get_lengths()[2]);
assert(num_head == o_batch_seq_nhead_hdim.get_lengths()[2]);
if constexpr(kStoreLSE)
if(store_lse)
{
assert(num_head == lse_batch_seq_nhead.get_lengths()[2]);
}
@@ -258,14 +258,14 @@ struct reference_no_group_hstu_attention_fwd
}
else
{
if constexpr(!kUseSoftmax)
if(!use_softmax)
locals.push_back(ck_tile::type_convert<CompDataType>(0.0f));
else
locals.push_back(-ck_tile::numeric<CompDataType>::infinity());
};
};
if constexpr(!kUseSoftmax)
if(!use_softmax)
{
// SiLu element-wise
for(CompDataType& elem : locals)
@@ -292,7 +292,7 @@ struct reference_no_group_hstu_attention_fwd
elem = std::exp(elem - m) / l;
}
if constexpr(kStoreLSE)
if(store_lse)
{
lse_batch_seq_nhead(i_batch, sq, i_head) = std::log(l) + m;
}
@@ -341,16 +341,13 @@ struct reference_no_group_hstu_attention_fwd
}
};
template <typename InOutDataType,
typename GemmAccDataType,
typename CompDataType,
bool kUseSoftmax,
bool kStoreLSE,
bool kUseCausal>
template <typename InOutDataType, typename GemmAccDataType, typename CompDataType, bool kUseCausal>
struct reference_group_hstu_attention_fwd
{
static void
Run(bool is_cross_attention,
bool use_softmax,
bool store_lse,
const HostTensor<InOutDataType>& q_batch_seq_nhead_hdim,
const HostTensor<InOutDataType>& k_batch_seq_nhead_hdim,
const HostTensor<InOutDataType>& v_batch_seq_nhead_hdim,
@@ -380,7 +377,7 @@ struct reference_group_hstu_attention_fwd
assert(k_batch_seq_nhead_hdim.get_lengths()[0] == 1);
assert(v_batch_seq_nhead_hdim.get_lengths()[0] == 1);
assert(o_batch_seq_nhead_hdim.get_lengths()[0] == 1);
if constexpr(kStoreLSE)
if(store_lse)
{
assert(lse_batch_seq_nhead.get_lengths()[0] == 1);
}
@@ -389,7 +386,7 @@ struct reference_group_hstu_attention_fwd
assert(q_batch_seq_nhead_hdim.get_lengths()[1] == k_batch_seq_nhead_hdim.get_lengths()[1]);
assert(q_batch_seq_nhead_hdim.get_lengths()[1] == v_batch_seq_nhead_hdim.get_lengths()[1]);
assert(q_batch_seq_nhead_hdim.get_lengths()[1] == o_batch_seq_nhead_hdim.get_lengths()[1]);
if constexpr(kStoreLSE)
if(store_lse)
{
assert(q_batch_seq_nhead_hdim.get_lengths()[1] == lse_batch_seq_nhead.get_lengths()[1]);
}
@@ -399,7 +396,7 @@ struct reference_group_hstu_attention_fwd
assert(num_head == k_batch_seq_nhead_hdim.get_lengths()[2]);
assert(num_head == v_batch_seq_nhead_hdim.get_lengths()[2]);
assert(num_head == o_batch_seq_nhead_hdim.get_lengths()[2]);
if constexpr(kStoreLSE)
if(store_lse)
{
assert(num_head == lse_batch_seq_nhead.get_lengths()[2]);
}
@@ -550,14 +547,14 @@ struct reference_group_hstu_attention_fwd
}
else
{
if constexpr(!kUseSoftmax)
if(!use_softmax)
locals.push_back(ck_tile::type_convert<CompDataType>(0.0f));
else
locals.push_back(-ck_tile::numeric<CompDataType>::infinity());
};
};
if constexpr(!kUseSoftmax)
if(!use_softmax)
{
// SiLu element-wise
for(CompDataType& elem : locals)
@@ -584,7 +581,7 @@ struct reference_group_hstu_attention_fwd
elem = std::exp(elem - m) / l;
}
if constexpr(kStoreLSE)
if(store_lse)
{
lse_batch_seq_nhead(i_batch, sq, i_head) = std::log(l) + m;
}