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

[HiCache] L3 storage prefetch lifecycle metrics and cross-tier attribution fixes (#37503)

Browse Source
This commit is contained in:
Zhiqiang Xie
2026-09-03 16:00:04 -07:00
committed by GitHub
parent 0610a6539d
commit a480f388b2
19 changed files with 828 additions and 143 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/unit/managers/test_prefill_adder.py
+50
View File
@@ -101,6 +101,11 @@ class TestPrefillAdder(CustomTestCase):
req.retracted_stain = False
req.host_hit_length = 0
req.storage_hit_length = 0
req.storage_hit_start = None
req.host_hit_is_storage = False
req.host_loaded_length = 0
req.materialized_host_hit_len.return_value = 0
req.fulfilled_storage_hit_len.return_value = 0
req.finished.return_value = False
req.needs_host_load_back.return_value = False
return req
@@ -120,6 +125,51 @@ class TestPrefillAdder(CustomTestCase):
defaults.update(kwargs)
return PrefillAdder(**defaults)
def test_storage_prefetch_fulfillment_resolves_at_admission(self):
adder = self.create_adder(self.create_running_batch())
req = self.create_mock_req("storage-hit", priority=0, max_new_tokens=1)
req.host_hit_length = 12
req.host_loaded_length = 4
req.storage_hit_length = 8
req.storage_hit_start = 4
req.materialized_host_hit_len.return_value = 4
req.fulfilled_storage_hit_len.return_value = 8
req.needs_host_load_back.return_value = True
adder._account_prefill_cache_admission(req, prefix_len=12)
self.mock_tree_cache.finish_storage_prefetch_admission.assert_called_once_with(
"storage-hit",
fulfilled_tokens=8,
reason=None,
)
self.assertEqual(adder.log_device_hit_tokens, 4)
self.assertEqual(adder.log_host_hit_tokens, 0)
self.assertEqual(adder.log_storage_hit_tokens, 8)
self.mock_tree_cache.finish_storage_prefetch_admission.reset_mock()
req.host_loaded_length = 0
req.materialized_host_hit_len.return_value = 0
req.fulfilled_storage_hit_len.return_value = 0
adder._account_prefill_cache_admission(req, prefix_len=0)
self.mock_tree_cache.finish_storage_prefetch_admission.assert_called_once_with(
"storage-hit", fulfilled_tokens=0, reason="device_capacity"
)
def test_retracted_storage_prefetch_accounting_is_omitted(self):
adder = self.create_adder(self.create_running_batch())
req = self.create_mock_req(
"retracted-storage-hit", priority=0, max_new_tokens=1
)
req.retracted_stain = True
adder._account_prefill_cache_admission(req, prefix_len=8)
self.mock_tree_cache.discard_storage_prefetch_accounting.assert_called_once_with(
"retracted-storage-hit"
)
self.mock_tree_cache.finish_storage_prefetch_admission.assert_not_called()
def test_preempt_success_high_priority_values_first(self):
params = [
("run1", 0, 50),
test/registered/unit/managers/test_schedule_batch_out_of_place.py
+63 -1
View File
@@ -11,7 +11,10 @@ from sglang.test.test_utils import maybe_stub_sgl_kernel
maybe_stub_sgl_kernel()
from sglang.srt.managers.schedule_batch import ScheduleBatch # noqa: E402
from sglang.srt.managers.schedule_batch import ( # noqa: E402
ScheduleBatch,
split_cached_prefix_by_tier,
)
from sglang.srt.model_executor.forward_batch_info import ForwardMode # noqa: E402
from sglang.srt.speculative.spec_info import SpeculativeAlgorithm # noqa: E402
from sglang.srt.utils.common import Range # noqa: E402
@@ -21,6 +24,65 @@ register_cpu_ci(est_time=5, suite="base-a-test-cpu")
AUTO_FILL_EXCLUDED_FIELDS = ["reqs"]
class TestCachedPrefixTierAttribution(unittest.TestCase):
def test_split_cached_prefix_by_tier(self):
cases = [
(
{"prefix_len": 100, "host_hit_len": 70, "storage_hit_len": 20},
(30, 50, 20),
),
(
{
"prefix_len": 80,
"host_hit_len": 50,
"storage_hit_len": 20,
"storage_hit_start": 80,
},
(30, 50, 0),
),
(
{
"prefix_len": 90,
"host_hit_len": 60,
"storage_hit_len": 20,
"storage_hit_start": 80,
},
(30, 50, 10),
),
(
{
"prefix_len": 100,
"host_hit_len": 70,
"storage_hit_len": 70,
"host_hit_is_storage": True,
},
(30, 0, 70),
),
(
{
"prefix_len": 100,
"host_hit_len": 0,
"storage_hit_len": 20,
"storage_hit_start": 80,
},
(80, 0, 20),
),
(
{
"prefix_len": 100,
"host_hit_len": 0,
"storage_hit_len": 20,
"storage_hit_start": 80,
"host_hit_is_storage": True,
},
(80, 0, 20),
),
]
for kwargs, expected in cases:
with self.subTest(**kwargs):
self.assertEqual(split_cached_prefix_by_tier(**kwargs), expected)
def make_schedule_batch(bs: int, **overrides) -> ScheduleBatch:
batch = ScheduleBatch(reqs=overrides.pop("reqs"))
# init_new always sets a SpeculativeAlgorithm enum, never None.
test/registered/unit/mem_cache/test_hicache_staged_write_back_dispatch.py
+30
View File
@@ -9,6 +9,7 @@ import torch
from sglang.srt.managers.cache_controller import CacheOperation, HiCacheController
from sglang.srt.mem_cache import l2_transfer as transfer_module
from sglang.srt.mem_cache.buffer_mode.pipeline import BufferModePipeline
from sglang.srt.mem_cache.hicache_storage import (
PoolHitPolicy,
PoolName,
@@ -262,6 +263,35 @@ class TestHiCacheStagedWriteBackDispatch(CustomTestCase):
controller._num_tokens_by_pool.assert_called_once_with(merged_op)
self.assertEqual(controller.ack_load_queue[0].node_ids, [7, 7])
def test_short_staged_swa_tail_resolves_device_covered_head(self):
pipeline = BufferModePipeline.__new__(BufferModePipeline)
pipeline._cache = mock.Mock()
pipeline.release_staged_hold = mock.Mock(return_value=True)
pipeline.staged_prefetches = {
"r": SimpleNamespace(
req_id="r",
key_tokens=list(range(8)),
extra_key=None,
cache_salt=None,
matched_len=2,
num_tokens=6,
occupied_tokens=6,
host_indices=_indices(0, 6),
aux_xfers=[
PoolTransfer(
name=PoolName.SWA,
host_indices=_indices(0, 4),
)
],
hash_values=[],
operation_id=1,
)
}
self.assertEqual(pipeline.plan_staged_splice("r", device_prefix_len=6), (0, 0))
pipeline._cache._resolve_storage_prefetch_tokens.assert_called_once_with("r", 4)
pipeline.release_staged_hold.assert_called_once_with("r", reason="shrunk")
def test_l2_transfer_maps_global_layers(self):
host_pool = mock.Mock()
transfer = L2Transfer(
test/registered/unit/mem_cache/test_mamba_path_state_cap.py
+4
View File
@@ -47,6 +47,10 @@ class _FakeTreeCore:
}
self.evicted = []
self.cascaded = []
# The real eviction helper reports unbacked FULL evictions to the
# write-through drop counter; this fake never tracks a walk.
self._is_tracking_unbacked_tokens = False
self._tracked_unbacked_tokens = 0
def _evict_component_and_detach_lru(self, node, component, *args, **kwargs):
self.evicted.append(node)
test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py
+59 -1
View File
@@ -3676,6 +3676,8 @@ class UnifiedRadixCacheSuite:
self.cfg, enable_kv_cache_events=True
)
self._init_buffer_hicache(cons, storage_dir)
cons.enable_storage_metrics = True
cons.storage_metrics_collector = mock.Mock()
cons.take_events()
avail0 = self._host_avail_sizes(cons)
dev_avail0 = cons.token_to_kv_pool_allocator.available_size()
@@ -3783,8 +3785,14 @@ class UnifiedRadixCacheSuite:
and e.medium == StorageMedium.CPU
]
self.assertEqual(cpu_events, [])
cons.storage_metrics_collector.log_storage_prefetch_hit_tokens.assert_called_once_with(
len(seq)
)
cons.storage_metrics_collector.log_prefetched_tokens.assert_called_once_with(
len(seq)
)
cons.storage_metrics_collector.log_storage_prefetch_unfulfilled_tokens.assert_not_called()
self.assertIn("occupancy_ratio", cons.prefetch_outcome_stats_snapshot())
cons.sanity_check()
def test_buffer_only_cache_salt_uses_the_request_namespace(self):
@@ -4148,6 +4156,8 @@ class UnifiedRadixCacheSuite:
cons, cons_alloc, cons_rtp = build_fixture(self.cfg)
self._init_buffer_hicache(cons, storage_dir)
cons.enable_storage_metrics = True
cons.storage_metrics_collector = mock.Mock()
avail0 = self._host_avail_sizes(cons)
req_id = "sibling-publish"
@@ -4188,6 +4198,7 @@ class UnifiedRadixCacheSuite:
k, v = self._snapshot_full_kv(cons_alloc, m.device_indices)
self.assertTrue(torch.equal(k, sib_kv[0]))
self.assertTrue(torch.equal(v, sib_kv[1]))
cons.storage_metrics_collector.log_storage_prefetch_unfulfilled_tokens.assert_not_called()
cons.sanity_check()
def test_buffer_only_load_back_drops_on_full_overlap_masked_by_swa_tombstone(
@@ -8448,6 +8459,7 @@ class TestPrefetchCommitOrdering(CustomTestCase):
cache.tree_core.insert_host.return_value = insert_result
operation = mock.MagicMock()
operation.request_id = "req"
operation.completed_tokens = 8
cache.ongoing_prefetch = {
operation.request_id: (
7,
@@ -8588,6 +8600,8 @@ class TestUnifiedRadixPrefetchCorruption(CustomTestCase):
"""The caller owns every completed buffer when host insertion drops."""
cache, allocator, _ = build_fixture(self.cfg)
self._init_hicache(cache)
cache.enable_storage_metrics = True
cache.storage_metrics_collector = mock.Mock()
parent_id = self._insert_device(
cache, allocator, list(range(1, 1 + 3 * self.ps))
@@ -8633,6 +8647,7 @@ class TestUnifiedRadixPrefetchCorruption(CustomTestCase):
anchor_lock_params,
comp_xfers,
)
cache._record_storage_prefetch_hit(req_id, completed_tokens)
cache.cache_controller.prefetch_tokens_occupied = completed_tokens
hashes = [f"h{i}" for i in range(completed_tokens // self.ps)]
operation.hash_value = hashes
@@ -8681,9 +8696,38 @@ class TestUnifiedRadixPrefetchCorruption(CustomTestCase):
)
self.assertIs(drop_releases[0].kwargs["extra_pools"][0], swa_transfer)
self.assertIs(drop_releases[0].kwargs["extra_pools"][1], mamba_transfer)
cache.storage_metrics_collector.log_storage_prefetch_hit_tokens.assert_called_once_with(
completed_tokens
)
cache.storage_metrics_collector.log_storage_prefetch_unfulfilled_tokens.assert_called_once_with(
completed_tokens, "dropped"
)
cache.sanity_check()
def test_write_through_eviction_counts_unbacked_tokens(self):
if _selected_tree_core_test_backend() == "rust":
# The unbacked-eviction tracker is a Python tree-core feature;
# UnifiedRadixCache only enables it for that backend.
self.skipTest(
"write-through unbacked-eviction tracking is Python-core only"
)
cache, allocator, _ = build_fixture(self.cfg)
self._init_hicache(cache)
cache.metrics_collector = mock.Mock()
seq = list(range(1, 1 + 2 * self.ps))
self._insert_device(cache, allocator, seq)
result = cache.evict(EvictParams(num_tokens=len(seq)))
self.assertGreaterEqual(result.num_tokens_evicted, len(seq))
cache.metrics_collector.increment_dropped_tokens.assert_called_once_with(
num_tokens=len(seq),
reason="write_through_unbacked_eviction",
pool=PoolName.KV.value,
)
cache.sanity_check()
def test_prefetch_refill_leaves_eviction_path_uncorrupted(self):
"""Write-through: eviction after such a refill must not corrupt the tree."""
cache, allocator, _ = build_fixture(self.cfg)
@@ -8946,6 +8990,20 @@ class TestAnchorLockOutcomePolicy(CustomTestCase):
self.assertEqual(pipeline.anchor_locks, {})
self.assertEqual(pipeline.anchor_locked_tokens_, 0)
def test_positive_hit_with_lost_anchor_is_reported_as_shrunk(self):
cache = UnifiedRadixCache.__new__(UnifiedRadixCache)
cache._storage_prefetch_missed_rids = set()
cache._finish_storage_prefetch = mock.Mock()
cache.revoke_pending_prefetch = mock.Mock()
cache._handle_storage_prefetch_anchor_loss(self._REQ)
cache._finish_storage_prefetch.assert_called_once_with(
self._REQ, fulfilled_tokens=0, reason="shrunk"
)
self.assertIn(self._REQ, cache._storage_prefetch_missed_rids)
cache.revoke_pending_prefetch.assert_called_once_with(self._REQ)
def test_over_cap_reports_cap_skip_before_matching(self):
cache = self._make_cache(live_match_len=len(self._PREFIX))
pipeline = self._make_pipeline(cache, cap_tokens=len(self._PREFIX) - 1)
test/registered/unit/observability/test_stat_loggers_di.py
+87
View File
@@ -42,6 +42,48 @@ from sglang.srt.observability.metrics_collector import (
from sglang.srt.runtime_context import get_context, reset_context
class _BoundRecordingMetric:
def __init__(self, metric, labels):
self.metric = metric
self.labels = labels
def inc(self, value=1):
self.metric.increments.append((self.labels, value))
def observe(self, value):
self.metric.observations.append((self.labels, value))
def set(self, value):
self.metric.sets.append((self.labels, value))
class _RecordingMetric:
"""Small prometheus_client-compatible metric that preserves labels."""
def __init__(self, *args, name=None, labelnames=(), **kwargs):
self.name = name if name is not None else args[0]
self.labelnames = tuple(labelnames)
self.increments = []
self.observations = []
self.sets = []
def labels(self, *values, **labels):
if values:
labels = dict(zip(self.labelnames, values, strict=True))
return _BoundRecordingMetric(self, labels)
class _RecordingTokenizerMetricsCollector(TokenizerMetricsCollector):
_counter_cls = _RecordingMetric
_gauge_cls = _RecordingMetric
_histogram_cls = _RecordingMetric
class _RecordingStorageMetricsCollector(StorageMetricsCollector):
_counter_cls = _RecordingMetric
_histogram_cls = _RecordingMetric
class TestCollectorClassAttrs(unittest.TestCase):
"""All five collectors expose four DI hook class attrs, all defaulting to
None so the existing prometheus_client backend is used unchanged."""
@@ -137,5 +179,50 @@ class TestDefaultBackend(unittest.TestCase):
)
class TestHiCacheMetrics(unittest.TestCase):
def test_cached_tokens_uses_literal_storage_source(self):
labels = {"model_name": "test"}
with get_context().override_server_args(
prompt_tokens_buckets=None, generation_tokens_buckets=None
):
collector = _RecordingTokenizerMetricsCollector(labels=labels)
collector.observe_one_finished_request(
labels=labels,
prompt_tokens=20,
generation_tokens=2,
cached_tokens=12,
e2e_latency=0.1,
has_grammar=False,
cached_tokens_details={
"device": 3,
"host": 4,
"storage": 5,
"storage_backend": "BackendShim",
},
)
by_source = {
metric_labels["cache_source"]: value
for metric_labels, value in collector.cached_tokens_total.increments
}
self.assertEqual(by_source, {"device": 3, "host": 4, "storage": 5})
def test_storage_prefetch_lifecycle_metrics(self):
labels = {"model_name": "test"}
collector = _RecordingStorageMetricsCollector(labels=labels)
collector.log_storage_prefetch_hit_tokens(21)
collector.log_storage_prefetch_unfulfilled_tokens(4, "storage_transfer")
self.assertEqual(
collector.storage_prefetch_hit_tokens_total.increments, [(labels, 21)]
)
self.assertEqual(
collector.storage_prefetch_unfulfilled_tokens_total.increments,
[({**labels, "reason": "storage_transfer"}, 4)],
)
if __name__ == "__main__":
unittest.main()