# Copyright 2025 - Oumi
#
# 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.
from __future__ import annotations
import copy
import dataclasses
import functools
import json
import math
import os
import re
import shutil
import subprocess
import warnings
from pathlib import Path
from types import SimpleNamespace
from typing import TYPE_CHECKING, cast, get_args
import torch
from typing_extensions import override
from oumi.builders import build_tokenizer
from oumi.core.configs import GenerationParams, InferenceConfig, ModelParams
from oumi.core.inference import BaseInferenceEngine
from oumi.core.types.conversation import Conversation, FinishReason, Message, Role
from oumi.core.types.tool_call import ToolCall, ToolDefinition
from oumi.utils.conversation_utils import (
create_list_of_message_json_dicts,
split_assistant_content_and_tool_calls,
)
from oumi.utils.logging import logger
from oumi.utils.model_caching import get_local_filepath_for_gguf
from oumi.utils.peft_utils import get_lora_rank
from oumi.utils.vllm_utils.gemma4_compat import ACTIVATION_ENV_VAR
if TYPE_CHECKING:
from vllm.outputs import RequestOutput # pyright: ignore[reportMissingImports]
try:
import vllm # pyright: ignore[reportMissingImports]
try:
from vllm.config import ( # pyright: ignore[reportMissingImports]
ModelDType, # pyright: ignore[reportAttributeAccessIssue]
)
except ImportError:
# For compatibility with newer vLLM versions
ModelDType = str # type: ignore
from vllm.entrypoints.chat_utils import ( # pyright: ignore[reportMissingImports]
ChatCompletionMessageParam,
)
from vllm.lora.request import LoRARequest # pyright: ignore[reportMissingImports]
from vllm.model_executor.layers.quantization import ( # pyright: ignore[reportMissingImports]
QuantizationMethods,
)
from vllm.sampling_params import ( # pyright: ignore[reportMissingImports]
SamplingParams,
)
from oumi.utils.packaging import is_vllm_v0_12_or_later
_VLLM_V0_12 = is_vllm_v0_12_or_later()
if _VLLM_V0_12:
from vllm.sampling_params import ( # pyright: ignore[reportMissingImports]
StructuredOutputsParams as VLLMGuidedDecodingParams, # pyright: ignore[reportAttributeAccessIssue]
)
else:
from vllm.sampling_params import ( # pyright: ignore[reportMissingImports]
GuidedDecodingParams as VLLMGuidedDecodingParams, # pyright: ignore[reportAttributeAccessIssue]
)
# Tool-call parsers ship at vllm.tool_parsers in 0.14+; earlier versions
# exposed them at vllm.entrypoints.openai.tool_parsers. Guard both paths.
try:
from vllm.tool_parsers import ( # pyright: ignore[reportMissingImports]
ToolParserManager,
)
_VLLM_TOOL_PARSERS_AVAILABLE = True
except ImportError:
try:
from vllm.entrypoints.openai.tool_parsers import ( # pyright: ignore[reportMissingImports]
ToolParserManager,
)
_VLLM_TOOL_PARSERS_AVAILABLE = True
except ImportError:
ToolParserManager = None # type: ignore[assignment]
_VLLM_TOOL_PARSERS_AVAILABLE = False
# Parser-engine tool parsers (vLLM 0.24+, e.g. Gemma 4) recover argument
# types from the request's tool schemas, which they only recognize as
# `ChatCompletionToolsParam` objects; plain dicts are skipped.
try:
from vllm.entrypoints.openai.chat_completion.protocol import ( # pyright: ignore[reportMissingImports]
ChatCompletionToolsParam,
)
from vllm.parser.engine.adapters import ( # pyright: ignore[reportMissingImports]
ParserEngineToolAdapter,
)
except ImportError:
ChatCompletionToolsParam = None # type: ignore[assignment]
ParserEngineToolAdapter = None # type: ignore[assignment]
except ModuleNotFoundError:
vllm = None
_VLLM_V0_12 = False
ToolParserManager = None # type: ignore[assignment]
_VLLM_TOOL_PARSERS_AVAILABLE = False
ChatCompletionToolsParam = None # type: ignore[assignment]
ParserEngineToolAdapter = None # type: ignore[assignment]
@dataclasses.dataclass(frozen=True)
class _ToolParserCapabilities:
"""Oumi-side behavior that is not declared by vLLM's parser interface."""
requires_raw_output: bool = False
requires_split_tool_narration: bool = False
_DEFAULT_TOOL_PARSER_CAPABILITIES = _ToolParserCapabilities()
_TOOL_PARSER_CAPABILITIES = {
# Gemma 4 marks native tool calls with special tokens. Its parser must see
# those tokens, even though they should remain hidden from displayed text.
# Its chat template also reorders narration and tool calls when they share
# one assistant message, so replay must split them while preserving order.
"gemma4": _ToolParserCapabilities(
requires_raw_output=True,
requires_split_tool_narration=True,
),
}
def _get_tool_parser_capabilities(
parser_name: str | None,
) -> _ToolParserCapabilities:
"""Return parser input requirements while preserving existing defaults."""
if parser_name is None:
return _DEFAULT_TOOL_PARSER_CAPABILITIES
return _TOOL_PARSER_CAPABILITIES.get(
parser_name,
_DEFAULT_TOOL_PARSER_CAPABILITIES,
)
def _build_parser_request_tools(
tool_parser: object, tools: list[ToolDefinition] | None
) -> list:
"""Return the conversation's tools in the form the tool parser expects.
Parser-engine adapters coerce each parsed argument to its schema type and
only find that schema on typed tool objects, so they get
`ChatCompletionToolsParam`s. Other parsers keep receiving plain dicts.
"""
tool_dicts = [t.model_dump(mode="json", exclude_none=True) for t in tools or []]
if (
ParserEngineToolAdapter is None
or ChatCompletionToolsParam is None
or not isinstance(tool_parser, ParserEngineToolAdapter)
):
return tool_dicts
try:
return [ChatCompletionToolsParam.model_validate(t) for t in tool_dicts]
except Exception:
logger.warning(
"Could not convert tools for %s; tool-call arguments may not be "
"coerced to their schema types.",
type(tool_parser).__name__,
exc_info=True,
)
return tool_dicts
def _parse_nvcc_release_version(nvcc_version_output: str) -> tuple[int, int] | None:
"""Parse the ``release X.Y`` field of ``nvcc --version`` output into (X, Y)."""
match = re.search(r"release (\d+)\.(\d+)", nvcc_version_output)
if match is None:
return None
return int(match.group(1)), int(match.group(2))
def _vllm_accepts_additional_config() -> bool:
"""Whether this vLLM build's ``EngineArgs`` accepts ``additional_config``.
Only newer vLLM accepts it — and that is also the only range with GDN
support, so on older builds there is nothing to fix.
"""
try:
from vllm.engine.arg_utils import ( # pyright: ignore[reportMissingImports]
EngineArgs,
)
return "additional_config" in {f.name for f in dataclasses.fields(EngineArgs)}
except Exception:
return False
@functools.lru_cache(maxsize=1)
def _cuda_toolkit_below_12_6() -> bool:
"""Whether the system CUDA toolkit is older than 12.6 (or absent).
flashinfer's ``gdn_prefill_sm90`` kernel JIT-builds against the system
toolkit and needs the ``cuda::ptx`` tensormap intrinsics added in 12.6.
"""
nvcc = shutil.which("nvcc")
if nvcc is None:
for cuda_home in (os.environ.get("CUDA_HOME"), "/usr/local/cuda"):
candidate = Path(cuda_home) / "bin" / "nvcc" if cuda_home else None
if candidate and candidate.is_file():
nvcc = str(candidate)
break
if nvcc is None:
return True # no toolkit at all → the flashinfer JIT can't build either
try:
output = subprocess.run(
[nvcc, "--version"], capture_output=True, text=True, timeout=10
).stdout
except (OSError, subprocess.SubprocessError):
return True
version = _parse_nvcc_release_version(output)
if version is None:
return False
return version < (12, 6)
# Model families using a Gated Delta Net attention architecture, whose
# flashinfer gdn_prefill_sm90 kernel JIT-builds against the CUDA toolkit.
_GDN_MODEL_TYPES = frozenset({"qwen3_5", "qwen3_6", "qwen3_next"})
@functools.cache
def _model_uses_gdn(model_name: str, trust_remote_code: bool) -> bool:
"""Best-effort check whether the model uses a GDN attention architecture.
Reads the HF config (resolves both Hub ids and local fine-tuned dirs).
Returns False if the config can't be read.
"""
try:
from transformers import AutoConfig
config = AutoConfig.from_pretrained(
model_name, trust_remote_code=trust_remote_code
)
return getattr(config, "model_type", None) in _GDN_MODEL_TYPES
except Exception:
return False
def _should_force_triton_gdn_backend(model_name: str, trust_remote_code: bool) -> bool:
"""Whether to force vLLM's Triton GDN-prefill backend for this model.
GDN models (e.g. Qwen3.5) JIT-compile a flashinfer ``gdn_prefill_sm90``
kernel that needs the system CUDA toolkit >= 12.6 (``cuda::ptx`` tensormap
intrinsics). On older toolkits the build fails and the engine hangs, so
fall back to the Triton/FLA GDN backend. No-op for non-GDN models and for
toolkits >= 12.6. Mirrors the intent of vllm-project/vllm#37507.
"""
return (
_vllm_accepts_additional_config()
and _cuda_toolkit_below_12_6()
and _model_uses_gdn(model_name, trust_remote_code)
)
LORA_ADAPTER_NAME = "oumi_lora_adapter"
[docs]
class VLLMInferenceEngine(BaseInferenceEngine):
"""Engine for running vLLM inference locally."""
def __init__(
self,
model_params: ModelParams,
*,
generation_params: GenerationParams | None = None,
tensor_parallel_size: int = -1,
quantization: str | None = None,
enable_prefix_caching: bool = True,
gpu_memory_utilization: float = 0.9,
enforce_eager: bool = True,
max_num_seqs: int | None = None,
tool_call_parser: str | None = None,
):
"""Initializes the inference Engine.
Args:
model_params: The model parameters to use for inference.
generation_params: The generation parameters to use for inference.
tensor_parallel_size: The number of tensor parallel processes to use.
If set to -1, we will use all the available GPUs.
quantization: The quantization method to use for inference.
enable_prefix_caching: Whether to enable prefix caching.
gpu_memory_utilization: The fraction of available GPU memory the model's
executor will use. It can range from 0 to 1. Defaults to 0.9, i.e.,
(90%) memory utilization.
enforce_eager: Whether to enforce eager execution. Defaults to True.
If False, will use eager mode and CUDA graph in hybrid mode.
max_num_seqs: Maximum number of sequences per iteration.
tool_call_parser: Optional name of a vLLM tool-call parser
(e.g. ``"hermes"``, ``"qwen3_xml"``, ``"llama4_pythonic"``,
``"mistral"``). When set, the engine parses the model's
output text into ``Message.tool_calls`` and sets
``finish_reason`` to ``tool_calls``. If left ``None``, falls
back to ``model_params.tool_call_parser``. Tied to vLLM
internals; available parsers depend on the installed
vLLM version.
"""
super().__init__(model_params=model_params, generation_params=generation_params)
if not vllm:
raise RuntimeError(
"vLLM is not installed. "
"Please install the GPU dependencies for this package."
)
final_vllm_kwargs = self.build_engine_kwargs(
model_params,
tensor_parallel_size=tensor_parallel_size,
quantization=quantization,
enable_prefix_caching=enable_prefix_caching,
gpu_memory_utilization=gpu_memory_utilization,
enforce_eager=enforce_eager,
max_num_seqs=max_num_seqs,
)
if final_vllm_kwargs["model"] != model_params.model_name:
# GGUF: build the tokenizer from the locally cached model file.
model_params = copy.deepcopy(model_params)
model_params.model_name = str(final_vllm_kwargs["model"])
self._lora_request = None
if model_params.adapter_model:
# ID should be unique for this adapter, but isn't enforced by vLLM.
self._lora_request = LoRARequest(
lora_name=LORA_ADAPTER_NAME,
lora_int_id=1,
lora_path=model_params.adapter_model,
)
logger.info(f"Loaded LoRA adapter: {model_params.adapter_model}")
self._tokenizer = build_tokenizer(model_params)
previous_compat_activation = os.environ.get(ACTIVATION_ENV_VAR)
os.environ[ACTIVATION_ENV_VAR] = "1"
try:
self._llm = vllm.LLM(**final_vllm_kwargs) # pyright: ignore[reportArgumentType, reportAttributeAccessIssue]
finally:
if previous_compat_activation is None:
os.environ.pop(ACTIVATION_ENV_VAR, None)
else:
os.environ[ACTIVATION_ENV_VAR] = previous_compat_activation
# Ensure the tokenizer is set properly.
# set_tokenizer() was deprecated in vLLM v0.12 and removed in v0.13; the
# tokenizer is already configured via the constructor's `tokenizer` parameter.
if not _VLLM_V0_12:
self._llm.set_tokenizer(self._tokenizer) # pyright: ignore[reportAttributeAccessIssue]
# Optional tool-call parser. Direct kwarg wins over model_params.
parser_name = tool_call_parser or model_params.tool_call_parser
self._tool_parser_name = parser_name
self._tool_parser = None
if parser_name is not None:
if not _VLLM_TOOL_PARSERS_AVAILABLE:
raise RuntimeError(
"vLLM tool parsers are not available in this vLLM version. "
"Upgrade vLLM or unset `tool_call_parser`."
)
try:
parser_cls = ToolParserManager.get_tool_parser(parser_name) # pyright: ignore[reportOptionalMemberAccess]
except KeyError as e:
raise ValueError(
f"Unknown vLLM tool_call_parser '{parser_name}'."
) from e
self._tool_parser = parser_cls(self._tokenizer) # pyright: ignore[reportArgumentType]
logger.info(f"VLLM engine will parse tool calls with '{parser_name}'.")
[docs]
@staticmethod
def build_engine_kwargs(
model_params: ModelParams,
*,
tensor_parallel_size: int = -1,
quantization: str | None = None,
enable_prefix_caching: bool = True,
gpu_memory_utilization: float = 0.9,
enforce_eager: bool = True,
max_num_seqs: int | None = None,
) -> dict[str, object]:
"""Builds the ``vllm.LLM`` kwargs the engine uses for ``model_params``.
Reads local state (GPU count, CUDA toolkit, LoRA adapter files), so call it
on the machine that runs vLLM.
Args:
model_params: The model parameters to use for inference.
tensor_parallel_size: See ``__init__``.
quantization: See ``__init__``.
enable_prefix_caching: See ``__init__``.
gpu_memory_utilization: See ``__init__``.
enforce_eager: See ``__init__``.
max_num_seqs: See ``__init__``.
Returns:
dict[str, object]: Keyword arguments for ``vllm.LLM``.
"""
if not (
math.isfinite(gpu_memory_utilization)
and gpu_memory_utilization > 0
and gpu_memory_utilization <= 1.0
):
raise ValueError(
"GPU memory utilization must be within (0, 1]. Got "
f"{gpu_memory_utilization}."
)
# Infer the `quantization` type from the model's kwargs.
if model_params.model_kwargs:
if not quantization:
# Check if quantization is BitsAndBytes.
bnb_quantization_kwargs = ["load_in_4bit", "load_in_8bit"]
for key in bnb_quantization_kwargs:
if model_params.model_kwargs.get(key):
quantization = "bitsandbytes"
break
# Check if quantization is MXFP4.
if not quantization and model_params.model_kwargs.get(
"quantization_config"
):
quant_config = model_params.model_kwargs.get("quantization_config")
if (
isinstance(quant_config, dict)
and quant_config.get("quant_method") == "mxfp4"
):
quantization = "mxfp4"
if not quantization and model_params.model_kwargs.get("filename"):
# Check if quantization is GGUF.
gguf_filename = str(model_params.model_kwargs.get("filename"))
if gguf_filename.lower().endswith(".gguf"):
quantization = "gguf"
if (
not model_params.tokenizer_name
or model_params.tokenizer_name == model_params.model_name
):
raise ValueError(
"GGUF quantization with the VLLM engine requires that you "
"explicitly set the `tokenizer_name` in `model_params`."
)
vllm_kwargs = {}
# Set the proper VLLM keys for the quantization type.
if quantization and quantization == "bitsandbytes":
vllm_kwargs["load_format"] = "bitsandbytes"
logger.info("VLLM engine loading a `bitsandbytes` quantized model.")
elif quantization and quantization == "mxfp4":
# logic may not be needed; to be cleaned up after the next vllm patch
# version release if possible
# For MXFP4, set quantization in vllm_kwargs and clear variable
# to avoid passing it twice
vllm_kwargs["quantization"] = "mxfp4"
quantization = None # Avoid double setting
logger.info("VLLM engine loading a `MXFP4` quantized model.")
elif quantization and quantization == "gguf":
# Download the GGUF file from HuggingFace to a local cache.
gguf_local_path = get_local_filepath_for_gguf(
repo_id=model_params.model_name,
filename=gguf_filename,
)
# Overwrite `model_name` with the locally cached GGUF model.
model_params = copy.deepcopy(model_params)
model_params.model_name = gguf_local_path
logger.info("VLLM engine loading a `GGUF` quantized model.")
if tensor_parallel_size <= 0:
if torch.cuda.device_count() > 1:
tensor_parallel_size = torch.cuda.device_count()
else:
tensor_parallel_size = 1
if model_params.adapter_model:
lora_rank = get_lora_rank(model_params.adapter_model)
vllm_kwargs["max_lora_rank"] = lora_rank
logger.info(f"Setting vLLM max LoRA rank to {lora_rank}")
if max_num_seqs is not None:
vllm_kwargs["max_num_seqs"] = max_num_seqs
supported_quantization_methods = list(get_args(QuantizationMethods))
if quantization and quantization not in supported_quantization_methods:
raise ValueError(
f"Unsupported quantization method: {quantization}. "
f"Supported methods are: {supported_quantization_methods}."
)
# Pass through selected vLLM kwargs from model_kwargs.
_VLLM_PASSTHROUGH_KWARGS = ("language_model_only", "hf_config_path")
if model_params.model_kwargs:
for key in _VLLM_PASSTHROUGH_KWARGS:
if key in model_params.model_kwargs:
vllm_kwargs[key] = model_params.model_kwargs[key]
final_vllm_kwargs: dict[str, object] = dict(
model=model_params.model_name,
tokenizer=model_params.tokenizer_name,
trust_remote_code=model_params.trust_remote_code,
dtype=cast(ModelDType, model_params.torch_dtype_str), # pyright: ignore[reportInvalidTypeForm]
# TODO: these params should be settable via config,
# but they don't belong to model_params
tensor_parallel_size=tensor_parallel_size,
enable_prefix_caching=enable_prefix_caching,
enable_lora=bool(model_params.adapter_model),
max_model_len=model_params.model_max_length,
gpu_memory_utilization=gpu_memory_utilization,
enforce_eager=enforce_eager,
**vllm_kwargs,
)
# Only add quantization if not already in vllm_kwargs and not None
if quantization is not None and "quantization" not in vllm_kwargs:
final_vllm_kwargs["quantization"] = quantization
# GDN-architecture models (e.g. Qwen3.5) hang when flashinfer's sm90
# GDN-prefill kernel can't JIT-build against an old (<12.6) CUDA
# toolkit; fall back to the Triton GDN backend. No-op for other models.
if _should_force_triton_gdn_backend(
model_params.model_name, model_params.trust_remote_code
):
final_vllm_kwargs["additional_config"] = { # pyright: ignore[reportArgumentType]
"gdn_prefill_backend": "triton"
}
logger.warning(
"System CUDA toolkit < 12.6: forcing vLLM Triton GDN-prefill "
"backend (the flashinfer sm90 GDN kernel would fail to build)."
)
return final_vllm_kwargs
[docs]
@classmethod
def build_serve_args(
cls, model_params: ModelParams, *, tool_call_parser: str | None = None
) -> list[str]:
"""Builds ``vllm serve`` arguments that load the model as the engine does.
Reads local state like ``build_engine_kwargs``. A LoRA adapter is served
under the model name ``LORA_ADAPTER_NAME``.
Args:
model_params: The model parameters to use for inference.
tool_call_parser: See ``__init__``.
Returns:
list[str]: Arguments for ``vllm serve``, starting with the model.
"""
engine_kwargs = cls.build_engine_kwargs(model_params)
args = [str(engine_kwargs.pop("model"))]
for key, value in engine_kwargs.items():
flag = key.replace("_", "-")
if value is None:
continue
if isinstance(value, bool):
args.append(f"--{flag}" if value else f"--no-{flag}")
elif isinstance(value, dict):
args += [f"--{flag}", json.dumps(value)]
else:
args += [f"--{flag}", str(value)]
if model_params.adapter_model:
args += [
"--lora-modules",
f"{LORA_ADAPTER_NAME}={model_params.adapter_model}",
]
parser_name = tool_call_parser or model_params.tool_call_parser
if parser_name:
args += ["--enable-auto-tool-choice", "--tool-call-parser", parser_name]
# Applied to every chat as the engine does; a request's own kwargs win.
if model_params.chat_template_kwargs:
args += [
"--default-chat-template-kwargs",
json.dumps(model_params.chat_template_kwargs),
]
# The engine sets sampling per request; don't let the model's
# generation_config.json fill in what a request leaves out.
args += ["--generation-config", "vllm"]
return args
@staticmethod
def _normalize_vllm_finish_reason(raw_reason: str | None) -> FinishReason | None:
"""Normalize vLLM finish_reason string to FinishReason enum."""
if raw_reason is None:
return None
mapping = {
"stop": FinishReason.STOP,
"length": FinishReason.LENGTH,
}
return mapping.get(raw_reason.lower(), FinishReason.UNKNOWN)
@staticmethod
def _extract_usage_from_vllm_output(
request_output: RequestOutput,
) -> dict[str, int] | None:
"""Extracts normalized token usage from a vLLM chat response.
Args:
request_output: A vLLM ``RequestOutput`` for a single conversation.
Returns:
A dict with prompt_tokens, completion_tokens, total_tokens, or None
if vLLM reported no token ids.
"""
prompt_tokens = len(request_output.prompt_token_ids or [])
completion_tokens = sum(
len(output.token_ids or []) for output in request_output.outputs
)
if not prompt_tokens and not completion_tokens:
return None
return {
"prompt_tokens": prompt_tokens,
"completion_tokens": completion_tokens,
"total_tokens": prompt_tokens + completion_tokens,
}
def _convert_conversation_to_vllm_input(
self, conversation: Conversation
) -> list[ChatCompletionMessageParam]:
"""Converts a conversation to a list of vllm input messages.
Args:
conversation: The conversation to convert.
Returns:
List[ChatCompletionMessageParam]: A list of vllm input messages.
"""
result: list[ChatCompletionMessageParam] = []
messages = conversation.messages
parser_capabilities = _get_tool_parser_capabilities(
getattr(self, "_tool_parser_name", None)
)
if parser_capabilities.requires_split_tool_narration:
messages = split_assistant_content_and_tool_calls(messages)
for json_dict in create_list_of_message_json_dicts(
messages, group_adjacent_same_role_turns=True
):
if "role" not in json_dict:
raise RuntimeError("The required field 'role' is missing!")
if "content" not in json_dict:
raise RuntimeError("The required field 'content' is missing!")
content = json_dict["content"]
# Assistant messages with only `tool_calls` legitimately have
# `content=None` per the OpenAI wire format.
if content is not None and not isinstance(content, str | list):
raise RuntimeError(
"The 'content' field must be `str`, `list`, or `None`. "
f"Actual: {type(content)}."
)
result.append(json_dict) # type: ignore[arg-type]
return result
def _infer(
self,
input: list[Conversation],
inference_config: InferenceConfig | None = None,
) -> list[Conversation]:
"""Runs model inference on the provided input.
Documentation: https://docs.vllm.ai/en/stable/dev/sampling_params.html
Args:
input: A list of conversations to run inference on.
inference_config: Parameters for inference.
Returns:
List[Conversation]: Inference output.
"""
generation_params = (
inference_config.generation
if inference_config and inference_config.generation
else self._generation_params
)
model_params = (
inference_config.model
if inference_config and inference_config.model
else self._model_params
)
if generation_params.guided_decoding is not None:
if _VLLM_V0_12:
# vLLM v0.12+ uses StructuredOutputsParams (direct construction)
guided_decoding = VLLMGuidedDecodingParams(
json=generation_params.guided_decoding.json,
regex=generation_params.guided_decoding.regex,
choice=generation_params.guided_decoding.choice,
)
else:
# vLLM <0.12 uses GuidedDecodingParams.from_optional()
guided_decoding = VLLMGuidedDecodingParams.from_optional( # type: ignore[attr-defined] # pyright: ignore[reportAttributeAccessIssue]
json=generation_params.guided_decoding.json,
regex=generation_params.guided_decoding.regex,
choice=generation_params.guided_decoding.choice,
)
else:
guided_decoding = None
# In vLLM v0.12+, the kwarg was renamed from 'guided_decoding'
# to 'structured_outputs'.
guided_decoding_kwarg = (
{"structured_outputs": guided_decoding}
if _VLLM_V0_12
else {"guided_decoding": guided_decoding}
)
sampling_params = SamplingParams(
n=1,
max_tokens=generation_params.max_new_tokens,
temperature=generation_params.temperature,
seed=generation_params.seed,
top_p=generation_params.top_p
if generation_params.top_p is not None
else 1.0,
frequency_penalty=generation_params.frequency_penalty,
presence_penalty=generation_params.presence_penalty,
stop=generation_params.stop_strings,
stop_token_ids=generation_params.stop_token_ids,
min_p=generation_params.min_p,
**guided_decoding_kwarg, # pyright: ignore[reportArgumentType]
skip_special_tokens=generation_params.skip_special_tokens,
)
output_conversations = []
vllm_conversations = []
non_skipped_conversations = []
for conversation in input:
if not conversation.messages:
logger.warning("Conversation must have at least one message.")
continue
vllm_input = self._convert_conversation_to_vllm_input(conversation)
vllm_conversations.append(vllm_input)
non_skipped_conversations.append(conversation)
if len(vllm_conversations) == 0:
return []
any_tools = any(c.tools for c in non_skipped_conversations)
if not any_tools:
# Fast path: single batched call, unchanged from prior behavior.
chat_responses = self._llm.chat(
vllm_conversations,
sampling_params=sampling_params,
lora_request=self._lora_request,
use_tqdm=(len(vllm_conversations) >= 2),
chat_template=None,
chat_template_content_format="auto",
chat_template_kwargs=model_params.chat_template_kwargs,
)
else:
# vLLM's `chat(tools=...)` applies one tool list to the whole
# batch, so we group conversations by identical tools and dispatch
# one chat() per group. Original input order is preserved.
chat_responses: list = [None] * len(non_skipped_conversations)
for indices, group_tools in self._group_by_tools(non_skipped_conversations):
group_inputs = [vllm_conversations[i] for i in indices]
group_responses = self._llm.chat(
group_inputs,
sampling_params=sampling_params,
lora_request=self._lora_request,
use_tqdm=(len(group_inputs) >= 2),
chat_template=None,
chat_template_content_format="auto",
chat_template_kwargs=model_params.chat_template_kwargs,
tools=group_tools,
)
for idx, resp in zip(indices, group_responses):
chat_responses[idx] = resp
for conversation, chat_response in zip(
non_skipped_conversations, chat_responses
):
assert chat_response is not None
new_messages, finish_reason_override = self._build_response_messages(
conversation, chat_response
)
messages = [
*conversation.messages,
*new_messages,
]
metadata = dict(conversation.metadata)
if chat_response.outputs:
if finish_reason_override is not None:
metadata["finish_reason"] = finish_reason_override.value
else:
raw_reason = chat_response.outputs[0].finish_reason
finish_reason = self._normalize_vllm_finish_reason(raw_reason)
if finish_reason is not None:
metadata["finish_reason"] = finish_reason.value
usage = self._extract_usage_from_vllm_output(chat_response)
if usage is not None:
metadata["usage"] = usage
new_conversation = Conversation(
messages=messages,
metadata=metadata,
conversation_id=conversation.conversation_id,
tools=conversation.tools,
)
self._save_conversation_to_scratch(
new_conversation,
inference_config.output_path if inference_config else None,
)
output_conversations.append(new_conversation)
return output_conversations
@staticmethod
def _group_by_tools(
conversations: list[Conversation],
) -> list[tuple[list[int], list[dict] | None]]:
"""Groups conversation indices by identical tool definitions.
Tools come from ``Conversation.tools`` (post-validation a list of
``ToolDefinition`` Pydantic models, or ``None``). vLLM's
``LLM.chat(tools=...)`` kwarg expects OpenAI-format dicts and is
batch-global, so we dump to dicts up front and split heterogeneous
batches.
Returns:
A list of ``(indices, tools)`` tuples in the order each unique
tools list is first seen.
"""
groups: list[tuple[list[int], list[dict] | None]] = []
key_to_group: dict[str, int] = {}
for i, conv in enumerate(conversations):
tools: list[dict] | None = (
None
if conv.tools is None
else [t.model_dump(mode="json", exclude_none=True) for t in conv.tools]
)
if tools is None:
key = "__none__"
else:
key = json.dumps(tools, sort_keys=True)
if key not in key_to_group:
key_to_group[key] = len(groups)
groups.append(([], tools))
groups[key_to_group[key]][0].append(i)
return groups
def _build_response_messages(
self,
conversation: Conversation,
chat_response,
) -> tuple[list[Message], FinishReason | None]:
"""Builds assistant messages from a vLLM chat response.
When ``self._tool_parser`` is set, uses its parsed content and tool calls.
On parser failure, falls back to clean display text.
"""
new_messages: list[Message] = []
finish_reason_override: FinishReason | None = None
parser_capabilities = _get_tool_parser_capabilities(
getattr(self, "_tool_parser_name", None)
)
for completion in chat_response.outputs:
display_text = completion.text
content: str | None = display_text
tool_calls_payload: list[ToolCall] | None = None
if self._tool_parser is not None:
parser_text = self._decode_tool_parser_text(completion)
# Some parsers read `request.tool_choice` from the
# non-streaming entry point; pass a stub so they don't crash.
# vLLM offline `LLM.chat()` has no `tool_choice` knob, so
# "auto" is the only honest value — it matches OpenAI's
# default-when-unset and the offline reality of "model
# decides based on the rendered tools". Other values would
# mislead the parser ("none" would drop real calls;
# "required" or a specific function would force/filter
# interpretations the model didn't actually make).
stub = SimpleNamespace(
tool_choice="auto",
tools=_build_parser_request_tools(
self._tool_parser, conversation.tools
),
)
try:
extracted = self._tool_parser.extract_tool_calls(
parser_text,
request=stub, # type: ignore[arg-type]
)
except Exception:
logger.exception(
"Tool-call parser %s failed; falling back to display text.",
type(self._tool_parser).__name__,
)
extracted = None
if extracted is not None:
tools_called = bool(getattr(extracted, "tools_called", False))
if tools_called:
content = extracted.content or None
if parser_capabilities.requires_raw_output:
content = self._clean_tool_parser_content(content)
tool_calls_payload = [
ToolCall.model_validate(tc.model_dump())
for tc in extracted.tool_calls
]
finish_reason_override = FinishReason.TOOL_CALLS
elif parser_capabilities.requires_raw_output:
# Lossless parser input may contain protocol tokens or
# otherwise differ from the user-facing representation.
content = display_text
else:
# Preserve the parser's normal clean-text behavior, but
# never let a parser miss erase a nonempty completion.
content = extracted.content or display_text
new_messages.append(
Message(
content=content,
role=Role.ASSISTANT,
tool_calls=tool_calls_payload,
)
)
return new_messages, finish_reason_override
def _decode_tool_parser_text(self, completion) -> str:
"""Build the text representation required by the active tool parser.
``completion.text`` follows the user-facing ``skip_special_tokens``
setting. Parsers use that clean text by default. Parsers that explicitly
require lossless output instead receive a second decode of the same
generated token IDs with special tokens preserved. The clean text remains
the source of truth for display and parser fallbacks.
"""
capabilities = _get_tool_parser_capabilities(
getattr(self, "_tool_parser_name", None)
)
if not capabilities.requires_raw_output:
return completion.text
token_ids = getattr(completion, "token_ids", None)
if token_ids is None or len(token_ids) == 0:
return completion.text
try:
decoded = self._tokenizer.decode(
token_ids,
skip_special_tokens=False,
clean_up_tokenization_spaces=False,
)
if isinstance(decoded, str):
return decoded
except Exception:
logger.exception(
"Failed to decode raw tokens for tool-call parsing; "
"falling back to display text."
)
return completion.text
def _clean_tool_parser_content(self, content: str | None) -> str | None:
"""Removes special tokens from content extracted from lossless parser input."""
if not content:
return None
try:
token_ids = self._tokenizer.encode(content, add_special_tokens=False)
decoded = self._tokenizer.decode(
token_ids,
skip_special_tokens=True,
clean_up_tokenization_spaces=False,
)
if isinstance(decoded, str):
return decoded if decoded.strip() else None
except Exception:
logger.exception(
"Failed to clean tool-call parser content; omitting content."
)
# Keep valid tool calls even when their accompanying text cannot be cleaned.
return None
[docs]
def infer_online(
self,
input: list[Conversation],
inference_config: InferenceConfig | None = None,
) -> list[Conversation]:
"""Runs model inference online.
Args:
input: A list of conversations to run inference on.
inference_config: Parameters for inference.
Returns:
List[Conversation]: Inference output.
"""
warnings.warn(
"infer_online() will be private in the future. Use infer() instead.",
DeprecationWarning,
stacklevel=2,
)
results = self._infer_online(input, inference_config)
if inference_config and inference_config.output_path:
self._save_conversations(results, inference_config.output_path)
return results
[docs]
def infer_from_file(
self,
input_filepath: str,
inference_config: InferenceConfig | None = None,
) -> list[Conversation]:
"""Runs model inference on inputs in the provided file.
This is a convenience method to prevent boilerplate from asserting the
existence of input_filepath in the generation_params.
Args:
input_filepath: Path to the input file containing prompts for
generation.
inference_config: Parameters for inference.
Returns:
List[Conversation]: Inference output.
"""
warnings.warn(
"infer_from_file() will be private in the future. Use infer() instead.",
DeprecationWarning,
stacklevel=2,
)
input = self._read_conversations(input_filepath)
results = self._infer(input, inference_config)
if inference_config and inference_config.output_path:
self._save_conversations(results, inference_config.output_path)
return results
@override
def _infer_online(
self,
input: list[Conversation],
inference_config: InferenceConfig | None = None,
) -> list[Conversation]:
"""Runs model inference online.
Args:
input: A list of conversations to run inference on.
inference_config: Parameters for inference.
Returns:
List[Conversation]: Inference output.
"""
return self._infer(input, inference_config)
[docs]
@override
def get_supported_params(self) -> set[str]:
"""Returns a set of supported generation parameters for this engine."""
return {
"frequency_penalty",
"guided_decoding",
"max_new_tokens",
"min_p",
"presence_penalty",
"seed",
"skip_special_tokens",
"stop_strings",
"stop_token_ids",
"temperature",
"top_p",
}
