A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[P/D][Nixl] Introduce `KVTransferMetrics` and aggregation strategy (#22188)
Signed-off-by: NickLucche <nlucches@redhat.com>
Nicolò Lucchesi committed
Sep 19, 2025 at 13:09 UTC
a3d087adecadd4f6f83b72181ade40ec2de92aef
11 files changed
+525
-25
tests/v1/kv_connector/unit/test_nixl_connector.py
+210
-1
@@ -18,12 +18,18 @@ import torch
18
19
from vllm import LLM
20
from vllm.config import KVTransferConfig
21
+from vllm.distributed.kv_transfer.kv_connector.utils import KVOutputAggregator
22
+from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
23
+ KVConnectorStats)
24
+from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import (
25
+ MultiKVConnectorStats)
26
from vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector import (
27
KVConnectorRole, NixlAgentMetadata, NixlConnector, NixlConnectorMetadata,
23
- NixlConnectorWorker)
28
+ NixlConnectorWorker, NixlKVConnectorStats)
29
from vllm.forward_context import ForwardContext
30
from vllm.sampling_params import SamplingParams
31
from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
32
+from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput
33
34
from .utils import create_request, create_scheduler, create_vllm_config
35
@@ -475,6 +481,209 @@ class TestNixlHandshake:
481
# NOTE: resource cleanup in mp backend is a bit finicky, so the order in which
482
# we put here is important. First run ray, it will clean up the resources, then
483
# the rest of the tests.
484
+@patch(
485
+ "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.NixlWrapper",
486
+ FakeNixlWrapper)
487
+def test_kv_connector_stats(dist_init):
488
+ """Test that KV transfer stats are properly recorded and retrieved."""
489
+ vllm_config = create_vllm_config()
490
+
491
+ # Test worker role in decode server.
492
+ connector = NixlConnector(vllm_config, KVConnectorRole.WORKER)
493
+ connector.connector_worker = FakeNixlConnectorWorker(vllm_config,
494
+ connector.engine_id,
495
+ hand_shake_latency=0)
496
+
497
+ # Verify that xfer_stats starts empty
498
+ initial_stats = connector.get_kv_connector_stats()
499
+ assert initial_stats is None
500
+
501
+ # Create transfer metadata
502
+ request_id = "test_req_for_stats"
503
+ metadata = NixlConnectorMetadata()
504
+ metadata.add_new_req(request_id=request_id,
505
+ local_block_ids=[1, 2, 3],
506
+ kv_transfer_params={
507
+ "remote_block_ids": [4, 5, 6],
508
+ "remote_engine_id":
509
+ FakeNixlConnectorWorker.REMOTE_ENGINE_ID,
510
+ "remote_host": "localhost",
511
+ "remote_port": 1234,
512
+ "remote_tp_size": 1,
513
+ })
514
+ connector.bind_connector_metadata(metadata)
515
+
516
+ # Start the transfer
517
+ dummy_ctx = ForwardContext(
518
+ no_compile_layers={},
519
+ attn_metadata={},
520
+ virtual_engine=0,
521
+ )
522
+ connector.start_load_kv(dummy_ctx)
523
+
524
+ # Verify stats are recorded after transfer is complete
525
+ max_iterations = 2
526
+ # Clear metadata before start_load_kv to prevent reprocessing same request
527
+ connector.bind_connector_metadata(NixlConnectorMetadata())
528
+ for _ in range(max_iterations):
529
+ # Need to call start_load_kv to process completed handshakes
530
+ connector.start_load_kv(dummy_ctx)
531
+ _, done_recving = connector.get_finished(finished_req_ids=set())
532
+ if len(done_recving) > 0 and request_id in done_recving:
533
+ break
534
+ time.sleep(
535
+ 0.1) # Small delay to allow background handshake to complete
536
+ else:
537
+ assert "Transfer did not complete within expected iterations"
538
+
539
+ # Now check that stats were recorded
540
+ stats_after_transfer = connector.get_kv_connector_stats()
541
+ assert isinstance(stats_after_transfer, NixlKVConnectorStats)
542
+
543
+ # Verify stats values are recorded
544
+ assert not stats_after_transfer.is_empty()
545
+ assert stats_after_transfer.data["num_successful_transfers"] == 1
546
+
547
+ # Verify stats are reset after retrieval
548
+ stats_after_reset = connector.get_kv_connector_stats()
549
+ assert stats_after_reset is None
550
+
551
+
552
+def test_kv_connector_stats_aggregation():
553
+ """
554
+ Test KV transfer stats aggregation across TP ranks using
555
+ KVOutputAggregator (used by MultiprocExecutor).
556
+ """
557
+
558
+ # Create KVOutputAggregator for 3 workers (simulating TP=3), same thing
559
+ # done in MultiprocExecutor.execute_model
560
+ aggregator = KVOutputAggregator(world_size=3)
561
+
562
+ # Create stats for multiple workers with different transfer patterns
563
+ worker1_stats = NixlKVConnectorStats()
564
+ worker2_stats = NixlKVConnectorStats()
565
+ worker3_stats = NixlKVConnectorStats()
566
+
567
+ # Record different transfers on each worker
568
+ # Worker 1: 2 transfers
569
+ worker1_stats.record_transfer()
570
+ worker1_stats.record_transfer()
571
+
572
+ # Worker 2: 1 transfer
573
+ worker2_stats.record_transfer()
574
+
575
+ # Worker 3: 3 transfers
576
+ worker3_stats.record_transfer()
577
+ worker3_stats.record_transfer()
578
+ worker3_stats.record_transfer()
579
+
580
+ # Create ModelRunnerOutput instances for each worker
581
+ worker_outputs = []
582
+ for i, worker_stats in enumerate(
583
+ [worker1_stats, worker2_stats, worker3_stats]):
584
+ output = ModelRunnerOutput(
585
+ req_ids=[f"req_{i}"],
586
+ req_id_to_index={f"req_{i}": 0},
587
+ sampled_token_ids=[[123]], # dummy token
588
+ logprobs=None,
589
+ prompt_logprobs_dict={},
590
+ pooler_output=[None],
591
+ kv_connector_output=KVConnectorOutput(
592
+ finished_sending=set([f"req_{i}_send"])
593
+ if i < 2 else None, # Workers 0,1 finished sending
594
+ finished_recving=set([f"req_{i}_recv"])
595
+ if i > 0 else None, # Workers 1,2 finished receiving
596
+ kv_connector_stats=worker_stats,
597
+ ))
598
+ worker_outputs.append(output)
599
+
600
+ # Use the real aggregation mechanism (like MultiprocExecutor.execute_model)
601
+ aggregated_output = aggregator.aggregate(worker_outputs, output_rank=0)
602
+ kv_connector_stats = \
603
+ aggregated_output.kv_connector_output.kv_connector_stats
604
+ assert isinstance(kv_connector_stats, NixlKVConnectorStats)
605
+ # Number of total transfers across all workers.
606
+ assert kv_connector_stats.data["num_successful_transfers"] == 6
607
+
608
+
609
+def test_multi_kv_connector_stats_aggregation():
610
+ """
611
+ Test MultiKVConnectorStats aggregation across TP ranks using
612
+ KVOutputAggregator (used by MultiprocExecutor).
613
+ """
614
+
615
+ aggregator = KVOutputAggregator(world_size=3)
616
+
617
+ from dataclasses import dataclass
618
+
619
+ @dataclass
620
+ class FooKVConnectorStats(KVConnectorStats):
621
+
622
+ def reset(self):
623
+ self.data = {"num_foo_transfers": 0}
624
+
625
+ def record_transfer(self):
626
+ if "num_foo_transfers" not in self.data:
627
+ self.data["num_foo_transfers"] = 0
628
+ self.data["num_foo_transfers"] += 1
629
+
630
+ def is_empty(self) -> bool:
631
+ return self.data["num_foo_transfers"] == 0
632
+
633
+ def aggregate(self,
634
+ other: "FooKVConnectorStats") -> "FooKVConnectorStats":
635
+ if not other.is_empty():
636
+ self.data["num_foo_transfers"] += other.data[
637
+ "num_foo_transfers"]
638
+ return self
639
+
640
+ def make_multi_stats(nixl_count: int,
641
+ foo_count: int) -> MultiKVConnectorStats:
642
+ data: dict[str, KVConnectorStats] = {}
643
+ if nixl_count > 0:
644
+ nixl_stats = NixlKVConnectorStats()
645
+ for _ in range(nixl_count):
646
+ nixl_stats.record_transfer()
647
+ data["NixlConnector"] = nixl_stats
648
+ if foo_count > 0:
649
+ foo_stats = FooKVConnectorStats()
650
+ for _ in range(foo_count):
651
+ foo_stats.record_transfer()
652
+ data["FooConnector"] = foo_stats
653
+ return MultiKVConnectorStats(data=data)
654
+
655
+ # Create heterogeneous stats across 3 workers
656
+ worker_patterns = [(2, 1), (3, 0), (0, 5)] # (Nixl, Foo)
657
+
658
+ worker_outputs: list[ModelRunnerOutput] = []
659
+ for i, (nixl, foo) in enumerate(worker_patterns):
660
+ stats = make_multi_stats(nixl, foo)
661
+ output = ModelRunnerOutput(
662
+ req_ids=[f"req_{i}"],
663
+ req_id_to_index={f"req_{i}": 0},
664
+ sampled_token_ids=[[123]],
665
+ logprobs=None,
666
+ prompt_logprobs_dict={},
667
+ pooler_output=[None],
668
+ kv_connector_output=KVConnectorOutput(
669
+ finished_sending=set([f"req_{i}_send"]) if i < 2 else None,
670
+ finished_recving=set([f"req_{i}_recv"]) if i > 0 else None,
671
+ kv_connector_stats=stats,
672
+ ),
673
+ )
674
+ worker_outputs.append(output)
675
+
676
+ aggregated_output = aggregator.aggregate(worker_outputs, output_rank=0)
677
+ kv_connector_stats = \
678
+ aggregated_output.kv_connector_output.kv_connector_stats
679
+ assert isinstance(kv_connector_stats, MultiKVConnectorStats)
680
+
681
+ # Validate per-connector totals across workers
682
+ assert kv_connector_stats["NixlConnector"].data[
683
+ "num_successful_transfers"] == 5
684
+ assert kv_connector_stats["FooConnector"].data["num_foo_transfers"] == 6
685
+
686
+
687
@pytest.mark.parametrize("distributed_executor_backend", ["ray", None])
688
@patch(
689
"vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.NixlWrapper",
vllm/distributed/kv_transfer/kv_connector/utils.py
+18
-3
@@ -129,7 +129,7 @@ class KVOutputAggregator:
129
def aggregate(self,
130
outputs: list[ModelRunnerOutput],
131
output_rank: int = 0) -> ModelRunnerOutput:
132
- # aggregate kv_connector_output from all workers
132
+ # Aggregate kv_connector_output from all workers
133
134
def update_finished_set(req_ids: Optional[set[str]],
135
remaining_count_dict: dict[str, int],
@@ -142,8 +142,9 @@ class KVOutputAggregator:
142
143
finished_sending = set[str]()
144
finished_recving = set[str]()
145
- for output in outputs:
146
- output = output.kv_connector_output
145
+ aggregated_kv_connector_stats = None
146
+ for model_runner_output in outputs:
147
+ output = model_runner_output.kv_connector_output
148
if not output:
149
continue
150
update_finished_set(output.finished_sending,
@@ -151,12 +152,26 @@ class KVOutputAggregator:
152
update_finished_set(output.finished_recving,
153
self._recv_remaining_count, finished_recving)
154
155
+ # Aggregate kv_connector_stats from all workers.
156
+ if aggregated_kv_connector_stats is None:
157
+ # Use the first worker's kv_connector_stats as accumulator.
158
+ aggregated_kv_connector_stats = output.kv_connector_stats
159
+ elif kv_connector_stats := output.kv_connector_stats:
160
+ if aggregated_kv_connector_stats is None:
161
+ aggregated_kv_connector_stats = kv_connector_stats
162
+ else:
163
+ assert isinstance(aggregated_kv_connector_stats,
164
+ type(kv_connector_stats))
165
+ aggregated_kv_connector_stats = \
166
+ aggregated_kv_connector_stats.aggregate(kv_connector_stats)
167
+
168
# select output of the worker specified by output_rank
169
output = outputs[output_rank]
170
171
output.kv_connector_output = KVConnectorOutput(
172
finished_sending=finished_sending or None,
173
finished_recving=finished_recving or None,
174
+ kv_connector_stats=aggregated_kv_connector_stats or None,
175
)
176
177
return output
vllm/distributed/kv_transfer/kv_connector/v1/base.py
+21
-1
@@ -49,6 +49,8 @@ if TYPE_CHECKING:
49
from vllm.attention.backends.abstract import AttentionMetadata
50
from vllm.config import VllmConfig
51
from vllm.distributed.kv_events import KVCacheEvent
52
+ from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
53
+ KVConnectorStats)
54
from vllm.forward_context import ForwardContext
55
from vllm.v1.core.kv_cache_manager import KVCacheBlocks
56
from vllm.v1.request import Request
@@ -235,6 +237,12 @@ class KVConnectorBase_V1(ABC):
237
"""
238
return None
239
240
+ def get_kv_connector_stats(self) -> Optional["KVConnectorStats"]:
241
+ """
242
+ Get the KV connector stats collected during the last interval.
243
+ """
244
+ return None
245
+
246
# ==============================
247
# Scheduler-side methods
248
# ==============================
@@ -365,4 +373,16 @@ class KVConnectorBase_V1(ABC):
373
int: expected sending or receiving completion count.
374
"""
375
368
- return None
\ No newline at end of file
376
+ return None
377
+
378
+ @classmethod
379
+ def build_kv_connector_stats(
380
+ cls,
381
+ data: Optional[dict[str,
382
+ Any]] = None) -> Optional["KVConnectorStats"]:
383
+ """
384
+ KVConnectorStats resolution method. This method allows dynamically
385
+ registered connectors to return their own KVConnectorStats object,
386
+ which can implement custom aggregation logic on the data dict.
387
+ """
388
+ return None
vllm/distributed/kv_transfer/kv_connector/v1/metrics.py
new
+100
@@ -0,0 +1,100 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from dataclasses import dataclass, field
4
+from typing import Any, Optional, Union
5
+
6
+from vllm.config.kv_transfer import KVTransferConfig
7
+from vllm.distributed.kv_transfer.kv_connector.factory import (
8
+ KVConnectorFactory)
9
+from vllm.distributed.kv_transfer.kv_transfer_state import (
10
+ has_kv_transfer_group)
11
+from vllm.logger import init_logger
12
+
13
+logger = init_logger(__name__)
14
+
15
+
16
+@dataclass
17
+class KVConnectorStats:
18
+ """
19
+ Base class for KV Connector Stats, a container for transfer performance
20
+ metrics or otherwise important telemetry from the connector.
21
+ All sub-classes need to be serializable as stats are sent from worker to
22
+ logger process.
23
+ """
24
+ data: dict[str, Any] = field(default_factory=dict)
25
+
26
+ def reset(self):
27
+ """Reset the stats, clear the state."""
28
+ raise NotImplementedError
29
+
30
+ def aggregate(self, other: "KVConnectorStats") -> "KVConnectorStats":
31
+ """
32
+ Aggregate stats with another `KVConnectorStats` object.
33
+ """
34
+ raise NotImplementedError
35
+
36
+ def reduce(self) -> dict[str, Union[int, float]]:
37
+ """
38
+ Reduce the observations collected during a time interval to one or
39
+ more representative values (eg avg/median/sum of the series).
40
+ This is meant to be called by the logger to produce a summary of the
41
+ stats for the last time interval.
42
+ """
43
+ raise NotImplementedError
44
+
45
+ def is_empty(self) -> bool:
46
+ """Return True if the stats are empty."""
47
+ raise NotImplementedError
48
+
49
+
50
+class KVConnectorLogging:
51
+
52
+ def __init__(self, kv_tranfer_config: KVTransferConfig):
53
+ # This should be called on frontend process.
54
+ assert not has_kv_transfer_group()
55
+ # Instantiate the connector's stats class.
56
+ if kv_tranfer_config and kv_tranfer_config.kv_connector:
57
+ self.connector_cls = KVConnectorFactory.get_connector_class(
58
+ kv_tranfer_config)
59
+ self.reset()
60
+
61
+ def reset(self):
62
+ self.transfer_stats_accumulator: Optional[KVConnectorStats] = None
63
+
64
+ def observe(self, transfer_stats_data: dict[str, Any]):
65
+ # Should not be called when a KVConnector is not configured.
66
+ assert self.connector_cls is not None
67
+ # Called periodically when connector syncs with the scheduler.
68
+ # Note that this is not the same as the logging interval.
69
+ # We expect transfer_stats_data to be aggregated across all workers and
70
+ # consist of observations from a single connector or a MultiConnector.
71
+ transfer_stats = self.connector_cls.build_kv_connector_stats(
72
+ transfer_stats_data)
73
+ if transfer_stats is None:
74
+ logger.warning_once(
75
+ "The connector %s is collecting stats but "
76
+ "does not implement the "
77
+ "`build_kv_connector_stats` method. "
78
+ "Stats will not be logged.", self.connector_cls)
79
+ return
80
+
81
+ if self.transfer_stats_accumulator is None:
82
+ self.transfer_stats_accumulator = transfer_stats
83
+ else:
84
+ # Accumulate last interval stats.
85
+ self.transfer_stats_accumulator = \
86
+ self.transfer_stats_accumulator.aggregate(transfer_stats)
87
+
88
+ def log(self, log_fn=logger.info):
89
+ """Log transfer metrics periodically, similar to throughput logging"""
90
+ if (self.transfer_stats_accumulator
91
+ and not self.transfer_stats_accumulator.is_empty()):
92
+ # Produce a single cumulative stats object for the last time
93
+ # interval from the recorded observations.
94
+ xfer_metrics = self.transfer_stats_accumulator.reduce()
95
+ xfer_metrics_str = ", ".join(f"{k}={v}"
96
+ for k, v in xfer_metrics.items())
97
+ log_fn("KV Transfer metrics: %s", xfer_metrics_str)
98
+
99
+ # Reset metrics for next interval
100
+ self.reset()
\ No newline at end of file
vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
+65
-3
@@ -9,19 +9,21 @@ import torch
9
10
from vllm.config import VllmConfig
11
from vllm.config.kv_transfer import KVTransferConfig
12
-from vllm.distributed.kv_events import KVCacheEvent
12
from vllm.distributed.kv_transfer.kv_connector.factory import (
13
KVConnectorFactory)
14
from vllm.distributed.kv_transfer.kv_connector.v1.base import (
15
KVConnectorBase_V1, KVConnectorMetadata, KVConnectorRole)
16
+from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
17
+ KVConnectorStats)
18
from vllm.logger import init_logger
18
-from vllm.v1.core.kv_cache_manager import KVCacheBlocks
19
from vllm.v1.core.sched.output import SchedulerOutput
20
from vllm.v1.outputs import KVConnectorOutput
21
22
if TYPE_CHECKING:
23
from vllm.attention.backends.abstract import AttentionMetadata
24
+ from vllm.distributed.kv_events import KVCacheEvent
25
from vllm.forward_context import ForwardContext
26
+ from vllm.v1.core.kv_cache_manager import KVCacheBlocks
27
from vllm.v1.request import Request
28
29
logger = init_logger(__name__)
@@ -33,6 +35,43 @@ class MultiKVConnectorMetadata(KVConnectorMetadata):
35
extra_async_saves: Optional[dict[str, int]] = None
36
37
38
+@dataclass
39
+class MultiKVConnectorStats(KVConnectorStats):
40
+ """
41
+ Maintain a dict of KVConnectorStats objects, one for each connector.
42
+ This is used to aggregate the stats from all connectors separately.
43
+ """
44
+
45
+ def aggregate(self, other: KVConnectorStats) -> KVConnectorStats:
46
+ for connector_id, stats in other.data.items():
47
+ if connector_id not in self.data:
48
+ self[connector_id] = stats
49
+ else:
50
+ assert isinstance(stats, type(self.data[connector_id]))
51
+ self[connector_id] = self[connector_id].aggregate(stats)
52
+ return self
53
+
54
+ def reset(self):
55
+ for stats in self.data.values():
56
+ stats.reset()
57
+
58
+ def reduce(self) -> dict[str, Any]:
59
+ # TODO (NickLucche) Adjust for logging on separate lines
60
+ return {
61
+ connector_id: stats.reduce()
62
+ for connector_id, stats in self.data.items()
63
+ }
64
+
65
+ def is_empty(self) -> bool:
66
+ return all(stats.is_empty() for stats in self.data.values())
67
+
68
+ def __getitem__(self, connector_id: str) -> KVConnectorStats:
69
+ return self.data[connector_id]
70
+
71
+ def __setitem__(self, connector_id: str, stats: KVConnectorStats):
72
+ self.data[connector_id] = stats
73
+
74
+
75
class MultiConnector(KVConnectorBase_V1):
76
"""
77
A wrapper for using multiple KVConnectors at the same time.
@@ -46,6 +85,7 @@ class MultiConnector(KVConnectorBase_V1):
85
def __init__(self, vllm_config: "VllmConfig", role: KVConnectorRole):
86
super().__init__(vllm_config=vllm_config, role=role)
87
self._connectors: list[KVConnectorBase_V1] = []
88
+ self._ktc_kv_transfer_config = []
89
ktcs = vllm_config.kv_transfer_config.kv_connector_extra_config.get(
90
"connectors")
91
assert ktcs is not None
@@ -57,6 +97,7 @@ class MultiConnector(KVConnectorBase_V1):
97
**ktc, engine_id=engine_id)
98
self._connectors.append(
99
KVConnectorFactory.create_connector(temp_config, role))
100
+ self._ktc_kv_transfer_config.append(temp_config.kv_transfer_config)
101
102
# A mapping from request id to the index of the connector chosen to
103
# load the request from (if any).
@@ -227,7 +268,7 @@ class MultiConnector(KVConnectorBase_V1):
268
269
return async_saves > 0, kv_txfer_params
270
230
- def take_events(self) -> Iterable[KVCacheEvent]:
271
+ def take_events(self) -> Iterable["KVCacheEvent"]:
272
for c in self._connectors:
273
yield from c.take_events()
274
@@ -264,3 +305,24 @@ class MultiConnector(KVConnectorBase_V1):
305
f"({', '.join(layouts) })."
306
f"All connectors must use the same layout.")
307
return next(iter(layouts), None)
308
+
309
+ @classmethod
310
+ def build_kv_connector_stats(
311
+ cls,
312
+ data: Optional[dict[str,
313
+ Any]] = None) -> Optional[KVConnectorStats]:
314
+ return MultiKVConnectorStats(data=data) if data is not None \
315
+ else MultiKVConnectorStats()
316
+
317
+ def get_kv_connector_stats(self) -> Optional[MultiKVConnectorStats]:
318
+ # Group connector stats by connector type.
319
+ stats_by_connector: Optional[MultiKVConnectorStats] = None
320
+ for c in self._connectors:
321
+ stats = c.get_kv_connector_stats()
322
+ if stats is None:
323
+ continue
324
+ if stats_by_connector is None:
325
+ # Lazy init to allow optional return value.
326
+ stats_by_connector = MultiKVConnectorStats()
327
+ stats_by_connector[c.__class__.__name__] = stats
328
+ return stats_by_connector
vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+65
-3
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import contextlib
4
+import copy
5
import logging
6
import math
7
import queue
@@ -11,7 +12,7 @@ from collections import defaultdict
12
from collections.abc import Iterator
13
from concurrent.futures import Future, ThreadPoolExecutor
14
from dataclasses import dataclass
14
-from typing import TYPE_CHECKING, Any, Optional
15
+from typing import TYPE_CHECKING, Any, Optional, Union
16
17
import msgspec
18
import numpy as np
@@ -23,6 +24,8 @@ from vllm.attention.selector import backend_name_to_enum, get_attn_backend
24
from vllm.config import VllmConfig
25
from vllm.distributed.kv_transfer.kv_connector.v1.base import (
26
CopyBlocksOp, KVConnectorBase_V1, KVConnectorMetadata, KVConnectorRole)
27
+from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
28
+ KVConnectorStats)
29
from vllm.distributed.parallel_state import (
30
get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size,
31
get_tp_group)
@@ -33,7 +36,6 @@ from vllm.platforms import _Backend, current_platform
36
from vllm.utils import make_zmq_path, make_zmq_socket
37
from vllm.v1.attention.backends.utils import get_kv_cache_layout
38
from vllm.v1.core.sched.output import SchedulerOutput
36
-from vllm.v1.request import RequestStatus
39
40
if TYPE_CHECKING:
41
from vllm.attention.backends.abstract import AttentionMetadata
@@ -206,6 +208,18 @@ class NixlConnector(KVConnectorBase_V1):
208
assert self.connector_worker is not None
209
return self.connector_worker.get_finished()
210
211
+ def get_kv_connector_stats(self) -> Optional[KVConnectorStats]:
212
+ assert self.connector_worker is not None
213
+ return self.connector_worker.get_kv_connector_stats()
214
+
215
+ @classmethod
216
+ def build_kv_connector_stats(
217
+ cls,
218
+ data: Optional[dict[str,
219
+ Any]] = None) -> Optional[KVConnectorStats]:
220
+ return NixlKVConnectorStats(data=data) if data is not None \
221
+ else NixlKVConnectorStats()
222
+
223
def start_load_kv(self, forward_context: "ForwardContext",
224
**kwargs) -> None:
225
assert self.connector_worker is not None
@@ -377,6 +391,7 @@ class NixlConnectorScheduler:
391
Once a request is finished, determine whether request blocks
392
should be freed now or will be sent asynchronously and freed later.
393
"""
394
+ from vllm.v1.request import RequestStatus
395
396
params = request.kv_transfer_params
397
logger.debug(
@@ -550,6 +565,7 @@ class NixlConnectorWorker:
565
# With heterogeneous TP, P must wait for all assigned D TP workers to
566
# finish reading before safely freeing the blocks.
567
self.consumer_notification_counts_by_req = defaultdict[ReqId, int](int)
568
+ self.xfer_stats = NixlKVConnectorStats()
569
570
def __del__(self):
571
"""Cleanup background threads on destruction."""
@@ -1097,6 +1113,8 @@ class NixlConnectorWorker:
1113
xfer_state = self.nixl_wrapper.check_xfer_state(handle)
1114
if xfer_state == "DONE":
1115
self.nixl_wrapper.release_xfer_handle(handle)
1116
+ # TODO (NickLucche) Get from NIXL telemetry once integrated
1117
+ self.xfer_stats.record_transfer()
1118
elif xfer_state == "PROC":
1119
in_progress = True
1120
continue
@@ -1248,7 +1266,6 @@ class NixlConnectorWorker:
1266
self.nixl_wrapper.transfer(handle)
1267
1268
# Use handle to check completion in future step().
1251
- # TODO (NickLucche) surface xfer elapsed time
1269
self._recving_transfers[request_id].append(
1270
(handle, time.perf_counter()))
1271
@@ -1300,6 +1317,15 @@ class NixlConnectorWorker:
1317
block_len = self.block_len
1318
return block_len
1319
1320
+ def get_kv_connector_stats(self) -> Optional[KVConnectorStats]:
1321
+ """
1322
+ Get the KV transfer stats for the connector.
1323
+ """
1324
+ # Clear stats for next iteration
1325
+ if not self.xfer_stats.is_empty():
1326
+ return self.xfer_stats.clone_and_reset()
1327
+ return None
1328
+
1329
1330
@contextlib.contextmanager
1331
def zmq_ctx(socket_type: Any, addr: str) -> Iterator[zmq.Socket]:
@@ -1318,3 +1344,39 @@ def zmq_ctx(socket_type: Any, addr: str) -> Iterator[zmq.Socket]:
1344
finally:
1345
if ctx is not None:
1346
ctx.destroy(linger=0)
1347
+
1348
+
1349
+@dataclass
1350
+class NixlKVConnectorStats(KVConnectorStats):
1351
+ """Container for transfer performance metrics"""
1352
+
1353
+ def __post_init__(self):
1354
+ if "num_successful_transfers" not in self.data:
1355
+ self.data["num_successful_transfers"] = 0
1356
+
1357
+ def reset(self):
1358
+ self.data = {"num_successful_transfers": 0}
1359
+
1360
+ def record_transfer(self):
1361
+ # TODO: record actual transfer stats when available
1362
+ self.data["num_successful_transfers"] += 1
1363
+
1364
+ def clone_and_reset(self) -> "NixlKVConnectorStats":
1365
+ old = copy.copy(self)
1366
+ self.reset()
1367
+ return old
1368
+
1369
+ def is_empty(self) -> bool:
1370
+ return self.data["num_successful_transfers"] == 0
1371
+
1372
+ def aggregate(self, other: KVConnectorStats) -> KVConnectorStats:
1373
+ if not other.is_empty():
1374
+ self.data["num_successful_transfers"] += other.data[
1375
+ "num_successful_transfers"]
1376
+ return self
1377
+
1378
+ def reduce(self) -> dict[str, Union[int, float]]:
1379
+ # TODO: reduce stats to a single value, calculate latency/throughput
1380
+ return {
1381
+ "num_successful_transfers": self.data["num_successful_transfers"]
1382
+ }
\ No newline at end of file
vllm/v1/core/sched/scheduler.py
+17
-10
@@ -15,6 +15,8 @@ from vllm.distributed.kv_transfer.kv_connector.factory import (
15
KVConnectorFactory)
16
from vllm.distributed.kv_transfer.kv_connector.v1 import (KVConnectorBase_V1,
17
KVConnectorRole)
18
+from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
19
+ KVConnectorStats)
20
from vllm.logger import init_logger
21
from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry
22
from vllm.v1.core.encoder_cache_manager import (EncoderCacheManager,
@@ -869,9 +871,12 @@ class Scheduler(SchedulerInterface):
871
num_scheduled_tokens = scheduler_output.num_scheduled_tokens
872
pooler_outputs = model_runner_output.pooler_output
873
num_nans_in_logits = model_runner_output.num_nans_in_logits
874
+ kv_connector_output = model_runner_output.kv_connector_output
875
876
outputs: dict[int, list[EngineCoreOutput]] = defaultdict(list)
877
spec_decoding_stats: Optional[SpecDecodingStats] = None
878
+ kv_connector_stats = (kv_connector_output.kv_connector_stats
879
+ if kv_connector_output else None)
880
881
# NOTE(woosuk): As len(num_scheduled_tokens) can be up to 1K or more,
882
# the below loop can be a performance bottleneck. We should do our best
@@ -1007,7 +1012,8 @@ class Scheduler(SchedulerInterface):
1012
finished_requests=finished_set)
1013
finished_req_ids.clear()
1014
1010
- if (stats := self.make_stats(spec_decoding_stats)) is not None:
1015
+ if (stats := self.make_stats(spec_decoding_stats,
1016
+ kv_connector_stats)) is not None:
1017
# Return stats to only one of the front-ends.
1018
if (eco := next(iter(engine_core_outputs.values()), None)) is None:
1019
# We must return the stats even if there are no request
@@ -1172,20 +1178,21 @@ class Scheduler(SchedulerInterface):
1178
def make_stats(
1179
self,
1180
spec_decoding_stats: Optional[SpecDecodingStats] = None,
1181
+ kv_connector_stats: Optional[KVConnectorStats] = None,
1182
) -> Optional[SchedulerStats]:
1183
if not self.log_stats:
1184
return None
1185
prefix_cache_stats = self.kv_cache_manager.make_prefix_cache_stats()
1186
assert prefix_cache_stats is not None
1180
- return SchedulerStats(
1181
- num_running_reqs=len(self.running),
1182
- num_waiting_reqs=len(self.waiting),
1183
- kv_cache_usage=self.kv_cache_manager.usage,
1184
- prefix_cache_stats=prefix_cache_stats,
1185
- spec_decoding_stats=spec_decoding_stats,
1186
- num_corrupted_reqs=sum(req.is_output_corrupted
1187
- for req in self.running),
1188
- )
1187
+ return SchedulerStats(num_running_reqs=len(self.running),
1188
+ num_waiting_reqs=len(self.waiting),
1189
+ kv_cache_usage=self.kv_cache_manager.usage,
1190
+ prefix_cache_stats=prefix_cache_stats,
1191
+ spec_decoding_stats=spec_decoding_stats,
1192
+ num_corrupted_reqs=sum(req.is_output_corrupted
1193
+ for req in self.running),
1194
+ kv_connector_stats=kv_connector_stats.data
1195
+ if kv_connector_stats else None)
1196
1197
def make_spec_decoding_stats(
1198
self,
vllm/v1/metrics/loggers.py
+7
-1
@@ -9,6 +9,8 @@ from typing import Callable, Optional, Union
9
import prometheus_client
10
11
from vllm.config import SupportsMetricsInfo, VllmConfig
12
+from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
13
+ KVConnectorLogging)
14
from vllm.logger import init_logger
15
from vllm.v1.core.kv_cache_utils import PrefixCachingMetrics
16
from vllm.v1.engine import FinishReason
@@ -59,6 +61,8 @@ class LoggingStatLogger(StatLoggerBase):
61
# TODO: Make the interval configurable.
62
self.prefix_caching_metrics = PrefixCachingMetrics()
63
self.spec_decoding_logging = SpecDecodingLogging()
64
+ kv_tranfer_config = self.vllm_config.kv_transfer_config
65
+ self.kv_transfer_logging = KVConnectorLogging(kv_tranfer_config)
66
self.last_prompt_throughput: float = 0.0
67
self.last_generation_throughput: float = 0.0
68
@@ -97,7 +101,8 @@ class LoggingStatLogger(StatLoggerBase):
101
if scheduler_stats.spec_decoding_stats is not None:
102
self.spec_decoding_logging.observe(
103
scheduler_stats.spec_decoding_stats)
100
-
104
+ if kv_connector_stats := scheduler_stats.kv_connector_stats:
105
+ self.kv_transfer_logging.observe(kv_connector_stats)
106
self.last_scheduler_stats = scheduler_stats
107
108
def log(self):
@@ -136,6 +141,7 @@ class LoggingStatLogger(StatLoggerBase):
141
self.prefix_caching_metrics.hit_rate * 100,
142
)
143
self.spec_decoding_logging.log(log_fn=log_fn)
144
+ self.kv_transfer_logging.log(log_fn=log_fn)
145
146
def log_engine_initialized(self):
147
if self.vllm_config.cache_config.num_gpu_blocks:
vllm/v1/metrics/stats.py
+2
-1
@@ -3,7 +3,7 @@
3
4
import time
5
from dataclasses import dataclass, field
6
-from typing import TYPE_CHECKING, Optional
6
+from typing import TYPE_CHECKING, Any, Optional
7
8
from vllm.v1.spec_decode.metrics import SpecDecodingStats
9
@@ -43,6 +43,7 @@ class SchedulerStats:
43
default_factory=PrefixCacheStats)
44
45
spec_decoding_stats: Optional[SpecDecodingStats] = None
46
+ kv_connector_stats: Optional[dict[str, Any]] = None
47
48
num_corrupted_reqs: int = 0
49
vllm/v1/outputs.py
+10
-1
@@ -3,10 +3,14 @@
3
4
from abc import ABC, abstractmethod
5
from dataclasses import dataclass
6
-from typing import NamedTuple, Optional
6
+from typing import TYPE_CHECKING, NamedTuple, Optional
7
8
import torch
9
10
+if TYPE_CHECKING:
11
+ from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
12
+ KVConnectorStats)
13
+
14
15
class LogprobsLists(NamedTuple):
16
@@ -77,6 +81,11 @@ class KVConnectorOutput:
81
# [req_ids]
82
finished_sending: Optional[set[str]] = None
83
finished_recving: Optional[set[str]] = None
84
+ kv_connector_stats: Optional["KVConnectorStats"] = None
85
+
86
+ def is_empty(self):
87
+ return (not self.finished_sending and not self.finished_recving
88
+ and not self.kv_connector_stats)
89
90
91
# ModelRunnerOutput is serialized and sent to the scheduler process.
vllm/v1/worker/kv_connector_model_runner_mixin.py
+10
-1
@@ -13,6 +13,8 @@ from vllm.distributed.kv_transfer import (ensure_kv_transfer_shutdown,
13
get_kv_transfer_group,
14
has_kv_transfer_group)
15
from vllm.distributed.kv_transfer.kv_connector.base import KVConnectorBase
16
+from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
17
+ KVConnectorStats)
18
from vllm.forward_context import get_forward_context, set_forward_context
19
from vllm.logger import init_logger
20
from vllm.v1.outputs import (EMPTY_MODEL_RUNNER_OUTPUT, KVConnectorOutput,
@@ -119,4 +121,11 @@ class KVConnectorModelRunnerMixin:
121
output.finished_sending, output.finished_recving = (
122
kv_connector.get_finished(scheduler_output.finished_req_ids))
123
122
- kv_connector.clear_connector_metadata()
124
+ output.kv_connector_stats = KVConnectorModelRunnerMixin.\
125
+ get_kv_connector_stats()
126
+
127
+ @staticmethod
128
+ def get_kv_connector_stats() -> Optional[KVConnectorStats]:
129
+ if has_kv_transfer_group():
130
+ return get_kv_transfer_group().get_kv_connector_stats()
131
+ return None