chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 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/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()