chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

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

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files

86
sgl-kernel/python/sgl_kernel/kvcacheio.py
View File

@@ -21,7 +21,7 @@ def transfer_kv_per_layer(
block_quota: int = 2,
num_warps_per_block: int = 16 if _is_hip else 32,
):
torch.ops.sgl_kernel.transfer_kv_per_layer(
torch.ops.sgl_kernel.transfer_kv_per_layer.default(
src_k,
dst_k,
src_v,
@@ -47,7 +47,7 @@ def transfer_kv_per_layer_pf_lf(
block_quota: int = 2,
num_warps_per_block: int = 16 if _is_hip else 32,
):
torch.ops.sgl_kernel.transfer_kv_per_layer_pf_lf(
torch.ops.sgl_kernel.transfer_kv_per_layer_pf_lf.default(
src_k,
dst_k,
src_v,
@@ -62,6 +62,38 @@ def transfer_kv_per_layer_pf_lf(
)
def transfer_kv_per_layer_ph_lf(
src_k: torch.Tensor,
dst_k: torch.Tensor,
src_v: torch.Tensor,
dst_v: torch.Tensor,
src_indices: torch.Tensor,
dst_indices: torch.Tensor,
layer_id: int,
item_size: int,
src_layout_dim: int,
page_size: int,
head_num: int,
block_quota: int = 2,
num_warps_per_block: int = 16 if _is_hip else 32,
):
torch.ops.sgl_kernel.transfer_kv_per_layer_ph_lf.default(
src_k,
dst_k,
src_v,
dst_v,
src_indices,
dst_indices,
layer_id,
item_size,
src_layout_dim,
page_size,
head_num,
block_quota,
num_warps_per_block,
)
def transfer_kv_all_layer(
src_k_layers: torch.Tensor,
dst_k_layers: torch.Tensor,
@@ -74,7 +106,7 @@ def transfer_kv_all_layer(
block_quota: int = 2,
num_warps_per_block: int = 16 if _is_hip else 32,
):
torch.ops.sgl_kernel.transfer_kv_all_layer(
torch.ops.sgl_kernel.transfer_kv_all_layer.default(
src_k_layers,
dst_k_layers,
src_v_layers,
@@ -101,7 +133,7 @@ def transfer_kv_all_layer_lf_pf(
block_quota: int = 2,
num_warps_per_block: int = 16 if _is_hip else 32,
):
torch.ops.sgl_kernel.transfer_kv_all_layer_lf_pf(
torch.ops.sgl_kernel.transfer_kv_all_layer_lf_pf.default(
src_k_layers,
dst_k,
src_v_layers,
@@ -116,6 +148,38 @@ def transfer_kv_all_layer_lf_pf(
)
def transfer_kv_all_layer_lf_ph(
src_k_layers: torch.Tensor,
dst_k: torch.Tensor,
src_v_layers: torch.Tensor,
dst_v: torch.Tensor,
src_indices: torch.Tensor,
dst_indices: torch.Tensor,
item_size: int,
dst_layout_dim: int,
num_layers: int,
page_size: int,
head_num: int,
block_quota: int = 2,
num_warps_per_block: int = 16 if _is_hip else 32,
):
torch.ops.sgl_kernel.transfer_kv_all_layer_lf_ph.default(
src_k_layers,
dst_k,
src_v_layers,
dst_v,
src_indices,
dst_indices,
item_size,
dst_layout_dim,
num_layers,
page_size,
head_num,
block_quota,
num_warps_per_block,
)
def transfer_kv_direct(
src_layers: List[torch.Tensor],
dst_layers: List[torch.Tensor],
@@ -123,7 +187,7 @@ def transfer_kv_direct(
dst_indices: torch.Tensor,
page_size: int,
):
torch.ops.sgl_kernel.transfer_kv_direct(
torch.ops.sgl_kernel.transfer_kv_direct.default(
src_layers, dst_layers, src_indices, dst_indices, page_size
)
@@ -136,7 +200,7 @@ def transfer_kv_per_layer_direct_pf_lf(
layer_id: int,
page_size: int,
):
torch.ops.sgl_kernel.transfer_kv_per_layer_direct_pf_lf(
torch.ops.sgl_kernel.transfer_kv_per_layer_direct_pf_lf.default(
src_ptrs, dst_ptrs, src_indices, dst_indices, layer_id, page_size
)
@@ -148,7 +212,7 @@ def transfer_kv_all_layer_direct_lf_pf(
dst_indices: torch.Tensor,
page_size: int,
):
torch.ops.sgl_kernel.transfer_kv_all_layer_direct_lf_pf(
torch.ops.sgl_kernel.transfer_kv_all_layer_direct_lf_pf.default(
src_ptrs, dst_ptrs, src_indices, dst_indices, page_size
)
@@ -162,7 +226,7 @@ def transfer_kv_per_layer_mla(
block_quota: int = 2,
num_warps_per_block: int = 16 if _is_hip else 32,
):
torch.ops.sgl_kernel.transfer_kv_per_layer_mla(
torch.ops.sgl_kernel.transfer_kv_per_layer_mla.default(
src,
dst,
src_indices,
@@ -184,7 +248,7 @@ def transfer_kv_per_layer_mla_pf_lf(
block_quota: int = 2,
num_warps_per_block: int = 16 if _is_hip else 32,
):
torch.ops.sgl_kernel.transfer_kv_per_layer_mla_pf_lf(
torch.ops.sgl_kernel.transfer_kv_per_layer_mla_pf_lf.default(
src,
dst,
src_indices,
@@ -207,7 +271,7 @@ def transfer_kv_all_layer_mla(
block_quota: int = 2,
num_warps_per_block: int = 16 if _is_hip else 32,
):
torch.ops.sgl_kernel.transfer_kv_all_layer_mla(
torch.ops.sgl_kernel.transfer_kv_all_layer_mla.default(
src_layers,
dst_layers,
src_indices,
@@ -230,7 +294,7 @@ def transfer_kv_all_layer_mla_lf_pf(
block_quota: int = 2,
num_warps_per_block: int = 16 if _is_hip else 32,
):
torch.ops.sgl_kernel.transfer_kv_all_layer_mla_lf_pf(
torch.ops.sgl_kernel.transfer_kv_all_layer_mla_lf_pf.default(
src_layers,
dst,
src_indices,