fix page first per layer pf2lf kernel (#8915)
Co-authored-by: Zhiqiang Xie <xiezhq@stanford.edu>
This commit is contained in:
@@ -210,6 +210,7 @@ void transfer_kv_per_layer_pf_lf(
|
||||
at::Tensor dst_v,
|
||||
const at::Tensor src_indices,
|
||||
const at::Tensor dst_indices,
|
||||
int64_t layer_id,
|
||||
int64_t item_size,
|
||||
int64_t src_layout_dim,
|
||||
int64_t block_quota,
|
||||
@@ -222,7 +223,7 @@ void transfer_kv_per_layer_pf_lf(
|
||||
dst_v,
|
||||
src_indices,
|
||||
dst_indices,
|
||||
0,
|
||||
layer_id,
|
||||
1,
|
||||
item_size,
|
||||
src_layout_dim,
|
||||
@@ -336,6 +337,7 @@ void transfer_kv_per_layer_mla_pf_lf(
|
||||
at::Tensor dst,
|
||||
const at::Tensor src_indices,
|
||||
const at::Tensor dst_indices,
|
||||
int64_t layer_id,
|
||||
int64_t item_size,
|
||||
int64_t src_layout_dim,
|
||||
int64_t block_quota,
|
||||
@@ -348,7 +350,7 @@ void transfer_kv_per_layer_mla_pf_lf(
|
||||
empty,
|
||||
src_indices,
|
||||
dst_indices,
|
||||
0,
|
||||
layer_id,
|
||||
1,
|
||||
item_size,
|
||||
src_layout_dim,
|
||||
|
||||
Reference in New Issue
Block a user