[HiCache] Label radix-cache metrics per rank and split the "shrunk" prefetch reason (#39280)
This commit is contained in:
@@ -862,7 +862,7 @@ class PrefillAdder:
|
|||||||
"device_capacity"
|
"device_capacity"
|
||||||
if req.needs_host_load_back()
|
if req.needs_host_load_back()
|
||||||
and req.host_loaded_length < req.host_hit_length
|
and req.host_loaded_length < req.host_hit_length
|
||||||
else "shrunk"
|
else "cache_admission_shortfall"
|
||||||
)
|
)
|
||||||
self.tree_cache.finish_storage_prefetch_admission(
|
self.tree_cache.finish_storage_prefetch_admission(
|
||||||
req.cache_request_handle,
|
req.cache_request_handle,
|
||||||
|
|||||||
@@ -25,9 +25,10 @@ from sglang.srt.mem_cache.unified_cache.component_type import ComponentType
|
|||||||
from sglang.srt.observability.metrics_collector import (
|
from sglang.srt.observability.metrics_collector import (
|
||||||
STAT_LOGGER_ROLE_RADIX_CACHE,
|
STAT_LOGGER_ROLE_RADIX_CACHE,
|
||||||
RadixCacheMetricsCollector,
|
RadixCacheMetricsCollector,
|
||||||
|
radix_cache_metric_labels,
|
||||||
resolve_collector_class,
|
resolve_collector_class,
|
||||||
)
|
)
|
||||||
from sglang.srt.runtime_context import get_observability
|
from sglang.srt.runtime_context import get_observability, get_parallel
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from sglang.srt.managers.cache_controller import HiCacheController
|
from sglang.srt.managers.cache_controller import HiCacheController
|
||||||
@@ -340,7 +341,11 @@ class BasePrefixCache(ABC, PrefixCacheTrait):
|
|||||||
kv_events: Optional[KVCacheEventRecorder] = None
|
kv_events: Optional[KVCacheEventRecorder] = None
|
||||||
|
|
||||||
def init_metrics_collector(self):
|
def init_metrics_collector(self):
|
||||||
labels = {"cache_type": self.__class__.__name__}
|
from sglang.srt.layers.dp_attention import is_dp_attention_enabled
|
||||||
|
|
||||||
|
labels = radix_cache_metric_labels(
|
||||||
|
self.__class__.__name__, get_parallel(), is_dp_attention_enabled()
|
||||||
|
)
|
||||||
if get_observability().extra_metric_labels:
|
if get_observability().extra_metric_labels:
|
||||||
labels.update(get_observability().extra_metric_labels)
|
labels.update(get_observability().extra_metric_labels)
|
||||||
radix_cache_cls = resolve_collector_class(
|
radix_cache_cls = resolve_collector_class(
|
||||||
|
|||||||
@@ -2418,7 +2418,7 @@ class UnifiedRadixCache(BasePrefixCache):
|
|||||||
def _handle_storage_prefetch_anchor_loss(self, request: CacheRequestHandle) -> None:
|
def _handle_storage_prefetch_anchor_loss(self, request: CacheRequestHandle) -> None:
|
||||||
operation = self.ongoing_prefetch[request].operation
|
operation = self.ongoing_prefetch[request].operation
|
||||||
storage_hit_end = operation.storage_start + operation.storage_hit_count
|
storage_hit_end = operation.storage_start + operation.storage_hit_count
|
||||||
self._finish_storage_prefetch(request, fulfilled_tokens=0, reason="shrunk")
|
self._finish_storage_prefetch(request, fulfilled_tokens=0, reason="anchor_lost")
|
||||||
self.revoke_pending_prefetch(request)
|
self.revoke_pending_prefetch(request)
|
||||||
self.storage_prefetch_retries.refetch(request.rid, storage_hit_end)
|
self.storage_prefetch_retries.refetch(request.rid, storage_hit_end)
|
||||||
|
|
||||||
|
|||||||
@@ -1936,8 +1936,16 @@ class StorageMetricsCollector(_StatLoggerDIMixin):
|
|||||||
|
|
||||||
self.storage_prefetch_unfulfilled_tokens_total = Counter(
|
self.storage_prefetch_unfulfilled_tokens_total = Counter(
|
||||||
name="sglang:storage_prefetch_unfulfilled_tokens_total",
|
name="sglang:storage_prefetch_unfulfilled_tokens_total",
|
||||||
documentation="Storage-hit tokens that did not become a usable "
|
documentation="Attempt-level storage-hit tokens that did not become "
|
||||||
"prefetch result, by terminal reason.",
|
"a usable prefetch result, by diagnostic reason. This is not a "
|
||||||
|
"final prefill cache-miss counter. anchor_lost means the original "
|
||||||
|
"device prefix was no longer present when the prefetch needed it; "
|
||||||
|
"aux_window_trim means a staged aux trailing window could not be "
|
||||||
|
"trimmed to the shorter splice, so the loaded span was released; "
|
||||||
|
"device_overlap means the live device prefix diverged from the "
|
||||||
|
"request's view at splice time, so the span beyond it was released; "
|
||||||
|
"cache_admission_shortfall means an L3-loaded L2 span was no longer "
|
||||||
|
"reusable when cache-mode admission ran.",
|
||||||
labelnames=list(labels.keys()) + ["reason"],
|
labelnames=list(labels.keys()) + ["reason"],
|
||||||
)
|
)
|
||||||
for reason in (
|
for reason in (
|
||||||
@@ -1946,7 +1954,10 @@ class StorageMetricsCollector(_StatLoggerDIMixin):
|
|||||||
"device_capacity",
|
"device_capacity",
|
||||||
"device_covered",
|
"device_covered",
|
||||||
"storage_transfer",
|
"storage_transfer",
|
||||||
"shrunk",
|
"anchor_lost",
|
||||||
|
"aux_window_trim",
|
||||||
|
"device_overlap",
|
||||||
|
"cache_admission_shortfall",
|
||||||
"dropped",
|
"dropped",
|
||||||
):
|
):
|
||||||
self.storage_prefetch_unfulfilled_tokens_total.labels(
|
self.storage_prefetch_unfulfilled_tokens_total.labels(
|
||||||
@@ -2105,6 +2116,25 @@ class ExpertDispatchCollector(_StatLoggerDIMixin):
|
|||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def radix_cache_metric_labels(
|
||||||
|
cache_type: str, parallel: Any, dp_attention_enabled: bool
|
||||||
|
) -> Dict[str, Any]:
|
||||||
|
# Every scheduler rank runs its own cache over its own KV shard; without
|
||||||
|
# rank labels the multiprocess registry sums ranks into TP x the count.
|
||||||
|
# Same rank keys as the storage collector (cache_controller's storage
|
||||||
|
# config), so one rank's L2 and L3 series line up.
|
||||||
|
if dp_attention_enabled:
|
||||||
|
tp_rank, dp_rank = parallel.attn_tp_rank, parallel.attn_dp_rank
|
||||||
|
else:
|
||||||
|
tp_rank, dp_rank = parallel.tp_rank, 0
|
||||||
|
return {
|
||||||
|
"cache_type": cache_type,
|
||||||
|
"tp_rank": tp_rank,
|
||||||
|
"pp_rank": parallel.pp_rank,
|
||||||
|
"dp_rank": dp_rank,
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
class RadixCacheMetricsCollector(_StatLoggerDIMixin):
|
class RadixCacheMetricsCollector(_StatLoggerDIMixin):
|
||||||
def __init__(
|
def __init__(
|
||||||
self,
|
self,
|
||||||
@@ -2225,7 +2255,9 @@ class RadixCacheMetricsCollector(_StatLoggerDIMixin):
|
|||||||
self.load_back_num_tokens = Counter(
|
self.load_back_num_tokens = Counter(
|
||||||
name="sglang:load_back_tokens_total",
|
name="sglang:load_back_tokens_total",
|
||||||
documentation="The number of tokens loaded back from local host "
|
documentation="The number of tokens loaded back from local host "
|
||||||
"DRAM (L2) to GPU, by host pool (kv, swa, mamba, ...).",
|
"DRAM (L2) to GPU, by host pool (kv, swa, mamba, ...). Every TP "
|
||||||
|
"rank reports the same logical count under its own rank labels; "
|
||||||
|
"read one rank rather than summing ranks.",
|
||||||
labelnames=list(labels.keys()) + ["pool"],
|
labelnames=list(labels.keys()) + ["pool"],
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -2243,9 +2275,11 @@ class RadixCacheMetricsCollector(_StatLoggerDIMixin):
|
|||||||
name="sglang:hicache_backup_bytes_total",
|
name="sglang:hicache_backup_bytes_total",
|
||||||
documentation="Bytes backed up from GPU to local host DRAM (L2), "
|
documentation="Bytes backed up from GPU to local host DRAM (L2), "
|
||||||
"all pools combined, including draft/sidecar transfers that the "
|
"all pools combined, including draft/sidecar transfers that the "
|
||||||
"token counter excludes. Divided by the rate of "
|
"token counter excludes. Each rank reports its own KV shard; sum "
|
||||||
"hicache_backup_duration_seconds_sum, gives the achieved D->H "
|
"ranks for the physical total. A window's byte delta divided by "
|
||||||
"bandwidth while transferring.",
|
"the same window's hicache_backup_duration_seconds sum is the "
|
||||||
|
"bytes-weighted active bandwidth; the plain rate is wall-clock "
|
||||||
|
"payload traffic that includes idle time.",
|
||||||
labelnames=labels.keys(),
|
labelnames=labels.keys(),
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -2253,9 +2287,11 @@ class RadixCacheMetricsCollector(_StatLoggerDIMixin):
|
|||||||
name="sglang:load_back_bytes_total",
|
name="sglang:load_back_bytes_total",
|
||||||
documentation="Bytes loaded back from local host DRAM (L2) to "
|
documentation="Bytes loaded back from local host DRAM (L2) to "
|
||||||
"GPU, all pools combined, including draft/sidecar transfers that "
|
"GPU, all pools combined, including draft/sidecar transfers that "
|
||||||
"the token counter excludes. Divided by the rate of "
|
"the token counter excludes. Each rank reports its own KV shard; "
|
||||||
"load_back_duration_seconds_sum, gives the achieved bandwidth "
|
"sum ranks for the physical total. A window's byte delta divided "
|
||||||
"over each merged H2D load operation.",
|
"by the same window's load_back_duration_seconds sum is the "
|
||||||
|
"bytes-weighted active bandwidth; the plain rate is wall-clock "
|
||||||
|
"payload traffic that includes idle time.",
|
||||||
labelnames=labels.keys(),
|
labelnames=labels.keys(),
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -2264,7 +2300,9 @@ class RadixCacheMetricsCollector(_StatLoggerDIMixin):
|
|||||||
documentation="The number of tokens backed up from GPU to local "
|
documentation="The number of tokens backed up from GPU to local "
|
||||||
"host DRAM (L2), by host pool (kv, swa, mamba, ...). Covers all "
|
"host DRAM (L2), by host pool (kv, swa, mamba, ...). Covers all "
|
||||||
"D->H backups regardless of --hicache-write-policy. Distinct from "
|
"D->H backups regardless of --hicache-write-policy. Distinct from "
|
||||||
"the host-to-storage (L3) sglang:backuped_tokens_total.",
|
"the host-to-storage (L3) sglang:backuped_tokens_total. Every TP "
|
||||||
|
"rank reports the same logical count under its own rank labels; "
|
||||||
|
"read one rank rather than summing ranks.",
|
||||||
labelnames=list(labels.keys()) + ["pool"],
|
labelnames=list(labels.keys()) + ["pool"],
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|||||||
@@ -291,7 +291,7 @@ class TestPrefillAdder(CustomTestCase):
|
|||||||
self.mock_tree_cache.finish_storage_prefetch_admission.assert_called_once_with(
|
self.mock_tree_cache.finish_storage_prefetch_admission.assert_called_once_with(
|
||||||
req.cache_request_handle,
|
req.cache_request_handle,
|
||||||
fulfilled_tokens=4,
|
fulfilled_tokens=4,
|
||||||
reason="shrunk",
|
reason="cache_admission_shortfall",
|
||||||
)
|
)
|
||||||
self.assertEqual(adder.log_device_hit_tokens, 8)
|
self.assertEqual(adder.log_device_hit_tokens, 8)
|
||||||
self.assertEqual(adder.log_host_hit_tokens, 0)
|
self.assertEqual(adder.log_host_hit_tokens, 0)
|
||||||
|
|||||||
@@ -10177,7 +10177,7 @@ class TestAnchorLockOutcomePolicy(CustomTestCase):
|
|||||||
cache.dec_host_lock_ref.assert_not_called()
|
cache.dec_host_lock_ref.assert_not_called()
|
||||||
self.assertEqual(controller.prefetch_tokens_occupied, 8)
|
self.assertEqual(controller.prefetch_tokens_occupied, 8)
|
||||||
|
|
||||||
def test_positive_hit_with_lost_anchor_is_reported_as_shrunk(self):
|
def test_positive_hit_with_lost_anchor_reports_anchor_lost(self):
|
||||||
cache = UnifiedRadixCache.__new__(UnifiedRadixCache)
|
cache = UnifiedRadixCache.__new__(UnifiedRadixCache)
|
||||||
cache.storage_prefetch_retries = StoragePrefetchRetries()
|
cache.storage_prefetch_retries = StoragePrefetchRetries()
|
||||||
cache.ongoing_prefetch = {
|
cache.ongoing_prefetch = {
|
||||||
@@ -10191,7 +10191,7 @@ class TestAnchorLockOutcomePolicy(CustomTestCase):
|
|||||||
cache._handle_storage_prefetch_anchor_loss(self._REQ)
|
cache._handle_storage_prefetch_anchor_loss(self._REQ)
|
||||||
|
|
||||||
cache._finish_storage_prefetch.assert_called_once_with(
|
cache._finish_storage_prefetch.assert_called_once_with(
|
||||||
self._REQ, fulfilled_tokens=0, reason="shrunk"
|
self._REQ, fulfilled_tokens=0, reason="anchor_lost"
|
||||||
)
|
)
|
||||||
# The eviction widened the span, so the request replans over it --
|
# The eviction widened the span, so the request replans over it --
|
||||||
# skipping the query, the prior hit having proved it stored, and the
|
# skipping the query, the prior hit having proved it stored, and the
|
||||||
|
|||||||
@@ -23,6 +23,7 @@ from sglang.test.ci.ci_register import register_cpu_ci
|
|||||||
register_cpu_ci(est_time=10, suite="base-a-test-cpu")
|
register_cpu_ci(est_time=10, suite="base-a-test-cpu")
|
||||||
|
|
||||||
import unittest
|
import unittest
|
||||||
|
from types import SimpleNamespace
|
||||||
|
|
||||||
import prometheus_client
|
import prometheus_client
|
||||||
|
|
||||||
@@ -37,6 +38,7 @@ from sglang.srt.observability.metrics_collector import (
|
|||||||
SchedulerMetricsCollector,
|
SchedulerMetricsCollector,
|
||||||
StorageMetricsCollector,
|
StorageMetricsCollector,
|
||||||
TokenizerMetricsCollector,
|
TokenizerMetricsCollector,
|
||||||
|
radix_cache_metric_labels,
|
||||||
resolve_collector_class,
|
resolve_collector_class,
|
||||||
)
|
)
|
||||||
from sglang.srt.runtime_context import get_context, reset_context
|
from sglang.srt.runtime_context import get_context, reset_context
|
||||||
@@ -179,6 +181,29 @@ class TestDefaultBackend(unittest.TestCase):
|
|||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
class TestRadixCacheMetricLabels(unittest.TestCase):
|
||||||
|
"""Radix-cache series must stay distinct per scheduler rank: an unlabeled
|
||||||
|
family is summed across local ranks by the multiprocess registry, which
|
||||||
|
reported TP x the logical token count in production. The rank keys follow
|
||||||
|
the storage collector's DP-aware convention so L2 and L3 series line up."""
|
||||||
|
|
||||||
|
def test_labels_follow_the_storage_collector_rank_keys(self):
|
||||||
|
parallel = SimpleNamespace(tp_rank=3, pp_rank=1, attn_tp_rank=1, attn_dp_rank=2)
|
||||||
|
self.assertEqual(
|
||||||
|
radix_cache_metric_labels("UnifiedRadixCache", parallel, True),
|
||||||
|
{
|
||||||
|
"cache_type": "UnifiedRadixCache",
|
||||||
|
"tp_rank": 1,
|
||||||
|
"pp_rank": 1,
|
||||||
|
"dp_rank": 2,
|
||||||
|
},
|
||||||
|
)
|
||||||
|
self.assertEqual(
|
||||||
|
radix_cache_metric_labels("RadixCache", parallel, False),
|
||||||
|
{"cache_type": "RadixCache", "tp_rank": 3, "pp_rank": 1, "dp_rank": 0},
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
class TestHiCacheMetrics(unittest.TestCase):
|
class TestHiCacheMetrics(unittest.TestCase):
|
||||||
def test_cached_tokens_uses_literal_storage_source(self):
|
def test_cached_tokens_uses_literal_storage_source(self):
|
||||||
labels = {"model_name": "test"}
|
labels = {"model_name": "test"}
|
||||||
|
|||||||
Reference in New Issue
Block a user