@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 48f309029a

Fork 0

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

[NIXL][Misc] Expose metrics from NIXL for logging to CLI (#25388)

Signed-off-by: NickLucche <nlucches@redhat.com>

Nicolò Lucchesi committed Oct 3, 2025 at 12:47 UTC 48f309029a3aa66b40b0613a9a7361b654b4865a

4 files changed +127 -28

requirements/kv_connectors.txt

+1 -1

@@ -1,2 +1,2 @@

1 lmcache

2 -nixl >= 0.5.1 # Required for disaggregated prefill

2 +nixl >= 0.6.0 # Required for disaggregated prefill

tests/v1/kv_connector/unit/test_nixl_connector.py

+54 -11

@@ -57,6 +57,26 @@ def clear_kv_transfer():

57 ensure_kv_transfer_shutdown()

60 +def get_default_xfer_telemetry(xferDurationS: float = 1,

61 + postDurationS: float = 1,

62 + totalBytes: int = 1,

63 + descCount: int = 1) -> dict:

64 +

65 + class AttributeDict(dict):

66 + __slots__ = ()

67 + __getattr__ = dict.__getitem__

68 + __setattr__ = dict.__setitem__ # type: ignore[assignment]

69 +

70 + # We can't instantiate nixlXferTelemetry because it's read only and

71 + # ray env does not have NIXL, so we must fake it

72 + return AttributeDict(

73 + xferDuration=xferDurationS * 1e6, # in us

74 + postDuration=postDurationS * 1e6, # in us

75 + totalBytes=totalBytes,

76 + descCount=descCount,

77 + )

78 +

79 +

80 class FakeNixlWrapper:

81 """Mock implementation of NixlWrapper for testing.

@@ -132,6 +152,9 @@ class FakeNixlWrapper:

152 def transfer(self, handle: int) -> str:

153 return "PROC"

154

155 + def get_xfer_telemetry(self, handle: int) -> dict:

156 + return get_default_xfer_telemetry()

157 +

158 ############################################################

159 # Follow are for changing the behavior during testing.

160 ############################################################

@@ -169,6 +192,11 @@ nixl_agent = FakeNixlWrapper

192 with open(os.path.join(pkg_root, "__init__.py"), "w") as f:

193 f.write(stub)

194

195 + # Mock nixlXferTelemetry class

196 + pkg_root2 = os.path.join(td, "nixl", "_bindings")

197 + os.makedirs(pkg_root2, exist_ok=True)

198 + with open(os.path.join(pkg_root2, "__init__.py"), "w") as f:

199 + f.write("class nixlXferTelemetry: pass")

200 # touch parent package

201 open(os.path.join(td, "nixl", "__init__.py"), "w").close()

202 yield td

@@ -575,7 +603,7 @@ def test_kv_connector_stats(dist_init):

603

604 # Verify stats values are recorded

605 assert not stats_after_transfer.is_empty()

578 - assert stats_after_transfer.data["num_successful_transfers"] == 1

606 + assert stats_after_transfer.num_successful_transfers == 1

607

608 # Verify stats are reset after retrieval

609 stats_after_reset = connector.get_kv_connector_stats()

@@ -599,16 +627,21 @@ def test_kv_connector_stats_aggregation():

627

628 # Record different transfers on each worker

629 # Worker 1: 2 transfers

602 - worker1_stats.record_transfer()

603 - worker1_stats.record_transfer()

630 + stats = get_default_xfer_telemetry()

631 + worker1_stats.record_transfer(stats)

632 + worker1_stats.record_transfer(stats)

633

634 # Worker 2: 1 transfer

606 - worker2_stats.record_transfer()

635 + worker2_stats.record_transfer(stats)

636

637 # Worker 3: 3 transfers

609 - worker3_stats.record_transfer()

610 - worker3_stats.record_transfer()

611 - worker3_stats.record_transfer()

638 + stats = get_default_xfer_telemetry(xferDurationS=2,

639 + postDurationS=2,

640 + totalBytes=2,

641 + descCount=2)

642 + worker3_stats.record_transfer(stats)

643 + worker3_stats.record_transfer(stats)

644 + worker3_stats.record_transfer(stats)

645

646 # Create ModelRunnerOutput instances for each worker

647 worker_outputs = []

@@ -636,7 +669,12 @@ def test_kv_connector_stats_aggregation():

669 aggregated_output.kv_connector_output.kv_connector_stats

670 assert isinstance(kv_connector_stats, NixlKVConnectorStats)

671 # Number of total transfers across all workers.

639 - assert kv_connector_stats.data["num_successful_transfers"] == 6

672 + assert kv_connector_stats.num_successful_transfers == 6

673 + # Logging proc, call reduce() to get CLI-friendly stats.

674 + cli_stats = kv_connector_stats.reduce()

675 + assert cli_stats["Avg xfer time (ms)"] == 1500.0

676 + assert cli_stats["Avg post time (ms)"] == 1500.0

677 + assert cli_stats["Avg number of descriptors"] == 1.5

678

679

680 def test_multi_kv_connector_stats_aggregation():

@@ -649,6 +687,7 @@ def test_multi_kv_connector_stats_aggregation():

687

688 from dataclasses import dataclass

689

690 + # Mock a KVConnectorStats class for testing aggregation over connectors.

691 @dataclass

692 class FooKVConnectorStats(KVConnectorStats):

693

@@ -676,7 +715,7 @@ def test_multi_kv_connector_stats_aggregation():

715 if nixl_count > 0:

716 nixl_stats = NixlKVConnectorStats()

717 for _ in range(nixl_count):

679 - nixl_stats.record_transfer()

718 + nixl_stats.record_transfer(get_default_xfer_telemetry())

719 data["NixlConnector"] = nixl_stats

720 if foo_count > 0:

721 foo_stats = FooKVConnectorStats()

@@ -712,8 +751,10 @@ def test_multi_kv_connector_stats_aggregation():

751 assert isinstance(kv_connector_stats, MultiKVConnectorStats)

752

753 # Validate per-connector totals across workers

715 - assert kv_connector_stats["NixlConnector"].data[

716 - "num_successful_transfers"] == 5

754 + assert isinstance(kv_connector_stats["NixlConnector"],

755 + NixlKVConnectorStats)

756 + assert kv_connector_stats["NixlConnector"].num_successful_transfers == 5

757 + assert isinstance(kv_connector_stats["FooConnector"], FooKVConnectorStats)

758 assert kv_connector_stats["FooConnector"].data["num_foo_transfers"] == 6

759

760

@@ -755,6 +796,8 @@ def test_abort_timeout_on_prefiller(monkeypatch, distributed_executor_backend):

796 "working_dir": working_dir, # ship fake nixl package

797 "env_vars": {

798 "VLLM_NIXL_ABORT_REQUEST_TIMEOUT": str(timeout),

799 + # TODO: for ray to carry over, remove once we set

800 + "NIXL_TELEMETRY_ENABLE": "1",

801 },

802 }

803 ray.init(runtime_env=runtime_env)

vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py

+69 -13

@@ -4,6 +4,7 @@ import contextlib

4 import copy

5 import logging

6 import math

7 +import os

8 import queue

9 import threading

10 import time

@@ -54,10 +55,12 @@ logger = init_logger(__name__)

55 # Lazy import nixl_wrapper to avoid loading nixl_bindings if nixl is not used

56 try:

57 from nixl._api import nixl_agent as NixlWrapper

58 + from nixl._bindings import nixlXferTelemetry

59 logger.info("NIXL is available")

60 except ImportError:

61 logger.warning("NIXL is not available")

62 NixlWrapper = None

63 + nixlXferTelemetry = None

65 try:

66 from nixl._api import nixl_agent_config

@@ -476,6 +479,9 @@ class NixlConnectorWorker:

479 self.nixl_backends = \

480 vllm_config.kv_transfer_config.get_from_extra_config(

481 "backends", ["UCX"])

482 + # TODO temporary, once nixl allows for telemetry flag in config

483 + # (next release), we can remove this env var.

484 + os.environ["NIXL_TELEMETRY_ENABLE"] = "1"

485 # Agent.

486 non_ucx_backends = [b for b in self.nixl_backends if b != "UCX"]

487 if nixl_agent_config is None:

@@ -1175,9 +1181,10 @@ class NixlConnectorWorker:

1181 for handle, _xfer_stime in handles:

1182 xfer_state = self.nixl_wrapper.check_xfer_state(handle)

1183 if xfer_state == "DONE":

1184 + # Get telemetry from NIXL

1185 + res = self.nixl_wrapper.get_xfer_telemetry(handle)

1186 + self.xfer_stats.record_transfer(res)

1187 self.nixl_wrapper.release_xfer_handle(handle)

1179 - # TODO (NickLucche) Get from NIXL telemetry once integrated

1180 - self.xfer_stats.record_transfer()

1188 elif xfer_state == "PROC":

1189 in_progress = True

1190 continue

@@ -1449,15 +1456,25 @@ class NixlKVConnectorStats(KVConnectorStats):

1456 """Container for transfer performance metrics"""

1457

1458 def __post_init__(self):

1452 - if "num_successful_transfers" not in self.data:

1453 - self.data["num_successful_transfers"] = 0

1459 + if not self.data:

1460 + # Empty container init, no data is passed in.

1461 + self.reset()

1462

1463 def reset(self):

1456 - self.data = {"num_successful_transfers": 0}

1464 + # Must be serializable

1465 + self.data: dict[str, list[float]] = {

1466 + "transfer_duration": [],

1467 + "post_duration": [],

1468 + "bytes_transferred": [],

1469 + "num_descriptors": [],

1470 + }

1471

1458 - def record_transfer(self):

1459 - # TODO: record actual transfer stats when available

1460 - self.data["num_successful_transfers"] += 1

1472 + def record_transfer(self, res: nixlXferTelemetry):

1473 + # Keep metrics units consistent with rest of the code: time us->s

1474 + self.data["transfer_duration"].append(res.xferDuration / 1e6)

1475 + self.data["post_duration"].append(res.postDuration / 1e6)

1476 + self.data["bytes_transferred"].append(res.totalBytes)

1477 + self.data["num_descriptors"].append(res.descCount)

1478

1479 def clone_and_reset(self) -> "NixlKVConnectorStats":

1480 old = copy.copy(self)

@@ -1465,16 +1482,55 @@ class NixlKVConnectorStats(KVConnectorStats):

1482 return old

1483

1484 def is_empty(self) -> bool:

1468 - return self.data["num_successful_transfers"] == 0

1485 + return self.num_successful_transfers == 0

1486

1487 def aggregate(self, other: KVConnectorStats) -> KVConnectorStats:

1488 if not other.is_empty():

1472 - self.data["num_successful_transfers"] += other.data[

1473 - "num_successful_transfers"]

1489 + for k, v in other.data.items():

1490 + accumulator = self.data[k]

1491 + assert isinstance(accumulator, list)

1492 + accumulator.extend(v)

1493 return self

1494

1495 def reduce(self) -> dict[str, Union[int, float]]:

1477 - # TODO: reduce stats to a single value, calculate latency/throughput

1496 + # Compute compact representative stats suitable for CLI logging

1497 + if self.is_empty():

1498 + return {

1499 + "Num successful transfers": 0,

1500 + "Avg xfer time (ms)": 0,

1501 + "P90 xfer time (ms)": 0,

1502 + "Avg post time (ms)": 0,

1503 + "P90 post time (ms)": 0,

1504 + "Avg MB per transfer": 0,

1505 + "Throughput (MB/s)": 0,

1506 + "Avg number of descriptors": 0,

1507 + }

1508 +

1509 + xfer_time = np.asarray(self.data["transfer_duration"])

1510 + post_time = np.asarray(self.data["post_duration"])

1511 + # Convert to MB for CLI logging.

1512 + mb = np.asarray(self.data["bytes_transferred"]) / 2**20

1513 + descs = np.asarray(self.data["num_descriptors"], dtype=np.uint32)

1514 + n = len(descs)

1515 + assert n == self.num_successful_transfers

1516 +

1517 + total_mb = mb.sum()

1518 + avg_mb = total_mb / n

1519 +

1520 + total_time_seconds = xfer_time.sum()

1521 + throughput_mb_s = total_mb / total_time_seconds

1522 +

1523 return {

1479 - "num_successful_transfers": self.data["num_successful_transfers"]

1524 + "Num successful transfers": n,

1525 + "Avg xfer time (ms)": round(xfer_time.mean() * 1e3, 3),

1526 + "P90 xfer time (ms)": round(np.percentile(xfer_time, 90) * 1e3, 3),

1527 + "Avg post time (ms)": round(post_time.mean() * 1e3, 3),

1528 + "P90 post time (ms)": round(np.percentile(post_time, 90) * 1e3, 3),

1529 + "Avg MB per transfer": round(avg_mb, 3),

1530 + "Throughput (MB/s)": round(throughput_mb_s, 3),

1531 + "Avg number of descriptors": round(descs.mean(), 1),

1532 }

1533 +

1534 + @property

1535 + def num_successful_transfers(self) -> int:

1536 + return len(self.data["transfer_duration"])

\ No newline at end of file

vllm/v1/metrics/loggers.py

+3 -3

@@ -62,7 +62,7 @@ class LoggingStatLogger(StatLoggerBase):

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)

65 + self.kv_connector_logging = KVConnectorLogging(kv_tranfer_config)

66 self.last_prompt_throughput: float = 0.0

67 self.last_generation_throughput: float = 0.0

@@ -101,7 +101,7 @@ class LoggingStatLogger(StatLoggerBase):

101 self.spec_decoding_logging.observe(

102 scheduler_stats.spec_decoding_stats)

103 if kv_connector_stats := scheduler_stats.kv_connector_stats:

104 - self.kv_transfer_logging.observe(kv_connector_stats)

104 + self.kv_connector_logging.observe(kv_connector_stats)

105 self.last_scheduler_stats = scheduler_stats

106

107 def log(self):

@@ -140,7 +140,7 @@ class LoggingStatLogger(StatLoggerBase):

140 self.prefix_caching_metrics.hit_rate * 100,

141 )

142 self.spec_decoding_logging.log(log_fn=log_fn)

143 - self.kv_transfer_logging.log(log_fn=log_fn)

143 + self.kv_connector_logging.log(log_fn=log_fn)

144

145 def log_engine_initialized(self):

146 if self.vllm_config.cache_config.num_gpu_blocks: