[4/N] Expose HiSparse cache metrics via KV connector stats (#56061)
Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai> Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: OpenAI Codex <codex@openai.com>
Nicolò Lucchesi committed
Sep 12, 2026 at 19:47 UTC
29332cf93628a67c49252ef7c89a8aa3185def4b
9 files changed
+308
-8
csrc/libtorch_stable/hisparse_kernels.cu
+19
-1
@@ -203,6 +203,7 @@ __global__ __launch_bounds__(1024) void hisparse_resolve_residency_kernel(
203
int32_t* __restrict__ miss_mask, // [num_rows, top_k] or nullptr
204
int32_t* __restrict__ device_global_indices, // [max_rows, region_stride]
205
int16_t* __restrict__ lru_slots, // [max_rows, hot_size]
206
+ unsigned long long* __restrict__ stats, // [2] hits,misses or nullptr
207
const int32_t* __restrict__ request_state_indices, // [num_requests] or
208
// nullptr
209
const int32_t request_state_count, const int64_t host_rows,
@@ -552,6 +553,11 @@ __global__ __launch_bounds__(1024) void hisparse_resolve_residency_kernel(
553
if (swap_counts != nullptr && tid == 0) {
554
swap_counts[batch_row] = total_misses;
555
}
556
+ if (stats != nullptr && tid == 0) {
557
+ atomicAdd(&stats[0], static_cast<unsigned long long>(total_hits));
558
+ atomicAdd(&stats[1], static_cast<unsigned long long>(total_misses));
559
+ }
560
+
561
// Phase 4: write back the LRU order: stale evictables at the front,
562
// freshly loaded misses next, then hits at MRU.
563
const int total_evictable = hot_size - total_hits;
@@ -779,6 +785,7 @@ void hisparse_resolve_residency(
785
std::optional<torch::stable::Tensor> const& request_state_indices,
786
int64_t region_stride,
787
std::optional<torch::stable::Tensor> const& miss_mask,
788
+ std::optional<torch::stable::Tensor> const& stats,
789
std::optional<torch::stable::Tensor> const& attention_indices,
790
int64_t attention_block_stride,
791
std::optional<torch::stable::Tensor> const& request_ids,
@@ -997,6 +1004,17 @@ void hisparse_resolve_residency(
1004
miss_mask_ptr = mm.mutable_data_ptr<int32_t>();
1005
}
1006
1007
+ unsigned long long* stats_ptr = nullptr;
1008
+ if (stats.has_value()) {
1009
+ auto const& st = stats.value();
1010
+ STD_TORCH_CHECK(
1011
+ st.is_cuda() && st.is_contiguous() && st.dim() == 1 &&
1012
+ st.scalar_type() == torch::headeronly::ScalarType::UInt64 &&
1013
+ st.numel() == 2,
1014
+ "stats must be a contiguous two-element uint64 CUDA tensor");
1015
+ stats_ptr = static_cast<unsigned long long*>(st.mutable_data_ptr());
1016
+ }
1017
+
1018
int32_t* attention_indices_ptr = nullptr;
1019
if (attention_indices.has_value()) {
1020
auto const& indices = attention_indices.value();
@@ -1046,7 +1064,7 @@ void hisparse_resolve_residency(
1064
swap_counts_ptr, hot_indices.mutable_data_ptr<int32_t>(),
1065
attention_indices_ptr, miss_mask_ptr,
1066
device_global_indices.mutable_data_ptr<int32_t>(),
1049
- lru_slots.mutable_data_ptr<int16_t>(), request_state_ptr,
1067
+ lru_slots.mutable_data_ptr<int16_t>(), stats_ptr, request_state_ptr,
1068
request_state_count, host_rows, hot_block_table.stride(0), hot_block_size,
1069
top_k, hot_size, hash_size, region_stride, attention_block_stride,
1070
source_bt_stride, source_num_reqs, source_num_blocks,
csrc/libtorch_stable/ops.h
+1
@@ -673,6 +673,7 @@ void hisparse_resolve_residency(
673
std::optional<torch::stable::Tensor> const& request_state_indices,
674
int64_t region_stride,
675
std::optional<torch::stable::Tensor> const& miss_mask,
676
+ std::optional<torch::stable::Tensor> const& stats,
677
std::optional<torch::stable::Tensor> const& attention_indices,
678
int64_t attention_block_stride,
679
std::optional<torch::stable::Tensor> const& request_ids,
csrc/libtorch_stable/torch_bindings.cpp
+7
-6
@@ -999,16 +999,17 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C_cache_ops, ops) {
999
" Tensor? request_state_indices,"
1000
" int region_stride,"
1001
" Tensor(a!)? miss_mask=None,"
1002
- " Tensor(b!)? attention_indices=None,"
1002
+ " Tensor(b!)? stats=None,"
1003
+ " Tensor(c!)? attention_indices=None,"
1004
" int attention_block_stride=0,"
1005
" Tensor? request_ids=None,"
1006
" Tensor? source_block_table=None,"
1007
" int source_block_size=0,"
1007
- " Tensor(c!)? resolved_global_indices=None,"
1008
- " Tensor(d!)? valid_counts=None,"
1009
- " Tensor(e!)? swap_host_physical_rows=None,"
1010
- " Tensor(f!)? swap_device_physical_rows=None,"
1011
- " Tensor(g!)? swap_counts=None,"
1008
+ " Tensor(d!)? resolved_global_indices=None,"
1009
+ " Tensor(e!)? valid_counts=None,"
1010
+ " Tensor(f!)? swap_host_physical_rows=None,"
1011
+ " Tensor(g!)? swap_device_physical_rows=None,"
1012
+ " Tensor(h!)? swap_counts=None,"
1013
" Tensor? resident_block_table=None,"
1014
" int resident_block_size=0,"
1015
" int resident_null_block=0) -> ()");
tests/v1/kv_connector/unit/test_hisparse_stats.py
new
+62
@@ -0,0 +1,62 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.connector import (
5
+ HiSparseConnector,
6
+)
7
+from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.stats import (
8
+ HiSparseKVConnectorStats,
9
+)
10
+
11
+
12
+def test_is_empty_on_fresh_stats():
13
+ stats = HiSparseKVConnectorStats()
14
+ assert stats.is_empty()
15
+
16
+
17
+def test_record_snapshot_and_reduce():
18
+ stats = HiSparseKVConnectorStats()
19
+ stats.record_snapshot(hits=7, misses=3, host_to_device_bytes=48)
20
+ stats.record_snapshot(hits=5, misses=1, host_to_device_bytes=16)
21
+ assert not stats.is_empty()
22
+
23
+ reduced = stats.reduce()
24
+ assert reduced["HiSparse hot-buffer hits"] == 12
25
+ assert reduced["HiSparse hot-buffer misses"] == 4
26
+ assert reduced["HiSparse host-to-device bytes"] == 64
27
+
28
+
29
+def test_aggregate_extends_snapshot_deltas():
30
+ first = HiSparseKVConnectorStats()
31
+ first.record_snapshot(hits=7, misses=3, host_to_device_bytes=48)
32
+ second = HiSparseKVConnectorStats()
33
+ second.record_snapshot(hits=5, misses=1, host_to_device_bytes=16)
34
+
35
+ first.aggregate(second)
36
+
37
+ assert first.data == {
38
+ "cache_hits": [7, 5],
39
+ "cache_misses": [3, 1],
40
+ "host_to_device_bytes": [48, 16],
41
+ }
42
+
43
+
44
+def test_aggregate_skips_empty_stats():
45
+ stats = HiSparseKVConnectorStats()
46
+ stats.record_snapshot(hits=2, misses=1, host_to_device_bytes=16)
47
+
48
+ stats.aggregate(HiSparseKVConnectorStats())
49
+
50
+ assert stats.data["cache_hits"] == [2]
51
+
52
+
53
+def test_build_kv_connector_stats_round_trip():
54
+ stats = HiSparseKVConnectorStats()
55
+ stats.record_snapshot(hits=12, misses=4, host_to_device_bytes=64)
56
+ payload = stats.to_dict()
57
+
58
+ rebuilt = HiSparseConnector.build_kv_connector_stats(data=payload)
59
+
60
+ assert rebuilt is not None
61
+ assert isinstance(rebuilt, HiSparseKVConnectorStats)
62
+ assert rebuilt.reduce() == stats.reduce()
tests/v1/worker/test_utils.py
+38
-1
@@ -48,10 +48,47 @@ def _make_hisparse_worker() -> HiSparseConnectorWorker:
48
worker.dma_stream = None
49
worker.shared_host_region = None
50
worker._completed_host_copy_dst_ids = []
51
- worker._slot_mapping_staging = None
51
+ worker._metrics_calls = 0
52
+ worker._metrics_event = MagicMock()
53
+ worker._metrics_pending = False
54
+ worker.leader_runtimes = []
55
return worker
56
57
58
+def test_hisparse_worker_get_kv_connector_stats_reads_completed_snapshot(monkeypatch):
59
+ worker = _make_hisparse_worker()
60
+ worker._metrics_calls = hisparse_worker_module._METRICS_INTERVAL - 1
61
+ worker._metrics_pending = False
62
+ worker._metrics_event = MagicMock()
63
+ worker._metrics_event.query.return_value = True
64
+ compute_stream = MagicMock()
65
+ group = SimpleNamespace(
66
+ swap_stats=torch.tensor([12, 4], dtype=torch.uint64),
67
+ swap_stats_host=torch.empty(2, dtype=torch.uint64),
68
+ stats_row_bytes=16,
69
+ copy_stream=MagicMock(),
70
+ )
71
+ worker.leader_runtimes = [SimpleNamespace(index_group=group)]
72
+ monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False)
73
+ monkeypatch.setattr(
74
+ hisparse_worker_module, "current_stream", lambda: compute_stream
75
+ )
76
+
77
+ assert worker.get_kv_connector_stats() is None
78
+ compute_stream.wait_stream.assert_called_once_with(group.copy_stream)
79
+ group.copy_stream.wait_stream.assert_called_once_with(compute_stream)
80
+ worker._metrics_event.record.assert_called_once_with()
81
+ assert group.swap_stats.tolist() == [0, 0]
82
+
83
+ stats = worker.get_kv_connector_stats()
84
+ assert stats is not None
85
+ assert stats.data == {
86
+ "cache_hits": [12],
87
+ "cache_misses": [4],
88
+ "host_to_device_bytes": [64],
89
+ }
90
+
91
+
92
def test_hisparse_row_mirrors_follow_runner_request_order():
93
first = SparseKVRowMirror((1,), 10, 1)
94
second = SparseKVRowMirror((2,), 20, 1)
vllm/distributed/kv_transfer/kv_connector/v1/hisparse/connector.py
+37
@@ -17,9 +17,19 @@ from vllm.distributed.kv_transfer.kv_connector.v1.base import (
17
KVConnectorWorkerMetadata,
18
SupportsHMA,
19
)
20
+from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.stats import (
21
+ HiSparseKVConnectorStats,
22
+ HiSparsePromMetrics,
23
+)
24
from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.worker import (
25
HiSparseConnectorWorker,
26
)
27
+from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
28
+ KVConnectorPromMetrics,
29
+ KVConnectorStats,
30
+ PromMetric,
31
+ PromMetricT,
32
+)
33
from vllm.v1.attention.backend import AttentionMetadata
34
from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy
35
from vllm.v1.core.sched.output import SchedulerOutput
@@ -244,6 +254,33 @@ class HiSparseConnector(KVConnectorBase_V1, SupportsHMA):
254
assert self.connector_worker is not None
255
self.connector_worker.reset_hot_state()
256
257
+ def get_kv_connector_stats(self) -> KVConnectorStats | None:
258
+ if self.connector_worker is None:
259
+ return None
260
+ return self.connector_worker.get_kv_connector_stats()
261
+
262
+ @classmethod
263
+ def build_kv_connector_stats(
264
+ cls, data: dict[str, Any] | None = None
265
+ ) -> KVConnectorStats | None:
266
+ return (
267
+ HiSparseKVConnectorStats(data=data)
268
+ if data is not None
269
+ else HiSparseKVConnectorStats()
270
+ )
271
+
272
+ @classmethod
273
+ def build_prom_metrics(
274
+ cls,
275
+ vllm_config: VllmConfig,
276
+ metric_types: dict[type[PromMetric], type[PromMetricT]],
277
+ labelnames: list[str],
278
+ per_engine_labelvalues: dict[int, list[object]],
279
+ ) -> KVConnectorPromMetrics:
280
+ return HiSparsePromMetrics(
281
+ vllm_config, metric_types, labelnames, per_engine_labelvalues
282
+ )
283
+
284
def start_load_kv(self, forward_context: ForwardContext, **kwargs: Any) -> None:
285
assert self.connector_worker is not None
286
metadata = self._get_connector_metadata()
vllm/distributed/kv_transfer/kv_connector/v1/hisparse/stats.py
new
+97
@@ -0,0 +1,97 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Stats and Prometheus metrics for the HiSparse connector."""
4
+
5
+from dataclasses import dataclass
6
+from typing import Any
7
+
8
+from vllm.config import VllmConfig
9
+from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
10
+ KVConnectorPromMetrics,
11
+ KVConnectorStats,
12
+ PromMetric,
13
+ PromMetricT,
14
+)
15
+from vllm.v1.metrics.utils import create_metric_per_engine
16
+
17
+_HISPARSE_COUNTERS: tuple[tuple[str, str], ...] = (
18
+ ("cache_hits", "Number of HiSparse device hot-buffer hits."),
19
+ ("cache_misses", "Number of HiSparse device hot-buffer misses."),
20
+ (
21
+ "host_to_device_bytes",
22
+ "Bytes transferred from host KV storage to HiSparse hot buffers.",
23
+ ),
24
+)
25
+
26
+
27
+@dataclass
28
+class HiSparseKVConnectorStats(KVConnectorStats):
29
+ """Container for HiSparse hot-buffer residency metrics.
30
+
31
+ Each list entry is the delta recorded over one device counter snapshot
32
+ interval, so a list can hold multiple snapshots per logging interval.
33
+ """
34
+
35
+ def __post_init__(self):
36
+ if not self.data:
37
+ # Empty container init, no data is passed in.
38
+ self.reset()
39
+
40
+ def reset(self):
41
+ # Must be serializable
42
+ self.data: dict[str, list[int]] = {
43
+ "cache_hits": [],
44
+ "cache_misses": [],
45
+ "host_to_device_bytes": [],
46
+ }
47
+
48
+ def record_snapshot(self, hits: int, misses: int, host_to_device_bytes: int):
49
+ self.data["cache_hits"].append(hits)
50
+ self.data["cache_misses"].append(misses)
51
+ self.data["host_to_device_bytes"].append(host_to_device_bytes)
52
+
53
+ def aggregate(self, other: KVConnectorStats) -> KVConnectorStats:
54
+ if not other.is_empty():
55
+ for k, v in other.data.items():
56
+ accumulator = self.data[k]
57
+ assert isinstance(accumulator, list)
58
+ accumulator.extend(v)
59
+ return self
60
+
61
+ def reduce(self) -> dict[str, int | float]:
62
+ # Compute compact representative stats suitable for CLI logging.
63
+ return {
64
+ "HiSparse hot-buffer hits": sum(self.data["cache_hits"]),
65
+ "HiSparse hot-buffer misses": sum(self.data["cache_misses"]),
66
+ "HiSparse host-to-device bytes": sum(self.data["host_to_device_bytes"]),
67
+ }
68
+
69
+ def is_empty(self) -> bool:
70
+ return all(len(values) == 0 for values in self.data.values())
71
+
72
+
73
+class HiSparsePromMetrics(KVConnectorPromMetrics):
74
+ def __init__(
75
+ self,
76
+ vllm_config: VllmConfig,
77
+ metric_types: dict[type[PromMetric], type[PromMetricT]],
78
+ labelnames: list[str],
79
+ per_engine_labelvalues: dict[int, list[object]],
80
+ ):
81
+ super().__init__(vllm_config, metric_types, labelnames, per_engine_labelvalues)
82
+
83
+ self.hisparse_counters: dict[str, Any] = {}
84
+ for name, documentation in _HISPARSE_COUNTERS:
85
+ counter = self._counter_cls(
86
+ name=f"vllm:hisparse_{name}",
87
+ documentation=documentation,
88
+ labelnames=labelnames,
89
+ )
90
+ self.hisparse_counters[name] = create_metric_per_engine(
91
+ counter, self.per_engine_labelvalues
92
+ )
93
+
94
+ def observe(self, transfer_stats_data: dict[str, Any], engine_idx: int = 0):
95
+ for name, counter in self.hisparse_counters.items():
96
+ for value in transfer_stats_data.get(name, []):
97
+ counter[engine_idx].inc(value)
vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py
+39
@@ -15,6 +15,9 @@ import torch
15
16
from vllm import _custom_ops as ops
17
from vllm.config import VllmConfig
18
+from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.stats import (
19
+ HiSparseKVConnectorStats,
20
+)
21
from vllm.distributed.parallel_state import (
22
get_tensor_model_parallel_rank,
23
get_tp_group,
@@ -68,6 +71,9 @@ def _allocate_dma_descriptors(size: int) -> _DMADescriptors:
71
)
72
73
74
+_METRICS_INTERVAL = 2000
75
+
76
+
77
def _get_hisparse_cache(
78
forward_context: dict[str, Any], layer_name: str
79
) -> HiSparseCacheHandle:
@@ -311,6 +317,9 @@ class HiSparseConnectorWorker:
317
self._pending_transfer_events: deque[tuple[torch.Event, tuple[int, ...]]] = (
318
deque()
319
)
320
+ self._metrics_calls = 0
321
+ self._metrics_event = torch.Event()
322
+ self._metrics_pending = False
323
self._init_dma()
324
if self.is_host_writer:
325
for layer_index, handle in enumerate(cache_handles):
@@ -496,6 +505,36 @@ class HiSparseConnectorWorker:
505
for runtime in self.leader_runtimes:
506
runtime.reset_hot_state()
507
508
+ def get_kv_connector_stats(self) -> HiSparseKVConnectorStats | None:
509
+ stats = None
510
+ if self._metrics_pending and self._metrics_event.query():
511
+ stats = HiSparseKVConnectorStats()
512
+ for runtime in self.leader_runtimes:
513
+ group = runtime.index_group
514
+ hits, misses = group.swap_stats_host.tolist()
515
+ if hits or misses:
516
+ stats.record_snapshot(hits, misses, misses * group.stats_row_bytes)
517
+ self._metrics_pending = False
518
+ if stats.is_empty():
519
+ stats = None
520
+
521
+ self._metrics_calls += 1
522
+ if (
523
+ self._metrics_calls % _METRICS_INTERVAL == 0
524
+ and not self._metrics_pending
525
+ and not torch.cuda.is_current_stream_capturing()
526
+ ):
527
+ compute_stream = current_stream()
528
+ for runtime in self.leader_runtimes:
529
+ group = runtime.index_group
530
+ compute_stream.wait_stream(group.copy_stream)
531
+ group.swap_stats_host.copy_(group.swap_stats, non_blocking=True)
532
+ group.swap_stats.zero_()
533
+ group.copy_stream.wait_stream(compute_stream)
534
+ self._metrics_event.record()
535
+ self._metrics_pending = True
536
+ return stats
537
+
538
def _release_completed_dma_descriptors(self) -> None:
539
pending = self._pending_dma_descriptors
540
while pending and pending[0][0].query():
vllm/v1/hisparse/runtime.py
+8
@@ -464,6 +464,11 @@ class HiSparseIndexGroup:
464
)
465
self.logical_topk_ready = logical_topk_ready
466
self.followers: list[HiSparseRuntime] = []
467
+ self.swap_stats = torch.zeros(2, dtype=torch.uint64, device=device)
468
+ self.swap_stats_host = torch.empty(
469
+ 2, dtype=torch.uint64, device="cpu", pin_memory=True
470
+ )
471
+ self.stats_row_bytes = 0
472
473
474
class HiSparseRuntime:
@@ -526,6 +531,7 @@ class HiSparseRuntime:
531
else:
532
index_group.followers.append(self)
533
self.index_group = index_group
534
+ index_group.stats_row_bytes += row_bytes
535
536
self.eager_host_mirror = config.eager_host_mirror
537
self.resident_source_index = -1
@@ -639,6 +645,7 @@ class HiSparseRuntime:
645
compute_stream.wait_stream(group.copy_stream)
646
group.device_global_indices.fill_(-1)
647
group.lru_slots.copy_(group.lru_init.expand_as(group.lru_slots))
648
+ group.swap_stats.zero_()
649
group.copy_stream.wait_stream(compute_stream)
650
651
def invalidate_sorted_slots(
@@ -741,6 +748,7 @@ class HiSparseRuntime:
748
request_state_indices,
749
self.region_stride,
750
None,
751
+ group.swap_stats,
752
physical_topk_indices,
753
hot.attention_block_stride,
754
req_id_per_token[:num_tokens].contiguous(),