87 lines
3.3 KiB
Python
87 lines
3.3 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""One sampler registration contract for token and post-sampling audio output."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Callable
|
|
from dataclasses import dataclass
|
|
from typing import Any
|
|
|
|
import torch
|
|
from vllm.v1.core.sched.output import GrammarOutput
|
|
from vllm.v1.worker.gpu.input_batch import InputBatch
|
|
from vllm.v1.worker.gpu.sample.output import SamplerOutput
|
|
from vllm.v1.worker.gpu.sample.sampler import Sampler
|
|
from vllm.v1.worker.gpu.states import RequestState
|
|
|
|
from vllm_omni.worker_v2.output_snapshot import RequestOutputSnapshot
|
|
|
|
StandardSample = Callable[
|
|
[torch.Tensor, InputBatch, GrammarOutput | None], tuple[SamplerOutput, torch.Tensor, torch.Tensor]
|
|
]
|
|
|
|
|
|
@dataclass
|
|
class OmniSamplingOutput:
|
|
sampler_output: SamplerOutput
|
|
num_sampled: torch.Tensor
|
|
num_rejected: torch.Tensor
|
|
# None preserves forward's payload. An empty dict explicitly replaces it.
|
|
# Tensors must remain valid until the runner finishes its output copy.
|
|
multimodal_outputs: dict[str, Any] | None = None
|
|
include_hidden_states: bool = True
|
|
owns_multimodal_outputs: bool = False
|
|
# Invoked on CPU-owned snapshots after D2H, with effective per-request counts.
|
|
finalize_multimodal: Callable[[dict[str, Any], list[int]], dict[str, Any] | RequestOutputSnapshot] | None = None
|
|
|
|
|
|
class OmniSampler:
|
|
"""Register through ModelState.custom_sampler, just like an upstream sampler.
|
|
|
|
Delegate request lifecycle/state to the original sampler without copying
|
|
its attributes. Models can override __call__ for logits-only adjustments,
|
|
or sample_step when sampling also produces audio. Both use one registered
|
|
sampler and one result contract; there is no separate ModelState step hook.
|
|
"""
|
|
|
|
def __init__(self, base_sampler: Sampler):
|
|
self.base_sampler = base_sampler
|
|
|
|
@property
|
|
def omni_static_staged_writes(self) -> bool:
|
|
# Preserve the runner's no-new-requests fast path for the stock sampler.
|
|
return type(self.base_sampler) is Sampler or bool(
|
|
getattr(self.base_sampler, "omni_static_staged_writes", False)
|
|
)
|
|
|
|
def __getattr__(self, name: str) -> Any:
|
|
return getattr(self.base_sampler, name)
|
|
|
|
def __call__(self, logits: torch.Tensor, input_batch: InputBatch) -> SamplerOutput:
|
|
return self.base_sampler(logits, input_batch)
|
|
|
|
def sample_step(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
input_batch: InputBatch,
|
|
req_states: RequestState,
|
|
grammar_output: GrammarOutput | None,
|
|
standard_sample: StandardSample,
|
|
) -> OmniSamplingOutput:
|
|
# standard_sample calls this adapter's __call__ through the runner.
|
|
return OmniSamplingOutput(*standard_sample(hidden_states, input_batch, grammar_output))
|
|
|
|
|
|
def sample_with_output(
|
|
sampler: Sampler | OmniSampler | None,
|
|
standard_sample: StandardSample,
|
|
hidden_states: torch.Tensor,
|
|
input_batch: InputBatch,
|
|
req_states: RequestState,
|
|
grammar_output: GrammarOutput | None,
|
|
) -> OmniSamplingOutput:
|
|
if isinstance(sampler, OmniSampler):
|
|
return sampler.sample_step(hidden_states, input_batch, req_states, grammar_output, standard_sample)
|
|
return OmniSamplingOutput(*standard_sample(hidden_states, input_batch, grammar_output))
|