[https://nvbugs/6713424][fix] Flush the HF streaming detokenizer on the final response; fix the AgentX aiperf install check

tensorrt_llm/_torch/models/checkpoints/mistral/tokenizer.py

@@ -293,6 +293,7 @@ def hf_decode_incrementally(
         prev_text: str | None = None,
         states: dict | None = None,
         *,
+        flush: bool = False,
         skip_special_tokens: bool = False,
         clean_up_tokenization_spaces: bool | None = None,
     ) -> tuple[str, dict]:

tensorrt_llm/tokenizer/tokenizer.py

@@ -469,6 +469,7 @@ def decode_incrementally(
                 token_ids,
                 prev_text,
                 states,
+                flush=flush,
                 skip_special_tokens=skip_special_tokens,
                 clean_up_tokenization_spaces=clean_up_tokenization_spaces)
 
@@ -536,16 +537,28 @@ def hf_decode_incrementally(
         prev_text: Optional[str] = None,
         states: Optional[dict] = None,
         *,
+        flush: bool = False,
         skip_special_tokens: bool = False,
         clean_up_tokenization_spaces: Optional[bool] = None
     ) -> Tuple[str, dict]:
+        """Incremental detokenization with the HF ``DecodeStream``.
+
+        ``DecodeStream.step`` returns ``None`` while a token cannot be decoded
+        on its own yet (e.g. the first byte-level piece of a multi-byte UTF-8
+        character), and has no way to release what it is holding. The ids it
+        holds are tracked in ``states['pending_ids']``, along with the last id
+        that did produce text, so that ``flush`` -- set on the final response
+        of a request -- can emit the held-back text the way a full decode of
+        the output would, instead of silently dropping it.
+        """
         if states is None:
             states = {
                 'decode_stream':
                 DecodeStream(skip_special_tokens=skip_special_tokens)
             }
 
         decode_stream = states.get('decode_stream')
+        pending_ids = states.setdefault('pending_ids', [])
         results = []
         for tid in token_ids:
             try:
@@ -561,9 +574,20 @@ def hf_decode_incrementally(
                 decode_stream = DecodeStream(
                     skip_special_tokens=skip_special_tokens)
                 states['decode_stream'] = decode_stream
+                pending_ids.clear()
+                states.pop('last_decoded_id', None)
                 result = decode_stream.step(self.tokenizer._tokenizer, tid)
-            if result is not None:
+            if result is None:
+                pending_ids.append(tid)
+            else:
                 results.append(result)
+                pending_ids.clear()
+                states['last_decoded_id'] = tid
+        if flush and pending_ids:
+            results.append(
+                self._decode_pending_tail(states.get('last_decoded_id'),
+                                          pending_ids, skip_special_tokens))
+            pending_ids.clear()
         curr_new_text = "".join(results)
         if clean_up_tokenization_spaces is None:
             clean_up_tokenization_spaces = self.clean_up_tokenization_spaces
@@ -575,6 +599,29 @@ def hf_decode_incrementally(
         else:
             return prev_text + curr_new_text, states
 
+    def _decode_pending_tail(self, last_decoded_id: Optional[int],
+                             pending_ids: List[int],
+                             skip_special_tokens: bool) -> str:
+        """Text a full decode would add for the ids ``DecodeStream`` still holds.
+
+        The ids are decoded after the last id that already produced text, and
+        only the part past that id's own text is returned, so decoders that
+        depend on the previous token (e.g. a leading-space rule) see the same
+        context the stream did. Space clean-up is left to the caller, which
+        applies it to the streamed text as a whole.
+        """
+        kwargs = {
+            'skip_special_tokens': skip_special_tokens,
+            'clean_up_tokenization_spaces': False,
+        }
+        if last_decoded_id is None:
+            return self.tokenizer.decode(pending_ids, **kwargs)
+        prefix = self.tokenizer.decode([last_decoded_id], **kwargs)
+        text = self.tokenizer.decode([last_decoded_id] + pending_ids, **kwargs)
+        if text.startswith(prefix):
+            return text[len(prefix):]
+        return self.tokenizer.decode(pending_ids, **kwargs)
+
 
 def tokenizer_factory(obj: Optional[Union[str, Path, PreTrainedTokenizerBase,
                                           TokenizerBase]] = None,

tests/integration/defs/perf/agentx_client.py

@@ -105,6 +105,11 @@
 # ("Feature type 'Json' not found"). 4.x+ is required.
 MIN_DATASETS_MAJOR = 4
 
+# The aiperf scenario that defines the AgentX workload. Only the agentx-harness
+# fork registers it, so its presence is also how the fork is told apart from the
+# upstream PyPI aiperf of the same import name.
+AGENTX_SCENARIO = "inferencex-agentx-mvp"
+
 # PyExecutor's HangDetector is hard-coded at 300 s with no env override, so a
 # multi-minute pip install with zero requests in flight idles the GEN worker
 # into MPI_Abort -- which then presents as a *client* connection failure. A
@@ -174,19 +179,43 @@ def _env_float(name: str, default: float) -> float:
 def aiperf_is_usable() -> bool:
     """Return whether the ambient interpreter already has a working agentx stack.
 
-    Checks importability *and* the datasets major version, since a container
-    with aiperf but datasets 3.x fails only later, when loading the corpus.
+    Checks importability, that the ``aiperf`` found is the agentx-harness fork
+    (it alone ships the ``inferencex-agentx-mvp`` scenario), *and* the datasets
+    major version, since a container with aiperf but datasets 3.x fails only
+    later, when loading the corpus. The fork check matters because
+    requirements-dev.txt pins the upstream PyPI ``aiperf``, which imports fine
+    but rejects ``--scenario`` at the start of the run.
+
+    The scenario is looked up through ``aiperf.common.scenario.registry``, the
+    module path in the pinned ``DEFAULT_AIPERF_REF`` build. The reason for a
+    failed probe is logged, so a fork that moves that module shows up as such
+    instead of as a reinstall on every run.
     """
     probe = (
-        "import aiperf, cyclopts, datasets, transformers, pydantic, msgspec, orjson; "
-        f"import sys; sys.exit(0 if int(datasets.__version__.split('.')[0]) >= {MIN_DATASETS_MAJOR} else 1)"
+        "import sys\n"
+        "try:\n"
+        "    import aiperf, cyclopts, datasets, transformers, pydantic, msgspec, orjson\n"
+        "except ImportError as e:\n"
+        "    sys.exit(f'missing module: {e}')\n"
+        "try:\n"
+        "    from aiperf.common.scenario.registry import get_scenario\n"
+        "except ImportError:\n"
+        "    sys.exit('aiperf has no aiperf.common.scenario.registry: not the agentx-harness fork, '\n"
+        "             'or the fork moved its scenario registry')\n"
+        "try:\n"
+        f"    get_scenario({AGENTX_SCENARIO!r})\n"
+        "except Exception as e:\n"
+        f"    sys.exit(f'aiperf does not provide the {AGENTX_SCENARIO} scenario: {{e}}')\n"
+        f"if int(datasets.__version__.split('.')[0]) < {MIN_DATASETS_MAJOR}:\n"
+        f"    sys.exit(f'datasets {{datasets.__version__}} < {MIN_DATASETS_MAJOR}')\n"
     )
-    return (
-        subprocess.call(
-            [sys.executable, "-c", probe], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL
-        )
-        == 0
+    result = subprocess.run(
+        [sys.executable, "-c", probe], stdout=subprocess.DEVNULL, stderr=subprocess.PIPE, text=True
     )
+    if result.returncode != 0:
+        lines = result.stderr.strip().splitlines()
+        _log(f"aiperf stack not usable: {lines[-1] if lines else f'exit {result.returncode}'}")
+    return result.returncode == 0
 
 
 def post_one_token(url: str, model: str, timeout: float) -> "str | None":
@@ -302,12 +331,19 @@ def build_aiperf_cmd(args: argparse.Namespace, url: str, artifact_dir: str) -> "
     Note ``--url`` carries scheme+host+port only; the request path is a separate
     ``--endpoint``. That is why this client takes ``--host``/``--port`` and
     composes them itself instead of having them appended to the argv.
+
+    aiperf runs as ``python -m aiperf`` under ``sys.executable``, the interpreter
+    that ``aiperf_is_usable`` checked and ``ensure_aiperf`` installed into. A bare
+    ``aiperf`` on PATH could be a different installation, such as the upstream
+    PyPI build that has no AgentX scenario.
     """
     cmd = [
+        sys.executable,
+        "-m",
         "aiperf",
         "profile",
         "--scenario",
-        "inferencex-agentx-mvp",
+        AGENTX_SCENARIO,
         "-m",
         args.model,
         "--tokenizer-trust-remote-code",
@@ -360,6 +396,17 @@ def build_aiperf_cmd(args: argparse.Namespace, url: str, artifact_dir: str) -> "
     if warmup > 0:
         cmd += ["--warmup-requests-per-lane", str(warmup)]
 
+    # The scenario forces ignore_eos, so the model can emit its end-of-sequence
+    # token as an ordinary output token. With the server's default
+    # skip_special_tokens it decodes to no text, and a one-token response
+    # (every warmup request, and every ctx_only request) then carries no
+    # content, which aiperf rejects as InvalidInferenceResultError. Keeping
+    # special tokens in the text makes every generated token visible. The token
+    # counts the metrics are built from stay the same only because the command
+    # also passes --use-server-token-count: without it aiperf would re-tokenize
+    # this text, special-token strings included, and count differently.
+    cmd += ["--extra-inputs", "skip_special_tokens:false"]
+
     # Plain os.environ.get with a default, mirroring the reference script's
     # `${VAR-default}`: an explicitly empty value must disable the flag, which
     # is how a production (>=900 s) run drops --unsafe-override.

tests/integration/test_lists/test-db/l0_cpu.yml

@@ -211,6 +211,7 @@ l0_cpu:
   - unittest/llmapi/test_serve_sleep_mode.py
   - unittest/llmapi/test_session_prefetcher.py
   - unittest/llmapi/test_tokenizer_aliases.py
+  - unittest/llmapi/test_tokenizer_decode.py
   - unittest/llmapi/test_trtllm_serve_config_schemas.py
   - unittest/llmapi/test_utils.py
   - unittest/llmapi/test_visual_gen_startup_shutdown.py