ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-13 18:51:13 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix the vector max

Browse Source
This commit is contained in:
Tianxing Wu
2025-10-23 11:42:15 +00:00
parent 3bcef59536
commit e03ed35944
1 changed files with 12 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -406,7 +406,18 @@ bool run_impl(const Problem& problem, const RunConfig& run_config)
args.seq_lens_ptr =reinterpret_cast<const ck_tile::index_t*>(seq_lens_buf.GetDeviceBuffer());
args.query_start_len_ptr =reinterpret_cast<const ck_tile::index_t*>(query_start_len_buf.GetDeviceBuffer());
int max_kv_len = std::max_element(eff_kv_lens.begin(), eff_kv_lens.end());
auto max_element = [&](const std::vector<ck_tile::index_t>& opt_vec) {
ck_tile::index_t max = opt_vec[0];
for (ck_tile::index_t i: opt_vec) {
if (i > max){
max = i;
}
}
return max;
};
ck_tile::index_t max_kv_len = max_element(eff_kv_lens);
ck_tile::index_t max_num_blocks_per_seq = (max_kv_len + problem.BLOCK_SIZE - 1) / problem.BLOCK_SIZE;