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 6, 14:30 UTC
[perf][connector] optimize build_connector_meta when host buffer transfer is not used (#37165)
Signed-off-by: youkaichao <youkaichao@gmail.com>
youkaichao committed
Mar 17, 2026 at 19:59 UTC
0fb142a454757ec2055000ca8a2607e797af3e71
1 file changed
+22
-12
vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+22
-12
@@ -815,20 +815,12 @@ class NixlConnectorScheduler:
815
# Only trigger 1 KV transfer per request.
816
params["do_remote_prefill"] = False
817
818
- def build_connector_meta(
818
+ def _build_save_meta(
819
self,
820
+ meta: NixlConnectorMetadata,
821
scheduler_output: SchedulerOutput,
821
- ) -> KVConnectorMetadata:
822
- meta = NixlConnectorMetadata()
823
-
824
- # Loop through scheduled reqs and convert to ReqMeta.
825
- for req_id, (req, block_ids) in self._reqs_need_recv.items():
826
- assert req.kv_transfer_params is not None
827
- meta.add_new_req_to_recv(
828
- request_id=req_id,
829
- local_block_ids=block_ids,
830
- kv_transfer_params=req.kv_transfer_params,
831
- )
822
+ ) -> None:
823
+ # only called when use_host_buffer is True to build the save metadata
824
825
# NOTE: For the prefill side, there might be a chance that an early added
826
# request is a chunked prefill, so we need to check if new blocks are added
@@ -858,6 +850,24 @@ class NixlConnectorScheduler:
850
# Therefore, only pop if `not is_partial`.
851
self._reqs_need_save.pop(req_id)
852
853
+ def build_connector_meta(
854
+ self,
855
+ scheduler_output: SchedulerOutput,
856
+ ) -> KVConnectorMetadata:
857
+ meta = NixlConnectorMetadata()
858
+
859
+ # Loop through scheduled reqs and convert to ReqMeta.
860
+ for req_id, (req, block_ids) in self._reqs_need_recv.items():
861
+ assert req.kv_transfer_params is not None
862
+ meta.add_new_req_to_recv(
863
+ request_id=req_id,
864
+ local_block_ids=block_ids,
865
+ kv_transfer_params=req.kv_transfer_params,
866
+ )
867
+
868
+ if self.use_host_buffer:
869
+ self._build_save_meta(meta, scheduler_output)
870
+
871
meta.reqs_to_send = self._reqs_need_send
872
meta.reqs_in_batch = self._reqs_in_batch
873
meta.reqs_not_processed = self._reqs_not_processed