diff --git a/rust/src/engine-core-client/src/metrics.rs b/rust/src/engine-core-client/src/metrics.rs index b392d52601..e8b0741066 100644 --- a/rust/src/engine-core-client/src/metrics.rs +++ b/rust/src/engine-core-client/src/metrics.rs @@ -6,14 +6,15 @@ use std::collections::BTreeSet; use std::time::{SystemTime, UNIX_EPOCH}; use vllm_metrics::{ - EngineLabels, EnginePositionLabels, F64Gauge, Family, HistogramMetric, LoraAdapterNames, - LoraInfoLabels, MooncakeOperationCounterFamily, MooncakeOperationHistogramFamily, - MooncakeOperationLabels, RequestMetrics, SchedulerLogStatsAccumulator, SchedulerMetrics, - U64Counter, U64Gauge, WaitingReasonLabels, + CacheHitSourceLabels, EngineLabels, EnginePositionLabels, F64Gauge, Family, HistogramMetric, + LoraAdapterNames, LoraInfoLabels, MooncakeOperationCounterFamily, + MooncakeOperationHistogramFamily, MooncakeOperationLabels, RequestMetrics, + SchedulerLogStatsAccumulator, SchedulerMetrics, U64Counter, U64Gauge, WaitingReasonLabels, }; use crate::protocol::stats::{ - KvConnectorStats, MooncakeStats, MultiConnectorStats, NixlStats, SchedulerStats, + CachedTokensBySource, KvConnectorStats, MooncakeStats, MultiConnectorStats, NixlStats, + SchedulerStats, }; use crate::transport::ConnectedEngine; @@ -60,6 +61,8 @@ struct SchedulerStatsHandles { prefix_cache_hits: U64Counter, external_prefix_cache_queries: U64Counter, external_prefix_cache_hits: U64Counter, + /// One handle per `CachedTokensBySource::SOURCES` entry. + prompt_tokens_cached_by_source: [U64Counter; 5], // Speculative decoding counters. spec_decode_num_drafts: U64Counter, @@ -165,6 +168,15 @@ fn resolve_scheduler_stats_handles( .external_prefix_cache_queries .get_or_create_owned(&labels), external_prefix_cache_hits: metrics.external_prefix_cache_hits.get_or_create_owned(&labels), + prompt_tokens_cached_by_source: CachedTokensBySource::SOURCES.map(|source| { + metrics + .prompt_tokens_cached_by_source + .get_or_create_owned(&CacheHitSourceLabels { + model_name: model_name.to_string(), + engine, + source, + }) + }), spec_decode_num_drafts: metrics.spec_decode_num_drafts.get_or_create_owned(&labels), spec_decode_num_draft_tokens: metrics .spec_decode_num_draft_tokens @@ -220,6 +232,8 @@ fn record_scheduler_stats_with_handles(handles: &SchedulerStatsHandles, stats: & // Prefix-cache counters, including the connector-backed external cache path. handles.prefix_cache_queries.inc_by(stats.prefix_cache_stats.base.queries); handles.prefix_cache_hits.inc_by(stats.prefix_cache_stats.base.hits); + // Local prefix-cache hits are the `device` tier. + handles.prompt_tokens_cached_by_source[0].inc_by(stats.prefix_cache_stats.base.hits); if let Some(connector_prefix_cache_stats) = &stats.connector_prefix_cache_stats { handles @@ -228,6 +242,10 @@ fn record_scheduler_stats_with_handles(handles: &SchedulerStatsHandles, stats: & handles .external_prefix_cache_hits .inc_by(connector_prefix_cache_stats.base.hits); + let counts = connector_prefix_cache_stats.hits_by_source.counts(); + for (handle, count) in handles.prompt_tokens_cached_by_source.iter().zip(counts) { + handle.inc_by(count); + } } // Speculative decoding counters. diff --git a/rust/src/engine-core-client/src/protocol/stats.rs b/rust/src/engine-core-client/src/protocol/stats.rs index aa1b9c69ab..7797ecf6e4 100644 --- a/rust/src/engine-core-client/src/protocol/stats.rs +++ b/rust/src/engine-core-client/src/protocol/stats.rs @@ -21,6 +21,34 @@ pub struct BaseCacheStats { pub hits: u64, } +/// Cache-hit tokens per source tier. Python sends a sparse +/// `dict[CacheHitSource, int]`, so absent tiers default to zero. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)] +#[serde(default)] +pub struct CachedTokensBySource { + pub device: u64, + pub host: u64, + pub p2p: u64, + pub disk: u64, + pub external_unspecified: u64, +} + +impl CachedTokensBySource { + /// Label values, in the same order as [`Self::counts`]. + pub const SOURCES: [&'static str; 5] = + ["device", "host", "p2p", "disk", "external_unspecified"]; + + pub fn counts(&self) -> [u64; 5] { + [ + self.device, + self.host, + self.p2p, + self.disk, + self.external_unspecified, + ] + } +} + /// Stores prefix cache hit statistics. /// - `reset`: Whether `reset_prefix_cache` was invoked. /// - `queries`: Refers to the number of tokens that were queried. @@ -38,6 +66,9 @@ pub struct PrefixCacheStats { pub preempted_queries: u64, /// The `hits` number for preempted requests. pub preempted_hits: u64, + /// `hits` split by the cache tier that supplied them (connector stats only). + #[serde(default)] + pub hits_by_source: CachedTokensBySource, } /// Single KV cache block eviction sample. @@ -304,3 +335,31 @@ pub struct SchedulerStats { /// Estimated MFU/performance stats, when enabled. pub perf_stats: Option, } + +#[cfg(test)] +mod tests { + use super::{CachedTokensBySource, PrefixCacheStats}; + + #[test] + fn prefix_cache_stats_hits_by_source_defaults_and_decodes() { + // Engines without per-source attribution omit the field entirely. + let legacy = serde_json::json!({ + "reset": false, "requests": 1, "queries": 8, "hits": 4, + "preempted_requests": 0, "preempted_queries": 0, "preempted_hits": 0 + }); + let wire = rmp_serde::to_vec_named(&legacy).unwrap(); + let stats: PrefixCacheStats = rmp_serde::from_slice(&wire).unwrap(); + assert_eq!(stats.base.hits, 4); + assert_eq!(stats.hits_by_source, CachedTokensBySource::default()); + + let payload = serde_json::json!({ + "reset": false, "requests": 1, "queries": 16, "hits": 12, + "preempted_requests": 0, "preempted_queries": 0, "preempted_hits": 0, + "hits_by_source": {"host": 8, "p2p": 4} + }); + let wire = rmp_serde::to_vec_named(&payload).unwrap(); + let stats: PrefixCacheStats = rmp_serde::from_slice(&wire).unwrap(); + assert_eq!(stats.hits_by_source.counts(), [0, 8, 4, 0, 0]); + assert_eq!(CachedTokensBySource::SOURCES[1], "host"); + } +} diff --git a/rust/src/metrics/src/scheduler.rs b/rust/src/metrics/src/scheduler.rs index cdc25a244d..64ca465418 100644 --- a/rust/src/metrics/src/scheduler.rs +++ b/rust/src/metrics/src/scheduler.rs @@ -86,6 +86,13 @@ pub struct EngineLabels { pub engine: u32, } +#[derive(Clone, Debug, Hash, PartialEq, Eq, EncodeLabelSet)] +pub struct CacheHitSourceLabels { + pub model_name: String, + pub engine: u32, + pub source: &'static str, +} + #[derive(Clone, Debug, Hash, PartialEq, Eq, EncodeLabelSet)] pub struct EnginePositionLabels { pub model_name: String, @@ -236,6 +243,7 @@ pub struct SchedulerMetrics { pub prefix_cache_hits: Family, pub external_prefix_cache_queries: Family, pub external_prefix_cache_hits: Family, + pub prompt_tokens_cached_by_source: Family, // Speculative decoding counters. pub spec_decode_num_drafts: Family, @@ -339,6 +347,13 @@ impl SchedulerMetrics { ); let external_prefix_cache_hits = Family::default(); + + let prompt_tokens_cached_by_source = Family::default(); + registry.register( + "vllm:prompt_tokens_cached_by_source", + "Prefix-cache hit tokens by the cache tier that supplied them: device (local HBM), host (offloaded to DRAM), disk, p2p (transferred from another vLLM instance) or external_unspecified. Counted at admission; sums to vllm:prefix_cache_hits + vllm:external_prefix_cache_hits.", + prompt_tokens_cached_by_source.clone(), + ); registry.register( "vllm:external_prefix_cache_hits", "External prefix cache hits from KV connector cross-instance cache sharing, in terms of number of cached tokens.", @@ -528,6 +543,7 @@ impl SchedulerMetrics { prefix_cache_hits, external_prefix_cache_queries, external_prefix_cache_hits, + prompt_tokens_cached_by_source, spec_decode_num_drafts, spec_decode_num_draft_tokens, spec_decode_num_accepted_tokens, diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index 66368a299d..f78db1917d 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -3,7 +3,7 @@ import dataclasses from concurrent.futures import Future from types import SimpleNamespace -from unittest.mock import Mock +from unittest.mock import Mock, patch import numpy as np import pytest @@ -60,6 +60,7 @@ from vllm.v1.kv_cache_interface import ( KVCacheGroupSpec, MambaSpec, ) +from vllm.v1.metrics.cache_hit_source import CacheHitSource from vllm.v1.outputs import ( DraftTokenIds, ECConnectorOutput, @@ -2613,6 +2614,65 @@ def test_kv_connector_honors_skip_reading_prefix_cache(): assert output.num_scheduled_tokens[scoring.request_id] == BLOCK_SIZE * 4 +@pytest.mark.parametrize("is_async", [False, True]) +def test_kv_connector_records_external_cache_hit_sources(monkeypatch, is_async): + block_size = 16 + num_matched_tokens = 2 * block_size + scheduler = create_scheduler( + enable_prefix_caching=True, + use_kv_connector=mock_kv( + matched_tokens=num_matched_tokens, + is_async=is_async, + ), + block_size=block_size, + ) + request = create_requests( + num_requests=1, + num_tokens=2 * num_matched_tokens, + block_size=block_size, + )[0] + assert scheduler.connector is not None + get_sources = Mock( + return_value={CacheHitSource.P2P: block_size, CacheHitSource.HOST: block_size} + ) + update_state = Mock(wraps=scheduler.connector.update_state_after_alloc) + calls = Mock() + calls.attach_mock(update_state, "allocated") + calls.attach_mock(get_sources, "sources") + monkeypatch.setattr(scheduler.connector, "update_state_after_alloc", update_state) + monkeypatch.setattr( + scheduler.connector, "get_external_cache_hit_sources", get_sources + ) + + scheduler.add_request(request) + # A failed allocation is not an admission: nothing is attributed. + with patch.object(scheduler.kv_cache_manager, "allocate_slots", return_value=None): + scheduler.schedule() + get_sources.assert_not_called() + update_state.assert_not_called() + + scheduler.schedule() + # Attribution runs once, after the connector has built its load plan. + get_sources.assert_called_once_with(request, num_matched_tokens) + assert [c[0] for c in calls.mock_calls] == ["allocated", "sources"] + connector_stats = scheduler.connector_prefix_cache_stats + assert connector_stats is not None + assert connector_stats.hits == num_matched_tokens + assert connector_stats.hits_by_source == { + CacheHitSource.P2P: block_size, + CacheHitSource.HOST: block_size, + } + + if is_async: + # Re-admission after the async load completes is not a new hit. + scheduler.make_stats() + _step_until_kv_transfer_finished(scheduler, [request.request_id]) + get_sources.assert_called_once() + connector_stats = scheduler.connector_prefix_cache_stats + assert connector_stats is not None + assert connector_stats.hits_by_source == {} + + @pytest.mark.parametrize("is_async", [False, True]) def test_kv_connector_basic(is_async: bool): """Test whether Scheduler with KVConnector schedules tokens, allocates diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py index 5b3bf1c8be..e71c3dbf18 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -77,6 +77,7 @@ from vllm.v1.kv_offload.base import ( make_offload_key, ) from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager +from vllm.v1.metrics.cache_hit_source import CacheHitSource from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput from vllm.v1.request import RequestStatus @@ -197,14 +198,18 @@ def test_swa_offload_window_covers_unaligned_hit( manager.complete_load(job.keys, state.req_context) -def _make_partial_tail_scheduler() -> OffloadingConnectorScheduler: +def _make_partial_tail_scheduler( + recurrent_only: bool = False, +) -> OffloadingConnectorScheduler: vllm_config = _make_vllm_config(extra_config={"self_describing_kv_events": True}) - vllm_config.cache_config.prefix_match_unit = 4 + vllm_config.cache_config.prefix_match_unit = None if recurrent_only else 4 vllm_config.speculative_config = None vllm_config.kv_events_config = KVEventsConfig( enable_kv_cache_events=True, publisher="null" ) kv_cache_config = _make_mamba_hybrid_kv_cache_config() + if recurrent_only: + kv_cache_config.kv_cache_groups = kv_cache_config.kv_cache_groups[1:] spec = MockOffloadingSpec(build_offloading_config(vllm_config, kv_cache_config)) return OffloadingConnectorScheduler(spec, vllm_config, kv_cache_config) @@ -576,6 +581,284 @@ def test_recurrent_group_unhashed_block_does_not_truncate_load_boundary(): assert 42 in loaded +@pytest.mark.parametrize( + ("recurrent_source", "expected"), + [ + ("host", {CacheHitSource.P2P: 16, CacheHitSource.HOST: 12}), + ("p2p", {CacheHitSource.P2P: 28}), + (None, {CacheHitSource.P2P: 16, CacheHitSource.HOST: 12}), + ], +) +def test_external_cache_hit_sources_preserve_partial_tail_origin( + recurrent_source, expected +): + """Recurrent state covers the whole prefix; the outermost tier wins.""" + scheduler = _make_partial_tail_scheduler() + request = _make_partial_tail_request(scheduler) + req_status = scheduler._req_status["req"] + req_status.update_offload_keys() + + scheduler.manager.lookup.return_value = LookupResult.HIT + assert scheduler._lookup(req_status) == 28 + + def get_load_source(key, req_context): + if get_offload_group_idx(key) == 1: + assert recurrent_source is not None + return recurrent_source + block_hash = get_offload_block_hash(key) + return "p2p" if block_hash == b"h3" else "host" + + scheduler.manager.get_load_source.side_effect = get_load_source + + recurrent_block = KVCacheBlock(41) + if recurrent_source is None: + recurrent_block.set_block_hash( + make_block_hash_with_group_id(BlockHash(b"cached"), 1), 12 + ) + scheduler.update_state_after_alloc( + request, + KVCacheBlocks( + ( + [KVCacheBlock(31), KVCacheBlock(32)], + [KVCacheBlock(0, is_null=True), recurrent_block], + ) + ), + num_external_tokens=28, + ) + assert req_status.partial_tail_boundary is None + assert scheduler.get_external_cache_hit_sources(request, 28) == expected + + +@pytest.mark.parametrize( + ("second_source", "expected"), + [ + ("host", {CacheHitSource.HOST: 4, CacheHitSource.DISK: 4}), + ("p2p", {CacheHitSource.P2P: 4, CacheHitSource.DISK: 4}), + ], +) +def test_external_cache_hit_sources_reconcile_different_group_chunk_sizes( + request_runner, + second_source, + expected, +): + """Groups with different chunk sizes overlap; the outermost tier wins.""" + block_size = 4 + groups = [ + KVCacheGroupSpec( + [f"layer{group_idx}"], + FullAttentionSpec( + block_size=block_size * (group_idx + 1), + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + ) + for group_idx in range(2) + ] + runner = request_runner( + block_size=block_size, + num_gpu_blocks=100, + async_scheduling=False, + kv_cache_groups=groups, + ) + scheduler = runner.connector_scheduler + request = MagicMock(request_id="mixed") + req_status = RequestOffloadState( + config=scheduler.config, + req=request, + req_context=ReqContext(req_id=request.request_id), + offloading_context=RequestOffloadingContext(policy=OffloadPolicy.CHUNK_LEVEL), + ) + for group_idx, group_state in enumerate(req_status.group_states): + group_state.offload_keys = [ + make_offload_key(f"h{chunk_idx}".encode(), group_idx) + for chunk_idx in range(2 // (group_idx + 1)) + ] + scheduler._req_status[request.request_id] = req_status + runner.manager.get_load_source.side_effect = lambda key, req_context: ( + ("host" if get_offload_block_hash(key) == b"h0" else "disk") + if get_offload_group_idx(key) == 0 + else second_source + ) + scheduler.update_state_after_alloc( + request, + KVCacheBlocks(([KVCacheBlock(1), KVCacheBlock(2)], [KVCacheBlock(3)])), + 2 * block_size, + ) + + assert scheduler.get_external_cache_hit_sources(request, 2 * block_size) == expected + + +@pytest.mark.parametrize( + ("ranges", "expected"), + [ + pytest.param( + [(8, 12, "host"), (0, 8, "host"), (4, 10, "host")], + {CacheHitSource.HOST: 8, CacheHitSource.EXTERNAL_UNSPECIFIED: 4}, + id="overlapping-same-source-and-adjacent-endpoints", + ), + pytest.param( + [(10, 20, "p2p"), (6, 10, "disk"), (0, 8, "host"), (4, 12, "host")], + {CacheHitSource.HOST: 2, CacheHitSource.DISK: 4, CacheHitSource.P2P: 6}, + id="nested-tiers-and-clipped-ranges", + ), + pytest.param( + [(0, 4, "disk"), (16, 20, "disk"), (8, 8, "disk")], + {CacheHitSource.EXTERNAL_UNSPECIFIED: 12}, + id="no-covering-ranges", + ), + ], +) +def test_external_cache_hit_sources_range_boundaries(ranges, expected): + scheduler = _make_partial_tail_scheduler() + request = _make_partial_tail_request(scheduler) + state = scheduler._req_status[request.request_id] + state.num_locally_computed_tokens = 4 + for idx, (lo, hi, source) in enumerate(ranges): + key = make_offload_key(str(idx).encode(), 0) + state.load_key_ranges.append((lo, hi, key)) + scheduler.manager.get_load_source.side_effect = lambda key, ctx: ranges[ + int(get_offload_block_hash(key)) + ][2] + + result = scheduler.get_external_cache_hit_sources(request, 12) + assert result == expected + assert sum(result.values()) == 12 + + +@pytest.mark.parametrize("source", ["host", "disk", "p2p", "external_unspecified"]) +def test_external_cache_hit_sources_recurrent_only_state(source): + scheduler = _make_partial_tail_scheduler(recurrent_only=True) + request = _make_partial_tail_request(scheduler) + request.num_prompt_tokens = request.num_tokens = 64 + request.block_hashes = [BlockHash(f"h{i}".encode()) for i in range(4)] + request.all_token_ids = list(range(64)) + state = scheduler._req_status[request.request_id] + state.update_offload_keys() + scheduler.manager.lookup.return_value = LookupResult.HIT + assert scheduler._lookup(state) == 48 + scheduler.manager.get_load_source.return_value = source + scheduler.update_state_after_alloc( + request, + KVCacheBlocks( + ( + [ + KVCacheBlock(0, is_null=True), + KVCacheBlock(0, is_null=True), + KVCacheBlock(41), + ], + ) + ), + 48, + ) + assert scheduler.get_external_cache_hit_sources(request, 48) == { + CacheHitSource(source): 48 + } + scheduler.manager.get_load_source.assert_called_once() + + +@pytest.mark.parametrize("sparse_kind", ["sliding", "chunked_local"]) +@pytest.mark.parametrize("blocks_per_chunk", [1, 2]) +@pytest.mark.parametrize("local_tokens", [0, 4]) +@pytest.mark.parametrize("full_attention", [False, True]) +@pytest.mark.parametrize( + "sparse_source", ["host", "disk", "p2p", "external_unspecified", "mixed"] +) +def test_external_cache_hit_sources_use_loaded_attention_ranges( + sparse_kind, + blocks_per_chunk, + local_tokens, + full_attention, + sparse_source, +): + """SWA/local attention affects only loaded positions, not omitted history.""" + block_size = 4 + chunk_size = block_size * blocks_per_chunk + spec_kwargs = dict( + block_size=block_size, num_kv_heads=1, head_size=1, dtype=torch.float32 + ) + sparse_spec = ( + SlidingWindowSpec(**spec_kwargs, sliding_window=2 * chunk_size) + if sparse_kind == "sliding" + else ChunkedLocalAttentionSpec( + **spec_kwargs, attention_chunk_size=2 * chunk_size + ) + ) + groups = [KVCacheGroupSpec(["sparse"], sparse_spec)] + if full_attention: + groups.insert(0, KVCacheGroupSpec(["full"], FullAttentionSpec(**spec_kwargs))) + vllm_config = _make_vllm_config(extra_config={"block_size": chunk_size}) + vllm_config.cache_config.block_size = block_size + vllm_config.speculative_config = None + kv_cache_config = KVCacheConfig( + num_blocks=100, + kv_cache_tensors=[], + kv_cache_groups=groups, + ) + spec = MockOffloadingSpec(build_offloading_config(vllm_config, kv_cache_config)) + scheduler = OffloadingConnectorScheduler(spec, vllm_config, kv_cache_config) + request = MagicMock(request_id="sparse") + req_status = RequestOffloadState( + config=scheduler.config, + req=request, + req_context=ReqContext(req_id=request.request_id), + offloading_context=RequestOffloadingContext(policy=OffloadPolicy.CHUNK_LEVEL), + num_locally_computed_tokens=local_tokens, + ) + scheduler._req_status[request.request_id] = req_status + block_groups = [] + for group_idx, state in enumerate(req_status.group_states): + state.offload_keys = [ + make_offload_key(str(i).encode(), group_idx) for i in range(4) + ] + block_groups.append( + [ + KVCacheBlock( + 0 + if group_idx == len(groups) - 1 and i < 2 * blocks_per_chunk + else 1 + group_idx * 100 + i, + is_null=group_idx == len(groups) - 1 and i < 2 * blocks_per_chunk, + ) + for i in range(4 * blocks_per_chunk) + ] + ) + + def get_load_source(key, req_context): + chunk_idx = int(get_offload_block_hash(key)) + if get_offload_group_idx(key) == len(groups) - 1: + assert chunk_idx >= 2 # Earlier sparse keys were not loaded. + if sparse_source == "mixed": + return "host" if chunk_idx == 2 else "disk" + return sparse_source + return "host" + + scheduler.manager.get_load_source.side_effect = get_load_source + count = 4 * chunk_size - local_tokens + scheduler.update_state_after_alloc( + request, KVCacheBlocks(tuple(block_groups)), count + ) + expected: dict[CacheHitSource, int] = {} + + def expect(source: CacheHitSource, num_tokens: int) -> None: + expected[source] = expected.get(source, 0) + num_tokens + + expect( + CacheHitSource.HOST if full_attention else CacheHitSource.EXTERNAL_UNSPECIFIED, + 2 * chunk_size - local_tokens, + ) + for chunk_idx in range(2, 4): + source = ( + ("host" if chunk_idx == 2 else "disk") + if sparse_source == "mixed" + else sparse_source + ) + expect(CacheHitSource(source), chunk_size) + result = scheduler.get_external_cache_hit_sources(request, count) + assert result == expected + assert sum(result.values()) == count + assert scheduler.get_external_cache_hit_sources(request, 0) == {} + + def test_partial_lookup_requires_every_cache_group(): scheduler = _make_partial_tail_scheduler() _make_partial_tail_request(scheduler) @@ -1400,7 +1683,6 @@ def test_two_groups_full_and_sliding_window(request_runner, async_scheduling: bo # are within the window → loads blocks 1,2 expected_loaded=((0, 0), (0, 1), (0, 2), (1, 1), (1, 2)), ) - runner.manager.touch.assert_not_called() # 3 blocks are hit on GPU [0, 1, 2] diff --git a/tests/v1/kv_connector/unit/offloading_connector/utils.py b/tests/v1/kv_connector/unit/offloading_connector/utils.py index cb766c6676..9621890f3c 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/utils.py +++ b/tests/v1/kv_connector/unit/offloading_connector/utils.py @@ -57,6 +57,7 @@ from vllm.v1.kv_offload.base import ( make_offload_key, ) from vllm.v1.kv_offload.config import OffloadingConfig +from vllm.v1.metrics.cache_hit_source import CacheHitSource from vllm.v1.request import Request from vllm.v1.structured_output import StructuredOutputManager @@ -129,6 +130,7 @@ class MockOffloadingSpec(OffloadingSpec): self.manager = MagicMock(spec=OffloadingManager) self.manager.prepare_load = lambda keys, req_context: MockLoadStoreSpec(keys) self.manager.lookup.return_value = LookupResult.MISS + self.manager.get_load_source.return_value = CacheHitSource.EXTERNAL_UNSPECIFIED self.manager.get_stats.return_value = None self.manager.on_new_request.return_value = RequestOffloadingContext() self.handler = MockOffloadingWorker() diff --git a/tests/v1/kv_connector/unit/test_cache_pollution_prevention.py b/tests/v1/kv_connector/unit/test_cache_pollution_prevention.py index 21b50a1379..0515994543 100644 --- a/tests/v1/kv_connector/unit/test_cache_pollution_prevention.py +++ b/tests/v1/kv_connector/unit/test_cache_pollution_prevention.py @@ -8,7 +8,6 @@ prefix cache hash table so future requests cannot match and reuse corrupted data """ from collections.abc import Callable -from unittest.mock import Mock import pytest @@ -16,6 +15,7 @@ from vllm.v1.core.sched.scheduler import Scheduler from vllm.v1.request import Request, RequestStatus from .utils import ( + create_mock_connector, create_model_runner_output, create_request, create_scheduler, @@ -73,7 +73,7 @@ def test_invalid_blocks_evicted_prevents_cache_pollution( } # mock connector indicating sync load - fail_scheduler.connector = Mock() + fail_scheduler.connector = create_mock_connector() fail_scheduler.connector.get_num_new_matched_tokens.side_effect = ( _make_get_num_new_matched_tokens(req_num_new_matched_tokens, False) ) diff --git a/tests/v1/kv_connector/unit/test_error_propagation.py b/tests/v1/kv_connector/unit/test_error_propagation.py index 35a3749e08..ded110d55b 100644 --- a/tests/v1/kv_connector/unit/test_error_propagation.py +++ b/tests/v1/kv_connector/unit/test_error_propagation.py @@ -2,7 +2,6 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Callable -from unittest.mock import Mock import pytest @@ -10,6 +9,7 @@ from vllm.v1.core.sched.scheduler import Scheduler from vllm.v1.request import FinishReason, Request, RequestStatus from .utils import ( + create_mock_connector, create_model_runner_output, create_request, create_scheduler, @@ -56,7 +56,7 @@ def test_error_propagation_sync_load(fail_scheduler: Scheduler): request.request_id: num_external_computed_tokens, } - fail_scheduler.connector = Mock() + fail_scheduler.connector = create_mock_connector() fail_scheduler.connector.get_num_new_matched_tokens.side_effect = ( _make_get_num_new_matched_tokens(req_num_new_matched_tokens, False) ) @@ -110,7 +110,7 @@ def test_error_propagation_async_load(fail_scheduler: Scheduler): request.request_id: num_external_computed_tokens, } - fail_scheduler.connector = Mock() + fail_scheduler.connector = create_mock_connector() fail_scheduler.connector.get_loaded_kv_cache_group_ids.return_value = (0,) fail_scheduler.connector.get_num_new_matched_tokens.side_effect = ( _make_get_num_new_matched_tokens(req_num_new_matched_tokens, True) diff --git a/tests/v1/kv_connector/unit/test_invalid_blocks_correctness.py b/tests/v1/kv_connector/unit/test_invalid_blocks_correctness.py index 7ea2170584..f328332696 100644 --- a/tests/v1/kv_connector/unit/test_invalid_blocks_correctness.py +++ b/tests/v1/kv_connector/unit/test_invalid_blocks_correctness.py @@ -11,7 +11,6 @@ These tests verify correct behavior in three scenarios: """ from collections.abc import Callable -from unittest.mock import Mock import pytest @@ -19,6 +18,7 @@ from vllm.v1.core.sched.scheduler import Scheduler from vllm.v1.request import FinishReason, Request, RequestStatus from .utils import ( + create_mock_connector, create_model_runner_output, create_request, create_scheduler, @@ -83,7 +83,7 @@ def test_sync_recompute_blocks_not_freed_for_running_requests( } # mock connector indicating sync load - recompute_scheduler.connector = Mock() + recompute_scheduler.connector = create_mock_connector() recompute_scheduler.connector.get_num_new_matched_tokens.side_effect = ( _make_get_num_new_matched_tokens(req_num_new_matched_tokens, False) ) @@ -206,7 +206,7 @@ def test_sync_fail_invalid_blocks_evicted(fail_scheduler: Scheduler): } # mock connector indicating sync load - fail_scheduler.connector = Mock() + fail_scheduler.connector = create_mock_connector() fail_scheduler.connector.get_num_new_matched_tokens.side_effect = ( _make_get_num_new_matched_tokens(req_num_new_matched_tokens, False) ) @@ -323,7 +323,7 @@ def test_async_recompute_blocks_not_cached_when_invalid( } # mock connector indicating async load - recompute_scheduler.connector = Mock() + recompute_scheduler.connector = create_mock_connector() recompute_scheduler.connector.get_num_new_matched_tokens.side_effect = ( _make_get_num_new_matched_tokens(req_num_new_matched_tokens, True) ) diff --git a/tests/v1/kv_connector/unit/test_kv_load_failure_recovery.py b/tests/v1/kv_connector/unit/test_kv_load_failure_recovery.py index 3dbbfdc6b2..5dce592576 100644 --- a/tests/v1/kv_connector/unit/test_kv_load_failure_recovery.py +++ b/tests/v1/kv_connector/unit/test_kv_load_failure_recovery.py @@ -2,7 +2,6 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Callable -from unittest.mock import Mock import pytest @@ -10,6 +9,7 @@ from vllm.v1.core.sched.scheduler import Scheduler from vllm.v1.request import Request, RequestStatus from .utils import ( + create_mock_connector, create_model_runner_output, create_request, create_scheduler, @@ -41,7 +41,7 @@ def test_failed_receive_completion_honors_failure_policy(policy): scheduler = create_scheduler(create_vllm_config(kv_load_failure_policy=policy)) request = create_request(num_tokens=3 * scheduler.block_size) scheduler.add_request(request) - scheduler.connector = Mock() + scheduler.connector = create_mock_connector() scheduler.connector.get_loaded_kv_cache_group_ids.return_value = (0,) scheduler.connector.get_num_new_matched_tokens.side_effect = [ (2 * scheduler.block_size, True), @@ -103,7 +103,7 @@ def test_async_load_failure( request3.request_id: num_external_computed_tokens, } - scheduler.connector = Mock() + scheduler.connector = create_mock_connector() scheduler.connector.get_loaded_kv_cache_group_ids.return_value = (0,) scheduler.connector.get_num_new_matched_tokens.side_effect = ( _make_get_num_new_matched_tokens(req_num_new_matched_tokens, async_load=True) @@ -181,7 +181,7 @@ def test_sync_load_failure( request3.request_id: num_external_computed_tokens, } - scheduler.connector = Mock() + scheduler.connector = create_mock_connector() scheduler.connector.get_num_new_matched_tokens.side_effect = ( _make_get_num_new_matched_tokens(req_num_new_matched_tokens, async_load=False) ) @@ -262,7 +262,7 @@ def test_sync_load_failure_with_shared_blocks( request1.request_id: num_external_computed_tokens, } - scheduler.connector = Mock() + scheduler.connector = create_mock_connector() scheduler.connector.get_num_new_matched_tokens.side_effect = ( _make_get_num_new_matched_tokens(req_num_new_matched_tokens, async_load=False) ) @@ -333,7 +333,7 @@ def test_async_progressive_load_failure( request.request_id: num_external_computed_tokens, } - scheduler.connector = Mock() + scheduler.connector = create_mock_connector() scheduler.connector.get_loaded_kv_cache_group_ids.return_value = (0,) scheduler.connector.get_num_new_matched_tokens.side_effect = ( _make_get_num_new_matched_tokens(req_num_new_matched_tokens, async_load=True) diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_connector.py b/tests/v1/kv_connector/unit/test_mooncake_store_connector.py index f3ffa55a54..824f806f48 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_connector.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_connector.py @@ -38,6 +38,7 @@ from vllm.v1.kv_cache_interface import ( KVCacheTensor, MambaSpec, ) +from vllm.v1.metrics.cache_hit_source import CacheHitSource from vllm.v1.outputs import KVConnectorOutput from .utils import create_vllm_config @@ -208,6 +209,29 @@ def test_scheduler_role_initializes_store_scheduler_only(): mock_scheduler.return_value.bind_gpu_block_pool.assert_called_once_with(block_pool) +def test_scheduler_reports_mooncake_cache_source(): + vllm_config = _make_vllm_config() + kv_cache_config = _make_kv_cache_config() + + with ( + set_current_vllm_config(vllm_config), + patch( + "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store." + "connector.MooncakeStoreScheduler" + ), + ): + connector = mooncake_store_connector.MooncakeStoreConnector( + vllm_config, KVConnectorRole.SCHEDULER, kv_cache_config + ) + + # Mooncake Store is an external cache service. Its internal memory and + # local-storage policy is intentionally hidden behind that stable label. + assert connector.get_external_cache_hit_sources( + None, # type: ignore[arg-type] + 32, + ) == {CacheHitSource.EXTERNAL_UNSPECIFIED: 32} + + def test_worker_methods_delegate_to_store_worker(): vllm_config = _make_vllm_config() kv_cache_config = _make_kv_cache_config() diff --git a/tests/v1/kv_connector/unit/test_multi_connector.py b/tests/v1/kv_connector/unit/test_multi_connector.py index 67cfdc1ab0..e576e3d0d9 100644 --- a/tests/v1/kv_connector/unit/test_multi_connector.py +++ b/tests/v1/kv_connector/unit/test_multi_connector.py @@ -34,6 +34,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.nixl import ( NixlKVConnectorStats, ) from vllm.v1.kv_cache_interface import KVCacheConfig +from vllm.v1.metrics.cache_hit_source import CacheHitSource from vllm.v1.outputs import KVConnectorOutput, KVConnectorWorkerMetadata MODEL_NAME = "meta-llama/Llama-3.2-1B-Instruct" @@ -246,6 +247,18 @@ def test_multi_connector_rejects_multiple_mem_pool_contexts(mc: MultiConnector): mc.get_mem_pool_context() +def test_cache_hit_sources_delegate_to_selected_connector(mc: MultiConnector): + request = MagicMock(request_id="request") + mc._requests_to_connector[request.request_id] = 1 + sources = {CacheHitSource.DISK: 32} + mc._connectors[1].get_external_cache_hit_sources.return_value = sources + + assert mc.get_external_cache_hit_sources(request, 32) is sources + mc._connectors[1].get_external_cache_hit_sources.assert_called_once_with( + request, 32 + ) + + # Helper function to compare directories recursively def _compare_directories(dir1: Path, dir2: Path) -> bool: """Compares two directories recursively for identical content.""" diff --git a/tests/v1/kv_connector/unit/test_nixl_simple_cpu_offload.py b/tests/v1/kv_connector/unit/test_nixl_simple_cpu_offload.py index bdab91a21f..5ae235c24d 100644 --- a/tests/v1/kv_connector/unit/test_nixl_simple_cpu_offload.py +++ b/tests/v1/kv_connector/unit/test_nixl_simple_cpu_offload.py @@ -36,6 +36,7 @@ from vllm.v1.kv_cache_interface import ( KVCacheTensor, SlidingWindowSpec, ) +from vllm.v1.metrics.cache_hit_source import CacheHitSource from vllm.v1.outputs import KVConnectorOutput from vllm.v1.request import Request from vllm.v1.simple_kv_offload.metadata import ( @@ -232,6 +233,9 @@ def test_nixl_wins_load_over_cpu_offload(): assert not sched_out.has_sync_kv_loads assert mc._requests_to_connector[request.request_id] == 0 + assert mc.get_external_cache_hit_sources(request, 2 * BLOCK_SIZE) == { + CacheHitSource.P2P: 2 * BLOCK_SIZE + } meta = sched_out.kv_connector_metadata assert isinstance(meta, MultiKVConnectorMetadata) @@ -302,6 +306,9 @@ def test_cpu_offload_wins_when_nixl_has_no_match(): assert hit_tokens is not None and hit_tokens > 0 assert mc._requests_to_connector[req2.request_id] == 1 assert is_async is True + assert mc.get_external_cache_hit_sources(req2, hit_tokens) == { + CacheHitSource.HOST: hit_tokens + } @pytest.mark.parametrize("swa_enabled", [False, True], ids=["fa_only", "fa_sw"]) diff --git a/tests/v1/kv_connector/unit/utils.py b/tests/v1/kv_connector/unit/utils.py index 6f27212281..86062fc478 100644 --- a/tests/v1/kv_connector/unit/utils.py +++ b/tests/v1/kv_connector/unit/utils.py @@ -6,6 +6,7 @@ from collections.abc import Callable from dataclasses import dataclass from itertools import chain, count from typing import Any, Literal +from unittest.mock import Mock import torch @@ -42,6 +43,7 @@ from vllm.v1.kv_cache_interface import ( MambaSpec, SlidingWindowSpec, ) +from vllm.v1.metrics.cache_hit_source import CacheHitSource from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput from vllm.v1.request import Request from vllm.v1.structured_output import StructuredOutputManager @@ -49,6 +51,14 @@ from vllm.v1.structured_output import StructuredOutputManager EOS_TOKEN_ID = 50256 +def create_mock_connector() -> Mock: + connector = Mock() + connector.get_external_cache_hit_sources.side_effect = lambda request, num_tokens: ( + {CacheHitSource.EXTERNAL_UNSPECIFIED: num_tokens} if num_tokens else {} + ) + return connector + + def assert_scheduler_empty(scheduler: Scheduler): """Confirm the scheduler is "empty" - i.e. no leaks.""" # Scheduler Metadata. diff --git a/tests/v1/kv_offload/tiering/test_factory.py b/tests/v1/kv_offload/tiering/test_factory.py index f91f54bb0d..262af6cb83 100644 --- a/tests/v1/kv_offload/tiering/test_factory.py +++ b/tests/v1/kv_offload/tiering/test_factory.py @@ -17,6 +17,7 @@ import pytest from vllm.v1.kv_offload.tiering.base import SecondaryTierManager from vllm.v1.kv_offload.tiering.example.manager import ExampleSecondaryTierManager from vllm.v1.kv_offload.tiering.factory import SecondaryTierFactory +from vllm.v1.metrics.cache_hit_source import CacheHitSource # --------------------------------------------------------------------------- # Fixtures @@ -43,12 +44,22 @@ def _make_mock_args(): def test_pre_registered_tiers_can_be_imported(): """CI sentinel: registered tiers import and yield SecondaryTierManager.""" + # Derived from each tier's ``medium``; P2P has none and sets its own. + expected_sources = { + "example": CacheHitSource.HOST, + "fs": CacheHitSource.DISK, + "obj": CacheHitSource.DISK, + "p2p": CacheHitSource.P2P, + } for tier_type in SecondaryTierFactory._registry: # KVCR is an optional external dependency and may not be installed. if tier_type == "kvcr" and find_spec("kvcr") is None: continue cls = SecondaryTierFactory._registry[tier_type]() assert issubclass(cls, SecondaryTierManager) + assert cls.cache_hit_source == expected_sources.get( + tier_type, CacheHitSource.EXTERNAL_UNSPECIFIED + ) def test_example_tier_registered(): @@ -95,7 +106,8 @@ def test_create_multiple_tiers(): assert all(isinstance(tier, ExampleSecondaryTierManager) for tier in tiers) -def test_register_new_tier_type(): +@pytest.mark.parametrize("tier_type", ["custom_tier", "cpu", "nvme"]) +def test_register_new_tier_type(tier_type): """Verify that new tier types can be registered and created. This is how external projects add custom secondary tiers @@ -103,20 +115,22 @@ def test_register_new_tier_type(): """ # Register a new tier type (reuse example manager for simplicity) SecondaryTierFactory.register_tier( - "custom_tier", + tier_type, "vllm.v1.kv_offload.tiering.example.manager", "ExampleSecondaryTierManager", ) primary_kv_view, offloading_spec = _make_mock_args() tier = SecondaryTierFactory.create_secondary_tier( - {"type": "custom_tier", "custom_param": 99}, + {"type": tier_type, "custom_param": 99}, primary_kv_view, offloading_spec, ) - assert tier.tier_type == "custom_tier" + assert tier.tier_type == tier_type assert isinstance(tier, ExampleSecondaryTierManager) + # Attribution follows the tier's medium, never its registration name. + assert tier.cache_hit_source is CacheHitSource.HOST # --------------------------------------------------------------------------- diff --git a/tests/v1/kv_offload/tiering/test_tiering_offloading.py b/tests/v1/kv_offload/tiering/test_tiering_offloading.py index 63b4c2c4d0..93b49afb4f 100644 --- a/tests/v1/kv_offload/tiering/test_tiering_offloading.py +++ b/tests/v1/kv_offload/tiering/test_tiering_offloading.py @@ -53,6 +53,7 @@ from vllm.v1.kv_offload.tiering.manager import ( TieringOffloadingManager, ) from vllm.v1.kv_offload.tiering.spec import TieringOffloadingSpec +from vllm.v1.metrics.cache_hit_source import CacheHitSource _CTX = ReqContext(req_id="test") _MOCK_OFFLOADING_SPEC = MagicMock() @@ -500,8 +501,13 @@ class TestTieringOffloadingManager: # Lookup should find all chunks in primary assert count_hits(self.manager, chunks) == 3 - def test_promotion_from_secondary(self, manager_setup): + @pytest.mark.parametrize( + "source", + [CacheHitSource.EXTERNAL_UNSPECIFIED, CacheHitSource.DISK, CacheHitSource.P2P], + ) + def test_promotion_from_secondary(self, manager_setup, monkeypatch, source): """Test promotion of chunks from secondary to primary tier.""" + monkeypatch.setattr(ExampleSecondaryTierManager, "cache_hit_source", source) chunks = to_keys(range(3)) # Manually add chunks to secondary tier (simulate previous cascade) @@ -525,6 +531,27 @@ class TestTieringOffloadingManager: # Next lookup should succeed assert count_hits(self.manager, chunks) == 3 + # The request that caused the promotion retains the secondary origin. + assert all( + self.manager.get_load_source(chunk, _CTX) == source for chunk in chunks + ) + + # A later request sees the chunks as ordinary host-memory primary hits. + later_context = ReqContext(req_id="later") + self.manager.on_new_request(later_context) + assert all( + self.manager.lookup(chunk, later_context) is LookupResult.HIT + for chunk in chunks + ) + assert all( + self.manager.get_load_source(chunk, later_context) == "host" + for chunk in chunks + ) + + @pytest.mark.parametrize( + "source", + [CacheHitSource.EXTERNAL_UNSPECIFIED, CacheHitSource.DISK, CacheHitSource.P2P], + ) @pytest.mark.parametrize( ("successful_indices", "expected_results"), [ @@ -540,8 +567,9 @@ class TestTieringOffloadingManager: ids=["partial", "legacy-full-failure"], ) def test_failed_promotion_keeps_only_successful_chunks( - self, manager_setup, successful_indices, expected_results + self, manager_setup, monkeypatch, source, successful_indices, expected_results ): + monkeypatch.setattr(ExampleSecondaryTierManager, "cache_hit_source", source) chunks = to_keys(range(3)) for chunk in chunks: self.secondary_tier1.chunks[chunk] = True @@ -571,6 +599,14 @@ class TestTieringOffloadingManager: assert [ self.primary_tier.lookup(chunk, _CTX) for chunk in chunks ] == expected_results + expected_source_keys = ( + set() + if successful_indices is None + else {chunks[i] for i in successful_indices} + ) + assert set(self.manager._request_load_sources[_CTX.req_id]) == ( + expected_source_keys + ) def test_lookup_reports_sync_delay_for_resolved_lookups(self, manager_setup): """Resolved lookups report one sync delay sample per tier and chunk.""" @@ -1315,6 +1351,7 @@ class TestTieringOffloadingManager: # Orchestrator state cleared. assert self.manager._jobs == {} assert self.manager._pending_load_submissions == {} + assert self.manager._request_load_sources == {} assert set(self.manager._req_state) == {_CTX.req_id, rl_ctx.req_id} assert self.manager._processed_jobs_this_step is False diff --git a/tests/v1/metrics/test_stats.py b/tests/v1/metrics/test_stats.py index 0de74f0faa..c51c223d28 100644 --- a/tests/v1/metrics/test_stats.py +++ b/tests/v1/metrics/test_stats.py @@ -1,10 +1,13 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project + from vllm.v1.core.sched.output import ScheduledEncoderInputStats, SchedulerOutput from vllm.v1.engine import EngineCoreOutputs, FinishReason +from vllm.v1.metrics.cache_hit_source import CacheHitSource from vllm.v1.metrics.stats import ( IterationStats, PrefillStats, + PrefixCacheStats, PromptTokenStats, RequestStateStats, SchedulerIterationDetails, @@ -45,6 +48,96 @@ def test_scheduler_iteration_details_serialization(): assert decoded.scheduler_stats.iteration_details == iteration_details +def test_prefix_cache_stats_hits_by_source_accumulate(): + stats = PrefixCacheStats() + stats.record( + num_tokens=8, + num_hits=5, + preempted=False, + hits_by_source={CacheHitSource.HOST: 3, CacheHitSource.DISK: 2}, + ) + stats.record( + num_tokens=8, + num_hits=6, + preempted=False, + hits_by_source={CacheHitSource.P2P: 5, CacheHitSource.HOST: 1}, + ) + # Zero counts never create an entry. + stats.record( + num_tokens=8, + num_hits=4, + preempted=False, + hits_by_source={CacheHitSource.DISK: 4, CacheHitSource.P2P: 0}, + ) + + assert stats.hits_by_source == { + CacheHitSource.HOST: 4, + CacheHitSource.DISK: 6, + CacheHitSource.P2P: 5, + } + assert sum(stats.hits_by_source.values()) == stats.hits == 15 + + +def test_prefix_cache_stats_record_hits_by_source(): + stats = PrefixCacheStats() + stats.record(num_tokens=32, num_hits=16, preempted=False) + stats.record( + num_tokens=32, + num_hits=24, + preempted=False, + hits_by_source={CacheHitSource.HOST: 16, CacheHitSource.DISK: 8}, + ) + # Preempted re-admissions never contribute to hits or the split. + stats.record( + num_tokens=32, + num_hits=32, + preempted=True, + hits_by_source={CacheHitSource.P2P: 32}, + ) + # A connector miscount is reported as unspecified, not dropped or raised. + stats.record( + num_tokens=8, + num_hits=8, + preempted=False, + hits_by_source={CacheHitSource.HOST: 4}, + ) + + assert stats.requests == 3 + assert stats.hits == 48 + assert stats.preempted_hits == 32 + assert stats.hits_by_source == { + CacheHitSource.HOST: 16, + CacheHitSource.DISK: 8, + CacheHitSource.EXTERNAL_UNSPECIFIED: 8, + } + assert sum(stats.hits_by_source.values()) == 48 - 16 + + +def test_prefix_cache_stats_hits_by_source_serialization(): + connector_stats = PrefixCacheStats() + connector_stats.record( + num_tokens=16, + num_hits=8, + preempted=False, + hits_by_source={CacheHitSource.P2P: 4, CacheHitSource.HOST: 4}, + ) + outputs = EngineCoreOutputs( + scheduler_stats=SchedulerStats( + connector_prefix_cache_stats=connector_stats, + ) + ) + + encoded = MsgpackEncoder().encode(outputs) + decoded = MsgpackDecoder(EngineCoreOutputs).decode(encoded) + + assert decoded.scheduler_stats is not None + assert decoded.scheduler_stats.connector_prefix_cache_stats is not None + assert decoded.scheduler_stats.connector_prefix_cache_stats.hits_by_source == { + CacheHitSource.P2P: 4, + CacheHitSource.HOST: 4, + } + + def test_compute_iteration_details_includes_encoder_stats(): scheduler_output = SchedulerOutput.make_empty() scheduler_output.scheduled_encoder_input_stats = ScheduledEncoderInputStats( diff --git a/tests/v1/simple_kv_offload/test_scheduler.py b/tests/v1/simple_kv_offload/test_scheduler.py index dee5446c61..08759bc1e3 100644 --- a/tests/v1/simple_kv_offload/test_scheduler.py +++ b/tests/v1/simple_kv_offload/test_scheduler.py @@ -23,6 +23,7 @@ from vllm.config import ( ) from vllm.config.cache import MambaCacheMode from vllm.config.kv_events import KVEventsConfig +from vllm.distributed.kv_transfer.kv_connector.v1 import KVConnectorRole from vllm.distributed.kv_transfer.kv_connector.v1.simple_cpu_offload_connector import ( SimpleCPUOffloadConnector, ) @@ -53,6 +54,7 @@ from vllm.v1.kv_cache_interface import ( MambaSpec, SlidingWindowSpec, ) +from vllm.v1.metrics.cache_hit_source import CacheHitSource from vllm.v1.outputs import KVConnectorOutput from vllm.v1.request import Request from vllm.v1.simple_kv_offload.manager import SimpleCPUOffloadScheduler @@ -213,6 +215,37 @@ class SchedulerFixture: num_groups: int = 1 +@pytest.mark.parametrize( + ("backend", "source"), + [("cpu", CacheHitSource.HOST), ("disk", CacheHitSource.DISK)], +) +def test_connector_reports_configured_cache_source( + tmp_path, backend: str, source: CacheHitSource +): + kv_cache_config = _make_kv_cache_config(num_blocks=4) + vllm_config = _make_vllm_config() + extra_config = { + "cpu_bytes_to_use_per_rank": _BYTES_PER_BLOCK * 4, + "kv_offload_backend": backend, + } + if backend == "disk": + extra_config.update( + disk_path=str(tmp_path / "kv-offload.bin"), + disk_capacity_bytes=_BYTES_PER_BLOCK * 4, + ) + assert vllm_config.kv_transfer_config is not None + vllm_config.kv_transfer_config.kv_connector_extra_config = extra_config + + connector = SimpleCPUOffloadConnector( + vllm_config, + KVConnectorRole.SCHEDULER, + kv_cache_config, + ) + + assert connector.get_external_cache_hit_sources(None, 32) == {source: 32} # type: ignore[arg-type] + assert connector.get_external_cache_hit_sources(None, 0) == {} # type: ignore[arg-type] + + def make_scheduler( num_cpu_blocks: int = 8, num_gpu_blocks: int = 16, diff --git a/tests/v1/worker/test_gpu_profiler.py b/tests/v1/worker/test_gpu_profiler.py index e272ed5106..4b5d2178c4 100644 --- a/tests/v1/worker/test_gpu_profiler.py +++ b/tests/v1/worker/test_gpu_profiler.py @@ -114,6 +114,7 @@ def test_torch_profiler_records_each_profile_round(tmp_path): profiler="torch", torch_profiler_dir=str(tmp_path), torch_profiler_dump_cuda_time_total=False, + torch_profiler_with_stack=False, ), worker_name="worker", local_rank=1, diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/base.py b/vllm/distributed/kv_transfer/kv_connector/v1/base.py index 61de1802d5..bb9d4597ce 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/base.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/base.py @@ -51,6 +51,7 @@ import torch from vllm.logger import init_logger from vllm.v1.attention.backend import AttentionMetadata from vllm.v1.core.sched.output import SchedulerOutput +from vllm.v1.metrics.cache_hit_source import CacheHitSource from vllm.v1.outputs import KVConnectorOutput if TYPE_CHECKING: @@ -179,6 +180,10 @@ class KVConnectorWorkerMetadata(ABC): class KVConnectorBase_V1(ABC): """Base class for KV connectors.""" + # Source label for this connector's external hits. Subclasses must set + # HOST/DISK/P2P or override get_external_cache_hit_sources(). + _cache_hit_source = CacheHitSource.EXTERNAL_UNSPECIFIED + @property def supports_divergent_local_hybrid_hits(self) -> bool: """Whether external hits can complete divergent local hybrid hits. @@ -527,6 +532,24 @@ class KVConnectorBase_V1(ABC): """ pass + def get_external_cache_hit_sources( + self, + request: "Request", + num_external_tokens: int, + ) -> dict[CacheHitSource, int]: + """Split ``num_external_tokens`` by the cache tier that supplied them. + + Called after :meth:`update_state_after_alloc`, so the load plan is + known. Counts must sum to ``num_external_tokens``; a mismatch is + reported as ``external_unspecified``. Only non-zero counts are + included. + + Default: all tokens under ``_cache_hit_source``. + """ + if num_external_tokens == 0: + return {} + return {self._cache_hit_source: num_external_tokens} + @abstractmethod def update_state_after_alloc( self, request: "Request", blocks: "KVCacheBlocks", num_external_tokens: int diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py index 172d41efdb..5f064259e8 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py @@ -67,6 +67,7 @@ from vllm.v1.kv_cache_interface import ( SlidingWindowMLASpec, SlidingWindowSpec, ) +from vllm.v1.metrics.cache_hit_source import CacheHitSource from vllm.v1.request import RequestStatus from vllm.v1.worker.block_table import BlockTable from vllm.v1.worker.utils import select_common_block_size @@ -700,6 +701,8 @@ class MooncakeConnectorMetadata(KVConnectorMetadata): class MooncakeConnector(KVConnectorBase_V1, SupportsHMA): + _cache_hit_source = CacheHitSource.P2P + def __init__( self, vllm_config: VllmConfig, diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py index 1dc6f8459a..167c7ef3c2 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py @@ -99,6 +99,7 @@ from vllm.v1.kv_cache_interface import ( SlidingWindowSpec, is_full_attention_spec, ) +from vllm.v1.metrics.cache_hit_source import CacheHitSource from vllm.v1.outputs import KVConnectorOutput from vllm.v1.request import RequestStatus @@ -222,6 +223,8 @@ def resolve_moriio_transfer_ack( class MoRIIOConnector(KVConnectorBase_V1, SupportsHMA): + _cache_hit_source = CacheHitSource.P2P + @property def supports_divergent_local_hybrid_hits(self) -> bool: # The READ path always transfers the recurrent-state slot, including diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py index 752acb01ec..1ded62adda 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py @@ -31,6 +31,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.metrics import ( from vllm.logger import init_logger from vllm.v1.attention.backend import AttentionMetadata from vllm.v1.core.sched.output import SchedulerOutput +from vllm.v1.metrics.cache_hit_source import CacheHitSource from vllm.v1.outputs import KVConnectorOutput if TYPE_CHECKING: @@ -451,6 +452,18 @@ class MultiConnector(KVConnectorBase_V1, SupportsHMA): to_return = (toks, load_async) return to_return + def get_external_cache_hit_sources( + self, + request: "Request", + num_external_tokens: int, + ) -> dict[CacheHitSource, int]: + chosen_connector = self._requests_to_connector.get(request.request_id) + if chosen_connector is None: + return super().get_external_cache_hit_sources(request, num_external_tokens) + return self._connectors[chosen_connector].get_external_cache_hit_sources( + request, num_external_tokens + ) + def update_state_after_alloc( self, request: "Request", blocks: "KVCacheBlocks", num_external_tokens: int ): diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py index 27b4d410f0..6e43356c0d 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py @@ -58,6 +58,7 @@ from vllm.forward_context import ForwardContext from vllm.logger import init_logger from vllm.v1.attention.backend import AttentionMetadata from vllm.v1.core.sched.output import SchedulerOutput +from vllm.v1.metrics.cache_hit_source import CacheHitSource from vllm.v1.outputs import KVConnectorOutput if TYPE_CHECKING: @@ -77,6 +78,8 @@ logger = init_logger(__name__) class NixlBaseConnector(KVConnectorBase_V1, SupportsHMA): """Base connector with common logic shared by pull and push modes.""" + _cache_hit_source = CacheHitSource.P2P + @property def supports_divergent_local_hybrid_hits(self) -> bool: return True diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py index 64b3ad5bb2..806267dcf3 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import time +from collections import Counter from collections.abc import Iterable, Sequence from dataclasses import dataclass, field, replace from itertools import chain, islice @@ -56,6 +57,7 @@ from vllm.v1.kv_offload.base import ( TierMatcher, make_offload_key, ) +from vllm.v1.metrics.cache_hit_source import CacheHitSource from vllm.v1.outputs import KVConnectorOutput from vllm.v1.request import Request, RequestStatus @@ -375,6 +377,9 @@ class RequestOffloadState: # Fine-grained token boundary selected beyond the last complete offload # chunk. It is consumed when the corresponding load is scheduled. partial_tail_boundary: int | None = None + # Loaded attention keys cover their token positions. Recurrent state + # supports the whole reused external prefix. + load_key_ranges: list[tuple[int, int, OffloadKey]] = field(default_factory=list) # True once on_request_finished has been signaled to the manager. finished_signaled: bool = False @@ -587,8 +592,9 @@ class OffloadingConnectorScheduler: sliding_window_groups.sort(key=_sliding_window_sort_key, reverse=True) # used by _lookup + self._full_attention_groups: tuple[int, ...] = tuple(full_attention_groups) self._sliding_window_groups: tuple[int, ...] = tuple(sliding_window_groups) - self._lookup_groups = tuple(full_attention_groups) + self._sliding_window_groups + self._lookup_groups = self._full_attention_groups + self._sliding_window_groups self._mamba_align_size: int | None = resolve_mamba_align_size( spec, kv_cache_config ) @@ -1090,6 +1096,53 @@ class OffloadingConnectorScheduler: return num_hit_tokens, bool(num_hit_tokens) + def get_external_cache_hit_sources( + self, + request: Request, + num_external_tokens: int, + ) -> dict[CacheHitSource, int]: + """Split the accepted external hit by the tier each loaded key came from. + + Attention keys cover only their loaded token ranges; recurrent state + covers the whole reused prefix. Overlaps report the outermost tier. + Uncovered tokens, including omitted sliding-window history, are + ``external_unspecified``. + """ + sources: dict[CacheHitSource, int] = {} + if num_external_tokens == 0: + return sources + req_status = self._req_status[request.request_id] + start = req_status.num_locally_computed_tokens + end = start + num_external_tokens + + events: list[tuple[int, int, CacheHitSource]] = [ + (end, 0, CacheHitSource.EXTERNAL_UNSPECIFIED) + ] + for lo, hi, key in req_status.load_key_ranges: + lo, hi = max(start, lo), min(end, hi) + if lo >= hi: + continue + source = CacheHitSource( + self.manager.get_load_source(key, req_status.req_context) + ) + events.extend(((lo, 1, source), (hi, -1, source))) + + # Sweep endpoints instead of rescanning every range for each interval. + active: Counter[CacheHitSource] = Counter() + position = start + for boundary, delta, source in sorted(events): + if boundary > position: + tiers = [tier for tier, count in active.items() if count > 0] + tier = ( + CacheHitSource.outermost(tiers) + if tiers + else CacheHitSource.EXTERNAL_UNSPECIFIED + ) + sources[tier] = sources.get(tier, 0) + boundary - position + active[source] += delta + position = boundary + return sources + def update_state_after_alloc( self, request: Request, blocks: KVCacheBlocks, num_external_tokens: int ): @@ -1105,6 +1158,7 @@ class OffloadingConnectorScheduler: assert partial_tail_boundary == num_cached_tokens keys_to_load: list[OffloadKey] = [] + req_status.load_key_ranges.clear() dst_block_ids: list[int] = [] # per group group_sizes: list[int] = [] @@ -1157,9 +1211,9 @@ class OffloadingConnectorScheduler: ) end_chunk_idx = num_chunks - (partial_tail_boundary is not None) assert len(offload_keys) >= end_chunk_idx - keys_to_load.extend(offload_keys[start_chunk_idx:end_chunk_idx]) + group_keys_to_load = offload_keys[start_chunk_idx:end_chunk_idx] if partial_tail_boundary is not None: - keys_to_load.append( + group_keys_to_load.append( self._make_boundary_key( request, group_config.group_idx, @@ -1167,6 +1221,24 @@ class OffloadingConnectorScheduler: req_status.req_context, ) ) + keys_to_load.extend(group_keys_to_load) + if isinstance(group_config.kv_cache_spec, MambaSpec): + req_status.load_key_ranges.extend( + (num_locally_computed_tokens, num_cached_tokens, key) + for key in group_keys_to_load + ) + else: + for chunk_idx, key in enumerate( + group_keys_to_load, start_chunk_idx + ): + range_start = max( + load_start_gpu_block_idx * tokens_per_block, + chunk_idx * tokens_per_chunk, + ) + range_end = min( + num_cached_tokens, (chunk_idx + 1) * tokens_per_chunk + ) + req_status.load_key_ranges.append((range_start, range_end, key)) dst_block_ids.extend( block.block_id diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py index def6e1e1f3..af3e113b5f 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py @@ -45,6 +45,7 @@ from vllm.v1.core.kv_cache_manager import KVCacheBlocks from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.kv_offload.factory import OffloadingSpecFactory +from vllm.v1.metrics.cache_hit_source import CacheHitSource from vllm.v1.outputs import KVConnectorOutput from vllm.v1.request import Request @@ -179,6 +180,16 @@ class OffloadingConnector(KVConnectorBase_V1, SupportsHMA): bound = min(per_group_hits[group_id] for group_id in self._bounding_group_ids) return max(0, bound - num_computed_tokens) + def get_external_cache_hit_sources( + self, + request: "Request", + num_external_tokens: int, + ) -> dict[CacheHitSource, int]: + assert self.connector_scheduler is not None + return self.connector_scheduler.get_external_cache_hit_sources( + request, num_external_tokens + ) + def update_state_after_alloc( self, request: "Request", blocks: "KVCacheBlocks", num_external_tokens: int ): diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py index 9124d9eac1..0bd8ebe3cd 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py @@ -23,6 +23,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.metrics import ( ) from vllm.logger import init_logger from vllm.v1.core.sched.output import SchedulerOutput +from vllm.v1.metrics.cache_hit_source import CacheHitSource from vllm.v1.outputs import KVConnectorOutput from vllm.v1.simple_kv_offload.manager import ( BoundaryStoreStats, @@ -109,6 +110,9 @@ class SimpleCPUOffloadConnector(KVConnectorBase_V1, SupportsHMA): f"expected one of {VALID_KV_OFFLOAD_BACKENDS}" ) disk_mode = kv_offload_backend == "disk" + self._cache_hit_source = ( + CacheHitSource.DISK if disk_mode else CacheHitSource.HOST + ) disk_path = extra_config.get("disk_path", None) or None disk_capacity_bytes = int( diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 0f2425312c..8bdf91da86 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -1261,10 +1261,20 @@ class Scheduler(SchedulerInterface): self.connector_prefix_cache_stats is not None and connector_prefix_cache_queries != 0 ): + preempted = request.num_preemptions > 0 self.connector_prefix_cache_stats.record( num_tokens=connector_prefix_cache_queries, num_hits=connector_prefix_cache_hits, - preempted=request.num_preemptions > 0, + preempted=preempted, + # The connector knows the tier only after the load + # plan above exists. + hits_by_source=( + self.connector.get_external_cache_hit_sources( + request, connector_prefix_cache_hits + ) + if connector_prefix_cache_hits and not preempted + else None + ), ) # Record at admission so unscheduled lookups are not counted. diff --git a/vllm/v1/kv_offload/base.py b/vllm/v1/kv_offload/base.py index f8e92abed6..398bc269e6 100644 --- a/vllm/v1/kv_offload/base.py +++ b/vllm/v1/kv_offload/base.py @@ -18,6 +18,7 @@ if TYPE_CHECKING: from vllm.v1.kv_hints import KvHintsEnvelope from vllm.v1.kv_offload.config import OffloadingConfig +from vllm.v1.metrics.cache_hit_source import CacheHitSource # `OffloadKey` identifies an offloaded block. It combines a block hash with # its KV cache group index, encoded as raw bytes to avoid tuple GC overhead. @@ -248,6 +249,18 @@ class OffloadingManager(ABC): """ pass + def get_load_source( + self, key: OffloadKey, req_context: ReqContext + ) -> CacheHitSource: + """Return the cache tier that supplied a successful lookup. + + This is queried only after ``lookup`` resolves to ``HIT``. Composing + managers can override it to retain the origin across asynchronous + promotion into a primary tier. The default keeps out-of-tree managers + compatible while making unknown provenance explicit. + """ + return CacheHitSource.EXTERNAL_UNSPECIFIED + @abstractmethod def prepare_load( self, diff --git a/vllm/v1/kv_offload/cpu/manager.py b/vllm/v1/kv_offload/cpu/manager.py index e42cb47f25..2c08ac897c 100644 --- a/vllm/v1/kv_offload/cpu/manager.py +++ b/vllm/v1/kv_offload/cpu/manager.py @@ -27,6 +27,7 @@ from vllm.v1.kv_offload.cpu.common import ( ) from vllm.v1.kv_offload.cpu.policies.base import CachePolicy, ChunkStatus from vllm.v1.kv_offload.cpu.policies.factory import CachePolicyFactory +from vllm.v1.metrics.cache_hit_source import CacheHitSource @dataclass(slots=True) @@ -193,6 +194,12 @@ class CPUOffloadingManager(OffloadingManager): return LookupResult.HIT_PENDING return LookupResult.HIT + @override + def get_load_source( + self, key: OffloadKey, req_context: ReqContext + ) -> CacheHitSource: + return CacheHitSource.HOST + @override def prepare_load( self, diff --git a/vllm/v1/kv_offload/tiering/base.py b/vllm/v1/kv_offload/tiering/base.py index a90d33eb11..b855b6b000 100644 --- a/vllm/v1/kv_offload/tiering/base.py +++ b/vllm/v1/kv_offload/tiering/base.py @@ -24,6 +24,7 @@ from vllm.v1.kv_offload.base import ( from vllm.v1.kv_offload.tiering.backpressure import ( BackpressureDetector, ) +from vllm.v1.metrics.cache_hit_source import CacheHitSource if TYPE_CHECKING: from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import ( @@ -141,6 +142,17 @@ class SecondaryTierManager(ABC): """ medium: ClassVar[Medium | None] = None + # Cache tier reported for hits served from this tier. Derived from + # ``medium`` unless a subclass sets it; ``None`` medium means unknown. + cache_hit_source: ClassVar[CacheHitSource] = CacheHitSource.EXTERNAL_UNSPECIFIED + + def __init_subclass__(cls, **kwargs: Any) -> None: + super().__init_subclass__(**kwargs) + if "cache_hit_source" not in cls.__dict__ and cls.medium is not None: + cls.cache_hit_source = { + Medium.CPU: CacheHitSource.HOST, + Medium.STORAGE: CacheHitSource.DISK, + }[cls.medium] def __init__( self, diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index ce518042b6..83ff2e2c12 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -54,6 +54,7 @@ from vllm.v1.kv_offload.tiering.base import ( TransferJob, ) from vllm.v1.kv_offload.tiering.metrics import TieringMetricsTracker +from vllm.v1.metrics.cache_hit_source import CacheHitSource logger = init_logger(__name__) @@ -232,6 +233,11 @@ class TieringOffloadingManager(OffloadingManager): # complete_store(), since complete_store() can still submit cascades. self._req_state: dict[str, RequestState] = {} + # Preserve the original tier for this request's cache-hit metrics, + # even after its KV blocks are promoted into host memory + # (otherwise a disk or P2P hit would be mislabeled as a host hit). + self._request_load_sources: dict[str, dict[OffloadKey, CacheHitSource]] = {} + # Cached ParentManager wrappers for each secondary tier. self._tier_parents: dict[SecondaryTierManager, _SecondaryTierFacingParent] = { tier: _SecondaryTierFacingParent(self, tier_idx) @@ -293,6 +299,13 @@ class TieringOffloadingManager(OffloadingManager): successful_keys = () failed_keys = transfer_job.keys + load_sources = self._request_load_sources.get(transfer_job.req_context.req_id) + if load_sources is not None: + source = self.secondary_tiers[job_metadata.tier_idx].cache_hit_source + for key in failed_keys: + if load_sources.get(key) == source: + del load_sources[key] + if successful_keys: self.primary_tier.complete_write( successful_keys, @@ -464,6 +477,15 @@ class TieringOffloadingManager(OffloadingManager): return LookupResult.RETRY return LookupResult.MISS + @override + def get_load_source( + self, key: OffloadKey, req_context: ReqContext + ) -> CacheHitSource: + load_sources = self._request_load_sources.get(req_context.req_id) + if load_sources is not None and key in load_sources: + return load_sources[key] + return self.primary_tier.get_load_source(key, req_context) + def _initiate_promotion( self, tier_idx: int, @@ -501,6 +523,10 @@ class TieringOffloadingManager(OffloadingManager): store_spec = primary_write_result.store_spec assert isinstance(store_spec, CPULoadStoreSpec) + load_sources = self._request_load_sources.setdefault(req_context.req_id, {}) + source = self.secondary_tiers[tier_idx].cache_hit_source + for promoted_key in primary_write_result.keys_to_store: + load_sources[promoted_key] = source # Defer submit_load to on_schedule_end(). Group by (tier, request) so # each request's chunks are submitted as one batched job per tier. tier_pending = self._pending_load_submissions.setdefault(tier_idx, {}) @@ -848,6 +874,7 @@ class TieringOffloadingManager(OffloadingManager): continue tier.on_request_finished(state.req_context) self._metrics.on_request_finished(state.req_context) + self._request_load_sources.pop(req_id, None) del self._req_state[req_id] @override @@ -930,6 +957,7 @@ class TieringOffloadingManager(OffloadingManager): # reset below invalidates; their submit_load() has not yet been # called so no tier I/O is touching that memory. self._pending_load_submissions.clear() + self._request_load_sources.clear() self._metrics.assert_idle() finished_req_ids = [] diff --git a/vllm/v1/kv_offload/tiering/p2p/manager.py b/vllm/v1/kv_offload/tiering/p2p/manager.py index 59b9c32069..f759a4481e 100644 --- a/vllm/v1/kv_offload/tiering/p2p/manager.py +++ b/vllm/v1/kv_offload/tiering/p2p/manager.py @@ -11,7 +11,7 @@ import time import uuid from collections.abc import Iterable, Sequence from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, ClassVar from typing_extensions import override @@ -35,6 +35,7 @@ from vllm.v1.kv_offload.tiering.base import ( from vllm.v1.kv_offload.tiering.p2p.control import ControlTransport, ZmqTransport from vllm.v1.kv_offload.tiering.p2p.data import DataTransport, NixlTransport from vllm.v1.kv_offload.tiering.p2p.session import P2PSession +from vllm.v1.metrics.cache_hit_source import CacheHitSource if TYPE_CHECKING: from vllm.v1.kv_offload.base import OffloadingSpec @@ -206,6 +207,8 @@ class P2PSecondaryTierManager(SecondaryTierManager): and existing sessions are polled even when no requests are scheduled. """ + cache_hit_source: ClassVar[CacheHitSource] = CacheHitSource.P2P + def __init__( self, offloading_spec: OffloadingSpec, diff --git a/vllm/v1/metrics/cache_hit_source.py b/vllm/v1/metrics/cache_hit_source.py new file mode 100644 index 0000000000..b6fc5270e7 --- /dev/null +++ b/vllm/v1/metrics/cache_hit_source.py @@ -0,0 +1,32 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +from collections.abc import Iterable +from enum import Enum + + +class CacheHitSource(str, Enum): + """Where cached prompt tokens' KV came from. + + Members are listed in the order blocks are offloaded: accelerator, then + host memory, then secondary tiers. Per-source token counts are plain + ``dict[CacheHitSource, int]`` mappings holding only non-zero entries. + """ + + DEVICE = "device" + HOST = "host" + P2P = "p2p" + DISK = "disk" + # Fallback: uninstrumented connector or remote store. Last, so mixing + # with a known tier never reports the known one. + EXTERNAL_UNSPECIFIED = "external_unspecified" + + @classmethod + def outermost(cls, sources: Iterable[CacheHitSource]) -> CacheHitSource: + """The tier in ``sources`` farthest from the accelerator.""" + return max(sources, key=_SOURCE_ORDER.__getitem__) + + +_SOURCE_ORDER = {source: rank for rank, source in enumerate(CacheHitSource)} diff --git a/vllm/v1/metrics/loggers.py b/vllm/v1/metrics/loggers.py index 681f80ce29..7ff1d2993c 100644 --- a/vllm/v1/metrics/loggers.py +++ b/vllm/v1/metrics/loggers.py @@ -23,6 +23,7 @@ from vllm.logger import init_logger from vllm.plugins import STAT_LOGGER_PLUGINS_GROUP, load_plugins_by_group from vllm.v1.engine import FinishReason from vllm.v1.metrics.buckets import histogram_buckets +from vllm.v1.metrics.cache_hit_source import CacheHitSource from vllm.v1.metrics.perf import PerfMetricsLogging, PerfMetricsProm from vllm.v1.metrics.prometheus import unregister_vllm_metrics from vllm.v1.metrics.stats import ( @@ -752,6 +753,23 @@ class PrometheusStatLogger(AggregateStatLoggerBase): counter_prompt_tokens_cached, per_engine_labelvalues ) + self.counter_prompt_tokens_cached_by_source = self._counter_cls( + name="vllm:prompt_tokens_cached_by_source", + documentation=( + "Prefix-cache hit tokens by the cache tier that supplied them: " + "device (local HBM), host (offloaded to DRAM), disk, p2p " + "(transferred from another vLLM instance) or " + "external_unspecified. Counted at admission; sums to " + "vllm:prefix_cache_hits + vllm:external_prefix_cache_hits." + ), + labelnames=labelnames + ["source"], + ) + for source in CacheHitSource: + for labelvalues in per_engine_labelvalues.values(): + self.counter_prompt_tokens_cached_by_source.labels( + *labelvalues, source.value + ) + counter_generation_tokens = self._counter_cls( name="vllm:generation_tokens", documentation="Number of generation tokens processed.", @@ -1111,13 +1129,23 @@ class PrometheusStatLogger(AggregateStatLoggerBase): scheduler_stats.prefix_cache_stats.hits ) - if scheduler_stats.connector_prefix_cache_stats is not None: + labelvalues = self.per_engine_labelvalues[engine_idx] + self.counter_prompt_tokens_cached_by_source.labels( + *labelvalues, CacheHitSource.DEVICE.value + ).inc(scheduler_stats.prefix_cache_stats.hits) + + connector_stats = scheduler_stats.connector_prefix_cache_stats + if connector_stats is not None: self.counter_connector_prefix_cache_queries[engine_idx].inc( - scheduler_stats.connector_prefix_cache_stats.queries + connector_stats.queries ) self.counter_connector_prefix_cache_hits[engine_idx].inc( - scheduler_stats.connector_prefix_cache_stats.hits + connector_stats.hits ) + for tier, num_tokens in connector_stats.hits_by_source.items(): + self.counter_prompt_tokens_cached_by_source.labels( + *labelvalues, tier.value + ).inc(num_tokens) if scheduler_stats.spec_decoding_stats is not None: self.spec_decoding_prom.observe( diff --git a/vllm/v1/metrics/stats.py b/vllm/v1/metrics/stats.py index 4db98884fe..0046dcefa9 100644 --- a/vllm/v1/metrics/stats.py +++ b/vllm/v1/metrics/stats.py @@ -8,9 +8,13 @@ from typing import TYPE_CHECKING, Any import vllm.envs as envs from vllm.compilation.cuda_graph import CUDAGraphStat +from vllm.logger import init_logger +from vllm.v1.metrics.cache_hit_source import CacheHitSource from vllm.v1.metrics.perf import PerfStats from vllm.v1.spec_decode.metrics import SpecDecodingStats +logger = init_logger(__name__) + if TYPE_CHECKING: from vllm.v1.engine import EngineCoreEvent, EngineCoreOutput, FinishReason @@ -128,19 +132,42 @@ class PrefixCacheStats(BaseCacheStats): preempted_hits: int = 0 """The `hits` number for preempted requests.""" + hits_by_source: dict[CacheHitSource, int] = field(default_factory=dict) + """`hits` split by the cache tier that supplied them (connector stats only).""" - def record(self, num_tokens: int, num_hits: int, preempted: bool) -> None: + def record( + self, + num_tokens: int, + num_hits: int, + preempted: bool, + hits_by_source: dict[CacheHitSource, int] | None = None, + ) -> None: """Aggregate request information into the stats.""" if preempted: # Previously preempted request self.preempted_requests += 1 self.preempted_queries += num_tokens self.preempted_hits += num_hits - else: - # New request - self.requests += 1 - self.queries += num_tokens - self.hits += num_hits + return + # New request + self.requests += 1 + self.queries += num_tokens + self.hits += num_hits + if hits_by_source is None: + return + if sum(hits_by_source.values()) != num_hits: + logger.warning_once( + "Connector attributed %s for %d cached tokens; reporting them as %s.", + str(hits_by_source), + num_hits, + CacheHitSource.EXTERNAL_UNSPECIFIED.value, + ) + hits_by_source = {CacheHitSource.EXTERNAL_UNSPECIFIED: num_hits} + for source, num_source_hits in hits_by_source.items(): + if num_source_hits: + self.hits_by_source[source] = ( + self.hits_by_source.get(source, 0) + num_source_hits + ) @dataclass