From 367dff06f3120cb686d3ece6ea6d720fa8aee6e8 Mon Sep 17 00:00:00 2001 From: laoyao0822 Date: Wed, 27 May 2026 05:46:34 +0800 Subject: [PATCH] Keep CP HiCache draft KV invisible until joint readiness CP HiCache now treats draft KV as a strict target-owned payload through pending write visibility, host eviction, and state-buffer registration. Host metadata created before async D2H ack is no longer request-visible, so match_prefix cannot select an in-flight host node. Draft host eviction now fails before target cleanup when draft metadata is missing, and prefill/decode share one helper for draft NSA state buffers so shared-KV mode cannot silently skip mismatched draft state. Constraint: CP shared KV + HiCache + EAGLE/MTP must not expose target-only host hits or skipped draft state as valid cache hits Rejected: Rely on event-loop ordering and lock_ref to hide in-flight writes | match_prefix does not consult lock_ref and can observe host_len/cp_hicache directly Rejected: Keep draft state mismatch as debug-only skip | it can poison speculative acceptance while looking like a successful cache hit Confidence: high Scope-risk: moderate Directive: Do not reintroduce silent draft/target fallback in CP shared-KV HiCache paths; malformed strong-sync metadata should fail fast Tested: python -m py_compile targeted modified files Tested: remote g0034 container pytest test/registered/unit/mem_cache/test_cp_hicache_metadata.py test/registered/unit/managers/test_hicache_controller_cp.py test/registered/unit/disaggregation/test_cp_shared_kv_transfer_mapping.py -q (91 passed) Not-tested: Full CP shared KV + HiCache + EAGLE/MTP ETE server run after this commit --- ...ll_cp_hicache_draft_kv_strong_sync_plan.md | 33 +++++ python/sglang/srt/disaggregation/decode.py | 39 +++--- python/sglang/srt/disaggregation/prefill.py | 39 +++--- python/sglang/srt/disaggregation/utils.py | 46 ++++++- .../sglang/srt/managers/cache_controller.py | 10 +- python/sglang/srt/mem_cache/hiradix_cache.py | 14 +- .../test_cp_shared_kv_transfer_mapping.py | 88 ++++++++++++ .../managers/test_hicache_controller_cp.py | 127 ++++++++++++++++++ .../mem_cache/test_cp_hicache_metadata.py | 94 +++++++++++++ 9 files changed, 448 insertions(+), 42 deletions(-) diff --git a/docs/advanced_features/nsa_prefill_cp_hicache_draft_kv_strong_sync_plan.md b/docs/advanced_features/nsa_prefill_cp_hicache_draft_kv_strong_sync_plan.md index 77cdfb59d..3eb04fdf6 100644 --- a/docs/advanced_features/nsa_prefill_cp_hicache_draft_kv_strong_sync_plan.md +++ b/docs/advanced_features/nsa_prefill_cp_hicache_draft_kv_strong_sync_plan.md @@ -35,6 +35,14 @@ eviction, or visibility decisions. zero-owned rank handling and frees the replay allocation on failure. - Phase 3 implemented: `HiRadixCache.reset()` clears target and draft host pools as one logical reset. +- Phase 5 implemented: CP host-backed predicate now treats in-flight + write-through nodes as pending, not request-visible host hits; duplicate + write-through launch is suppressed while a node is in + `ongoing_write_through`. +- Phase 6 implemented: controller tests now cover target+draft host eviction, + zero-owned draft metadata, and draft-before-target load readiness; CP draft + state-buffer registration is centralized and fails fast on target/draft state + mismatch when `SGLANG_CP_DRAFT_SHARED_KV=1`. - Adjacent controller invariant guard restored for CPU/fake-test tensors only: page-shaped host/device indices are validated in unit tests without reintroducing CUDA tensor truth-value checks or hot-path host sync. @@ -574,6 +582,14 @@ the explicit-ready path before per-layer backup. Pending write is a valid lifecycle state, so it may return “not ready”; malformed metadata must still raise. +Implemented decision: use the explicit-ready path. `_node_backuped()` still +validates CP/draft metadata first, but returns `False` while the node id is in +`ongoing_write_through`. This keeps in-flight D2H metadata out of request-visible +host-hit admission until `writing_check()` observes the joint target+draft ack +and removes the node from the pending table. `_inc_hit_count()` treats pending +write-through as “already launched” and does not start a second write for the +same node. + ### Phase 6: Keep internal transfer split, but expose one logical operation Files: @@ -615,6 +631,23 @@ Add controller tests: an error. It should not silently skip draft state buffers because that creates a target/draft cache-hit mismatch. +Implemented status: + +- `evict_cp_host()` fails before freeing target host slots if draft HiCache is + attached but metadata lacks `draft_host_indices`; otherwise target and draft + host slots are freed as one logical operation. +- `_write_cp()` zero-owned ranks keep explicit empty `draft_host_indices`, so + zero-owned target metadata is still draft-valid. +- `start_loading()` is test-protected to submit draft H2D before completing the + corresponding target layer readiness event. +- `append_cp_draft_state_buffers()` is the single helper used by prefill and + decode queue setup. It appends draft NSA state buffers to the target state + stream, preserves legacy skip behavior when draft shared KV is disabled, and + raises on target/draft state mismatch when draft shared KV is enabled. +- Code-read confirmation: PD transfer page indices remain target-driven; + draft KV buffers are extra buffers behind the same page-index stream, not an + independent draft page-index stream. + ### Phase 7: Remote validation Run targeted tests in the remote container because this repo's runtime deps and diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 09bcf81ac..9d03c645c 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -41,6 +41,7 @@ from sglang.srt.disaggregation.utils import ( MetadataBuffers, ReqToMetadataIdxAllocator, TransferBackend, + append_cp_draft_state_buffers, get_kv_class, is_mla_backend, poll_and_all_reduce, @@ -458,25 +459,25 @@ class DecodePreallocQueue: draft_state_item_lens, ) = _state_buf_infos(self.draft_token_to_kv_pool) - kv_args.draft_state_type = draft_state_type - kv_args.draft_state_buffer_start = len(kv_args.state_data_ptrs) - kv_args.draft_state_buffer_count = 0 - if draft_state_data_ptrs: - if kv_args.state_type == "nsa" and draft_state_type == "nsa": - kv_args.state_data_ptrs += draft_state_data_ptrs - kv_args.state_data_lens += draft_state_data_lens - kv_args.state_item_lens += draft_state_item_lens - kv_args.draft_state_buffer_count = len(draft_state_data_ptrs) - else: - _cp_draft_shared_kv_debug( - "decode_draft_state_skipped cp_rank=%s target_state_type=%s " - "draft_state_type=%s draft_state_bufs=%s " - "reason=unsupported_state_type", - self.tp_rank, - kv_args.state_type, - draft_state_type, - len(draft_state_data_ptrs), - ) + draft_state_registered = append_cp_draft_state_buffers( + kv_args, + draft_state_type, + draft_state_data_ptrs, + draft_state_data_lens, + draft_state_item_lens, + role="decode", + cp_rank=self.tp_rank, + ) + if draft_state_data_ptrs and not draft_state_registered: + _cp_draft_shared_kv_debug( + "decode_draft_state_skipped cp_rank=%s target_state_type=%s " + "draft_state_type=%s draft_state_bufs=%s " + "reason=unsupported_state_type", + self.tp_rank, + kv_args.state_type, + draft_state_type, + len(draft_state_data_ptrs), + ) if envs.SGLANG_CP_DRAFT_SHARED_KV_DEBUG.get(): _cp_draft_shared_kv_debug( diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index d3dc71721..a972c7d92 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -36,6 +36,7 @@ from sglang.srt.disaggregation.utils import ( MetadataBuffers, ReqToMetadataIdxAllocator, TransferBackend, + append_cp_draft_state_buffers, get_kv_class, is_mla_backend, kv_to_page_num, @@ -299,25 +300,25 @@ class PrefillBootstrapQueue: draft_state_item_lens, ) = _state_buf_infos(self.draft_token_to_kv_pool) - kv_args.draft_state_type = draft_state_type - kv_args.draft_state_buffer_start = len(kv_args.state_data_ptrs) - kv_args.draft_state_buffer_count = 0 - if draft_state_data_ptrs: - if kv_args.state_type == "nsa" and draft_state_type == "nsa": - kv_args.state_data_ptrs += draft_state_data_ptrs - kv_args.state_data_lens += draft_state_data_lens - kv_args.state_item_lens += draft_state_item_lens - kv_args.draft_state_buffer_count = len(draft_state_data_ptrs) - else: - _cp_draft_shared_kv_debug( - "prefill_draft_state_skipped cp_rank=%s target_state_type=%s " - "draft_state_type=%s draft_state_bufs=%s " - "reason=unsupported_state_type", - self.tp_rank, - kv_args.state_type, - draft_state_type, - len(draft_state_data_ptrs), - ) + draft_state_registered = append_cp_draft_state_buffers( + kv_args, + draft_state_type, + draft_state_data_ptrs, + draft_state_data_lens, + draft_state_item_lens, + role="prefill", + cp_rank=self.tp_rank, + ) + if draft_state_data_ptrs and not draft_state_registered: + _cp_draft_shared_kv_debug( + "prefill_draft_state_skipped cp_rank=%s target_state_type=%s " + "draft_state_type=%s draft_state_bufs=%s " + "reason=unsupported_state_type", + self.tp_rank, + kv_args.state_type, + draft_state_type, + len(draft_state_data_ptrs), + ) if envs.SGLANG_CP_DRAFT_SHARED_KV_DEBUG.get(): _cp_draft_shared_kv_debug( diff --git a/python/sglang/srt/disaggregation/utils.py b/python/sglang/srt/disaggregation/utils.py index 9a2b545e4..4552a0090 100644 --- a/python/sglang/srt/disaggregation/utils.py +++ b/python/sglang/srt/disaggregation/utils.py @@ -5,7 +5,7 @@ import random from collections import deque from contextlib import nullcontext from enum import Enum -from typing import TYPE_CHECKING, Literal, Optional, Tuple, Type, overload +from typing import TYPE_CHECKING, Any, Literal, Optional, Tuple, Type, overload import numpy as np import torch @@ -85,6 +85,50 @@ def poll_and_all_reduce_attn_cp_tp_group( ######################### +def append_cp_draft_state_buffers( + kv_args: Any, + draft_state_type: str, + draft_state_data_ptrs, + draft_state_data_lens, + draft_state_item_lens, + *, + role: str, + cp_rank: int, +) -> bool: + """Append draft NSA state buffers as payload attached to target state buffers. + + Under CP draft shared KV, draft state is not allowed to silently diverge + from target state. If draft exposes state buffers, target and draft must both + be NSA state so the transfer still has one target-driven state stream. + """ + + target_state_type = getattr(kv_args, "state_type", "none") + draft_state_bufs = len(draft_state_data_ptrs) + can_append = target_state_type == "nsa" and draft_state_type == "nsa" + + if draft_state_bufs and not can_append: + if envs.SGLANG_CP_DRAFT_SHARED_KV.get(): + raise RuntimeError( + "CP draft shared KV state mismatch during " + f"{role}: cp_rank={cp_rank} target_state_type={target_state_type} " + f"draft_state_type={draft_state_type} " + f"draft_state_bufs={draft_state_bufs}" + ) + + kv_args.draft_state_type = draft_state_type + kv_args.draft_state_buffer_start = len(kv_args.state_data_ptrs) + kv_args.draft_state_buffer_count = 0 + + if not draft_state_bufs or not can_append: + return False + + kv_args.state_data_ptrs += draft_state_data_ptrs + kv_args.state_data_lens += draft_state_data_lens + kv_args.state_item_lens += draft_state_item_lens + kv_args.draft_state_buffer_count = draft_state_bufs + return True + + class ReqToMetadataIdxAllocator: """A memory pool that maps a request to its first output token location.""" diff --git a/python/sglang/srt/managers/cache_controller.py b/python/sglang/srt/managers/cache_controller.py index 1e2d6e220..a531626d1 100644 --- a/python/sglang/srt/managers/cache_controller.py +++ b/python/sglang/srt/managers/cache_controller.py @@ -1261,9 +1261,15 @@ class HiCacheController: if not backup_only: raise ValueError("Other eviction policies are not supported yet.") - freed = self.evict_host(metadata.host_indices, backup_only=backup_only) draft_host_indices = getattr(metadata, "draft_host_indices", None) - if self.has_draft_hicache and draft_host_indices is not None: + if self.has_draft_hicache and draft_host_indices is None: + raise RuntimeError( + "CP HiCache draft KV host evict requested but node metadata " + "does not contain draft_host_indices" + ) + + freed = self.evict_host(metadata.host_indices, backup_only=backup_only) + if self.has_draft_hicache: self.draft_mem_pool_host.free(draft_host_indices) return freed diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 3792502b7..e89112c02 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -944,11 +944,20 @@ class HiRadixCache(RadixCache): ) return True + def _node_host_write_pending(self, node: TreeNode) -> bool: + return getattr(node, "id", None) in getattr(self, "ongoing_write_through", {}) + + def _node_host_write_ready(self, node: TreeNode) -> bool: + return not self._node_host_write_pending(node) + def _node_backuped(self, node: TreeNode) -> bool: if self._uses_cp_hicache: if node.host_len <= 0: return False - return self._node_has_required_draft_hicache(node) + return ( + self._node_has_required_draft_hicache(node) + and self._node_host_write_ready(node) + ) return node.host_value is not None def _node_host_len(self, node: TreeNode) -> int: @@ -1077,6 +1086,9 @@ class HiRadixCache(RadixCache): return node.hit_count += 1 + if self._uses_cp_hicache and self._node_host_write_pending(node): + return + if not self._node_backuped(node): if node.hit_count >= self.write_through_threshold: logger.info( diff --git a/test/registered/unit/disaggregation/test_cp_shared_kv_transfer_mapping.py b/test/registered/unit/disaggregation/test_cp_shared_kv_transfer_mapping.py index 9379230f2..bd7e78f53 100644 --- a/test/registered/unit/disaggregation/test_cp_shared_kv_transfer_mapping.py +++ b/test/registered/unit/disaggregation/test_cp_shared_kv_transfer_mapping.py @@ -1,8 +1,11 @@ import unittest +from types import SimpleNamespace +from unittest.mock import patch import numpy as np from sglang.srt.disaggregation.utils import ( + append_cp_draft_state_buffers, filter_kv_pages_for_cp_shared_kv, select_pages_by_request_positions, ) @@ -56,6 +59,91 @@ class TestCPSharedKVTransferMapping(unittest.TestCase): self.assertEqual(logical_positions.tolist(), [2, 6]) self.assertEqual(selected_decode_pages.tolist(), [43, 47]) + def test_append_cp_draft_state_buffers_registers_nsa_state_as_target_payload(self): + kv_args = SimpleNamespace( + state_type="nsa", + state_data_ptrs=[11], + state_data_lens=[12], + state_item_lens=[13], + ) + + appended = append_cp_draft_state_buffers( + kv_args, + draft_state_type="nsa", + draft_state_data_ptrs=[21, 22], + draft_state_data_lens=[23, 24], + draft_state_item_lens=[25, 26], + role="prefill", + cp_rank=2, + ) + + self.assertTrue(appended) + self.assertEqual(kv_args.draft_state_type, "nsa") + self.assertEqual(kv_args.draft_state_buffer_start, 1) + self.assertEqual(kv_args.draft_state_buffer_count, 2) + self.assertEqual(kv_args.state_data_ptrs, [11, 21, 22]) + self.assertEqual(kv_args.state_data_lens, [12, 23, 24]) + self.assertEqual(kv_args.state_item_lens, [13, 25, 26]) + + def test_append_cp_draft_state_buffers_rejects_state_mismatch_under_shared_kv(self): + kv_args = SimpleNamespace( + state_type="nsa", + state_data_ptrs=[11], + state_data_lens=[12], + state_item_lens=[13], + ) + + with patch( + "sglang.srt.disaggregation.utils.envs.SGLANG_CP_DRAFT_SHARED_KV.get", + return_value=True, + ): + with self.assertRaisesRegex( + RuntimeError, + "prefill.*cp_rank=2.*target_state_type=nsa.*draft_state_type=mamba", + ): + append_cp_draft_state_buffers( + kv_args, + draft_state_type="mamba", + draft_state_data_ptrs=[21], + draft_state_data_lens=[22], + draft_state_item_lens=[23], + role="prefill", + cp_rank=2, + ) + + self.assertFalse(hasattr(kv_args, "draft_state_buffer_count")) + self.assertEqual(kv_args.state_data_ptrs, [11]) + + def test_append_cp_draft_state_buffers_keeps_legacy_skip_when_shared_kv_disabled( + self, + ): + kv_args = SimpleNamespace( + state_type="nsa", + state_data_ptrs=[11], + state_data_lens=[12], + state_item_lens=[13], + ) + + with patch( + "sglang.srt.disaggregation.utils.envs.SGLANG_CP_DRAFT_SHARED_KV.get", + return_value=False, + ): + appended = append_cp_draft_state_buffers( + kv_args, + draft_state_type="mamba", + draft_state_data_ptrs=[21], + draft_state_data_lens=[22], + draft_state_item_lens=[23], + role="decode", + cp_rank=3, + ) + + self.assertFalse(appended) + self.assertEqual(kv_args.draft_state_type, "mamba") + self.assertEqual(kv_args.draft_state_buffer_start, 1) + self.assertEqual(kv_args.draft_state_buffer_count, 0) + self.assertEqual(kv_args.state_data_ptrs, [11]) + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/managers/test_hicache_controller_cp.py b/test/registered/unit/managers/test_hicache_controller_cp.py index ed46c88f2..473b38dee 100644 --- a/test/registered/unit/managers/test_hicache_controller_cp.py +++ b/test/registered/unit/managers/test_hicache_controller_cp.py @@ -251,6 +251,24 @@ class DummyLayerDoneCounter: return 0 +class RecordingProducerEvent: + def __init__(self, order): + self.start_event = DummyEvent() + self.finish_event = DummyEvent() + self.order = order + + def complete(self, layer_id): + self.order.append(("complete", layer_id)) + + +class RecordingLayerDoneCounter: + def __init__(self, order): + self.events = [RecordingProducerEvent(order)] + + def update_producer(self): + return 0 + + class TestHiCacheControllerCPWrite(CustomTestCase): def setUp(self): self.device_module_patcher = patch( @@ -344,6 +362,26 @@ class TestHiCacheControllerCPWrite(CustomTestCase): self.assertEqual(host_pool.alloc_calls, []) self.assertEqual(len(controller.ack_write_queue), 1) + def test_cp_write_zero_owned_with_draft_returns_empty_draft_metadata(self): + host_pool = FakeHostPool(torch.tensor([], dtype=torch.int64)) + draft_host_pool = FakeHostPool(torch.tensor([], dtype=torch.int64)) + controller = self.make_controller( + host_pool, + cp_rank=3, + draft_host_pool=draft_host_pool, + draft_mem_pool_device=FakeDevicePool("draft"), + ) + logical_locs = torch.arange(4, 8, dtype=torch.int64) + + result = controller.write(logical_locs, node_id=18) + + self.assertEqual(result.metadata.logical_len, 4) + self.assertEqual(result.metadata.host_indices.tolist(), []) + self.assertEqual(result.metadata.draft_host_indices.tolist(), []) + self.assertEqual(host_pool.alloc_calls, []) + self.assertEqual(draft_host_pool.alloc_calls, []) + self.assertEqual(len(controller.ack_write_queue), 1) + def test_cp_write_allocation_failure_reports_required_host_slots(self): host_pool = FakeHostPool(None) controller = self.make_controller(host_pool, cp_rank=1) @@ -519,6 +557,50 @@ class TestHiCacheControllerCPLoad(TestHiCacheControllerCPWrite): self.assertEqual(len(controller.ack_load_queue), 1) self.assertEqual(controller.ack_load_queue[0].node_ids, [14]) + def test_cp_start_loading_loads_draft_before_target_layer_ready(self): + order = [] + host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64)) + draft_host_pool = FakeHostPool( + torch.tensor([200, 201, 202, 203], dtype=torch.int64) + ) + allocator = FakeAllocator(alloc_result=torch.arange(64, 80, dtype=torch.int64)) + controller = self.make_controller( + host_pool, + allocator=allocator, + cp_rank=1, + draft_host_pool=draft_host_pool, + draft_mem_pool_device=FakeDevicePool("draft"), + ) + controller.layer_done_counter = RecordingLayerDoneCounter(order) + + def record_target_load( + device_pool, host_indices, device_indices, layer_id, io_backend + ): + order.append(("target", layer_id)) + + def record_draft_load( + device_pool, host_indices, device_indices, layer_id, io_backend + ): + order.append(("draft", layer_id)) + + host_pool.load_to_device_per_layer = record_target_load + draft_host_pool.load_to_device_per_layer = record_draft_load + 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), + draft_host_indices=torch.tensor([200, 201, 202, 203], dtype=torch.int64), + page_owners=torch.tensor([3, 0, 1, 2], dtype=torch.int8), + page_size=4, + ) + + controller.load_cp([node], node_id=114) + controller.start_loading() + + self.assertEqual(order, [("draft", 0), ("target", 0), ("complete", 0)]) + def test_cp_load_zero_owned_returns_full_logical_locs_and_noop_ack(self): host_pool = FakeHostPool(torch.tensor([], dtype=torch.int64)) allocator = FakeAllocator(alloc_result=torch.arange(64, 68, dtype=torch.int64)) @@ -634,6 +716,51 @@ class TestHiCacheControllerCPLoad(TestHiCacheControllerCPWrite): with self.assertRaisesRegex(ValueError, "host_indices.*contiguous page spans"): controller.load_cp([node], node_id=32) + def test_cp_evict_host_frees_target_and_draft_host_indices(self): + host_pool = FakeHostPool(torch.tensor([], dtype=torch.int64)) + draft_host_pool = FakeHostPool(torch.tensor([], dtype=torch.int64)) + controller = self.make_controller( + host_pool, + draft_host_pool=draft_host_pool, + draft_mem_pool_device=FakeDevicePool("draft"), + ) + metadata = CpHiCacheNodeMetadata( + logical_len=16, + owned_positions=torch.tensor([0, 1, 2, 3], 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.tensor([0, 1, 2, 3], dtype=torch.int8), + page_size=4, + ) + + freed = controller.evict_cp_host(metadata) + + self.assertEqual(freed, 4) + self.assertEqual(host_pool.frees[0].tolist(), [100, 101, 102, 103]) + self.assertEqual(draft_host_pool.frees[0].tolist(), [200, 201, 202, 203]) + + def test_cp_evict_host_rejects_missing_draft_metadata_before_target_free(self): + host_pool = FakeHostPool(torch.tensor([], dtype=torch.int64)) + draft_host_pool = FakeHostPool(torch.tensor([], dtype=torch.int64)) + controller = self.make_controller( + host_pool, + draft_host_pool=draft_host_pool, + draft_mem_pool_device=FakeDevicePool("draft"), + ) + metadata = CpHiCacheNodeMetadata( + logical_len=16, + owned_positions=torch.tensor([0, 1, 2, 3], dtype=torch.int64), + host_indices=torch.tensor([100, 101, 102, 103], dtype=torch.int64), + page_owners=torch.tensor([0, 1, 2, 3], dtype=torch.int8), + page_size=4, + ) + + with self.assertRaisesRegex(RuntimeError, "draft.*evict.*draft_host_indices"): + controller.evict_cp_host(metadata) + + self.assertEqual(host_pool.frees, []) + self.assertEqual(draft_host_pool.frees, []) + if __name__ == "__main__": main() 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 21d0b4a74..4ca621fa2 100644 --- a/test/registered/unit/mem_cache/test_cp_hicache_metadata.py +++ b/test/registered/unit/mem_cache/test_cp_hicache_metadata.py @@ -620,6 +620,29 @@ class TestHiRadixCacheCPBackup(CustomTestCase): self.assertTrue(cache._node_backuped(node)) + def test_node_backuped_excludes_inflight_cp_write_until_ack(self): + cache = HiRadixCache.__new__(HiRadixCache) + cache._uses_cp_hicache = True + cache.cache_controller = types.SimpleNamespace(has_draft_hicache=True) + cache.ongoing_write_through = {} + node = TreeNode() + node.id = 125 + node.host_len = 64 + node.cp_hicache = CpHiCacheNodeMetadata( + logical_len=64, + owned_positions=torch.empty((0,), dtype=torch.int64), + host_indices=torch.empty((0,), dtype=torch.int64), + draft_host_indices=torch.empty((0,), dtype=torch.int64), + page_owners=torch.tensor([0], dtype=torch.int8), + page_size=64, + ) + + cache.ongoing_write_through[node.id] = node + self.assertFalse(cache._node_backuped(node)) + + cache.ongoing_write_through.clear() + self.assertTrue(cache._node_backuped(node)) + def test_single_node_write_lock_updates_device_evictable_leaf_set(self): cache = HiRadixCache.__new__(HiRadixCache) cache.disable = False @@ -672,6 +695,34 @@ class TestHiRadixCacheCPBackup(CustomTestCase): self.assertEqual(node.hit_count, 1) + def test_inc_hit_count_does_not_duplicate_inflight_cp_write(self): + cache = HiRadixCache.__new__(HiRadixCache) + cache._uses_cp_hicache = True + cache.write_through_threshold = 1 + cache.ongoing_write_through = {} + cache.cache_controller = type( + "Controller", (), {"write_policy": "write_through"} + )() + cache.write_backup = lambda node: (_ for _ in ()).throw( + AssertionError("must not launch a second write") + ) + node = TreeNode() + node.id = 127 + node.host_len = 4 + node.cp_hicache = CpHiCacheNodeMetadata( + logical_len=4, + owned_positions=torch.tensor([], dtype=torch.int64), + host_indices=torch.tensor([], dtype=torch.int64), + draft_host_indices=torch.tensor([], dtype=torch.int64), + page_owners=torch.zeros(max(4, 0), dtype=torch.int8), + page_size=1, + ) + cache.ongoing_write_through[node.id] = node + + cache._inc_hit_count(node) + + self.assertEqual(node.hit_count, 1) + def test_write_backup_retries_by_required_physical_slots(self): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = True @@ -1217,6 +1268,49 @@ class TestHiRadixCacheCPLoadBack(CustomTestCase): self.assertIs(result.last_host_node, node) self.assertIs(result.last_device_node, root) + def test_cp_match_prefix_does_not_admit_inflight_write_as_host_hit(self): + cache = HiRadixCache.__new__(HiRadixCache) + cache._uses_cp_hicache = True + cache.device = "cpu" + cache.disable = False + cache.page_size = 1 + cache.ongoing_write_through = {} + cache.cache_controller = types.SimpleNamespace(has_draft_hicache=True) + cache.get_child_key_fn = lambda key: key.token_ids[0] + cache.key_match_fn = lambda child_key, key: sum( + 1 for lhs, rhs in zip(child_key.token_ids, key.token_ids) if lhs == rhs + ) + cache.maybe_bigram_convert = lambda key: (key, None) + root = TreeNode() + root.key = RadixKey([]) + root.value = torch.empty((0,), dtype=torch.int64) + root.host_len = 0 + cache.root_node = root + node = TreeNode() + node.id = 126 + node.parent = root + node.key = RadixKey(list(range(8))) + node.value = torch.arange(8, dtype=torch.int64) + node.host_value = None + node.host_len = 8 + node.cp_hicache = CpHiCacheNodeMetadata( + logical_len=8, + owned_positions=torch.tensor([0, 1], dtype=torch.int64), + host_indices=torch.tensor([50, 51], dtype=torch.int64), + draft_host_indices=torch.tensor([150, 151], dtype=torch.int64), + page_owners=torch.zeros(max(8, 0), dtype=torch.int8), + page_size=1, + ) + root.children[0] = node + cache.ongoing_write_through[node.id] = node + + result = cache.match_prefix(MatchPrefixParams(key=RadixKey(list(range(8))))) + + self.assertEqual(result.device_indices.tolist(), list(range(8))) + self.assertEqual(result.host_hit_length, 0) + self.assertIs(result.last_device_node, node) + self.assertIs(result.last_host_node, root) + def test_cp_match_prefix_shorter_than_page_returns_empty_root_match(self): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = True