diff --git a/docs/source/user_guide/examples/experimental-server.md b/docs/source/user_guide/examples/experimental-server.md
index f0f001427..6a8d0a158 100644
--- a/docs/source/user_guide/examples/experimental-server.md
+++ b/docs/source/user_guide/examples/experimental-server.md
@@ -348,13 +348,32 @@ API.
## Runtime Concurrency
-The current high-level runtime has one mutable generation state. The server
-therefore admits one request at a time and uses a bounded async queue configured
+The current high-level runtime has one mutable generation state. By default the
+server admits one request at a time and uses a bounded async queue configured
by `--max-queued-requests` and `--queue-timeout`. Queue overflow and timeout
return HTTP 429 (Anthropic 529). Streaming disconnects cancel the native channel
immediately, wait for the native worker to exit, and then release the runtime
lease. Engines stay resident across HTTP connections; graceful server shutdown
drains active work and releases the runtime and its device resources.
+`--enable-batching` merges concurrent non-streaming requests into one runtime
+call. Requests join a batch only when their generation-level settings match
+(sampling parameters, `max_tokens`, chat-template and thinking flags,
+speculative-decoding and context-cache options); each request keeps its own
+messages, media, stop strings and logit bias. A request waits at most
+`--batch-timeout-ms` (default 10) for others to arrive, and one call carries at
+most `--max-queue-batch-size` requests, capped by the engine's
+`--max-batch-size`. Decode is memory-bandwidth bound, so a batch of four costs
+little more than a single request and aggregate throughput scales almost
+linearly. Streaming requests still run one at a time; a streaming request in
+flight delays batched requests until it completes. `/health` reports the
+active batching settings.
+
+`--max-verify-tree-size` and `--max-draft-tree-size` bound the speculative
+tree the compiled bundle supports. The builder defaults both to 60; for linear
+MTP drafting a bound of one more than the largest `num_speculative_tokens` you
+intend to run is sufficient, and the runtime's spec-verify state buffers
+scale with it (about 0.9 GB per position for a 27B hybrid model at batch 4).
+
Continuous batching, chunked prefill scheduling, and tensor parallelism require
additional native scheduler support and are rejected at launch.
diff --git a/experimental/builder/core/artifacts/chat_template.py b/experimental/builder/core/artifacts/chat_template.py
index 663dde55b..3942a949f 100644
--- a/experimental/builder/core/artifacts/chat_template.py
+++ b/experimental/builder/core/artifacts/chat_template.py
@@ -92,8 +92,23 @@ def try_write_packaged_chat_template(model_dir: str,
def _format_chat(tokenizer, messages, **kwargs) -> str:
- """Apply a tokenizer chat template and return text."""
- return tokenizer.apply_chat_template(messages, tokenize=False, **kwargs)
+ """Apply a tokenizer chat template and return text.
+
+ Probes render with thinking disabled unless told otherwise: templates
+ that rewrite the conversation in thinking mode (Qwen3.8 injects reasoning
+ instructions into the system block and treats an undefined flag as
+ enabled) would otherwise yield prefixes that never match a request.
+ """
+ kwargs.setdefault("enable_thinking", False)
+ try:
+ return tokenizer.apply_chat_template(messages,
+ tokenize=False,
+ **kwargs)
+ except TypeError:
+ kwargs.pop("enable_thinking", None)
+ return tokenizer.apply_chat_template(messages,
+ tokenize=False,
+ **kwargs)
def _format_chat_generation(tokenizer, messages, enable_thinking: bool) -> str:
@@ -213,10 +228,17 @@ def process_chat_template(model_dir: str,
generation_prompt_thinking = None
try:
+ # Slice against a thinking-mode baseline: a template may render a
+ # longer conversation in thinking mode, so the non-thinking
+ # baseline offset would cut into the wrong place.
+ thinking_base = _format_chat(tokenizer,
+ [system_message, user_message],
+ add_generation_prompt=False,
+ enable_thinking=True)
thinking_formatted = _format_chat_generation(
tokenizer, [system_message, user_message],
enable_thinking=True)
- candidate = thinking_formatted[len(user_formatted):]
+ candidate = thinking_formatted[len(thinking_base):]
if candidate != generation_prompt:
generation_prompt_thinking = candidate
except (TypeError, ValueError, KeyError):
diff --git a/experimental/builder/core/weights.py b/experimental/builder/core/weights.py
index a1e63567f..ac3a31f9d 100644
--- a/experimental/builder/core/weights.py
+++ b/experimental/builder/core/weights.py
@@ -188,6 +188,9 @@ def _resolve(self, name: str, required: bool = True) -> Optional[str]:
if not candidate.endswith("embed_tokens.weight")
]
candidates = list(dict.fromkeys(candidates))
+ owned = self._sidecar_in_owner_namespace(name)
+ if owned is not None:
+ candidates = [owned]
for candidate in candidates:
if self.store.has(candidate):
return candidate
@@ -195,6 +198,36 @@ def _resolve(self, name: str, required: bool = True) -> Optional[str]:
raise KeyError(f"checkpoint tensor not found: {name!r}")
return None
+ # Per-module tensors that must live beside the module's weight. A draft
+ # resolves its modules in a separate checkpoint namespace (Qwen3.5 MTP
+ # ``mtp.``); without this pin, a sidecar missing there would be answered
+ # by the base model's identically named module.
+ _MODULE_SIDECARS = (".weight_scale", ".weight_scale_2",
+ ".weight_global_scale", ".input_scale",
+ ".input_global_scale", ".bias", ".pre_quant_scale",
+ ".qzeros", ".scales", ".g_idx", ".q_scale", ".k_scale",
+ ".v_scale")
+ _PRIMARY_WEIGHTS = (".weight", ".qweight", ".weight_packed")
+
+ def _sidecar_in_owner_namespace(self, name: str) -> Optional[str]:
+ """Concrete key for a module sidecar, pinned to its weight's namespace.
+
+ Returns None when *name* is not a module sidecar or the module has no
+ resolvable primary weight, leaving ordinary candidate resolution to
+ run.
+ """
+ for suffix in self._MODULE_SIDECARS:
+ if name.endswith(suffix) and len(name) > len(suffix):
+ prefix = name[:-len(suffix)]
+ break
+ else:
+ return None
+ for weight_suffix in self._PRIMARY_WEIGHTS:
+ key = self._resolve(prefix + weight_suffix, required=False)
+ if key is not None:
+ return key[:-len(weight_suffix)] + suffix
+ return None
+
def checkpoint_key(self, name: str) -> str:
"""Return the concrete checkpoint key backing a model tensor name."""
return self._resolve(name, required=False) or name
@@ -241,15 +274,36 @@ def checkpoint_locations(self, names: Sequence[str]) -> dict:
locations[name] = location
return locations
+ _QUANT_SIDECAR_SUFFIXES = (".weight_scale", ".weight_scale_2",
+ ".weight_global_scale", ".qweight",
+ ".weight_packed", ".scales")
+
def module_quant_type(self,
name: str,
*,
tie_word_embeddings: bool = False) -> str:
"""Return the checkpoint precision owned by one model projection."""
+ lookup = name
normalize = getattr(self.conversion, "normalize_checkpoint_name", None)
if normalize is not None:
- name = normalize(name)
- return self.quant.module_type(name, tie_word_embeddings)
+ lookup = normalize(lookup)
+ quant_type = self.quant.module_type(lookup, tie_word_embeddings)
+ # A draft resolves its projections in a separate checkpoint namespace
+ # (Qwen3.5 MTP stores them under ``mtp.``), where the same short name
+ # may be unquantized even though the base layer carries a quantized
+ # override. The tensor actually read decides: a bare weight with no
+ # quantization sidecar is FP16. ``lm_head`` is skipped because its
+ # resolution consults this method for embedding tying.
+ if (quant_type != quantization.QUANT_FP16 and name != "lm_head"
+ and self._has_plain_weight(name)):
+ return quantization.QUANT_FP16
+ return quant_type
+
+ def _has_plain_weight(self, name: str) -> bool:
+ if not self.has(name + ".weight"):
+ return False
+ return not any(
+ self.has(name + suffix) for suffix in self._QUANT_SIDECAR_SUFFIXES)
def parameter_spec(self,
name: str,
diff --git a/experimental/builder/models/dflash/modeling_dflash_draft.py b/experimental/builder/models/dflash/modeling_dflash_draft.py
index 116381513..734b783fb 100644
--- a/experimental/builder/models/dflash/modeling_dflash_draft.py
+++ b/experimental/builder/models/dflash/modeling_dflash_draft.py
@@ -163,13 +163,15 @@ def __init__(self, ctx, lm_head=None) -> None:
def input_tensors(self) -> Dict[str, object]:
cfg = self.cfg
+ kv_dtype = (trt.DataType.FP8
+ if cfg.kv_cache_quant == "fp8" else trt.float16)
target_layers = cfg.dflash_target_layer_ids or [1, 8, 15, 22, 29]
return {
"inputs_embeds":
self.add_input("inputs_embeds", trt.float16,
(-1, -1, cfg.hidden_size)),
"past_key_values": [
- self.add_input(f"past_key_values_{index}", trt.float16,
+ self.add_input(f"past_key_values_{index}", kv_dtype,
(2, -1, F.KV_PAGE_SIZE, cfg.num_key_value_heads,
cfg.head_dim))
for index in range(cfg.num_hidden_layers)
diff --git a/experimental/builder/models/dspark/modeling_dspark_draft.py b/experimental/builder/models/dspark/modeling_dspark_draft.py
index 41b0189ce..7fa164102 100644
--- a/experimental/builder/models/dspark/modeling_dspark_draft.py
+++ b/experimental/builder/models/dspark/modeling_dspark_draft.py
@@ -128,13 +128,15 @@ def __init__(self, ctx) -> None:
def input_tensors(self) -> Dict[str, object]:
cfg = self.cfg
+ kv_dtype = (trt.DataType.FP8
+ if cfg.kv_cache_quant == "fp8" else trt.float16)
target_layers = cfg.dspark_target_layer_ids
return {
"inputs_embeds":
self.add_input("inputs_embeds", trt.float16,
(-1, -1, cfg.hidden_size)),
"past_key_values": [
- self.add_input(f"past_key_values_{index}", trt.float16,
+ self.add_input(f"past_key_values_{index}", kv_dtype,
(2, -1, F.KV_PAGE_SIZE, cfg.num_key_value_heads,
cfg.head_dim))
for index in range(cfg.num_hidden_layers)
diff --git a/experimental/builder/models/eagle3/modeling_eagle3_draft.py b/experimental/builder/models/eagle3/modeling_eagle3_draft.py
index 5ce89f1c8..6797d4da6 100644
--- a/experimental/builder/models/eagle3/modeling_eagle3_draft.py
+++ b/experimental/builder/models/eagle3/modeling_eagle3_draft.py
@@ -102,6 +102,8 @@ def __init__(self, ctx, lm_head=None) -> None:
def input_tensors(self) -> Dict[str, object]:
cfg = self.cfg
+ kv_dtype = (trt.DataType.FP8
+ if cfg.kv_cache_quant == "fp8" else trt.float16)
target_hidden = int(cfg.target_hidden_size or cfg.raw_component.get(
"eagle3_target_hidden_size", cfg.hidden_size))
target_layers = len(cfg.eagle3_target_layer_ids)
@@ -112,7 +114,7 @@ def input_tensors(self) -> Dict[str, object]:
self.add_input("inputs_embeds", trt.float16,
(-1, -1, cfg.hidden_size)),
"past_key_values": [
- self.add_input(f"past_key_values_{index}", trt.float16,
+ self.add_input(f"past_key_values_{index}", kv_dtype,
(2, -1, F.KV_PAGE_SIZE, cfg.num_key_value_heads,
cfg.head_dim))
for index in range(cfg.num_hidden_layers)
diff --git a/experimental/builder/models/qwen3_5/modeling_qwen3_5_mtp.py b/experimental/builder/models/qwen3_5/modeling_qwen3_5_mtp.py
index 9b58938d0..ab496de1a 100644
--- a/experimental/builder/models/qwen3_5/modeling_qwen3_5_mtp.py
+++ b/experimental/builder/models/qwen3_5/modeling_qwen3_5_mtp.py
@@ -77,12 +77,14 @@ def __init__(self, ctx) -> None:
def input_tensors(self) -> Dict[str, object]:
cfg = self.cfg
+ kv_dtype = (trt.DataType.FP8
+ if cfg.kv_cache_quant == "fp8" else trt.float16)
return {
"inputs_embeds":
self.add_input("inputs_embeds", trt.float16,
(-1, -1, cfg.hidden_size)),
"past_key_values": [
- self.add_input(f"past_key_values_{index}", trt.float16,
+ self.add_input(f"past_key_values_{index}", kv_dtype,
(2, -1, F.KV_PAGE_SIZE, cfg.num_key_value_heads,
cfg.head_dim))
for index in range(cfg.num_hidden_layers)
diff --git a/experimental/builder/models/qwen3_omni/modeling_qwen3_omni_mtp.py b/experimental/builder/models/qwen3_omni/modeling_qwen3_omni_mtp.py
index 483f9738d..f6a9918d4 100644
--- a/experimental/builder/models/qwen3_omni/modeling_qwen3_omni_mtp.py
+++ b/experimental/builder/models/qwen3_omni/modeling_qwen3_omni_mtp.py
@@ -127,12 +127,14 @@ def __init__(self, ctx) -> None:
def input_tensors(self) -> Dict[str, object]:
cfg = self.cfg
+ kv_dtype = (trt.DataType.FP8
+ if cfg.kv_cache_quant == "fp8" else trt.float16)
return {
"inputs_embeds":
self.add_input("inputs_embeds", trt.float16,
(-1, -1, cfg.hidden_size)),
"past_key_values": [
- self.add_input(f"past_key_values_{index}", trt.float16,
+ self.add_input(f"past_key_values_{index}", kv_dtype,
(2, -1, F.KV_PAGE_SIZE, cfg.num_key_value_heads,
cfg.head_dim))
for index in range(cfg.num_hidden_layers)
diff --git a/experimental/builder/models/qwen3_omni_next/modeling_qwen3_omni_next_mtp.py b/experimental/builder/models/qwen3_omni_next/modeling_qwen3_omni_next_mtp.py
index f4dd93d6c..00166d5e2 100644
--- a/experimental/builder/models/qwen3_omni_next/modeling_qwen3_omni_next_mtp.py
+++ b/experimental/builder/models/qwen3_omni_next/modeling_qwen3_omni_next_mtp.py
@@ -85,12 +85,14 @@ def __init__(self, ctx) -> None:
def input_tensors(self) -> Dict[str, object]:
cfg = self.cfg
+ kv_dtype = (trt.DataType.FP8
+ if cfg.kv_cache_quant == "fp8" else trt.float16)
return {
"inputs_embeds":
self.add_input("inputs_embeds", trt.float16,
(-1, -1, cfg.hidden_size)),
"past_key_values": [
- self.add_input(f"past_key_values_{index}", trt.float16,
+ self.add_input(f"past_key_values_{index}", kv_dtype,
(2, -1, F.KV_PAGE_SIZE, cfg.num_key_value_heads,
cfg.head_dim))
for index in range(cfg.num_hidden_layers)
diff --git a/experimental/server/api/routes.py b/experimental/server/api/routes.py
index f4648a6b2..37ed5d130 100644
--- a/experimental/server/api/routes.py
+++ b/experimental/server/api/routes.py
@@ -50,6 +50,7 @@ async def health(request: Request):
"model": client.model_name,
"active_requests": client.active_requests,
"queued_requests": client.queued_requests,
+ "batching": client.batching,
"capabilities": {
"chat": caps.chat,
"transcription": caps.transcription,
diff --git a/experimental/server/config.py b/experimental/server/config.py
index 7be49bc83..7f46a2a63 100644
--- a/experimental/server/config.py
+++ b/experimental/server/config.py
@@ -174,6 +174,11 @@ class ModelConfig:
draft_top_k: Optional[int] = None
draft_step: Optional[int] = None
verify_tree_size: Optional[int] = None
+ max_verify_tree_size: Optional[int] = None
+ max_draft_tree_size: Optional[int] = None
+ enable_batching: bool = False
+ batch_timeout_ms: float = 10.0
+ max_queue_batch_size: Optional[int] = None
speculative_config: Optional[SpeculativeConfig] = None
context_cache_config: ContextCacheConfig = field(
default_factory=ContextCacheConfig)
@@ -184,6 +189,10 @@ def __post_init__(self) -> None:
or self.engine_cache_max_size_gb <= 0):
raise ServerConfigError(
"engine_cache_max_size_gb must be positive")
+ if (isinstance(self.batch_timeout_ms, bool)
+ or not math.isfinite(self.batch_timeout_ms)
+ or self.batch_timeout_ms < 0):
+ raise ServerConfigError("batch_timeout_ms must be non-negative")
def llm_kwargs(self) -> Dict[str, Any]:
return {
@@ -197,6 +206,11 @@ def llm_kwargs(self) -> Dict[str, Any]:
"draft_top_k": self.draft_top_k,
"draft_step": self.draft_step,
"verify_tree_size": self.verify_tree_size,
+ "max_verify_tree_size": self.max_verify_tree_size,
+ "max_draft_tree_size": self.max_draft_tree_size,
+ "enable_batching": self.enable_batching,
+ "batch_timeout_ms": self.batch_timeout_ms,
+ "max_queue_batch_size": self.max_queue_batch_size,
"speculative_config": self.speculative_config,
"context_cache_config": self.context_cache_config,
}
@@ -254,6 +268,13 @@ def _positive_float(value: str) -> float:
return parsed
+def _non_negative_float(value: str) -> float:
+ parsed = float(value)
+ if not math.isfinite(parsed) or parsed < 0:
+ raise argparse.ArgumentTypeError("must be non-negative")
+ return parsed
+
+
def _bounded_port(value: str) -> int:
parsed = int(value)
if not 1 <= parsed <= 65535:
@@ -325,7 +346,39 @@ def create_argument_parser() -> argparse.ArgumentParser:
model.add_argument("--draft-top-k", type=_positive_int)
model.add_argument("--draft-step", type=_positive_int)
model.add_argument("--verify-tree-size", type=_positive_int)
+ model.add_argument(
+ "--max-verify-tree-size",
+ type=_positive_int,
+ help="Largest base verification input the compiled bundle supports. "
+ "Spec-verify state buffers scale with it; the builder default is 60.",
+ )
+ model.add_argument(
+ "--max-draft-tree-size",
+ type=_positive_int,
+ help="Largest draft proposal the compiled bundle supports; the "
+ "builder default is 60.",
+ )
model.add_argument("--speculative-config", default="")
+ model.add_argument(
+ "--enable-batching",
+ action="store_true",
+ help="Merge concurrent non-streaming requests with matching sampling "
+ "settings into one runtime call, up to the engine's max batch size. "
+ "Streaming requests still run one at a time.",
+ )
+ model.add_argument(
+ "--batch-timeout-ms",
+ type=_non_negative_float,
+ default=10.0,
+ help="How long a request waits for compatible requests to join its "
+ "batch before running.",
+ )
+ model.add_argument(
+ "--max-queue-batch-size",
+ type=_positive_int,
+ help="Cap on requests merged per runtime call; defaults to the "
+ "engine's max batch size.",
+ )
model.add_argument(
"--enable-context-reuse",
action="store_true",
@@ -377,6 +430,11 @@ def parse_server_config(argv: Optional[Sequence[str]] = None) -> ServerConfig:
draft_top_k=args.draft_top_k,
draft_step=draft_step,
verify_tree_size=args.verify_tree_size,
+ max_verify_tree_size=args.max_verify_tree_size,
+ max_draft_tree_size=args.max_draft_tree_size,
+ enable_batching=args.enable_batching,
+ batch_timeout_ms=args.batch_timeout_ms,
+ max_queue_batch_size=args.max_queue_batch_size,
speculative_config=speculative,
context_cache_config=context_cache,
)
diff --git a/experimental/server/runtime/batching.py b/experimental/server/runtime/batching.py
new file mode 100644
index 000000000..bcc4441eb
--- /dev/null
+++ b/experimental/server/runtime/batching.py
@@ -0,0 +1,289 @@
+# 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.
+"""Micro-batching of compatible non-streaming generation requests.
+
+The C++ runtime processes one ``LLMGenerationRequest`` at a time, but that
+request may carry several rows (``requests``) that decode together up to the
+engine's ``max_batch_size``. Decode is memory-bandwidth bound, so rows in one
+call cost little more than a single row. The batcher collects concurrent
+HTTP requests whose generation-level settings match, merges their rows into
+one runtime call, and splits the response back per caller.
+"""
+
+import logging
+import threading
+import time
+from concurrent.futures import Future
+from dataclasses import dataclass
+from typing import Any, Callable, List, Optional, Tuple
+
+logger = logging.getLogger("edgellm.batching")
+
+
+class BatcherOverflow(RuntimeError):
+ """The pending queue is at capacity; the caller should shed load."""
+
+
+# Rows may only share one runtime call when these request-level settings
+# match. Per-row inputs (messages, media, stop strings, logit bias) live under
+# ``LLMGenerationRequest.requests`` and are free to differ.
+BATCH_COMPATIBILITY_FIELDS = (
+ "temperature",
+ "top_p",
+ "top_k",
+ "max_generate_length",
+ "lora_weights_name",
+ "save_system_prompt_kv_cache",
+ "apply_chat_template",
+ "add_generation_prompt",
+ "enable_thinking",
+ "disable_spec_decode",
+ "recurrent_capture_interval",
+ "num_logprobs",
+ "context_cache_lookup_policy",
+ "context_cache_commit_policy",
+ "context_cache_replay_tail_length",
+)
+
+
+@dataclass
+class ResponseSlice:
+ """The rows of one runtime response that belong to one submitted request.
+
+ Exposes the same attribute names as ``LLMGenerationResponse`` so callers
+ can consume either interchangeably.
+ """
+
+ output_texts: List[str]
+ output_ids: List[List[int]]
+ finish_reasons: List[Any]
+ logprobs: List[Any]
+ prompt_token_counts: List[int]
+
+
+@dataclass
+class _QueuedRequest:
+ request: Any
+ future: Future
+
+
+def resolve_batch_size(engine_max_batch_size: int,
+ max_queue_batch_size: Optional[int]) -> int:
+ """Effective micro-batch size: the requested cap, bounded by the engine."""
+ if max_queue_batch_size is not None:
+ if 0 < engine_max_batch_size < max_queue_batch_size:
+ logger.warning(
+ "Capping max_queue_batch_size=%d to engine max_batch_size=%d",
+ max_queue_batch_size, engine_max_batch_size)
+ return engine_max_batch_size
+ return max_queue_batch_size
+ return engine_max_batch_size or 1
+
+
+def _batch_key(request) -> Tuple[Any, ...]:
+ return tuple(
+ getattr(request, field) for field in BATCH_COMPATIBILITY_FIELDS)
+
+
+def _is_batchable(request, video_requires_singleton: bool) -> bool:
+ """Video rows on runners that enqueue per-clip tubelet shapes run alone."""
+ if not video_requires_singleton:
+ return True
+ try:
+ for row in request.requests:
+ for buf in row.image_buffers:
+ if buf.is_video:
+ return False
+ except Exception:
+ # A request that cannot be introspected runs alone.
+ return False
+ return True
+
+
+def _copy_batch_settings(source, target) -> None:
+ for field in BATCH_COMPATIBILITY_FIELDS:
+ setattr(target, field, getattr(source, field))
+ target.stream_channels = []
+
+
+def _slice_response(response, start: int, count: int) -> ResponseSlice:
+ end = start + count
+ logprobs = getattr(response, "logprobs", None) or []
+ prompt_tokens = getattr(response, "prompt_token_counts", None) or []
+ return ResponseSlice(
+ output_texts=list(response.output_texts[start:end]),
+ output_ids=[list(ids) for ids in response.output_ids[start:end]],
+ finish_reasons=list(response.finish_reasons[start:end]),
+ logprobs=list(logprobs[start:end]),
+ prompt_token_counts=list(prompt_tokens[start:end]),
+ )
+
+
+class RequestBatcher:
+ """Merge compatible requests and serialize runtime calls on one worker.
+
+ ``runtime_handler`` receives a merged ``LLMGenerationRequest`` and returns
+ the runtime response; it is expected to hold whatever lock the runtime
+ needs. ``submit`` blocks the calling thread until its slice is available.
+ """
+
+ def __init__(
+ self,
+ runtime_handler: Callable[[Any], Any],
+ max_batch_size: int,
+ timeout_ms: float,
+ max_pending: Optional[int] = None,
+ video_requires_singleton: bool = False,
+ ) -> None:
+ if max_batch_size < 1:
+ raise ValueError("max_batch_size must be positive")
+ if timeout_ms < 0:
+ raise ValueError("timeout_ms must be non-negative")
+ if max_pending is not None and max_pending < 1:
+ raise ValueError("max_pending must be positive")
+
+ self._runtime_handler = runtime_handler
+ self._max_batch_size = max_batch_size
+ self._timeout_s = timeout_ms / 1000.0
+ self._max_pending = max_pending
+ self._video_requires_singleton = video_requires_singleton
+ self._cv = threading.Condition()
+ self._queue: List[_QueuedRequest] = []
+ self._closed = False
+ self._worker = threading.Thread(target=self._run,
+ name="edgellm-request-batcher",
+ daemon=True)
+ self._worker.start()
+
+ @property
+ def max_batch_size(self) -> int:
+ return self._max_batch_size
+
+ @property
+ def timeout_ms(self) -> float:
+ return self._timeout_s * 1000.0
+
+ @property
+ def pending(self) -> int:
+ with self._cv:
+ return len(self._queue)
+
+ def submit(self, request) -> ResponseSlice:
+ future: Future = Future()
+ with self._cv:
+ if self._closed:
+ raise RuntimeError("Request batcher is closed")
+ if (self._max_pending is not None
+ and len(self._queue) >= self._max_pending):
+ raise BatcherOverflow(
+ f"batcher queue full ({self._max_pending} pending)")
+ self._queue.append(_QueuedRequest(request=request, future=future))
+ self._cv.notify()
+ return future.result()
+
+ def close(self) -> None:
+ """Stop accepting work; queued requests still complete."""
+ with self._cv:
+ self._closed = True
+ self._cv.notify_all()
+ self._worker.join(timeout=5.0)
+
+ def _run(self) -> None:
+ # The worker must outlive any single failure: a dead worker would
+ # leave every pending Future, and the caller thread parked on it,
+ # blocked forever.
+ while True:
+ try:
+ batch = self._take_batch()
+ except Exception:
+ logger.exception("Batch selection failed; continuing")
+ continue
+ if batch is None:
+ return
+ if batch:
+ self._process_batch(batch)
+
+ def _take_batch(self) -> Optional[List[_QueuedRequest]]:
+ with self._cv:
+ while not self._queue and not self._closed:
+ self._cv.wait()
+ if not self._queue:
+ return None
+
+ first = self._queue.pop(0)
+ try:
+ key = _batch_key(first.request)
+ except Exception as exc:
+ first.future.set_exception(exc)
+ return []
+ batch = [first]
+ if not _is_batchable(first.request,
+ self._video_requires_singleton):
+ return batch
+ deadline = time.monotonic() + self._timeout_s
+ while len(batch) < self._max_batch_size:
+ self._move_compatible_locked(batch, key)
+ if len(batch) >= self._max_batch_size or self._closed:
+ break
+ remaining = deadline - time.monotonic()
+ if remaining <= 0:
+ break
+ self._cv.wait(remaining)
+ return batch
+
+ def _move_compatible_locked(self, batch: List[_QueuedRequest],
+ key: Tuple[Any, ...]) -> None:
+ idx = 0
+ while idx < len(self._queue) and len(batch) < self._max_batch_size:
+ item = self._queue[idx]
+ try:
+ compatible = (_batch_key(
+ item.request) == key and _is_batchable(
+ item.request, self._video_requires_singleton))
+ except Exception:
+ # Unkeyable requests are popped first on a later round and
+ # fail there on their own.
+ compatible = False
+ if compatible:
+ batch.append(self._queue.pop(idx))
+ else:
+ idx += 1
+
+ def _process_batch(self, batch: List[_QueuedRequest]) -> None:
+ try:
+ response = self._runtime_handler(self._merge(batch))
+ offset = 0
+ for item in batch:
+ count = len(item.request.requests)
+ item.future.set_result(_slice_response(response, offset,
+ count))
+ offset += count
+ except Exception as exc:
+ logger.exception("Batched inference failed")
+ for item in batch:
+ item.future.set_exception(exc)
+
+ @staticmethod
+ def _merge(batch: List[_QueuedRequest]):
+ if len(batch) == 1:
+ return batch[0].request
+ first = batch[0].request
+ merged = type(first)()
+ _copy_batch_settings(first, merged)
+ rows = []
+ for item in batch:
+ rows.extend(item.request.requests)
+ merged.requests = rows
+ return merged
diff --git a/experimental/server/runtime/engine.py b/experimental/server/runtime/engine.py
index 1b92752df..2d3cf428c 100644
--- a/experimental/server/runtime/engine.py
+++ b/experimental/server/runtime/engine.py
@@ -51,6 +51,7 @@
validate_tool_request)
from ..parsing.tool_chat_template import (ToolChatTemplateFormatter,
needs_tool_chat_template)
+from .batching import RequestBatcher, resolve_batch_size
from .engine_layout import BundleLayout, EngineType, inspect_bundle
logger = logging.getLogger("edgellm.server")
@@ -658,6 +659,11 @@ def __init__(
draft_top_k: Optional[int] = None,
draft_step: Optional[int] = None,
verify_tree_size: Optional[int] = None,
+ max_verify_tree_size: Optional[int] = None,
+ max_draft_tree_size: Optional[int] = None,
+ enable_batching: bool = False,
+ batch_timeout_ms: float = 10.0,
+ max_queue_batch_size: Optional[int] = None,
build_options: Optional["BuildOptions"] = None,
speculative_config: Optional[Any] = None,
context_cache_config: Optional[Union[ContextCacheConfig,
@@ -671,12 +677,19 @@ def __init__(
raise ValueError("engine_cache_max_size_gb must be positive")
for name, value in (("draft_top_k", draft_top_k),
("draft_step", draft_step), ("verify_tree_size",
- verify_tree_size)):
+ verify_tree_size),
+ ("max_verify_tree_size", max_verify_tree_size),
+ ("max_draft_tree_size", max_draft_tree_size),
+ ("max_queue_batch_size", max_queue_batch_size)):
if value is None:
continue
if (isinstance(value, bool) or not isinstance(value, int)
or value <= 0):
raise ValueError(f"{name} must be a positive integer")
+ if (isinstance(batch_timeout_ms, bool)
+ or not math.isfinite(batch_timeout_ms)
+ or batch_timeout_ms < 0):
+ raise ValueError("batch_timeout_ms must be non-negative")
self._model_id = _derive_model_id(model)
self._draft_top_k = draft_top_k or _DEFAULT_DRAFT_TOP_K
@@ -704,6 +717,7 @@ def __init__(
self._prev_ctx_admitted_sequences = 0
self._closed = False
self._runtime = None
+ self._batcher: Optional[RequestBatcher] = None
from .engine_build import BuildOptions, cache_root, prepare_model
@@ -711,6 +725,8 @@ def __init__(
max_input_len=max_input_len,
max_batch_size=max_batch_size,
max_kv_cache_capacity=max_kv_cache_capacity,
+ max_verify_tree_size=max_verify_tree_size,
+ max_draft_tree_size=max_draft_tree_size,
)
spec_method = options.spec_type
num_speculative_tokens = None
@@ -758,6 +774,9 @@ def __init__(
self._init_from_bundle(prepared.bundle_dir)
self._load_runtime()
+ if enable_batching:
+ self._batcher = self._make_batcher(batch_timeout_ms,
+ max_queue_batch_size)
# ------------------------------------------------------------------
# Initialization
@@ -1136,7 +1155,7 @@ def _complete_prepared_request(
tool_parser: str = "auto",
reasoning_parser: str = "none",
) -> CompletionOutput:
- response = self._handle_request(request)
+ response = self._run_generation(request)
text = response.output_texts[0] if response.output_texts else ""
token_ids = response.output_ids[0] if response.output_ids else []
prompt_tokens = (response.prompt_token_counts[0]
@@ -1223,6 +1242,38 @@ def _log_context_reuse_metrics(self) -> None:
int(cc.hybrid_restores),
)
+ def _make_batcher(self, batch_timeout_ms: float,
+ max_queue_batch_size: Optional[int]) -> RequestBatcher:
+ try:
+ video_singleton = self._video_model_family() == "nemotron"
+ except Exception:
+ video_singleton = False
+ batcher = RequestBatcher(
+ runtime_handler=self._handle_request,
+ max_batch_size=resolve_batch_size(self._max_batch_size,
+ max_queue_batch_size),
+ timeout_ms=batch_timeout_ms,
+ video_requires_singleton=video_singleton,
+ )
+ logger.info(
+ "Request batching enabled (max_batch_size=%d, timeout_ms=%.1f)",
+ batcher.max_batch_size, batcher.timeout_ms)
+ return batcher
+
+ @property
+ def batcher(self) -> Optional[RequestBatcher]:
+ """Active request batcher, or None when requests run one at a time."""
+ return getattr(self, "_batcher", None)
+
+ def _run_generation(self, request):
+ """Complete one non-streaming request, merged with concurrent ones
+ when batching is enabled."""
+ batcher = self.batcher
+ if batcher is not None and not getattr(request, "stream_channels",
+ None):
+ return batcher.submit(request)
+ return self._handle_request(request)
+
def _handle_request(self, request):
"""Serialized entry to the C++ runtime."""
with self._infer_guard():
@@ -1240,6 +1291,12 @@ def close(self) -> None:
with self._close_lock:
if self._closed:
return
+ # Drain queued batches before the runtime goes away; the worker
+ # routes them through _handle_request, which takes the same locks.
+ batcher = self.batcher
+ if batcher is not None:
+ batcher.close()
+ self._batcher = None
self._closed = True
with self._admission_sem:
with self._infer_lock:
diff --git a/experimental/server/runtime/engine_build.py b/experimental/server/runtime/engine_build.py
index 3bc177994..b12f80afe 100644
--- a/experimental/server/runtime/engine_build.py
+++ b/experimental/server/runtime/engine_build.py
@@ -48,6 +48,8 @@ class BuildOptions:
max_input_len: Optional[int] = None
max_kv_cache_capacity: Optional[int] = None
max_batch_size: Optional[int] = None
+ max_verify_tree_size: Optional[int] = None
+ max_draft_tree_size: Optional[int] = None
max_image_tokens: Optional[int] = None
max_image_tokens_per_image: Optional[int] = None
tp_size: Optional[int] = None
@@ -85,6 +87,8 @@ def to_argv(self, model_dir: str, bundle_dir: str) -> List[str]:
("--max-input-len", self.max_input_len),
("--max-kv-cache-capacity", self.max_kv_cache_capacity),
("--max-batch-size", self.max_batch_size),
+ ("--max-verify-tree-size", self.max_verify_tree_size),
+ ("--max-draft-tree-size", self.max_draft_tree_size),
("--max-image-tokens", self.max_image_tokens),
("--max-image-tokens-per-image", self.max_image_tokens_per_image),
("--tp-size", self.tp_size),
diff --git a/experimental/server/runtime/engine_client.py b/experimental/server/runtime/engine_client.py
index 68ae7f7d3..dd8a86fe8 100644
--- a/experimental/server/runtime/engine_client.py
+++ b/experimental/server/runtime/engine_client.py
@@ -49,10 +49,21 @@ class EngineCapabilities:
class _AdmissionController:
- """Bounded queue for the runtime's single generation slot."""
+ """Bounded queue in front of the runtime's generation slots.
- def __init__(self, max_queued_requests: int, timeout: float) -> None:
- self._semaphore = asyncio.Semaphore(1)
+ One slot unless request batching is enabled, in which case up to the
+ micro-batch size of requests may hold leases at once so their rows can
+ share a runtime call.
+ """
+
+ def __init__(self,
+ max_queued_requests: int,
+ timeout: float,
+ max_active: int = 1) -> None:
+ if max_active < 1:
+ raise ValueError("max_active must be positive")
+ self._max_active = max_active
+ self._semaphore = asyncio.Semaphore(max_active)
self._max_queued = max_queued_requests
self._timeout = timeout
self._active = 0
@@ -73,7 +84,8 @@ async def reserve(self) -> "_AdmissionLease":
acquired = False
if self._closing:
raise ServerUnavailableError()
- if self._active + self._waiting >= self._max_queued + 1:
+ if (self._active + self._waiting
+ >= self._max_queued + self._max_active):
raise ServerOverloadedError()
self._waiting += 1
@@ -93,7 +105,7 @@ async def reserve(self) -> "_AdmissionLease":
self._semaphore.release()
acquired = False
raise ServerUnavailableError()
- self._active = 1
+ self._active += 1
return _AdmissionLease(self)
except BaseException:
if acquired:
@@ -101,7 +113,7 @@ async def reserve(self) -> "_AdmissionLease":
raise
def release(self) -> None:
- self._active = 0
+ self._active -= 1
self._semaphore.release()
async def close(self) -> None:
@@ -110,7 +122,8 @@ async def close(self) -> None:
if self._drained:
return
self._closing = True
- await self._semaphore.acquire()
+ for _ in range(self._max_active):
+ await self._semaphore.acquire()
self._drained = True
@@ -147,6 +160,11 @@ def _read_json(path: str) -> Dict[str, Any]:
return {}
+def _batch_size_for(llm: Union[LLM, TTS]) -> int:
+ batcher = getattr(llm, "batcher", None)
+ return batcher.max_batch_size if batcher is not None else 1
+
+
def _capabilities_for(llm: Union[LLM, TTS]) -> EngineCapabilities:
bundle_dir = llm.bundle_dir
layout = llm.bundle_layout
@@ -190,7 +208,7 @@ def _capabilities_for(llm: Union[LLM, TTS]) -> EngineCapabilities:
builder.get("max_input_len"), int) else None,
max_batch_size=builder.get("max_batch_size") if isinstance(
builder.get("max_batch_size"), int) else None,
- max_num_seqs=1,
+ max_num_seqs=_batch_size_for(llm),
kv_cache_dtype=str(config.get("kv_cache_dtype", "unknown")),
speculative_decoding=llm.has_draft_model,
speculative_method=str(config.get("spec_decode_type", "none")),
@@ -260,6 +278,7 @@ def __init__(self,
self._admission = _AdmissionController(
self._api_config.max_queued_requests,
self._api_config.queue_timeout,
+ max_active=_batch_size_for(llm),
)
self._capabilities = _capabilities_for(llm)
self._close_lock = asyncio.Lock()
@@ -286,6 +305,15 @@ def active_requests(self) -> int:
def queued_requests(self) -> int:
return self._admission.waiting
+ @property
+ def batching(self) -> Dict[str, Any]:
+ batcher = getattr(self._llm, "batcher", None)
+ return {
+ "enabled": batcher is not None,
+ "max_batch_size": batcher.max_batch_size if batcher else 1,
+ "timeout_ms": batcher.timeout_ms if batcher else 0.0,
+ }
+
async def close(self) -> None:
"""Drain request ownership, then release the model runtime once."""
async with self._close_lock:
diff --git a/tests/python-unittests/test_direct_builder_chat_template.py b/tests/python-unittests/test_direct_builder_chat_template.py
new file mode 100644
index 000000000..56aaa2b00
--- /dev/null
+++ b/tests/python-unittests/test_direct_builder_chat_template.py
@@ -0,0 +1,94 @@
+# 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.
+"""Direct-builder chat template extraction for thinking-aware templates."""
+
+import json
+
+import pytest
+
+transformers = pytest.importorskip("transformers")
+
+from experimental.builder.core.artifacts.chat_template import \
+ process_chat_template
+
+_REASONING = ("Reasoning effort is set to xhigh. Please think carefully "
+ "through the task.")
+
+
+class _Qwen38LikeTokenizer:
+ """Mimics Qwen3.8: an undefined ``enable_thinking`` means thinking on,
+ and thinking mode injects reasoning instructions into the system block."""
+
+ chat_template = "fake"
+ bos_token = None
+
+ def apply_chat_template(self, messages, **kwargs):
+ del kwargs["tokenize"]
+ add_generation_prompt = kwargs.get("add_generation_prompt", False)
+ thinking = kwargs.get("enable_thinking", None) is not False
+ instructions = _REASONING if thinking else ""
+
+ output = ""
+ rest = messages
+ if messages[0]["role"] == "system":
+ content = messages[0]["content"]
+ if instructions:
+ content = instructions + "\n\n" + content
+ output += f"<|im_start|>system\n{content}<|im_end|>\n"
+ rest = messages[1:]
+ elif instructions:
+ output += f"<|im_start|>system\n{instructions}<|im_end|>\n"
+ for message in rest:
+ output += f"<|im_start|>{message['role']}\n"
+ if message["role"] == "assistant":
+ output += "\n\n\n\n"
+ output += f"{message['content']}<|im_end|>\n"
+ if add_generation_prompt:
+ output += "<|im_start|>assistant\n"
+ output += "\n" if thinking else "\n\n\n\n"
+ return output
+
+
+def test_thinking_aware_template_yields_clean_role_prefixes(
+ monkeypatch, tmp_path):
+ model_dir = tmp_path / "model"
+ out_dir = tmp_path / "out"
+ model_dir.mkdir()
+ (model_dir / "config.json").write_text(
+ json.dumps({"model_type": "qwen3_5"}))
+ monkeypatch.setattr(transformers.AutoTokenizer, "from_pretrained",
+ lambda *args, **kwargs: _Qwen38LikeTokenizer())
+
+ process_chat_template(str(model_dir), str(out_dir))
+
+ data = json.loads((out_dir / "processed_chat_template.json").read_text())
+ assert data["roles"]["system"] == {
+ "prefix": "<|im_start|>system\n",
+ "suffix": "<|im_end|>\n",
+ }
+ assert data["roles"]["user"] == {
+ "prefix": "<|im_start|>user\n",
+ "suffix": "<|im_end|>\n",
+ }
+ assert data["roles"]["assistant"]["prefix"] == (
+ "<|im_start|>assistant\n\n\n\n\n")
+ assert data["generation_prompt"] == (
+ "<|im_start|>assistant\n\n\n\n\n")
+ assert data[
+ "generation_prompt_thinking"] == "<|im_start|>assistant\n\n"
+ assert data["default_system_prompt"] == ""
+ serialized = json.dumps(data)
+ assert "placeholder" not in serialized
+ assert _REASONING not in serialized
diff --git a/tests/python-unittests/test_direct_builder_draft_quant.py b/tests/python-unittests/test_direct_builder_draft_quant.py
new file mode 100644
index 000000000..ed435b0f4
--- /dev/null
+++ b/tests/python-unittests/test_direct_builder_draft_quant.py
@@ -0,0 +1,102 @@
+# 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.
+"""Draft-namespace quantization lookup in the direct builder."""
+
+import numpy as np
+import pytest
+
+safetensors_numpy = pytest.importorskip("safetensors.numpy")
+
+from experimental.builder.core import quantization
+from experimental.builder.core.weights import Weights
+from experimental.builder.models.qwen3_5 import weights as qwen3_5_weights
+
+_BASE = "model.language_model.layers.0.mlp."
+
+
+def _nvfp4(prefix):
+ return {
+ prefix + "weight": np.zeros((8, 4), dtype=np.uint8),
+ prefix + "weight_scale": np.ones((8, 1), dtype=np.float16),
+ prefix + "weight_scale_2": np.ones((), dtype=np.float32),
+ prefix + "input_scale": np.ones((), dtype=np.float32),
+ }
+
+
+@pytest.fixture
+def checkpoint(tmp_path):
+ tensors = {}
+ tensors.update(_nvfp4(_BASE + "gate_proj."))
+ tensors.update(_nvfp4(_BASE + "up_proj."))
+ tensors.update(_nvfp4("lm_head."))
+ # Unquantized MTP draft MLP: plain weight, no quantization sidecar.
+ tensors["mtp.layers.0.mlp.gate_proj.weight"] = np.zeros((8, 8),
+ dtype=np.float16)
+ # Quantized MTP draft projection keeps its declared precision.
+ tensors.update(_nvfp4("mtp.layers.0.mlp.up_proj."))
+ safetensors_numpy.save_file(tensors, str(tmp_path / "model.safetensors"))
+ return str(tmp_path)
+
+
+_QUANT = quantization.QuantConfig(
+ quant_type=quantization.QUANT_NVFP4,
+ group_size=16,
+ layer_overrides={
+ "layers.0.mlp.gate_proj": quantization.QUANT_NVFP4,
+ "layers.0.mlp.up_proj": quantization.QUANT_NVFP4,
+ "lm_head": quantization.QUANT_NVFP4,
+ },
+ is_mixed_precision=True,
+)
+
+
+def test_unquantized_draft_projection_is_fp16_despite_base_override(
+ checkpoint):
+ weights = Weights(checkpoint,
+ quant=_QUANT,
+ conversion=qwen3_5_weights,
+ spec_type="mtp",
+ spec_role="draft")
+ try:
+ assert weights.checkpoint_key("layers.0.mlp.gate_proj.weight") == (
+ "mtp.layers.0.mlp."
+ "gate_proj.weight")
+ assert weights.module_quant_type(
+ "layers.0.mlp.gate_proj") == quantization.QUANT_FP16
+ # Sidecars stay pinned to the draft namespace: the base layer's
+ # scales must not be borrowed, so the module loads as a plain weight.
+ assert not weights.has("layers.0.mlp.gate_proj.weight_scale")
+ assert not weights.is_nvfp4("layers.0.mlp.gate_proj")
+ weight, bias = weights.linear_fp16("layers.0.mlp.gate_proj")
+ assert weight.shape == (8, 8) and bias is None
+ assert weights.module_quant_type(
+ "layers.0.mlp.up_proj") == quantization.QUANT_NVFP4
+ assert weights.is_nvfp4("layers.0.mlp.up_proj")
+ # The draft shares the base lm_head, which stays quantized.
+ assert weights.module_quant_type("lm_head") == quantization.QUANT_NVFP4
+ assert weights.is_nvfp4("lm_head")
+ finally:
+ weights.close()
+
+
+def test_base_projection_keeps_its_override(checkpoint):
+ weights = Weights(checkpoint, quant=_QUANT, conversion=qwen3_5_weights)
+ try:
+ assert weights.module_quant_type(
+ "layers.0.mlp.gate_proj") == quantization.QUANT_NVFP4
+ assert weights.module_quant_type(
+ "layers.0.mlp.up_proj") == quantization.QUANT_NVFP4
+ finally:
+ weights.close()
diff --git a/tests/python-unittests/test_server_batching.py b/tests/python-unittests/test_server_batching.py
new file mode 100644
index 000000000..682256c89
--- /dev/null
+++ b/tests/python-unittests/test_server_batching.py
@@ -0,0 +1,319 @@
+# 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 non-streaming request batching in the experimental server."""
+
+import asyncio
+import threading
+import time
+from concurrent.futures import ThreadPoolExecutor
+
+import pytest
+
+from experimental.server.api.errors import ServerOverloadedError
+from experimental.server.config import parse_server_config
+from experimental.server.runtime.batching import (BATCH_COMPATIBILITY_FIELDS,
+ BatcherOverflow,
+ RequestBatcher,
+ resolve_batch_size)
+from experimental.server.runtime.engine import LLM
+from experimental.server.runtime.engine_client import _AdmissionController
+
+
+class _Row:
+
+ def __init__(self, text: str) -> None:
+ self.text = text
+ self.image_buffers = []
+
+
+class _Request:
+ """Duck-typed stand-in for the pybind ``LLMGenerationRequest``."""
+
+ def __init__(self, *rows: str, **settings) -> None:
+ for field in BATCH_COMPATIBILITY_FIELDS:
+ setattr(self, field, settings.get(field, 0))
+ self.requests = [_Row(text) for text in rows]
+ self.stream_channels = []
+
+
+class _Response:
+
+ def __init__(self, rows) -> None:
+ self.output_texts = [f"echo:{row.text}" for row in rows]
+ self.output_ids = [[len(row.text)] for row in rows]
+ self.finish_reasons = ["stop" for _ in rows]
+ self.logprobs = [None for _ in rows]
+ self.prompt_token_counts = [1 for _ in rows]
+
+
+class _Handler:
+
+ def __init__(self, delay: float = 0.0, gate=None) -> None:
+ self.calls = []
+ self.lock = threading.Lock()
+ self.delay = delay
+ self.gate = gate
+
+ def __call__(self, request):
+ with self.lock:
+ self.calls.append(list(request.requests))
+ if self.gate is not None:
+ self.gate.wait(5.0)
+ if self.delay:
+ time.sleep(self.delay)
+ return _Response(request.requests)
+
+
+def _submit_all(batcher, requests):
+ with ThreadPoolExecutor(max_workers=len(requests)) as pool:
+ return list(pool.map(batcher.submit, requests))
+
+
+def test_resolve_batch_size_caps_to_engine():
+ assert resolve_batch_size(4, None) == 4
+ assert resolve_batch_size(4, 2) == 2
+ assert resolve_batch_size(4, 8) == 4
+ assert resolve_batch_size(0, None) == 1
+
+
+def test_compatible_requests_share_one_runtime_call():
+ gate = threading.Event()
+ handler = _Handler(gate=gate)
+ batcher = RequestBatcher(handler, max_batch_size=4, timeout_ms=200)
+ try:
+ requests = [_Request("a"), _Request("b"), _Request("c")]
+ with ThreadPoolExecutor(max_workers=3) as pool:
+ futures = [pool.submit(batcher.submit, r) for r in requests]
+ # Let all three enqueue before the worker's timeout elapses.
+ deadline = time.monotonic() + 2.0
+ while batcher.pending + len(handler.calls) < 1:
+ assert time.monotonic() < deadline
+ time.sleep(0.005)
+ gate.set()
+ results = [f.result(timeout=5.0) for f in futures]
+ finally:
+ batcher.close()
+
+ assert len(handler.calls) == 1
+ assert [row.text for row in handler.calls[0]] == ["a", "b", "c"]
+ assert [r.output_texts for r in results] == [["echo:a"], ["echo:b"],
+ ["echo:c"]]
+ assert [r.output_ids for r in results] == [[[1]], [[1]], [[1]]]
+ assert all(r.finish_reasons == ["stop"] for r in results)
+ assert all(r.prompt_token_counts == [1] for r in results)
+
+
+def test_incompatible_settings_run_in_separate_calls():
+ handler = _Handler(delay=0.02)
+ batcher = RequestBatcher(handler, max_batch_size=4, timeout_ms=100)
+ try:
+ results = _submit_all(batcher, [
+ _Request("a", temperature=0.0),
+ _Request("b", temperature=0.7),
+ _Request("c", temperature=0.0),
+ ])
+ finally:
+ batcher.close()
+
+ assert len(handler.calls) == 2
+ batched = sorted(len(call) for call in handler.calls)
+ assert batched == [1, 2]
+ assert sorted(r.output_texts[0]
+ for r in results) == ["echo:a", "echo:b", "echo:c"]
+
+
+def test_batch_size_bounds_rows_per_call():
+ handler = _Handler(delay=0.02)
+ batcher = RequestBatcher(handler, max_batch_size=2, timeout_ms=100)
+ try:
+ _submit_all(batcher, [_Request(str(i)) for i in range(4)])
+ finally:
+ batcher.close()
+
+ assert all(len(call) <= 2 for call in handler.calls)
+ assert sum(len(call) for call in handler.calls) == 4
+
+
+def test_multi_row_requests_are_sliced_back_correctly():
+ handler = _Handler(delay=0.02)
+ batcher = RequestBatcher(handler, max_batch_size=8, timeout_ms=100)
+ try:
+ results = _submit_all(
+ batcher,
+ [_Request("a", "b"),
+ _Request("c"),
+ _Request("d", "e", "f")])
+ finally:
+ batcher.close()
+
+ texts = sorted(tuple(r.output_texts) for r in results)
+ assert texts == [("echo:a", "echo:b"), ("echo:c", ),
+ ("echo:d", "echo:e", "echo:f")]
+
+
+def test_pending_queue_overflow_is_reported():
+ gate = threading.Event()
+ handler = _Handler(gate=gate)
+ batcher = RequestBatcher(handler,
+ max_batch_size=1,
+ timeout_ms=0,
+ max_pending=1)
+ try:
+ with ThreadPoolExecutor(max_workers=2) as pool:
+ running = pool.submit(batcher.submit, _Request("a"))
+ deadline = time.monotonic() + 2.0
+ while not handler.calls:
+ assert time.monotonic() < deadline
+ time.sleep(0.005)
+ queued = pool.submit(batcher.submit, _Request("b"))
+ deadline = time.monotonic() + 2.0
+ while batcher.pending < 1:
+ assert time.monotonic() < deadline
+ time.sleep(0.005)
+ with pytest.raises(BatcherOverflow):
+ batcher.submit(_Request("c"))
+ gate.set()
+ assert running.result(5.0).output_texts == ["echo:a"]
+ assert queued.result(5.0).output_texts == ["echo:b"]
+ finally:
+ batcher.close()
+
+
+def test_runtime_failure_propagates_to_every_member():
+
+ def failing(_request):
+ raise RuntimeError("boom")
+
+ batcher = RequestBatcher(failing, max_batch_size=4, timeout_ms=50)
+ try:
+ with ThreadPoolExecutor(max_workers=2) as pool:
+ futures = [
+ pool.submit(batcher.submit, _Request("a")),
+ pool.submit(batcher.submit, _Request("b")),
+ ]
+ for future in futures:
+ with pytest.raises(RuntimeError, match="boom"):
+ future.result(5.0)
+ finally:
+ batcher.close()
+ # The worker survives a failed batch and keeps serving.
+ batcher2 = RequestBatcher(_Handler(), max_batch_size=1, timeout_ms=0)
+ try:
+ assert batcher2.submit(_Request("z")).output_texts == ["echo:z"]
+ finally:
+ batcher2.close()
+ with pytest.raises(RuntimeError, match="closed"):
+ batcher2.submit(_Request("late"))
+
+
+class _Runtime:
+
+ def __init__(self) -> None:
+ self.calls = []
+
+ def handle_request(self, request):
+ self.calls.append(list(request.requests))
+ time.sleep(0.02)
+ return _Response(request.requests)
+
+
+def _bare_llm(runtime):
+ llm = LLM.__new__(LLM)
+ llm._runtime = runtime
+ llm._admission_sem = threading.Semaphore(1)
+ llm._infer_lock = threading.Lock()
+ llm._close_lock = threading.Lock()
+ llm._closed = False
+ llm._batcher = None
+ return llm
+
+
+def test_llm_runs_generation_through_batcher_and_closes_it():
+ runtime = _Runtime()
+ llm = _bare_llm(runtime)
+ llm._batcher = RequestBatcher(llm._handle_request,
+ max_batch_size=4,
+ timeout_ms=100)
+
+ with ThreadPoolExecutor(max_workers=2) as pool:
+ futures = [
+ pool.submit(llm._run_generation, _Request("a")),
+ pool.submit(llm._run_generation, _Request("b")),
+ ]
+ results = [f.result(5.0) for f in futures]
+
+ assert len(runtime.calls) == 1
+ assert sorted(r.output_texts[0] for r in results) == ["echo:a", "echo:b"]
+
+ # Streaming requests bypass the batcher even when it is enabled.
+ streaming = _Request("s")
+ streaming.stream_channels = [object()]
+ assert llm._run_generation(streaming).output_texts == ["echo:s"]
+ assert len(runtime.calls) == 2
+
+ llm.close()
+ assert llm.batcher is None
+ assert llm._runtime is None
+
+
+def test_llm_without_batcher_calls_runtime_directly():
+ runtime = _Runtime()
+ llm = _bare_llm(runtime)
+ assert llm._run_generation(_Request("a")).output_texts == ["echo:a"]
+ assert len(runtime.calls) == 1
+
+
+def test_admission_allows_batch_size_concurrent_leases():
+
+ async def exercise():
+ admission = _AdmissionController(max_queued_requests=1,
+ timeout=1.0,
+ max_active=2)
+ first = await admission.reserve()
+ second = await admission.reserve()
+ assert admission.active == 2
+ waiting = asyncio.create_task(admission.reserve())
+ while admission.waiting != 1:
+ await asyncio.sleep(0)
+ with pytest.raises(ServerOverloadedError):
+ await admission.reserve()
+ first.release()
+ third = await waiting
+ assert admission.active == 2
+ second.release()
+ third.release()
+ assert admission.active == 0
+ closing = asyncio.create_task(admission.close())
+ await closing
+ with pytest.raises(Exception):
+ await admission.reserve()
+
+ asyncio.run(exercise())
+
+
+def test_server_config_parses_batching_flags():
+ config = parse_server_config([
+ "/model", "--enable-batching", "--batch-timeout-ms", "5",
+ "--max-queue-batch-size", "2", "--max-batch-size", "4"
+ ])
+ kwargs = config.model.llm_kwargs()
+ assert kwargs["enable_batching"] is True
+ assert kwargs["batch_timeout_ms"] == 5.0
+ assert kwargs["max_queue_batch_size"] == 2
+ assert kwargs["max_batch_size"] == 4
+
+ default = parse_server_config(["/model"]).model.llm_kwargs()
+ assert default["enable_batching"] is False
+ assert default["max_queue_batch_size"] is None