ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-14 02:57:45 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix

Browse Source
This commit is contained in:
Juuso Korhonen
2025-12-05 10:52:15 +00:00
parent f01e964e22
commit 0501d37efe
2 changed files with 53 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files

69
example/ck_tile/01_unified_attention/unified_attention_impl.hpp
View File

@@ -177,38 +177,9 @@ float unified_attention_kernel_launch(const unified_attention_args& args,
assert(BLOCK_Q == BLOCK_M / args.num_queries_per_kv &&
"BLOCK_Q must equal BLOCK_M / num_queries_per_kv");
index_t total_num_q_blocks = args.num_tokens / BLOCK_Q + args.num_seqs;
auto kargs =
auto kargs_lambda =
[&]() {
if(Kernel::QuantEnum == UnifiedAttentionQuantScaleEnum::NO_SCALE)
{
return Kernel::MakeKargs(args.q_ptr,
args.k_ptr,
args.v_ptr,
args.o_ptr,
args.num_blks,
args.num_head_q,
args.num_queries_per_kv,
args.scale_s,
total_num_q_blocks,
args.query_stride_0,
args.query_stride_1,
args.stride_k_cache_0,
args.stride_k_cache_1,
args.stride_k_cache_2,
args.stride_k_cache_3,
args.stride_v_cache_0,
args.stride_v_cache_1,
args.stride_v_cache_2,
args.stride_v_cache_3,
args.output_stride_0,
args.output_stride_1,
args.block_tables_ptr,
args.block_table_stride,
args.seq_lens_ptr,
args.query_start_len_ptr,
args.num_seqs);
}
else
if constexpr(Kernel::kIsQuantized)
{
return Kernel::MakeKargs(args.q_ptr,
args.k_ptr,
@@ -241,8 +212,40 @@ float unified_attention_kernel_launch(const unified_attention_args& args,
args.query_start_len_ptr,
args.num_seqs);
}
} dim3 grids =
Kernel::GridSize2D(args.num_head_q / args.num_queries_per_kv, total_num_q_blocks);
else
{
return Kernel::MakeKargs(args.q_ptr,
args.k_ptr,
args.v_ptr,
args.o_ptr,
args.num_blks,
args.num_head_q,
args.num_queries_per_kv,
args.scale_s,
total_num_q_blocks,
args.query_stride_0,
args.query_stride_1,
args.stride_k_cache_0,
args.stride_k_cache_1,
args.stride_k_cache_2,
args.stride_k_cache_3,
args.stride_v_cache_0,
args.stride_v_cache_1,
args.stride_v_cache_2,
args.stride_v_cache_3,
args.output_stride_0,
args.output_stride_1,
args.block_tables_ptr,
args.block_table_stride,
args.seq_lens_ptr,
args.query_start_len_ptr,
args.num_seqs);
}
};
auto kargs = kargs_lambda();
dim3 grids = Kernel::GridSize2D(args.num_head_q / args.num_queries_per_kv, total_num_q_blocks);
constexpr dim3 blocks = Kernel::BlockSize();
constexpr index_t kBlockPerCu = Kernel::kBlockPerCu;