Implementation:Sgl project Sglang KVCache IO

Knowledge Sources	Sgl_project_Sglang
Domains	GPU Kernels, KV Cache, Memory Management
Last Updated	2026-02-10 00:00 GMT

Overview

Python interface for high-performance KV cache transfer operations between memory locations, supporting multiple memory layouts and both standard multi-head and MLA (Multi-head Latent Attention) architectures.

Description

kvcacheio.py provides a comprehensive set of CUDA kernel wrappers for transferring KV cache data between different memory locations on the GPU. These operations are essential for prefix caching, radix tree eviction, disaggregated serving, and KV cache migration between memory pools.

The module is organized into several categories based on transfer scope and memory layout:

Per-Layer Transfers (transfer data for a single transformer layer):

• transfer_kv_per_layer -- Basic per-layer KV transfer between source and destination tensors using index-based addressing. Copies key and value data at positions specified by src_indices and dst_indices.
• transfer_kv_per_layer_pf_lf -- Per-layer transfer with page-first, layer-first source layout. Uses a layer_id to index into the layer dimension and src_layout_dim to describe the source memory layout.
• transfer_kv_per_layer_ph_lf -- Per-layer transfer with page-head, layer-first source layout. Additionally requires page_size and head_num parameters for the paged layout.

All-Layer Transfers (transfer data for all layers simultaneously):

• transfer_kv_all_layer -- Transfers KV data across all layers at once, taking lists of layer tensors and a num_layers count.
• transfer_kv_all_layer_lf_pf -- All-layer transfer from layer-first source to page-first destination layout with dst_layout_dim.
• transfer_kv_all_layer_lf_ph -- All-layer transfer from layer-first source to page-head destination, with page_size and head_num for the destination.

MLA-Specific Transfers (for Multi-head Latent Attention architectures like DeepSeek):

• transfer_kv_per_layer_mla -- MLA per-layer transfer operating on a single combined KV tensor (since MLA compresses K and V into a single latent representation).
• transfer_kv_per_layer_mla_pf_lf -- MLA per-layer with page-first, layer-first layout.
• transfer_kv_all_layer_mla -- MLA all-layer transfer.
• transfer_kv_all_layer_mla_lf_pf -- MLA all-layer from layer-first to page-first.

Direct Transfers (using lists of per-layer tensor pointers):

• transfer_kv_direct -- Direct KV transfer using Python lists of per-layer tensors, with page-size-based addressing.
• transfer_kv_per_layer_direct_pf_lf -- Direct per-layer transfer with page-first to layer-first conversion.
• transfer_kv_all_layer_direct_lf_pf -- Direct all-layer transfer from layer-first to page-first.

All functions share common tuning parameters:

• block_quota -- Number of CUDA blocks allocated per transfer operation (default: 2)
• num_warps_per_block -- Warps per CUDA block, with ROCm-aware defaults (16 for HIP vs 32 for CUDA)
• item_size -- Size of each KV cache element in bytes

The module detects ROCm at import time via torch.version.hip and adjusts the default warp count accordingly.

Usage

Use these functions when KV cache data needs to be relocated -- for example, during prefix cache eviction, radix tree reorganization, disaggregated prefill-decode separation, or when migrating KV cache between different memory pools. Choose the appropriate variant based on the source and destination memory layouts.

Code Reference

Source Location

• Repository: Sgl_project_Sglang
• File: sgl-kernel/python/sgl_kernel/kvcacheio.py
• Lines: 1-307

Signature

def transfer_kv_per_layer(
    src_k: torch.Tensor, dst_k: torch.Tensor,
    src_v: torch.Tensor, dst_v: torch.Tensor,
    src_indices: torch.Tensor, dst_indices: torch.Tensor,
    item_size: int, block_quota: int = 2,
    num_warps_per_block: int = 16 if _is_hip else 32,
):

def transfer_kv_per_layer_pf_lf(
    src_k, dst_k, src_v, dst_v, src_indices, dst_indices,
    layer_id: int, item_size: int, src_layout_dim: int,
    block_quota: int = 2, num_warps_per_block: int = ...,
):

def transfer_kv_per_layer_ph_lf(
    src_k, dst_k, src_v, dst_v, src_indices, dst_indices,
    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 = ...,
):

def transfer_kv_all_layer(
    src_k_layers, dst_k_layers, src_v_layers, dst_v_layers,
    src_indices, dst_indices, item_size: int, num_layers: int,
    block_quota: int = 2, num_warps_per_block: int = ...,
):

def transfer_kv_all_layer_lf_pf(
    src_k_layers, dst_k, src_v_layers, dst_v,
    src_indices, dst_indices, item_size: int,
    dst_layout_dim: int, num_layers: int,
    block_quota: int = 2, num_warps_per_block: int = ...,
):

def transfer_kv_all_layer_lf_ph(
    src_k_layers, dst_k, src_v_layers, dst_v,
    src_indices, dst_indices, 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 = ...,
):

def transfer_kv_per_layer_mla(src, dst, src_indices, dst_indices, item_size, ...):
def transfer_kv_per_layer_mla_pf_lf(src, dst, src_indices, dst_indices, layer_id, item_size, src_layout_dim, ...):
def transfer_kv_all_layer_mla(src_layers, dst_layers, src_indices, dst_indices, item_size, num_layers, ...):
def transfer_kv_all_layer_mla_lf_pf(src_layers, dst, src_indices, dst_indices, item_size, dst_layout_dim, num_layers, ...):

def transfer_kv_direct(src_layers: List[torch.Tensor], dst_layers: List[torch.Tensor],
    src_indices, dst_indices, page_size: int):
def 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, dst_ptrs, src_indices, dst_indices, page_size):

Import

from sgl_kernel import transfer_kv_per_layer, transfer_kv_all_layer
from sgl_kernel import transfer_kv_per_layer_mla, transfer_kv_direct

I/O Contract

Inputs

Name	Type	Required	Description
src_k, src_v	torch.Tensor	Yes	Source key/value cache tensors
dst_k, dst_v	torch.Tensor	Yes	Destination key/value cache tensors
src_indices	torch.Tensor	Yes	Source page/slot indices for the transfer
dst_indices	torch.Tensor	Yes	Destination page/slot indices for the transfer
item_size	int	Yes	Size of each KV cache element in bytes
layer_id	int	Yes (pf_lf variants)	Layer index for per-layer transfers in multi-layer layouts
num_layers	int	Yes (all_layer variants)	Total number of transformer layers
page_size	int	Yes (ph/direct variants)	Page size for paged KV cache layouts
head_num	int	Yes (ph variants)	Number of attention heads
block_quota	int	No	CUDA blocks per transfer (default: 2)
num_warps_per_block	int	No	Warps per block (default: 32 CUDA, 16 ROCm)

Outputs

Name	Type	Description
(in-place)	void	All functions modify destination tensors in-place

Usage Examples

from sgl_kernel import transfer_kv_per_layer, transfer_kv_all_layer_mla

# Transfer KV cache for a single layer
transfer_kv_per_layer(
    src_k=old_k_cache,
    dst_k=new_k_cache,
    src_v=old_v_cache,
    dst_v=new_v_cache,
    src_indices=source_page_ids,  # which pages to read from
    dst_indices=dest_page_ids,    # which pages to write to
    item_size=2,  # bfloat16 = 2 bytes
)

# Transfer MLA cache across all layers
transfer_kv_all_layer_mla(
    src_layers=old_mla_cache_layers,
    dst_layers=new_mla_cache_layers,
    src_indices=src_idx,
    dst_indices=dst_idx,
    item_size=2,
    num_layers=32,
)

Related Pages

• Environment:Sgl_project_Sglang_CUDA_GPU_Runtime