# NOTE: This module is imported at the top level of furiosa_llm, so it must not
# import packages that require native packages (e.g.,
# furiosa-llm-native, furiosa-native-llm-common) at
# module scope. Use lazy imports inside methods or TYPE_CHECKING blocks instead,
# so that `import furiosa_llm` works without those packages installed.
# This is necessary because npu-tools and download CI environments depend only
# on furiosa-models without native packages.
from collections.abc import AsyncGenerator, Sequence
import json
import logging
import os
from pathlib import Path
from typing import (
TYPE_CHECKING,
Any,
Literal,
Protocol,
cast,
)
import weakref
from pydantic_core import to_json
from furiosa_llm.vllm_compat import (
PromptType,
fit_prompt_to_context,
get_score_prompt,
preprocess_prompt,
prompt_to_str,
resolve_score_truncation_kwargs,
resolve_truncation_side,
)
if TYPE_CHECKING:
from furiosa.llm_native.llm import VisionPatches
from openai.types.chat import ChatCompletionMessageParam
import uuid
import torch
from transformers.tokenization_python import PreTrainedTokenizer
from transformers.tokenization_utils_base import BatchEncoding
from transformers.tokenization_utils_tokenizers import PreTrainedTokenizerFast
from furiosa_llm.metadata.tasks import POOLING_TASKS, PoolingTask, TaskType
from furiosa_llm.server.utils import is_list_of, random_uuid
from furiosa_llm.version import FURIOSA_LLM_VERSION
from .device import resolve_devices
from .errors import validate_context_length
from .metadata.config_types import LoggerConfig, SchedulerConfig
from .metadata.runtime import ModelMetadata
from .outputs import (
EmbeddingRequestOutput,
NativeOutputConverter,
PoolingOutput,
PoolingRequestOutput,
RequestOutput,
RequestOutputKind,
ScoringRequestOutput,
)
from .sampling_params import PoolingParams, SamplingParams
from .tokenizer import encode_auto, get_tokenizer
from .utils import get_logger_with_tz
logger = get_logger_with_tz(logging.getLogger(__name__))
def _resolve_fxb(resolved, model_id_or_path, revision):
"""Resolve the FXB to run when none was given explicitly via ``--fxb``.
Order: (1) an ``.fxb`` shipped in the model's own repo (``discover_fxb``,
already filtered to the running build's revision), then (2) the local FXB
cache. Cache compatibility is decided by the model fingerprint
(architecture plus the kernel-affecting ``config.json`` fields — hidden_size,
intermediate_size, quantization, etc.), matched across *every* cached bundle
regardless of which repo it came from; the source repo id is only a ranking
tiebreaker. The cache is therefore consulted only when ``model_id_or_path``
is a repo id (not a local path), because building that fingerprint requires
fetching the target's ``config.json`` from the Hub by repo id.
``revision`` is the already-resolved HuggingFace revision (the caller has
applied the furiosa-llm-version default), used to fetch that target
``config.json`` so the fingerprint reflects the same repo snapshot as the
model download.
From the cache we auto-use only a fingerprint-compatible bundle whose
FuriosaIR (npu-ir) revision matches the running build; if the cache holds
only compatible-but-stale (revision-mismatched) bundles, we refuse and tell
the user to pass ``--fxb`` explicitly rather than silently loading a bundle
built for a different compiler revision.
"""
from furiosa_llm.utils import _looks_like_path, discover_fxb
try:
return discover_fxb(resolved)
except FileNotFoundError as repo_err:
if _looks_like_path(model_id_or_path):
raise
from furiosa.llm_native import fxb as _fxb
try:
outcome = _fxb.check(str(model_id_or_path), revision=revision, cache_dir=None)
except Exception as check_err: # noqa: BLE001 - no network / not a real repo / no config
logger.debug("FXB cache lookup failed: %s", check_err)
raise repo_err from None
recommended = next(
(m for m in outcome.matches if m.recommended and m.npu_tools_match), None
)
if recommended is not None:
logger.info(
"No FXB in '%s'; using compatible cached FXB %s",
model_id_or_path,
recommended.entry.path,
)
return recommended.entry.path
if outcome.matches:
# Compatible bundles exist but none matches the running FuriosaIR
# revision; require an explicit choice.
raise FileNotFoundError(
f"No FXB found in '{model_id_or_path}', and the local cache only has "
f"FuriosaIR-revision-mismatched (stale) matches. Pass --fxb <path> to use one "
f"explicitly (e.g. {outcome.matches[0].entry.path})."
)
raise repo_err
def _detect_fxb_task(model_dir: Path, hf_config: Any) -> tuple[TaskType, bool]:
"""Infer (task, use_binary_seq_class) for an FXB-served model from its
sentence-transformers metadata, mirroring the Rust `detect_convert_task`
(furiosa-generator/src/fxb/metadata.rs): Qwen3 reuses `Qwen3ForCausalLM` for
its embedding and reranker repos, so the HF architecture can't tell them
apart. Gate on `modules.json` (absent -> generative); otherwise
`config_sentence_transformers.json`'s `model_type` `"CrossEncoder"` -> score
(binary seq-class head), anything else -> embed.
Cross-encoder rerankers ship no sentence-transformers metadata at all
(BAAI/bge-reranker-v2-m3), so before falling back to generative, read the HF
architecture: a one-label sequence-classification head is a binary scorer."""
from furiosa_llm.metadata.tasks import EMBED, GENERATE, SCORE
if not (model_dir / "modules.json").exists():
architectures = getattr(hf_config, "architectures", None) or []
if (
any(arch.endswith("ForSequenceClassification") for arch in architectures)
and getattr(hf_config, "num_labels", None) == 1
):
return SCORE, True
return GENERATE, False
model_type = None
st_config = model_dir / "config_sentence_transformers.json"
if st_config.exists():
try:
model_type = json.loads(st_config.read_text(encoding="utf-8")).get("model_type")
except (OSError, ValueError):
model_type = None
return (SCORE, True) if model_type == "CrossEncoder" else (EMBED, False)
# Default index of the padding block when paged attention model is used.
DEFAULT_PAGED_ATTENTION_PADDING_BLOCK_IDX = 0
TokenizerModeType = Literal["auto", "slow"]
ChatTemplateContentFormatOption = Literal["string"]
RAY_LOG_PREFIX = "[furiosa-llm]"
class NativeEngineLike(Protocol):
@property
def max_model_len(self) -> int: ...
@property
def pooling_type(self) -> str | None: ...
def generate(self, inputs: Any, sampling_params: Any) -> Any: ...
def stream_generate(
self,
inputs: Any,
sampling_params: Any,
request_id: str | None = None,
) -> AsyncGenerator[Any, None]: ...
def acquire_vision(self, media_id: int) -> list[int] | None: ...
async def encode_vision(
self,
media_ids: list[int],
patches: list["VisionPatches"],
request_id: str | None = None,
) -> list[list[int]]: ...
async def encode(
self,
inputs: Any,
pooling_params: Any,
request_id: str | None = None,
) -> Any: ...
def abort_request(self, request_id: str) -> None: ...
def is_alive(self) -> bool: ...
def wait_until_terminated(self) -> None: ...
def shutdown(self) -> None: ...
def reset_prefix_cache(self) -> bool: ...
def get_prometheus_metrics(self) -> str: ...
def _llm_finalize(engine: "NativeEngineLike | None", tmp_dir: Any) -> None:
if engine is not None:
engine.shutdown()
if tmp_dir is not None:
tmp_dir.cleanup()
def validate_speculative_args(
draft_model_id_or_path: "str | os.PathLike | None",
draft_fxb_path: "str | os.PathLike | None",
num_speculative_tokens: "int | None",
) -> None:
"""Validate the mutual constraints between the speculative-decoding arguments.
Shared by ``LLM.__init__`` and ``EngineArgs.__post_init__`` so the two stay in sync.
"""
if draft_model_id_or_path is None:
if num_speculative_tokens is not None:
raise ValueError(
"draft_model_id_or_path is required when num_speculative_tokens is set"
)
if draft_fxb_path is not None:
raise ValueError("draft_model_id_or_path is required when draft_fxb_path is set")
elif num_speculative_tokens is None:
raise ValueError("num_speculative_tokens is required when draft_model_id_or_path is set")
elif num_speculative_tokens < 1:
raise ValueError(f"num_speculative_tokens must be >= 1, got {num_speculative_tokens}")
def _validate_speculative_tokenizers(
target_tokenizer: PreTrainedTokenizer | PreTrainedTokenizerFast,
draft_tokenizer: PreTrainedTokenizer | PreTrainedTokenizerFast,
) -> None:
"""Require target and draft tokenizers to use the same token-id contract."""
target_vocab = target_tokenizer.get_vocab()
draft_vocab = draft_tokenizer.get_vocab()
if target_vocab != draft_vocab:
token = min(
token
for token in target_vocab.keys() | draft_vocab.keys()
if target_vocab.get(token) != draft_vocab.get(token)
)
raise ValueError(
"speculative decoding requires identical target and draft token-to-id mappings; "
f"first mismatch: token={token!r}, target id={target_vocab.get(token)!r}, "
f"draft id={draft_vocab.get(token)!r} "
f"(entries: target={len(target_vocab)}, draft={len(draft_vocab)})"
)
[docs]
class LLM:
"""An LLM for generating texts from given prompts and sampling parameters."""
max_model_len: int
engine: NativeEngineLike
def __init__(
self,
model_id_or_path: str | os.PathLike,
*,
# V3 engine Configuration
fxb: str | os.PathLike | None = None,
draft_model_id_or_path: str | os.PathLike | None = None,
draft_fxb_path: str | os.PathLike | None = None,
num_speculative_tokens: int | None = None,
# Repo Configuration
revision: str | None = None,
# Runtime Configuration
devices: str | Sequence[str] | None = None,
data_parallel_size: int | None = None,
pipeline_parallel_size: int | None = None,
max_io_memory_mb: int | None = None,
max_model_len: int | None = None,
# Pipeline selection related Configs
scheduler_config: SchedulerConfig | None = None,
# Observability related Configs
logger_config: LoggerConfig | None = None,
# Structured outputs related Configuration
structured_outputs_backend: Literal["auto", "guidance", "xgrammar"] = "auto",
# Other Configuration
tokenizer: str | PreTrainedTokenizer | PreTrainedTokenizerFast | None = None,
tokenizer_mode: TokenizerModeType = "auto",
seed: int | None = None,
served_model_name: str | None = None,
kv_transfer_config: dict[str, Any] | None = None,
**kwargs,
) -> None:
"""Instantiate LLM from saved artifacts or from a v3 engine (FXB).
When ``fxb`` is provided, the v3 engine path is used: ``model_id_or_path`` is
treated as a HuggingFace model id or directory, and model metadata is derived from it.
When ``fxb`` is omitted, ``model_id_or_path`` is resolved via HuggingFace Hub
download if needed. If the resolved directory contains ``artifact.json``, the
legacy artifact path is used. Otherwise, the v3 engine path is used and ``.fxb``
files are auto-discovered from the model directory.
Args:
model_id_or_path: A path to furiosa llm engine artifact or a HuggingFace model id.
fxb: Path to the FXB artifact file. When provided, the v3 engine is used
and ``model_id_or_path`` is treated as a HuggingFace model id or directory.
When omitted, FXB files are auto-discovered from the model directory if
no ``artifact.json`` is present.
draft_model_id_or_path: HuggingFace model id or local path for the draft model used by
speculative decoding.
draft_fxb_path: Optional FXB file for the draft model in v3 engine mode. When omitted,
draft FXB files are auto-discovered from the draft model directory.
num_speculative_tokens: Maximum draft tokens to propose per speculative decoding round.
revision: The revision of the model, if `model_id_or_path` is a HuggingFace model id.
devices: The devices to run the model. It can be a single device or a list of devices.
Each device can be either "npu:X" or "npu:X:\\*" where X is a specific device index.
If not given, all available devices will be used.
data_parallel_size: The size of the data parallelism group. If not given, it will be inferred from
total available PEs and other parallelism degrees.
pipeline_parallel_size: The size of the pipeline parallelism.
If not given, it will use the value from artifact.
num_blocks_per_pp_stage: The number of transformer blocks per each pipeline parallelism stage.
If only `pipeline_parallel_size` is provided, transformer blocks will be distributed equally.
max_io_memory_mb: NPU runtime I/O buffer reserve per device pool, from 1 to 49152 MiB.
Co-located speculative target and draft lanes share one reserve. If unspecified,
the runtime estimates the required size from the model's bucket shapes.
max_model_len: Cap on the context length (prompt + generated tokens) a single request
may use. It may only lower the model's own limit: a value above either the
model's ``max_position_embeddings`` or the context length the artifact was
compiled for is rejected before any device is opened. If not given, the
artifact's own limit is used.
scheduler_config: Configuration for the scheduler, allowing to maximum number of tasks which can
be queued to HW, maximum number of samples that can be processed by the scheduler, and ratio
of spare blocks that are reserved by scheduler. If this is not given, scheduler config saved
in the artifacts will be used.
logger_config: Observability configuration (e.g. disabling the periodic per-DP stats
log). If not given, default logging behavior is used.
structured_outputs_backend: The backend for structured outputs. "auto" will automatically select the
best backend based on the model. "guidance" will use the guidance library. "xgrammar" will
use the xgrammar library.
tokenizer: The name or path of a HuggingFace Transformers tokenizer.
tokenizer_mode: The tokenizer mode. "auto" will use the fast tokenizer
if available, and "slow" will always use the slow tokenizer.
seed: The seed to initialize the random number generator for sampling.
skip_engine: If True, the native runtime engine will not be initialized. This is useful when you need
the pipelines for other purposes than running them with the engine.
served_model_name: The model name used in metrics and API responses.
If not specified, defaults to ``model_id_or_path``.
kv_transfer_config: Configuration for the remote KV connector used in PD-disaggregation.
"""
# The engine rejects an out-of-range cap, but `skip_engine=True` never builds one.
if max_model_len is not None and max_model_len <= 0:
raise ValueError(f"max_model_len must be greater than 0, got {max_model_len}")
validate_speculative_args(draft_model_id_or_path, draft_fxb_path, num_speculative_tokens)
self.num_speculative_tokens = num_speculative_tokens
# Resolve model_id_or_path and determine v2 (artifact) vs v3 (FXB) engine.
# - fxb explicitly provided → v3 engine
# - artifact.json exists → legacy v2 artifact
# - otherwise → v3 engine with FXB auto-discovery
from furiosa_llm.utils import get_path_or_hf_download, resolve_effective_revision
# When no explicit revision is given, default to the furiosa-llm version
# tag for furiosa-ai repos (same rule as the legacy artifact path). This
# effective revision is reused for both the model download below and the
# FXB cache lookup, so they always target the same repo snapshot. When
# the default tag was auto-applied and the repo has no such tag yet, fall
# back to the default branch (main).
effective_revision, allow_main_fallback = resolve_effective_revision(
model_id_or_path, revision, FURIOSA_LLM_VERSION
)
resolved = get_path_or_hf_download(
model_id_or_path,
effective_revision,
fallback_to_main_on_missing=allow_main_fallback,
)
if draft_model_id_or_path is not None:
if os.fspath(draft_model_id_or_path) == os.fspath(model_id_or_path):
draft_effective_revision = effective_revision
resolved_draft = resolved
else:
draft_effective_revision, draft_allow_main_fallback = resolve_effective_revision(
draft_model_id_or_path, None, FURIOSA_LLM_VERSION
)
resolved_draft = get_path_or_hf_download(
draft_model_id_or_path,
draft_effective_revision,
fallback_to_main_on_missing=draft_allow_main_fallback,
)
else:
draft_effective_revision = None
resolved_draft = None
resolved_draft_fxb_path = draft_fxb_path
fxb_path = fxb or _resolve_fxb(resolved, model_id_or_path, effective_revision)
if resolved_draft is not None and resolved_draft_fxb_path is None:
if Path(resolved_draft).resolve() == Path(resolved).resolve():
resolved_draft_fxb_path = fxb_path
else:
resolved_draft_fxb_path = _resolve_fxb(
resolved_draft, draft_model_id_or_path, draft_effective_revision
)
self._init_from_v3_engine(
model_id_or_path=str(resolved),
fxb_path=fxb_path,
draft_model_id_or_path=resolved_draft,
draft_fxb_path=resolved_draft_fxb_path,
num_speculative_tokens=num_speculative_tokens,
devices=devices,
data_parallel_size=data_parallel_size,
pipeline_parallel_size=pipeline_parallel_size,
max_io_memory_mb=max_io_memory_mb,
max_model_len=max_model_len,
scheduler_config=scheduler_config,
logger_config=logger_config,
structured_outputs_backend=structured_outputs_backend,
tokenizer=tokenizer,
tokenizer_mode=tokenizer_mode,
served_model_name=served_model_name,
kv_transfer_config=kv_transfer_config,
)
self._finalizer = weakref.finalize(
self,
_llm_finalize,
getattr(self, "engine", None),
getattr(self, "tmp_dir", None),
)
def _init_from_v3_engine(
self,
model_id_or_path: str | os.PathLike,
fxb_path: str | os.PathLike,
*,
draft_model_id_or_path: str | os.PathLike | None = None,
draft_fxb_path: str | os.PathLike | None = None,
num_speculative_tokens: int | None = None,
devices: str | Sequence[str] | None = None,
data_parallel_size: int | None = None,
pipeline_parallel_size: int | None = None,
max_io_memory_mb: int | None = None,
max_model_len: int | None = None,
scheduler_config: SchedulerConfig | None = None,
logger_config: LoggerConfig | None = None,
structured_outputs_backend: Literal["auto", "guidance", "xgrammar"] = "auto",
tokenizer: str | PreTrainedTokenizer | PreTrainedTokenizerFast | None = None,
tokenizer_mode: TokenizerModeType = "auto",
served_model_name: str | None = None,
kv_transfer_config: dict[str, Any] | None = None,
) -> None:
from transformers import AutoConfig
from furiosa_llm.generation_config import get_diff_sampling_params
from furiosa_llm.metadata.tasks import GENERATION_TASKS
model_id_or_path = str(model_id_or_path)
fxb_path = str(fxb_path)
draft_model_path = (
str(draft_model_id_or_path) if draft_model_id_or_path is not None else None
)
draft_fxb = str(draft_fxb_path) if draft_fxb_path is not None else None
# Load tokenizer from HF model directory
self.tokenizer = get_tokenizer(model_id_or_path, tokenizer, tokenizer_mode=tokenizer_mode)
if draft_model_path is not None:
draft_tokenizer = get_tokenizer(draft_model_path, tokenizer_mode=tokenizer_mode)
_validate_speculative_tokenizers(self.tokenizer, draft_tokenizer)
# Recover the pooling task from sentence-transformers metadata; an FXB
# may be an embedding/reranker bundle despite a generative HF arch, and
# treating it as generative would 400 the pooling endpoints.
hf_config = AutoConfig.from_pretrained(model_id_or_path)
hf_configs = hf_config.to_dict()
task, use_binary_seq_class = _detect_fxb_task(Path(model_id_or_path), hf_config)
model_metadata = ModelMetadata(
model_type=hf_config.model_type,
task=task,
hf_configs=hf_configs,
use_binary_seq_class=use_binary_seq_class,
)
# Derive sequence length limits from HF config
max_position_embeddings = hf_configs.get("max_position_embeddings") or hf_configs.get(
"text_config", {}
).get("max_position_embeddings")
if max_position_embeddings is None:
raise ValueError("'max_position_embeddings' not found in the model config.")
# Set instance attributes
# `model_id_or_path` is a local snapshot dir (already resolved
# via get_path_or_hf_download in __init__). The MM helpers in
# furiosa_llm.multimodal feed it to AutoProcessor.from_pretrained.
self.model_id_or_path = model_id_or_path
self.served_model_name = served_model_name or model_id_or_path
self.model_metadata = model_metadata
self.model_config = hf_configs
self.artifact_id = Path(fxb_path).name
self.is_generative_model = task in GENERATION_TASKS
self.default_generation_config: dict[str, Any] = get_diff_sampling_params(model_id_or_path)
devices = resolve_devices(devices)
from furiosa.llm_native.llm import NativeLLMEngine
self.engine = NativeLLMEngine(
model_id_or_path,
draft_model_path,
devices,
data_parallel_size,
pipeline_parallel_size,
max_io_memory_mb,
self._serialize_obj(scheduler_config or SchedulerConfig()),
self._serialize_obj(logger_config or LoggerConfig()),
structured_outputs_backend,
num_speculative_tokens,
self.served_model_name,
fxb_path, # fxb_path — signals v3 engine
draft_fxb,
self._serialize_obj(kv_transfer_config) if kv_transfer_config else None,
max_model_len,
)
# Read back the value the engine resolved, not the one requested.
self.max_model_len = self.engine.max_model_len
[docs]
def get_default_sampling_params(self) -> SamplingParams:
"""Return SamplingParams reflecting model's generation_config defaults.
If the model has no generation_config or it matches HF defaults,
returns a default SamplingParams().
"""
if self.default_generation_config:
return SamplingParams.from_optional(**self.default_generation_config)
return SamplingParams()
@classmethod
def _serialize_obj(
cls,
obj: Any,
) -> str:
return to_json(obj).decode("utf-8")
[docs]
def generate(
self,
prompts: str | list[str],
sampling_params: SamplingParams | None = None,
prompt_token_ids: BatchEncoding | None = None,
tokenizer_kwargs: dict[str, Any] | None = None,
) -> RequestOutput | list[RequestOutput]:
"""Generate texts from given prompts and sampling parameters.
Args:
prompts: The prompts to generate texts.
sampling_params: The sampling parameters for generating texts.
If None, model's generation config defaults are used.
prompt_token_ids: Pre-tokenized prompt input as a `BatchEncoding` object.
If not provided, the prompt will be tokenized internally using the tokenizer.
tokenizer_kwargs: Additional keyword arguments passed to the tokenizer's
`encode` method, such as `{"use_special_tokens": True}`.
Returns:
A list of `RequestOutput` objects containing the generated
completions in the same order as the input prompts.
"""
if sampling_params is None:
sampling_params = self.get_default_sampling_params()
if not self.is_generative_model:
raise ValueError("generate API can only be used for generative models.")
if prompt_token_ids is None:
if tokenizer_kwargs is None:
tokenizer_kwargs = {}
prompt_token_ids = encode_auto(self.tokenizer, prompts, **tokenizer_kwargs)
input_ids = prompt_token_ids.input_ids
if input_ids and isinstance(input_ids[0], list):
longest_prompt_len = max(len(prompt) for prompt in input_ids)
else:
longest_prompt_len = len(input_ids)
validate_context_length(
max_model_len=self.max_model_len,
prompt_tokens=longest_prompt_len,
max_completion_tokens=sampling_params.max_tokens,
min_completion_tokens=max(1, sampling_params.min_tokens),
)
if isinstance(prompts, list):
prompt_list = prompts
batch_encodings = [
BatchEncoding({key: value[index] for key, value in prompt_token_ids.items()})
for index in range(len(prompts))
]
else:
prompt_list = [prompts]
batch_encodings = [prompt_token_ids]
async def generate_one(prompt: str, batch_encoding: BatchEncoding) -> RequestOutput:
final_output: RequestOutput | None = None
async for output in self._convert_generation_stream(
prompt,
batch_encoding,
sampling_params,
RequestOutputKind.FINAL,
):
final_output = output
if final_output is None:
raise RuntimeError("Native generation stream completed without producing output.")
return final_output
from furiosa_llm.utils import async_gather, run_sync
outputs = run_sync(
async_gather(
*(
generate_one(prompt, batch_encoding)
for prompt, batch_encoding in zip(prompt_list, batch_encodings)
)
)
)
return outputs if isinstance(prompts, list) else outputs[0]
[docs]
def chat(
self,
messages: list["ChatCompletionMessageParam"] | list[list["ChatCompletionMessageParam"]],
sampling_params: SamplingParams | None = None,
chat_template: str | None = None,
chat_template_content_format: ChatTemplateContentFormatOption = "string",
add_generation_prompt: bool = True,
continue_final_message: bool = False,
tools: list[dict[str, Any]] | None = None,
chat_template_kwargs: dict[str, Any] | None = None,
) -> list[RequestOutput]:
"""
Generate responses for a chat conversation.
The chat conversation is converted into a text prompt using the
tokenizer and calls the :meth:`generate` method to generate the
responses.
Args:
messages: A list of conversations or a single conversation.
- Each conversation is represented as a list of messages.
- Each message is a dictionary with 'role' and 'content' keys.
sampling_params: The sampling parameters for text generation.
chat_template: The template to use for structuring the chat.
If not provided, the model's default chat template will be used.
chat_template_content_format: The format to render message content.
Currently only "string" is supported.
add_generation_prompt: If True, adds a generation template
to each message.
continue_final_message: If True, continues the final message in
the conversation instead of starting a new one. Cannot be
``True`` if ``add_generation_prompt`` is also ``True``.
tools: Optional list of tools to use in the chat.
chat_template_kwargs: Additional keyword arguments to pass to the
chat template rendering function.
Returns:
A list of ``RequestOutput`` objects containing the generated
responses in the same order as the input messages.
"""
if sampling_params is None:
sampling_params = self.get_default_sampling_params()
if continue_final_message and add_generation_prompt:
raise ValueError(
"continue_final_message cannot be True when add_generation_prompt is True."
)
messages_list: list[list[ChatCompletionMessageParam]]
if is_list_of(messages, list):
messages_list = cast(list[list["ChatCompletionMessageParam"]], messages)
else:
messages_list = [cast(list["ChatCompletionMessageParam"], messages)]
_chat_template_kwargs: dict[str, Any] = dict(
chat_template=chat_template,
add_generation_prompt=add_generation_prompt,
continue_final_message=continue_final_message,
tools=tools,
)
_chat_template_kwargs.update(chat_template_kwargs or {})
rendered_prompts = self.tokenizer.apply_chat_template(
messages_list, # type: ignore[arg-type]
tokenize=False,
**_chat_template_kwargs,
)
return self.generate(
rendered_prompts, sampling_params, tokenizer_kwargs={"add_special_tokens": False} # type: ignore
)
[docs]
async def stream_generate(
self,
prompt: str,
sampling_params: SamplingParams | None = None,
prompt_token_ids: BatchEncoding | None = None,
tokenizer_kwargs: dict[str, Any] | None = None,
is_demo: bool = False,
) -> AsyncGenerator[str, None]:
"""Generate texts from given prompt and sampling parameters.
Args:
prompt: The prompt to generate texts. Note that unlike `generate`,
this API supports only a single prompt.
sampling_params: The sampling parameters for generating texts.
prompt_token_ids: Pre-tokenized prompt input as a `BatchEncoding` object.
If not provided, the prompt will be tokenized internally using the tokenizer.
tokenizer_kwargs: Additional keyword arguments passed to the tokenizer's
`encode` method, such as `{"use_special_tokens": True}`.
Returns:
A stream of generated output tokens.
"""
if sampling_params is None:
sampling_params = self.get_default_sampling_params()
if not self.is_generative_model:
raise ValueError("generate API can only be used for generative models.")
if not isinstance(prompt, str):
raise ValueError("prompt must be a single string.")
if prompt_token_ids is None:
if tokenizer_kwargs is None:
tokenizer_kwargs = {}
prompt_token_ids = encode_auto(self.tokenizer, prompt, **tokenizer_kwargs)
input_ids = prompt_token_ids.input_ids
if input_ids and isinstance(input_ids[0], list):
max_prompt_len = max(len(p) for p in input_ids)
else:
max_prompt_len = len(input_ids)
validate_context_length(
max_model_len=self.max_model_len,
prompt_tokens=max_prompt_len,
max_completion_tokens=sampling_params.max_tokens,
min_completion_tokens=max(1, sampling_params.min_tokens),
)
# FIXME: LLM.__init__() should take max_tokens to determine the maximum sequence length through bucket generations
# and use the config value to raise an error.
if is_demo and len(prompt_token_ids.input_ids) > 1024: # type: ignore
raise ValueError("The length of the prompt is larger than 1024 tokens")
output_stream = self._convert_generation_stream(
prompt,
prompt_token_ids,
sampling_params,
RequestOutputKind.DELTA,
)
async for request_output in output_stream:
for completion_output in request_output.outputs:
if completion_output.text:
yield completion_output.text
def _convert_generation_stream(
self,
prompt: str,
prompt_token_ids: BatchEncoding,
sampling_params: SamplingParams,
output_kind: RequestOutputKind,
) -> AsyncGenerator[RequestOutput, None]:
request_id = f"llm-engine-{random_uuid()}"
converter = NativeOutputConverter(
self.tokenizer,
sampling_params.n,
output_kind,
sampling_params.skip_special_tokens,
sampling_params.stop,
sampling_params.include_stop_str_in_output,
self.engine.abort_request,
request_id,
prompt,
prompt_token_ids.input_ids,
min_tokens=sampling_params.min_tokens,
)
native_sampling_params = sampling_params.clone()
native_sampling_params.stop = None
native_output_generator = self.engine.stream_generate(
prompt_token_ids,
native_sampling_params,
request_id,
)
return converter.convert_stream(native_output_generator)
# XXX(n0gu):
# More pooling APIs should be implemented - classify, reward, score.
# However as of 2025.11 only embed API will be used; We support Qwen3-Reranker,
# but this model uses slightly different scoring logic, thus not supported by vLLM's LLM.score() too.
# See: https://huggingface.co/Qwen/Qwen3-Reranker-0.6B#vllm-usage
[docs]
def encode(
self,
prompts: PromptType | Sequence[PromptType],
pooling_params: PoolingParams | Sequence[PoolingParams] | None = None,
*,
pooling_task: PoolingTask | None = None,
) -> list[PoolingRequestOutput]:
"""
Apply pooling to the hidden states corresponding to the input prompts.
Args:
prompts: The prompts to the LLM. You may pass a sequence of prompts
for batch inference.
pooling_params: The pooling parameters for pooling.
pooling_task: Override the pooling task to use.
Returns:
A list of `PoolingRequestOutput` objects containing the
pooled hidden states in the same order as the input prompts.
"""
task_type_from_model = self.model_metadata.task
if task_type_from_model not in POOLING_TASKS:
raise ValueError("Pooling API is not supported by this model.")
if not isinstance(prompts, list):
prompt_list = [cast(PromptType, prompts)]
else:
prompt_list = prompts
coroutines = []
prompt_token_ids = []
for i, prompt in enumerate(prompt_list):
param: PoolingParams
if isinstance(pooling_params, (list, tuple)):
param = pooling_params[i]
elif pooling_params is None:
param = PoolingParams()
elif isinstance(pooling_params, PoolingParams):
param = pooling_params
else:
raise TypeError(
f"pooling_params must be PoolingParams, Sequence[PoolingParams], or None, "
f"got {type(pooling_params)}"
)
batch_encoding, _ = preprocess_prompt(prompt, self.tokenizer)
fit_prompt_to_context(
batch_encoding,
truncate_prompt_tokens=param.truncate_prompt_tokens,
truncation_side=resolve_truncation_side(self.engine),
max_model_len=self.max_model_len,
)
# Set pooling task by precedence (highest to lowest).
# 1. Use the `pooling_task` argument if provided.
# 2. Otherwise, use the task already set in `params.task`.
# 3. If neither is set, infer the task from the model metadata.
from furiosa_llm.utils import coalesce
param.task = coalesce(pooling_task, param.task, cast(PoolingTask, task_type_from_model))
assert param.task is not None, "pooling task must be set at this point."
coroutines.append(
self.engine.encode(
batch_encoding,
param,
None, # TODO: set request id
)
)
prompt_token_ids.append(batch_encoding.input_ids)
request_id = uuid.uuid4().__str__()
from furiosa_llm.utils import async_gather, run_sync
native_outputs_list = run_sync(async_gather(*coroutines))
return [
PoolingRequestOutput(
request_id=request_id,
prompt_token_ids=prompt_token_ids[i],
outputs=PoolingOutput(data=torch.Tensor(native_outputs[0].data)),
finished=True,
)
for i, native_outputs in enumerate(native_outputs_list)
]
[docs]
def embed(
self,
prompts: PromptType | Sequence[PromptType],
pooling_params: PoolingParams | Sequence[PoolingParams] | None = None,
) -> list[EmbeddingRequestOutput]:
"""
Generate an embedding vector for each prompt. Only applicable to embedding models.
Args:
prompts: The prompts to the LLM. You may pass a sequence of prompts
for batch embedding.
pooling_params: The pooling parameters for pooling.
Returns:
A list of `EmbeddingRequestOutput` objects containing the
embedding vectors in the same order as the input prompts.
"""
if "embed" != self.model_metadata.task:
raise ValueError("Embedding API is not supported by this model.")
items = self.encode(
prompts,
pooling_params=pooling_params,
pooling_task="embed",
)
return [EmbeddingRequestOutput.from_base(item) for item in items]
[docs]
def score(
self,
data_1: PromptType | Sequence[PromptType],
data_2: PromptType | Sequence[PromptType],
/,
*,
truncate_prompt_tokens: int | None = None,
pooling_params: PoolingParams | None = None,
chat_template: str | None = None,
) -> list[ScoringRequestOutput]:
"""Generate similarity scores for all pairs `<text,text_pair>`.
The inputs can be `1 -> 1`, `1 -> N` or `N -> N`.
In the `1 - N` case the `data_1` input will be replicated `N`
times to pair with the `data_2` inputs.
Args:
data_1: Can be a single prompt or a list of prompts.
When a list, it must have the same length as the `data_2` list.
data_2: The data to pair with the query to form the input to
the LLM.
truncate_prompt_tokens: The number of tokens to truncate the prompt to.
pooling_params: The pooling parameters for pooling. If None, we
use the default pooling parameters.
chat_template: The chat template to use for the scoring. If None, the model's
own format is used. The template receives exactly two messages, `query` and
`document`, matching vLLM's template contract.
Returns:
A list of `ScoringRequestOutput` objects containing the
generated scores in the same order as the input prompts.
"""
model_metadata = self.model_metadata
if model_metadata.task != "score" and not model_metadata.use_binary_seq_class:
raise ValueError("LLM.score() is only supported for binary classification models.")
# Validate inputs and create pairs
# Convert single prompts to lists for uniform processing
is_data_1_list = isinstance(data_1, list)
is_data_2_list = isinstance(data_2, list)
# Normalize inputs to List[PromptType]
data_1_list: list[PromptType]
data_2_list: list[PromptType]
if not is_data_1_list and not is_data_2_list:
# 1 -> 1 case
data_1_list = [cast(PromptType, data_1)]
data_2_list = [cast(PromptType, data_2)]
elif not is_data_1_list and is_data_2_list:
# 1 -> N case: replicate data_1
data_2_list = cast(list[PromptType], data_2)
data_1_list = [cast(PromptType, data_1)] * len(data_2_list)
elif is_data_1_list and is_data_2_list:
# N -> N case: must have same length
data_1_list = cast(list[PromptType], data_1)
data_2_list = cast(list[PromptType], data_2)
if len(data_1_list) != len(data_2_list):
raise ValueError(
f"When both data_1 and data_2 are lists, they must have the same length. "
f"Got {len(data_1_list)} and {len(data_2_list)}."
)
else:
# data_1 is list, data_2 is not - this is not a standard case for scoring
raise ValueError(
"Invalid input combination. data_1 is a list but data_2 is not. "
"Expected patterns: (single, single), (single, list), or (list, list)."
)
# Normalize inputs to list[str]
data_1_strs: list[str] = [prompt_to_str(d, self.tokenizer) for d in data_1_list]
data_2_strs: list[str] = [prompt_to_str(d, self.tokenizer) for d in data_2_list]
# Set up pooling parameters with truncation if specified
if pooling_params is None:
pooling_params = PoolingParams()
if truncate_prompt_tokens is not None:
pooling_params.truncate_prompt_tokens = truncate_prompt_tokens
tokenization_kwargs = resolve_score_truncation_kwargs(
truncate_prompt_tokens=pooling_params.truncate_prompt_tokens,
max_model_len=self.max_model_len,
)
# Construct prompts for each pair
prompts: list[PromptType] = []
for str_1, str_2 in zip(data_1_strs, data_2_strs):
_, prompt = get_score_prompt(
self.tokenizer,
str_1,
str_2,
hf_configs=model_metadata.hf_configs,
score_template=chat_template,
tokenization_kwargs=tokenization_kwargs,
)
prompts.append(prompt)
# Call encode with the constructed prompts
items = self.encode(
prompts,
pooling_params=pooling_params,
pooling_task="score",
)
return [ScoringRequestOutput.from_base(item) for item in items]
[docs]
def shutdown(self):
"""Shutdown the LLM engine gracefully. Idempotent."""
finalizer = getattr(self, "_finalizer", None)
if finalizer is not None:
finalizer()
def __enter__(self) -> "LLM":
return self
def __exit__(self, exc_type, exc_val, exc_tb) -> None:
self.shutdown()
