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

Force batch mode when invoking appendkv & splitkv apis

Browse Source
This commit is contained in:
PoYen, Chen
2024-08-18 06:05:42 +00:00
parent cc52587bcc
commit 05157bf3a3
1 changed files with 7 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
example/ck_tile/01_fmha/fmha_fwd.cpp
View File

@@ -257,9 +257,9 @@ int override_num_splits_if_necessary(
template <typename DataType>
bool run(const ck_tile::ArgParser& arg_parser)
{
std::string data_type = arg_parser.get_str("prec");
int do_validation = arg_parser.get_int("v");
std::string data_type = arg_parser.get_str("prec");
int do_validation = arg_parser.get_int("v");
auto mode = static_cast<mode_enum>(arg_parser.get_uint32("mode"));
ck_tile::index_t batch = arg_parser.get_int("b");
ck_tile::index_t nhead = arg_parser.get_int("h");
ck_tile::index_t nhead_k = arg_parser.get_int("h_k");
@@ -323,11 +323,11 @@ bool run(const ck_tile::ArgParser& arg_parser)
<< std::endl;
use_cache_batch_idx = false;
}
auto mode = static_cast<mode_enum>(arg_parser.get_uint32("mode"));
if((use_cache_batch_idx || 0 < page_block_size) && mode != mode_enum::batch)
// the input layout we use for kvcache is same as batch mode
if((0 < seqlen_knew || 0 < rotary_dim || use_cache_batch_idx || 0 < page_block_size) &&
mode != mode_enum::batch)
{
std::cerr << "both kvcache & split-kv enabled. ignoring the 'mode' option" << std::endl;
std::cerr << "kvcache enabled. ignoring the 'mode' option" << std::endl;
mode = mode_enum::batch;
}