[HiCache]Page head layout IO kernel (#11615)

This commit is contained in:
huangtingwei
2025-10-26 15:53:50 +08:00
committed by GitHub
parent 6371f7af27
commit 3e6281d0aa
5 changed files with 574 additions and 31 deletions

View File

@@ -562,6 +562,21 @@ void transfer_kv_per_layer_pf_lf(
int64_t block_quota,
int64_t num_warps_per_block);
void transfer_kv_per_layer_ph_lf(
const at::Tensor src_k,
at::Tensor dst_k,
const at::Tensor src_v,
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 page_size,
int64_t head_num,
int64_t block_quota,
int64_t num_warps_per_block);
void transfer_kv_all_layer(
const at::Tensor src_k_layers,
const at::Tensor dst_k_layers,
@@ -587,6 +602,21 @@ void transfer_kv_all_layer_lf_pf(
int64_t block_quota,
int64_t num_warps_per_block);
void transfer_kv_all_layer_lf_ph(
const at::Tensor src_k_layers,
at::Tensor dst_k,
const at::Tensor src_v_layers,
at::Tensor dst_v,
const at::Tensor src_indices,
const at::Tensor dst_indices,
int64_t item_size,
int64_t dst_layout_dim,
int64_t num_layers,
int64_t page_size,
int64_t head_num,
int64_t block_quota,
int64_t num_warps_per_block);
void transfer_kv_per_layer_mla(
const at::Tensor src,
at::Tensor dst,