iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Metrics] Expose cached prompt tokens by cache tier (#56318)

Browse Source 
Signed-off-by: Cam Quilici <cjquilici@gmail.com>
Signed-off-by: Cam Quilici <cameron@semianalysis.com>
Co-authored-by: Cam Quilici <cameron@semianalysis.com>
Co-authored-by: Nick Hill <nickhill123@gmail.com>
Co-authored-by: Yifan Qiao <yifanqiao@inferact.ai>
This commit is contained in:
Cameron Quilici
2026-10-06 14:52:05 -04:00
committed by GitHub
co-authored by Cam Quilici Nick Hill Yifan Qiao
parent ad0f67a7ec
commit e37e51dd24
36 changed files with 1005 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rust/src/engine-core-client/src/metrics.rs
+23 -5
View File
@@ -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.
rust/src/engine-core-client/src/protocol/stats.rs
+59
View File
@@ -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<PerfStats>,
}
#[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");
}
}
rust/src/metrics/src/scheduler.rs
+16
View File
@@ -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<EngineLabels, U64Counter>,
pub external_prefix_cache_queries: Family<EngineLabels, U64Counter>,
pub external_prefix_cache_hits: Family<EngineLabels, U64Counter>,
pub prompt_tokens_cached_by_source: Family<CacheHitSourceLabels, U64Counter>,
// Speculative decoding counters.
pub spec_decode_num_drafts: Family<EngineLabels, U64Counter>,
@@ -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,
tests/v1/core/test_scheduler.py
+61 -1
View File
@@ -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
tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py
+285 -3
View File
@@ -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]
tests/v1/kv_connector/unit/offloading_connector/utils.py
+2
View File
@@ -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()
tests/v1/kv_connector/unit/test_cache_pollution_prevention.py
+2 -2
View File
@@ -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)
)
tests/v1/kv_connector/unit/test_error_propagation.py
+3 -3
View File
@@ -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)
tests/v1/kv_connector/unit/test_invalid_blocks_correctness.py
+4 -4
View File
@@ -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)
)
tests/v1/kv_connector/unit/test_kv_load_failure_recovery.py
+6 -6
View File
@@ -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)
tests/v1/kv_connector/unit/test_mooncake_store_connector.py
+24
View File
@@ -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()
tests/v1/kv_connector/unit/test_multi_connector.py
+13
View File
@@ -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."""
tests/v1/kv_connector/unit/test_nixl_simple_cpu_offload.py
+7
View File
@@ -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"])
tests/v1/kv_connector/unit/utils.py
+10
View File
@@ -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.
tests/v1/kv_offload/tiering/test_factory.py
+18 -4
View File
@@ -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
# ---------------------------------------------------------------------------
tests/v1/kv_offload/tiering/test_tiering_offloading.py
+39 -2
View File
@@ -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
tests/v1/metrics/test_stats.py
+93
View File
@@ -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(
tests/v1/simple_kv_offload/test_scheduler.py
+33
View File
@@ -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,
tests/v1/worker/test_gpu_profiler.py
+1
View File
@@ -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,
vllm/distributed/kv_transfer/kv_connector/v1/base.py
+23
View File
@@ -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
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py
+3
View File
@@ -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,
vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py
+3
View File
@@ -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
vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
+13
View File
@@ -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
):
vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py
+3
View File
@@ -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
vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py
+75 -3
View File
@@ -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
vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
+11
View File
@@ -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
):
vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py
+4
View File
@@ -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(
vllm/v1/core/sched/scheduler.py
+11 -1
View File
@@ -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.
vllm/v1/kv_offload/base.py
+13
View File
@@ -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,
vllm/v1/kv_offload/cpu/manager.py
+7
View File
@@ -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,
vllm/v1/kv_offload/tiering/base.py
+12
View File
@@ -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,
vllm/v1/kv_offload/tiering/manager.py
+28
View File
@@ -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 = []
vllm/v1/kv_offload/tiering/p2p/manager.py
+4 -1
View File
@@ -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,
vllm/v1/metrics/cache_hit_source.py
+32
View File
@@ -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)}
vllm/v1/metrics/loggers.py
+31 -3
View File
@@ -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(
vllm/v1/metrics/stats.py
+33 -6
View File
@@ -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