fix page first per layer pf2lf kernel (#8915)

Co-authored-by: Zhiqiang Xie <xiezhq@stanford.edu>
This commit is contained in:
huangtingwei
2025-08-10 08:16:11 +08:00
committed by GitHub
parent 5c31b35db2
commit 86497d99f2
5 changed files with 15 additions and 5 deletions

View File

@@ -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,