vllm-cpu-amxbf16 0.11.2.post2__cp310-cp310-manylinux_2_17_x86_64.whl

vllm/v1/pool/metadata.py

@@ -0,0 +1,82 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from dataclasses import dataclass
+
+import torch
+
+from vllm.pooling_params import PoolingParams
+from vllm.utils.platform_utils import is_pin_memory_available
+
+pin_memory = is_pin_memory_available()
+
+
+@dataclass
+class PoolingCursor:
+    index: list[int]
+    first_token_indices_gpu: torch.Tensor
+    last_token_indices_gpu: torch.Tensor
+    prompt_lens_cpu: torch.Tensor
+    num_scheduled_tokens_cpu: torch.Tensor
+
+    def __getitem__(self, indices: slice):
+        return PoolingCursor(
+            index=self.index[indices],
+            first_token_indices_gpu=self.first_token_indices_gpu[indices],
+            last_token_indices_gpu=self.last_token_indices_gpu[indices],
+            prompt_lens_cpu=self.prompt_lens_cpu[indices],
+            num_scheduled_tokens_cpu=self.num_scheduled_tokens_cpu[indices],
+        )
+
+    def is_partial_prefill(self):
+        return not torch.all(self.prompt_lens_cpu == self.num_scheduled_tokens_cpu)
+
+
+@dataclass
+class PoolingMetadata:
+    """Tensors for pooling."""
+
+    prompt_lens: torch.Tensor  # CPU Tensor
+    prompt_token_ids: torch.Tensor | None
+    pooling_params: list[PoolingParams]
+    pooling_cursor: PoolingCursor | None = None
+
+    def __getitem__(self, indices: slice):
+        return PoolingMetadata(
+            prompt_lens=self.prompt_lens[indices],
+            prompt_token_ids=None
+            if self.prompt_token_ids is None
+            else self.prompt_token_ids[indices],
+            pooling_params=self.pooling_params[indices],
+            pooling_cursor=None
+            if self.pooling_cursor is None
+            else self.pooling_cursor[indices],
+        )
+
+    def build_pooling_cursor(
+        self, num_scheduled_tokens: list[int], device: torch.device
+    ):
+        self.pooling_cursor = build_pooling_cursor(
+            num_scheduled_tokens, self.prompt_lens, device
+        )
+
+
+def build_pooling_cursor(
+    num_scheduled_tokens: list[int], prompt_lens: torch.Tensor, device: torch.device
+):
+    assert len(prompt_lens) == len(num_scheduled_tokens)
+
+    n_seq = len(num_scheduled_tokens)
+    index = list(range(n_seq))
+    num_scheduled_tokens = torch.tensor(num_scheduled_tokens, device="cpu")
+    cumsum = torch.zeros(
+        n_seq + 1, dtype=torch.int64, pin_memory=pin_memory, device="cpu"
+    )
+    torch.cumsum(num_scheduled_tokens, dim=0, out=cumsum[1:])
+    cumsum = cumsum.to(device, non_blocking=True)
+    return PoolingCursor(
+        index=index,
+        first_token_indices_gpu=cumsum[:n_seq],
+        last_token_indices_gpu=cumsum[1:] - 1,
+        prompt_lens_cpu=prompt_lens,
+        num_scheduled_tokens_cpu=num_scheduled_tokens,
+    )

vllm/v1/request.py

@@ -0,0 +1,259 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import enum
+import time
+from collections.abc import Callable, Mapping
+from functools import partial
+from typing import TYPE_CHECKING, Any, Optional
+
+import torch
+
+from vllm.multimodal.inputs import MultiModalFeatureSpec
+from vllm.pooling_params import PoolingParams
+from vllm.sampling_params import SamplingParams
+from vllm.utils import length_from_prompt_token_ids_or_embeds
+from vllm.v1.engine import (
+    EngineCoreEvent,
+    EngineCoreEventType,
+    EngineCoreRequest,
+    FinishReason,
+)
+from vllm.v1.structured_output.request import StructuredOutputRequest
+from vllm.v1.utils import ConstantList
+
+if TYPE_CHECKING:
+    from vllm.lora.request import LoRARequest
+    from vllm.v1.core.kv_cache_utils import BlockHash
+
+
+class Request:
+    def __init__(
+        self,
+        request_id: str,
+        prompt_token_ids: list[int] | None,
+        sampling_params: SamplingParams | None,
+        pooling_params: PoolingParams | None,
+        eos_token_id: int | None,
+        client_index: int = 0,
+        arrival_time: float | None = None,
+        prompt_embeds: torch.Tensor | None = None,
+        mm_features: list[MultiModalFeatureSpec] | None = None,
+        lora_request: Optional["LoRARequest"] = None,
+        cache_salt: str | None = None,
+        priority: int = 0,
+        trace_headers: Mapping[str, str] | None = None,
+        block_hasher: Callable[["Request"], list["BlockHash"]] | None = None,
+    ) -> None:
+        self.request_id = request_id
+        self.client_index = client_index
+        self.priority = priority
+        self.sampling_params = sampling_params
+        self.pooling_params = pooling_params
+        # Because of LoRA, the eos token id can be different for each request.
+        self.eos_token_id = eos_token_id
+        self.lora_request = lora_request
+        self.structured_output_request = StructuredOutputRequest.from_sampling_params(
+            sampling_params
+        )
+        self.arrival_time = arrival_time if arrival_time is not None else time.time()
+
+        self.status = RequestStatus.WAITING
+        self.events: list[EngineCoreEvent] = []
+        self.stop_reason: int | str | None = None
+
+        # P/D: Connector-specific KV transfer parameters.
+        self.kv_transfer_params: dict[str, Any] | None = None
+
+        if pooling_params is not None:
+            # Pooling models.
+            self.max_tokens = 1
+        elif sampling_params is not None:
+            # Generative models.
+            assert sampling_params.max_tokens is not None
+            self.max_tokens = sampling_params.max_tokens
+            if self.structured_output_request is not None:
+                self.status = RequestStatus.WAITING_FOR_FSM
+
+            if sampling_params.extra_args is not None:
+                self.kv_transfer_params = sampling_params.extra_args.get(
+                    "kv_transfer_params"
+                )
+        else:
+            raise ValueError("sampling_params and pooling_params can't both be unset")
+
+        self.prompt_token_ids = prompt_token_ids
+        self.prompt_embeds = prompt_embeds
+        self.num_prompt_tokens = length_from_prompt_token_ids_or_embeds(
+            prompt_token_ids, prompt_embeds
+        )
+        self._output_token_ids: list[int] = []
+        self._all_token_ids: list[int] = (
+            self.prompt_token_ids.copy()
+            if self.prompt_token_ids is not None
+            else [0] * self.num_prompt_tokens
+        )
+        self.num_output_placeholders = 0  # Used in async scheduling.
+        self.spec_token_ids: list[int] = []
+        self.num_computed_tokens = 0
+        self.cache_salt: str | None = cache_salt
+
+        # Multi-modal related
+        self.mm_features = mm_features or []
+        self.num_encoder_inputs = len(self.mm_features)
+        self.has_encoder_inputs = self.num_encoder_inputs > 0
+
+        # Read-only views
+        # Prevent directly appending to these lists since
+        # they should also be updated simultaneously.
+        self.output_token_ids = ConstantList(self._output_token_ids)
+        self.all_token_ids = ConstantList(self._all_token_ids)
+        # trace_headers
+        self.trace_headers = trace_headers
+        # State
+        # The number of tokens with prefix cache hits.
+        self.num_cached_tokens = -1
+
+        # The number of NaNs in logits. A value greater than 0
+        # indicates that the output is corrupted
+        self.num_nans_in_logits = 0
+
+        # The number of requests being preempted by the scheduler
+        self.num_preemptions = 0
+
+        self.block_hashes: list[BlockHash] = []
+        self.get_hash_new_full_blocks: Callable[[], list[BlockHash]] | None = None
+        if block_hasher is not None:
+            self.get_hash_new_full_blocks = partial(block_hasher, self)
+            self.block_hashes = self.get_hash_new_full_blocks()
+
+        self.skip_reading_prefix_cache = self.get_skip_reading_prefix_cache()
+
+    @classmethod
+    def from_engine_core_request(
+        cls,
+        request: EngineCoreRequest,
+        block_hasher: Callable[["Request"], list["BlockHash"]] | None,
+    ) -> "Request":
+        return cls(
+            request_id=request.request_id,
+            client_index=request.client_index,
+            prompt_token_ids=request.prompt_token_ids,
+            prompt_embeds=request.prompt_embeds,
+            mm_features=request.mm_features,
+            sampling_params=request.sampling_params,
+            pooling_params=request.pooling_params,
+            eos_token_id=request.eos_token_id,
+            arrival_time=request.arrival_time,
+            lora_request=request.lora_request,
+            cache_salt=request.cache_salt,
+            priority=request.priority,
+            trace_headers=request.trace_headers,
+            block_hasher=block_hasher,
+        )
+
+    def append_output_token_ids(
+        self,
+        token_ids: int | list[int],
+    ) -> None:
+        if isinstance(token_ids, int):
+            self._output_token_ids.append(token_ids)
+            self._all_token_ids.append(token_ids)
+        else:
+            self._output_token_ids.extend(token_ids)
+            self._all_token_ids.extend(token_ids)
+
+        if self.get_hash_new_full_blocks is not None:
+            self.block_hashes.extend(self.get_hash_new_full_blocks())
+
+    @property
+    def use_structured_output(self) -> bool:
+        return self.structured_output_request is not None
+
+    @property
+    def num_tokens(self) -> int:
+        return len(self._all_token_ids)
+
+    @property
+    def num_tokens_with_spec(self) -> int:
+        return len(self._all_token_ids) + len(self.spec_token_ids)
+
+    @property
+    def num_output_tokens(self) -> int:
+        return len(self._output_token_ids)
+
+    def get_skip_reading_prefix_cache(self) -> bool:
+        if (
+            self.sampling_params is not None
+            and self.sampling_params.skip_reading_prefix_cache is not None
+        ):
+            return self.sampling_params.skip_reading_prefix_cache
+        elif (
+            self.pooling_params is not None
+            and self.pooling_params.skip_reading_prefix_cache is not None
+        ):
+            return self.pooling_params.skip_reading_prefix_cache
+        return False
+
+    def is_finished(self) -> bool:
+        return RequestStatus.is_finished(self.status)
+
+    def get_finished_reason(self) -> FinishReason | None:
+        return RequestStatus.get_finished_reason(self.status)
+
+    def get_num_encoder_tokens(self, input_id: int) -> int:
+        assert input_id < len(self.mm_features)
+        num_tokens = self.mm_features[input_id].mm_position.length
+        return num_tokens
+
+    def record_event(
+        self,
+        event_type: EngineCoreEventType,
+        timestamp: float | None = None,
+    ) -> None:
+        self.events.append(EngineCoreEvent.new_event(event_type, timestamp))
+
+    def take_events(self) -> list[EngineCoreEvent] | None:
+        if not self.events:
+            return None
+        events, self.events = self.events, []
+        return events
+
+
+class RequestStatus(enum.IntEnum):
+    """Status of a request."""
+
+    WAITING = enum.auto()
+    WAITING_FOR_FSM = enum.auto()
+    WAITING_FOR_REMOTE_KVS = enum.auto()
+    RUNNING = enum.auto()
+    PREEMPTED = enum.auto()
+    # Note: anything after PREEMPTED will be considered
+    # as a finished status.
+    FINISHED_STOPPED = enum.auto()
+    FINISHED_LENGTH_CAPPED = enum.auto()
+    FINISHED_ABORTED = enum.auto()
+    FINISHED_IGNORED = enum.auto()
+
+    def __str__(self):
+        return self.name
+
+    @staticmethod
+    def is_finished(status: "RequestStatus") -> bool:
+        return status > RequestStatus.PREEMPTED
+
+    @staticmethod
+    def get_finished_reason(status: "RequestStatus") -> FinishReason | None:
+        return _FINISHED_REASON_MAP.get(status)
+
+
+# Mapping of finished statuses to their finish reasons.
+# NOTE: The ignored requests are the requests whose prompt lengths
+# are longer than the model's length cap. Therefore, the stop
+# reason should also be "length" as in OpenAI API.
+_FINISHED_REASON_MAP = {
+    RequestStatus.FINISHED_STOPPED: FinishReason.STOP,
+    RequestStatus.FINISHED_LENGTH_CAPPED: FinishReason.LENGTH,
+    RequestStatus.FINISHED_ABORTED: FinishReason.ABORT,
+    RequestStatus.FINISHED_IGNORED: FinishReason.LENGTH,
+}

vllm/v1/sample/logits_processor/__init__.py

@@ -0,0 +1,352 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import importlib
+import inspect
+import itertools
+from abc import abstractmethod
+from collections.abc import Sequence
+from functools import lru_cache, partial
+from typing import TYPE_CHECKING
+
+import torch
+
+from vllm.logger import init_logger
+from vllm.logits_process import LogitsProcessor as RequestLogitsProcessor
+from vllm.sampling_params import SamplingParams
+from vllm.utils.torch_utils import guard_cuda_initialization
+from vllm.v1.sample.logits_processor.builtin import (
+    LogitBiasLogitsProcessor,
+    MinPLogitsProcessor,
+    MinTokensLogitsProcessor,
+    process_dict_updates,
+)
+from vllm.v1.sample.logits_processor.interface import (
+    BatchUpdate,
+    LogitsProcessor,
+    MoveDirectionality,
+)
+from vllm.v1.sample.logits_processor.state import BatchUpdateBuilder, LogitsProcessors
+
+if TYPE_CHECKING:
+    from vllm.config import VllmConfig
+
+logger = init_logger(__name__)
+
+# Error message when the user tries to initialize vLLM with a pooling model
+# and custom logitsproces
+STR_POOLING_REJECTS_LOGITSPROCS = (
+    "Pooling models do not support custom logits processors."
+)
+
+# Error message when the user tries to initialize vLLM with a speculative
+# decoding enabled and custom logitsproces
+STR_SPEC_DEC_REJECTS_LOGITSPROCS = (
+    "Custom logits processors are not supported when speculative decoding is enabled."
+)
+
+LOGITSPROCS_GROUP = "vllm.logits_processors"
+
+BUILTIN_LOGITS_PROCESSORS: list[type[LogitsProcessor]] = [
+    MinTokensLogitsProcessor,
+    LogitBiasLogitsProcessor,
+    MinPLogitsProcessor,
+]
+
+
+def _load_logitsprocs_plugins() -> list[type[LogitsProcessor]]:
+    """Load all installed logit processor plugins"""
+
+    from importlib.metadata import entry_points
+
+    installed_logitsprocs_plugins = entry_points(group=LOGITSPROCS_GROUP)
+    if len(installed_logitsprocs_plugins) == 0:
+        logger.debug("No logitsprocs plugins installed (group %s).", LOGITSPROCS_GROUP)
+        return []
+
+    # Load logitsprocs plugins
+    logger.debug("Loading installed logitsprocs plugins (group %s):", LOGITSPROCS_GROUP)
+    classes: list[type[LogitsProcessor]] = []
+    for entrypoint in installed_logitsprocs_plugins:
+        try:
+            logger.debug(
+                "- Loading logitproc plugin entrypoint=%s target=%s",
+                entrypoint.name,
+                entrypoint.value,
+            )
+            with guard_cuda_initialization():
+                classes.append(entrypoint.load())
+        except Exception as e:
+            logger.error("Failed to load LogitsProcessor plugin %s: %s", entrypoint, e)
+            raise RuntimeError(
+                f"Failed to load LogitsProcessor plugin {entrypoint}"
+            ) from e
+    return classes
+
+
+def _load_logitsprocs_by_fqcns(
+    logits_processors: Sequence[str | type[LogitsProcessor]] | None,
+) -> list[type[LogitsProcessor]]:
+    """Load logit processor types, identifying them by fully-qualified class
+    names (FQCNs).
+
+    Effectively, a mixed list of logitproc types and FQCN strings is converted
+    into a list of entirely logitproc types, by loading from the FQCNs.
+
+    FQCN syntax is <module>:<type> i.e. x.y.z:CustomLogitProc
+
+    Already-loaded logitproc types must be subclasses of LogitsProcessor
+
+    Args:
+      logits_processors: Potentially mixed list of logitsprocs types and FQCN
+                         strings for logitproc types
+
+    Returns:
+      List of logitproc types
+
+    """
+    if not logits_processors:
+        return []
+
+    logger.debug(
+        "%s additional custom logits processors specified, checking whether "
+        "they need to be loaded.",
+        len(logits_processors),
+    )
+
+    classes: list[type[LogitsProcessor]] = []
+    for ldx, logitproc in enumerate(logits_processors):
+        if isinstance(logitproc, type):
+            logger.debug(" - Already-loaded logit processor: %s", logitproc.__name__)
+            if not issubclass(logitproc, LogitsProcessor):
+                raise ValueError(
+                    f"{logitproc.__name__} is not a subclass of LogitsProcessor"
+                )
+            classes.append(logitproc)
+            continue
+
+        logger.debug("- Loading logits processor %s", logitproc)
+        module_path, qualname = logitproc.split(":")
+
+        try:
+            # Load module
+            with guard_cuda_initialization():
+                module = importlib.import_module(module_path)
+        except Exception as e:
+            logger.error(
+                "Failed to load %sth LogitsProcessor plugin %s: %s",
+                ldx,
+                logitproc,
+                e,
+            )
+            raise RuntimeError(
+                f"Failed to load {ldx}th LogitsProcessor plugin {logitproc}"
+            ) from e
+
+        # Walk down dotted name to get logitproc class
+        obj = module
+        for attr in qualname.split("."):
+            obj = getattr(obj, attr)
+        if not isinstance(obj, type):
+            raise ValueError("Loaded logit processor must be a type.")
+        if not issubclass(obj, LogitsProcessor):
+            raise ValueError(f"{obj.__name__} must be a subclass of LogitsProcessor")
+        classes.append(obj)
+
+    return classes
+
+
+def _load_custom_logitsprocs(
+    logits_processors: Sequence[str | type[LogitsProcessor]] | None,
+) -> list[type[LogitsProcessor]]:
+    """Load all custom logits processors.
+
+    * First load all installed logitproc plugins
+    * Second load custom logitsprocs pass by the user at initialization time
+
+    Args:
+      logits_processors: potentially mixed list of logitproc types and
+                         logitproc type fully-qualified names (FQCNs)
+                         which need to be loaded
+
+    Returns:
+      A list of all loaded logitproc types
+    """
+    from vllm.platforms import current_platform
+
+    if current_platform.is_tpu():
+        # No logitsprocs specified by caller
+        # TODO(andy) - vLLM V1 on TPU does not support custom logitsprocs
+        return []
+
+    return _load_logitsprocs_plugins() + _load_logitsprocs_by_fqcns(logits_processors)
+
+
+def build_logitsprocs(
+    vllm_config: "VllmConfig",
+    device: torch.device,
+    is_pin_memory: bool,
+    is_pooling_model: bool,
+    custom_logitsprocs: Sequence[str | type[LogitsProcessor]] = (),
+) -> LogitsProcessors:
+    if is_pooling_model:
+        if custom_logitsprocs:
+            raise ValueError(STR_POOLING_REJECTS_LOGITSPROCS)
+        logger.debug(
+            "Skipping logits processor loading because pooling models"
+            " do not support logits processors."
+        )
+        return LogitsProcessors()
+
+    # Check if speculative decoding is enabled.
+    if vllm_config.speculative_config:
+        if custom_logitsprocs:
+            raise ValueError(STR_SPEC_DEC_REJECTS_LOGITSPROCS)
+        logger.warning(
+            "min_p, logit_bias, and min_tokens parameters won't currently work "
+            "with speculative decoding enabled."
+        )
+        return LogitsProcessors()
+
+    custom_logitsprocs_classes = _load_custom_logitsprocs(custom_logitsprocs)
+    return LogitsProcessors(
+        ctor(vllm_config, device, is_pin_memory)
+        for ctor in itertools.chain(
+            BUILTIN_LOGITS_PROCESSORS, custom_logitsprocs_classes
+        )
+    )
+
+
+cached_load_custom_logitsprocs = lru_cache(_load_custom_logitsprocs)
+
+
+def validate_logits_processors_parameters(
+    logits_processors: Sequence[str | type[LogitsProcessor]] | None,
+    sampling_params: SamplingParams,
+):
+    logits_processors = (
+        tuple(logits_processors) if logits_processors is not None else None
+    )
+    for logits_procs in cached_load_custom_logitsprocs(logits_processors):
+        logits_procs.validate_params(sampling_params)
+
+
+class AdapterLogitsProcessor(LogitsProcessor):
+    """Wrapper for per-request logits processors
+
+    To wrap a specific per-request logits processor,
+    * Subclass `AdapterLogitsProcessor`
+    * Implement `self.is_argmax_invariant()` base-class method
+    * Implement `self.new_req_logits_processor(params)`
+
+    `self.__init__(vllm_config, device, is_pin_memory)` does not need to be
+    overridden in general. However, to implement custom constructor behavior -
+    especially any logic which operates on or stores `vllm_config`, `device`,
+    or `is_pin_memory` - `self.__init__(vllm_config, device, is_pin_memory)`
+    must be overridden and the override must call
+    `super().__init__(vllm_config, device, is_pin_memory)`
+    """
+
+    def __init__(
+        self, vllm_config: "VllmConfig", device: torch.device, is_pin_memory: bool
+    ):
+        """Subclass must invoke
+        `super().__init__(vllm_config, device, is_pin_memory)`.
+
+        Subclass constructor may find it useful to utilize the `vllm_config`,
+        `device` and `is_pin_memory` argument. However regardless of whether
+        these arguments are used, the vLLM logits processor interface requires
+        all three arguments to be present.
+        """
+
+        # Map req index -> logits processor state
+        #
+        # State representation is a partial[Tensor] comprising a request-level
+        # logits processor with the output token ids argument and (if required)
+        # the prompt token ids argument pre-populated
+        #
+        # Note that the partial carries a *reference* to output token ids, and
+        # will thus always operate on the list as it is currently, not as it
+        # was when the partial was created.
+        self.req_info: dict[int, partial[torch.Tensor]] = {}
+
+    @abstractmethod
+    def new_req_logits_processor(
+        self,
+        params: SamplingParams,
+    ) -> RequestLogitsProcessor | None:
+        """Consume request info; return a per-request logits processor.
+
+        Return None if logits processor does not need to be applied to request
+
+        Args:
+          params: request sampling params
+
+        Returns:
+          None if logits processor should not be applied to request; otherwise
+          returns a `RequestLogitsProcessor` instance
+
+        """
+        raise NotImplementedError
+
+    def _new_state(
+        self,
+        params: SamplingParams,
+        prompt_ids: list[int] | None,
+        output_ids: list[int],
+    ) -> partial[torch.Tensor] | None:
+        """Return state representation for new request
+
+        Returns None if logits processor is not applicable to request
+
+        Args:
+          params: request sampling params
+          prompt_ids: request prompt token ids
+          output_ids: decoded tokens so far for this request
+
+        Returns:
+          logits processor partial[Tensor] or None
+
+        """
+        if req_lp := self.new_req_logits_processor(params):
+            args = (
+                [prompt_ids, output_ids]
+                if (len(inspect.signature(req_lp).parameters) == 3)
+                else [output_ids]
+            )
+            return partial(req_lp, *args)
+        return None
+
+    def update_state(self, batch_update: BatchUpdate | None):
+        process_dict_updates(
+            self.req_info,
+            batch_update,
+            self._new_state,
+        )
+
+    def apply(self, logits: torch.Tensor) -> torch.Tensor:
+        if self.req_info:
+            # Apply per-request logits processors to corresponding rows of
+            # logits tensor
+            for req_idx, req_lp in self.req_info.items():
+                req_logits = logits[req_idx]
+                new_logits = req_lp(req_logits)
+                if new_logits is not req_logits:
+                    # Modify logits tensor row in-place if necessary
+                    logits[req_idx] = new_logits
+        return logits
+
+
+__all__ = [
+    "LogitsProcessor",
+    "LogitBiasLogitsProcessor",
+    "MinPLogitsProcessor",
+    "MinTokensLogitsProcessor",
+    "BatchUpdate",
+    "BatchUpdateBuilder",
+    "MoveDirectionality",
+    "LogitsProcessors",
+    "build_logitsprocs",
+    "STR_POOLING_REJECTS_LOGITSPROCS",
+    "LOGITSPROCS_GROUP",
+    "AdapterLogitsProcessor",
+]