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

fix

Browse Source
This commit is contained in:
Juuso Korhonen
2025-11-24 10:20:04 +00:00
parent f552cd7841
commit f2fbc44b7b
3 changed files with 23 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -310,18 +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; // _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,10 +897,11 @@ struct UnifiedAttentionPipeline
auto fmha_mask = [&](auto sp_reg_idx) {
if constexpr(FmhaMask::IsMasking)
{
bool need_perpixel_check = mask.IsEdgeTile(q_origin.at(number<0>{}),
i_total_loops * BLOCK_SIZE,
number<BLOCK_M>{},
number<BLOCK_SIZE>{});
bool need_perpixel_check = false;
// 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,