# 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.mm_options import validate_mm_processor_kwargs
from furiosa_llm.mm_types import MediaId
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 .speculative_args import (
SpeculativeConfigLike,
SpeculativeMethod,
resolve_speculative_config,
)
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)
def _load_fxb_hf_config(model_dir: str | os.PathLike) -> Any:
"""Load serving metadata even when Transformers lacks the FXB architecture.
FXB execution only needs the serialized config fields; its model implementation
lives in the native runtime. Permit a newly compiled architecture before its config
class reaches the pinned Transformers release. Registered-config errors remain intact.
"""
from transformers import AutoConfig
config_path = Path(model_dir) / "config.json"
try:
config = json.loads(config_path.read_text(encoding="utf-8"))
except (OSError, ValueError):
return AutoConfig.from_pretrained(model_dir)
model_type = config.get("model_type")
if not isinstance(model_type, str) or not model_type:
return AutoConfig.from_pretrained(model_dir)
try:
AutoConfig.for_model(model_type)
except (KeyError, ValueError):
pass
else:
return AutoConfig.from_pretrained(model_dir)
from transformers import PretrainedConfig
logger.warning(
"Transformers has no matching config class for FXB model_type '%s'; "
"using config.json metadata",
model_type,
)
metadata_config_type = type(
"_FxbMetadataConfig",
(PretrainedConfig,),
{"model_type": model_type},
)
return cast(type[PretrainedConfig], metadata_config_type).from_dict(config)
# 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: MediaId) -> list[int] | None: ...
async def encode_vision(
self,
media_ids: list[MediaId],
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, target: str = "all") -> 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_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,
speculative_config: SpeculativeConfigLike | None = None,
speculative_method: SpeculativeMethod | 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,
npu_io_cpus: str | None = None,
ssm_state_cache_size: int | None = None,
ssm_full_memory_ratio: float | 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",
served_model_name: str | None = None,
kv_transfer_config: dict[str, Any] | None = None,
mm_processor_kwargs: dict[str, Any] | None = None,
**kwargs,
) -> None:
"""Instantiate LLM from a Hugging Face model and a compatible FXB.
``model_id_or_path`` is resolved through the Hugging Face Hub when needed and provides
the model weights, configuration, and tokenizer. When ``fxb`` is omitted, a compatible
``.fxb`` file is discovered from the model directory or the local FXB cache.
Args:
model_id_or_path: A path to furiosa llm engine artifact or a HuggingFace model id.
fxb: Optional path to a compatible FXB. When omitted, FXB files are discovered
from the model directory or the local FXB cache.
speculative_config: Canonical vLLM-compatible speculative configuration.
It cannot be combined with the flat speculative arguments.
speculative_method: Explicit method for the flat speculative arguments: ``draft_model``
or ``mtp``.
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. 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.
The flat API also requires ``speculative_method``.
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 the loaded FXB.
npu_io_cpus: One Linux logical CPU list for all NPU runtime I/O threads,
e.g. "8-20" or "8-12,16,18". CPUs are assigned in list order to selected
clusters in ascending NPU/cluster order. The list must contain at least
as many distinct CPUs as selected clusters. The next CPU, if available,
pins the shared I/O completion thread (reaper). Remaining CPU IDs in the
supplied list are not used for pinning and are logged at debug level.
None leaves CPU placement unchanged.
This does not set a NUMA memory policy.
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.
ssm_state_cache_size: Number of state cache slots for SSM
models; ignored for softmax-only models. If unspecified, the pool
is sized by `ssm_full_memory_ratio`.
ssm_full_memory_ratio: Size the SSM state cache pool relative to the
full-attention KV cache instead of a fixed slot count (state : KV = ratio : 1).
Mutually exclusive with `ssm_state_cache_size`. Defaults to 1.0
when neither option is specified.
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 covered by the loaded
FXB is rejected before any device is opened. If not given, the loaded FXB's limit
is used.
scheduler_config: Configuration for the scheduler, including the maximum number of tasks
queued to hardware and the maximum number of processing samples. If omitted, the
default scheduler configuration is 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.
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.
mm_processor_kwargs: Processor defaults, e.g. {"min_pixels": 1024}, with optional
HF-style images_kwargs. Request keys replace top-level defaults,
including whole images_kwargs objects. An option must not appear both at the top level and in images_kwargs.
"""
if mm_processor_kwargs is not None:
mm_processor_kwargs = validate_mm_processor_kwargs(mm_processor_kwargs)
self.mm_processor_kwargs = mm_processor_kwargs
# 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}")
resolved_speculative_config = resolve_speculative_config(
speculative_config,
speculative_method=speculative_method,
draft_model_id_or_path=draft_model_id_or_path,
draft_fxb_path=draft_fxb_path,
num_speculative_tokens=num_speculative_tokens,
)
if resolved_speculative_config is None:
speculative_method = None
draft_model_id_or_path = None
draft_fxb_path = None
num_speculative_tokens = None
else:
speculative_method = resolved_speculative_config.method
draft_model_id_or_path = resolved_speculative_config.model
draft_fxb_path = resolved_speculative_config.draft_fxb_path
num_speculative_tokens = resolved_speculative_config.num_speculative_tokens
self.num_speculative_tokens = num_speculative_tokens
# Resolve the model snapshot and compatible FXB. An explicit FXB takes precedence;
# otherwise, discover one from the model directory or local FXB cache.
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. 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,
speculative_method=speculative_method,
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,
npu_io_cpus=npu_io_cpus,
ssm_state_cache_size=ssm_state_cache_size,
ssm_full_memory_ratio=ssm_full_memory_ratio,
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,
speculative_method: SpeculativeMethod | 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,
npu_io_cpus: str | None = None,
ssm_state_cache_size: int | None = None,
ssm_full_memory_ratio: float | 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 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 = _load_fxb_hf_config(model_id_or_path)
hf_configs = hf_config.to_dict()
if self.mm_processor_kwargs is not None:
self.mm_processor_kwargs = validate_mm_processor_kwargs(
self.mm_processor_kwargs, model_type=hf_config.model_type
)
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,
speculative_method,
num_speculative_tokens,
self.served_model_name,
fxb_path,
draft_fxb,
self._serialize_obj(kv_transfer_config) if kv_transfer_config else None,
max_model_len,
npu_io_cpus=npu_io_cpus,
ssm_state_cache_size=ssm_state_cache_size,
ssm_full_memory_ratio=ssm_full_memory_ratio,
)
# 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()
