diff --git a/python/sglang/srt/mem_cache/unified_cache_components/swa_component.py b/python/sglang/srt/mem_cache/unified_cache_components/swa_component.py index 6c1c52c8b..1f2a68f40 100644 --- a/python/sglang/srt/mem_cache/unified_cache_components/swa_component.py +++ b/python/sglang/srt/mem_cache/unified_cache_components/swa_component.py @@ -594,14 +594,13 @@ class SWAComponent(TreeComponent): ] if phase == CacheTransferPhase.PREFETCH: - num_pages = min( - prefetch_tokens // self.cache.page_size, - (self.sliding_window_size + self.cache.page_size - 1) - // self.cache.page_size, - ) - if num_pages == 0: + # Require a full sliding window. + sw_pages = ( + self.sliding_window_size + self.cache.page_size - 1 + ) // self.cache.page_size + if sw_pages == 0 or prefetch_tokens // self.cache.page_size < sw_pages: return None - num_tokens = num_pages * self.cache.page_size + num_tokens = sw_pages * self.cache.page_size host_indices = self._swa_kv_pool_host.alloc(num_tokens) if host_indices is None: self.cache.evict_host(num_tokens, ComponentType.SWA) @@ -612,7 +611,7 @@ class SWAComponent(TreeComponent): PoolTransfer( name=PoolName.SWA, host_indices=host_indices, - keys=["__placeholder__"] * num_pages, + keys=["__placeholder__"] * sw_pages, hit_policy=PoolHitPolicy.TRAILING_PAGES, ) ] @@ -694,33 +693,38 @@ class SWAComponent(TreeComponent): insert_result: Optional[InsertResult] = None, pool_storage_result: Optional[PoolTransferResult] = None, ) -> None: - """Distribute the prefetched SWA buffer onto the leaf→anchor path. + """Fill the prefetched SWA window onto the leaf→anchor path. - The buffer holds the trailing ``loaded_pages`` of the completed KV - prefix, mapped to token range ``[loaded_start, total_len)``. We walk - upward from ``inserted_host_node`` to ``anchor`` and, for each node - whose token range overlaps the buffer: - - SWA tombstone (host_value is None) → fill from buffer (split if - the node only partially overlaps at the buffer's left edge) - - already has SWA host_value → release the corresponding slice - Any leftover buffer beyond the walked range is also released. + All-or-nothing over one full window: ``loaded_pages`` is the cross-rank + MIN, so ``loaded_pages < window_pages`` drops the whole window (keeps the + tree identical across TP ranks). Otherwise map the buffer to token range + ``[loaded_start, total_len)`` and walk leaf→anchor, filling SWA + tombstones and releasing slices that already have host_value. """ if not transfers: return ct = self.component_type + page_size = self.cache.page_size host_indices = transfers[0].host_indices + window_require_pages = ( + host_indices.numel() // page_size if host_indices is not None else 0 + ) loaded_pages = ( pool_storage_result.extra_pool_hit_pages.get(PoolName.SWA, 0) if pool_storage_result else 0 ) target = insert_result.inserted_host_node if insert_result else None - if not loaded_pages or target is None: + if ( + target is None + or window_require_pages == 0 + or loaded_pages < window_require_pages + ): self._release_swa_host(host_indices) return # Buffer covers token range [loaded_start, total_len). - loaded_start = insert_result.total_len - loaded_pages * self.cache.page_size + loaded_start = insert_result.total_len - window_require_pages * page_size # Walk leaf → anchor; ``pos`` is the right edge of ``cur`` in tokens. pos, cur = insert_result.total_len, target diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index 5017ab3af..e4a6ca264 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -1890,16 +1890,23 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): operation ) min_completed_tokens = completed_tokens + hit_pages = operation.pool_storage_result.extra_pool_hit_pages if self.tp_world_size > 1: - completed_tokens_tensor = torch.tensor( - min_completed_tokens, dtype=torch.int + # Reduce full completed tokens together with the sidecar pools that + # this prefetch actually transferred, in one all_reduce. + sidecar_pools = [t.name for xfers in comp_xfers.values() for t in xfers] + packed = torch.tensor( + [completed_tokens] + [hit_pages.get(p, 0) for p in sidecar_pools], + dtype=torch.int, ) torch.distributed.all_reduce( - completed_tokens_tensor, + packed, op=torch.distributed.ReduceOp.MIN, group=self.tp_group, ) - min_completed_tokens = int(completed_tokens_tensor.item()) + min_completed_tokens = int(packed[0].item()) + for i, p in enumerate(sidecar_pools, start=1): + hit_pages[p] = int(packed[i].item()) fetched_key = prefetch_key[:min_completed_tokens] insert_result = self._insert_helper_host( diff --git a/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_dsv4.py b/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_dsv4.py index d73c68f02..c19b7928f 100644 --- a/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_dsv4.py +++ b/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_dsv4.py @@ -1,5 +1,10 @@ +import os +import shutil +import tempfile import unittest +from test_unified_radix_cache_kl_nightly import AccuracyTwoPassMixin + from sglang.srt.utils import kill_process_tree from sglang.test.ci.ci_register import register_cuda_ci from sglang.test.kits.unified_radix_cache_kit import UnifiedRadixTreeTestMixin @@ -100,5 +105,63 @@ class TestUnifiedDeepSeekV4FlashHiCachePageFirstDirect( hicache_mem_layout = "layer_first" +# ─── DeepSeek V4 Flash + HiCache L3 (file backend) ────────────────────── + + +class TestUnifiedDeepSeekV4FlashHiCacheL3(AccuracyTwoPassMixin, CustomTestCase): + """DeepSeek V4 Flash FP8 + HiCache L3 (file backend) + UnifiedRadixCache.""" + + @classmethod + def setUpClass(cls): + cls.model = DSV4_FLASH_MODEL + cls.base_url = DEFAULT_URL_FOR_TEST + cls.hicache_dir = tempfile.mkdtemp(prefix="hicache_l3_dsv4_") + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DSV4_FLASH_LAUNCH_TIMEOUT, + other_args=[ + "--trust-remote-code", + "--tp-size", + "4", + "--attention-backend", + "compressed", + "--page-size", + "256", + "--chunked-prefill-size", + "8192", + "--mem-fraction-static", + "0.9", + "--disable-shared-experts-fusion", + "--enable-hierarchical-cache", + "--hicache-ratio", + "2", + "--hicache-write-policy", + "write_through", + "--hicache-storage-prefetch-policy", + "wait_complete", + "--hicache-io-backend", + "direct", + "--hicache-mem-layout", + "page_first_direct", + "--hicache-storage-backend", + "file", + "--swa-full-tokens-ratio", + "0.25", + ], + env={ + "SGLANG_DSV4_FP4_EXPERTS": "0", + "SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1", + "SGLANG_HICACHE_FILE_BACKEND_STORAGE_DIR": cls.hicache_dir, + }, + ) + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.process.pid) + if os.path.isdir(cls.hicache_dir): + shutil.rmtree(cls.hicache_dir, ignore_errors=True) + + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py index af971d9af..acb922559 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py +++ b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py @@ -1935,7 +1935,16 @@ class UnifiedRadixCacheSuite: storage_dir = tempfile.mkdtemp() self.addCleanup(shutil.rmtree, storage_dir, ignore_errors=True) - seq = self._make_seq(1, 4) + num_pages = 4 + if self.cfg.has_swa: + # SWA L3 prefetch is all-or-nothing over one full sliding window. + # Keep the generic L3 round trip valid for SWA configs whose window + # is larger than the old fixed 4-page request. + sw_pages = ( + self.cfg.sliding_window_size + self.cfg.page_size - 1 + ) // self.cfg.page_size + num_pages = max(num_pages, sw_pages + 1) + seq = self._make_seq(1, num_pages) # --- Producer tree: fill KV, backup D->H, offload H->L3. --- prod, prod_alloc, prod_rtp = build_fixture(self.cfg) @@ -1985,6 +1994,145 @@ class UnifiedRadixCacheSuite: self.assertTrue(torch.equal(loaded_v, expected_v)) cons.sanity_check() + # ---------- TP consistency for SWA prefetch (all-or-nothing) ---------- + + def _patch_tp_all_reduce(self, tree, drop_swa: bool): + """Fake all_reduce so check_prefetch_progress runs the tp>1 path.""" + import torch.distributed as dist + + min_sizes = [] + + def swa_packed_index(): + # Packed tensor is [completed_tokens, *sidecar_hits]; sidecar order + # matches comp_xfers stored in ongoing_prefetch (one live entry). + for info in tree.ongoing_prefetch.values(): + comp_xfers = info[-1] + names = [t.name for xfers in comp_xfers.values() for t in xfers] + if PoolName.SWA in names: + return 1 + names.index(PoolName.SWA) + return None + + def fake(tensor, op=None, group=None): + if op == dist.ReduceOp.MIN: + min_sizes.append(tensor.numel()) + if drop_swa: + idx = swa_packed_index() + if idx is not None and idx < tensor.numel(): + tensor[idx] = 0 + return None + + p = mock.patch.object(dist, "all_reduce", side_effect=fake) + p.start() + self.addCleanup(p.stop) + return min_sizes + + def _swa_host_on_path(self, tree, seq): + m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq)))) + node = m.last_host_node + while node is not tree.root_node: + if node.component_data[ComponentType.SWA].host_value is not None: + return True + node = node.parent + return False + + def _l3_produce(self, storage_dir, seq): + prod, prod_alloc, prod_rtp = build_fixture(self.cfg) + self._init_hicache( + prod, storage_backend="file", storage_dir=storage_dir, prefetch_threshold=1 + ) + self._insert(prod, prod_alloc, prod_rtp, seq) + leaf = prod.match_prefix( + MatchPrefixParams(key=RadixKey(array("q", seq))) + ).last_device_node + self._backup_node(prod, leaf) + self._write_path_to_l3(prod, leaf) + self._flush_l3_backups(prod) + + def _l3_consumer(self, storage_dir): + cons, _, _ = build_fixture(self.cfg) + self._init_hicache( + cons, storage_backend="file", storage_dir=storage_dir, prefetch_threshold=1 + ) + return cons + + def _consume_prefetch(self, cons, seq, req_id): + cons.prefetch_from_storage(req_id, cons.root_node, array("q", seq), None, None) + self._run_prefetch_to_completion(cons, req_id) + cons.drain_storage_control_queues() + + def _setup_swa_tp_prefetch(self): + """Skip non-SWA fixtures; produce one full SWA window+1 page to L3. + + Returns (storage_dir, seq) or None when this fixture cannot exercise + SWA L3 prefetch (then the caller skips). + """ + if not self.cfg.has_swa or self.cfg.has_mamba: + self.skipTest("SWA-only fixture required") + if self._skip_unsupported_hicache_test(): + return None + sw_pages = ( + self.cfg.sliding_window_size + self.cfg.page_size - 1 + ) // self.cfg.page_size + seq = self._make_seq(1, sw_pages + 1) + + storage_dir = tempfile.mkdtemp() + self.addCleanup(shutil.rmtree, storage_dir, ignore_errors=True) + self._l3_produce(storage_dir, seq) + + # Baseline (single rank) must actually adopt SWA, else the TP assertions + # below would be vacuous -> skip. + base = self._l3_consumer(storage_dir) + self._consume_prefetch(base, seq, "base") + if not self._swa_host_on_path(base, seq): + self.skipTest("fixture does not exercise SWA L3 prefetch") + return storage_dir, seq + + def test_tp_swa_prefetch_dropped_when_peer_misses(self): + """A peer rank missing the SWA window drops the whole prefetch result + on every rank (TP-consistent all-or-nothing).""" + setup = self._setup_swa_tp_prefetch() + if setup is None: + return + storage_dir, seq = setup + + cons = self._l3_consumer(storage_dir) + cons.tp_world_size = 2 + min_sizes = self._patch_tp_all_reduce(cons, drop_swa=True) + self._consume_prefetch(cons, seq, "drop") + + m = cons.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq)))) + self.assertEqual(m.host_hit_length, 0) + self.assertFalse( + self._swa_host_on_path(cons, seq), "SWA must be dropped when a peer misses" + ) + # Full + sidecars must be synced through a packed MIN all_reduce. The + # poll loop may observe more than one completed check, so do not pin the + # exact number of reductions. + self.assertIn(2, min_sizes) + cons.sanity_check() + + def test_tp_swa_prefetch_adopted_when_peer_present(self): + """When every rank has the full SWA window, SWA is adopted under tp>1 + (the single packed all_reduce path still works).""" + setup = self._setup_swa_tp_prefetch() + if setup is None: + return + storage_dir, seq = setup + + cons = self._l3_consumer(storage_dir) + cons.tp_world_size = 2 + min_sizes = self._patch_tp_all_reduce(cons, drop_swa=False) # peer == local + self._consume_prefetch(cons, seq, "keep") + + m = cons.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq)))) + self.assertEqual(m.host_hit_length, len(seq)) + self.assertTrue( + self._swa_host_on_path(cons, seq), + "SWA must be adopted when all ranks have it", + ) + self.assertIn(2, min_sizes) + cons.sanity_check() + def _skip_unsupported_hicache_test(self): if self.cfg.has_swa and self.cfg.has_mamba: self.skipTest("HiCache unit fixture does not support SWA + Mamba stacks")