Reduce CP HiCache capacity synchronization to owner-lane logic

CP shared KV and HiCache now use owner-lane metadata as the
authoritative capacity view for host write admission and GPU load-back
planning. This removes the debug scalar capacity env and keeps CP load-back
from relying on a rank-wide scalar collective when per-owner availability is
already known. The load-back planner also accounts for evicting child leaves
that unlock ancestor device residency, which fixes small lane deficits despite
large aggregate evictable capacity.

The commit also adds gated CPU timing logs for CP shared-KV MLA/index
prefetch and a CUDA microbenchmark for comparing dense all-reduce with
owner-packed all-gather layouts. The timing logs are intentionally behind the
existing MLA prefetch log env and should not be enabled for throughput
measurements.

Constraint: CP shared KV owner lanes require target/draft capacity decisions to preserve page_owners rather than total-token scalars
Constraint: CUDA collective benchmarks must run on target GPU hosts, not locally
Rejected: Keep SGLANG_CP_HICACHE_CAPACITY_DEBUG observer env | owner-lane admission now replaces that scalar debug path
Rejected: Add a silent scalar-allreduce fallback | unexpected owner-lane mismatch should fail fast or log loudly
Confidence: medium
Scope-risk: moderate
Directive: Do not reintroduce CP capacity collectives on the scheduler hot path without proving the owner-lane metadata is insufficient
Directive: Disable SGLANG_CP_SHARED_KV_LOG_MLA_PREFETCH for end-to-end performance runs; it is diagnostic and high-volume
Tested: git diff --check
Tested: python -m py_compile on changed runtime/test/benchmark Python files
Tested: remote pytest -q test/registered/unit/mem_cache/test_cp_hicache_load_back_owner_lanes.py test/registered/unit/mem_cache/test_cp_hicache_metadata.py (81 passed, 5 warnings)
Not-tested: CUDA benchmark benchmark/hicache/bench_cp_shared_kv_prefetch_collective.py
Not-tested: full GLM5 E2E throughput after this commit
This commit is contained in:
laoyao0822
2026-05-28 08:31:49 +08:00
parent ff33446787
commit 25f2147677
8 changed files with 996 additions and 201 deletions

View File

@@ -272,6 +272,47 @@ class TestCpHiCacheLoadBackOwnerLanes(CustomTestCase):
self.assertEqual(target.value.tolist(), loaded.tolist())
self.assertIn(target.id, cache.ongoing_load_back)
def test_load_back_evicts_blocking_leaf_to_unlock_parent_owner_lane(self):
allocator = _make_allocator()
# Only owner lane 1 is initially available. The target needs owner lane
# 0. The current leaf itself belongs to owner 1, but evicting it makes
# its parent evictable and the parent frees owner 0.
allocator.free_pages = torch.tensor([2], dtype=torch.int64)
cache = _make_cache(allocator)
parent = _make_node(
40,
600,
[0],
value=torch.tensor([4, 5, 6, 7], dtype=torch.int64),
priority=0,
)
_attach_child(cache, cache.root_node, parent)
child = _make_node(
41,
700,
[1],
value=torch.tensor([8, 9, 10, 11], dtype=torch.int64),
priority=0,
)
_attach_child(cache, parent, child)
cache.evictable_leaves.add(child)
cache.evictable_size_ = len(parent.key) + len(child.key)
target = _make_node(42, 800, [0], value=None, priority=10)
_attach_child(cache, cache.root_node, target)
loaded = cache.load_back(target, mem_quota=100)
self.assertIsNotNone(loaded)
self.assertEqual(cache.cache_controller.load_calls, 1)
self.assertEqual(
[indices.tolist() for indices in cache.cache_controller.evicted_device_indices],
[[8, 9, 10, 11], [4, 5, 6, 7]],
)
self.assertEqual(target.value.tolist(), loaded.tolist())
self.assertIn(target.id, cache.ongoing_load_back)
def test_load_back_failure_leaves_node_unassigned_and_unlocked(self):
allocator = _make_allocator()
allocator.free_pages = torch.tensor([1, 2], dtype=torch.int64)

View File

@@ -64,6 +64,7 @@ except (ImportError, RuntimeError):
"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",
"moe_fused_gate(Tensor input, Tensor bias, int num_expert_group, int topk_group, int topk, int num_fused_shared_experts, float routed_scaling_factor, bool apply_routed_scaling_factor_on_output) -> Tensor[]",
):
try:
_sgl_kernel_lib.define(_schema)
@@ -83,6 +84,7 @@ for _schema in (
"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",
"moe_fused_gate(Tensor input, Tensor bias, int num_expert_group, int topk_group, int topk, int num_fused_shared_experts, float routed_scaling_factor, bool apply_routed_scaling_factor_on_output) -> Tensor[]",
):
try:
_sgl_kernel_lib.define(_schema)
@@ -629,6 +631,20 @@ class FakeTokenAllocator:
def available_size(self):
return 0
def compute_owner_lane_stats(self, page_owners):
if len(page_owners) == 0:
return [], [], []
cp_size = max(int(owner) for owner in page_owners) + 1
required = [0] * cp_size
for owner in page_owners:
required[int(owner)] += 1
available = list(required)
deficits = [0] * cp_size
return required, available, deficits
def allocator_state_str(self):
return "FakeTokenAllocator"
class TestHiRadixCacheCPBackup(CustomTestCase):
def test_session_aware_cache_forwards_cp_hicache_prepare(self):
@@ -1243,7 +1259,7 @@ class TestHiRadixCacheCPBackup(CustomTestCase):
with self.assertRaisesRegex(
RuntimeError,
"planner predicted no eviction need",
"owner-lane admission predicted no host deficit",
):
cache.write_backup(node)
@@ -1868,6 +1884,8 @@ class TestHiRadixCacheCPLoadBack(CustomTestCase):
cache._uses_cp_hicache = True
cache.root_node = TreeNode()
cache.device = "cpu"
cache.page_size = 1
cache.token_to_kv_pool_allocator = FakeTokenAllocator()
cache.load_back_threshold = 1
cache.evictable_size_ = 0
cache.metrics_collector = None
@@ -1911,6 +1929,8 @@ class TestHiRadixCacheCPLoadBack(CustomTestCase):
cache._uses_cp_hicache = True
cache.root_node = TreeNode()
cache.root_node.value = torch.empty((0,), dtype=torch.int64)
cache.page_size = 1
cache.token_to_kv_pool_allocator = FakeTokenAllocator()
cache.load_back_threshold = 5
cache.evictable_size_ = 0
cache.metrics_collector = None