[#21460][perf] Skip no-op guided-decoding work in eager steps without guided requests

tensorrt_llm/_torch/pyexecutor/guided_decoder.py

@@ -538,6 +538,11 @@ def __init__(self,
         # self.requests_hostfunc should be accessed by hostfunc (CUDA callback).
         self.requests_hostfunc: Optional[GuidedRequests] = None
         self.queue = Queue()
+        # State for skipping no-op eager steps, see _skip_unguided_step.
+        self._no_guided_requests = False
+        self._step_skip: Optional[bool] = None
+        self._cur_step_eager = False
+        self._prev_step_eager = False
         # The batches handed over through self.queue are numbered on the
         # host thread (GuidedRequests.batch_id), and the build host function
         # tags the failures it reports with that number; see
@@ -595,6 +600,12 @@ def add_batch(self,
         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)
+        self._prev_step_eager, self._cur_step_eager = (self._cur_step_eager,
+                                                       False)
+        self._no_guided_requests = not any(
+            req.guided_decoding_params is not None
+            for req in self.requests.requests)
+        self._step_skip = None
         # self.token_event.record() should be called inside CUDA graph capturing;
         # currently, it is in DecoderInputPreparer.preprocess_inputs.
 
@@ -604,6 +615,38 @@ def _hand_over_batch(self, has_new_tokens: bool) -> None:
         self.requests.batch_id = self._num_added_batches
         self.queue.put((self.requests, has_new_tokens))
 
+    def _skip_unguided_step(self) -> bool:
+        """Whether this eager step's guided-decoding calls are no-ops.
+
+        A batch without guided requests needs no grammar work: ``_build`` sets
+        ``token_mask`` only for guided rows, so the bitmask kernel leaves the
+        logits unchanged, and the host functions only touch guided requests.
+        Skipping also drops the ``_drain_host_functions`` waits, which empty
+        the GPU queue in every eager step.
+
+        Skipping is allowed only when the previous step's calls also ran
+        eagerly, so no host node of an earlier CUDA-graph step can still be
+        pending while this step continues without draining. The first eager
+        step after graph replays keeps the normal path.
+
+        The first eager call of a step decides for the whole step. A skipped
+        step consumes its own queue entry here (the newest one; older entries
+        belong to CUDA-graph host nodes that have not run yet and pop from the
+        front) instead of in ``fetch_batch``, so the queue stays aligned with
+        the steps.
+        """
+        if torch.cuda.is_current_stream_capturing():
+            return False
+        self._cur_step_eager = True
+        if self._step_skip is None:
+            self._step_skip = (self._no_guided_requests
+                               and self._prev_step_eager)
+            if self._step_skip:
+                with self.queue.mutex:
+                    if self.queue.queue:
+                        self.queue.queue.pop()
+        return self._step_skip
+
     @hostfunc
     def fetch_batch(self) -> None:
         # CUDA graph warmup calls model forward for multiple times for one prepared inputs
@@ -667,6 +710,8 @@ def execute(self,
         build runs as a host function, and the executor collects its
         failures with take_failed_requests.
         """
+        if self._skip_unguided_step():
+            return
         with torch.cuda.stream(self.stream):
             torch.cuda.current_stream().wait_event(self.token_event)
             self.fetch_batch()
@@ -697,6 +742,8 @@ def add_draft_batch(self,
                         new_tokens: torch.Tensor,
                         num_accepted_tokens: torch.Tensor,
                         draft_step: int = 0) -> None:
+        if self._skip_unguided_step():
+            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)
@@ -739,6 +786,8 @@ def add_accepted_batch(self, num_accepted_tokens: torch.Tensor) -> None:
         Args:
             num_accepted_tokens: Int32 tensor of shape [batch_size].
         """
+        if self._skip_unguided_step():
+            return
         batch_size = len(self.requests)
         assert num_accepted_tokens.size(0) == batch_size
         self.num_accepted_tokens[:batch_size].copy_(num_accepted_tokens,
@@ -775,6 +824,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_unguided_step():
+            return
         self.add_accepted_batch(num_accepted_tokens)
         with torch.cuda.stream(self.stream):
             torch.cuda.current_stream().wait_event(self.token_event)
@@ -790,6 +841,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_unguided_step():
+            return
         with torch.cuda.stream(self.stream):
             torch.cuda.current_stream().wait_event(self.token_event)
             self.fetch_draft_batch(draft_step=draft_step)

tests/unittest/_torch/speculative/hw_agnostic/test_guided_hostfunc_drain.py

@@ -30,12 +30,18 @@
 worker-side guided-decoding path goes through the drained methods.
 """
 
+from queue import Queue
+from types import SimpleNamespace
 from unittest.mock import MagicMock, patch
 
 import pytest
 import torch
 
-from tensorrt_llm._torch.pyexecutor.guided_decoder import CapturableGuidedDecoder, GuidedDecoder
+from tensorrt_llm._torch.pyexecutor.guided_decoder import (
+    CapturableGuidedDecoder,
+    GuidedDecoder,
+    GuidedRequests,
+)
 from tensorrt_llm._torch.speculative.interface import SpecWorkerBase
 
 
@@ -73,6 +79,14 @@ def _decoder(cls):
     decoder.bitmask_event = MagicMock(name="bitmask_event")
     decoder.requests = []
     decoder.max_num_draft_tokens = 1
+    decoder.queue = Queue()
+    # Skip state as __init__ / add_batch leave it: a batch with a guided
+    # request, so the drain tests exercise the normal path.
+    decoder._no_guided_requests = False
+    decoder._step_skip = None
+    decoder._cur_step_eager = False
+    decoder._prev_step_eager = False
+    decoder._num_added_batches = 0
     for name in (
         "fetch_batch",
         "init_disagg_gen_requests",
@@ -150,6 +164,195 @@ def test_no_drain_for_the_non_capturable_decoder():
     decoder.bitmask_event.synchronize.assert_not_called()
 
 
+def _unguided_eager_step(decoder, *, prev_step_eager=True):
+    """State after ``add_batch`` for a batch without guided requests."""
+    decoder._no_guided_requests = True
+    decoder._step_skip = None
+    decoder._cur_step_eager = False
+    decoder._prev_step_eager = prev_step_eager
+
+
+@pytest.mark.parametrize("method", _ENQUEUE_METHODS)
+def test_unguided_eager_step_after_eager_step_skips(method):
+    """No guided request and nothing pending from a graph step: no grammar work, no drain."""
+    decoder = _decoder(CapturableGuidedDecoder)
+    _unguided_eager_step(decoder)
+    _run(_call(decoder, method), capturing=False)
+
+    decoder.bitmask_event.synchronize.assert_not_called()
+    decoder.build.assert_not_called()
+    decoder.apply_bitmask.assert_not_called()
+    decoder.add_accepted_batch.assert_not_called()
+    decoder.fetch_accepted_batch.assert_not_called()
+    decoder.rollback_rejected_tokens.assert_not_called()
+
+
+@pytest.mark.parametrize("method", _ENQUEUE_METHODS)
+def test_first_eager_step_after_graph_step_still_drains(method):
+    """Host nodes of the previous (captured) step may still be pending, so keep the drain."""
+    decoder = _decoder(CapturableGuidedDecoder)
+    _unguided_eager_step(decoder, prev_step_eager=False)
+    _run(_call(decoder, method), capturing=False)
+
+    decoder.bitmask_event.synchronize.assert_called_once_with()
+
+
+_CAPTURED_CALLS = {
+    "execute": (
+        "fetch_batch",
+        "init_disagg_gen_requests",
+        "build",
+        "copy_bitmask",
+        "apply_bitmask",
+    ),
+    "execute_draft_batch": (
+        "fetch_draft_batch",
+        "rollback_rejected_tokens",
+        "build",
+        "rollback_draft_tokens",
+        "copy_bitmask",
+        "apply_bitmask",
+    ),
+    "rollback_rejected_batch": (
+        "add_accepted_batch",
+        "fetch_accepted_batch",
+        "rollback_rejected_tokens",
+    ),
+}
+
+
+@pytest.mark.parametrize("method", _ENQUEUE_METHODS)
+def test_no_skip_while_capturing(method):
+    """A captured step never skips: its host nodes run at replay and need their queue entry."""
+    decoder = _decoder(CapturableGuidedDecoder)
+    decoder.queue.put("captured-step")
+    _unguided_eager_step(decoder)
+    _run(_call(decoder, method), capturing=True)
+
+    assert decoder._step_skip is None
+    assert decoder._cur_step_eager is False
+    assert list(decoder.queue.queue) == ["captured-step"]
+    for name in _CAPTURED_CALLS[method]:
+        getattr(decoder, name).assert_called_once()
+    decoder.bitmask_event.record.assert_called_once_with()
+
+
+def _add_batch(decoder, *, guided):
+    """Submit one batch through the real ``add_batch``; only the request snapshot is stubbed."""
+    params = object() if guided else None
+    requests = SimpleNamespace(requests=[SimpleNamespace(guided_decoding_params=params)])
+    with patch.object(GuidedRequests, "from_scheduled_requests", return_value=requests):
+        decoder.add_batch(MagicMock(name="scheduled_requests"))
+
+
+def _eager_execute(decoder):
+    _run(_call(decoder, "execute"), capturing=False)
+
+
+def _decoder_after_an_eager_step():
+    decoder = _decoder(CapturableGuidedDecoder)
+    _add_batch(decoder, guided=False)
+    _eager_execute(decoder)
+    # The first eager step follows a step that wasn't eager, so it keeps the normal path.
+    decoder.build.assert_called_once()
+    decoder.apply_bitmask.assert_called_once()
+    decoder.bitmask_event.synchronize.assert_called_once_with()
+    decoder.build.reset_mock()
+    decoder.apply_bitmask.reset_mock()
+    decoder.bitmask_event.synchronize.reset_mock()
+    return decoder
+
+
+def test_add_batch_unguided_batch_after_eager_step_skips():
+    decoder = _decoder_after_an_eager_step()
+    _add_batch(decoder, guided=False)
+    _eager_execute(decoder)
+
+    decoder.build.assert_not_called()
+    decoder.apply_bitmask.assert_not_called()
+    decoder.bitmask_event.synchronize.assert_not_called()
+
+
+def test_add_batch_guided_batch_after_eager_step_runs_the_grammar():
+    decoder = _decoder_after_an_eager_step()
+    _add_batch(decoder, guided=True)
+    _eager_execute(decoder)
+
+    decoder.build.assert_called_once()
+    decoder.apply_bitmask.assert_called_once()
+    decoder.bitmask_event.synchronize.assert_called_once_with()
+
+
+_BATCH_ADD_METHODS = {
+    "add_draft_batch": lambda decoder: CapturableGuidedDecoder.add_draft_batch(
+        decoder, torch.zeros(1), torch.ones(1), draft_step=0
+    ),
+    # The fixture mocks ``add_accepted_batch`` for the rollback tests, so call the real one.
+    "add_accepted_batch": lambda decoder: CapturableGuidedDecoder.add_accepted_batch(
+        decoder, torch.ones(1)
+    ),
+}
+
+
+def _batch_add_decoder():
+    decoder = _decoder(CapturableGuidedDecoder)
+    decoder.requests = [object()]
+    decoder.new_tokens = MagicMock(name="new_tokens")
+    decoder.num_accepted_tokens = MagicMock(name="num_accepted_tokens")
+    return decoder
+
+
+@pytest.mark.parametrize("method", _BATCH_ADD_METHODS)
+def test_batch_add_on_a_skipped_step_copies_nothing(method):
+    decoder = _batch_add_decoder()
+    _unguided_eager_step(decoder)
+    _run(lambda: _BATCH_ADD_METHODS[method](decoder), capturing=False)
+
+    decoder.new_tokens.__getitem__.assert_not_called()
+    decoder.num_accepted_tokens.__getitem__.assert_not_called()
+    decoder.token_event.record.assert_not_called()
+
+
+@pytest.mark.parametrize("method", _BATCH_ADD_METHODS)
+def test_batch_add_while_capturing_copies_and_records(method):
+    """A captured step never skips, even when its batch is unguided."""
+    decoder = _batch_add_decoder()
+    _unguided_eager_step(decoder)
+    _run(lambda: _BATCH_ADD_METHODS[method](decoder), capturing=True)
+
+    decoder.num_accepted_tokens.__getitem__.return_value.copy_.assert_called_once()
+    if method == "add_draft_batch":
+        decoder.new_tokens.__getitem__.return_value.copy_.assert_called_once()
+    decoder.token_event.record.assert_called_once_with()
+    assert decoder._step_skip is None
+
+
+@pytest.mark.parametrize("method", _BATCH_ADD_METHODS)
+def test_batch_add_on_a_guided_step_copies_and_records(method):
+    decoder = _batch_add_decoder()
+    _run(lambda: _BATCH_ADD_METHODS[method](decoder), capturing=False)
+
+    decoder.num_accepted_tokens.__getitem__.return_value.copy_.assert_called_once()
+    if method == "add_draft_batch":
+        decoder.new_tokens.__getitem__.return_value.copy_.assert_called_once()
+    decoder.token_event.record.assert_called_once_with()
+
+
+def test_skipped_step_consumes_only_its_own_queue_entry():
+    """A skipped step drops the newest entry; older ones belong to pending graph host nodes."""
+    decoder = _decoder(CapturableGuidedDecoder)
+    decoder.queue.put("graph-step")
+    decoder.queue.put("this-step")
+    _unguided_eager_step(decoder)
+
+    _run(lambda: decoder.execute(MagicMock(name="logits")), capturing=False)
+    _run(
+        lambda: decoder.execute_draft_batch(MagicMock(name="logits"), draft_step=0), capturing=False
+    )
+
+    assert list(decoder.queue.queue) == ["graph-step"]
+
+
 def test_worker_executes_the_guided_decoder():
     """The worker helper delegates to ``execute``, which drains internally."""
     decoder = MagicMock(name="guided_decoder")