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

Count evictable device cache when gating HiCache load-back

Browse Source 
HiCache host hits can be skipped before load-back when the quota gate only counts immediately free KV allocator space. Under CP shared-KV pressure most reusable capacity may be represented as evictable radix-cache leaves, so the gate can incorrectly reject a host hit and leave prefill with cached-token zero despite host residency. Count device evictable cache in the quota estimate while leaving actual owner-lane allocation and eviction checks in the load path.

Constraint: CP HiCache load-back still has to respect owner-lane allocation and allocator eviction semantics.

Rejected: Force load-back regardless of quota | would bypass the scheduler pressure signal and increase OOM risk.

Rejected: Treat cache-hit zero as a transfer issue | logs showed host hits were found but skipped by quota before transfer.

Confidence: medium

Scope-risk: moderate

Directive: Do not remove evictable cache from load-back capacity accounting without checking CP HiCache host-hit behavior under device pressure.

Tested: git diff --check

Tested: remote g0034 container pytest -q test/registered/unit/managers/test_prefill_adder.py test/registered/unit/managers/test_hicache_controller_cp.py test/registered/unit/mem_cache/test_cp_hicache_metadata.py test/registered/unit/mem_cache/test_alloc_pages_with_owners.py (90 passed, 3 warnings)

Not-tested: Full ETE GLM5 CP+HiCache+EAGLE pressure run after this quota change

Co-authored-by: OmX <omx@oh-my-codex.dev>
This commit is contained in:
laoyao0822
2026-05-27 02:41:45 +08:00
co-authored by OmX
parent e5982dcceb
commit d14c02b0dc
4 changed files with 206 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/unit/managers/test_hicache_controller_cp.py
+70 -20
View File
@@ -91,6 +91,20 @@ if "sgl_kernel.kvcacheio" not in sys.modules:
setattr(kvcacheio_stub, name, lambda *args, **kwargs: None)
sys.modules["sgl_kernel.kvcacheio"] = kvcacheio_stub
_sgl_kernel_lib = torch.library.Library("sgl_kernel", "FRAGMENT")
for _schema in (
"sgl_per_token_group_quant_8bit(Tensor input, Tensor(a!) output_q, Tensor(b!) output_s, int group_size, float eps, float fp8_min, float fp8_max, bool scale_ue8m0) -> ()",
"sgl_per_token_group_quant_fp8(Tensor input, Tensor(a!) output_q, Tensor(b!) output_s, int group_size, float eps, float fp8_min, float fp8_max, bool scale_ue8m0) -> ()",
"sgl_per_token_quant_fp8(Tensor input, Tensor(a!) output_q, Tensor(b!) output_s) -> ()",
"fp8_scaled_mm(Tensor mat_a, Tensor mat_b, Tensor scales_a, Tensor scales_b, ScalarType out_dtype, Tensor? bias=None) -> Tensor",
"fp8_blockwise_scaled_mm(Tensor mat_a, Tensor mat_b, Tensor scales_a, Tensor scales_b, ScalarType out_dtype) -> Tensor",
):
try:
_sgl_kernel_lib.define(_schema)
except RuntimeError as exc:
if "already" not in str(exc).lower() and "duplicate" not in str(exc).lower():
raise
from sglang.srt.managers.cache_controller import HiCacheController
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
from sglang.srt.mem_cache.hiradix_cache import CpHiCacheNodeMetadata
@@ -150,9 +164,11 @@ class FakeAllocator:
def __init__(self, alloc_result=None):
self.alloc_result = alloc_result
self.alloc_calls = []
self.owner_alloc_calls = []
self.frees = []
self.cp_size = 4
self.cp_rank = 1
self.page_size = 4
self.device_pool = FakeDevicePool()
def get_kvcache(self):
@@ -164,6 +180,14 @@ class FakeAllocator:
return None
return self.alloc_result[:need_size].clone()
def alloc_pages_with_owners(self, page_owners):
owners = list(page_owners)
self.owner_alloc_calls.append(owners)
if self.alloc_result is None:
return None
need_size = len(owners) * self.page_size
return self.alloc_result[:need_size].clone()
def free(self, indices):
self.frees.append(indices.clone())
return len(indices)
@@ -286,7 +310,7 @@ class TestHiCacheControllerCPWrite(CustomTestCase):
logical_locs = torch.tensor([8, 9, 10], dtype=torch.int64)
with self.assertRaisesRegex(
ValueError, "(host_indices|physical_device_indices).*whole pages"
ValueError, "_write_cp expects page-aligned device_indices"
):
controller.write(logical_locs, node_id=21)
@@ -425,17 +449,40 @@ class TestHiCacheControllerCPLoad(TestHiCacheControllerCPWrite):
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.zeros(max(16, 0), dtype=torch.int8),
page_size=1,
page_owners=torch.tensor([3, 0, 1, 2], dtype=torch.int8),
page_size=4,
)
device_indices = controller.load_cp([node], node_id=11)
controller.start_loading()
self.assertEqual(device_indices.tolist(), list(range(64, 80)))
self.assertEqual(allocator.alloc_calls, [16])
self.assertEqual(allocator.alloc_calls, [])
self.assertEqual(allocator.owner_alloc_calls, [[3, 0, 1, 2]])
self.assertEqual(host_pool.loads[0][1].tolist(), [20, 21, 22, 23])
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))
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,
)
with self.assertRaisesRegex(
RuntimeError, "alloc_pages_with_owners returned unexpected length"
):
controller.load_cp([node], node_id=111)
self.assertEqual(allocator.owner_alloc_calls, [[3, 0, 1, 2]])
self.assertEqual(allocator.frees[0].tolist(), list(range(64, 76)))
def test_cp_load_with_draft_pool_restores_target_and_draft_locs(self):
host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
draft_host_pool = FakeHostPool(
@@ -457,14 +504,15 @@ class TestHiCacheControllerCPLoad(TestHiCacheControllerCPWrite):
owned_positions=torch.tensor([4, 5, 6, 7], dtype=torch.int64),
host_indices=torch.tensor([100, 101, 102, 103], dtype=torch.int64),
draft_host_indices=torch.tensor([200, 201, 202, 203], dtype=torch.int64),
page_owners=torch.zeros(max(16, 0), dtype=torch.int8),
page_size=1,
page_owners=torch.tensor([3, 0, 1, 2], dtype=torch.int8),
page_size=4,
)
device_indices = controller.load_cp([node], node_id=14)
controller.start_loading()
self.assertEqual(device_indices.tolist(), list(range(64, 80)))
self.assertEqual(allocator.owner_alloc_calls, [[3, 0, 1, 2]])
self.assertEqual(host_pool.loads[0][1].tolist(), [20, 21, 22, 23])
self.assertEqual(draft_host_pool.loads[0][1].tolist(), [20, 21, 22, 23])
self.assertIs(draft_host_pool.loads[0][3], draft_device_pool)
@@ -481,14 +529,15 @@ class TestHiCacheControllerCPLoad(TestHiCacheControllerCPWrite):
logical_len=4,
owned_positions=torch.empty((0,), dtype=torch.int64),
host_indices=torch.empty((0,), dtype=torch.int64),
page_owners=torch.zeros(max(4, 0), dtype=torch.int8),
page_size=1,
page_owners=torch.tensor([0], dtype=torch.int8),
page_size=4,
)
device_indices = controller.load_cp([node], node_id=12)
controller.start_loading()
self.assertEqual(device_indices.tolist(), [64, 65, 66, 67])
self.assertEqual(allocator.owner_alloc_calls, [[0]])
self.assertEqual(host_pool.loads, [])
self.assertEqual(len(controller.ack_load_queue), 1)
@@ -499,14 +548,15 @@ class TestHiCacheControllerCPLoad(TestHiCacheControllerCPWrite):
host_indices = HostIndicesTensor(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
node = TreeNode()
node.host_len = 16
node.cp_hicache = type(
"CpHiCacheMetadataStub",
(),
{
"owned_positions": torch.tensor([4, 5, 6, 7], dtype=torch.int64),
"host_indices": host_indices,
},
)()
metadata = 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,
)
metadata.host_indices = host_indices
node.cp_hicache = metadata
controller.load_cp([node], node_id=13)
@@ -530,8 +580,8 @@ class TestHiCacheControllerCPLoad(TestHiCacheControllerCPWrite):
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.zeros(max(16, 0), dtype=torch.int8),
page_size=1,
page_owners=torch.tensor([3, 0, 1, 2], dtype=torch.int8),
page_size=4,
)
with self.assertRaisesRegex(
@@ -549,8 +599,8 @@ class TestHiCacheControllerCPLoad(TestHiCacheControllerCPWrite):
logical_len=16,
owned_positions=torch.tensor([4, 5, 6, 7], dtype=torch.int64),
host_indices=torch.tensor([100, 101, 103, 102], dtype=torch.int64),
page_owners=torch.zeros(max(16, 0), dtype=torch.int8),
page_size=1,
page_owners=torch.tensor([3, 0, 1, 2], dtype=torch.int8),
page_size=4,
)
with self.assertRaisesRegex(ValueError, "host_indices.*contiguous page spans"):