feat: add data parallel rank to KVEventBatch (#18925)

2025-06-03 17:14:20 -07:00
parent a8da78eac9
commit b712be98c7
6 changed files with 359 additions and 83 deletions
--- a/vllm/v1/core/sched/scheduler.py
+++ b/vllm/v1/core/sched/scheduler.py
@@ -80,7 +80,9 @@ class Scheduler(SchedulerInterface):
                config=self.vllm_config, role=KVConnectorRole.SCHEDULER)

        self.kv_event_publisher = EventPublisherFactory.create(
-            self.kv_events_config)
+            self.kv_events_config,
+            vllm_config.parallel_config.data_parallel_rank,
+        )

        num_gpu_blocks = self.cache_config.num_gpu_blocks
        assert num_gpu_blocks is not None and num_gpu_blocks > 0