Repository: huggingface/transformers
Pull request: [serving] Fix continuous batching JSON response serialization

Change model_dump_json() to model_dump() to avoid double JSON encoding. When using continuous batching with stream=false, the response was being double-encoded as a string instead of returning a proper JSON object.


Added a UV script to run GPQA using transformer-serve:

```
hf jobs uv run examples/pytorch/transformers_serve_cb_eval_job.py \
--model {model} \
--flavor l4x1 \
--timeout 1h \
--secrets HF_TOKEN \
-e TRANSFORMERS_SERVE_API_KEY="1234"
```


| Model | Job ID | Issue | Output Directory |
|-------|--------|-------|------------------|
| meta-llama/Llama-3.1-8B-Instruct | `69c695adf900226fc14ae1ab` | ❌ `AttributeError: 'NoneType' object has no attribute 'generated_tokens'` - Server starts but all requests fail | [hf/SaylorTwift/llama-3.1-8b-instruct-cb](https://huggingface.co/spaces/SaylorTwift/llama-3.1-8b-instruct-cb) |
| Qwen/Qwen2.5-0.5B-Instruct | `69c696bcbf20ec90acee2ffa` | ✅ Working - Evaluation running successfully | [hf/SaylorTwift/qwen2.5-0.5b-instruct-cb](https://huggingface.co/spaces/SaylorTwift/qwen2.5-0.5b-instruct-cb) |


Changed files:
filename: examples/pytorch/transformers_serve_cb_eval_job.py
patch: @@ -0,0 +1,192 @@
+# /// script
+# dependencies = [
+#     "inspect-ai",
+#     "inspect-evals",
+#     "huggingface-hub",
+#     "transformers[serving] @ git+https://github.com/huggingface/transformers.git@main",
+#     "openai>=2.26.0",
+#     "torchvision",
+#     "kernels",
+# ]
+# ///
+import argparse
+import subprocess
+import sys
+import time
+
+from inspect_ai import eval
+from inspect_ai.log import bundle_log_dir
+from inspect_evals.gpqa import gpqa_diamond
+
+
+def wait_for_server_up(server_process, timeout=600):
+    start_time = time.time()
+
+    import urllib.error
+    import urllib.request
+
+    while time.time() - start_time < timeout:
+        try:
+            req = urllib.request.urlopen("http://127.0.0.1:8000/health", timeout=2)
+            if req.status == 200:
+                elapsed = time.time() - start_time
+                print("\n" + "=" * 70)
+                print(f"✓ Server is ready! (took {elapsed:.1f} seconds)")
+                print("=" * 70)
+                sys.stdout.flush()
+                break
+        except (
+            urllib.error.URLError,
+            ConnectionRefusedError,
+            TimeoutError,
+            OSError,
+        ):
+            elapsed = time.time() - start_time
+            print(f"[{elapsed:.0f}s] Still waiting...", flush=True)
+            time.sleep(5)
+    else:
+        print("\n" + "=" * 70)
+        print("✗ Server failed to start within timeout")
+        print("=" * 70)
+        sys.stdout.flush()
+        if server_process:
+            server_process.terminate()
+        sys.exit(1)
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Run transformers serve with continuous batching and evaluate with inspect-ai"
+    )
+    parser.add_argument(
+        "--model",
+        type=str,
+        required=True,
+        help="Model to serve (e.g., meta-llama/Llama-3.1-8B-Instruct)",
+    )
+    parser.add_argument(
+        "--no-continuous-batching",
+        action="store_true",
+        help="Disable continuous batching (enabled by default)",
+    )
+    parser.add_argument(
+        "--limit",
+        type=int,
+        default=10,
+        help="Number of evaluation samples to run (default: 5)",
+    )
+    parser.add_argument(
+        "--max-connections",
+        type=int,
+        default=10,
+        help="Maximum concurrent connections for evaluation (default: 2)",
+    )
+    parser.add_argument(
+        "--log-dir",
+        type=str,
+        default="./gpqa-diamond-transformers-serve",
+        help="Directory to save evaluation logs (default: ./gpqa-diamond-transformers-serve)",
+    )
+    parser.add_argument(
+        "--output-space",
+        type=str,
+        default="SaylorTwift/transformers-CB",
+        help="Output directory for bundled logs (will be prefixed with 'hf/', default: SaylorTwift/transformers-CB)",
+    )
+
+    # Continuous batching configuration arguments
+    parser.add_argument(
+        "--cb-block-size",
+        type=int,
+        default=256,
+        help="KV cache block size in tokens for continuous batching (default: 256)",
+    )
+    parser.add_argument(
+        "--cb-num-blocks",
+        type=int,
+        default=None,
+        help="Number of KV cache blocks for continuous batching (default: auto-calculated)",
+    )
+    parser.add_argument(
+        "--cb-max-batch-tokens",
+        type=int,
+        default=None,
+        help="Maximum tokens per batch for continuous batching (default: auto-calculated)",
+    )
+    parser.add_argument(
+        "--cb-max-memory-percent",
+        type=float,
+        default=0.8,
+        help="Maximum GPU memory percentage to use for KV cache (0.0-1.0, default: 0.8)",
+    )
+    parser.add_argument(
+        "--cb-use-cuda-graph",
+        action="store_true",
+        help="Enable CUDA graphs for continuous batching performance",
+    )
+
+    args = parser.parse_args()
+
+    server_process = None
+
+    # Build transformers serve command
+    serve_cmd = [
+        "transformers",
+        "serve",
+    ]
+
+    # Add continuous batching if not disabled
+    if not args.no_continuous_batching:
+        serve_cmd.append("--continuous-batching")
+
+        # Add continuous batching configuration arguments
+        serve_cmd.extend(["--cb-block-size", str(args.cb_block_size)])
+
+        if args.cb_num_blocks is not None:
+            serve_cmd.extend(["--cb-num-blocks", str(args.cb_num_blocks)])
+
+        if args.cb_max_batch_tokens is not None:
+            serve_cmd.extend(["--cb-max-batch-tokens", str(args.cb_max_batch_tokens)])
+
+        serve_cmd.extend(["--cb-max-memory-percent", str(args.cb_max_memory_percent)])
+
+        serve_cmd.append("--cb-use-cuda-graph")
+
+    # Always use sdpa attention implementation
+    serve_cmd.extend(["--attn-implementation", "kernels-community/flash-attn2"])
+
+    print("Starting transformers serve with continuous batching...")
+    print(f"Model: {args.model}")
+    if not args.no_continuous_batching:
+        print(f"CB Block Size: {args.cb_block_size}")
+        print(f"CB Num Blocks: {args.cb_num_blocks if args.cb_num_blocks else 'auto'}")
+        print(f"CB Max Batch Tokens: {args.cb_max_batch_tokens if args.cb_max_batch_tokens else 'auto'}")
+        print(f"CB Max Memory: {args.cb_max_memory_percent * 100}%")
+        print(f"CB CUDA Graph: {args.cb_use_cuda_graph}")
+    print(f"Command: {' '.join(serve_cmd)}")
+    print("=" * 70)
+    print("SERVER OUTPUT:")
+    print("=" * 70)
+
+    # Start server with output going directly to stdout/stderr
+    server_process = subprocess.Popen(serve_cmd, stdout=None, stderr=None)
+
+    wait_for_server_up(server_process, timeout=600)
+
+    eval(
+        gpqa_diamond,
+        model=f"openai-api/transformers-serve/{args.model}",
+        log_dir=args.log_dir,
+        model_base_url="http://localhost:8000/v1",
+        display="plain",
+        limit=args.limit,
+        model_args={"stream": False},
+        max_connections=args.max_connections,
+        max_tokens=2048,
+    )
+
+    bundle_log_dir(args.log_dir, output_dir=f"hf/{args.output_space}", overwrite=True)
+
+
+if __name__ == "__main__":
+    main()

filename: src/transformers/cli/serve.py
patch: @@ -561,7 +561,35 @@ def __init__(
         self,
         continuous_batching: Annotated[
             bool | None,
-            typer.Option(help="Whether to use continuous batching for chat completions."),
+            typer.Option(help="Whether to use continuous batching for chat completions. Configure with --cb-* flags."),
+        ] = None,
+        cb_block_size: Annotated[
+            int | None,
+            typer.Option(help="Size of each KV cache block in tokens for continuous batching. Default: 256."),
+        ] = None,
+        cb_num_blocks: Annotated[
+            int | None,
+            typer.Option(
+                help="Number of blocks in KV cache for continuous batching. Default: auto-inferred from GPU memory."
+            ),
+        ] = None,
+        cb_max_batch_tokens: Annotated[
+            int | None,
+            typer.Option(
+                help="Maximum number of tokens in a batch for continuous batching. Default: auto-inferred from GPU memory."
+            ),
+        ] = None,
+        cb_max_memory_percent: Annotated[
+            float | None,
+            typer.Option(
+                help="Maximum percentage of free GPU memory to use for KV cache in continuous batching (0.0-1.0). Default: 0.8."
+            ),
+        ] = None,
+        cb_use_cuda_graph: Annotated[
+            bool | None,
+            typer.Option(
+                help="Enable CUDA graphs for continuous batching performance. Default: auto-inferred based on attention implementation."
+            ),
         ] = None,
         device: Annotated[
             str,
@@ -637,6 +665,13 @@ def __init__(
         self.force_model = force_model
         self.non_blocking = non_blocking
 
+        # Continuous batching configuration arguments
+        self.cb_block_size = cb_block_size
+        self.cb_num_blocks = cb_num_blocks
+        self.cb_max_batch_tokens = cb_max_batch_tokens
+        self.cb_max_memory_percent = cb_max_memory_percent
+        self.cb_use_cuda_graph = cb_use_cuda_graph
+
         # Seed
         if default_seed is not None:
             set_torch_seed(default_seed)
@@ -660,6 +695,8 @@ def __init__(
         # Thread-safety for load_model_and_processor / load_audio_model_and_processor
         self.model_locks: dict[str, threading.Lock] = {}
         self.model_locks_guard = threading.Lock()
+        # Thread-safety for continuous batching manager init/teardown
+        self._cb_manager_lock = threading.Lock()
 
         # 2. preserves information about the last call and last KV cache, to determine whether we can reuse the KV
         # cache and avoid re-running prefill
@@ -1102,45 +1139,64 @@ def continuous_batching_chat_completion(self, req: dict, request_id: str) -> "St
         """
 
         model_id_and_revision = self.process_model_name(req["model"])
-        must_discard_cache = model_id_and_revision != self.last_model
 
-        self.last_model = model_id_and_revision
+        with self._cb_manager_lock:
+            must_discard_cache = model_id_and_revision != self.last_model
+            self.last_model = model_id_and_revision
 
-        # When switching models, terminate a continuous batching manager if it is running.
-        if must_discard_cache:
-            if self.running_continuous_batching_manager is not None:
-                self.running_continuous_batching_manager.stop(block=True, timeout=2)
-                self.running_continuous_batching_manager = None
+            # When switching models, terminate a continuous batching manager if it is running.
+            if must_discard_cache:
+                if self.running_continuous_batching_manager is not None:
+                    self.running_continuous_batching_manager.stop(block=True, timeout=2)
+                    self.running_continuous_batching_manager = None
 
-        model, processor = self.load_model_and_processor(model_id_and_revision)
+            model, processor = self.load_model_and_processor(model_id_and_revision)
 
-        # Continuous batching only supports text-only models
-        if self.get_model_modality(model, processor=processor) != Modality.LLM:
-            logger.warning_once(
-                "Continuous batching is not supported for non-text-only models. Falling back to regular generate."
+            # Continuous batching only supports text-only models
+            if self.get_model_modality(model, processor=processor) != Modality.LLM:
+                logger.warning_once(
+                    "Continuous batching is not supported for non-text-only models. Falling back to regular generate."
+                )
+                return self.generate_chat_completion(req)
+
+            tokenizer = processor.tokenizer if hasattr(processor, "tokenizer") else processor
+
+            generation_config = create_generation_config_from_req(
+                req,
+                model_generation_config=model.generation_config,
+                eos_token_id=tokenizer.eos_token_id,
+                pad_token_id=tokenizer.pad_token_id,
+                use_cache=False,
+                do_sample=False,
+                scheduler="fifo",
             )
-            return self.generate_chat_completion(req)
-
-        tokenizer = processor.tokenizer if hasattr(processor, "tokenizer") else processor
-
-        generation_config = create_generation_config_from_req(
-            req,
-            model_generation_config=model.generation_config,
-            eos_token_id=tokenizer.eos_token_id,
-            pad_token_id=tokenizer.pad_token_id,
-            use_cache=False,
-            do_sample=False,
-            scheduler="fifo",
-        )
 
-        if self.running_continuous_batching_manager is None:
-            self.running_continuous_batching_manager = model.init_continuous_batching(
-                generation_config=generation_config
-            )
+            if self.running_continuous_batching_manager is None:
+                from transformers import ContinuousBatchingConfig
+
+                # Build continuous batching config from CLI arguments
+                cb_config_kwargs = {}
+                if self.cb_block_size is not None:
+                    cb_config_kwargs["block_size"] = self.cb_block_size
+                if self.cb_num_blocks is not None:
+                    cb_config_kwargs["num_blocks"] = self.cb_num_blocks
+                if self.cb_max_batch_tokens is not None:
+                    cb_config_kwargs["max_batch_tokens"] = self.cb_max_batch_tokens
+                if self.cb_max_memory_percent is not None:
+                    cb_config_kwargs["max_memory_percent"] = self.cb_max_memory_percent
+                if self.cb_use_cuda_graph is not None:
+                    cb_config_kwargs["use_cuda_graph"] = self.cb_use_cuda_graph
+
+                cb_config = ContinuousBatchingConfig(**cb_config_kwargs) if cb_config_kwargs else None
+
+                self.running_continuous_batching_manager = model.init_continuous_batching(
+                    generation_config=generation_config,
+                    continuous_batching_config=cb_config,
+                )
 
-            # TODO (Joao, Lysandre): the logits processors should be fixed in continuous batching and correctly applied in non-cb
-            self.running_continuous_batching_manager.logit_processor = LogitsProcessorList()
-            self.running_continuous_batching_manager.start()
+                # TODO (Joao, Lysandre): the logits processors should be fixed in continuous batching and correctly applied in non-cb
+                self.running_continuous_batching_manager.logit_processor = LogitsProcessorList()
+                self.running_continuous_batching_manager.start()
 
         # TODO (Joao, Lysandre): this should also work with tool support
         modality = self.get_model_modality(model, processor=processor)
@@ -1209,6 +1265,9 @@ def buffer_chat_completion(_request_id):
             while self.running_continuous_batching_manager.is_running() and result is None:
                 result = self.running_continuous_batching_manager.get_result(request_id=_request_id, timeout=1)
 
+            if result is None:
+                raise RuntimeError(f"Request {_request_id} failed: generation loop stopped before producing a result.")
+
             content = tokenizer.decode(result.generated_tokens)
 
             chat_completion_result = ChatCompletion(
@@ -1257,7 +1316,7 @@ def cancellation_wrapper_buffer(_request_id):
             return StreamingResponse(cancellation_wrapper_stream(request_id), media_type="text/event-stream")
         else:
             chunk = cancellation_wrapper_buffer(request_id)
-            json_chunk = chunk.model_dump_json(exclude_none=True)
+            json_chunk = chunk.model_dump(exclude_none=True)
             return JSONResponse(json_chunk, media_type="application/json")
 
     @staticmethod
@@ -1633,9 +1692,7 @@ def generate_response(self, req: dict) -> Generator[str, None, None]:
         else:
             raise TypeError("inputs should be a list, dict, or str")
 
-        inputs = processor.apply_chat_template(
-            inputs, add_generation_prompt=True, return_tensors="pt", return_dict=True
-        )["input_ids"]
+        inputs = processor.apply_chat_template(inputs, tokenize=True, add_generation_prompt=True, return_tensors="pt")
         inputs = inputs.to(model.device)
         request_id = req.get("previous_response_id", "req_0")
 

filename: tests/cli/test_serve.py
patch: @@ -729,6 +729,38 @@ def test_request_cancellation(self):
             f"(last seen at {last_seen}). Check cancellation propagation.",
         )
 
+    def test_non_streaming_response_json_format(self):
+        """
+        Tests that non-streaming continuous batching responses return proper JSON objects,
+        not double-encoded JSON strings (regression test for JSON serialization fix).
+        """
+        client = OpenAI(base_url=f"http://localhost:{self.port}/v1", api_key="<KEY>")
+
+        # Make a non-streaming request
+        response = client.chat.completions.create(
+            model="Qwen/Qwen2.5-0.5B-Instruct",
+            messages=[{"role": "user", "content": "Say hello"}],
+            stream=False,
+            max_tokens=5,
+        )
+
+        # Verify response is a proper ChatCompletion object (not a string)
+        self.assertIsNotNone(response)
+        self.assertIsNotNone(response.id)
+        self.assertIsNotNone(response.choices)
+        self.assertEqual(len(response.choices), 1)
+
+        # Verify the choice has proper structure
+        choice = response.choices[0]
+        self.assertIsNotNone(choice.message)
+        self.assertIsNotNone(choice.message.content)
+        self.assertEqual(choice.message.role, "assistant")
+
+        # Verify content is a string, not a serialized JSON
+        content = choice.message.content
+        self.assertIsInstance(content, str)
+        self.assertTrue(len(content) > 0)
+
 
 @require_openai
 class ServeResponsesMixin: