mirror of
https://github.com/ROCm/composable_kernel.git
synced 2026-07-11 17:51:40 +00:00
Adapt fmha_bwd_runner.cpp to new q, kv sequence padding
Add backward q/kv sequence padding unit tests.
This commit is contained in:
Jeff Huang
@@ -131,8 +131,16 @@ bwd_result fmha_bwd_run(mode_enum mode,
|
||||
mode == mode_enum::group && (!seqlen_kpads.empty() && seqlen_kpads[0] != -1);
|
||||
|
||||
#if 0
|
||||
std::cout << "use_qpadding: " << use_qpadding << std::endl;
|
||||
std::cout << "use_kpadding: " << use_kpadding << std::endl;
|
||||
std::cout << "seqlen_qs: " << seqlen_qs << std::endl;
|
||||
std::cout << "seqlen_ks: " << seqlen_ks << std::endl;
|
||||
if (use_qpadding) {
|
||||
std::cout << "seqlen_qpads: " << seqlen_qpads << std::endl;
|
||||
}
|
||||
if (use_kpadding) {
|
||||
std::cout << "seqlen_kpads: " << seqlen_kpads << std::endl;
|
||||
}
|
||||
#endif
|
||||
|
||||
mask_info mask = mask_info::decode(mask_str, seqlen_qs[0], seqlen_ks[0]);
|
||||
@@ -185,8 +193,11 @@ bwd_result fmha_bwd_run(mode_enum mode,
|
||||
{
|
||||
for(ck_tile::index_t wb = 0; wb < batch; ++wb)
|
||||
{
|
||||
const int32_t real_seqlen_q = seqstart_q_host[wb + 1] - seqstart_q_host[wb];
|
||||
const int32_t real_seqlen_k = seqstart_k_host[wb + 1] - seqstart_k_host[wb];
|
||||
// When padding is enabled, use logical lengths for flop/bandwidth calculation
|
||||
const int32_t real_seqlen_q =
|
||||
use_qpadding ? seqlen_qs[wb] : (seqstart_q_host[wb + 1] - seqstart_q_host[wb]);
|
||||
const int32_t real_seqlen_k =
|
||||
use_kpadding ? seqlen_ks[wb] : (seqstart_k_host[wb + 1] - seqstart_k_host[wb]);
|
||||
|
||||
if(max_seqlen_q < real_seqlen_q)
|
||||
{
|
||||
@@ -575,8 +586,18 @@ bwd_result fmha_bwd_run(mode_enum mode,
|
||||
|
||||
for(ck_tile::index_t wb = 0; wb < batch; ++wb)
|
||||
{
|
||||
const ck_tile::index_t real_seqlen_q = seqstart_q_host[wb + 1] - seqstart_q_host[wb];
|
||||
const ck_tile::index_t real_seqlen_k = seqstart_k_host[wb + 1] - seqstart_k_host[wb];
|
||||
// When padding is enabled, use logical lengths instead of computing from padded
|
||||
// prefix-sum
|
||||
const ck_tile::index_t real_seqlen_q =
|
||||
use_qpadding ? seqlen_qs[wb] : (seqstart_q_host[wb + 1] - seqstart_q_host[wb]);
|
||||
const ck_tile::index_t real_seqlen_k =
|
||||
use_kpadding ? seqlen_ks[wb] : (seqstart_k_host[wb + 1] - seqstart_k_host[wb]);
|
||||
|
||||
// Skip forward reference computation for batches with zero length sequences
|
||||
if(real_seqlen_q == 0 || real_seqlen_k == 0)
|
||||
{
|
||||
continue;
|
||||
}
|
||||
|
||||
// adjust matrix index according to the mode
|
||||
const ck_tile::index_t b = (mode == mode_enum::batch ? wb : 0);
|
||||
@@ -821,10 +842,23 @@ bwd_result fmha_bwd_run(mode_enum mode,
|
||||
dv_buf.FromDevice(dv_host.data());
|
||||
dbias_buf.FromDevice(dbias_host.data());
|
||||
|
||||
// Track the index into reference vectors (may differ from wb if batches were skipped)
|
||||
ck_tile::index_t ref_idx = 0;
|
||||
|
||||
for(ck_tile::index_t wb = 0; wb < batch; ++wb)
|
||||
{
|
||||
const ck_tile::index_t real_seqlen_q = seqstart_q_host[wb + 1] - seqstart_q_host[wb];
|
||||
const ck_tile::index_t real_seqlen_k = seqstart_k_host[wb + 1] - seqstart_k_host[wb];
|
||||
// When padding is enabled, use logical lengths instead of computing from padded
|
||||
// prefix-sum
|
||||
const ck_tile::index_t real_seqlen_q =
|
||||
use_qpadding ? seqlen_qs[wb] : (seqstart_q_host[wb + 1] - seqstart_q_host[wb]);
|
||||
const ck_tile::index_t real_seqlen_k =
|
||||
use_kpadding ? seqlen_ks[wb] : (seqstart_k_host[wb + 1] - seqstart_k_host[wb]);
|
||||
|
||||
// Skip validation for batches with zero length sequences
|
||||
if(real_seqlen_q == 0 || real_seqlen_k == 0)
|
||||
{
|
||||
continue;
|
||||
}
|
||||
|
||||
// adjust matrix index according to the mode
|
||||
const ck_tile::index_t b = (mode == mode_enum::batch ? wb : 0);
|
||||
@@ -857,14 +891,14 @@ bwd_result fmha_bwd_run(mode_enum mode,
|
||||
|
||||
// dP = dO@V x Z w/ dropout
|
||||
// dP = dO@V w/o dropout
|
||||
auto v_t_host_ref = v_host_refs[wb].transpose({0, 2, 1}); // v_g_o_n -> v_g_n_o
|
||||
auto v_t_host_ref = v_host_refs[ref_idx].transpose({0, 2, 1}); // v_g_o_n -> v_g_n_o
|
||||
ck_tile::reference_batched_gemm<OGradDataType, VDataType, AccDataType, AccDataType>(
|
||||
do_host_ref, v_t_host_ref, dp_hp_host_ref); // dp_g_m_n = do_g_m_o@v_g_n_o
|
||||
|
||||
if(p_drop > 0)
|
||||
{
|
||||
ck_tile::reference_batched_dropout(
|
||||
dp_hp_host_ref, randval_host_refs[wb], p_undrop_in_uint8_t, rp_undrop);
|
||||
dp_hp_host_ref, randval_host_refs[ref_idx], p_undrop_in_uint8_t, rp_undrop);
|
||||
}
|
||||
|
||||
// dS_i_j = P_i_j .* (dP_i_j - dO_i dot O_i)
|
||||
@@ -873,11 +907,13 @@ bwd_result fmha_bwd_run(mode_enum mode,
|
||||
AccDataType do_dot_o = 0;
|
||||
for(int o = 0; o < hdim_v; o++)
|
||||
{
|
||||
do_dot_o += ck_tile::type_convert<AccDataType>(do_host_ref(i0, i1, o)) *
|
||||
ck_tile::type_convert<AccDataType>(o_host_refs[wb](i0, i1, o));
|
||||
do_dot_o +=
|
||||
ck_tile::type_convert<AccDataType>(do_host_ref(i0, i1, o)) *
|
||||
ck_tile::type_convert<AccDataType>(o_host_refs[ref_idx](i0, i1, o));
|
||||
}
|
||||
ds_hp_host_ref(i0, i1, i2) = ck_tile::type_convert<AccDataType>(
|
||||
p_hp_host_refs[wb](i0, i1, i2) * (dp_hp_host_ref(i0, i1, i2) - do_dot_o));
|
||||
ds_hp_host_ref(i0, i1, i2) =
|
||||
ck_tile::type_convert<AccDataType>(p_hp_host_refs[ref_idx](i0, i1, i2) *
|
||||
(dp_hp_host_ref(i0, i1, i2) - do_dot_o));
|
||||
},
|
||||
ds_hp_host_ref.mDesc.get_lengths()[0],
|
||||
ds_hp_host_ref.mDesc.get_lengths()[1],
|
||||
@@ -893,14 +929,14 @@ bwd_result fmha_bwd_run(mode_enum mode,
|
||||
// dV = P_drop^T@dO^T
|
||||
// dV = P^T@dO^T w/o dropout
|
||||
auto p_t_lp_host_ref =
|
||||
p_lp_host_refs[wb].transpose({0, 2, 1}); // p_lp_g_m_n -> p_lp_g_n_m
|
||||
p_lp_host_refs[ref_idx].transpose({0, 2, 1}); // p_lp_g_m_n -> p_lp_g_n_m
|
||||
auto do_t_host_ref = do_host_ref.transpose({0, 2, 1}); // do_g_m_o -> do_g_o_m
|
||||
ck_tile::
|
||||
reference_batched_gemm<GemmDataType, OGradDataType, AccDataType, VGradDataType>(
|
||||
p_t_lp_host_ref, do_t_host_ref, dv_host_ref); // dv_g_n_o = p_lp_g_n_m@do_g_o_m
|
||||
|
||||
// dQ = scale * dS@K^T
|
||||
auto k_t_host_ref = k_host_refs[wb].transpose({0, 2, 1}); // k_g_n_k -> k_g_k_n
|
||||
auto k_t_host_ref = k_host_refs[ref_idx].transpose({0, 2, 1}); // k_g_n_k -> k_g_k_n
|
||||
ck_tile::reference_batched_gemm<GemmDataType, KDataType, AccDataType, QGradDataType>(
|
||||
ds_lp_host_ref,
|
||||
k_t_host_ref,
|
||||
@@ -910,8 +946,8 @@ bwd_result fmha_bwd_run(mode_enum mode,
|
||||
ck_tile::scales(scale)); // dq_g_m_k = ds_g_m_n@k_g_k_n
|
||||
|
||||
// dK = scale * dS^T@Q^T
|
||||
auto ds_t_lp_host_ref = ds_lp_host_ref.transpose({0, 2, 1}); // ds_g_m_n -> ds_g_n_m
|
||||
auto q_t_host_ref = q_host_refs[wb].transpose({0, 2, 1}); // q_g_m_k -> q_g_k_m
|
||||
auto ds_t_lp_host_ref = ds_lp_host_ref.transpose({0, 2, 1}); // ds_g_m_n -> ds_g_n_m
|
||||
auto q_t_host_ref = q_host_refs[ref_idx].transpose({0, 2, 1}); // q_g_m_k -> q_g_k_m
|
||||
ck_tile::reference_batched_gemm<GemmDataType, QDataType, AccDataType, KGradDataType>(
|
||||
ds_t_lp_host_ref,
|
||||
q_t_host_ref,
|
||||
@@ -985,6 +1021,9 @@ bwd_result fmha_bwd_run(mode_enum mode,
|
||||
|
||||
break;
|
||||
}
|
||||
|
||||
// Increment reference vector index for successfully validated batches
|
||||
ref_idx++;
|
||||
}
|
||||
|
||||
std::cout << ", valid:" << (pass ? "y" : "n") << std::flush << std::endl;
|
||||
|
||||
@@ -6,10 +6,6 @@ endif()
|
||||
set(FMHA_BWD_INSTANCES "tile_fmha_bwd_instances")
|
||||
set(FMHA_FWD_INSTANCES "tile_fmha_fwd_instances")
|
||||
|
||||
add_gtest_executable(test_ck_tile_fmha_bwd_kernels test_fmha_bwd_kernel_padding.cpp)
|
||||
target_link_libraries(test_ck_tile_fmha_bwd_kernels PRIVATE ${FMHA_BWD_INSTANCES})
|
||||
|
||||
|
||||
set(TEST_NAME "test_ck_tile_fmha")
|
||||
|
||||
function(add_gtest_fwd test_group)
|
||||
|
||||
@@ -248,3 +248,741 @@ INSTANTIATE_TEST_SUITE_P(TestCkTileFmhaBwd,
|
||||
Values(true) // deterministic
|
||||
));
|
||||
TEST_P(Deterministic, DataTypeConfig) { fmha_bwd_test(GetParam()); }
|
||||
|
||||
// ============================================================================
|
||||
// Q/KV Padding Tests - High Priority
|
||||
// ============================================================================
|
||||
|
||||
// 1. BasicQPadding: Test Q padding only (K/V have no padding)
|
||||
class BasicQPadding : public TestWithParam<FmhaBwdTestParam>
|
||||
{
|
||||
};
|
||||
|
||||
INSTANTIATE_TEST_SUITE_P(
|
||||
TestCkTileFmhaBwd,
|
||||
BasicQPadding,
|
||||
Combine(Values(mode_enum::group), // Only group mode supports padding
|
||||
HDimValues,
|
||||
Values(std::tuple{true, true}), // perm
|
||||
Values("n"), // no bias for basic test
|
||||
Values(false), // use_dbias
|
||||
Values(0.0f), // no dropout
|
||||
Values(std::tuple{0, 0, false}), // seed/offset/prefs
|
||||
ValuesIn([]() {
|
||||
// Define test cases with Q padding: seqlen_q < seqlen_qpad
|
||||
// Format: {batch, nhead, nhead_k, seqlen_q, seqlen_k, mask_str}
|
||||
// Note: Will set seqlen_qpad separately in the test
|
||||
std::vector<FmhaBwdDimsMaskParam> test_cases;
|
||||
|
||||
// Small padding: logical length close to physical
|
||||
test_cases.push_back(std::tuple{2, 2, 2, 127, 128, "0"}); // Q: 127->128
|
||||
test_cases.push_back(std::tuple{3, 4, 2, 250, 256, "0"}); // Q: 250->256
|
||||
|
||||
// Medium padding: ~20-30% padding
|
||||
test_cases.push_back(std::tuple{2, 2, 1, 180, 256, "0"}); // Q: 180->256
|
||||
test_cases.push_back(std::tuple{3, 3, 3, 350, 512, "1"}); // Q: 350->512, causal
|
||||
|
||||
// Large padding: ~50% padding
|
||||
test_cases.push_back(std::tuple{2, 4, 2, 128, 256, "0"}); // Q: 128->256
|
||||
test_cases.push_back(std::tuple{2, 2, 2, 200, 512, "2"}); // Q: 200->512, causal
|
||||
|
||||
return test_cases;
|
||||
}()),
|
||||
Values(false) // deterministic
|
||||
));
|
||||
|
||||
TEST_P(BasicQPadding, DataTypeConfig)
|
||||
{
|
||||
auto [mode, hdims, perm, bias_str, use_dbias, p_drop, drop_misc, dims_mask, det] = GetParam();
|
||||
auto [hdim_q, hdim_v] = hdims;
|
||||
auto [i_perm, o_perm] = perm;
|
||||
auto [drop_seed, drop_offset, drop_prefs] = drop_misc;
|
||||
auto [batch, nhead, nhead_k, seqlen_q, seqlen_k, mask_str] = dims_mask;
|
||||
|
||||
// Set up Q padding: physical length larger than logical
|
||||
std::vector<ck_tile::index_t> seqlen_qs(batch, seqlen_q);
|
||||
std::vector<ck_tile::index_t> seqlen_ks(batch, seqlen_k);
|
||||
|
||||
// Calculate physical Q length (padded)
|
||||
ck_tile::index_t seqlen_qpad = ((seqlen_q + 63) / 64) * 64; // Round up to multiple of 64
|
||||
if(seqlen_q > 256)
|
||||
seqlen_qpad = ((seqlen_q + 127) / 128) * 128; // Larger alignment for longer sequences
|
||||
|
||||
std::vector<ck_tile::index_t> seqlen_qpads(batch, seqlen_qpad);
|
||||
std::vector<ck_tile::index_t> seqlen_kpads(batch, seqlen_k); // No K padding
|
||||
|
||||
auto result = fmha_bwd_run<DataTypeConfig>(
|
||||
mode,
|
||||
batch,
|
||||
nhead,
|
||||
nhead_k,
|
||||
seqlen_qs,
|
||||
seqlen_ks,
|
||||
seqlen_qpads,
|
||||
seqlen_kpads,
|
||||
hdim_q,
|
||||
hdim_v,
|
||||
i_perm,
|
||||
o_perm,
|
||||
0, // scale
|
||||
bias_str,
|
||||
use_dbias,
|
||||
p_drop,
|
||||
drop_seed,
|
||||
drop_offset,
|
||||
drop_prefs,
|
||||
mask_str,
|
||||
det,
|
||||
init_method,
|
||||
static_cast<uint32_t>(ck_tile::EnvValue(CK_TILE_ENV(CK_TILE_TEST_SEED))),
|
||||
1,
|
||||
stream_config);
|
||||
|
||||
if(result == bwd_result::no_instance)
|
||||
GTEST_SKIP() << "No instance for Q padding with hdim_q=" << hdim_q;
|
||||
ASSERT_EQ(result, bwd_result::success);
|
||||
}
|
||||
|
||||
// 2. BasicKVPadding: Test K/V padding only (Q has no padding)
|
||||
class BasicKVPadding : public TestWithParam<FmhaBwdTestParam>
|
||||
{
|
||||
};
|
||||
|
||||
INSTANTIATE_TEST_SUITE_P(
|
||||
TestCkTileFmhaBwd,
|
||||
BasicKVPadding,
|
||||
Combine(Values(mode_enum::group),
|
||||
HDimValues,
|
||||
Values(std::tuple{true, true}),
|
||||
Values("n"),
|
||||
Values(false),
|
||||
Values(0.0f),
|
||||
Values(std::tuple{0, 0, false}),
|
||||
ValuesIn([]() {
|
||||
std::vector<FmhaBwdDimsMaskParam> test_cases;
|
||||
|
||||
// Small K/V padding
|
||||
test_cases.push_back(std::tuple{2, 2, 2, 128, 127, "0"}); // K: 127->128
|
||||
test_cases.push_back(std::tuple{3, 4, 2, 256, 250, "0"}); // K: 250->256
|
||||
|
||||
// Medium K/V padding
|
||||
test_cases.push_back(std::tuple{2, 2, 1, 256, 180, "0"}); // K: 180->256
|
||||
test_cases.push_back(std::tuple{3, 3, 3, 512, 350, "1"}); // K: 350->512
|
||||
|
||||
// Large K/V padding
|
||||
test_cases.push_back(std::tuple{2, 4, 2, 256, 128, "0"}); // K: 128->256
|
||||
test_cases.push_back(std::tuple{2, 2, 2, 512, 200, "2"}); // K: 200->512
|
||||
|
||||
return test_cases;
|
||||
}()),
|
||||
Values(false)));
|
||||
|
||||
TEST_P(BasicKVPadding, DataTypeConfig)
|
||||
{
|
||||
auto [mode, hdims, perm, bias_str, use_dbias, p_drop, drop_misc, dims_mask, det] = GetParam();
|
||||
auto [hdim_q, hdim_v] = hdims;
|
||||
auto [i_perm, o_perm] = perm;
|
||||
auto [drop_seed, drop_offset, drop_prefs] = drop_misc;
|
||||
auto [batch, nhead, nhead_k, seqlen_q, seqlen_k, mask_str] = dims_mask;
|
||||
|
||||
std::vector<ck_tile::index_t> seqlen_qs(batch, seqlen_q);
|
||||
std::vector<ck_tile::index_t> seqlen_ks(batch, seqlen_k);
|
||||
|
||||
// No Q padding
|
||||
std::vector<ck_tile::index_t> seqlen_qpads(batch, seqlen_q);
|
||||
|
||||
// Set up K/V padding
|
||||
ck_tile::index_t seqlen_kpad = ((seqlen_k + 63) / 64) * 64;
|
||||
if(seqlen_k > 256)
|
||||
seqlen_kpad = ((seqlen_k + 127) / 128) * 128;
|
||||
std::vector<ck_tile::index_t> seqlen_kpads(batch, seqlen_kpad);
|
||||
|
||||
auto result = fmha_bwd_run<DataTypeConfig>(
|
||||
mode,
|
||||
batch,
|
||||
nhead,
|
||||
nhead_k,
|
||||
seqlen_qs,
|
||||
seqlen_ks,
|
||||
seqlen_qpads,
|
||||
seqlen_kpads,
|
||||
hdim_q,
|
||||
hdim_v,
|
||||
i_perm,
|
||||
o_perm,
|
||||
0,
|
||||
bias_str,
|
||||
use_dbias,
|
||||
p_drop,
|
||||
drop_seed,
|
||||
drop_offset,
|
||||
drop_prefs,
|
||||
mask_str,
|
||||
det,
|
||||
init_method,
|
||||
static_cast<uint32_t>(ck_tile::EnvValue(CK_TILE_ENV(CK_TILE_TEST_SEED))),
|
||||
1,
|
||||
stream_config);
|
||||
|
||||
if(result == bwd_result::no_instance)
|
||||
GTEST_SKIP() << "No instance for K/V padding with hdim_q=" << hdim_q;
|
||||
ASSERT_EQ(result, bwd_result::success);
|
||||
}
|
||||
|
||||
// 3. QKVPadding: Test both Q and K/V padding simultaneously
|
||||
class QKVPadding : public TestWithParam<FmhaBwdTestParam>
|
||||
{
|
||||
};
|
||||
|
||||
INSTANTIATE_TEST_SUITE_P(
|
||||
TestCkTileFmhaBwd,
|
||||
QKVPadding,
|
||||
Combine(Values(mode_enum::group),
|
||||
HDimValues,
|
||||
Values(std::tuple{true, true}),
|
||||
Values("n"),
|
||||
Values(false),
|
||||
Values(0.0f),
|
||||
Values(std::tuple{0, 0, false}),
|
||||
ValuesIn([]() {
|
||||
std::vector<FmhaBwdDimsMaskParam> test_cases;
|
||||
|
||||
// Both Q and K have small padding
|
||||
test_cases.push_back(std::tuple{2, 2, 2, 120, 125, "0"}); // Q:120->128, K:125->128
|
||||
|
||||
// Both Q and K have medium padding
|
||||
test_cases.push_back(std::tuple{2, 4, 2, 180, 200, "0"}); // Q:180->256, K:200->256
|
||||
test_cases.push_back(std::tuple{3, 3, 3, 300, 350, "1"}); // Q:300->320, K:350->384
|
||||
|
||||
// Both Q and K have large padding
|
||||
test_cases.push_back(std::tuple{2, 2, 1, 150, 180, "0"}); // Q:150->256, K:180->256
|
||||
test_cases.push_back(std::tuple{2, 4, 2, 256, 300, "2"}); // Q:256->384, K:300->384
|
||||
|
||||
// Asymmetric padding (Q more padded than K)
|
||||
test_cases.push_back(std::tuple{2, 2, 2, 100, 200, "0"}); // Q:100->128, K:200->256
|
||||
|
||||
// Asymmetric padding (K more padded than Q)
|
||||
test_cases.push_back(std::tuple{2, 3, 1, 200, 100, "0"}); // Q:200->256, K:100->128
|
||||
|
||||
return test_cases;
|
||||
}()),
|
||||
Values(false)));
|
||||
|
||||
TEST_P(QKVPadding, DataTypeConfig)
|
||||
{
|
||||
auto [mode, hdims, perm, bias_str, use_dbias, p_drop, drop_misc, dims_mask, det] = GetParam();
|
||||
auto [hdim_q, hdim_v] = hdims;
|
||||
auto [i_perm, o_perm] = perm;
|
||||
auto [drop_seed, drop_offset, drop_prefs] = drop_misc;
|
||||
auto [batch, nhead, nhead_k, seqlen_q, seqlen_k, mask_str] = dims_mask;
|
||||
|
||||
std::vector<ck_tile::index_t> seqlen_qs(batch, seqlen_q);
|
||||
std::vector<ck_tile::index_t> seqlen_ks(batch, seqlen_k);
|
||||
|
||||
// Set up both Q and K/V padding
|
||||
ck_tile::index_t seqlen_qpad = ((seqlen_q + 63) / 64) * 64;
|
||||
if(seqlen_q > 256)
|
||||
seqlen_qpad = ((seqlen_q + 127) / 128) * 128;
|
||||
|
||||
ck_tile::index_t seqlen_kpad = ((seqlen_k + 63) / 64) * 64;
|
||||
if(seqlen_k > 256)
|
||||
seqlen_kpad = ((seqlen_k + 127) / 128) * 128;
|
||||
|
||||
std::vector<ck_tile::index_t> seqlen_qpads(batch, seqlen_qpad);
|
||||
std::vector<ck_tile::index_t> seqlen_kpads(batch, seqlen_kpad);
|
||||
|
||||
auto result = fmha_bwd_run<DataTypeConfig>(
|
||||
mode,
|
||||
batch,
|
||||
nhead,
|
||||
nhead_k,
|
||||
seqlen_qs,
|
||||
seqlen_ks,
|
||||
seqlen_qpads,
|
||||
seqlen_kpads,
|
||||
hdim_q,
|
||||
hdim_v,
|
||||
i_perm,
|
||||
o_perm,
|
||||
0,
|
||||
bias_str,
|
||||
use_dbias,
|
||||
p_drop,
|
||||
drop_seed,
|
||||
drop_offset,
|
||||
drop_prefs,
|
||||
mask_str,
|
||||
det,
|
||||
init_method,
|
||||
static_cast<uint32_t>(ck_tile::EnvValue(CK_TILE_ENV(CK_TILE_TEST_SEED))),
|
||||
1,
|
||||
stream_config);
|
||||
|
||||
if(result == bwd_result::no_instance)
|
||||
GTEST_SKIP() << "No instance for Q+K/V padding with hdim_q=" << hdim_q;
|
||||
ASSERT_EQ(result, bwd_result::success);
|
||||
}
|
||||
|
||||
// 4. ZeroLengthPadding: Test zero-length sequences with padding
|
||||
class ZeroLengthPadding : public TestWithParam<FmhaBwdTestParam>
|
||||
{
|
||||
};
|
||||
|
||||
INSTANTIATE_TEST_SUITE_P(TestCkTileFmhaBwd,
|
||||
ZeroLengthPadding,
|
||||
Combine(Values(mode_enum::group),
|
||||
Values(std::tuple{64, -1},
|
||||
std::tuple{128, -1}), // Limited hdim for edge cases
|
||||
Values(std::tuple{true, true}),
|
||||
Values("n"),
|
||||
Values(false),
|
||||
Values(0.0f),
|
||||
Values(std::tuple{0, 0, false}),
|
||||
Values(
|
||||
// Test case 1: First batch has zero Q length
|
||||
std::tuple{3, 2, 2, 0, 128, "0"},
|
||||
// Test case 2: Middle batch has zero Q length (multi-batch)
|
||||
std::tuple{3, 2, 1, 100, 128, "0"},
|
||||
// Test case 3: Last batch has zero Q length
|
||||
std::tuple{3, 3, 3, 150, 200, "0"},
|
||||
// Test case 4: Zero K length (first batch)
|
||||
std::tuple{3, 2, 2, 128, 0, "0"},
|
||||
// Test case 5: Mixed zero lengths with padding
|
||||
std::tuple{4, 2, 2, 80, 100, "0"}),
|
||||
Values(false)));
|
||||
|
||||
TEST_P(ZeroLengthPadding, DataTypeConfig)
|
||||
{
|
||||
auto [mode, hdims, perm, bias_str, use_dbias, p_drop, drop_misc, dims_mask, det] = GetParam();
|
||||
auto [hdim_q, hdim_v] = hdims;
|
||||
auto [i_perm, o_perm] = perm;
|
||||
auto [drop_seed, drop_offset, drop_prefs] = drop_misc;
|
||||
auto [batch, nhead, nhead_k, seqlen_q, seqlen_k, mask_str] = dims_mask;
|
||||
|
||||
// Create varied sequence lengths with some zero-length sequences
|
||||
std::vector<ck_tile::index_t> seqlen_qs;
|
||||
std::vector<ck_tile::index_t> seqlen_ks;
|
||||
std::vector<ck_tile::index_t> seqlen_qpads;
|
||||
std::vector<ck_tile::index_t> seqlen_kpads;
|
||||
|
||||
for(int b = 0; b < batch; ++b)
|
||||
{
|
||||
// Create pattern with zero-length sequences
|
||||
ck_tile::index_t q_len, k_len;
|
||||
|
||||
if(seqlen_q == 0 && b == 1) // Middle batch zero Q
|
||||
{
|
||||
q_len = (b == 1) ? 0 : ((b == 0) ? 100 : 80);
|
||||
k_len = seqlen_k;
|
||||
}
|
||||
else if(seqlen_k == 0 && b == 0) // First batch zero K
|
||||
{
|
||||
q_len = seqlen_q;
|
||||
k_len = (b == 0) ? 0 : 100;
|
||||
}
|
||||
else
|
||||
{
|
||||
// Varied lengths
|
||||
q_len = (b == 0 && seqlen_q == 0) ? 0 : (seqlen_q + b * 10);
|
||||
k_len = seqlen_k + b * 15;
|
||||
}
|
||||
|
||||
seqlen_qs.push_back(q_len);
|
||||
seqlen_ks.push_back(k_len);
|
||||
|
||||
// Add padding for non-zero lengths
|
||||
ck_tile::index_t qpad = (q_len == 0) ? 0 : ((q_len + 63) / 64) * 64;
|
||||
ck_tile::index_t kpad = (k_len == 0) ? 0 : ((k_len + 63) / 64) * 64;
|
||||
|
||||
seqlen_qpads.push_back(qpad);
|
||||
seqlen_kpads.push_back(kpad);
|
||||
}
|
||||
|
||||
auto result = fmha_bwd_run<DataTypeConfig>(
|
||||
mode,
|
||||
batch,
|
||||
nhead,
|
||||
nhead_k,
|
||||
seqlen_qs,
|
||||
seqlen_ks,
|
||||
seqlen_qpads,
|
||||
seqlen_kpads,
|
||||
hdim_q,
|
||||
hdim_v,
|
||||
i_perm,
|
||||
o_perm,
|
||||
0,
|
||||
bias_str,
|
||||
use_dbias,
|
||||
p_drop,
|
||||
drop_seed,
|
||||
drop_offset,
|
||||
drop_prefs,
|
||||
mask_str,
|
||||
det,
|
||||
init_method,
|
||||
static_cast<uint32_t>(ck_tile::EnvValue(CK_TILE_ENV(CK_TILE_TEST_SEED))),
|
||||
1,
|
||||
stream_config);
|
||||
|
||||
if(result == bwd_result::no_instance)
|
||||
GTEST_SKIP() << "No instance for zero-length padding";
|
||||
ASSERT_EQ(result, bwd_result::success);
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Q/KV Padding Tests - Medium Priority
|
||||
// ============================================================================
|
||||
|
||||
// 5. VariedPaddingRatios: Test different padding ratios (waste ratios)
|
||||
class VariedPaddingRatios : public TestWithParam<FmhaBwdTestParam>
|
||||
{
|
||||
};
|
||||
|
||||
INSTANTIATE_TEST_SUITE_P(
|
||||
TestCkTileFmhaBwd,
|
||||
VariedPaddingRatios,
|
||||
Combine(Values(mode_enum::group),
|
||||
HDimValues,
|
||||
Values(std::tuple{true, true}),
|
||||
Values("n"),
|
||||
Values(false),
|
||||
Values(0.0f),
|
||||
Values(std::tuple{0, 0, false}),
|
||||
ValuesIn([]() {
|
||||
std::vector<FmhaBwdDimsMaskParam> test_cases;
|
||||
|
||||
// Minimal waste: ~1-5% padding (logical ≈ physical - small delta)
|
||||
test_cases.push_back(
|
||||
std::tuple{2, 2, 2, 127, 127, "0"}); // Q:127->128 (~0.8%), K:127->128
|
||||
test_cases.push_back(
|
||||
std::tuple{2, 4, 2, 252, 250, "0"}); // Q:252->256 (~1.6%), K:250->256
|
||||
test_cases.push_back(std::tuple{2, 2, 1, 509, 505, "1"}); // Q:509->512, K:505->512
|
||||
|
||||
// Low waste: ~10-20% padding
|
||||
test_cases.push_back(
|
||||
std::tuple{2, 3, 3, 220, 210, "0"}); // Q:220->256 (~16%), K:210->256
|
||||
test_cases.push_back(
|
||||
std::tuple{3, 2, 2, 440, 420, "0"}); // Q:440->512 (~16%), K:420->512
|
||||
test_cases.push_back(std::tuple{2, 4, 2, 350, 340, "1"}); // Q:350->384, K:340->384
|
||||
|
||||
// Medium waste: ~30-40% padding
|
||||
test_cases.push_back(
|
||||
std::tuple{2, 2, 2, 180, 170, "0"}); // Q:180->256 (~42%), K:170->256
|
||||
test_cases.push_back(
|
||||
std::tuple{2, 3, 1, 320, 310, "0"}); // Q:320->384 (~20%), K:310->384
|
||||
test_cases.push_back(std::tuple{3, 2, 2, 350, 340, "2"}); // Q:350->512, K:340->512
|
||||
|
||||
// High waste: ~50%+ padding
|
||||
test_cases.push_back(
|
||||
std::tuple{2, 2, 2, 130, 130, "0"}); // Q:130->256 (~97%), K:130->256
|
||||
test_cases.push_back(
|
||||
std::tuple{2, 4, 2, 260, 260, "0"}); // Q:260->512 (~97%), K:260->512
|
||||
test_cases.push_back(
|
||||
std::tuple{2, 2, 1, 200, 200, "1"}); // Q:200->256 (~28%), K:200->256
|
||||
|
||||
// Extreme waste: very small logical vs large physical
|
||||
test_cases.push_back(std::tuple{2, 2, 2, 65, 70, "0"}); // Q:65->128, K:70->128
|
||||
test_cases.push_back(std::tuple{2, 3, 3, 100, 90, "0"}); // Q:100->128, K:90->128
|
||||
|
||||
return test_cases;
|
||||
}()),
|
||||
Values(false)));
|
||||
|
||||
TEST_P(VariedPaddingRatios, DataTypeConfig)
|
||||
{
|
||||
auto [mode, hdims, perm, bias_str, use_dbias, p_drop, drop_misc, dims_mask, det] = GetParam();
|
||||
auto [hdim_q, hdim_v] = hdims;
|
||||
auto [i_perm, o_perm] = perm;
|
||||
auto [drop_seed, drop_offset, drop_prefs] = drop_misc;
|
||||
auto [batch, nhead, nhead_k, seqlen_q, seqlen_k, mask_str] = dims_mask;
|
||||
|
||||
std::vector<ck_tile::index_t> seqlen_qs(batch, seqlen_q);
|
||||
std::vector<ck_tile::index_t> seqlen_ks(batch, seqlen_k);
|
||||
|
||||
// Calculate padding based on common alignment strategies
|
||||
auto calc_pad = [](ck_tile::index_t len) -> ck_tile::index_t {
|
||||
if(len <= 64)
|
||||
return 64;
|
||||
else if(len <= 128)
|
||||
return 128;
|
||||
else if(len <= 256)
|
||||
return 256;
|
||||
else if(len <= 384)
|
||||
return 384;
|
||||
else if(len <= 512)
|
||||
return 512;
|
||||
else
|
||||
return ((len + 127) / 128) * 128;
|
||||
};
|
||||
|
||||
std::vector<ck_tile::index_t> seqlen_qpads(batch, calc_pad(seqlen_q));
|
||||
std::vector<ck_tile::index_t> seqlen_kpads(batch, calc_pad(seqlen_k));
|
||||
|
||||
auto result = fmha_bwd_run<DataTypeConfig>(
|
||||
mode,
|
||||
batch,
|
||||
nhead,
|
||||
nhead_k,
|
||||
seqlen_qs,
|
||||
seqlen_ks,
|
||||
seqlen_qpads,
|
||||
seqlen_kpads,
|
||||
hdim_q,
|
||||
hdim_v,
|
||||
i_perm,
|
||||
o_perm,
|
||||
0,
|
||||
bias_str,
|
||||
use_dbias,
|
||||
p_drop,
|
||||
drop_seed,
|
||||
drop_offset,
|
||||
drop_prefs,
|
||||
mask_str,
|
||||
det,
|
||||
init_method,
|
||||
static_cast<uint32_t>(ck_tile::EnvValue(CK_TILE_ENV(CK_TILE_TEST_SEED))),
|
||||
1,
|
||||
stream_config);
|
||||
|
||||
if(result == bwd_result::no_instance)
|
||||
GTEST_SKIP() << "No instance for varied padding ratios";
|
||||
ASSERT_EQ(result, bwd_result::success);
|
||||
}
|
||||
|
||||
// 6. PaddingWithMask: Test padding combined with various mask types
|
||||
class PaddingWithMask : public TestWithParam<FmhaBwdTestParam>
|
||||
{
|
||||
};
|
||||
|
||||
INSTANTIATE_TEST_SUITE_P(
|
||||
TestCkTileFmhaBwd,
|
||||
PaddingWithMask,
|
||||
Combine(Values(mode_enum::group),
|
||||
Values(std::tuple{64, -1}, std::tuple{128, -1}), // Focus on common sizes
|
||||
Values(std::tuple{true, true}),
|
||||
Values("n"),
|
||||
Values(false),
|
||||
Values(0.0f),
|
||||
Values(std::tuple{0, 0, false}),
|
||||
ValuesIn([]() {
|
||||
std::vector<FmhaBwdDimsMaskParam> test_cases;
|
||||
|
||||
// No mask with padding (baseline)
|
||||
test_cases.push_back(std::tuple{2, 2, 2, 200, 180, "0"});
|
||||
|
||||
// Causal mask (top-left) with Q padding
|
||||
test_cases.push_back(std::tuple{2, 2, 2, 200, 256, "1"}); // Q padded, K exact
|
||||
test_cases.push_back(std::tuple{2, 4, 2, 180, 200, "t"}); // Both padded, causal
|
||||
|
||||
// Causal mask (bottom-right) with K/V padding
|
||||
test_cases.push_back(std::tuple{2, 2, 1, 256, 180, "2"}); // K padded, Q exact
|
||||
test_cases.push_back(
|
||||
std::tuple{2, 3, 3, 200, 180, "b"}); // Both padded, bottom-right
|
||||
|
||||
// Sliding window attention with padding
|
||||
test_cases.push_back(std::tuple{2, 2, 2, 200, 190, "t:64,32"}); // SWA + padding
|
||||
test_cases.push_back(std::tuple{2, 4, 2, 180, 170, "b:32,64"}); // SWA + padding
|
||||
test_cases.push_back(std::tuple{3, 2, 1, 220, 210, "t:100,50"}); // Larger window
|
||||
|
||||
// Sliding window with asymmetric padding
|
||||
test_cases.push_back(std::tuple{2, 2, 2, 150, 250, "t:80,40"}); // Q more padded
|
||||
test_cases.push_back(std::tuple{2, 3, 3, 250, 150, "b:50,70"}); // K more padded
|
||||
|
||||
// Mixed scenarios
|
||||
test_cases.push_back(std::tuple{2, 4, 2, 190, 185, "t:50,50"}); // Symmetric window
|
||||
test_cases.push_back(std::tuple{3, 2, 2, 300, 280, "1"}); // Multi-batch causal
|
||||
|
||||
return test_cases;
|
||||
}()),
|
||||
Values(false)));
|
||||
|
||||
TEST_P(PaddingWithMask, DataTypeConfig)
|
||||
{
|
||||
auto [mode, hdims, perm, bias_str, use_dbias, p_drop, drop_misc, dims_mask, det] = GetParam();
|
||||
auto [hdim_q, hdim_v] = hdims;
|
||||
auto [i_perm, o_perm] = perm;
|
||||
auto [drop_seed, drop_offset, drop_prefs] = drop_misc;
|
||||
auto [batch, nhead, nhead_k, seqlen_q, seqlen_k, mask_str] = dims_mask;
|
||||
|
||||
std::vector<ck_tile::index_t> seqlen_qs(batch, seqlen_q);
|
||||
std::vector<ck_tile::index_t> seqlen_ks(batch, seqlen_k);
|
||||
|
||||
// Apply padding
|
||||
ck_tile::index_t seqlen_qpad = ((seqlen_q + 63) / 64) * 64;
|
||||
ck_tile::index_t seqlen_kpad = ((seqlen_k + 63) / 64) * 64;
|
||||
|
||||
if(seqlen_q > 256)
|
||||
seqlen_qpad = ((seqlen_q + 127) / 128) * 128;
|
||||
if(seqlen_k > 256)
|
||||
seqlen_kpad = ((seqlen_k + 127) / 128) * 128;
|
||||
|
||||
std::vector<ck_tile::index_t> seqlen_qpads(batch, seqlen_qpad);
|
||||
std::vector<ck_tile::index_t> seqlen_kpads(batch, seqlen_kpad);
|
||||
|
||||
auto result = fmha_bwd_run<DataTypeConfig>(
|
||||
mode,
|
||||
batch,
|
||||
nhead,
|
||||
nhead_k,
|
||||
seqlen_qs,
|
||||
seqlen_ks,
|
||||
seqlen_qpads,
|
||||
seqlen_kpads,
|
||||
hdim_q,
|
||||
hdim_v,
|
||||
i_perm,
|
||||
o_perm,
|
||||
0,
|
||||
bias_str,
|
||||
use_dbias,
|
||||
p_drop,
|
||||
drop_seed,
|
||||
drop_offset,
|
||||
drop_prefs,
|
||||
mask_str,
|
||||
det,
|
||||
init_method,
|
||||
static_cast<uint32_t>(ck_tile::EnvValue(CK_TILE_ENV(CK_TILE_TEST_SEED))),
|
||||
1,
|
||||
stream_config);
|
||||
|
||||
if(result == bwd_result::no_instance)
|
||||
GTEST_SKIP() << "No instance for padding with mask";
|
||||
ASSERT_EQ(result, bwd_result::success);
|
||||
}
|
||||
|
||||
// 7. MultiBatchPadding: Test multiple batches with different padding configurations
|
||||
class MultiBatchPadding : public TestWithParam<FmhaBwdTestParam>
|
||||
{
|
||||
};
|
||||
|
||||
INSTANTIATE_TEST_SUITE_P(TestCkTileFmhaBwd,
|
||||
MultiBatchPadding,
|
||||
Combine(Values(mode_enum::group),
|
||||
Values(std::tuple{64, -1}, std::tuple{128, -1}),
|
||||
Values(std::tuple{true, true}),
|
||||
Values("n"),
|
||||
Values(false),
|
||||
Values(0.0f),
|
||||
Values(std::tuple{0, 0, false}),
|
||||
Values(
|
||||
// 3 batches with varied Q/K lengths and padding
|
||||
std::tuple{3, 2, 2, 150, 200, "0"},
|
||||
// 4 batches with different patterns
|
||||
std::tuple{4, 3, 3, 180, 220, "0"},
|
||||
// 5 batches with mixed scenarios
|
||||
std::tuple{5, 2, 1, 120, 160, "1"},
|
||||
// 3 batches with causal mask
|
||||
std::tuple{3, 4, 2, 200, 180, "t"},
|
||||
// 4 batches with sliding window
|
||||
std::tuple{4, 2, 2, 160, 140, "t:50,30"}),
|
||||
Values(false)));
|
||||
|
||||
TEST_P(MultiBatchPadding, DataTypeConfig)
|
||||
{
|
||||
auto [mode, hdims, perm, bias_str, use_dbias, p_drop, drop_misc, dims_mask, det] = GetParam();
|
||||
auto [hdim_q, hdim_v] = hdims;
|
||||
auto [i_perm, o_perm] = perm;
|
||||
auto [drop_seed, drop_offset, drop_prefs] = drop_misc;
|
||||
auto [batch, nhead, nhead_k, base_seqlen_q, base_seqlen_k, mask_str] = dims_mask;
|
||||
|
||||
// Create varied sequence lengths for each batch
|
||||
std::vector<ck_tile::index_t> seqlen_qs;
|
||||
std::vector<ck_tile::index_t> seqlen_ks;
|
||||
std::vector<ck_tile::index_t> seqlen_qpads;
|
||||
std::vector<ck_tile::index_t> seqlen_kpads;
|
||||
|
||||
for(int b = 0; b < batch; ++b)
|
||||
{
|
||||
// Generate varied lengths across batches
|
||||
// Pattern: decreasing, increasing, or random variation
|
||||
ck_tile::index_t q_len, k_len;
|
||||
|
||||
switch(b % 3)
|
||||
{
|
||||
case 0: // Decreasing
|
||||
q_len = base_seqlen_q - b * 20;
|
||||
k_len = base_seqlen_k - b * 25;
|
||||
break;
|
||||
case 1: // Increasing
|
||||
q_len = base_seqlen_q + b * 15;
|
||||
k_len = base_seqlen_k + b * 20;
|
||||
break;
|
||||
case 2: // Mixed
|
||||
q_len = base_seqlen_q + (b % 2 == 0 ? 10 : -10) * b;
|
||||
k_len = base_seqlen_k + (b % 2 == 0 ? -15 : 15) * b;
|
||||
break;
|
||||
}
|
||||
|
||||
// Ensure positive lengths
|
||||
q_len = std::max<ck_tile::index_t>(64, q_len);
|
||||
k_len = std::max<ck_tile::index_t>(64, k_len);
|
||||
|
||||
seqlen_qs.push_back(q_len);
|
||||
seqlen_ks.push_back(k_len);
|
||||
|
||||
// Calculate different padding strategies per batch
|
||||
ck_tile::index_t qpad, kpad;
|
||||
|
||||
if(b % 4 == 0)
|
||||
{
|
||||
// Tight padding (minimal waste)
|
||||
qpad = ((q_len + 31) / 32) * 32;
|
||||
kpad = ((k_len + 31) / 32) * 32;
|
||||
}
|
||||
else if(b % 4 == 1)
|
||||
{
|
||||
// Medium padding
|
||||
qpad = ((q_len + 63) / 64) * 64;
|
||||
kpad = ((k_len + 63) / 64) * 64;
|
||||
}
|
||||
else if(b % 4 == 2)
|
||||
{
|
||||
// Loose padding
|
||||
qpad = ((q_len + 127) / 128) * 128;
|
||||
kpad = ((k_len + 127) / 128) * 128;
|
||||
}
|
||||
else
|
||||
{
|
||||
// Mixed: Q tight, K loose
|
||||
qpad = ((q_len + 31) / 32) * 32;
|
||||
kpad = ((k_len + 127) / 128) * 128;
|
||||
}
|
||||
|
||||
seqlen_qpads.push_back(qpad);
|
||||
seqlen_kpads.push_back(kpad);
|
||||
}
|
||||
|
||||
auto result = fmha_bwd_run<DataTypeConfig>(
|
||||
mode,
|
||||
batch,
|
||||
nhead,
|
||||
nhead_k,
|
||||
seqlen_qs,
|
||||
seqlen_ks,
|
||||
seqlen_qpads,
|
||||
seqlen_kpads,
|
||||
hdim_q,
|
||||
hdim_v,
|
||||
i_perm,
|
||||
o_perm,
|
||||
0,
|
||||
bias_str,
|
||||
use_dbias,
|
||||
p_drop,
|
||||
drop_seed,
|
||||
drop_offset,
|
||||
drop_prefs,
|
||||
mask_str,
|
||||
det,
|
||||
init_method,
|
||||
static_cast<uint32_t>(ck_tile::EnvValue(CK_TILE_ENV(CK_TILE_TEST_SEED))),
|
||||
1,
|
||||
stream_config);
|
||||
|
||||
if(result == bwd_result::no_instance)
|
||||
GTEST_SKIP() << "No instance for multi-batch padding";
|
||||
ASSERT_EQ(result, bwd_result::success);
|
||||
}
|
||||
|
||||
@@ -1,715 +0,0 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
|
||||
|
||||
#include <hip/hip_runtime.h>
|
||||
#include <cmath>
|
||||
#include <functional>
|
||||
#include <initializer_list>
|
||||
#include <vector>
|
||||
#include "ck_tile/host.hpp"
|
||||
#include "ck_tile/host/device_memory.hpp"
|
||||
#include "ck_tile/host/kernel_launch.hpp"
|
||||
#include "example/ck_tile/01_fmha/fmha_bwd.hpp"
|
||||
#include "example/ck_tile/01_fmha/fmha_bwd_runner.hpp" // for get_elimit
|
||||
#include "gtest/gtest.h"
|
||||
|
||||
namespace {
|
||||
|
||||
using bf16 = ck_tile::bf16_t;
|
||||
using ck_tile::DeviceMem;
|
||||
|
||||
const ck_tile::stream_config kStreamConfig{
|
||||
nullptr, // stream_id_
|
||||
false, // time_kernel_
|
||||
1, // log_level_
|
||||
0, // cold_niters_
|
||||
1, // nrepeat_
|
||||
true, // is_gpu_timer_
|
||||
false, // flush_cache_
|
||||
1, // rotating_count_
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
std::vector<T> MakeVectorFromFunction(size_t count, std::function<float(size_t)> fn)
|
||||
{
|
||||
std::vector<T> data(count);
|
||||
for(size_t i = 0; i < count; ++i)
|
||||
{
|
||||
data[i] = static_cast<T>(fn(i));
|
||||
}
|
||||
return data;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
std::vector<float> ToFloatVector(const std::vector<T>& src)
|
||||
{
|
||||
std::vector<float> dst(src.size());
|
||||
for(size_t i = 0; i < src.size(); ++i)
|
||||
{
|
||||
dst[i] = ck_tile::type_convert<float>(src[i]);
|
||||
}
|
||||
return dst;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
std::vector<T> CopyDeviceToHost(const ck_tile::DeviceMem& dev, size_t element_count)
|
||||
{
|
||||
std::vector<T> host(element_count);
|
||||
if(element_count > 0)
|
||||
{
|
||||
dev.FromDevice(host.data());
|
||||
}
|
||||
return host;
|
||||
}
|
||||
|
||||
float SentinelValue() { return -999.f; }
|
||||
|
||||
} // namespace
|
||||
|
||||
// Typed tests over {fp32, fp16, bf16}
|
||||
template <typename DataTypeConfig>
|
||||
class FmhaBwdKernelPaddingTyped : public ::testing::Test
|
||||
{
|
||||
};
|
||||
|
||||
using KernelPaddingTypes = ::testing::Types<FmhaBwdFp32, FmhaBwdFp16, FmhaBwdBf16>;
|
||||
TYPED_TEST_SUITE(FmhaBwdKernelPaddingTyped, KernelPaddingTypes);
|
||||
|
||||
TYPED_TEST(FmhaBwdKernelPaddingTyped, OGradDotO_GroupPaddingRespectsLogicalLengths)
|
||||
{
|
||||
constexpr ck_tile::index_t batch = 2;
|
||||
constexpr ck_tile::index_t nhead = 1;
|
||||
constexpr ck_tile::index_t hdim = 128;
|
||||
constexpr ck_tile::index_t phys_rows0 = 8;
|
||||
constexpr ck_tile::index_t phys_rows1 = 8;
|
||||
constexpr ck_tile::index_t max_phys = phys_rows0; // both batches equal
|
||||
|
||||
const std::vector<int32_t> seqstart_q_host{0, phys_rows0, phys_rows0 + phys_rows1};
|
||||
const std::vector<int32_t> seqlen_q_host{5, 3};
|
||||
|
||||
const ck_tile::index_t total_rows = seqstart_q_host.back();
|
||||
|
||||
// Types per config
|
||||
using TypeConfig = FmhaBwdTypeConfig<TypeParam>;
|
||||
using OType = typename TypeConfig::ODataType;
|
||||
using DOType = typename TypeConfig::OGradDataType;
|
||||
using DType = typename TypeConfig::DDataType; // float under bf16 config
|
||||
using AccType = typename TypeConfig::AccDataType; // float
|
||||
|
||||
// Host tensors laid out as [b, h, s, d] with b=1, h=1 under group mode
|
||||
ck_tile::HostTensor<OType> o_host({1, nhead, total_rows, hdim});
|
||||
ck_tile::HostTensor<DOType> do_host({1, nhead, total_rows, hdim});
|
||||
ck_tile::HostTensor<DType> d_init_host({1, nhead, total_rows});
|
||||
|
||||
// Initialize O/dO with constants using FillConstant (no manual bf16 casts)
|
||||
const float o_const = 0.25f;
|
||||
const float do_const = 0.5f;
|
||||
ck_tile::FillConstant<OType>{ck_tile::type_convert<OType>(o_const)}(o_host);
|
||||
ck_tile::FillConstant<DOType>{ck_tile::type_convert<DOType>(do_const)}(do_host);
|
||||
ck_tile::FillConstant<DType>{ck_tile::type_convert<DType>(SentinelValue())}(d_init_host);
|
||||
|
||||
// Prepare expected D via runner-style CPU reference, sentinel elsewhere
|
||||
std::vector<float> expected(static_cast<size_t>(total_rows), SentinelValue());
|
||||
for(ck_tile::index_t b = 0; b < batch; ++b)
|
||||
{
|
||||
const ck_tile::index_t start = seqstart_q_host[b];
|
||||
const ck_tile::index_t len = seqlen_q_host[b];
|
||||
for(ck_tile::index_t row = 0; row < len; ++row)
|
||||
{
|
||||
AccType acc = 0;
|
||||
for(ck_tile::index_t c = 0; c < hdim; ++c)
|
||||
{
|
||||
// o_host/do_host are [1, nhead, s, d]
|
||||
const auto o_val = ck_tile::type_convert<AccType>(o_host(0, 0, start + row, c));
|
||||
const auto do_val = ck_tile::type_convert<AccType>(do_host(0, 0, start + row, c));
|
||||
acc += do_val * o_val;
|
||||
}
|
||||
expected[start + row] = ck_tile::type_convert<float>(acc);
|
||||
}
|
||||
}
|
||||
std::vector<float> sentinel_ref(static_cast<size_t>(total_rows), SentinelValue());
|
||||
|
||||
// Device buffers
|
||||
ck_tile::DeviceMem o_dev(o_host.get_element_space_size_in_bytes());
|
||||
ck_tile::DeviceMem do_dev(do_host.get_element_space_size_in_bytes());
|
||||
ck_tile::DeviceMem d_dev(d_init_host.get_element_space_size_in_bytes());
|
||||
ck_tile::DeviceMem seqstart_dev(seqstart_q_host.size() * sizeof(int32_t));
|
||||
ck_tile::DeviceMem seqlen_dev(seqlen_q_host.size() * sizeof(int32_t));
|
||||
|
||||
o_dev.ToDevice(o_host.data());
|
||||
do_dev.ToDevice(do_host.data());
|
||||
d_dev.ToDevice(d_init_host.data());
|
||||
seqstart_dev.ToDevice(seqstart_q_host.data());
|
||||
seqlen_dev.ToDevice(seqlen_q_host.data());
|
||||
|
||||
fmha_bwd_args args{};
|
||||
args.q_ptr = nullptr;
|
||||
args.k_ptr = nullptr;
|
||||
args.v_ptr = nullptr;
|
||||
args.bias_ptr = nullptr;
|
||||
args.o_ptr = o_dev.GetDeviceBuffer();
|
||||
args.lse_ptr = nullptr;
|
||||
args.do_ptr = do_dev.GetDeviceBuffer();
|
||||
args.d_ptr = d_dev.GetDeviceBuffer();
|
||||
args.rand_val_ptr = nullptr;
|
||||
args.dq_ptr = nullptr;
|
||||
args.dk_ptr = nullptr;
|
||||
args.dv_ptr = nullptr;
|
||||
args.dbias_ptr = nullptr;
|
||||
args.dq_acc_ptr = nullptr;
|
||||
args.seqstart_q_ptr = seqstart_dev.GetDeviceBuffer();
|
||||
args.seqstart_k_ptr = nullptr;
|
||||
args.seqlen_k_ptr = nullptr;
|
||||
args.seqlen_q_ptr = seqlen_dev.GetDeviceBuffer();
|
||||
args.seqlen_q = 0;
|
||||
args.seqlen_k = 0;
|
||||
args.batch = batch;
|
||||
args.max_seqlen_q = max_phys;
|
||||
args.max_seqlen_k = 0;
|
||||
args.hdim_q = hdim;
|
||||
args.hdim_v = hdim;
|
||||
args.nhead_q = nhead;
|
||||
args.nhead_k = nhead;
|
||||
args.scale = 1.0f;
|
||||
args.stride_q = 0;
|
||||
args.stride_k = 0;
|
||||
args.stride_v = 0;
|
||||
args.stride_bias = 0;
|
||||
args.stride_o = hdim;
|
||||
args.stride_randval = 0;
|
||||
args.stride_do = hdim;
|
||||
args.stride_dq_acc = 0;
|
||||
args.stride_dq = 0;
|
||||
args.stride_dk = 0;
|
||||
args.stride_dv = 0;
|
||||
args.stride_dbias = 0;
|
||||
args.nhead_stride_q = 0;
|
||||
args.nhead_stride_k = 0;
|
||||
args.nhead_stride_v = 0;
|
||||
args.nhead_stride_bias = 0;
|
||||
args.nhead_stride_o = max_phys * hdim;
|
||||
args.nhead_stride_randval = 0;
|
||||
args.nhead_stride_do = max_phys * hdim;
|
||||
args.nhead_stride_lsed = max_phys;
|
||||
args.nhead_stride_dq_acc = 0;
|
||||
args.nhead_stride_dq = 0;
|
||||
args.nhead_stride_dk = 0;
|
||||
args.nhead_stride_dv = 0;
|
||||
args.nhead_stride_dbias = 0;
|
||||
args.batch_stride_q = 0;
|
||||
args.batch_stride_k = 0;
|
||||
args.batch_stride_v = 0;
|
||||
args.batch_stride_bias = 0;
|
||||
args.batch_stride_o = 0;
|
||||
args.batch_stride_randval = 0;
|
||||
args.batch_stride_do = 0;
|
||||
args.batch_stride_lsed = 0;
|
||||
args.batch_stride_dq_acc = 0;
|
||||
args.batch_stride_dq = 0;
|
||||
args.batch_stride_dk = 0;
|
||||
args.batch_stride_dv = 0;
|
||||
args.batch_stride_dbias = 0;
|
||||
args.split_stride_dq_acc = 0;
|
||||
args.window_size_left = -1;
|
||||
args.window_size_right = 0;
|
||||
args.mask_type = static_cast<ck_tile::index_t>(mask_enum::no_mask);
|
||||
args.p_drop = 0.0f;
|
||||
args.p_undrop = 1.0f;
|
||||
args.drop_seed_offset = std::make_pair(uint64_t{0}, uint64_t{0});
|
||||
|
||||
using DotTileTraits = ck_tile::TileFmhaBwdOGradDotOTraits<true, true, 2>;
|
||||
using DotProblem =
|
||||
ck_tile::BlockFmhaBwdOGradDotOPipelineProblem<typename TypeConfig::ODataType,
|
||||
typename TypeConfig::OGradDataType,
|
||||
typename TypeConfig::DDataType,
|
||||
64,
|
||||
hdim,
|
||||
true,
|
||||
DotTileTraits>;
|
||||
using DotPipeline = ck_tile::BlockFmhaBwdOGradDotO<DotProblem>;
|
||||
using DotKernel = ck_tile::FmhaBwdOGradDotOKernel<DotPipeline>;
|
||||
|
||||
auto [dot_kargs, dot_grids] = fmha_bwd_dot_do_o_create_kargs_and_grids<DotKernel>(args);
|
||||
const dim3 dot_blocks = DotKernel::BlockSize();
|
||||
constexpr ck_tile::index_t kDotBlockPerCu = DotKernel::kBlockPerCu;
|
||||
auto dot_kernel =
|
||||
ck_tile::make_kernel<kDotBlockPerCu>(DotKernel{}, dot_grids, dot_blocks, 0, dot_kargs);
|
||||
dot_kernel(kStreamConfig);
|
||||
ASSERT_EQ(hipDeviceSynchronize(), hipSuccess);
|
||||
|
||||
auto d_result_host = CopyDeviceToHost<float>(d_dev, total_rows);
|
||||
|
||||
auto [rtol_doto, atol_doto] = get_elimit<TypeParam>(hdim, hdim);
|
||||
for(size_t i = 0; i < d_result_host.size(); ++i)
|
||||
{
|
||||
SCOPED_TRACE(::testing::Message() << "index=" << i);
|
||||
if(std::fabs(expected[i] - sentinel_ref[i]) < 1e-6f)
|
||||
{
|
||||
EXPECT_FLOAT_EQ(d_result_host[i], sentinel_ref[i]);
|
||||
}
|
||||
else
|
||||
{
|
||||
EXPECT_NEAR(d_result_host[i], expected[i], static_cast<float>(atol_doto));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
TYPED_TEST(FmhaBwdKernelPaddingTyped, OGradDotO_VariedPhysicalAndZeroLogical)
|
||||
{
|
||||
constexpr ck_tile::index_t batch = 3;
|
||||
constexpr ck_tile::index_t nhead = 1;
|
||||
constexpr ck_tile::index_t hdim = 64;
|
||||
constexpr ck_tile::index_t phys_r0 = 5;
|
||||
constexpr ck_tile::index_t phys_r1 = 7;
|
||||
constexpr ck_tile::index_t phys_r2 = 4;
|
||||
constexpr ck_tile::index_t max_phys = phys_r1;
|
||||
|
||||
const std::vector<int32_t> seqstart_q_host{
|
||||
0, phys_r0, phys_r0 + phys_r1, phys_r0 + phys_r1 + phys_r2};
|
||||
const std::vector<int32_t> seqlen_q_host{3, 0, 4};
|
||||
const ck_tile::index_t total_rows = seqstart_q_host.back();
|
||||
|
||||
using TypeConfig = FmhaBwdTypeConfig<TypeParam>;
|
||||
using OType = typename TypeConfig::ODataType;
|
||||
using DOType = typename TypeConfig::OGradDataType;
|
||||
using DType = typename TypeConfig::DDataType;
|
||||
|
||||
ck_tile::HostTensor<OType> o_host({1, nhead, total_rows, hdim});
|
||||
ck_tile::HostTensor<DOType> do_host({1, nhead, total_rows, hdim});
|
||||
ck_tile::HostTensor<DType> d_init_host({1, nhead, total_rows});
|
||||
ck_tile::FillConstant<OType>{ck_tile::type_convert<OType>(1.0f)}(o_host);
|
||||
ck_tile::FillConstant<DOType>{ck_tile::type_convert<DOType>(2.0f)}(do_host);
|
||||
ck_tile::FillConstant<DType>{ck_tile::type_convert<DType>(SentinelValue())}(d_init_host);
|
||||
|
||||
std::vector<float> expected(static_cast<size_t>(total_rows), SentinelValue());
|
||||
const float dot = 2.0f * 1.0f * static_cast<float>(hdim);
|
||||
for(ck_tile::index_t b = 0; b < batch; ++b)
|
||||
{
|
||||
const ck_tile::index_t start = seqstart_q_host[b];
|
||||
const ck_tile::index_t len = seqlen_q_host[b];
|
||||
for(ck_tile::index_t row = 0; row < len; ++row)
|
||||
expected[start + row] = dot;
|
||||
}
|
||||
std::vector<float> sentinel_ref(static_cast<size_t>(total_rows), SentinelValue());
|
||||
|
||||
ck_tile::DeviceMem o_dev(o_host.get_element_space_size_in_bytes());
|
||||
ck_tile::DeviceMem do_dev(do_host.get_element_space_size_in_bytes());
|
||||
ck_tile::DeviceMem d_dev(d_init_host.get_element_space_size_in_bytes());
|
||||
ck_tile::DeviceMem seqstart_dev(seqstart_q_host.size() * sizeof(int32_t));
|
||||
ck_tile::DeviceMem seqlen_dev(seqlen_q_host.size() * sizeof(int32_t));
|
||||
o_dev.ToDevice(o_host.data());
|
||||
do_dev.ToDevice(do_host.data());
|
||||
d_dev.ToDevice(d_init_host.data());
|
||||
seqstart_dev.ToDevice(seqstart_q_host.data());
|
||||
seqlen_dev.ToDevice(seqlen_q_host.data());
|
||||
|
||||
fmha_bwd_args args{};
|
||||
args.o_ptr = o_dev.GetDeviceBuffer();
|
||||
args.do_ptr = do_dev.GetDeviceBuffer();
|
||||
args.d_ptr = d_dev.GetDeviceBuffer();
|
||||
args.seqstart_q_ptr = seqstart_dev.GetDeviceBuffer();
|
||||
args.seqlen_q_ptr = seqlen_dev.GetDeviceBuffer();
|
||||
args.batch = batch;
|
||||
args.max_seqlen_q = max_phys;
|
||||
args.hdim_v = hdim;
|
||||
args.nhead_q = nhead;
|
||||
args.nhead_k = nhead;
|
||||
args.stride_o = hdim;
|
||||
args.stride_do = hdim;
|
||||
args.nhead_stride_o = max_phys * hdim;
|
||||
args.nhead_stride_do = max_phys * hdim;
|
||||
args.nhead_stride_lsed = max_phys;
|
||||
args.p_undrop = 1.0f;
|
||||
|
||||
using DotTileTraits = ck_tile::TileFmhaBwdOGradDotOTraits<true, true, 2>;
|
||||
using DotProblem =
|
||||
ck_tile::BlockFmhaBwdOGradDotOPipelineProblem<typename TypeConfig::ODataType,
|
||||
typename TypeConfig::OGradDataType,
|
||||
typename TypeConfig::DDataType,
|
||||
64,
|
||||
hdim,
|
||||
true,
|
||||
DotTileTraits>;
|
||||
using DotPipeline = ck_tile::BlockFmhaBwdOGradDotO<DotProblem>;
|
||||
using DotKernel = ck_tile::FmhaBwdOGradDotOKernel<DotPipeline>;
|
||||
|
||||
auto [dot_kargs, dot_grids] = fmha_bwd_dot_do_o_create_kargs_and_grids<DotKernel>(args);
|
||||
const dim3 dot_blocks = DotKernel::BlockSize();
|
||||
constexpr ck_tile::index_t kDotBlockPerCu = DotKernel::kBlockPerCu;
|
||||
auto dot_kernel =
|
||||
ck_tile::make_kernel<kDotBlockPerCu>(DotKernel{}, dot_grids, dot_blocks, 0, dot_kargs);
|
||||
dot_kernel(kStreamConfig);
|
||||
ASSERT_EQ(hipDeviceSynchronize(), hipSuccess);
|
||||
|
||||
auto d_result_host = CopyDeviceToHost<float>(d_dev, total_rows);
|
||||
auto [rtol, atol] = get_elimit<TypeParam>(hdim, hdim);
|
||||
for(size_t i = 0; i < d_result_host.size(); ++i)
|
||||
{
|
||||
SCOPED_TRACE(::testing::Message() << "index=" << i);
|
||||
if(std::fabs(expected[i] - sentinel_ref[i]) < 1e-6f)
|
||||
EXPECT_FLOAT_EQ(d_result_host[i], sentinel_ref[i]);
|
||||
else
|
||||
EXPECT_NEAR(d_result_host[i], expected[i], static_cast<float>(atol));
|
||||
}
|
||||
}
|
||||
|
||||
TYPED_TEST(FmhaBwdKernelPaddingTyped, OGradDotO_VariedPhysical_NoLogicalPtr)
|
||||
{
|
||||
constexpr ck_tile::index_t batch = 3;
|
||||
constexpr ck_tile::index_t nhead = 1;
|
||||
constexpr ck_tile::index_t hdim = 64;
|
||||
constexpr ck_tile::index_t phys_r0 = 5;
|
||||
constexpr ck_tile::index_t phys_r1 = 7;
|
||||
constexpr ck_tile::index_t phys_r2 = 4;
|
||||
constexpr ck_tile::index_t max_phys = phys_r1;
|
||||
|
||||
const std::vector<int32_t> seqstart_q_host{
|
||||
0, phys_r0, phys_r0 + phys_r1, phys_r0 + phys_r1 + phys_r2};
|
||||
const ck_tile::index_t total_rows = seqstart_q_host.back();
|
||||
|
||||
using TypeConfig = FmhaBwdTypeConfig<TypeParam>;
|
||||
using OType = typename TypeConfig::ODataType;
|
||||
using DOType = typename TypeConfig::OGradDataType;
|
||||
using DType = typename TypeConfig::DDataType;
|
||||
|
||||
ck_tile::HostTensor<OType> o_host({1, nhead, total_rows, hdim});
|
||||
ck_tile::HostTensor<DOType> do_host({1, nhead, total_rows, hdim});
|
||||
ck_tile::HostTensor<DType> d_init_host({1, nhead, total_rows});
|
||||
ck_tile::FillConstant<OType>{ck_tile::type_convert<OType>(1.0f)}(o_host);
|
||||
ck_tile::FillConstant<DOType>{ck_tile::type_convert<DOType>(2.0f)}(do_host);
|
||||
ck_tile::FillConstant<DType>{ck_tile::type_convert<DType>(SentinelValue())}(d_init_host);
|
||||
|
||||
std::vector<float> expected(static_cast<size_t>(total_rows), SentinelValue());
|
||||
const float dot = 2.0f * 1.0f * static_cast<float>(hdim);
|
||||
// seqlen_q_ptr is null; logical lengths equal physical lengths per group
|
||||
for(int r = 0; r < phys_r0; ++r)
|
||||
expected[0 + r] = dot;
|
||||
for(int r = 0; r < phys_r1; ++r)
|
||||
expected[phys_r0 + r] = dot;
|
||||
for(int r = 0; r < phys_r2; ++r)
|
||||
expected[phys_r0 + phys_r1 + r] = dot;
|
||||
std::vector<float> sentinel_ref(static_cast<size_t>(total_rows), SentinelValue());
|
||||
|
||||
ck_tile::DeviceMem o_dev(o_host.get_element_space_size_in_bytes());
|
||||
ck_tile::DeviceMem do_dev(do_host.get_element_space_size_in_bytes());
|
||||
ck_tile::DeviceMem d_dev(d_init_host.get_element_space_size_in_bytes());
|
||||
ck_tile::DeviceMem seqstart_dev(seqstart_q_host.size() * sizeof(int32_t));
|
||||
o_dev.ToDevice(o_host.data());
|
||||
do_dev.ToDevice(do_host.data());
|
||||
d_dev.ToDevice(d_init_host.data());
|
||||
seqstart_dev.ToDevice(seqstart_q_host.data());
|
||||
|
||||
fmha_bwd_args args{};
|
||||
args.o_ptr = o_dev.GetDeviceBuffer();
|
||||
args.do_ptr = do_dev.GetDeviceBuffer();
|
||||
args.d_ptr = d_dev.GetDeviceBuffer();
|
||||
args.seqstart_q_ptr = seqstart_dev.GetDeviceBuffer();
|
||||
args.seqlen_q_ptr = nullptr; // no logical len ptr
|
||||
args.batch = batch;
|
||||
args.max_seqlen_q = max_phys;
|
||||
args.hdim_v = hdim;
|
||||
args.nhead_q = nhead;
|
||||
args.nhead_k = nhead;
|
||||
args.stride_o = hdim;
|
||||
args.stride_do = hdim;
|
||||
args.nhead_stride_o = max_phys * hdim;
|
||||
args.nhead_stride_do = max_phys * hdim;
|
||||
args.nhead_stride_lsed = max_phys;
|
||||
args.p_undrop = 1.0f;
|
||||
|
||||
using DotTileTraits = ck_tile::TileFmhaBwdOGradDotOTraits<true, true, 2>;
|
||||
using DotProblem =
|
||||
ck_tile::BlockFmhaBwdOGradDotOPipelineProblem<typename TypeConfig::ODataType,
|
||||
typename TypeConfig::OGradDataType,
|
||||
typename TypeConfig::DDataType,
|
||||
64,
|
||||
hdim,
|
||||
true,
|
||||
DotTileTraits>;
|
||||
using DotPipeline = ck_tile::BlockFmhaBwdOGradDotO<DotProblem>;
|
||||
using DotKernel = ck_tile::FmhaBwdOGradDotOKernel<DotPipeline>;
|
||||
|
||||
auto [dot_kargs, dot_grids] = fmha_bwd_dot_do_o_create_kargs_and_grids<DotKernel>(args);
|
||||
const dim3 dot_blocks = DotKernel::BlockSize();
|
||||
constexpr ck_tile::index_t kDotBlockPerCu = DotKernel::kBlockPerCu;
|
||||
auto dot_kernel =
|
||||
ck_tile::make_kernel<kDotBlockPerCu>(DotKernel{}, dot_grids, dot_blocks, 0, dot_kargs);
|
||||
dot_kernel(kStreamConfig);
|
||||
ASSERT_EQ(hipDeviceSynchronize(), hipSuccess);
|
||||
|
||||
auto d_result_host = CopyDeviceToHost<float>(d_dev, total_rows);
|
||||
auto [rtol, atol] = get_elimit<TypeParam>(hdim, hdim);
|
||||
for(size_t i = 0; i < d_result_host.size(); ++i)
|
||||
{
|
||||
SCOPED_TRACE(::testing::Message() << "index=" << i);
|
||||
if(std::fabs(expected[i] - sentinel_ref[i]) < 1e-6f)
|
||||
EXPECT_FLOAT_EQ(d_result_host[i], sentinel_ref[i]);
|
||||
else
|
||||
EXPECT_NEAR(d_result_host[i], expected[i], static_cast<float>(atol));
|
||||
}
|
||||
}
|
||||
|
||||
TYPED_TEST(FmhaBwdKernelPaddingTyped, ConvertQGrad_GroupPaddingAndZeroLength)
|
||||
{
|
||||
constexpr ck_tile::index_t batch = 3;
|
||||
constexpr ck_tile::index_t nhead = 1;
|
||||
constexpr ck_tile::index_t hdim = 128;
|
||||
|
||||
const std::vector<int32_t> seqstart_q_host{0, 6, 6, 10}; // physical lengths: 6,0,4
|
||||
const std::vector<int32_t> seqlen_q_host{4, 0, 3};
|
||||
const std::vector<int32_t> seqstart_k_host{0, 7, 15, 18};
|
||||
const std::vector<int32_t> seqlen_k_host{5, 8, 3};
|
||||
|
||||
const ck_tile::index_t total_rows_q = seqstart_q_host.back();
|
||||
|
||||
using TypeConfigC = FmhaBwdTypeConfig<TypeParam>;
|
||||
using AccType = typename TypeConfigC::AccDataType; // float
|
||||
using QGradType = typename TypeConfigC::QGradDataType; // bf16
|
||||
|
||||
ck_tile::HostTensor<AccType> dq_acc_host({1, nhead, total_rows_q, hdim});
|
||||
ck_tile::HostTensor<QGradType> dq_host_init({1, nhead, total_rows_q, hdim});
|
||||
|
||||
const float dq_acc_const = 1.25f;
|
||||
ck_tile::FillConstant<AccType>{ck_tile::type_convert<AccType>(dq_acc_const)}(dq_acc_host);
|
||||
ck_tile::FillConstant<QGradType>{ck_tile::type_convert<QGradType>(SentinelValue())}(
|
||||
dq_host_init);
|
||||
|
||||
const float dq_sentinel_val =
|
||||
ck_tile::type_convert<float>(ck_tile::type_convert<QGradType>(SentinelValue()));
|
||||
std::vector<float> dq_sentinel_ref(static_cast<size_t>(total_rows_q * hdim), dq_sentinel_val);
|
||||
std::vector<float> expected = dq_sentinel_ref;
|
||||
for(ck_tile::index_t b = 0; b < batch; ++b)
|
||||
{
|
||||
const ck_tile::index_t q_start = seqstart_q_host[b];
|
||||
const ck_tile::index_t q_len = seqlen_q_host[b];
|
||||
for(ck_tile::index_t row = 0; row < q_len; ++row)
|
||||
{
|
||||
for(ck_tile::index_t c = 0; c < hdim; ++c)
|
||||
{
|
||||
const size_t idx = (q_start + row) * hdim + c;
|
||||
// dq_acc_host is [1, nhead, s, d]
|
||||
expected[idx] = ck_tile::type_convert<float>(dq_acc_host(0, 0, q_start + row, c));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
ck_tile::DeviceMem dq_acc_dev(dq_acc_host.get_element_space_size_in_bytes());
|
||||
ck_tile::DeviceMem dq_dev(dq_host_init.get_element_space_size_in_bytes());
|
||||
ck_tile::DeviceMem seqstart_q(seqstart_q_host.size() * sizeof(int32_t));
|
||||
ck_tile::DeviceMem seqstart_k(seqstart_k_host.size() * sizeof(int32_t));
|
||||
ck_tile::DeviceMem seqlen_q_dev(seqlen_q_host.size() * sizeof(int32_t));
|
||||
ck_tile::DeviceMem seqlen_k_dev(seqlen_k_host.size() * sizeof(int32_t));
|
||||
|
||||
dq_acc_dev.ToDevice(dq_acc_host.data());
|
||||
dq_dev.ToDevice(dq_host_init.data());
|
||||
seqstart_q.ToDevice(seqstart_q_host.data());
|
||||
seqstart_k.ToDevice(seqstart_k_host.data());
|
||||
seqlen_q_dev.ToDevice(seqlen_q_host.data());
|
||||
seqlen_k_dev.ToDevice(seqlen_k_host.data());
|
||||
|
||||
fmha_bwd_args args{};
|
||||
args.dq_acc_ptr = dq_acc_dev.GetDeviceBuffer();
|
||||
args.dq_ptr = dq_dev.GetDeviceBuffer();
|
||||
args.seqstart_q_ptr = seqstart_q.GetDeviceBuffer();
|
||||
args.seqstart_k_ptr = seqstart_k.GetDeviceBuffer();
|
||||
args.seqlen_q_ptr = seqlen_q_dev.GetDeviceBuffer();
|
||||
args.seqlen_k_ptr = seqlen_k_dev.GetDeviceBuffer();
|
||||
args.batch = batch;
|
||||
args.nhead_q = nhead;
|
||||
args.nhead_k = nhead;
|
||||
args.hdim_q = hdim;
|
||||
args.hdim_v = hdim;
|
||||
args.max_seqlen_q = 6;
|
||||
args.max_seqlen_k = 8;
|
||||
args.stride_dq_acc = hdim;
|
||||
args.stride_dq = hdim;
|
||||
args.nhead_stride_dq_acc = hdim * args.max_seqlen_q;
|
||||
args.nhead_stride_dq = hdim * args.max_seqlen_q;
|
||||
args.nhead_stride_q = 0;
|
||||
args.nhead_stride_k = 0;
|
||||
args.nhead_stride_v = 0;
|
||||
args.nhead_stride_o = 0;
|
||||
args.batch_stride_dq_acc = 0;
|
||||
args.batch_stride_dq = 0;
|
||||
args.split_stride_dq_acc = args.max_seqlen_q * args.stride_dq_acc;
|
||||
args.window_size_left = -1;
|
||||
args.window_size_right = 0;
|
||||
args.mask_type = static_cast<ck_tile::index_t>(mask_enum::no_mask);
|
||||
args.p_drop = 0.0f;
|
||||
args.p_undrop = 1.0f;
|
||||
args.drop_seed_offset = std::make_pair(uint64_t{0}, uint64_t{0});
|
||||
|
||||
using TypeConfig = FmhaBwdTypeConfig<TypeParam>;
|
||||
using ConvertTileTraits = ck_tile::TileFmhaBwdConvertQGradTraits<true, true, 2>;
|
||||
using ConvertProblem =
|
||||
ck_tile::BlockFmhaBwdConvertQGradPipelineProblem<typename TypeConfig::AccDataType,
|
||||
typename TypeConfig::QGradDataType,
|
||||
256,
|
||||
64,
|
||||
0,
|
||||
hdim,
|
||||
true,
|
||||
false,
|
||||
ConvertTileTraits>;
|
||||
using ConvertPipeline = ck_tile::BlockFmhaBwdConvertQGrad<ConvertProblem>;
|
||||
using ConvertKernel = ck_tile::FmhaBwdConvertQGradKernel<ConvertPipeline>;
|
||||
|
||||
auto [convert_kargs, convert_grids] =
|
||||
fmha_bwd_convert_dq_create_kargs_and_grids<ConvertKernel>(args);
|
||||
const dim3 convert_blocks = ConvertKernel::BlockSize();
|
||||
constexpr ck_tile::index_t kConvertBlockPerCu = ConvertKernel::kBlockPerCu;
|
||||
auto convert_kernel = ck_tile::make_kernel<kConvertBlockPerCu>(
|
||||
ConvertKernel{}, convert_grids, convert_blocks, 0, convert_kargs);
|
||||
convert_kernel(kStreamConfig);
|
||||
ASSERT_EQ(hipDeviceSynchronize(), hipSuccess);
|
||||
|
||||
using QGradOutT = typename TypeConfigC::QGradDataType;
|
||||
auto dq_result_host_t = CopyDeviceToHost<QGradOutT>(dq_dev, total_rows_q * hdim);
|
||||
auto dq_result_host = ToFloatVector(dq_result_host_t);
|
||||
|
||||
auto [rtol_gpad, atol_gpad] = get_elimit<TypeParam>(hdim, hdim);
|
||||
for(size_t i = 0; i < dq_result_host.size(); ++i)
|
||||
{
|
||||
SCOPED_TRACE(::testing::Message() << "index=" << i);
|
||||
if(std::fabs(expected[i] - dq_sentinel_ref[i]) < 1e-6f)
|
||||
{
|
||||
EXPECT_FLOAT_EQ(dq_result_host[i], dq_sentinel_ref[i]);
|
||||
}
|
||||
else
|
||||
{
|
||||
EXPECT_NEAR(dq_result_host[i], expected[i], static_cast<float>(atol_gpad));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
TYPED_TEST(FmhaBwdKernelPaddingTyped, ConvertQGrad_DeterministicPaddingUsesLogicalLengths)
|
||||
{
|
||||
constexpr ck_tile::index_t batch = 1;
|
||||
constexpr ck_tile::index_t nhead = 1;
|
||||
constexpr ck_tile::index_t hdim = 128;
|
||||
constexpr ck_tile::index_t phys_rows = 8;
|
||||
constexpr ck_tile::index_t logical_rows = 5;
|
||||
constexpr ck_tile::index_t phys_k = 24;
|
||||
constexpr ck_tile::index_t logical_k = 20;
|
||||
constexpr ck_tile::index_t kN0 = 16;
|
||||
constexpr ck_tile::index_t nsplits = (logical_k + kN0 - 1) / kN0;
|
||||
|
||||
const std::vector<int32_t> seqstart_q_host{0, phys_rows};
|
||||
const std::vector<int32_t> seqlen_q_host{logical_rows};
|
||||
const std::vector<int32_t> seqstart_k_host{0, phys_k};
|
||||
const std::vector<int32_t> seqlen_k_host{logical_k};
|
||||
const ck_tile::index_t total_rows_q = seqstart_q_host.back();
|
||||
|
||||
using TypeConfigD = FmhaBwdTypeConfig<TypeParam>;
|
||||
using AccTypeDet = typename TypeConfigD::AccDataType; // float
|
||||
using QGradTypeD = typename TypeConfigD::QGradDataType; // bf16
|
||||
|
||||
ck_tile::HostTensor<AccTypeDet> dq_acc_host({nsplits, 1, nhead, phys_rows, hdim});
|
||||
dq_acc_host.ForEach([&](auto& self, auto idx) {
|
||||
const float s = static_cast<float>(idx[0]);
|
||||
// Use split-dependent constant to avoid per-element variance and rounding interplay
|
||||
self(idx) = ck_tile::type_convert<AccTypeDet>(1.0f + 0.1f * s);
|
||||
});
|
||||
|
||||
const float dq_sentinel_val_det =
|
||||
ck_tile::type_convert<float>(ck_tile::type_convert<QGradTypeD>(SentinelValue()));
|
||||
std::vector<float> expected(total_rows_q * hdim, dq_sentinel_val_det);
|
||||
// Expected is the sum over splits of the constant (1.0 + 0.1*s)
|
||||
for(ck_tile::index_t row = 0; row < logical_rows; ++row)
|
||||
for(ck_tile::index_t c = 0; c < hdim; ++c)
|
||||
{
|
||||
float acc = 0.f;
|
||||
for(ck_tile::index_t s = 0; s < nsplits; ++s)
|
||||
{
|
||||
acc += (1.0f + 0.1f * static_cast<float>(s));
|
||||
}
|
||||
expected[row * hdim + c] = acc;
|
||||
}
|
||||
|
||||
ck_tile::HostTensor<QGradTypeD> dq_init({1, nhead, total_rows_q, hdim});
|
||||
ck_tile::FillConstant<QGradTypeD>{ck_tile::type_convert<QGradTypeD>(SentinelValue())}(dq_init);
|
||||
|
||||
DeviceMem dq_acc_dev(dq_acc_host.get_element_space_size_in_bytes());
|
||||
DeviceMem dq_dev(dq_init.get_element_space_size_in_bytes());
|
||||
DeviceMem seqstart_q(seqstart_q_host.size() * sizeof(int32_t));
|
||||
DeviceMem seqstart_k(seqstart_k_host.size() * sizeof(int32_t));
|
||||
DeviceMem seqlen_q_dev(seqlen_q_host.size() * sizeof(int32_t));
|
||||
DeviceMem seqlen_k_dev(seqlen_k_host.size() * sizeof(int32_t));
|
||||
|
||||
dq_acc_dev.ToDevice(dq_acc_host.data());
|
||||
dq_dev.ToDevice(dq_init.data());
|
||||
seqstart_q.ToDevice(seqstart_q_host.data());
|
||||
seqstart_k.ToDevice(seqstart_k_host.data());
|
||||
seqlen_q_dev.ToDevice(seqlen_q_host.data());
|
||||
seqlen_k_dev.ToDevice(seqlen_k_host.data());
|
||||
|
||||
fmha_bwd_args args{};
|
||||
args.dq_acc_ptr = dq_acc_dev.GetDeviceBuffer();
|
||||
args.dq_ptr = dq_dev.GetDeviceBuffer();
|
||||
args.seqstart_q_ptr = seqstart_q.GetDeviceBuffer();
|
||||
args.seqstart_k_ptr = seqstart_k.GetDeviceBuffer();
|
||||
args.seqlen_q_ptr = seqlen_q_dev.GetDeviceBuffer();
|
||||
args.seqlen_k_ptr = seqlen_k_dev.GetDeviceBuffer();
|
||||
args.batch = batch;
|
||||
args.nhead_q = nhead;
|
||||
args.nhead_k = nhead;
|
||||
args.hdim_q = hdim;
|
||||
args.hdim_v = hdim;
|
||||
args.max_seqlen_q = phys_rows;
|
||||
args.max_seqlen_k = phys_k;
|
||||
args.stride_dq_acc = hdim;
|
||||
args.stride_dq = hdim;
|
||||
args.nhead_stride_dq_acc = phys_rows * hdim;
|
||||
args.nhead_stride_dq = phys_rows * hdim;
|
||||
args.split_stride_dq_acc = phys_rows * hdim;
|
||||
args.window_size_left = -1;
|
||||
args.window_size_right = 0;
|
||||
args.mask_type = static_cast<ck_tile::index_t>(mask_enum::no_mask);
|
||||
args.p_drop = 0.0f;
|
||||
args.p_undrop = 1.0f;
|
||||
args.drop_seed_offset = std::make_pair(uint64_t{0}, uint64_t{0});
|
||||
|
||||
using TypeConfig = FmhaBwdTypeConfig<TypeParam>;
|
||||
using TileTraitsDet = ck_tile::TileFmhaBwdConvertQGradTraits<true, true, 2>;
|
||||
using PipelineProblemDet =
|
||||
ck_tile::BlockFmhaBwdConvertQGradPipelineProblem<typename TypeConfig::AccDataType,
|
||||
typename TypeConfig::QGradDataType,
|
||||
256,
|
||||
64,
|
||||
kN0,
|
||||
hdim,
|
||||
true,
|
||||
true,
|
||||
TileTraitsDet>;
|
||||
using PipelineDet = ck_tile::BlockFmhaBwdConvertQGrad<PipelineProblemDet>;
|
||||
using ConvertKernelDet = ck_tile::FmhaBwdConvertQGradKernel<PipelineDet>;
|
||||
|
||||
auto [convert_kargs, convert_grids] =
|
||||
fmha_bwd_convert_dq_create_kargs_and_grids<ConvertKernelDet>(args);
|
||||
const dim3 convert_blocks = ConvertKernelDet::BlockSize();
|
||||
constexpr ck_tile::index_t kConvertBlockPerCu = ConvertKernelDet::kBlockPerCu;
|
||||
auto convert_kernel = ck_tile::make_kernel<kConvertBlockPerCu>(
|
||||
ConvertKernelDet{}, convert_grids, convert_blocks, 0, convert_kargs);
|
||||
convert_kernel(kStreamConfig);
|
||||
ASSERT_EQ(hipDeviceSynchronize(), hipSuccess);
|
||||
|
||||
using QGradOutTD = typename TypeConfigD::QGradDataType;
|
||||
auto dq_result_host_t = CopyDeviceToHost<QGradOutTD>(dq_dev, total_rows_q * hdim);
|
||||
auto dq_result_host = ToFloatVector(dq_result_host_t);
|
||||
|
||||
const float dq_sentinel_val2 = dq_sentinel_val_det;
|
||||
|
||||
auto [rtol_det, atol_det] = get_elimit<TypeParam>(hdim, hdim);
|
||||
for(size_t i = 0; i < dq_result_host.size(); ++i)
|
||||
{
|
||||
SCOPED_TRACE(::testing::Message() << "index=" << i);
|
||||
if(std::fabs(expected[i] - dq_sentinel_val2) < 1e-6f)
|
||||
{
|
||||
EXPECT_FLOAT_EQ(dq_result_host[i], dq_sentinel_val2);
|
||||
}
|
||||
else
|
||||
{
|
||||
EXPECT_NEAR(dq_result_host[i], expected[i], static_cast<float>(atol_det));
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user