From cc7caf4d7dfbde0368c5c6c385b462f3aa109b76 Mon Sep 17 00:00:00 2001 From: Tianxing Wu Date: Tue, 25 Nov 2025 09:27:40 +0000 Subject: [PATCH] correct results --- .../example_unified_attention.cpp | 40 +++++++++++++++---- .../kernel/unified_attention_kernel.hpp | 19 +++++---- .../pipeline/unified_attention_pipeline.hpp | 9 ++--- 3 files changed, 45 insertions(+), 23 deletions(-) diff --git a/example/ck_tile/01_unified_attention/example_unified_attention.cpp b/example/ck_tile/01_unified_attention/example_unified_attention.cpp index 765c55a5ce..74f0e3f80f 100644 --- a/example/ck_tile/01_unified_attention/example_unified_attention.cpp +++ b/example/ck_tile/01_unified_attention/example_unified_attention.cpp @@ -233,14 +233,14 @@ CK_TILE_HOST void fmha_fwd(const ck_tile::HostTensor& q_bshd, ck_tile::reference_batched_gemm( 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( - // -1, - // 0, - // seqlen_q, - // seqlen_kv, - // true)); + ck_tile::reference_batched_masking( + s_host_ref, + ck_tile::make_generic_attention_mask_from_lr_window( + -1, + 0, + seqlen_q, + seqlen_kv, + true)); ck_tile::reference_batched_softmax( s_host_ref, p_host_ref, ck_tile::identity{}); ck_tile::reference_batched_gemm( @@ -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(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(o_ref(tok, h, d)) << " "; + // } + // std::cout << "\n"; + // } + // } + return ck_tile::check_err(o, o_ref, std::string("found incorrect results!"), rtol, atol); } diff --git a/include/ck_tile/ops/unified_attention/kernel/unified_attention_kernel.hpp b/include/ck_tile/ops/unified_attention/kernel/unified_attention_kernel.hpp index 90433da0fb..735a8c4252 100644 --- a/include/ck_tile/ops/unified_attention/kernel/unified_attention_kernel.hpp +++ b/include/ck_tile/ops/unified_attention/kernel/unified_attention_kernel.hpp @@ -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 diff --git a/include/ck_tile/ops/unified_attention/pipeline/unified_attention_pipeline.hpp b/include/ck_tile/ops/unified_attention/pipeline/unified_attention_pipeline.hpp index 3fbfb0cd9e..d661c342f9 100644 --- a/include/ck_tile/ops/unified_attention/pipeline/unified_attention_pipeline.hpp +++ b/include/ck_tile/ops/unified_attention/pipeline/unified_attention_pipeline.hpp @@ -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{}, - // number{}); + bool need_perpixel_check = mask.IsEdgeTile(q_origin.at(number<0>{}), + i_total_loops * BLOCK_SIZE, + number{}, + number{}); if(need_perpixel_check) { set_tile_if(sp(sp_reg_idx).sp_compute,