[https://nvbugs/6482606][fix] Enforce output length at disaggregated generation first-token handoff

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -85,7 +85,7 @@
 from .resource_manager import (NoFreeSlotsError, ResourceManager,
                                ResourceManagerType, request_context)
 from .sampler import (AsyncWorkerMixin, Sampler, SamplerEvent, SampleState,
-                      SampleStateTensors, TRTLLMSampler)
+                      SampleStateTensors, TorchSampler, TRTLLMSampler)
 from .scheduler import (RequestScheduler, ScheduledRequests,
                         SerializableSchedulerOutput, WaitingQueue,
                         create_waiting_queue)
@@ -5875,8 +5875,57 @@ def _prepare_disagg_gen_transmission_complete(self, scheduled_batch):
                 for beam in range(0, beam_width):
                     req.add_new_token(first_gen_tokens[beam], beam)
 
+                # The context-side first token is a real generated token and can
+                # already satisfy a stop criterion (e.g. max_tokens=1, or the token
+                # is the end id / completes a stop word). In aggregated serving the
+                # sampler applies this exact check to the context step's first token
+                # and finishes the request before any generation step is scheduled.
+                # The disagg gen path injects the first token without that check, so a
+                # generation (MTP) step still runs and emits an extra token, and the
+                # request returns OSL+1 tokens (nvbugs/6482606). Apply the same stop
+                # criteria here, before this iteration's forward, so no extra token is
+                # produced (the sampler skips GENERATION_COMPLETE requests).
+                self._apply_disagg_gen_first_token_stop_criteria(
+                    req, first_gen_tokens, beam_width)
+
                 self._maybe_prepend_logprobs_and_logits(req, beam_width)
 
+    def _apply_disagg_gen_first_token_stop_criteria(self, req, first_gen_tokens,
+                                                    beam_width: int) -> None:
+        """Enforce stop criteria on the context-side first token of a disagg
+        generation request, mirroring the aggregated context-step behavior.
+
+        Marks the request GENERATION_COMPLETE once all beams have met a stop
+        criterion, so the immediately following generation step does not emit an
+        extra token. See nvbugs/6482606. Reuses the same primitives the aggregated
+        sampler applies in ``TorchSampler._handle_stop_criteria`` (END_ID, then
+        max-token / OSL, then stop-words), in the same order.
+        """
+        # max_seq_len is Optional on PyExecutor; when unset there is no seq-len cap
+        # to apply, only the per-request max_new_tokens budget.
+        max_seq_len = self.max_seq_len if self.max_seq_len is not None else (
+            1 << 62)
+        beam_finish_reasons: List[Optional[FinishReason]] = []
+        for beam in range(0, beam_width):
+            token = first_gen_tokens[beam]
+            if token == req.py_end_id:
+                reason = FinishReason.END_ID
+            elif TorchSampler._meet_max_token_stop_criteria(
+                    req, max_seq_len, beam):
+                reason = FinishReason.LENGTH
+            elif TorchSampler._meet_stop_token_criteria(req, token, beam):
+                reason = FinishReason.STOP_WORDS
+            else:
+                reason = None
+            beam_finish_reasons.append(reason)
+
+        # Complete the request only when every beam has stopped, matching the agg
+        # semantics in TorchSampler._handle_finish_reasons_impl.
+        if all(reason is not None for reason in beam_finish_reasons):
+            req.state = LlmRequestState.GENERATION_COMPLETE
+            for beam, reason in enumerate(beam_finish_reasons):
+                req.set_finished_reason(reason, beam)
+
     def _update_sampler_state_for_disagg_gen_request(self, req, beam_width,
                                                      first_gen_tokens) -> bool:
         """Update beam sampler state with context-side first-token data."""
@@ -6676,6 +6725,14 @@ def _handle_first_token_response(self, scheduled_batch):
         new_responses = []
         for req in scheduled_batch.generation_requests:
             if req.py_decoding_iter == 1:
+                # A disagg gen request whose context-side first token already met a
+                # stop criterion (e.g. max_tokens=1) is finished at injection time.
+                # Its single, final response is emitted by _handle_responses; emitting
+                # a first-token response here would produce a second final response and
+                # double-count the token (nvbugs/6482606). Mirror the is_finished skip
+                # in _emit_first_token_responses.
+                if req.is_finished:
+                    continue
                 logger.debug(
                     f'Send first token response for request {req.py_request_id}'
                 )

tests/unittest/_torch/executor/test_disagg_gen_osl_stop.py

@@ -0,0 +1,186 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Unit tests for OSL enforcement at the disaggregated context->gen handoff.
+
+Regression coverage for nvbugs/6482606: under disagg + MTP/spec-decode, a request
+with ``max_tokens=1`` received 2 output tokens because the context-side first token
+was injected on the gen server without any stop-criteria check, so an extra
+generation (MTP) step ran and emitted a second token before the sampler's stop
+check fired. These tests exercise the two fix surfaces directly, without a GPU:
+
+  * ``PyExecutor._apply_disagg_gen_first_token_stop_criteria`` — applies the same
+    stop criteria (END_ID / max-token / stop-words) the aggregated sampler applies
+    to the context step's first token, finishing the request before the gen step.
+  * ``PyExecutor._handle_first_token_response`` — must skip an already-finished
+    request so it does not emit a (final) first-token response in addition to the
+    final response from ``_handle_responses`` (which would double-count the token).
+"""
+
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState
+from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor
+from tensorrt_llm.bindings.executor import FinishReason
+
+
+def _make_executor(max_seq_len=8192):
+    executor = object.__new__(PyExecutor)
+    executor.max_seq_len = max_seq_len
+    return executor
+
+
+def _make_disagg_gen_req(
+    *, num_tokens, orig_prompt_len, max_new_tokens, end_id=-1, beam_width=1, stop_words=None
+):
+    """Stub of a disagg gen request *after* the context first token was injected.
+
+    ``num_tokens`` is the post-injection token count (prompt + generated). The same
+    value is returned for every beam, which matches the state right after the
+    per-beam ``add_new_token`` loop in ``_prepare_disagg_gen_transmission_complete``.
+    """
+    req = SimpleNamespace()
+    req.py_orig_prompt_len = orig_prompt_len
+    req.py_max_new_tokens = max_new_tokens
+    req.py_end_id = end_id
+    req.py_stop_words_list = stop_words
+    req.py_beam_width = beam_width
+    req.state = LlmRequestState.GENERATION_IN_PROGRESS
+    req.get_num_tokens = lambda beam_idx=0: num_tokens
+    req.set_finished_reason = Mock()
+    return req
+
+
+def test_first_token_completes_request_when_max_tokens_is_one():
+    # The core nvbugs/6482606 case: with max_tokens=1 the injected context token
+    # already exhausts the budget, so the request must finish here (LENGTH) before
+    # the immediately-following MTP gen step can emit a second token.
+    executor = _make_executor()
+    req = _make_disagg_gen_req(num_tokens=11, orig_prompt_len=10, max_new_tokens=1)
+
+    PyExecutor._apply_disagg_gen_first_token_stop_criteria(executor, req, [42], 1)
+
+    assert req.state == LlmRequestState.GENERATION_COMPLETE
+    req.set_finished_reason.assert_called_once_with(FinishReason.LENGTH, 0)
+
+
+def test_first_token_does_not_complete_when_budget_remains():
+    # max_tokens=2: one more token is allowed, so the request must NOT be finished
+    # by the first token and should proceed to a normal generation step.
+    executor = _make_executor()
+    req = _make_disagg_gen_req(num_tokens=11, orig_prompt_len=10, max_new_tokens=2)
+
+    PyExecutor._apply_disagg_gen_first_token_stop_criteria(executor, req, [42], 1)
+
+    assert req.state == LlmRequestState.GENERATION_IN_PROGRESS
+    req.set_finished_reason.assert_not_called()
+
+
+def test_first_token_end_id_completes_request():
+    # An end-id first token finishes the request even when the token budget is not
+    # exhausted, matching TorchSampler._handle_stop_criteria's END_ID branch.
+    executor = _make_executor()
+    req = _make_disagg_gen_req(num_tokens=11, orig_prompt_len=10, max_new_tokens=16, end_id=42)
+
+    PyExecutor._apply_disagg_gen_first_token_stop_criteria(executor, req, [42], 1)
+
+    assert req.state == LlmRequestState.GENERATION_COMPLETE
+    req.set_finished_reason.assert_called_once_with(FinishReason.END_ID, 0)
+
+
+def test_first_token_stop_word_completes_request():
+    # A single-token stop word completed by the first token finishes the request.
+    # This is the parity case that a bare max-token/END_ID check would miss.
+    executor = _make_executor()
+    req = _make_disagg_gen_req(
+        num_tokens=11, orig_prompt_len=10, max_new_tokens=16, stop_words=([99], [1])
+    )
+
+    PyExecutor._apply_disagg_gen_first_token_stop_criteria(executor, req, [99], 1)
+
+    assert req.state == LlmRequestState.GENERATION_COMPLETE
+    req.set_finished_reason.assert_called_once_with(FinishReason.STOP_WORDS, 0)
+
+
+def test_none_max_seq_len_still_enforces_output_length():
+    # max_seq_len is Optional on PyExecutor; when unset, the per-request
+    # max_new_tokens budget must still be enforced.
+    executor = _make_executor(max_seq_len=None)
+    req = _make_disagg_gen_req(num_tokens=11, orig_prompt_len=10, max_new_tokens=1)
+
+    PyExecutor._apply_disagg_gen_first_token_stop_criteria(executor, req, [42], 1)
+
+    assert req.state == LlmRequestState.GENERATION_COMPLETE
+    req.set_finished_reason.assert_called_once_with(FinishReason.LENGTH, 0)
+
+
+def test_all_beams_complete_together_for_length():
+    # With max_tokens=1 every beam is exhausted by its injected token, so the whole
+    # request completes and each beam gets a finish reason.
+    executor = _make_executor()
+    req = _make_disagg_gen_req(num_tokens=11, orig_prompt_len=10, max_new_tokens=1, beam_width=2)
+
+    PyExecutor._apply_disagg_gen_first_token_stop_criteria(executor, req, [42, 43], 2)
+
+    assert req.state == LlmRequestState.GENERATION_COMPLETE
+    assert req.set_finished_reason.call_count == 2
+
+
+def test_request_incomplete_until_all_beams_stop():
+    # Beam 0's first token is the end id but beam 1's is not and the budget is not
+    # exhausted, so the request must remain in progress (agg completes a request
+    # only when every beam has stopped).
+    executor = _make_executor()
+    req = _make_disagg_gen_req(
+        num_tokens=11, orig_prompt_len=10, max_new_tokens=16, end_id=42, beam_width=2
+    )
+
+    PyExecutor._apply_disagg_gen_first_token_stop_criteria(executor, req, [42, 43], 2)
+
+    assert req.state == LlmRequestState.GENERATION_IN_PROGRESS
+    req.set_finished_reason.assert_not_called()
+
+
+def _make_first_token_req(*, is_finished, request_id):
+    req = SimpleNamespace()
+    req.py_decoding_iter = 1
+    req.is_finished = is_finished
+    req.py_request_id = request_id
+    req.create_response = Mock(return_value=SimpleNamespace(result=SimpleNamespace()))
+    return req
+
+
+def test_handle_first_token_response_skips_finished_request():
+    # A request finished at injection must not emit a first-token response here:
+    # _handle_responses emits its single final response, and emitting one here too
+    # would produce a second final response and double-count the token
+    # (nvbugs/6482606). Only the still-in-progress request gets a first-token
+    # response.
+    executor = object.__new__(PyExecutor)
+    executor._enqueue_responses = Mock()
+    executor._has_prepended_logits = Mock(return_value=False)
+    executor.dist = SimpleNamespace(rank=0)
+
+    finished = _make_first_token_req(is_finished=True, request_id=1)
+    unfinished = _make_first_token_req(is_finished=False, request_id=2)
+    scheduled = SimpleNamespace(generation_requests=[finished, unfinished])
+
+    PyExecutor._handle_first_token_response(executor, scheduled)
+
+    finished.create_response.assert_not_called()
+    unfinished.create_response.assert_called_once()
+    executor._enqueue_responses.assert_called_once()
+    enqueued = executor._enqueue_responses.call_args.args[0]
+    assert [request_id for request_id, _ in enqueued] == [2]