quan-yi-aiandHongsheng Liu
65f5a1d5b3
Signed-off-by: quan-yi-ai <82917002+quan-yi-ai@users.noreply.github.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
43 lines
1.7 KiB
Python
43 lines
1.7 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
from __future__ import annotations
|
|
|
|
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
|
|
|
|
|
def apply_declared_extra_args(
|
|
sampling_params: OmniDiffusionSamplingParams,
|
|
declared_params: frozenset[str],
|
|
user_kwargs: dict[str, object],
|
|
) -> None:
|
|
"""Route pipeline-declared request params into ``sampling_params.extra_args``.
|
|
|
|
Both online serving and offline examples call this so that model-specific
|
|
keys (e.g. ``cfg_text_scale`` for BAGEL) end up in ``extra_args`` instead
|
|
of being silently dropped.
|
|
|
|
This is a no-op when no declared params are present in ``user_kwargs``, so
|
|
it is safe to call on non-diffusion (e.g. AR) sampling params whose
|
|
``extra_args`` defaults to ``None``.
|
|
"""
|
|
declared = {key: user_kwargs[key] for key in declared_params if user_kwargs.get(key) is not None}
|
|
if not declared:
|
|
return
|
|
sampling_params.extra_args = {**(sampling_params.extra_args or {}), **declared}
|
|
|
|
|
|
def ar_grid_max_tokens(ar_width: int, ar_height: int) -> int | None:
|
|
"""Return the AR-stage ``max_tokens`` for MammothModa2-style visual-token grids.
|
|
|
|
The AR stage emits one visual token per grid cell, one EOL token per row,
|
|
and one final look-ahead token, so the generation budget is
|
|
``ar_height * (ar_width + 1) + 1``. Shared by the online serving path and
|
|
the offline example so a grid-contract change cannot desynchronize one
|
|
copy. Returns ``None`` when the grid is absent or non-positive, letting
|
|
callers keep their existing default.
|
|
"""
|
|
if ar_width <= 0 or ar_height <= 0:
|
|
return None
|
|
return ar_height * (ar_width + 1) + 1
|