From ec6ee1e33684df0c08098f14efd3171364f80f03 Mon Sep 17 00:00:00 2001 From: Juan Pablo Zuluaga <46724788+JuanPZuluaga@users.noreply.github.com> Date: Mon, 23 Mar 2026 18:09:04 +0100 Subject: [PATCH] [Bugfix][Chunk Transfer Adapter] deque mutated fix (#2102) Signed-off-by: JuanPZuluaga --- .../omni_connectors/test_chunk_transfer_adapter.py | 3 ++- .../distributed/omni_connectors/transfer_adapter/base.py | 4 ++++ .../transfer_adapter/chunk_transfer_adapter.py | 4 ++-- vllm_omni/worker/gpu_model_runner.py | 6 +++++- 4 files changed, 13 insertions(+), 4 deletions(-) diff --git a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py index fb9e22b84..dddf49a05 100644 --- a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py +++ b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py @@ -51,6 +51,7 @@ def build_adapter(monkeypatch, mocker: MockerFixture): self.config = config self._pending_load_reqs = deque() self._finished_load_reqs = set() + self._cancelled_load_reqs = set() self._pending_save_reqs = deque() self._finished_save_reqs = set() self.stop_event = threading.Event() @@ -212,7 +213,7 @@ def test_cleanup_clears_all_state(build_adapter): assert req_id not in adapter.get_req_chunk assert req_id not in adapter.requests_with_ready_chunks assert req_id not in adapter.request_ids_mapping - assert all(getattr(r, "request_id", None) != req_id for r in adapter._pending_load_reqs) + assert req_id in adapter._cancelled_load_reqs assert req_id not in adapter._finished_load_reqs assert ext_id not in adapter.put_req_chunk diff --git a/vllm_omni/distributed/omni_connectors/transfer_adapter/base.py b/vllm_omni/distributed/omni_connectors/transfer_adapter/base.py index 9dde0c626..bbaca48e0 100644 --- a/vllm_omni/distributed/omni_connectors/transfer_adapter/base.py +++ b/vllm_omni/distributed/omni_connectors/transfer_adapter/base.py @@ -25,6 +25,7 @@ class OmniTransferAdapterBase: self._pending_load_reqs = deque() # Requests that have successfully retrieved data self._finished_load_reqs = set() + self._cancelled_load_reqs: set[str] = set() # Requests that are waiting to be saved self._pending_save_reqs = deque() @@ -60,6 +61,9 @@ class OmniTransferAdapterBase: break request = self._pending_load_reqs.popleft() request_id = request.request_id + if request_id in self._cancelled_load_reqs: + self._cancelled_load_reqs.discard(request_id) + continue self.request_ids_mapping[request_id] = request.external_req_id try: is_success = self._poll_single_request(request) diff --git a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py index e190ab6b7..6b5319885 100644 --- a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py +++ b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py @@ -94,6 +94,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): return if not hasattr(request, "additional_information"): request.additional_information = None + self._cancelled_load_reqs.discard(request.request_id) self._pending_load_reqs.append(request) with self._recv_cond: self._recv_cond.notify() @@ -272,8 +273,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): self.requests_with_ready_chunks.discard(request_id) self.request_ids_mapping.pop(request_id, None) - remaining = deque(r for r in self._pending_load_reqs if getattr(r, "request_id", None) != request_id) - self._pending_load_reqs = remaining + self._cancelled_load_reqs.add(request_id) self._finished_load_reqs.discard(request_id) self.put_req_chunk.pop(external_req_id, None) diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index 396567002..a7441f674 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -1323,9 +1323,13 @@ class OmniGPUModelRunner(GPUModelRunner): use_cascade_attn=False, ) # Force eager for unwrapped code predictors (AR loops / multinomial). + # When talker_mtp is not a CUDAGraphWrapper, it manages its own CUDA + # graphs internally (code_predictor has its own bucket sizes). if not isinstance(self.talker_mtp, CUDAGraphWrapper): _cudagraph_mode = CUDAGraphMode.NONE - num_tokens_padded = batch_desc.num_tokens + num_tokens_padded = decode_batch_size + else: + num_tokens_padded = batch_desc.num_tokens req_input_ids = self.talker_mtp_input_ids.gpu[:num_tokens_padded] req_embeds = self.talker_mtp_inputs_embeds.gpu[:num_tokens_padded] last_talker_hidden = self.last_talker_hidden.gpu[:num_tokens_padded]