0565f1fdec · tduncan/Nvidia-Optimization-Stack-VLLM

@@ -3,7 +3,6 @@

3 import asyncio

4 import threading

5 import time

6 -import uuid

6 from collections import defaultdict

7 from concurrent.futures import ThreadPoolExecutor

8 from dataclasses import dataclass

@@ -83,28 +82,10 @@ class RecvReqMeta:

82 @dataclass

83 class SendBlockMeta:

84 local_block_ids: list[int]

86 - ready: threading.Event

85 + ready: asyncio.Event

86 expire_time: float = float("inf")

87

88

90 -@dataclass

91 -class SendReqMeta:

92 - reqs: dict[ReqId, SendBlockMeta]

93 - lock: threading.Lock

94 -

95 -

96 -@dataclass

97 -class FinishedSendReqSet:

98 - set: set[ReqId]

99 - lock: threading.Lock

100 -

101 -

102 -@dataclass

103 -class FinishedReceiveReqSet:

104 - set: set[ReqId]

105 - lock: asyncio.Lock

106 -

107 -

89 class MooncakeConnectorMetadata(KVConnectorMetadata):

90 def __init__(self):

91 self.reqs_to_recv: dict[ReqId, RecvReqMeta] = {}

@@ -437,39 +418,50 @@ class MooncakeConnectorWorker:

418

419 assert vllm_config.kv_transfer_config

420 self.kv_role = vllm_config.kv_transfer_config.kv_role

440 - self.num_workers = vllm_config.kv_transfer_config.kv_connector_extra_config.get(

441 - "num_workers", 10

421 + self.num_sender_workers = (

422 + vllm_config.kv_transfer_config.kv_connector_extra_config.get(

423 + "num_workers", 10

424 + )

425 )

426 + # Create more tasks than workers to keep the thread pool saturated.

427 + # Tasks can await async events, so a surplus (2x is a robust heuristic)

428 + # prevents workers from idling.

429 + self.num_sender_tasks = self.num_sender_workers * 2

430

431 self.kv_caches_base_addr: list[int] = []

432 self.device_kv_caches: dict[str, torch.Tensor] = {}

446 - self.reqs_need_send: SendReqMeta = SendReqMeta(reqs={}, lock=threading.Lock())

433 + self.reqs_need_send: dict[ReqId, SendBlockMeta] = {}

434

435 # For kv_both, we will act both prefiller and decoder.

436 if self.kv_role != "kv_consumer":

450 - # Background thread for sending kvcaches to D.

451 - self._mooncake_sender_t: threading.Thread | None = None

452 - # Background thread for processing new sending requests.

437 + # Background threads for sending kvcaches to D.

438 self._sender_executor = ThreadPoolExecutor(

454 - max_workers=self.num_workers, thread_name_prefix="vllm-mooncake-sender"

439 + max_workers=self.num_sender_workers,

440 + thread_name_prefix="vllm-mooncake-sender",

441 )

442 logger.debug(

457 - "Mooncake Prefiller: use %d workers to send kvcaches", self.num_workers

443 + "Mooncake Prefiller: use %d workers to send kvcaches",

444 + self.num_sender_workers,

445 )

446 + # An asyncio queue to buffer incoming requests for the sender

447 + self.sender_worker_queue = asyncio.Queue[tuple[bytes, bytes]]()

448 + self.sender_loop = asyncio.new_event_loop()

449 + # Background thread for processing new sending requests.

450 + self._sender_listener_t = threading.Thread(

451 + target=_async_loop, args=(self.sender_loop,), daemon=True

452 + )

453 + self._sender_listener_t.start()

454 +

455 if self.kv_role != "kv_producer":

456 self.receiver_loop = asyncio.new_event_loop()

457 self._mooncake_receiver_t = threading.Thread(

462 - target=self._receiver_loop, args=(self.receiver_loop,), daemon=True

458 + target=_async_loop, args=(self.receiver_loop,), daemon=True

459 )

460 self._mooncake_receiver_t.start()

461 logger.debug("Mooncake Decoder: start receiver thread")

462

467 - self.finished_sending_reqs: FinishedSendReqSet = FinishedSendReqSet(

468 - set(), threading.Lock()

469 - )

470 - self.finished_recving_reqs: FinishedReceiveReqSet = FinishedReceiveReqSet(

471 - set(), asyncio.Lock()

472 - )

463 + self.finished_sending_reqs: set[ReqId] = set()

464 + self.finished_recving_reqs: set[ReqId] = set()

465

466 self.block_size = vllm_config.cache_config.block_size

467 self.model_config = vllm_config.model_config

@@ -500,7 +492,6 @@ class MooncakeConnectorWorker:

492 attn_backend=backend,

493 )

494

503 - self.zmq_ctx = zmq.Context()

495 self.async_zmq_ctx = zmq.asyncio.Context()

496 self._encoder = msgspec.msgpack.Encoder()

497 self._decoder = msgspec.msgpack.Decoder(MooncakeAgentMetadata)

@@ -510,21 +501,17 @@ class MooncakeConnectorWorker:

501

502 def shutdown(self):

503 """Cleanup background threads on destruction."""

513 - self.zmq_ctx.term()

504 self.async_zmq_ctx.term()

505 if self.kv_role != "kv_consumer":

506 self._sender_executor.shutdown(wait=False)

517 - if self._mooncake_sender_t:

518 - self._mooncake_sender_t.join()

507 + if self.sender_loop.is_running():

508 + self.sender_loop.call_soon_threadsafe(self.sender_loop.stop)

509 + self._sender_listener_t.join()

510 if self.kv_role != "kv_producer" and self.receiver_loop.is_running():

511 self.receiver_loop.call_soon_threadsafe(self.receiver_loop.stop)

512 self._mooncake_receiver_t.join()

513

523 - def _receiver_loop(self, loop: asyncio.AbstractEventLoop):

524 - asyncio.set_event_loop(loop)

525 - loop.run_forever()

526 -

527 - def _mooncake_sender(

514 + async def _mooncake_sender_listener(

515 self, ready_event: threading.Event, base_port: int, tp_rank: int

516 ):

517 """

@@ -532,93 +519,86 @@ class MooncakeConnectorWorker:

519 to a thread pool, and sends acknowledgments upon completion.

520 """

521

535 - frontend_path = make_zmq_path("tcp", self.hostname, base_port + tp_rank)

536 - frontend = make_zmq_socket(self.zmq_ctx, frontend_path, zmq.ROUTER)

537 - logger.debug("Mooncake sender starting listening on path: %s", frontend_path)

538 -

539 - backend_path = make_zmq_path("inproc", str(uuid.uuid4()))

540 - backend = make_zmq_socket(self.zmq_ctx, backend_path, zmq.PULL)

522 + path = make_zmq_path("tcp", self.hostname, base_port + tp_rank)

523 + sock = make_zmq_socket(self.async_zmq_ctx, path, zmq.ROUTER)

524 + logger.debug("Mooncake sender starting listening on path: %s", path)

525

542 - poller = zmq.Poller()

543 - poller.register(frontend, zmq.POLLIN)

544 - poller.register(backend, zmq.POLLIN)

526 + # Create async worker tasks that process items from the queue

527 + sender_tasks = [

528 + asyncio.create_task(self._sender_worker(sock))

529 + for _ in range(self.num_sender_tasks)

530 + ]

531

532 ready_event.set()

533

534 try:

535 while True:

550 - sockets = dict(poller.poll())

551 -

552 - if frontend in sockets:

553 - identity, _, metadata_bytes = frontend.recv_multipart()

554 - self._sender_executor.submit(

555 - self._sender_worker,

556 - identity,

557 - metadata_bytes,

558 - backend_path,

559 - )

560 -

561 - if backend in sockets:

562 - identity, status = backend.recv_multipart()

563 - frontend.send_multipart((identity, b"", status))

564 -

536 + identity, _, metadata_bytes = await sock.recv_multipart()

537 + await self.sender_worker_queue.put((identity, metadata_bytes))

538 except zmq.ContextTerminated:

539 logger.debug("ZMQ context terminated, exiting Mooncake sender thread.")

540 except Exception as e:

541 logger.error("Error in Mooncake sender thread: %s. Exiting thread.", str(e))

542 finally:

570 - frontend.close()

571 - backend.close()

572 -

573 - def _sender_worker(

574 - self, identity: bytes, metadata_bytes: bytes, worker_channel_path: str

575 - ):

576 - status = TRANS_ERROR

543 + # Clean up worker tasks

544 + for task in sender_tasks:

545 + task.cancel()

546 + await asyncio.gather(*sender_tasks, return_exceptions=True)

547 + sock.close()

548

578 - try:

579 - metadata = self._decoder.decode(metadata_bytes)

580 - self.send_kv_to_decode(metadata)

581 - status = TRANS_DONE

582 - except Exception as e:

583 - logger.error("Error processing Mooncake handshake: %s", e)

584 - finally:

585 - pusher = make_zmq_socket(self.zmq_ctx, worker_channel_path, zmq.PUSH)

549 + async def _sender_worker(self, sock: zmq.asyncio.Socket):

550 + while True:

551 try:

587 - pusher.send_multipart((identity, status))

588 - except zmq.ZMQError as e:

589 - logger.warning(

590 - "Internal error, maybe the server is shutting down. Error: %s",

591 - e,

592 - )

593 - finally:

594 - pusher.close()

595 -

596 - def send_kv_to_decode(self, meta: MooncakeAgentMetadata):

552 + identity, metadata_bytes = await self.sender_worker_queue.get()

553 + try:

554 + metadata = self._decoder.decode(metadata_bytes)

555 + await self.send_kv_to_decode(metadata)

556 + await sock.send_multipart((identity, b"", TRANS_DONE))

557 + except Exception as e:

558 + logger.error("Error processing Mooncake xfer request: %s", e)

559 + await sock.send_multipart((identity, b"", TRANS_ERROR))

560 + finally:

561 + self.sender_worker_queue.task_done()

562 + except asyncio.CancelledError:

563 + break

564 + except Exception as e:

565 + logger.error("Error in _sender_worker: %s", e)

566 +

567 + async def send_kv_to_decode(self, meta: MooncakeAgentMetadata):

568 send_reqs: list[tuple[ReqId, SendBlockMeta]] = []

598 - with self.reqs_need_send.lock:

599 - for req_id in meta.request_ids:

600 - send_meta = self.reqs_need_send.reqs.get(req_id)

601 - if send_meta is None:

602 - logger.warning("Request %s not found in reqs_need_send", req_id)

603 - return

604 - # Mark it as not expired. We will send it now.

605 - send_meta.expire_time = float("inf")

606 - send_reqs.append((req_id, send_meta))

569 + for req_id in meta.request_ids:

570 + send_meta = self.reqs_need_send.get(req_id)

571 + if send_meta is None:

572 + logger.warning("Request %s not found in reqs_need_send", req_id)

573 + return

574 + # Mark it as not expired. We will send it now.

575 + send_meta.expire_time = float("inf")

576 + send_reqs.append((req_id, send_meta))

577 +

578 + src_ptrs, dst_ptrs, lengths = await self._build_transfer_params(send_reqs, meta)

579 + remote_session = f"{meta.remote_hostname}:{meta.remote_port}"

580 + ret_value = await self.sender_loop.run_in_executor(

581 + self._sender_executor,

582 + self._send_blocks,

583 + remote_session,

584 + src_ptrs,

585 + dst_ptrs,

586 + lengths,

587 + )

588

608 - self._send_blocks(send_reqs, meta)

589 + if ret_value != 0:

590 + raise RuntimeError(f"Error in batch_transfer_sync_write: {ret_value}")

591

610 - with self.reqs_need_send.lock:

611 - for req_id in meta.request_ids:

612 - del self.reqs_need_send.reqs[req_id]

592 + for req_id in meta.request_ids:

593 + del self.reqs_need_send[req_id]

594

614 - with self.finished_sending_reqs.lock:

615 - self.finished_sending_reqs.set.update(meta.request_ids)

595 + self.finished_sending_reqs.update(meta.request_ids)

596

617 - def _send_blocks(

597 + async def _build_transfer_params(

598 self,

599 send_reqs: list[tuple[ReqId, SendBlockMeta]],

600 agent_meta: MooncakeAgentMetadata,

621 - ):

601 + ) -> tuple[list[int], list[int], list[int]]:

602 src_ptrs = []

603 dst_ptrs = []

604 lengths = []

@@ -631,7 +611,7 @@ class MooncakeConnectorWorker:

611 for (req_id, send_meta), remote_block_ids in zip(

612 send_reqs, agent_meta.block_ids

613 ):

634 - send_meta.ready.wait()

614 + await send_meta.ready.wait()

615

616 num_remote_blocks = len(remote_block_ids)

617 if num_remote_blocks == 0:

@@ -670,18 +650,26 @@ class MooncakeConnectorWorker:

650 remote_session,

651 )

652

653 + return src_ptrs, dst_ptrs, lengths

654 +

655 + def _send_blocks(

656 + self,

657 + remote_session: str,

658 + src_ptrs: list[int],

659 + dst_ptrs: list[int],

660 + lengths: list[int],

661 + ) -> int:

662 start_time = time.perf_counter()

663 ret_value = self.engine.batch_transfer_sync_write(

664 remote_session, src_ptrs, dst_ptrs, lengths

665 )

677 - if ret_value != 0:

678 - raise RuntimeError(f"Error in batch_transfer_sync_write: {ret_value}")

679 -

680 - logger.debug(

681 - "Sending to %s done, took %s",

682 - remote_session,

683 - time.perf_counter() - start_time,

684 - )

666 + if ret_value == 0:

667 + logger.debug(

668 + "Sending to %s done, took %s",

669 + remote_session,

670 + time.perf_counter() - start_time,

671 + )

672 + return ret_value

673

674 def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):

675 """Register the KV Cache data in mooncake."""

@@ -740,41 +728,63 @@ class MooncakeConnectorWorker:

728 return

729

730 ready_event = threading.Event()

743 - self._mooncake_sender_t = threading.Thread(

744 - target=self._mooncake_sender,

745 - args=(ready_event, self.side_channel_port, self.tp_rank),

746 - daemon=True,

747 - name="mooncake_sender",

731 + asyncio.run_coroutine_threadsafe(

732 + self._mooncake_sender_listener(

733 + ready_event, self.side_channel_port, self.tp_rank

734 + ),

735 + self.sender_loop,

736 )

749 - self._mooncake_sender_t.start()

737 ready_event.wait() # Wait for listener ZMQ socket to be ready.

738

739 async def fetch_finished_recving_reqs(self) -> set[ReqId]:

753 - async with self.finished_recving_reqs.lock:

754 - finished_recving_reqs = self.finished_recving_reqs.set

755 - self.finished_recving_reqs.set = set()

740 + finished_recving_reqs = self.finished_recving_reqs

741 + self.finished_recving_reqs = set()

742 return finished_recving_reqs

743

744 + async def fetch_finished_sending_reqs(self) -> set[ReqId]:

745 + finished_sending_reqs = self.finished_sending_reqs

746 + self.finished_sending_reqs = set()

747 +

748 + # Handle timeout to avoid stranding blocks on remote.

749 + now = time.perf_counter()

750 + expired_reqs = [

751 + req_id

752 + for req_id, send_meta in self.reqs_need_send.items()

753 + if send_meta.expire_time < now

754 + ]

755 + for req_id in expired_reqs:

756 + logger.warning(

757 + "Request %s timed out after %d seconds without "

758 + "being sent. Freeing its blocks on the producer side.",

759 + req_id,

760 + envs.VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT,

761 + )

762 + del self.reqs_need_send[req_id]

763 + if expired_reqs:

764 + finished_sending_reqs.update(expired_reqs)

765 +

766 + return finished_sending_reqs

767 +

768 def get_finished(self) -> tuple[set[str] | None, set[str] | None]:

769 """

770 Get requests that are done sending or recving on this specific worker.

771 The scheduler process (via the MultiprocExecutor) will use this output

772 to track which workers are done.

773 """

764 - fut = None

774 + recv_fut = None

775 + send_fut = None

776 if self.kv_role != "kv_producer":

766 - fut = asyncio.run_coroutine_threadsafe(

777 + recv_fut = asyncio.run_coroutine_threadsafe(

778 self.fetch_finished_recving_reqs(), self.receiver_loop

779 )

780

781 if self.kv_role != "kv_consumer":

771 - with self.finished_sending_reqs.lock:

772 - finished_sending_reqs = self.finished_sending_reqs.set

773 - self.finished_sending_reqs.set = set()

774 - else:

775 - finished_sending_reqs = set()

782 + send_fut = asyncio.run_coroutine_threadsafe(

783 + self.fetch_finished_sending_reqs(), self.sender_loop

784 + )

785

777 - finished_recving_reqs = fut.result() if fut else set()

786 + finished_recving_reqs = recv_fut.result() if recv_fut else set()

787 + finished_sending_reqs = send_fut.result() if send_fut else set()

788

789 if finished_sending_reqs or finished_recving_reqs:

790 logger.debug(

@@ -785,25 +795,6 @@ class MooncakeConnectorWorker:

795 len(finished_recving_reqs),

796 )

797

788 - # Handle timeout to avoid stranding blocks on remote.

789 - now = time.perf_counter()

790 - with self.reqs_need_send.lock:

791 - expired_reqs = [

792 - req_id

793 - for req_id, send_meta in self.reqs_need_send.reqs.items()

794 - if send_meta.expire_time < now

795 - ]

796 - for req_id in expired_reqs:

797 - logger.warning(

798 - "Request %s timed out after %d seconds without "

799 - "being sent. Freeing its blocks on the producer side.",

800 - req_id,

801 - envs.VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT,

802 - )

803 - del self.reqs_need_send.reqs[req_id]

804 - if expired_reqs:

805 - finished_sending_reqs.update(expired_reqs)

806 -

798 return finished_sending_reqs or None, finished_recving_reqs or None

799

800 async def receive_kv(self, path: str, req_blocks: list[tuple[str, list[int]]]):

@@ -844,8 +835,7 @@ class MooncakeConnectorWorker:

835 finally:

836 sock.close()

837

847 - async with self.finished_recving_reqs.lock:

848 - self.finished_recving_reqs.set.update(req_ids)

838 + self.finished_recving_reqs.update(req_ids)

839

840 logger.debug("pulling kv_caches for %s finished", req_ids)

841

@@ -865,6 +855,24 @@ class MooncakeConnectorWorker:

855

856 return kv_pulls

857

858 + async def record_send_reqs(self, metadata: MooncakeConnectorMetadata):

859 + for req_id, block_ids in metadata.reqs_to_send.items():

860 + if block_ids:

861 + # Already gone through request_finished()

862 + send_meta = self.reqs_need_send[req_id]

863 + send_meta.local_block_ids = block_ids

864 + send_meta.expire_time = (

865 + time.perf_counter() + envs.VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT

866 + )

867 + send_meta.ready.set()

868 + else:

869 + # From update_state_after_alloc(),

870 + # but not reach request_finished() yet

871 + self.reqs_need_send[req_id] = SendBlockMeta(

872 + local_block_ids=[],

873 + ready=asyncio.Event(),

874 + )

875 +

876 def start_load_kv(self, metadata: MooncakeConnectorMetadata):

877 if self.kv_role != "kv_producer":

878 kv_pulls = self.group_kv_pull(metadata)

@@ -874,23 +882,9 @@ class MooncakeConnectorWorker:

882 )

883

884 if self.kv_role != "kv_consumer":

877 - with self.reqs_need_send.lock:

878 - for req_id, block_ids in metadata.reqs_to_send.items():

879 - if block_ids:

880 - # Already gone through request_finished()

881 - send_meta = self.reqs_need_send.reqs[req_id]

882 - send_meta.local_block_ids = block_ids

883 - send_meta.ready.set()

884 - send_meta.expire_time = (

885 - time.perf_counter()

886 - + envs.VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT

887 - )

888 - else:

889 - # From update_state_after_alloc(),

890 - # but not reach request_finished() yet

891 - self.reqs_need_send.reqs[req_id] = SendBlockMeta(

892 - local_block_ids=[], ready=threading.Event()

893 - )

885 + asyncio.run_coroutine_threadsafe(

886 + self.record_send_reqs(metadata), self.sender_loop

887 + )

888

889

890 def group_concurrent_contiguous(

@@ -917,3 +911,8 @@ def get_mooncake_side_channel_port(vllm_config: VllmConfig) -> int:

911 + vllm_config.parallel_config.data_parallel_index

912 * vllm_config.parallel_config.tensor_parallel_size

913 )

914 +

915 +

916 +def _async_loop(loop: asyncio.AbstractEventLoop):

917 + asyncio.set_event_loop(loop)

918 + loop.run_forever()

[P/D] Refactor mooncake connector sender thread using async coroutines (#31573)