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

fix to example

Browse Source
This commit is contained in:
Juuso Korhonen
2025-11-17 07:33:10 +00:00
parent db7224e067
commit 4a13749f7f
1 changed files with 3 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -30,9 +30,9 @@ auto parse_cmd_args(int argc, char* argv[]) -> std::pair<bool, ck_tile::ArgParse
ck_tile::ArgParser arg_parser;
arg_parser.insert("prec", "fp16", "data type. fp16/bf16")
.insert("b", "3", "batch size")
.insert("h", "8", "num of head, for q")
.insert("h", "32", "num of head, for q")
.insert("h_k",
"-1",
"8",
"num of head, for k/v, -1 means equal to h\n"
"if not equal to h, then this is GQA/MQA case")
.insert("s", "1024", "max_seqlen_q")
@@ -337,7 +337,7 @@ bool run_impl(const Problem& problem, const RunConfig& run_config)
// args.kv_lens = problem.kv_lens
args.q_ptr = q_buf.GetDeviceBuffer();
args.query_stride_0 = problem.hdim * problem.nhead_q;
args.query_stride_0 = problem.hdim;
args.query_stride_1 = problem.hdim;
args.k_ptr = k_buf.GetDeviceBuffer();