minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[UnifiedTree]: Sync sidecar component hits across TP ranks and make SWA prefetch all-or-nothing (#27264)

Browse Source
This commit is contained in:
Zhangheng
2026-06-05 09:23:54 +08:00
committed by GitHub
parent 5af02c18ae
commit 631db6c757
4 changed files with 246 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/mem_cache/unified_cache_components/swa_component.py
+23 -19
View File
@@ -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
python/sglang/srt/mem_cache/unified_radix_cache.py
+11 -4
View File
@@ -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(
test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_dsv4.py
+63
View File
@@ -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()
test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py
+149 -1
View File
@@ -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")