ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-17 09:08:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

correct results

Browse Source
This commit is contained in:
Tianxing Wu
2025-11-25 09:27:40 +00:00
parent b3c5cd0c76
commit cc7caf4d7d
3 changed files with 45 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files

40
example/ck_tile/01_unified_attention/example_unified_attention.cpp
View File

@@ -233,14 +233,14 @@ CK_TILE_HOST void fmha_fwd(const ck_tile::HostTensor<QDataType>& q_bshd,
ck_tile::reference_batched_gemm<QDataType, KDataType, AccDataType>(
q_host_ref, k_host_ref, s_host_ref, q_element_op, k_element_op, s_acc_element_op);
// ck_tile::reference_batched_masking(
// s_host_ref,
// ck_tile::make_generic_attention_mask_from_lr_window<FmhaMasks::CausalMask>(
// -1,
// 0,
// seqlen_q,
// seqlen_kv,
// true));
ck_tile::reference_batched_masking(
s_host_ref,
ck_tile::make_generic_attention_mask_from_lr_window<FmhaMasks::CausalMask>(
-1,
0,
seqlen_q,
seqlen_kv,
true));
ck_tile::reference_batched_softmax<AccDataType, AccDataType>(
s_host_ref, p_host_ref, ck_tile::identity{});
ck_tile::reference_batched_gemm<PDataType, VDataType, AccDataType>(
@@ -556,6 +556,30 @@ bool run_impl(const Problem& problem, const RunConfig& run_config)
<< nonzero << " / " << total
<< " (" << percent << "%)\n";
// std::cout << "\n=== Complete Output Tensor (o) ===\n";
// for (int tok = 0; tok < problem.num_tokens; ++tok) {
// std::cout << "Token " << tok << ":\n";
// for (int h = 0; h < problem.nhead_q; ++h) {
// std::cout << " Head " << h << ": ";
// for (int d = 0; d < problem.hdim; ++d) {
// std::cout << static_cast<float>(o(tok, h, d)) << " ";
// }
// std::cout << "\n";
// }
// }
// std::cout << "\n=== Complete Reference Tensor (o_ref) ===\n";
// for (int tok = 0; tok < problem.num_tokens; ++tok) {
// std::cout << "Token " << tok << ":\n";
// for (int h = 0; h < problem.nhead_q; ++h) {
// std::cout << " Head " << h << ": ";
// for (int d = 0; d < problem.hdim; ++d) {
// std::cout << static_cast<float>(o_ref(tok, h, d)) << " ";
// }
// std::cout << "\n";
// }
// }
return ck_tile::check_err(o, o_ref, std::string("found incorrect results!"), rtol, atol);
}

19
include/ck_tile/ops/unified_attention/kernel/unified_attention_kernel.hpp
View File

@@ -310,19 +310,18 @@ struct UnifiedAttentionKernel
const index_t query_pos = amd_wave_read_first_lane(q_block_local_idx * BLOCK_Q);
const index_t seq_len = kargs.seq_lens_ptr[seq_idx];
// const index_t context_len = amd_wave_read_first_lane(seq_len - cur_batch_query_len);
const index_t context_len = amd_wave_read_first_lane(seq_len - cur_batch_query_len);
// index_t _max_seq_prefix_len =
// amd_wave_read_first_lane((context_len + q_block_local_idx * BLOCK_Q + (BLOCK_M - 1)
// + 1));
index_t _max_seq_prefix_len =
amd_wave_read_first_lane((context_len + q_block_local_idx * BLOCK_Q + (BLOCK_M - 1)
+ 1));
// if(seq_len < _max_seq_prefix_len)
// {
// _max_seq_prefix_len = seq_len;
// }
if(seq_len < _max_seq_prefix_len)
{
_max_seq_prefix_len = seq_len;
}
// const auto max_seq_prefix_len = _max_seq_prefix_len;
const auto max_seq_prefix_len = seq_len;
const auto max_seq_prefix_len = _max_seq_prefix_len;
const index_t num_blocks = amd_wave_read_first_lane((max_seq_prefix_len + BLOCK_SIZE - 1) / BLOCK_SIZE);
// TODO sliding window

9
include/ck_tile/ops/unified_attention/pipeline/unified_attention_pipeline.hpp
View File

@@ -897,11 +897,10 @@ struct UnifiedAttentionPipeline
auto fmha_mask = [&](auto sp_reg_idx) {
if constexpr(FmhaMask::IsMasking)
{
bool need_perpixel_check = false;
// mask.IsEdgeTile(q_origin.at(number<0>{}),
// i_total_loops * BLOCK_SIZE,
// number<BLOCK_M>{},
// number<BLOCK_SIZE>{});
bool need_perpixel_check = mask.IsEdgeTile(q_origin.at(number<0>{}),
i_total_loops * BLOCK_SIZE,
number<BLOCK_M>{},
number<BLOCK_SIZE>{});
if(need_perpixel_check)
{
set_tile_if(sp(sp_reg_idx).sp_compute,