minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

[HiCache] TMA-staged host<->device KV transfer kernel (sm_90+) (#40278)

Browse Source
This commit is contained in:
cctry
2026-09-21 10:38:23 -07:00
committed by GitHub
parent 0cb37c018c
commit 7ad55e4386
7 changed files with 1060 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/kernels/benchmark/kvcache/bench_hicache.py
+90 -14
View File
@@ -24,8 +24,11 @@ from sgl_kernel import transfer_kv_all_layer, transfer_kv_per_layer
from sglang.kernels.jit.benchmark import marker
from sglang.kernels.jit.benchmark.utils import get_benchmark_range
from sglang.kernels.ops.kvcache.hicache import (
transfer_hicache_all_layer,
transfer_hicache_one_layer,
DEFAULT_BLOCK_QUOTA,
TMA_BLOCK_QUOTA,
_default_unroll,
_jit_hicache_module,
_jit_hicache_tma_module,
)
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
@@ -108,15 +111,19 @@ def sglang_jit_transfer_one(
indices_src: torch.Tensor,
element_dim: int,
) -> None:
"""SGL JIT Kernel for single layer transfer."""
transfer_hicache_one_layer(
k_cache_dst,
v_cache_dst,
"""SGL JIT register kernel for single layer transfer (bypasses TMA routing)."""
element_size = element_dim * k_cache_dst.element_size()
_jit_hicache_module(
element_size=element_size,
unroll=_default_unroll(element_size),
block_quota=DEFAULT_BLOCK_QUOTA,
).launch_one(
k_cache_dst.view(-1, element_dim),
v_cache_dst.view(-1, element_dim),
indices_dst,
k_cache_src,
v_cache_src,
k_cache_src.view(-1, element_dim),
v_cache_src.view(-1, element_dim),
indices_src,
element_dim=element_dim,
)
@@ -153,17 +160,58 @@ def sglang_jit_transfer_all(
stride_bytes: int,
element_size: int,
) -> None:
"""SGL JIT Kernel for all layer transfer."""
transfer_hicache_all_layer(
"""SGL JIT register kernel for all layer transfer (bypasses TMA routing)."""
_jit_hicache_module(
element_size=element_size,
unroll=_default_unroll(element_size),
block_quota=DEFAULT_BLOCK_QUOTA,
).launch_all(
k_ptrs_dst,
v_ptrs_dst,
indices_dst,
k_ptrs_src,
v_ptrs_src,
indices_src,
kv_cache_src_stride_bytes=stride_bytes,
kv_cache_dst_stride_bytes=stride_bytes,
element_size=element_size,
stride_bytes,
stride_bytes,
)
def sglang_tma_transfer_one(
k_cache_dst: torch.Tensor,
v_cache_dst: torch.Tensor,
indices_dst: torch.Tensor,
k_cache_src: torch.Tensor,
v_cache_src: torch.Tensor,
indices_src: torch.Tensor,
) -> None:
"""SGL TMA staging kernel for single layer transfer."""
_jit_hicache_tma_module(block_quota=TMA_BLOCK_QUOTA).launch_one(
k_cache_dst, v_cache_dst, indices_dst, k_cache_src, v_cache_src, indices_src
)
def sglang_tma_transfer_all(
k_ptrs_dst: torch.Tensor,
v_ptrs_dst: torch.Tensor,
indices_dst: torch.Tensor,
k_ptrs_src: torch.Tensor,
v_ptrs_src: torch.Tensor,
indices_src: torch.Tensor,
stride_bytes: int,
element_size: int,
) -> None:
"""SGL TMA staging kernel for all layer transfer."""
_jit_hicache_tma_module(block_quota=TMA_BLOCK_QUOTA).launch_all(
k_ptrs_dst,
v_ptrs_dst,
indices_dst,
k_ptrs_src,
v_ptrs_src,
indices_src,
stride_bytes,
stride_bytes,
element_size,
)
@@ -191,6 +239,13 @@ ELEMENT_SIZE_RANGE = get_benchmark_range(
LINE_VALS = ["aot", "jit", "torch"]
if DISABLE_TORCH:
LINE_VALS.remove("torch")
# The TMA staging kernel needs sm_90+ (cp.async.bulk); skip the line elsewhere.
if (
torch.cuda.is_available()
and torch.version.hip is None
and torch.cuda.get_device_capability()[0] >= 9
):
LINE_VALS.insert(2, "tma")
# =============================================================================
@@ -246,6 +301,17 @@ def benchmark_one_layer_h2d(element_size: int, batch_size: int, provider: str):
)
for i in range(NUM_LAYERS)
],
"tma": lambda: [
sglang_tma_transfer_one(
k_cache_dst[i],
v_cache_dst[i],
indices_dst_gpu,
k_cache_src[i],
v_cache_src[i],
indices_src_gpu,
)
for i in range(NUM_LAYERS)
],
"torch": lambda: [
pytorch_transfer(
k_cache_dst[i],
@@ -329,6 +395,16 @@ def benchmark_all_layer_d2h(element_size: int, batch_size: int, provider: str):
element_bytes,
element_bytes,
),
"tma": lambda: sglang_tma_transfer_all(
k_ptrs_dst,
v_ptrs_dst,
indices_dst_gpu,
k_ptrs_src,
v_ptrs_src,
indices_src_gpu,
element_bytes,
element_bytes,
),
"torch": lambda: [
pytorch_transfer(
k_caches_dst[i],
test/registered/kernels/ops/kvcache/test_hicache_tma.py
+173
View File
@@ -0,0 +1,173 @@
import sys
import pytest
import torch
from sglang.kernels.ops.kvcache.hicache import _jit_hicache_tma_module
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=60, stage="base-b-kernel-unit", runner_config="1-gpu-large")
pytestmark = pytest.mark.skipif(
not torch.cuda.is_available()
or torch.version.hip is not None
or torch.cuda.get_device_capability()[0] < 9,
reason="HiCache TMA kernel requires SM90+",
)
POOL_TOKENS = 8192
NUM_LAYERS = 3
ROW_DIM = 256 # 512-byte bf16 rows: below the register kernel's 128 B unit width x 4
def _token_indices(num_tokens: int, page_size: int, dtype: torch.dtype, seed: int):
gen = torch.Generator().manual_seed(seed)
pages = torch.randperm(POOL_TOKENS // page_size, generator=gen)[
: num_tokens // page_size
]
idx = (pages[:, None] * page_size + torch.arange(page_size)).reshape(-1)
return idx.to(device="cuda", dtype=dtype)
def _fill(t: torch.Tensor, seed: int) -> None:
t.view(torch.int16).copy_(
torch.randint(
0,
30000,
t.shape,
dtype=torch.int16,
generator=torch.Generator().manual_seed(seed),
)
)
def _host_view(layout: str, layer: int):
if layout == "layer_first":
return torch.empty(POOL_TOKENS, ROW_DIM, dtype=torch.bfloat16, pin_memory=True)
# page_first: [tokens, layers, dim]; a per-layer view has strided rows
return torch.empty(
POOL_TOKENS, NUM_LAYERS, ROW_DIM, dtype=torch.bfloat16, pin_memory=True
)[:, layer]
@pytest.mark.parametrize("host_layout", ["layer_first", "page_first"])
@pytest.mark.parametrize("index_dtype", [torch.int64, torch.int32])
@pytest.mark.parametrize("page_size", [128, 1])
def test_one_layer_roundtrip(
host_layout: str, index_dtype: torch.dtype, page_size: int
) -> None:
"""H2D then D2H of one layer; page runs take the single-op paths (bulk copy,
tensor-map box, bulk store), scattered rows take the per-row paths, and the
odd token count leaves a partial tail chunk."""
module = _jit_hicache_tma_module(block_quota=2)
num_tokens = 2048 + (96 if page_size == 1 else 0)
k_host, v_host = _host_view(host_layout, 1), _host_view(host_layout, 2)
k_dev = torch.zeros(POOL_TOKENS, ROW_DIM, dtype=torch.bfloat16, device="cuda")
v_dev = torch.zeros_like(k_dev)
_fill(k_host, 1)
_fill(v_host, 2)
host_idx = _token_indices(num_tokens, page_size, index_dtype, seed=3)
dev_idx = _token_indices(num_tokens, page_size, index_dtype, seed=4)
module.launch_one(k_dev, v_dev, dev_idx, k_host, v_host, host_idx)
torch.cuda.synchronize()
assert torch.equal(k_dev[dev_idx.long()].cpu(), k_host[host_idx.cpu().long()])
assert torch.equal(v_dev[dev_idx.long()].cpu(), v_host[host_idx.cpu().long()])
untouched = torch.ones(POOL_TOKENS, dtype=torch.bool, device="cuda")
untouched[dev_idx.long()] = False
assert not k_dev[untouched].any() and not v_dev[untouched].any()
_fill(k_dev, 5)
_fill(v_dev, 6)
k_host.zero_()
v_host.zero_()
module.launch_one(k_host, v_host, host_idx, k_dev, v_dev, dev_idx)
torch.cuda.synchronize()
assert torch.equal(k_host[host_idx.cpu().long()], k_dev[dev_idx.long()].cpu())
assert torch.equal(v_host[host_idx.cpu().long()], v_dev[dev_idx.long()].cpu())
def _ptr_table(tensors) -> torch.Tensor:
return torch.tensor(
[t.data_ptr() for t in tensors], dtype=torch.uint64, device="cuda"
)
def test_all_layer_tables_lf_to_pf() -> None:
"""All-layer D2H through per-layer pointer tables into a page-first host pool
(strided destination rows), the write-back shape."""
module = _jit_hicache_tma_module(block_quota=2)
k_dev = [
torch.empty(POOL_TOKENS, ROW_DIM, dtype=torch.bfloat16, device="cuda")
for _ in range(NUM_LAYERS)
]
v_dev = [torch.empty_like(k_dev[0]) for _ in range(NUM_LAYERS)]
for i, t in enumerate(k_dev + v_dev):
_fill(t, 10 + i)
k_host = torch.zeros(
POOL_TOKENS, NUM_LAYERS, ROW_DIM, dtype=torch.bfloat16, pin_memory=True
)
v_host = torch.zeros_like(k_host).pin_memory()
host_idx = _token_indices(2048, 128, torch.int64, seed=7)
dev_idx = _token_indices(2048, 128, torch.int64, seed=8)
row_bytes = ROW_DIM * 2
module.launch_all(
_ptr_table([k_host[:, l] for l in range(NUM_LAYERS)]),
_ptr_table([v_host[:, l] for l in range(NUM_LAYERS)]),
host_idx,
_ptr_table(k_dev),
_ptr_table(v_dev),
dev_idx,
row_bytes,
NUM_LAYERS * row_bytes,
row_bytes,
)
torch.cuda.synchronize()
for l in range(NUM_LAYERS):
assert torch.equal(k_host[host_idx.cpu(), l], k_dev[l][dev_idx].cpu())
assert torch.equal(v_host[host_idx.cpu(), l], v_dev[l][dev_idx].cpu())
untouched = torch.ones(POOL_TOKENS, dtype=torch.bool)
untouched[host_idx.cpu()] = False
assert not k_host[untouched].any() and not v_host[untouched].any()
def test_mla_single_buffer() -> None:
"""MLA rows (576 x bf16 = 1152 B, not a multiple of 128 B) through the
single-buffer entry points, one layer and all layers."""
module = _jit_hicache_tma_module(block_quota=2)
dim = 576
dev = [
torch.empty(POOL_TOKENS, dim, dtype=torch.bfloat16, device="cuda")
for _ in range(NUM_LAYERS)
]
for i, t in enumerate(dev):
_fill(t, 20 + i)
host = torch.zeros(
POOL_TOKENS, NUM_LAYERS, dim, dtype=torch.bfloat16, pin_memory=True
)
host_idx = _token_indices(1024, 128, torch.int64, seed=9)
dev_idx = _token_indices(1024, 128, torch.int64, seed=10)
row_bytes = dim * 2
module.launch_all_mla(
_ptr_table([host[:, l] for l in range(NUM_LAYERS)]),
host_idx,
_ptr_table(dev),
dev_idx,
row_bytes,
NUM_LAYERS * row_bytes,
row_bytes,
)
torch.cuda.synchronize()
for l in range(NUM_LAYERS):
assert torch.equal(host[host_idx.cpu(), l], dev[l][dev_idx].cpu())
dev[0].zero_()
module.launch_one_mla(dev[0], dev_idx, host[:, 0], host_idx)
torch.cuda.synchronize()
assert torch.equal(dev[0][dev_idx].cpu(), host[host_idx.cpu(), 0])
if __name__ == "__main__":
sys.exit(pytest.main([__file__, "-v", "-s"]))