From ee25fc40c8e47041298180ef5d68f28f725be6f2 Mon Sep 17 00:00:00 2001
From: Aswin Visva <31215515+aswinvisva@users.noreply.github.com>
Date: Fri, 6 Nov 2026 14:54:27 -0800
Subject: [PATCH 1/4] [None][fix] Drain captured guided-decoding host functions
 before multimodal encoders and skip guided work in unconstrained eager
 forwards

Signed-off-by: Aswin Visva <31215515+aswinvisva@users.noreply.github.com>
---
 tensorrt_llm/_torch/hostfunc.py               |  33 ++++++
 .../models/modeling_multimodal_mixin.py       |   5 +
 .../models/modeling_multimodal_utils.py       |   6 +
 .../_torch/pyexecutor/cuda_graph_runner.py    |   9 ++
 .../_torch/pyexecutor/engine/multimodal.py    |   5 +
 .../engine/runners/decoder/inputs.py          |   3 +
 .../engine/runners/decoder/runner.py          |   6 +
 .../_torch/pyexecutor/guided_decoder.py       |  38 +++++-
 .../misc/test_guided_decoder_bitmask.py       |   1 +
 .../test_mm_encoder_cross_iter_prefetch.py    |  22 ++++
 .../pyexecutor/engine/test_multimodal.py      |  34 ++++++
 .../hw_agnostic/test_guided_hostfunc_drain.py | 111 +++++++++++++++++-
 .../hw_agnostic/test_sa_guided_rollback.py    |   7 +-
 13 files changed, 274 insertions(+), 6 deletions(-)

diff --git a/tensorrt_llm/_torch/hostfunc.py b/tensorrt_llm/_torch/hostfunc.py
index 39174e3d3d28..d90cf20086f7 100644
--- a/tensorrt_llm/_torch/hostfunc.py
+++ b/tensorrt_llm/_torch/hostfunc.py
@@ -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,35 @@ 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 the last replay's host functions may still be pending."""
+    return _replay_event is not None and not _replay_event.query()
+
+
+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):
diff --git a/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py b/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py
index 242c6cedda90..f49d491472d0 100644
--- a/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py
+++ b/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py
@@ -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. While host functions of a replayed
+    # graph are pending, these requests are encoded when scheduled instead, after a drain.
+    if hostfuncs_pending():
+        return 0
 
     aux_stream = _get_mm_aux_stream(max_prefetch_ahead)
     if aux_stream is None:
diff --git a/tensorrt_llm/_torch/models/modeling_multimodal_utils.py b/tensorrt_llm/_torch/models/modeling_multimodal_utils.py
index fcd32b77aa4a..72f68ecb0a7b 100644
--- a/tensorrt_llm/_torch/models/modeling_multimodal_utils.py
+++ b/tensorrt_llm/_torch/models/modeling_multimodal_utils.py
@@ -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],
diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py
index 1a8c1c9a60ed..ca2967a55014 100644
--- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py
+++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py
@@ -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,6 +915,7 @@ 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
@@ -919,6 +923,8 @@ def run_warmup():
                 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()
diff --git a/tensorrt_llm/_torch/pyexecutor/engine/multimodal.py b/tensorrt_llm/_torch/pyexecutor/engine/multimodal.py
index 67a9b2d017e6..e2044b82054a 100644
--- a/tensorrt_llm/_torch/pyexecutor/engine/multimodal.py
+++ b/tensorrt_llm/_torch/pyexecutor/engine/multimodal.py
@@ -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
diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py
index df05cc46e71e..3757888ca34c 100644
--- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py
+++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py
@@ -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
diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py
index 7419e578f422..60cb6bf1a45e 100644
--- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py
+++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py
@@ -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),
diff --git a/tensorrt_llm/_torch/pyexecutor/guided_decoder.py b/tensorrt_llm/_torch/pyexecutor/guided_decoder.py
index a8f8fcfe4ae3..81f6049ace4b 100644
--- a/tensorrt_llm/_torch/pyexecutor/guided_decoder.py
+++ b/tensorrt_llm/_torch/pyexecutor/guided_decoder.py
@@ -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)
diff --git a/tests/unittest/_torch/misc/test_guided_decoder_bitmask.py b/tests/unittest/_torch/misc/test_guided_decoder_bitmask.py
index dee5a6c6ce04..8a579d584f97 100644
--- a/tests/unittest/_torch/misc/test_guided_decoder_bitmask.py
+++ b/tests/unittest/_torch/misc/test_guided_decoder_bitmask.py
@@ -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)
diff --git a/tests/unittest/_torch/multimodal/test_mm_encoder_cross_iter_prefetch.py b/tests/unittest/_torch/multimodal/test_mm_encoder_cross_iter_prefetch.py
index 987dee4ba9b4..961994fa5e14 100644
--- a/tests/unittest/_torch/multimodal/test_mm_encoder_cross_iter_prefetch.py
+++ b/tests/unittest/_torch/multimodal/test_mm_encoder_cross_iter_prefetch.py
@@ -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())
 
diff --git a/tests/unittest/_torch/pyexecutor/engine/test_multimodal.py b/tests/unittest/_torch/pyexecutor/engine/test_multimodal.py
index 5941744adb41..f7c580cec0b8 100644
--- a/tests/unittest/_torch/pyexecutor/engine/test_multimodal.py
+++ b/tests/unittest/_torch/pyexecutor/engine/test_multimodal.py
@@ -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):
diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_guided_hostfunc_drain.py b/tests/unittest/_torch/speculative/hw_agnostic/test_guided_hostfunc_drain.py
index 7437308c915a..e1f998680db7 100644
--- a/tests/unittest/_torch/speculative/hw_agnostic/test_guided_hostfunc_drain.py
+++ b/tests/unittest/_torch/speculative/hw_agnostic/test_guided_hostfunc_drain.py
@@ -28,15 +28,27 @@
 host nodes do not run until replay), and for the non-capturable decoder, which
 runs its build inline and enqueues no callback. They also pin that every
 worker-side guided-decoding path goes through the drained methods.
+
+Host functions of a replayed CUDA graph are drained before multimodal encoder
+work. An eager forward with no constrained request enqueues none, still drains
+a replayed graph, and cannot be captured.
 """
 
+from queue import Queue
 from unittest.mock import MagicMock, patch
 
 import pytest
 import torch
 
-from tensorrt_llm._torch.pyexecutor.guided_decoder import CapturableGuidedDecoder, GuidedDecoder
+from tensorrt_llm._torch import hostfunc
+from tensorrt_llm._torch.models import modeling_multimodal_utils
+from tensorrt_llm._torch.pyexecutor.guided_decoder import (
+    CapturableGuidedDecoder,
+    GuidedDecoder,
+    GuidedRequests,
+)
 from tensorrt_llm._torch.speculative.interface import SpecWorkerBase
+from tensorrt_llm.inputs.multimodal import MultimodalParams
 
 
 class _Worker(SpecWorkerBase):
@@ -72,6 +84,7 @@ def _decoder(cls):
     decoder.token_event = MagicMock(name="token_event")
     decoder.bitmask_event = MagicMock(name="bitmask_event")
     decoder.requests = []
+    decoder.active = True
     decoder.max_num_draft_tokens = 1
     for name in (
         "fetch_batch",
@@ -142,6 +155,102 @@ def test_no_drain_while_capturing_a_cuda_graph(method):
     decoder.bitmask_event.synchronize.assert_not_called()
 
 
+@pytest.mark.parametrize("method", _ENQUEUE_METHODS)
+def test_inactive_forward_enqueues_nothing(method):
+    """An inactive forward enqueues no host function but still drains a replayed graph."""
+    decoder = _decoder(CapturableGuidedDecoder)
+    decoder.active = False
+
+    with patch("tensorrt_llm._torch.pyexecutor.guided_decoder.drain_captured_hostfuncs") as drain:
+        result = _run(_call(decoder, method), capturing=False)
+
+    for name in ("fetch_batch", "fetch_draft_batch", "fetch_accepted_batch", "build"):
+        getattr(decoder, name).assert_not_called()
+    decoder.bitmask_event.record.assert_not_called()
+    decoder.bitmask_event.synchronize.assert_not_called()
+    drain.assert_called_once_with()
+    assert result in (None, [])
+
+
+@pytest.mark.parametrize("method", _ENQUEUE_METHODS)
+def test_capturing_an_inactive_forward_raises(method):
+    """A graph captured from an inactive forward would replay without guided decoding."""
+    decoder = _decoder(CapturableGuidedDecoder)
+    decoder.active = False
+
+    with pytest.raises(
+        RuntimeError, match="Cannot capture a CUDA graph while guided decoding is inactive"
+    ):
+        _run(_call(decoder, method), capturing=True)
+
+    decoder.fetch_batch.assert_not_called()
+
+
+@pytest.mark.parametrize(
+    "is_cuda_graph,constrained",
+    [(False, False), (False, True), (True, False), (True, True)],
+)
+def test_add_batch_activates_for_cuda_graphs_and_constrained_batches(is_cuda_graph, constrained):
+    """A graph replays the guided work it captured, whatever the batch holds."""
+    decoder = object.__new__(CapturableGuidedDecoder)
+    decoder.max_num_draft_tokens = 1
+    decoder.queue = Queue()
+    snapshot = MagicMock(name="guided_requests")
+    snapshot.valid_requests.side_effect = lambda: iter([object()] if constrained else [])
+
+    with patch.object(GuidedRequests, "from_scheduled_requests", return_value=snapshot):
+        decoder.add_batch(MagicMock(name="scheduled_requests"), is_cuda_graph=is_cuda_graph)
+
+    expected = is_cuda_graph or constrained
+    assert decoder.active == expected
+    assert decoder.queue.qsize() == int(expected)
+
+
+@pytest.mark.parametrize("capturing", [False, True])
+def test_drain_waits_for_the_last_replay_outside_capture(monkeypatch, capturing):
+    """The drain waits on the replay event, except during stream capture."""
+    event = MagicMock(name="replay_event")
+    monkeypatch.setattr(hostfunc, "_replay_event", event)
+
+    with patch("torch.cuda.is_current_stream_capturing", return_value=capturing):
+        hostfunc.drain_captured_hostfuncs()
+
+    assert event.synchronize.call_count == int(not capturing)
+    assert hostfunc._replay_event is (event if capturing else None)
+
+
+@pytest.mark.parametrize(
+    "replayed,completed,expected",
+    [(False, False, False), (True, False, True), (True, True, False)],
+)
+def test_hostfuncs_pending_until_the_replay_completes(monkeypatch, replayed, completed, expected):
+    event = MagicMock(name="replay_event")
+    event.query.return_value = completed
+    monkeypatch.setattr(hostfunc, "_replay_event", event if replayed else None)
+
+    assert hostfunc.hostfuncs_pending() is expected
+
+
+_IMAGE = {"image": {"pixel_values": torch.zeros(1)}}
+_EMBEDDING = {"multimodal_embedding": torch.zeros(1, 4)}
+
+
+@pytest.mark.parametrize(
+    "multimodal_data,expected",
+    [
+        ([], False),
+        ([{}], False),
+        ([_EMBEDDING], False),
+        ([_IMAGE], True),
+        ([_EMBEDDING, _IMAGE], True),
+    ],
+)
+def test_needs_multimodal_encoder(multimodal_data, expected):
+    """Only a param with content and no attached embedding needs the encoder."""
+    params = [MultimodalParams(multimodal_data=data) for data in multimodal_data]
+    assert modeling_multimodal_utils.needs_multimodal_encoder(params) is expected
+
+
 def test_no_drain_for_the_non_capturable_decoder():
     """The base decoder builds inline and enqueues no callback, so it owes no drain."""
     decoder = _decoder(GuidedDecoder)
diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_sa_guided_rollback.py b/tests/unittest/_torch/speculative/hw_agnostic/test_sa_guided_rollback.py
index 8706be6de981..028f23a33d8a 100644
--- a/tests/unittest/_torch/speculative/hw_agnostic/test_sa_guided_rollback.py
+++ b/tests/unittest/_torch/speculative/hw_agnostic/test_sa_guided_rollback.py
@@ -87,9 +87,9 @@ def _json_request(**overrides) -> GuidedRequest:
 def _target_step(decoder: CapturableGuidedDecoder, request: GuidedRequest, new_tokens=None):
     """One target verification for a hand-built request.
 
-    Mirrors ``CapturableGuidedDecoder.add_batch`` (requests snapshot, new-token
-    column, queue hand-off) and then runs ``execute`` on zero logits, so the
-    applied mask is readable as ``-inf`` entries.
+    Mirrors ``CapturableGuidedDecoder.add_batch`` (requests snapshot, active
+    flag, new-token column, queue hand-off) and then runs ``execute`` on zero
+    logits, so the applied mask is readable as ``-inf`` entries.
     """
     requests = GuidedRequests(
         [request],
@@ -98,6 +98,7 @@ def _target_step(decoder: CapturableGuidedDecoder, request: GuidedRequest, new_t
         max_num_draft_tokens=K,
     )
     decoder.requests = requests
+    decoder.active = True
     if new_tokens is not None:
         decoder.new_tokens[:, request.seq_slot].copy_(torch.tensor(new_tokens, dtype=torch.int32))
     decoder.queue.put((requests, new_tokens is not None))

From 12d26c84b128a4e5846b97bd8e012dc8f716d4f1 Mon Sep 17 00:00:00 2001
From: Aswin Visva <31215515+aswinvisva@users.noreply.github.com>
Date: Mon, 9 Nov 2026 00:36:13 -0800
Subject: [PATCH 2/4] [None][fix] Set the host-function graph keys on
 hand-built CUDA graph runners and expect is_cuda_graph in the add_batch
 assertion

Signed-off-by: Aswin Visva <31215515+aswinvisva@users.noreply.github.com>
---
 .../_torch/pyexecutor/engine/runners/decoder/test_inputs.py     | 2 +-
 tests/unittest/_torch/pyexecutor/test_cuda_graph_runner.py      | 1 +
 tests/unittest/_torch/pyexecutor/test_step_inputs.py            | 1 +
 3 files changed, 3 insertions(+), 1 deletion(-)

diff --git a/tests/unittest/_torch/pyexecutor/engine/runners/decoder/test_inputs.py b/tests/unittest/_torch/pyexecutor/engine/runners/decoder/test_inputs.py
index 5add16dd533a..4d737149d6b8 100644
--- a/tests/unittest/_torch/pyexecutor/engine/runners/decoder/test_inputs.py
+++ b/tests/unittest/_torch/pyexecutor/engine/runners/decoder/test_inputs.py
@@ -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
diff --git a/tests/unittest/_torch/pyexecutor/test_cuda_graph_runner.py b/tests/unittest/_torch/pyexecutor/test_cuda_graph_runner.py
index b593ab930904..ceb534767e7b 100644
--- a/tests/unittest/_torch/pyexecutor/test_cuda_graph_runner.py
+++ b/tests/unittest/_torch/pyexecutor/test_cuda_graph_runner.py
@@ -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")
diff --git a/tests/unittest/_torch/pyexecutor/test_step_inputs.py b/tests/unittest/_torch/pyexecutor/test_step_inputs.py
index a00652c45547..a38489334b9c 100644
--- a/tests/unittest/_torch/pyexecutor/test_step_inputs.py
+++ b/tests/unittest/_torch/pyexecutor/test_step_inputs.py
@@ -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 = {

From 775cb31ea1d780df0e40b7fec08259d3ffe16af9 Mon Sep 17 00:00:00 2001
From: Aswin Visva <31215515+aswinvisva@users.noreply.github.com>
Date: Mon, 9 Nov 2026 10:57:00 -0800
Subject: [PATCH 3/4] [None][fix] Base the multimodal prefetch skip on
 undrained graph replays rather than device progress so tensor-parallel ranks
 agree

Signed-off-by: Aswin Visva <31215515+aswinvisva@users.noreply.github.com>
---
 tensorrt_llm/_torch/hostfunc.py                           | 8 ++++++--
 tensorrt_llm/_torch/models/modeling_multimodal_mixin.py   | 4 ++--
 .../speculative/hw_agnostic/test_guided_hostfunc_drain.py | 6 ++++--
 3 files changed, 12 insertions(+), 6 deletions(-)

diff --git a/tensorrt_llm/_torch/hostfunc.py b/tensorrt_llm/_torch/hostfunc.py
index d90cf20086f7..700739ed6cb5 100644
--- a/tensorrt_llm/_torch/hostfunc.py
+++ b/tensorrt_llm/_torch/hostfunc.py
@@ -39,8 +39,12 @@ def record_hostfunc_graph_replay() -> None:
 
 
 def hostfuncs_pending() -> bool:
-    """Return True if the last replay's host functions may still be pending."""
-    return _replay_event is not None and not _replay_event.query()
+    """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:
diff --git a/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py b/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py
index f49d491472d0..48a9f43dc84d 100644
--- a/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py
+++ b/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py
@@ -1811,8 +1811,8 @@ 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. While host functions of a replayed
-    # graph are pending, these requests are encoded when scheduled instead, after a drain.
+    # 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
 
diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_guided_hostfunc_drain.py b/tests/unittest/_torch/speculative/hw_agnostic/test_guided_hostfunc_drain.py
index e1f998680db7..77ec0f1c2a74 100644
--- a/tests/unittest/_torch/speculative/hw_agnostic/test_guided_hostfunc_drain.py
+++ b/tests/unittest/_torch/speculative/hw_agnostic/test_guided_hostfunc_drain.py
@@ -221,14 +221,16 @@ def test_drain_waits_for_the_last_replay_outside_capture(monkeypatch, capturing)
 
 @pytest.mark.parametrize(
     "replayed,completed,expected",
-    [(False, False, False), (True, False, True), (True, True, False)],
+    [(False, False, False), (True, False, True), (True, True, True)],
 )
-def test_hostfuncs_pending_until_the_replay_completes(monkeypatch, replayed, completed, expected):
+def test_hostfuncs_pending_until_the_replay_is_drained(monkeypatch, replayed, completed, expected):
+    """Device progress is not consulted, so every rank answers the same way."""
     event = MagicMock(name="replay_event")
     event.query.return_value = completed
     monkeypatch.setattr(hostfunc, "_replay_event", event if replayed else None)
 
     assert hostfunc.hostfuncs_pending() is expected
+    event.query.assert_not_called()
 
 
 _IMAGE = {"image": {"pixel_values": torch.zeros(1)}}

From 09cdaf58a5f8d7507dbe2983d6a6e469fd5ad34d Mon Sep 17 00:00:00 2001
From: Aswin Visva <31215515+aswinvisva@users.noreply.github.com>
Date: Mon, 9 Nov 2026 14:13:28 -0800
Subject: [PATCH 4/4] [None][fix] Set the batch counter on the hand-built
 decoder in the add_batch activation test

Signed-off-by: Aswin Visva <31215515+aswinvisva@users.noreply.github.com>
---
 .../_torch/speculative/hw_agnostic/test_guided_hostfunc_drain.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_guided_hostfunc_drain.py b/tests/unittest/_torch/speculative/hw_agnostic/test_guided_hostfunc_drain.py
index 77ec0f1c2a74..a634f803a870 100644
--- a/tests/unittest/_torch/speculative/hw_agnostic/test_guided_hostfunc_drain.py
+++ b/tests/unittest/_torch/speculative/hw_agnostic/test_guided_hostfunc_drain.py
@@ -195,6 +195,7 @@ def test_add_batch_activates_for_cuda_graphs_and_constrained_batches(is_cuda_gra
     decoder = object.__new__(CapturableGuidedDecoder)
     decoder.max_num_draft_tokens = 1
     decoder.queue = Queue()
+    decoder._num_added_batches = 0
     snapshot = MagicMock(name="guided_requests")
     snapshot.valid_requests.side_effect = lambda: iter([object()] if constrained else [])