Skip to content

vllm_omni.worker_v2.omni_sampler

One sampler registration contract for token and post-sampling audio output.

StandardSample module-attribute 

StandardSample = Callable[
    [torch.Tensor, InputBatch, GrammarOutput | None],
    tuple[SamplerOutput, torch.Tensor, torch.Tensor],
]

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.

base_sampler instance-attribute 

base_sampler = base_sampler

omni_static_staged_writes property 

omni_static_staged_writes: bool

sample_step 

sample_step(
    hidden_states: Tensor,
    input_batch: InputBatch,
    req_states: RequestState,
    grammar_output: GrammarOutput | None,
    standard_sample: StandardSample,
) -> OmniSamplingOutput

OmniSamplingOutput dataclass

finalize_multimodal class-attribute instance-attribute 

finalize_multimodal: (
    Callable[
        [dict[str, Any], list[int]],
        dict[str, Any] | RequestOutputSnapshot,
    ]
    | None
) = None

include_hidden_states class-attribute instance-attribute 

include_hidden_states: bool = True

multimodal_outputs class-attribute instance-attribute 

multimodal_outputs: dict[str, Any] | None = None

num_rejected instance-attribute 

num_rejected: Tensor

num_sampled instance-attribute 

num_sampled: Tensor

owns_multimodal_outputs class-attribute instance-attribute 

owns_multimodal_outputs: bool = False

sampler_output instance-attribute 

sampler_output: SamplerOutput

sample_with_output 

sample_with_output(
    sampler: Sampler | OmniSampler | None,
    standard_sample: StandardSample,
    hidden_states: Tensor,
    input_batch: InputBatch,
    req_states: RequestState,
    grammar_output: GrammarOutput | None,
) -> OmniSamplingOutput