diff --git a/docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_parallel_workstreams_zh.md b/docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_parallel_workstreams_zh.md index b747a1f42..9ac67d172 100644 --- a/docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_parallel_workstreams_zh.md +++ b/docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_parallel_workstreams_zh.md @@ -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 +``` diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 293eb5dfc..431e7f49e 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -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) diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 58cf7187a..d085022f5 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -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", ""), - ) - 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", ""), + ) + 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 diff --git a/python/sglang/srt/mem_cache/session_aware_cache.py b/python/sglang/srt/mem_cache/session_aware_cache.py index edc22775d..da6721e84 100644 --- a/python/sglang/srt/mem_cache/session_aware_cache.py +++ b/python/sglang/srt/mem_cache/session_aware_cache.py @@ -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() diff --git a/test/registered/unit/mem_cache/test_cp_hicache_metadata.py b/test/registered/unit/mem_cache/test_cp_hicache_metadata.py index ab2adf7ff..dbe5b2f2a 100644 --- a/test/registered/unit/mem_cache/test_cp_hicache_metadata.py +++ b/test/registered/unit/mem_cache/test_cp_hicache_metadata.py @@ -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