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

Batch CP HiCache host admission before write reservation

Browse Source 
CP shared-KV prefill batching can present multiple requests to HiCache write-through in one forward. Keeping host free-room admission per request causes repeated small L2 evictions and defeats the free-room policy, while request metadata and radix attachment still need per-request ownership.\n\nThis change adds a batch prepare path that aggregates required host tokens by CP owner lane, performs one host admission/eviction step for the batch, and then reserves/submits each request independently. SessionAwareCache and the scheduler prefer the batch API when available.\n\nConstraint: Radix nodes, host slots, draft slots, rollback, and ack semantics remain per request.\nRejected: Merge reservations or radix nodes across requests | would complicate rollback and writing_check ack ownership.\nRejected: Add collective synchronization for host admission | local owner-lane logic already provides deterministic capacity accounting.\nConfidence: medium\nScope-risk: moderate\nDirective: Do not reintroduce per-request host free-room eviction in the scheduler path without profiling host-full workloads.\nTested: local py_compile for touched Python files\nTested: remote g0034 docker PYTHONPATH=python python -m pytest -q test/registered/unit/mem_cache/test_cp_hicache_metadata.py => 117 passed, 5 warnings\nNot-tested: full ETE bs>1 HiCache run under production traffic
This commit is contained in:
laoyao0822
2026-06-04 04:32:40 +08:00
parent 108e7d866d
commit 02af370e87
5 changed files with 342 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_parallel_workstreams_zh.md
+58
View File
@@ -935,3 +935,61 @@ scheduler admission 不做 bf16/fp8 字节级估算。原因是实际 KV pool
1. `cp_shared_kv_prefill_max_total_extend_tokens` 首版使用 page-aligned extend 累计，和 allocator 的最小 page 单位一致；日志/指标里如果要展示用户侧 token，需要另加 valid-token 统计。
2. 该 gate 只控制 scheduler 组 batch，不解决所有 bs>1 runtime correctness；打开 ETE 时仍应保留现有 fail-fast，特别是 draft/EAGLE、logprob、hidden capture、compute padding 路径。
3. 如果 ETE 中 `KVCapacityWaitError` 频繁出现，下一步应把 owner-lane free pages/evictable pages 的 batch precheck 前移到 scheduler，而不是引入 all-reduce。
## 19. 2026-06-04 L2 host write admission batch 化
### 背景
打开 scheduler bs>1 gate 后，CP HiCache write-through 的 radix node / host reservation 仍然必须保持 per request 独立；但 host 侧容量 admission 如果仍按 request 串行执行，会在同一个 prefill batch 内反复触发 L2 host free-room eviction：
- correctness 没问题，因为 `reserve_write_cp()` 最终会消费真实 host allocator slot；
- 性能和碎片化较差，因为每个 request 都可能单独计算一次 `required + free_room_target - available`，在 host 接近满载时容易多次小额 eviction；
- 这与 L1/L2 free-room 的目标相反：应尽量一次 evict 到目标余量，降低 eviction 频率并给后续 allocate 更连续的空间。
### 实现原则
只把 **host write admission / eviction planning** 聚合到 batch level，不合并 request 生命周期：
- `CpHiCacheNodeMetadata`、host slot reservation、draft host slot、radix attach、rollback、ack 仍 per request；
- batch admission 只聚合多个 request 的 `required_by_owner`，用一次 `_cp_build_write_admission_from_required()` 计算 host free-room deficit；
- 如果 batch admission 需要 eviction，则只在 forward 前触发一次 `_evict_cp_host_for_write_admission()`；
- admission 成功后，每个 request 继续调用 `reserve_write_cp()` 拿自己的 host slots，但跳过重复的 per-request initial admission；
- 如果 batch admission 无法满足 deficit，会显式打 `[CP_HICACHE_FALLBACK][prepare_write_backups_batch_admission_failed]`，再回到 per-request reservation 路径，不 silent fallback；
- 不新增 collective，不改变 CP owner-lane 分布，不改变 TAI kernel 接口。
### 完成状态
代码路径：
- `Scheduler._prepare_hicache_write_backups_before_forward()` 优先调用 `prepare_write_backups_for_reqs(batch.reqs)`；
- `SessionAwareCache.prepare_write_backups_for_reqs()` 透传到 inner cache，缺失 batch API 时才逐 req fallback；
- `HiRadixCache.prepare_write_backups_for_reqs()`：
- 先构造每个 request 的 `CpWriteBackupCandidate`；
- 聚合所有 candidate 的 `_cp_required_host_tokens_by_rank(kv_indices)`；
- 对聚合后的 `required_by_owner` 做一次 host write admission / eviction；
- 再逐 request 做实际 reservation 和 per-layer backup registration。
新增测试：
- `test_cp_host_write_batch_admission_uses_aggregate_required`
- 验证 batch admission 对聚合 required 使用 free-room 公式。
- `test_prepare_write_backups_for_reqs_runs_one_batch_admission`
- 验证两个 request 只触发一次 batch admission / eviction，最终仍各自得到 prepared backup，并提交两个独立 reservation。
验证：
```text
python -m py_compile \
python/sglang/srt/mem_cache/hiradix_cache.py \
python/sglang/srt/mem_cache/session_aware_cache.py \
python/sglang/srt/managers/scheduler.py \
test/registered/unit/mem_cache/test_cp_hicache_metadata.py
PYTHONPATH=python python -m pytest -q \
test/registered/unit/mem_cache/test_cp_hicache_metadata.py::TestCpHiCacheFreeRoom
=> 8 passed, 3 warnings
PYTHONPATH=python python -m pytest -q \
test/registered/unit/mem_cache/test_cp_hicache_metadata.py
=> 117 passed, 5 warnings
```
python/sglang/srt/managers/scheduler.py
+12
View File
@@ -2654,6 +2654,18 @@ class Scheduler(
):
return
prepare_batch_fn = getattr(
self.tree_cache, "prepare_write_backups_for_reqs", None
)
if prepare_batch_fn is None:
inner_cache = getattr(self.tree_cache, "inner", None)
prepare_batch_fn = getattr(
inner_cache, "prepare_write_backups_for_reqs", None
)
if prepare_batch_fn is not None:
prepare_batch_fn(batch.reqs)
return
prepare_fn = getattr(self.tree_cache, "prepare_write_backup_for_req", None)
if prepare_fn is None:
inner_cache = getattr(self.tree_cache, "inner", None)
python/sglang/srt/mem_cache/hiradix_cache.py
+145 -27
View File
@@ -313,6 +313,13 @@ class PreparedCpHiCacheBackup:
attached: bool = False
@dataclass(frozen=True)
class CpWriteBackupCandidate:
req: object
kv_indices: torch.Tensor
node_id: int
@dataclass(frozen=True)
class CpHiCacheCapacitySnapshot:
target_capacity: Tuple[int, ...]
@@ -1407,10 +1414,10 @@ class HiRadixCache(RadixCache):
)
return EvictResult(num_tokens_evicted=num_evicted)
def _cp_build_write_admission(
self, device_indices: torch.Tensor, *, node_id: int, phase: str
def _cp_build_write_admission_from_required(
self, required: Tuple[int, ...], *, node_id: int, phase: str
) -> CpWriteAdmission:
required = self._cp_required_host_tokens_by_rank(device_indices)
required = tuple(int(v) for v in required)
snapshot = self._cp_host_capacity_snapshot()
target_available = self._cp_host_available_tokens_by_rank(snapshot)
draft_available = self._cp_host_available_tokens_by_rank(snapshot, draft=True)
@@ -1418,12 +1425,13 @@ class HiRadixCache(RadixCache):
trigger_ratio = float(
getattr(self, "hicache_host_free_room_trigger_ratio", 0.0) or 0.0
)
page_size = int(getattr(self, "page_size", 1) or 1)
target_deficit = tuple(
_free_room_deficit(
required=req,
available=avail,
capacity=capacity,
page_size=self.page_size,
page_size=page_size,
target_ratio=target_ratio,
trigger_ratio=trigger_ratio,
)
@@ -1438,7 +1446,7 @@ class HiRadixCache(RadixCache):
required=req,
available=avail,
capacity=capacity,
page_size=self.page_size,
page_size=page_size,
target_ratio=target_ratio,
trigger_ratio=trigger_ratio,
)
@@ -1463,6 +1471,16 @@ class HiRadixCache(RadixCache):
eviction_plan=eviction_plan,
)
def _cp_build_write_admission(
self, device_indices: torch.Tensor, *, node_id: int, phase: str
) -> CpWriteAdmission:
required = self._cp_required_host_tokens_by_rank(device_indices)
return self._cp_build_write_admission_from_required(
required,
node_id=node_id,
phase=phase,
)
def shutdown(self):
"""Best-effort auto-detach of storage backend on process shutdown.
@@ -2038,18 +2056,25 @@ class HiRadixCache(RadixCache):
return True
def _reserve_write_cp_indices_no_collective(
self, device_indices: torch.Tensor, node_id: int
self,
device_indices: torch.Tensor,
node_id: int,
*,
admission_checked: bool = False,
):
admission = self._cp_build_write_admission(
device_indices, node_id=node_id, phase="initial"
)
if any(v > 0 for v in admission.deficit_by_owner):
if not self._evict_cp_host_for_write_admission(
admission, node_id=node_id, phase="initial"
):
return HiCacheWriteFailure(
required_host_slots=max(admission.eviction_plan.remaining_deficit)
)
if not admission_checked:
admission = self._cp_build_write_admission(
device_indices, node_id=node_id, phase="initial"
)
if any(v > 0 for v in admission.deficit_by_owner):
if not self._evict_cp_host_for_write_admission(
admission, node_id=node_id, phase="initial"
):
return HiCacheWriteFailure(
required_host_slots=max(
admission.eviction_plan.remaining_deficit
)
)
result = self.cache_controller.reserve_write_cp(
device_indices=device_indices,
@@ -2416,16 +2441,9 @@ class HiRadixCache(RadixCache):
if hasattr(self, "evictable_host_leaves"):
self._update_host_leaf_status(parent)
def prepare_write_backup_for_req(self, req) -> None:
if self.disable or not self._uses_cp_hicache:
return
if self.cache_controller.write_policy == "write_back":
self._warn_cp_hicache_fallback(
"prepare_write_backup_skipped",
"write_back_policy",
rid=getattr(req, "rid", "<unknown>"),
)
return
def _build_prepare_write_backup_candidate_for_req(
self, req
) -> Optional[CpWriteBackupCandidate]:
if getattr(req, "cp_hicache_prepared_backup", None) is not None:
prepared = getattr(req, "cp_hicache_prepared_backup", None)
self._warn_cp_hicache_fallback(
@@ -2485,7 +2503,32 @@ class HiRadixCache(RadixCache):
node_id = TreeNode.counter
TreeNode.counter += 1
result = self._reserve_write_cp_indices_no_collective(kv_indices, node_id)
return CpWriteBackupCandidate(
req=req,
kv_indices=kv_indices,
node_id=node_id,
)
def _submit_prepare_write_backup_candidate(
self,
candidate: CpWriteBackupCandidate,
*,
admission_checked: bool = False,
) -> None:
req = candidate.req
kv_indices = candidate.kv_indices
node_id = candidate.node_id
if admission_checked:
result = self._reserve_write_cp_indices_no_collective(
kv_indices,
node_id,
admission_checked=True,
)
else:
result = self._reserve_write_cp_indices_no_collective(
kv_indices,
node_id,
)
if isinstance(result, HiCacheWriteFailure):
self._warn_cp_hicache_fallback(
"prepare_write_backup_reservation_failed",
@@ -2519,6 +2562,81 @@ class HiRadixCache(RadixCache):
result.metadata.owned_positions.numel(),
)
def prepare_write_backup_for_req(self, req) -> None:
if self.disable or not self._uses_cp_hicache:
return
if self.cache_controller.write_policy == "write_back":
self._warn_cp_hicache_fallback(
"prepare_write_backup_skipped",
"write_back_policy",
rid=getattr(req, "rid", "<unknown>"),
)
return
candidate = self._build_prepare_write_backup_candidate_for_req(req)
if candidate is None:
return
self._submit_prepare_write_backup_candidate(candidate)
def prepare_write_backups_for_reqs(self, reqs) -> None:
if self.disable or not self._uses_cp_hicache:
return
reqs = list(reqs)
if len(reqs) == 0:
return
if self.cache_controller.write_policy == "write_back":
self._warn_cp_hicache_fallback(
"prepare_write_backups_skipped",
"write_back_policy",
batch_size=len(reqs),
)
return
candidates: List[CpWriteBackupCandidate] = []
for req in reqs:
candidate = self._build_prepare_write_backup_candidate_for_req(req)
if candidate is not None:
candidates.append(candidate)
if len(candidates) == 0:
return
admission_checked = False
if len(candidates) > 1:
required = self._cp_zero_counts()
for candidate in candidates:
required = self._cp_add_counts(
required,
self._cp_required_host_tokens_by_rank(candidate.kv_indices),
)
admission = self._cp_build_write_admission_from_required(
required,
node_id=candidates[0].node_id,
phase="batch_prepare",
)
if any(v > 0 for v in admission.deficit_by_owner):
if self._evict_cp_host_for_write_admission(
admission,
node_id=candidates[0].node_id,
phase="batch_prepare",
):
admission_checked = True
else:
self._warn_cp_hicache_fallback(
"prepare_write_backups_batch_admission_failed",
"host_reservation_failed",
batch_size=len(candidates),
required_by_owner=required,
remaining_deficit=admission.eviction_plan.remaining_deficit,
)
else:
admission_checked = True
for candidate in candidates:
self._submit_prepare_write_backup_candidate(
candidate,
admission_checked=admission_checked,
)
def _node_host_len(self, node: TreeNode) -> int:
if self._uses_cp_hicache:
return node.host_len
python/sglang/srt/mem_cache/session_aware_cache.py
+8
View File
@@ -337,6 +337,14 @@ class SessionAwareCache(BasePrefixCache):
return prepare_fn(req)
return None
def prepare_write_backups_for_reqs(self, reqs) -> None:
prepare_fn = getattr(self.inner, "prepare_write_backups_for_reqs", None)
if prepare_fn is not None:
return prepare_fn(reqs)
for req in reqs:
self.prepare_write_backup_for_req(req)
return None
def check_hicache_events(self):
return self.inner.check_hicache_events()
test/registered/unit/mem_cache/test_cp_hicache_metadata.py
+119
View File
@@ -110,6 +110,7 @@ from sglang.srt.mem_cache.hiradix_cache import (
CpHiCacheCapacitySnapshot,
CpHiCacheEvictionPlan,
CpHiCacheNodeMetadata,
CpWriteAdmission,
HiRadixCache,
PendingHiCacheBackup,
PreparedCpHiCacheBackup,
@@ -732,6 +733,110 @@ class TestCpHiCacheFreeRoom(CustomTestCase):
self.assertEqual(admission.draft_available_by_owner, (64, 256))
self.assertEqual(admission.deficit_by_owner, (64, 0))
def test_cp_host_write_batch_admission_uses_aggregate_required(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache.page_size = 64
cache.hicache_host_free_room_ratio = 0.5
cache.hicache_host_free_room_trigger_ratio = 0.25
cache._cp_host_capacity_snapshot = lambda: CpHiCacheCapacitySnapshot(
target_capacity=(256, 256),
draft_capacity=None,
committed_target=(160, 0),
committed_draft=(0, 0),
pending_target=(0, 0),
pending_draft=(0, 0),
)
cache._plan_cp_host_evictions = lambda deficit: CpHiCacheEvictionPlan(
victims=(),
planned_freed=tuple(0 for _ in deficit),
remaining_deficit=tuple(0 for _ in deficit),
)
admission = cache._cp_build_write_admission_from_required(
(128, 0),
node_id=603,
phase="batch_unit",
)
self.assertEqual(admission.required_by_owner, (128, 0))
self.assertEqual(admission.target_available_by_owner, (96, 256))
# required=128, target_room=128, available=96.
self.assertEqual(admission.deficit_by_owner, (160, 0))
def test_prepare_write_backups_for_reqs_runs_one_batch_admission(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache.disable = False
cache._uses_cp_hicache = True
cache.is_eagle = False
cache.page_size = 1
cache.req_to_token_pool = types.SimpleNamespace(
req_to_token=torch.arange(16, dtype=torch.int64).view(2, 8)
)
cache.cache_controller = FakeReserveWriteController(
[
lambda device_indices, node_id: make_write_reservation(
device_indices, node_id=node_id, host_start=700
),
lambda device_indices, node_id: make_write_reservation(
device_indices, node_id=node_id, host_start=800
),
]
)
cache._probe_existing_radix_prefix_len_no_split = lambda key: 0
cache._cp_zero_counts = lambda cp_size=None: (0,)
cache._cp_add_counts = lambda lhs, rhs: (lhs[0] + rhs[0],)
cache._cp_required_host_tokens_by_rank = lambda indices: (int(indices.numel()),)
batch_admissions = []
evicted_admissions = []
def build_batch_admission(required, *, node_id, phase):
batch_admissions.append((required, node_id, phase))
return CpWriteAdmission(
node_id=node_id,
phase=phase,
required_by_owner=required,
target_available_by_owner=(0,),
draft_available_by_owner=(0,),
deficit_by_owner=(1,),
eviction_plan=CpHiCacheEvictionPlan(
victims=(),
planned_freed=(1,),
remaining_deficit=(0,),
),
)
cache._cp_build_write_admission_from_required = build_batch_admission
cache._evict_cp_host_for_write_admission = (
lambda admission, **_: evicted_admissions.append(admission) or True
)
req0 = types.SimpleNamespace(
rid="batch-rid-0",
fill_ids=list(range(4)),
cache_protected_len=0,
req_pool_idx=0,
is_chunked=0,
cp_hicache_prepared_backup=None,
)
req1 = types.SimpleNamespace(
rid="batch-rid-1",
fill_ids=list(range(6)),
cache_protected_len=0,
req_pool_idx=1,
is_chunked=0,
cp_hicache_prepared_backup=None,
)
cache.prepare_write_backups_for_reqs([req0, req1])
self.assertEqual(len(batch_admissions), 1)
self.assertEqual(batch_admissions[0][0], (10,))
self.assertEqual(batch_admissions[0][2], "batch_prepare")
self.assertEqual(len(evicted_admissions), 1)
self.assertIsNotNone(req0.cp_hicache_prepared_backup)
self.assertIsNotNone(req1.cp_hicache_prepared_backup)
self.assertEqual(len(cache.cache_controller.submitted), 2)
class FakeWriteFailure:
metadata = None
@@ -1030,6 +1135,20 @@ class TestHiRadixCacheCPBackup(CustomTestCase):
self.assertEqual(calls, [req])
def test_session_aware_cache_forwards_cp_hicache_batch_prepare(self):
calls = []
class Inner:
def prepare_write_backups_for_reqs(self, reqs):
calls.append(list(reqs))
wrapper = SessionAwareCache(Inner())
reqs = [object(), object()]
wrapper.prepare_write_backups_for_reqs(reqs)
self.assertEqual(calls, [reqs])
def test_node_backuped_uses_cp_metadata(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache._uses_cp_hicache = True