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

Stabilize CP HiCache page-first direct transfers on CUDA 13

Browse Source 
CP HiCache direct/page_first_direct all-layer backup was still able to enter sgl-kernel's stale cudaMemcpyBatchAsync path, which segfaults under CUDA 13 before Python can surface an error. The SGLang route now avoids that all-layer sgl-kernel op for page_first_direct backup and uses the TAI per-layer direct LF->PF op for MHA, MLA, and NSA indexer data.\n\nThe load-back path also prepares NSA indexer page indices once per load op and reuses them across per-layer H2D loads, preserving per-layer overlap while removing redundant page-index derivation.\n\nConstraint: Remote runtime is CUDA 13.0 where sgl-kernel's all-layer direct LF->PF op uses the wrong cudaMemcpyBatchAsync ABI.\nRejected: Patch sgl-kernel in this branch | we are converging production HiCache direct/page_first_direct paths onto tai-kernel and do not want to maintain another CUDA-ABI-sensitive copy path here.\nRejected: Collapse H2D load-back into one all-layer op | that would reduce submit count but lose per-layer completion visibility and forward overlap.\nConfidence: medium\nScope-risk: moderate\nDirective: Do not reintroduce sgl_kernel.transfer_kv_all_layer_direct_lf_pf for direct/page_first_direct HiCache backup without CUDA 13 ABI verification.\nTested: g0034 container: PYTHONPATH=python python -m pytest -q -s test/registered/unit/mem_cache/test_nsa_pool_host_unit.py -> 10 passed, 3 warnings.\nTested: g0034 container: PYTHONPATH=python python -m pytest -q test/registered/unit/managers/test_hicache_controller_cp.py -> 61 passed, 3 warnings.\nTested: python -m py_compile python/sglang/srt/mem_cache/memory_pool_host.py python/sglang/srt/managers/cache_controller.py test/registered/unit/mem_cache/test_nsa_pool_host_unit.py test/registered/unit/managers/test_hicache_controller_cp.py\nNot-tested: Full ETE prefill/decode traffic after this commit.\nNot-tested: sgl-kernel implementation itself remains unchanged.
This commit is contained in:
laoyao0822
2026-05-31 01:26:07 +08:00
parent b328baec7c
commit 251a48fb0a
5 changed files with 843 additions and 63 deletions
Download Patch File Download Diff File Expand all files Collapse all files

106
python/sglang/srt/managers/cache_controller.py
View File

@@ -1517,6 +1517,18 @@ class HiCacheController:
if len(self.load_queue) == 0 and len(self.draft_load_queue) == 0:
return -1
def begin_load_op(mem_pool_host, host_indices, device_indices) -> bool:
begin = getattr(mem_pool_host, "begin_load_to_device_op", None)
if begin is None:
return False
begin(host_indices, device_indices, self.io_backend)
return True
def end_load_op(mem_pool_host) -> None:
end = getattr(mem_pool_host, "end_load_to_device_op", None)
if end is not None:
end()
producer_id = self.layer_done_counter.update_producer()
op = CacheOperation.merge_ops(self.load_queue) if self.load_queue else None
draft_op = (
@@ -1540,44 +1552,62 @@ class HiCacheController:
producer_event = self.layer_done_counter.events[producer_id]
producer_event.start_event.record()
with device_module.stream(self.load_stream):
producer_event.start_event.wait(self.load_stream)
draft_layer_num = (
self.draft_mem_pool_device.layer_num if draft_op is not None else 0
)
for i in range(max(self.layer_num, draft_layer_num)):
if draft_op is not None and i < draft_layer_num:
if len(draft_host_indices) > 0:
self.draft_mem_pool_host.load_to_device_per_layer(
self.draft_mem_pool_device,
draft_host_indices,
draft_device_indices,
i,
self.io_backend,
)
if op is not None and i < self.layer_num:
if len(host_indices) > 0:
self.mem_pool_host.load_to_device_per_layer(
self.mem_pool_device,
host_indices,
device_indices,
i,
self.io_backend,
)
producer_event.complete(i)
elif op is None and i < self.layer_num:
producer_event.complete(i)
# NOTE: We must save the host indices and device indices here,
# this is because we need to guarantee that these tensors are
# still alive when the load stream is executing.
if host_indices is not None and host_indices.is_cuda:
host_indices.record_stream(self.load_stream)
if device_indices is not None and device_indices.is_cuda:
device_indices.record_stream(self.load_stream)
if draft_host_indices is not None and draft_host_indices.is_cuda:
draft_host_indices.record_stream(self.load_stream)
if draft_device_indices is not None and draft_device_indices.is_cuda:
draft_device_indices.record_stream(self.load_stream)
target_load_op_started = False
draft_load_op_started = False
try:
with device_module.stream(self.load_stream):
producer_event.start_event.wait(self.load_stream)
if op is not None:
target_load_op_started = begin_load_op(
self.mem_pool_host, host_indices, device_indices
)
if draft_op is not None:
draft_load_op_started = begin_load_op(
self.draft_mem_pool_host,
draft_host_indices,
draft_device_indices,
)
draft_layer_num = (
self.draft_mem_pool_device.layer_num if draft_op is not None else 0
)
for i in range(max(self.layer_num, draft_layer_num)):
if draft_op is not None and i < draft_layer_num:
if len(draft_host_indices) > 0:
self.draft_mem_pool_host.load_to_device_per_layer(
self.draft_mem_pool_device,
draft_host_indices,
draft_device_indices,
i,
self.io_backend,
)
if op is not None and i < self.layer_num:
if len(host_indices) > 0:
self.mem_pool_host.load_to_device_per_layer(
self.mem_pool_device,
host_indices,
device_indices,
i,
self.io_backend,
)
producer_event.complete(i)
elif op is None and i < self.layer_num:
producer_event.complete(i)
# NOTE: We must save the host indices and device indices here,
# this is because we need to guarantee that these tensors are
# still alive when the load stream is executing.
if host_indices is not None and host_indices.is_cuda:
host_indices.record_stream(self.load_stream)
if device_indices is not None and device_indices.is_cuda:
device_indices.record_stream(self.load_stream)
if draft_host_indices is not None and draft_host_indices.is_cuda:
draft_host_indices.record_stream(self.load_stream)
if draft_device_indices is not None and draft_device_indices.is_cuda:
draft_device_indices.record_stream(self.load_stream)
finally:
if draft_load_op_started:
end_load_op(self.draft_mem_pool_host)
if target_load_op_started:
end_load_op(self.mem_pool_host)
self.ack_load_queue.append(
HiCacheAck(

96
python/sglang/srt/mem_cache/memory_pool_host.py
View File

@@ -266,6 +266,14 @@ class HostKVCache(abc.ABC):
"""
raise NotImplementedError()
def begin_load_to_device_op(
self, host_indices: torch.Tensor, device_indices: torch.Tensor, io_backend: str
) -> None:
"""Prepare layer-invariant metadata for one host->device load op."""
def end_load_to_device_op(self) -> None:
"""Release metadata prepared by ``begin_load_to_device_op``."""
@abc.abstractmethod
def backup_from_device_per_layer(
self, device_pool, host_indices, device_indices, layer_id, io_backend
@@ -605,13 +613,23 @@ class MHATokenToKVPoolHost(HostKVCache):
page_size=self.page_size,
)
elif self.layout == "page_first_direct":
transfer_kv_all_layer_direct_lf_pf(
src_ptrs=device_pool.k_buffer + device_pool.v_buffer,
dst_ptrs=[self.k_buffer, self.v_buffer],
src_indices=device_indices,
dst_indices=host_indices,
page_size=self.page_size,
)
# Avoid sgl-kernel's all-layer direct LF->PF path. CUDA 13
# changes the cudaMemcpyBatchAsync ABI, and the stale all-layer
# op can segfault before Python sees an error. Reuse the TAI
# per-layer direct op, which owns the version-specific ABI.
tai_transfer = _load_tai_transfer_kv_per_layer_direct_lf_pf()
for layer_id in range(self.layer_num):
tai_transfer(
src_ptrs=[
device_pool.k_buffer[layer_id],
device_pool.v_buffer[layer_id],
],
dst_ptrs=[self.k_buffer, self.v_buffer],
src_indices=device_indices,
dst_indices=host_indices,
layer_id=layer_id,
page_size=self.page_size,
)
else:
raise ValueError(f"Unsupported layout: {self.layout}")
elif io_backend == "kernel_ascend":
@@ -1090,13 +1108,20 @@ class MLATokenToKVPoolHost(HostKVCache):
page_size=self.page_size,
)
elif self.layout == "page_first_direct":
transfer_kv_all_layer_direct_lf_pf(
src_ptrs=device_pool.kv_buffer,
dst_ptrs=[self.kv_buffer],
src_indices=device_indices,
dst_indices=host_indices,
page_size=self.page_size,
)
# Avoid sgl-kernel's all-layer direct LF->PF path. CUDA 13
# changes the cudaMemcpyBatchAsync ABI, and the stale all-layer
# op can segfault before Python sees an error. Reuse the TAI
# per-layer direct op, which owns the version-specific ABI.
tai_transfer = _load_tai_transfer_kv_per_layer_direct_lf_pf()
for layer_id in range(self.layer_num):
tai_transfer(
src_ptrs=[device_pool.kv_buffer[layer_id]],
dst_ptrs=[self.kv_buffer],
src_indices=device_indices,
dst_indices=host_indices,
layer_id=layer_id,
page_size=self.page_size,
)
else:
raise ValueError(f"Unsupported layout: {self.layout}")
elif io_backend == "kernel_ascend":
@@ -1391,12 +1416,27 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost):
)
return host_page_indices, device_page_indices
def begin_load_to_device_op(
self, host_indices: torch.Tensor, device_indices: torch.Tensor, io_backend: str
) -> None:
self._active_load_indexer_page_indices = None
self._active_load_indexer_page_indices = self._get_indexer_page_indices(
host_indices, device_indices
)
def end_load_to_device_op(self) -> None:
self._active_load_indexer_page_indices = None
def _load_indexer_to_device_per_layer(
self, device_pool, host_indices, device_indices, layer_id, io_backend
):
host_page_indices, device_page_indices = self._get_indexer_page_indices(
host_indices, device_indices
)
prepared_indices = getattr(self, "_active_load_indexer_page_indices", None)
if prepared_indices is None:
host_page_indices, device_page_indices = self._get_indexer_page_indices(
host_indices, device_indices
)
else:
host_page_indices, device_page_indices = prepared_indices
use_kernel = io_backend == "kernel" and self.indexer_page_stride_size % 8 == 0
if use_kernel:
if self.layout == "layer_first":
@@ -1481,13 +1521,21 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost):
page_size=1,
)
elif self.layout == "page_first_direct":
transfer_kv_all_layer_direct_lf_pf(
src_ptrs=device_pool.index_k_with_scale_buffer,
dst_ptrs=[self.index_k_with_scale_buffer],
src_indices=device_page_indices,
dst_indices=host_page_indices,
page_size=1,
)
# Avoid sgl-kernel's all-layer direct LF->PF path here. On CUDA 13
# that op still uses the pre-v2 cudaMemcpyBatchAsync ABI and can
# segfault before surfacing a Python exception. Route NSA indexer
# page-first-direct all-layer backup through the TAI per-layer op,
# which owns the CUDA-version-specific memcpy-batch ABI handling.
tai_transfer = _load_tai_transfer_kv_per_layer_direct_lf_pf()
for layer_id in range(self.layer_num):
tai_transfer(
src_ptrs=[device_pool.index_k_with_scale_buffer[layer_id]],
dst_ptrs=[self.index_k_with_scale_buffer],
src_indices=device_page_indices,
dst_indices=host_page_indices,
layer_id=layer_id,
page_size=1,
)
else:
raise ValueError(f"Unsupported layout: {self.layout}")
else: