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 11, 05:00 UTC
[NIXL][BUG FIX] Fix both failing issue and accuracy issue with nixl + host_buffer on CUDA (#30419)
Signed-off-by: Chendi Xue <chendi.xue@intel.com> Signed-off-by: Chendi.Xue <chendi.xue@intel.com> Co-authored-by: Nicolò Lucchesi <nlucches@redhat.com>
Chendi.Xue committed
Dec 18, 2025 at 16:10 UTC
6ca74bc11afc1a985ab80f4b94b2d3bcda764630
3 files changed
+53
-36
vllm/distributed/kv_transfer/kv_connector/utils.py
+21
@@ -4,6 +4,7 @@
4
KV cache helper for store.
5
"""
6
7
+from collections.abc import Iterator
8
from dataclasses import dataclass
9
from typing import TYPE_CHECKING, Literal
10
@@ -203,6 +204,26 @@ def copy_kv_blocks(
204
copy_fn(src_tensor, dst_tensor, src_indices, dst_indices)
205
206
207
+def yield_req_data(
208
+ scheduler_output,
209
+) -> Iterator[tuple[str, tuple[list[int], ...], bool]]:
210
+ """
211
+ Yields:
212
+ (req_id, new_block_id_groups, preempted)
213
+ """
214
+ # new requests
215
+ for req_data in scheduler_output.scheduled_new_reqs:
216
+ yield req_data.req_id, req_data.block_ids, False
217
+
218
+ # cached requests
219
+ cached_reqs = scheduler_output.scheduled_cached_reqs
220
+ yield from zip(
221
+ cached_reqs.req_ids,
222
+ cached_reqs.new_block_ids,
223
+ (req_id in cached_reqs.resumed_req_ids for req_id in cached_reqs.req_ids),
224
+ )
225
+
226
+
227
@dataclass
228
class TpKVTopology:
229
"""
vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+30
-15
@@ -23,7 +23,11 @@ from vllm import envs
23
from vllm.attention.backends.abstract import AttentionMetadata
24
from vllm.attention.selector import get_attn_backend
25
from vllm.config import VllmConfig
26
-from vllm.distributed.kv_transfer.kv_connector.utils import EngineId, TpKVTopology
26
+from vllm.distributed.kv_transfer.kv_connector.utils import (
27
+ EngineId,
28
+ TpKVTopology,
29
+ yield_req_data,
30
+)
31
from vllm.distributed.kv_transfer.kv_connector.v1.base import (
32
CopyBlocksOp,
33
KVConnectorBase_V1,
@@ -481,7 +485,7 @@ class NixlConnectorScheduler:
485
# New requests are added by update_state_after_alloc in
486
# the scheduler. Used to make metadata passed to Worker.
487
self._reqs_need_recv: dict[ReqId, tuple[Request, list[int]]] = {}
484
- self._reqs_need_save: dict[ReqId, tuple[Request, list[int]]] = {}
488
+ self._reqs_need_save: dict[ReqId, Request] = {}
489
# Reqs to send and their expiration time
490
self._reqs_need_send: dict[ReqId, float] = {}
491
self._reqs_in_batch: set[ReqId] = set()
@@ -627,16 +631,7 @@ class NixlConnectorScheduler:
631
if self.use_host_buffer and params.get("do_remote_decode"):
632
# NOTE: when accelerator is not directly supported by Nixl,
633
# prefilled blocks need to be saved to host memory before transfer.
630
-
631
- # save all blocks
632
- block_ids = blocks.get_block_ids()[0]
633
- # TODO: skip the blocks that are already in the host xfer buffer.
634
- # Currently, the host xfer buffer block is 1-to-1 mapped to device
635
- # kv blocks, so host blocks won't be flushed as long as its device
636
- # block is not overwritten; and it will be safe to skip saving them
637
- # to host xfer buffer.
638
- if block_ids:
639
- self._reqs_need_save[request.request_id] = (request, block_ids)
634
+ self._reqs_need_save[request.request_id] = request
635
elif params.get("do_remote_prefill"):
636
if params.get("remote_block_ids"):
637
if all(
@@ -688,13 +683,32 @@ class NixlConnectorScheduler:
683
kv_transfer_params=req.kv_transfer_params,
684
)
685
691
- for req_id, (req, block_ids) in self._reqs_need_save.items():
686
+ # NOTE: For the prefill side, there might be a chance that an early added
687
+ # request is a chunked prefill, so we need to check if new blocks are added
688
+ for req_id, new_block_id_groups, _ in yield_req_data(scheduler_output):
689
+ req_to_save = self._reqs_need_save.get(req_id)
690
+ if req_to_save is None or new_block_id_groups is None:
691
+ continue
692
+ req = req_to_save
693
+
694
assert req.kv_transfer_params is not None
695
meta.add_new_req_to_save(
696
request_id=req_id,
695
- local_block_ids=block_ids,
697
+ local_block_ids=new_block_id_groups[0],
698
kv_transfer_params=req.kv_transfer_params,
699
)
700
+ assert scheduler_output.num_scheduled_tokens is not None
701
+ num_scheduled_tokens = scheduler_output.num_scheduled_tokens[req_id]
702
+ is_partial = (
703
+ req.num_computed_tokens + num_scheduled_tokens
704
+ ) < req.num_prompt_tokens
705
+ if not is_partial:
706
+ # For non-partial prefills, once new req_meta is scheduled, it
707
+ # can be removed from _reqs_need_save.
708
+ # For partial prefill case, we will retain the request in
709
+ # _reqs_need_save until all blocks are scheduled with req_meta.
710
+ # Therefore, only pop if `not is_partial`.
711
+ self._reqs_need_save.pop(req_id)
712
713
meta.reqs_to_send = self._reqs_need_send
714
meta.reqs_in_batch = self._reqs_in_batch
@@ -702,7 +716,6 @@ class NixlConnectorScheduler:
716
717
# Clear the list once workers start the transfers
718
self._reqs_need_recv.clear()
705
- self._reqs_need_save.clear()
719
self._reqs_in_batch = set()
720
self._reqs_not_processed = set()
721
self._reqs_need_send = {}
@@ -748,6 +761,8 @@ class NixlConnectorScheduler:
761
# Also include the case of a P/D Prefill request with immediate
762
# block free (eg abort). Stop tracking this request.
763
self._reqs_not_processed.add(request.request_id)
764
+ # Clear _reqs_need_save if a request is aborted as partial prefill.
765
+ self._reqs_need_save.pop(request.request_id, None)
766
return False, None
767
768
# TODO: check whether block_ids actually ever be 0. If not we could
vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
+2
-21
@@ -1,7 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from collections import defaultdict
4
-from collections.abc import Iterable, Iterator
4
+from collections.abc import Iterable
5
from dataclasses import dataclass
6
from itertools import islice
7
from typing import Any, ClassVar
@@ -12,6 +12,7 @@ from vllm.attention.backends.abstract import AttentionBackend, AttentionMetadata
12
from vllm.attention.layer import Attention
13
from vllm.config import VllmConfig, get_layers_from_vllm_config
14
from vllm.distributed.kv_events import BlockRemoved, BlockStored, KVCacheEvent
15
+from vllm.distributed.kv_transfer.kv_connector.utils import yield_req_data
16
from vllm.distributed.kv_transfer.kv_connector.v1 import (
17
KVConnectorBase_V1,
18
KVConnectorRole,
@@ -516,23 +517,3 @@ class OffloadingConnectorWorker:
517
del self._store_jobs[req_id]
518
519
return finished_sending, finished_recving
519
-
520
-
521
-def yield_req_data(
522
- scheduler_output,
523
-) -> Iterator[tuple[str, tuple[list[int], ...], bool]]:
524
- """
525
- Yields:
526
- (req_id, new_block_id_groups, preempted)
527
- """
528
- # new requests
529
- for req_data in scheduler_output.scheduled_new_reqs:
530
- yield req_data.req_id, req_data.block_ids, False
531
-
532
- # cached requests
533
- cached_reqs = scheduler_output.scheduled_cached_reqs
534
- yield from zip(
535
- cached_reqs.req_ids,
536
- cached_reqs.new_block_ids,
537
- (req_id in cached_reqs.resumed_req_ids for req_id in cached_reqs.req_ids),
538
- )