aswinvisva · Nov 6, 2026 · Nov 9, 2026 · Nov 9, 2026 · Nov 9, 2026
@@ -8,6 +8,10 @@
 
 _low_latency_dispatch: bool = False
 
+# Event recorded after the last replay of a CUDA graph that contains host
+# functions; None once a drain has waited for it.
+_replay_event: torch.cuda.Event | None = None
+
 
 def set_low_latency_dispatch(enabled: bool) -> None:
     global _low_latency_dispatch
@@ -27,6 +31,39 @@ def launch_hostfunc(hostfunc, *args, **kwargs):
     return handle
 
 
+def record_hostfunc_graph_replay() -> None:
+    """Record an event after replaying a graph that contains host functions."""
+    global _replay_event
+    _replay_event = torch.cuda.Event()
+    _replay_event.record()
+
+
+def hostfuncs_pending() -> bool:
+    """Return True if a graph with host functions has replayed since the last drain.
+
+    Only host-side state is read, not device progress, so ranks that issue the
+    same replays and drains reach the same answer.
+    """
+    return _replay_event is not None
+
+
+def drain_captured_hostfuncs() -> None:
+    """Wait until the host functions of the last replayed graph have run.
+
+    A host function takes the GIL when it runs. A call that holds the GIL while
+    it waits on the device, such as a multimodal encoder that builds a cuDNN
+    plan, would otherwise deadlock with a host function that has not run yet.
+    The wait is skipped during stream capture, where host synchronization is
+    not allowed.
+    """
+    global _replay_event
+    if _replay_event is None or torch.cuda.is_current_stream_capturing():
+        return
+    # Event.synchronize releases the GIL, which lets the host functions run.
+    _replay_event.synchronize()
+    _replay_event = None
+
+
 def hostfunc(hostfunc):
 
     def wrapper(*args, **kwargs):

@@ -49,6 +49,7 @@
 )
 from tensorrt_llm.logger import logger
 
+from ..hostfunc import hostfuncs_pending
 from ..nccl_window_tensor_scope import discard_nccl_window_tensor_outputs
 from ..pyexecutor.runtime_hooks import DEFAULT_RUNTIME_HOOKS, RuntimeHooks, runtime_hooks_of
 from .modeling_multimodal_utils import (
@@ -1810,6 +1811,10 @@ def maybe_prefetch_mm_encoder_for_next_iter(
         max_prefetch_ahead = 0
     if max_prefetch_ahead <= 0:
         return 0
+    # The encoder can hold the GIL while it waits on the device. Until the host functions of a
+    # replayed graph have been drained, these requests are encoded when scheduled instead.
+    if hostfuncs_pending():
+        return 0
 
     aux_stream = _get_mm_aux_stream(max_prefetch_ahead)
     if aux_stream is None:

@@ -308,6 +308,12 @@ def _get_uncached_multimodal_params(
     return params_to_run
 
 
+def needs_multimodal_encoder(
+        multimodal_params: Optional[List[MultimodalParams]]) -> bool:
+    """Return True if any param has multimodal content but no embedding."""
+    return bool(_get_uncached_multimodal_params(multimodal_params or []))
+
+
 def _store_chunked_prefill_embeddings(
     multimodal_params: List[MultimodalParams],
     embeddings: List[torch.Tensor],

@@ -20,6 +20,7 @@
 
 from ..attention.backends.trtllm import TrtllmAttentionMetadata
 from ..distributed import Distributed
+from ..hostfunc import HOSTFUNC_USER_DATA_HANDLES, record_hostfunc_graph_replay
 from ..memory_buffer_utils import Buffers, get_memory_buffers
 from ..modules.multi_stream_utils import with_multi_stream
 from ..moe.expert_statistic import ExpertStatistic
@@ -289,6 +290,8 @@ def __init__(self, config: CUDAGraphRunnerConfig):
         self.graph_outputs: Dict[KeyType,
                                  Callable[[], Optional[torch.Tensor]]] = {}
         self.graph_metadata: Dict[KeyType, Dict[str, Any]] = {}
+        # Keys of the graphs that captured host functions.
+        self._hostfunc_graph_keys: set[KeyType] = set()
         self.memory_pool = None
         self.padding_dummy_requests: Dict[int, LlmRequest] = {}
         self.dynamic_draft_len_mapping = config.dynamic_draft_len_mapping
@@ -912,13 +915,16 @@ def run_warmup():
             # setup/capture; release its reference before entering.
             output = None
             graph_pool = self.memory_pool or torch.cuda.graph_pool_handle()
+            num_hostfuncs = len(HOSTFUNC_USER_DATA_HANDLES)
             with nccl_window_graph_capture(graph, graph_pool):
                 # A pre-replay hook reads the static inputs before the replay, so their copies must land before it.
                 prologue_pack = (None if self._pre_replay_hooks else
                                  step_inputs.capture_prologue())
                 if prologue_pack is not None:
                     self._step_input_prologue_keys[key] = prologue_pack
                 output = forward_fn(capture_inputs)
+            if len(HOSTFUNC_USER_DATA_HANDLES) > num_hostfuncs:
+                self._hostfunc_graph_keys.add(key)
             # The addresses the graph baked in, for the strict buffer check.
             self._record_strict_buffer_snapshot(key)
             _restore_spec_decode_capture_state(attn_metadata,
@@ -1078,6 +1084,8 @@ def replay(
         finally:
             if prologue_pack is not None:
                 prologue_pack.on_graph_replay()
+        if key in self._hostfunc_graph_keys:
+            record_hostfunc_graph_replay()
         output_ref = self.graph_outputs[key]
 
         return output_ref
@@ -1596,6 +1604,7 @@ def clear(self, *, release_nccl_window_owners: bool = True):
         self.graphs.clear()
         self.graph_outputs.clear()
         self.graph_metadata.clear()
+        self._hostfunc_graph_keys.clear()
         self._pre_replay_position_ids = None
         self._step_input_prologue_keys.clear()
         self.reset_padding_dummies()

@@ -15,6 +15,7 @@
 import torch
 from torch import nn
 
+from tensorrt_llm._torch.hostfunc import drain_captured_hostfuncs
 from tensorrt_llm._torch.models.modeling_multimodal_encoder import MultimodalEncoderMixin
 from tensorrt_llm._torch.models.modeling_multimodal_mixin import (
     MultimodalEncoderContractError,
@@ -582,6 +583,10 @@ def requests_using_cache_keys(cache_keys: set[Hashable]) -> set[int]:
                 target_keywords=getattr(self.model, "multimodal_data_device_paths", None),
             )
 
+        # The encoder can hold the GIL while it waits on the device, so the host functions of the
+        # previous graph replay must run first.
+        drain_captured_hostfuncs()
+
         # The vision tower runs outside ModelCaller, so its TP all-reduces can
         # return NCCL window tensors here. The cache stores clones, so every
         # window lease produced by this batch ends with the scope. Request-level

@@ -987,10 +987,13 @@ def prepare_inputs(
 
         # Must be before the update of py_batch_idx
         if self.guided_decoder is not None:
+            # Prefill CUDA graphs can also capture guided decoding.
             self.guided_decoder.add_batch(
                 scheduled_requests,
                 new_tokens=new_tokens_device,
                 runtime_draft_len=runtime_draft_len,
+                is_cuda_graph=maybe_graph
+                or self._config.prefill_cuda_graph_backend != PrefillCudaGraphBackend.DISABLED,
             )
 
         spec_config = self._config.spec_config if enable_spec_decode else None

@@ -18,6 +18,8 @@
 )
 from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata
 from tensorrt_llm._torch.distributed import Distributed
+from tensorrt_llm._torch.hostfunc import drain_captured_hostfuncs
+from tensorrt_llm._torch.models.modeling_multimodal_utils import needs_multimodal_encoder
 from tensorrt_llm._torch.models.modeling_utils import DecoderModelForCausalLM
 from tensorrt_llm._torch.modules.context_logits import context_logits_heads, context_logits_scope
 from tensorrt_llm._torch.moe.fused_moe.moe_load_balancer import (
@@ -829,6 +831,10 @@ def postprocess_inputs(model_inputs: dict[str, Any]) -> None:
 
             graph_key = key if can_run_graph else None
             captures_model_body = self._graphs.captures_model_body(graph_key)
+            # The encoder can hold the GIL while it waits on the device, so the host functions of
+            # the previous graph replay must run first.
+            if needs_multimodal_encoder(inputs.get("multimodal_params")):
+                drain_captured_hostfuncs()
             with (
                 self._maybe_bypass_torch_compile(bypass=context_plan is not None),
                 context_logits_scope(context_plan),

@@ -14,7 +14,7 @@
 from ...bindings.executor import GuidedDecodingParams
 from ...bindings.internal.batch_manager import LlmRequestType
 from ...logger import logger
-from ..hostfunc import hostfunc
+from ..hostfunc import drain_captured_hostfuncs, hostfunc
 from .grammar_matcher import (GrammarMatcher, LLGuidanceMatcherFactory,
                               XGrammarMatcherFactory)
 from .llm_request import LlmRequest
@@ -545,6 +545,8 @@ def __init__(self,
         self._num_added_batches: int = 0
         self._failed_requests_queue: Queue = Queue()
         self._pending_failed_requests: List[Tuple[int, int, str]] = []
+        # Whether the current forward runs guided decoding; set by add_batch.
+        self.active = False
 
         self.new_tokens = torch.empty(self.max_num_draft_tokens + 1,
                                       self.max_num_sequences,
@@ -595,17 +597,41 @@ def reset_warmup_state(self) -> None:
         self.requests = None
         self.requests_hostfunc = None
 
+    def _skip_inactive(self) -> bool:
+        """Return True if the current forward has no guided work.
+
+        An inactive forward still waits for the host functions of a previously
+        replayed graph, so the caller does not enter a GIL-holding call while
+        they are pending. Capturing an inactive forward raises an error, since
+        the graph would replay without guided decoding.
+        """
+        if self.active:
+            return False
+        if torch.cuda.is_current_stream_capturing():
+            raise RuntimeError(
+                "Cannot capture a CUDA graph while guided decoding is "
+                "inactive; pass is_cuda_graph=True to add_batch.")
+        drain_captured_hostfuncs()
+        return True
+
     @nvtx_range("GuidedDecoder.add_batch")
     def add_batch(self,
                   scheduled_requests: ScheduledRequests,
                   new_tokens: Optional[torch.Tensor] = None,
-                  runtime_draft_len: Optional[int] = None) -> None:
+                  runtime_draft_len: Optional[int] = None,
+                  is_cuda_graph: bool = False) -> None:
         # See GuidedDecoder.add_batch: the layout must follow the runtime draft
         # length so the captured graph's bitmask matches the target logits.
         num_draft_tokens = (self.max_num_draft_tokens
                             if runtime_draft_len is None else runtime_draft_len)
         self.requests = GuidedRequests.from_scheduled_requests(
             scheduled_requests, num_draft_tokens)
+        # A CUDA graph replays its captured guided work for every batch, so
+        # graph forwards stay active. An eager forward skips guided decoding
+        # when no request is constrained.
+        self.active = is_cuda_graph or any(self.requests.valid_requests())
+        if not self.active:
+            return
         if new_tokens is not None:
             self.new_tokens.copy_(new_tokens.squeeze(-1), non_blocking=True)
         self._hand_over_batch(has_new_tokens=new_tokens is not None)
@@ -681,6 +707,8 @@ def execute(self,
         build runs as a host function, and the executor collects its
         failures with take_failed_requests.
         """
+        if self._skip_inactive():
+            return
         with torch.cuda.stream(self.stream):
             torch.cuda.current_stream().wait_event(self.token_event)
             self.fetch_batch()
@@ -711,6 +739,8 @@ def add_draft_batch(self,
                         new_tokens: torch.Tensor,
                         num_accepted_tokens: torch.Tensor,
                         draft_step: int = 0) -> None:
+        if self._skip_inactive():
+            return
         batch_size = len(self.requests)
         assert new_tokens.size(0) == batch_size
         self.new_tokens[0, :batch_size].copy_(new_tokens, non_blocking=True)
@@ -789,6 +819,8 @@ def rollback_rejected_batch(self,
             num_accepted_tokens: Int32 tensor of shape [batch_size], including
                 the bonus token in each verification count.
         """
+        if self._skip_inactive():
+            return
         self.add_accepted_batch(num_accepted_tokens)
         with torch.cuda.stream(self.stream):
             torch.cuda.current_stream().wait_event(self.token_event)
@@ -804,6 +836,8 @@ def execute_draft_batch(self,
                             d2t: Optional[torch.Tensor] = None,
                             draft_step: int = 0) -> None:
         """Enqueue the guided decoding of one draft step; see execute."""
+        if self._skip_inactive():
+            return
         with torch.cuda.stream(self.stream):
             torch.cuda.current_stream().wait_event(self.token_event)
             self.fetch_draft_batch(draft_step=draft_step)

@@ -358,6 +358,7 @@ def _run_capturable_step(
 ) -> torch.Tensor:
     """Hand `requests` to the host functions like add_batch does and run one target step."""
     decoder.requests = requests
+    decoder.active = True
     decoder._hand_over_batch(has_new_tokens=False)
     logits = torch.zeros(requests.num_bitmask_tokens, _VOCAB_SIZE, device="cuda")
     decoder.execute(logits)

@@ -257,6 +257,28 @@ def test_cross_iter_prefetch_max_ahead_config_override(monkeypatch):
     assert [req.py_request_id for req, _, _ in dispatches[0]] == [1]
 
 
+@pytest.mark.parametrize("pending", [False, True])
+def test_cross_iter_prefetch_skips_while_hostfuncs_pending(monkeypatch, pending):
+    """Prefetch is skipped while host functions of a replayed graph are pending."""
+    monkeypatch.setattr(mm_mixin, "_get_mm_aux_stream", lambda *args, **kwargs: object())
+    monkeypatch.setattr(mm_mixin, "hostfuncs_pending", lambda: pending)
+
+    dispatches = []
+    monkeypatch.setattr(
+        mm_mixin,
+        "_dispatch_cross_iter_prefetch",
+        lambda model, candidates, aux_stream: dispatches.append(candidates),
+    )
+
+    model = _StubModel(hidden_size=8, tokens_per_image=4)
+    req = _make_request(request_id=0, num_tokens=4)
+
+    n = maybe_prefetch_mm_encoder_for_next_iter(model, [req], max_prefetch_ahead=1)
+
+    assert n == int(not pending)
+    assert len(dispatches) == int(not pending)
+
+
 def test_cross_iter_prefetch_max_ahead_counts_pending_event(monkeypatch):
     monkeypatch.setattr(mm_mixin, "_get_mm_aux_stream", lambda *args, **kwargs: object())
 

@@ -286,7 +286,7 @@ def test_prepare_inputs_returns_steady_gen_fast_inputs(self) -> None:
         self.assertIs(result.kwargs, fast_inputs.kwargs)
         self.assertEqual(result.runtime_draft_len, 3)
         inputs.guided_decoder.add_batch.assert_called_once_with(
-            batch, new_tokens=new_tokens, runtime_draft_len=3
+            batch, new_tokens=new_tokens, runtime_draft_len=3, is_cuda_graph=False
         )
         inputs._can_use_steady_gen_fast_prepare.assert_called_once_with(
             batch, new_tokens, None, None, True

@@ -541,6 +541,40 @@ def forward_multimodal_encoder_items(
         mm_item_scheduler.forward_items([request], {request.request_id: [0]})
 
 
+def test_item_encoder_drains_replayed_host_functions_first() -> None:
+    """The item path drains replayed host functions before it runs the encoder."""
+    calls = []
+
+    class _Model(MultimodalModelMixin):
+        def prepare_multimodal_encoder_inputs(
+            self, _: Sequence[tuple[MultimodalParams, int]]
+        ) -> list[tuple[SimpleNamespace, list[int], str]]:
+            encoder_input = SimpleNamespace(to_device=lambda *_args, **_kwargs: None)
+            return [(encoder_input, [1], "image")]
+
+        def forward_multimodal_encoder_items(
+            self, _: Sequence[tuple[MultimodalParams, Sequence[int], str]]
+        ) -> list[torch.Tensor]:
+            calls.append("encode")
+            return []
+
+    mm_item_scheduler = bare_mm_item_scheduler(_Model())
+    request = make_mm_request(1, [4])
+    _bind_items(mm_item_scheduler, request)
+
+    with (
+        patch(
+            "tensorrt_llm._torch.pyexecutor.engine.multimodal.drain_captured_hostfuncs",
+            side_effect=lambda: calls.append("drain"),
+        ),
+        # The empty output ends the forward right after the encoder call.
+        pytest.raises(MultimodalEncoderRequestError, match="one output per item"),
+    ):
+        mm_item_scheduler.forward_items([request], {request.request_id: [0]})
+
+    assert calls == ["drain", "encode"]
+
+
 @pytest.mark.parametrize("failure_stage", ["prepare", "forward"])
 def test_item_encoder_translates_model_contract_errors(failure_stage: str) -> None:
     class _Model(MultimodalModelMixin):

@@ -666,6 +666,7 @@ def test_replay_defers_copies_only_to_its_captured_pack(monkeypatch, same_pack,
     runner._get_num_tokens_for_key = lambda key: 1
     runner._step_input_prologue_keys = {key: captured}
     runner._pre_replay_hooks = []
+    runner._hostfunc_graph_keys = set()
 
     def replay():
         queued.append("replay")

@@ -983,6 +983,7 @@ def test_replay_with_a_different_pack_issues_its_copies_eagerly(early):
     runner._get_num_tokens_for_key = lambda key: 16
     runner._step_input_prologue_keys = {key: captured}
     runner._pre_replay_hooks = []
+    runner._hostfunc_graph_keys = set()
     runner.graphs = {key: graph}
     runner.graph_outputs = {key: out}
     inputs = {