fix page first per layer pf2lf kernel (#8915)

Co-authored-by: Zhiqiang Xie <xiezhq@stanford.edu>
2025-08-10 08:16:11 +08:00
parent 5c31b35db2
commit 86497d99f2
5 changed files with 15 additions and 5 deletions
--- a/sgl-kernel/include/sgl_kernel_ops.h
+++ b/sgl-kernel/include/sgl_kernel_ops.h
@@ -419,6 +419,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,
@@ -463,6 +464,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,