[Bugfix][Chunk Transfer Adapter] deque mutated fix (#2102)

Signed-off-by: JuanPZuluaga <juanz9312@gmail.com>

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

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)

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)

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]