Page first direct IO kernel (#10060)

Co-authored-by: Zhiqiang Xie <xiezhq@stanford.edu>
This commit is contained in:
huangtingwei
2025-09-10 13:35:34 +08:00
committed by GitHub
parent 737d73ed5b
commit 5be8c2f7f7
5 changed files with 358 additions and 2 deletions

View File

@@ -128,6 +128,31 @@ def transfer_kv_direct(
)
def transfer_kv_per_layer_direct_pf_lf(
src_ptrs: List[torch.Tensor],
dst_ptrs: List[torch.Tensor],
src_indices: torch.Tensor,
dst_indices: torch.Tensor,
layer_id: int,
page_size: int,
):
torch.ops.sgl_kernel.transfer_kv_per_layer_direct_pf_lf(
src_ptrs, dst_ptrs, src_indices, dst_indices, layer_id, page_size
)
def transfer_kv_all_layer_direct_lf_pf(
src_ptrs: List[torch.Tensor],
dst_ptrs: List[torch.Tensor],
src_indices: torch.Tensor,
dst_indices: torch.Tensor,
page_size: int,
):
torch.ops.sgl_kernel.transfer_kv_all_layer_direct_lf_pf(
src_ptrs, dst_ptrs, src_indices, dst_indices, page_size
)
def transfer_kv_per_layer_mla(
src: torch.Tensor,
dst: torch.Tensor,