Check input + max_new_tokens against max_out_tokens in InferenceEngine.generate

deepspeed/inference/engine.py

@@ -630,12 +630,20 @@ def _generate(self, *inputs, **kwargs):
                                       "add your request to: https://github.com/deepspeedai/DeepSpeed/issues/2506")
 
         if ("input_ids" in kwargs) and (kwargs["input_ids"].dim() == 2):
+            # max_out_tokens bounds the KV-cache workspace for the whole generation
+            # (input + newly generated tokens), so requested new tokens must be
+            # included here too, not just the input length.
+            max_new_tokens = kwargs.get("max_new_tokens", None)
+            if max_new_tokens is None and "generation_config" in kwargs:
+                max_new_tokens = getattr(kwargs["generation_config"], "max_new_tokens", None)
             for input_tensor in kwargs["input_ids"]:
                 tensor_length = input_tensor.shape[-1]
-                if tensor_length > self._config.max_out_tokens:
+                total_length = tensor_length if max_new_tokens is None else tensor_length + max_new_tokens
+                if total_length > self._config.max_out_tokens:
                     raise RuntimeError(
-                        f"Input with size {tensor_length} exceeds maximum length of {self._config.max_out_tokens}. Please increase max_tokens in the DeepSpeed Inference Config."
-                    )
+                        f"Input with size {tensor_length} and max_new_tokens {max_new_tokens} together exceed "
+                        f"maximum length of {self._config.max_out_tokens}. Please increase max_tokens in the "
+                        "DeepSpeed Inference Config.")
 
         return self.module.generate(*inputs, **kwargs)
 

tests/unit/inference/test_inference_config.py

@@ -21,6 +21,28 @@ def test_max_out_tokens_must_be_positive(field, value):
         DeepSpeedInferenceConfig(**{field: value})
 
 
+@pytest.mark.inference
+def test_generate_rejects_input_plus_max_new_tokens_over_budget():
+    # Regression test for https://github.com/deepspeedai/DeepSpeed/issues/3081.
+    # max_out_tokens sizes the KV-cache workspace for the whole generation (input
+    # tokens plus newly generated ones), but the length check in
+    # InferenceEngine._generate only compared the input length against it, so a
+    # request whose input alone fit could still push input + max_new_tokens past
+    # the workspace size and reach the wrapped module's generate() unguarded.
+    class GenerateStub(torch.nn.Module):
+
+        def generate(self, *args, **kwargs):
+            return "reached-generate"
+
+    engine = deepspeed.init_inference(GenerateStub(), config={"max_out_tokens": 100, "dtype": torch.float32})
+
+    with pytest.raises(RuntimeError, match="exceed"):
+        engine.generate(input_ids=torch.zeros((1, 90), dtype=torch.long), max_new_tokens=20)
+
+    # An input that fits alongside its requested new tokens must still be let through.
+    assert engine.generate(input_ids=torch.zeros((1, 58), dtype=torch.long), max_new_tokens=20) == "reached-generate"
+
+
 @pytest.mark.inference
 @pytest.mark.skipif(not get_accelerator().is_available(), reason="requires accelerator")
 class TestInferenceCudaGraphConfig: