Add block_table kernel args for appendkv kernel

This commit is contained in:
PoYen, Chen
2024-08-07 04:27:15 +00:00
parent 15d0034a64
commit 443a528adc
2 changed files with 20 additions and 2 deletions

View File

@@ -502,6 +502,9 @@ auto fmha_fwd_appendkv_create_kargs_and_grids(fmha_fwd_appendkv_args args)
args.rotary_cos_ptr,
args.rotary_sin_ptr,
args.rotary_dim,
args.block_table_ptr,
args.batch_stride_block_table,
args.page_block_size,
args.stride_q,
args.stride_k,
args.stride_knew,
@@ -532,6 +535,9 @@ auto fmha_fwd_appendkv_create_kargs_and_grids(fmha_fwd_appendkv_args args)
args.rotary_cos_ptr,
args.rotary_sin_ptr,
args.rotary_dim,
args.block_table_ptr,
args.batch_stride_block_table,
args.page_block_size,
args.stride_q,
args.stride_k,
args.stride_knew,