From f03f3b0899206fbe34458494d86b880a417c7e05 Mon Sep 17 00:00:00 2001
From: w_lihongyan <w_lihongyan@kingsoft.com>
Date: Tue, 27 Oct 2026 11:30:43 +0800
Subject: [PATCH] fix: raise when the engine rejects a request in LLM.generate
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

add_request() returns None when it drops a request with invalid token
ids, but LLM._add_request() ignored the return value, so generate()
could return fewer results than the input prompts — silently misaligning
results[i] with prompts[i]. Fail loudly with a clear error instead.

Co-Authored-By: Kscc <noreply@owtffssent.com>
---
 deepinfer/entrypoints/llm.py              | 12 ++++++-
 tests/entrypoints/test_llm_add_request.py | 38 +++++++++++++++++++++++
 2 files changed, 49 insertions(+), 1 deletion(-)
 create mode 100644 tests/entrypoints/test_llm_add_request.py

diff --git a/deepinfer/entrypoints/llm.py b/deepinfer/entrypoints/llm.py
index aa8d5e8..45759db 100644
--- a/deepinfer/entrypoints/llm.py
+++ b/deepinfer/entrypoints/llm.py
@@ -157,13 +157,23 @@ class LLM:
         prompt_token_ids = np.asarray(prompt_token_ids, dtype=np.int32)
         if multi_modal_data is None and self.llm_engine.multi_modal_processor is not None:
             multi_modal_data = dict(image_bytes=[], image_sizes=[])
-        return self.llm_engine.add_request(
+        sequence = self.llm_engine.add_request(
             user_id="0",  # HACK 随便加一个 user_id
             request_id=request_id,
             sampling_params=sampling_params,
             prompt_token_ids=prompt_token_ids,
             multi_modal_data=multi_modal_data,
         )
+        # add_request returns None when it rejects the request (e.g. out-of-range
+        # prompt token ids, see the engine log for details). Silently dropping it
+        # would make generate() return fewer results than the input prompts,
+        # misaligning results[i] with prompts[i], so fail loudly instead.
+        if sequence is None:
+            raise ValueError(
+                f"Request {request_id} was rejected by the engine and returned None "
+                "(see engine log for the reason, e.g. out-of-range prompt token ids)"
+            )
+        return sequence
 
     def _run_engine(self, use_tqdm: bool) -> list[list[RequestOutput] | str]:
         # Initialize tqdm.
diff --git a/tests/entrypoints/test_llm_add_request.py b/tests/entrypoints/test_llm_add_request.py
new file mode 100644
index 0000000..4034ab5
--- /dev/null
+++ b/tests/entrypoints/test_llm_add_request.py
@@ -0,0 +1,38 @@
+import itertools
+from types import SimpleNamespace
+
+import pytest
+
+from deepinfer.entrypoints.llm import LLM
+from deepinfer.sampling_params import SamplingParams
+
+SENTINEL_SEQUENCE = object()
+
+
+def _llm(add_request_result: object) -> LLM:
+    """A bare LLM whose engine add_request always returns the given result."""
+    llm = object.__new__(LLM)
+    llm.request_counter = itertools.count()
+    llm.llm_engine = SimpleNamespace(
+        multi_modal_processor=None,
+        add_request=lambda **kwargs: add_request_result,
+    )
+    return llm
+
+
+def test_add_request_returns_engine_sequence() -> None:
+    llm = _llm(SENTINEL_SEQUENCE)
+
+    sequence = llm._add_request(prompt=None, sampling_params=SamplingParams(), prompt_token_ids=[1, 2, 3])
+
+    assert sequence is SENTINEL_SEQUENCE
+
+
+def test_add_request_raises_when_engine_rejects_request() -> None:
+    # add_request returns None when the engine drops the request (e.g. out-of-range
+    # prompt token ids). generate() must fail loudly instead of silently returning
+    # fewer outputs than input prompts, which would misalign results with prompts.
+    llm = _llm(None)
+
+    with pytest.raises(ValueError, match="rejected by the engine"):
+        llm._add_request(prompt=None, sampling_params=SamplingParams(), prompt_token_ids=[1, 2, 3])