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
test/registered/unit/managers/test_hicache_controller_cp.py
+104
View File
@@ -537,6 +537,61 @@ class TestPageFirstPerLayerBackupTaiKernel(CustomTestCase):
self.assertEqual(kwargs["layer_id"], 1)
self.assertEqual(kwargs["page_size"], 1)
def test_nsa_indexer_load_reuses_precomputed_page_indices_across_layers(self):
calls = []
def fake_direct(src_ptrs, dst_ptrs, src_indices, dst_indices, **kwargs):
calls.append(
(src_ptrs, dst_ptrs, src_indices.clone(), dst_indices.clone(), kwargs)
)
host_pool = NSATokenToKVPoolHost.__new__(NSATokenToKVPoolHost)
host_pool.layout = "page_first_direct"
host_pool.page_size = 4
host_pool.index_k_with_scale_buffer = torch.empty(
(8, 3, 1, 32), dtype=torch.uint8
)
device_pool = type("DevicePool", (), {})()
device_pool.index_k_with_scale_buffer = torch.empty(
(3, 8, 32), dtype=torch.uint8
)
host_indices = torch.tensor([4, 5, 6, 7], dtype=torch.int64)
device_indices = torch.tensor([12, 13, 14, 15], dtype=torch.int64)
original_getter = host_pool._get_indexer_page_indices
getter_calls = []
def counting_getter(h, d):
getter_calls.append((h.clone(), d.clone()))
return original_getter(h, d)
host_pool._get_indexer_page_indices = counting_getter
with patch(
"sglang.srt.mem_cache.memory_pool_host._load_tai_transfer_kv_per_layer_direct_pf_lf",
return_value=fake_direct,
):
host_pool.begin_load_to_device_op(
host_indices, device_indices, io_backend="direct"
)
try:
for layer_id in range(3):
host_pool._load_indexer_to_device_per_layer(
device_pool,
host_indices,
device_indices,
layer_id=layer_id,
io_backend="direct",
)
finally:
host_pool.end_load_to_device_op()
self.assertEqual(len(getter_calls), 1)
self.assertEqual(len(calls), 3)
self.assertEqual([call[4]["layer_id"] for call in calls], [0, 1, 2])
self.assertEqual([call[2].tolist() for call in calls], [[1], [1], [1]])
self.assertEqual([call[3].tolist() for call in calls], [[3], [3], [3]])
class FakeAllocator:
def __init__(self, alloc_result=None):
@@ -1109,6 +1164,55 @@ class TestHiCacheControllerCPLoad(TestHiCacheControllerCPWrite):
self.assertEqual(host_pool.loads[0][0].tolist(), [100, 101, 102, 103])
self.assertEqual(host_pool.loads[0][1].tolist(), [20, 21, 22, 23])
def test_start_loading_prepares_load_op_once_for_all_layers(self):
class PreparingFakeHostPool(FakeHostPool):
def __init__(self, alloc_result):
super().__init__(alloc_result)
self.begin_calls = []
self.end_calls = 0
self.active = False
def begin_load_to_device_op(self, host_indices, device_indices, io_backend):
self.begin_calls.append(
(host_indices.clone(), device_indices.clone(), io_backend)
)
self.active = True
def end_load_to_device_op(self):
self.end_calls += 1
self.active = False
def load_to_device_per_layer(
self, device_pool, host_indices, device_indices, layer_id, io_backend
):
assert self.active
super().load_to_device_per_layer(
device_pool, host_indices, device_indices, layer_id, io_backend
)
host_pool = PreparingFakeHostPool(
torch.tensor([100, 101, 102, 103], dtype=torch.int64)
)
allocator = FakeAllocator(alloc_result=torch.arange(64, 80, dtype=torch.int64))
allocator.device_pool = FakeDevicePool(layer_num=3)
controller = self.make_controller(host_pool, allocator=allocator, cp_rank=1)
node = TreeNode()
node.host_len = 16
node.cp_hicache = CpHiCacheNodeMetadata(
logical_len=16,
owned_positions=torch.tensor([4, 5, 6, 7], dtype=torch.int64),
host_indices=torch.tensor([100, 101, 102, 103], dtype=torch.int64),
page_owners=torch.tensor([3, 0, 1, 2], dtype=torch.int8),
page_size=4,
)
controller.load_cp([node], node_id=113)
controller.start_loading()
self.assertEqual(len(host_pool.begin_calls), 1)
self.assertEqual(host_pool.end_calls, 1)
self.assertEqual([load[2] for load in host_pool.loads], [0, 1, 2])
def test_cp_load_frees_unexpected_owner_allocator_length(self):
host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
allocator = FakeAllocator(alloc_result=torch.arange(64, 76, dtype=torch.int64))
test/registered/unit/mem_cache/test_nsa_pool_host_unit.py
+152 -1
View File
@@ -1,10 +1,13 @@
import unittest
from unittest.mock import patch
import torch
from sglang.srt.mem_cache.memory_pool import NSATokenToKVPool
from sglang.srt.mem_cache.memory_pool_host import (
ALLOC_MEMORY_FUNCS,
MHATokenToKVPoolHost,
MLATokenToKVPoolHost,
NSATokenToKVPoolHost,
alloc_with_pin_memory,
)
@@ -101,13 +104,16 @@ class TestNSAHiCacheTransfer(CustomTestCase):
host_pool.backup_from_device_all_layer(
device_pool, host_indices, device_indices, io_backend
)
torch.cuda.synchronize()
for layer_id in range(layer_num):
for host_page, device_page in zip(
host_pages.tolist(), device_pages.tolist()
):
if layout == "page_first_direct":
got = host_pool.index_k_with_scale_buffer[host_page, layer_id].cpu()
got = host_pool.index_k_with_scale_buffer[
host_page, layer_id, 0
].cpu()
else:
got = host_pool.index_k_with_scale_buffer[layer_id][host_page].cpu()
expected = device_pool.index_k_with_scale_buffer[layer_id][
@@ -143,6 +149,102 @@ class TestNSAHiCacheTransfer(CustomTestCase):
)
class TestPageFirstDirectAllLayerBackupRoute(CustomTestCase):
def test_mla_page_first_direct_all_layer_backup_uses_tai_per_layer_route(self):
host_pool = MLATokenToKVPoolHost.__new__(MLATokenToKVPoolHost)
host_pool.layout = "page_first_direct"
host_pool.page_size = 4
host_pool.layer_num = 2
host_pool.kv_buffer = "host-mla-page-first"
device_pool = type(
"FakeDevicePool",
(),
{"kv_buffer": ["device-mla-layer-0", "device-mla-layer-1"]},
)()
calls = []
def fake_tai_transfer(**kwargs):
calls.append(kwargs)
with (
patch(
"sglang.srt.mem_cache.memory_pool_host.transfer_kv_all_layer_direct_lf_pf",
side_effect=AssertionError(
"MLA page_first_direct all-layer backup must not use sgl-kernel direct LF->PF"
),
),
patch(
"sglang.srt.mem_cache.memory_pool_host._load_tai_transfer_kv_per_layer_direct_lf_pf",
return_value=fake_tai_transfer,
),
):
host_pool.backup_from_device_all_layer(
device_pool,
torch.tensor([0, 1, 2, 3], dtype=torch.int64),
torch.tensor([8, 9, 10, 11], dtype=torch.int64),
"direct",
)
self.assertEqual([call["layer_id"] for call in calls], [0, 1])
for layer_id, call in enumerate(calls):
self.assertEqual(call["src_ptrs"], [f"device-mla-layer-{layer_id}"])
self.assertEqual(call["dst_ptrs"], ["host-mla-page-first"])
self.assertEqual(call["src_indices"].tolist(), [8, 9, 10, 11])
self.assertEqual(call["dst_indices"].tolist(), [0, 1, 2, 3])
self.assertEqual(call["page_size"], 4)
def test_mha_page_first_direct_all_layer_backup_uses_tai_per_layer_route(self):
host_pool = MHATokenToKVPoolHost.__new__(MHATokenToKVPoolHost)
host_pool.layout = "page_first_direct"
host_pool.page_size = 4
host_pool.layer_num = 2
host_pool.kv_buffer = ["host-k-page-first", "host-v-page-first"]
device_pool = type(
"FakeDevicePool",
(),
{
"k_buffer": ["device-k-layer-0", "device-k-layer-1"],
"v_buffer": ["device-v-layer-0", "device-v-layer-1"],
},
)()
calls = []
def fake_tai_transfer(**kwargs):
calls.append(kwargs)
with (
patch(
"sglang.srt.mem_cache.memory_pool_host.transfer_kv_all_layer_direct_lf_pf",
side_effect=AssertionError(
"MHA page_first_direct all-layer backup must not use sgl-kernel direct LF->PF"
),
),
patch(
"sglang.srt.mem_cache.memory_pool_host._load_tai_transfer_kv_per_layer_direct_lf_pf",
return_value=fake_tai_transfer,
),
):
host_pool.backup_from_device_all_layer(
device_pool,
torch.tensor([0, 1, 2, 3], dtype=torch.int64),
torch.tensor([8, 9, 10, 11], dtype=torch.int64),
"direct",
)
self.assertEqual([call["layer_id"] for call in calls], [0, 1])
for layer_id, call in enumerate(calls):
self.assertEqual(
call["src_ptrs"],
[f"device-k-layer-{layer_id}", f"device-v-layer-{layer_id}"],
)
self.assertEqual(
call["dst_ptrs"], ["host-k-page-first", "host-v-page-first"]
)
self.assertEqual(call["src_indices"].tolist(), [8, 9, 10, 11])
self.assertEqual(call["dst_indices"].tolist(), [0, 1, 2, 3])
self.assertEqual(call["page_size"], 4)
class TestNSAIndexerPageIndices(CustomTestCase):
def make_host_pool_stub(self, page_size: int):
host_pool = NSATokenToKVPoolHost.__new__(NSATokenToKVPoolHost)
@@ -189,6 +291,55 @@ class TestNSAIndexerPageIndices(CustomTestCase):
torch.tensor([16, 17, 18, 19], dtype=torch.int64),
)
def test_page_first_direct_all_layer_indexer_backup_uses_tai_per_layer_route(self):
host_pool = self.make_host_pool_stub(page_size=4)
host_pool.layout = "page_first_direct"
host_pool.indexer_page_stride_size = 8
host_pool.layer_num = 3
host_pool.index_k_with_scale_buffer = "host-page-first-indexer"
device_pool = type(
"FakeDevicePool",
(),
{
"index_k_with_scale_buffer": [
"device-layer-0",
"device-layer-1",
"device-layer-2",
]
},
)()
calls = []
def fake_tai_transfer(**kwargs):
calls.append(kwargs)
with (
patch(
"sglang.srt.mem_cache.memory_pool_host.transfer_kv_all_layer_direct_lf_pf",
side_effect=AssertionError(
"NSA page_first_direct indexer all-layer backup must not use sgl-kernel direct LF->PF"
),
),
patch(
"sglang.srt.mem_cache.memory_pool_host._load_tai_transfer_kv_per_layer_direct_lf_pf",
return_value=fake_tai_transfer,
),
):
host_pool._backup_indexer_from_device_all_layer(
device_pool,
torch.tensor([0, 1, 2, 3, 4, 5, 6, 7], dtype=torch.int64),
torch.tensor([8, 9, 10, 11, 12, 13, 14, 15], dtype=torch.int64),
"direct",
)
self.assertEqual([call["layer_id"] for call in calls], [0, 1, 2])
for layer_id, call in enumerate(calls):
self.assertEqual(call["src_ptrs"], [f"device-layer-{layer_id}"])
self.assertEqual(call["dst_ptrs"], ["host-page-first-indexer"])
self.assertEqual(call["src_indices"].tolist(), [2, 3])
self.assertEqual(call["dst_indices"].tolist(), [0, 1])
self.assertEqual(call["page_size"], 1)
if __name__ == "__main__":
unittest.main()