From eeffd2717ae5c4ed1c3be793270e2d873ea17374 Mon Sep 17 00:00:00 2001 From: Jeff Huang Date: Sat, 18 Oct 2025 22:51:04 +0800 Subject: [PATCH] Adapt fmha_bwd_runner.cpp to new q, kv sequence padding Add backward q/kv sequence padding unit tests. --- example/ck_tile/01_fmha/fmha_bwd_runner.hpp | 71 +- test/ck_tile/fmha/CMakeLists.txt | 4 - test/ck_tile/fmha/test_fmha_bwd.cpp | 738 ++++++++++++++++++ .../fmha/test_fmha_bwd_kernel_padding.cpp | 715 ----------------- 4 files changed, 793 insertions(+), 735 deletions(-) delete mode 100644 test/ck_tile/fmha/test_fmha_bwd_kernel_padding.cpp diff --git a/example/ck_tile/01_fmha/fmha_bwd_runner.hpp b/example/ck_tile/01_fmha/fmha_bwd_runner.hpp index 365656a629..0f34ef851f 100644 --- a/example/ck_tile/01_fmha/fmha_bwd_runner.hpp +++ b/example/ck_tile/01_fmha/fmha_bwd_runner.hpp @@ -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( 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(do_host_ref(i0, i1, o)) * - ck_tile::type_convert(o_host_refs[wb](i0, i1, o)); + do_dot_o += + ck_tile::type_convert(do_host_ref(i0, i1, o)) * + ck_tile::type_convert(o_host_refs[ref_idx](i0, i1, o)); } - ds_hp_host_ref(i0, i1, i2) = ck_tile::type_convert( - 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(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( 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( 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( 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; diff --git a/test/ck_tile/fmha/CMakeLists.txt b/test/ck_tile/fmha/CMakeLists.txt index bbd9fc3d14..e769a79c08 100644 --- a/test/ck_tile/fmha/CMakeLists.txt +++ b/test/ck_tile/fmha/CMakeLists.txt @@ -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) diff --git a/test/ck_tile/fmha/test_fmha_bwd.cpp b/test/ck_tile/fmha/test_fmha_bwd.cpp index 710069febe..3eea02f888 100644 --- a/test/ck_tile/fmha/test_fmha_bwd.cpp +++ b/test/ck_tile/fmha/test_fmha_bwd.cpp @@ -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 +{ +}; + +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 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 seqlen_qs(batch, seqlen_q); + std::vector 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 seqlen_qpads(batch, seqlen_qpad); + std::vector seqlen_kpads(batch, seqlen_k); // No K padding + + auto result = fmha_bwd_run( + 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(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 +{ +}; + +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 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 seqlen_qs(batch, seqlen_q); + std::vector seqlen_ks(batch, seqlen_k); + + // No Q padding + std::vector 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 seqlen_kpads(batch, seqlen_kpad); + + auto result = fmha_bwd_run( + 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(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 +{ +}; + +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 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 seqlen_qs(batch, seqlen_q); + std::vector 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 seqlen_qpads(batch, seqlen_qpad); + std::vector seqlen_kpads(batch, seqlen_kpad); + + auto result = fmha_bwd_run( + 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(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 +{ +}; + +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 seqlen_qs; + std::vector seqlen_ks; + std::vector seqlen_qpads; + std::vector 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( + 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(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 +{ +}; + +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 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 seqlen_qs(batch, seqlen_q); + std::vector 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 seqlen_qpads(batch, calc_pad(seqlen_q)); + std::vector seqlen_kpads(batch, calc_pad(seqlen_k)); + + auto result = fmha_bwd_run( + 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(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 +{ +}; + +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 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 seqlen_qs(batch, seqlen_q); + std::vector 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 seqlen_qpads(batch, seqlen_qpad); + std::vector seqlen_kpads(batch, seqlen_kpad); + + auto result = fmha_bwd_run( + 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(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 +{ +}; + +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 seqlen_qs; + std::vector seqlen_ks; + std::vector seqlen_qpads; + std::vector 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(64, q_len); + k_len = std::max(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( + 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(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); +} diff --git a/test/ck_tile/fmha/test_fmha_bwd_kernel_padding.cpp b/test/ck_tile/fmha/test_fmha_bwd_kernel_padding.cpp deleted file mode 100644 index 6d0fe6e397..0000000000 --- a/test/ck_tile/fmha/test_fmha_bwd_kernel_padding.cpp +++ /dev/null @@ -1,715 +0,0 @@ -// SPDX-License-Identifier: MIT -// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. - -#include -#include -#include -#include -#include -#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 -std::vector MakeVectorFromFunction(size_t count, std::function fn) -{ - std::vector data(count); - for(size_t i = 0; i < count; ++i) - { - data[i] = static_cast(fn(i)); - } - return data; -} - -template -std::vector ToFloatVector(const std::vector& src) -{ - std::vector dst(src.size()); - for(size_t i = 0; i < src.size(); ++i) - { - dst[i] = ck_tile::type_convert(src[i]); - } - return dst; -} - -template -std::vector CopyDeviceToHost(const ck_tile::DeviceMem& dev, size_t element_count) -{ - std::vector 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 -class FmhaBwdKernelPaddingTyped : public ::testing::Test -{ -}; - -using KernelPaddingTypes = ::testing::Types; -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 seqstart_q_host{0, phys_rows0, phys_rows0 + phys_rows1}; - const std::vector seqlen_q_host{5, 3}; - - const ck_tile::index_t total_rows = seqstart_q_host.back(); - - // Types per config - using TypeConfig = FmhaBwdTypeConfig; - 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 o_host({1, nhead, total_rows, hdim}); - ck_tile::HostTensor do_host({1, nhead, total_rows, hdim}); - ck_tile::HostTensor 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{ck_tile::type_convert(o_const)}(o_host); - ck_tile::FillConstant{ck_tile::type_convert(do_const)}(do_host); - ck_tile::FillConstant{ck_tile::type_convert(SentinelValue())}(d_init_host); - - // Prepare expected D via runner-style CPU reference, sentinel elsewhere - std::vector expected(static_cast(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(o_host(0, 0, start + row, c)); - const auto do_val = ck_tile::type_convert(do_host(0, 0, start + row, c)); - acc += do_val * o_val; - } - expected[start + row] = ck_tile::type_convert(acc); - } - } - std::vector sentinel_ref(static_cast(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(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; - using DotProblem = - ck_tile::BlockFmhaBwdOGradDotOPipelineProblem; - using DotPipeline = ck_tile::BlockFmhaBwdOGradDotO; - using DotKernel = ck_tile::FmhaBwdOGradDotOKernel; - - auto [dot_kargs, dot_grids] = fmha_bwd_dot_do_o_create_kargs_and_grids(args); - const dim3 dot_blocks = DotKernel::BlockSize(); - constexpr ck_tile::index_t kDotBlockPerCu = DotKernel::kBlockPerCu; - auto dot_kernel = - ck_tile::make_kernel(DotKernel{}, dot_grids, dot_blocks, 0, dot_kargs); - dot_kernel(kStreamConfig); - ASSERT_EQ(hipDeviceSynchronize(), hipSuccess); - - auto d_result_host = CopyDeviceToHost(d_dev, total_rows); - - auto [rtol_doto, atol_doto] = get_elimit(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(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 seqstart_q_host{ - 0, phys_r0, phys_r0 + phys_r1, phys_r0 + phys_r1 + phys_r2}; - const std::vector seqlen_q_host{3, 0, 4}; - const ck_tile::index_t total_rows = seqstart_q_host.back(); - - using TypeConfig = FmhaBwdTypeConfig; - using OType = typename TypeConfig::ODataType; - using DOType = typename TypeConfig::OGradDataType; - using DType = typename TypeConfig::DDataType; - - ck_tile::HostTensor o_host({1, nhead, total_rows, hdim}); - ck_tile::HostTensor do_host({1, nhead, total_rows, hdim}); - ck_tile::HostTensor d_init_host({1, nhead, total_rows}); - ck_tile::FillConstant{ck_tile::type_convert(1.0f)}(o_host); - ck_tile::FillConstant{ck_tile::type_convert(2.0f)}(do_host); - ck_tile::FillConstant{ck_tile::type_convert(SentinelValue())}(d_init_host); - - std::vector expected(static_cast(total_rows), SentinelValue()); - const float dot = 2.0f * 1.0f * static_cast(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 sentinel_ref(static_cast(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; - using DotProblem = - ck_tile::BlockFmhaBwdOGradDotOPipelineProblem; - using DotPipeline = ck_tile::BlockFmhaBwdOGradDotO; - using DotKernel = ck_tile::FmhaBwdOGradDotOKernel; - - auto [dot_kargs, dot_grids] = fmha_bwd_dot_do_o_create_kargs_and_grids(args); - const dim3 dot_blocks = DotKernel::BlockSize(); - constexpr ck_tile::index_t kDotBlockPerCu = DotKernel::kBlockPerCu; - auto dot_kernel = - ck_tile::make_kernel(DotKernel{}, dot_grids, dot_blocks, 0, dot_kargs); - dot_kernel(kStreamConfig); - ASSERT_EQ(hipDeviceSynchronize(), hipSuccess); - - auto d_result_host = CopyDeviceToHost(d_dev, total_rows); - auto [rtol, atol] = get_elimit(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(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 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; - using OType = typename TypeConfig::ODataType; - using DOType = typename TypeConfig::OGradDataType; - using DType = typename TypeConfig::DDataType; - - ck_tile::HostTensor o_host({1, nhead, total_rows, hdim}); - ck_tile::HostTensor do_host({1, nhead, total_rows, hdim}); - ck_tile::HostTensor d_init_host({1, nhead, total_rows}); - ck_tile::FillConstant{ck_tile::type_convert(1.0f)}(o_host); - ck_tile::FillConstant{ck_tile::type_convert(2.0f)}(do_host); - ck_tile::FillConstant{ck_tile::type_convert(SentinelValue())}(d_init_host); - - std::vector expected(static_cast(total_rows), SentinelValue()); - const float dot = 2.0f * 1.0f * static_cast(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 sentinel_ref(static_cast(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; - using DotProblem = - ck_tile::BlockFmhaBwdOGradDotOPipelineProblem; - using DotPipeline = ck_tile::BlockFmhaBwdOGradDotO; - using DotKernel = ck_tile::FmhaBwdOGradDotOKernel; - - auto [dot_kargs, dot_grids] = fmha_bwd_dot_do_o_create_kargs_and_grids(args); - const dim3 dot_blocks = DotKernel::BlockSize(); - constexpr ck_tile::index_t kDotBlockPerCu = DotKernel::kBlockPerCu; - auto dot_kernel = - ck_tile::make_kernel(DotKernel{}, dot_grids, dot_blocks, 0, dot_kargs); - dot_kernel(kStreamConfig); - ASSERT_EQ(hipDeviceSynchronize(), hipSuccess); - - auto d_result_host = CopyDeviceToHost(d_dev, total_rows); - auto [rtol, atol] = get_elimit(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(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 seqstart_q_host{0, 6, 6, 10}; // physical lengths: 6,0,4 - const std::vector seqlen_q_host{4, 0, 3}; - const std::vector seqstart_k_host{0, 7, 15, 18}; - const std::vector seqlen_k_host{5, 8, 3}; - - const ck_tile::index_t total_rows_q = seqstart_q_host.back(); - - using TypeConfigC = FmhaBwdTypeConfig; - using AccType = typename TypeConfigC::AccDataType; // float - using QGradType = typename TypeConfigC::QGradDataType; // bf16 - - ck_tile::HostTensor dq_acc_host({1, nhead, total_rows_q, hdim}); - ck_tile::HostTensor dq_host_init({1, nhead, total_rows_q, hdim}); - - const float dq_acc_const = 1.25f; - ck_tile::FillConstant{ck_tile::type_convert(dq_acc_const)}(dq_acc_host); - ck_tile::FillConstant{ck_tile::type_convert(SentinelValue())}( - dq_host_init); - - const float dq_sentinel_val = - ck_tile::type_convert(ck_tile::type_convert(SentinelValue())); - std::vector dq_sentinel_ref(static_cast(total_rows_q * hdim), dq_sentinel_val); - std::vector 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(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(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; - using ConvertTileTraits = ck_tile::TileFmhaBwdConvertQGradTraits; - using ConvertProblem = - ck_tile::BlockFmhaBwdConvertQGradPipelineProblem; - using ConvertPipeline = ck_tile::BlockFmhaBwdConvertQGrad; - using ConvertKernel = ck_tile::FmhaBwdConvertQGradKernel; - - auto [convert_kargs, convert_grids] = - fmha_bwd_convert_dq_create_kargs_and_grids(args); - const dim3 convert_blocks = ConvertKernel::BlockSize(); - constexpr ck_tile::index_t kConvertBlockPerCu = ConvertKernel::kBlockPerCu; - auto convert_kernel = ck_tile::make_kernel( - 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(dq_dev, total_rows_q * hdim); - auto dq_result_host = ToFloatVector(dq_result_host_t); - - auto [rtol_gpad, atol_gpad] = get_elimit(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(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 seqstart_q_host{0, phys_rows}; - const std::vector seqlen_q_host{logical_rows}; - const std::vector seqstart_k_host{0, phys_k}; - const std::vector seqlen_k_host{logical_k}; - const ck_tile::index_t total_rows_q = seqstart_q_host.back(); - - using TypeConfigD = FmhaBwdTypeConfig; - using AccTypeDet = typename TypeConfigD::AccDataType; // float - using QGradTypeD = typename TypeConfigD::QGradDataType; // bf16 - - ck_tile::HostTensor dq_acc_host({nsplits, 1, nhead, phys_rows, hdim}); - dq_acc_host.ForEach([&](auto& self, auto idx) { - const float s = static_cast(idx[0]); - // Use split-dependent constant to avoid per-element variance and rounding interplay - self(idx) = ck_tile::type_convert(1.0f + 0.1f * s); - }); - - const float dq_sentinel_val_det = - ck_tile::type_convert(ck_tile::type_convert(SentinelValue())); - std::vector 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(s)); - } - expected[row * hdim + c] = acc; - } - - ck_tile::HostTensor dq_init({1, nhead, total_rows_q, hdim}); - ck_tile::FillConstant{ck_tile::type_convert(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(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; - using TileTraitsDet = ck_tile::TileFmhaBwdConvertQGradTraits; - using PipelineProblemDet = - ck_tile::BlockFmhaBwdConvertQGradPipelineProblem; - using PipelineDet = ck_tile::BlockFmhaBwdConvertQGrad; - using ConvertKernelDet = ck_tile::FmhaBwdConvertQGradKernel; - - auto [convert_kargs, convert_grids] = - fmha_bwd_convert_dq_create_kargs_and_grids(args); - const dim3 convert_blocks = ConvertKernelDet::BlockSize(); - constexpr ck_tile::index_t kConvertBlockPerCu = ConvertKernelDet::kBlockPerCu; - auto convert_kernel = ck_tile::make_kernel( - 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(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(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(atol_det)); - } - } -}