anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

Merge origin/main into dev/vllm-align

Browse Source 
Resolve conflicts:
- docker/Dockerfile.ci: Keep wheel install section, update hash to 6042e66cd5
- vllm_omni/platforms/npu/worker/npu_model_runner.py: Accept main's refactoring

Made-with: Cursor

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
This commit is contained in:
tzhouam committed 2026-02-26 08:19:53 +00:00
2 parents a252918008 caed95fb45
commit 8d0c46487e
58 files changed
+4358 -1530
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

.buildkite/test-nightly.yml
+2 -2
View File
@@ -1,6 +1,6 @@
steps:
- label: "Omni Model Test with H100"
timeout_in_minutes: 60
timeout_in_minutes: 90
depends_on: upload-nightly-pipeline
if: build.env("NIGHTLY") == "1"
commands:
@@ -41,7 +41,6 @@ steps:
depends_on: upload-nightly-pipeline
if: build.env("NIGHTLY") == "1"
commands:
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/examples/online_serving/test_qwen2_5_omni.py -m "advanced_model" --run-level "advanced_model"
agents:
@@ -50,6 +49,7 @@ steps:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
shm-size: "8gb"
propagate-environment: true
environment:
- "HF_HOME=/fsx/hf_cache"
.buildkite/test-ready.yml
+1 -1
View File
@@ -157,7 +157,7 @@ steps:
# type: DirectoryOrCreate
- label: "Omni Model Test"
timeout_in_minutes: 15
timeout_in_minutes: 17
depends_on: upload-ready-pipeline
commands:
- export VLLM_LOGGING_LEVEL=DEBUG
docker/Dockerfile.ci
+4 -19
View File
@@ -1,5 +1,5 @@
ARG VLLM_BASE_IMAGE=vllm/vllm-openai
ARG VLLM_BASE_TAG=v0.15.0
ARG VLLM_BASE_TAG=v0.16.0
FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG}
ARG APP_DIR=/workspace/vllm-omni
WORKDIR ${APP_DIR}
@@ -8,38 +8,23 @@ COPY . .
# Install system dependencies
RUN apt-get update && \
apt-get install -y ffmpeg sox libsox-fmt-all && \
apt-get install -y ffmpeg sox libsox-fmt-all jq && \
apt-get clean && \
rm -rf /var/lib/apt/lists/*
# Install vLLM wheel for the specific commit
# URL format: https://wheels.vllm.ai/<commit_hash>/
# Current vLLM wheel commit: 24650715105a26b41c88f4164777b5e1a0f3200b
# Current vLLM wheel commit: 6042e66cd5304fc043d96aaa0c22d56f939af320
# Note: --prerelease=allow is required because nightly/commit wheels use
# pre-release version numbers (e.g. 2.10.0.dev...) which uv ignores by default,
# causing it to fall back to the older stable release from PyPI.
RUN uv pip install --system --upgrade --force-reinstall --prerelease=allow vllm --extra-index-url https://wheels.vllm.ai/24650715105a26b41c88f4164777b5e1a0f3200b
RUN uv pip install --system --upgrade --force-reinstall --prerelease=allow vllm --extra-index-url https://wheels.vllm.ai/6042e66cd5304fc043d96aaa0c22d56f939af320
RUN uv pip uninstall --system nvidia-cublas-cu12
# Install vllm-omni into the same uv-managed Python environment used by the base image.
# Use bash -c so that $(python3 -c ...) is expanded inside the container.
RUN uv pip install --system --no-cache-dir ".[dev]"
# pplx_kernels is optional; the base image may already ship a version.
# Reinstall attempt may fail if no compatible wheel exists for the current
# PyTorch/CUDA combination – allow the build to continue regardless.
RUN pip install --upgrade --force-reinstall "pplx_kernels" || true
RUN /bin/bash -c 'FLASHINFER_CUDA_TAG="$(python3 -c "import torch; print((torch.version.cuda or \"12.4\").replace(\".\", \"\"))")" && \
uv pip install --system --upgrade --force-reinstall \
"flashinfer-python==0.6.3" \
"flashinfer-cubin==0.6.3" \
"flashinfer-jit-cache==0.6.3" \
--extra-index-url "https://flashinfer.ai/whl/cu${FLASHINFER_CUDA_TAG}"'
# Pin nvidia-cublas-cu12 after vllm-omni install so dependency resolution does not overwrite it
RUN uv pip install --system --upgrade --force-reinstall "nvidia-cublas-cu12==12.9.1.4"
RUN uv pip install --system --upgrade --force-reinstall "numpy==2.2.6"
RUN ln -sf /usr/bin/python3 /usr/bin/python
ENTRYPOINT []
docker/Dockerfile.npu
+3 -2
View File
@@ -7,8 +7,9 @@ WORKDIR ${APP_DIR}
COPY . .
# Install vllm-omni with dev dependencies
RUN pip install --no-cache-dir -e . --no-build-isolation
RUN sed -i -E 's/^([[:space:]]*)"fa3-fwd==0\.0\.1",/\1# "fa3-fwd==0.0.1",/' pyproject.toml \
&& sed -i -E 's/\bonnxruntime\b/onnxruntime-cann/g' pyproject.toml
RUN pip install -v -e .
ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
docker/Dockerfile.npu.a3
+3 -2
View File
@@ -7,8 +7,9 @@ WORKDIR ${APP_DIR}
COPY . .
# Install vllm-omni with dev dependencies
RUN pip install --no-cache-dir -e . --no-build-isolation
RUN sed -i -E 's/^([[:space:]]*)"fa3-fwd==0\.0\.1",/\1# "fa3-fwd==0.0.1",/' pyproject.toml \
&& sed -i -E 's/\bonnxruntime\b/onnxruntime-cann/g' pyproject.toml
RUN pip install -v -e .
ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
docs/configuration/stage_configs.md
+12
View File
@@ -273,3 +273,15 @@ Default: `True`
Penalty applied to tokens that have already appeared in the generated sequence. Values greater than 1.0 discourage repetition, while values less than 1.0 encourage it. A value of 1.0 applies no penalty.
Default: `1.1`
### `tts_args` (TTS stages only)
Configuration for Text-to-Speech specific parameters. This section is only applicable to TTS model stages (e.g., `qwen3_tts`).
#### `tts_args.max_instructions_length`
Maximum character length for voice style/emotion instructions. Instructions exceeding this limit will be rejected with a validation error.
Default: `500`
This value can be overridden at runtime using the `--tts-max-instructions-length` CLI parameter when starting the server.
docs/getting_started/installation/npu.md
+6 -2
View File
@@ -24,6 +24,10 @@ vLLM-Omni supports NPU through the vLLM Ascend Plugin (vllm-ascend). This is a c
### Build wheel from source
=== "NPU"
=== "NPU release"
--8<-- "docs/getting_started/installation/npu/npu.inc.md:installation"
--8<-- "docs/getting_started/installation/npu/npu.inc.md:installation-release"
=== "NPU from main"
--8<-- "docs/getting_started/installation/npu/npu.inc.md:installation-main"
docs/getting_started/installation/npu/npu.inc.md
+38 -15
View File
@@ -3,14 +3,11 @@
For detailed hardware and software requirements, please refer to the [vllm-ascend installation documentation](https://docs.vllm.ai/projects/ascend/en/latest/installation.html).
# --8<-- [end:requirements]
# --8<-- [start:installation]
# --8<-- [start:installation-release]
The recommended way to use vLLM-Omni on NPU is through the vllm-ascend pre-built Docker images:
```bash
# Update DEVICE according to your NPUs (/dev/davinci[0-7])
export DEVICE0=/dev/davinci0
export DEVICE1=/dev/davinci1
# Update the vllm-ascend image
# Atlas A2:
# export IMAGE=quay.io/ascend/vllm-ascend:v0.14.0rc1
@@ -20,8 +17,10 @@ export IMAGE=quay.io/ascend/vllm-ascend:v0.14.0rc1
docker run --rm \
--name vllm-omni-npu \
--shm-size=1g \
--device $DEVICE0 \
--device $DEVICE1 \
--device /dev/davinci0 \
--device /dev/davinci1 \
--device /dev/davinci2 \
--device /dev/davinci3 \
--device /dev/davinci_manager \
--device /dev/devmm_svm \
--device /dev/hisi_hdc \
@@ -39,8 +38,9 @@ cd /vllm-workspace
git clone -b v0.14.0 https://github.com/vllm-project/vllm-omni.git
cd vllm-omni
VLLM_OMNI_TARGET_DEVICE=npu pip install -v -e .
# OR pip install -v -e . --no-build-isolation
sed -i -E 's/^([[:space:]]*)"fa3-fwd==0\.0\.1",/\1# "fa3-fwd==0.0.1",/' pyproject.toml \
&& sed -i -E 's/\bonnxruntime\b/onnxruntime-cann/g' pyproject.toml
pip install -v -e .
export VLLM_WORKER_MULTIPROC_METHOD=spawn
```
@@ -50,11 +50,38 @@ For other installation methods (pip installation, building from source, custom D
We are keeping [issue #997](https://github.com/vllm-project/vllm-omni/issues/997) up to date with the aligned versions of vLLM, vLLM-Ascend, and vLLM-Omni, and also outlining the Q1 roadmap there.
# --8<-- [end:installation]
# --8<-- [end:installation-release]
# --8<-- [start:installation-main]
You can also build vLLM-Omni from the latest main branch if you want to use the latest features or bug fixes. (But sometimes it will break for a while. You can check [issue #997](https://github.com/vllm-project/vllm-omni/issues/997) for the status of the latest commit of vLLM-Omni main branch on NPU.)
```bash
# Pin vLLM version to 0.16.0
cd /vllm-workspace/vllm
git pull origin main
git fetch origin --tags
git checkout v0.16.0
# Because vllm-ascend has not yet entered continuous development and has not been officially released, we need to pin it to a specific commit. Please note that this commit may change over time.
cd ../vllm-ascend
git pull origin main
git fetch origin --tags
git checkout e2175d9c7e62b437391dfee996b1375674ba7c18
# Install vLLM-Omni from the latest main branch
cd ../vllm-omni
git clone https://github.com/vllm-project/vllm-omni.git
pip install -v -e . --no-build-isolation
# or VLLM_OMNI_TARGET_DEVICE=npu pip install -v -e .
export VLLM_WORKER_MULTIPROC_METHOD=spawn
```
# --8<-- [end:installation-main]
# --8<-- [start:pre-built-images]
`vllm-ascend` offers Docker images for deployment. You can just pull the **prebuilt image** from the image repository [ascend/vllm-ascend](https://quay.io/repository/ascend/vllm-ascend?tab=tags) and run it with bash.
vLLM-Omni offers Docker images for Ascend NPU deployment. You can just pull the **prebuilt image** from the image repository [ascend/vllm-omni](https://quay.io/repository/ascend/vllm-omni?tab=tags) and run it with bash.
Supported images as following.
@@ -63,9 +90,7 @@ Supported images as following.
| image-tag | Atlas A2 | Ubuntu |
| image-tag-a3 | Atlas A3 | Ubuntu |
vLLM-Omni offers Docker images for Ascend NPU deployment. You can just pull the **prebuilt image** from the image repository [ascend/vllm-ascend](https://quay.io/repository/ascend/vllm-ascend?tab=tags) and run it with bash.
Here's an example deployment command that has been verified on 6 x NPUs:
Here's an example deployment command that has been verified on 4 x NPUs:
```bash
# Atlas A2:
@@ -80,8 +105,6 @@ docker run --rm \
--device /dev/davinci1 \
--device /dev/davinci2 \
--device /dev/davinci3 \
--device /dev/davinci4 \
--device /dev/davinci5 \
--device /dev/davinci_manager \
--device /dev/devmm_svm \
--device /dev/hisi_hdc \
docs/models/supported_models.md
+1
View File
@@ -40,6 +40,7 @@ th {
|`Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` |
|`Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` |
|`Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-0.6B-Base` |
|`NextStep11Pipeline` | NextStep-1.1 | `stepfun-ai/NextStep-1.1` |
## List of Supported Models for NPU
docs/user_guide/diffusion_acceleration.md
+1
View File
@@ -65,6 +65,7 @@ The following table shows which models are currently supported by each accelerat
| **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ |
| **Bagel** | `ByteDance-Seed/BAGEL-7B-MoT` | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
| **FLUX.1-dev** | `black-forest-labs/FLUX.1-dev` | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ |
| **NextStep-1.1** | `stepfun-ai/NextStep-1.1` | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ✅ |
### VideoGen
examples/offline_inference/text_to_image/README.md
+31 -3
View File
@@ -1,6 +1,6 @@
# Text-To-Image
This folder provides several entrypoints for experimenting with `Qwen/Qwen-Image` `Qwen/Qwen-Image-2512` `Tongyi-MAI/Z-Image-Turbo` using vLLM-Omni:
This folder provides several entrypoints for experimenting with `Qwen/Qwen-Image` `Qwen/Qwen-Image-2512` `Tongyi-MAI/Z-Image-Turbo` `stepfun-ai/NextStep-1.1` using vLLM-Omni, note that NextStep-1.1 has different architecture so we treat it differently regarding running arguments and pipeline.
- `text_to_image.py`: command-line script for single image generation with advanced options.
- `web_demo.py`: lightweight Gradio UI for interactive prompt/seed/CFG exploration.
@@ -74,6 +74,8 @@ if __name__ == "__main__":
## Local CLI Usage
### Qwen/Tongyi Models
```bash
python text_to_image.py \
--model Tongyi-MAI/Z-Image-Turbo \
@@ -87,7 +89,26 @@ python text_to_image.py \
--output outputs/coffee.png
```
Key arguments:
### NextStep Models
NextStep-1.1 can have extra arguments
```bash
python text_to_image.py \
--model stepfun-ai/NextStep-1.1 \
--prompt "A baby panda wearing an Iron Man mask, holding a board with 'NextStep-1' written on it" \
--height 512 \
--width 512 \
--num-inference-steps 28 \
--guidance-scale 7.5 \
--guidance-scale-2 1.0 \
--cfg-schedule constant \
--output nextstep_output.png \
--seed 42
```
### Key Arguments
**Common arguments:**
- `--prompt`: text description (string).
- `--seed`: integer seed for deterministic sampling.
@@ -98,8 +119,15 @@ Key arguments:
- `--output`: path to save the generated PNG.
- `--vae-use-slicing`: enable VAE slicing for memory optimization.
- `--vae-use-tiling`: enable VAE tiling for memory optimization.
- `--cfg-parallel-size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](../../../docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
- `--cfg-parallel-size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](../../../docs/user_guide/diffusion_acceleration.md#using-cfg-parallel).
- `--enable-cpu-offload`: enable CPU offloading for diffusion models.
- `--guidance-scale`: classifier-free guidance scale.
**NextStep-1.1 specific:**
- `--guidance-scale-2`: secondary guidance scale, e.g. image-level CFG (default: 1.0).
- `--timesteps-shift`: timesteps shift parameter for sampling (default: 1.0).
- `--cfg-schedule`: CFG schedule type, "constant" or "linear" (default: "constant").
- `--use-norm`: apply layer normalization to sampled tokens.
> ℹ️ If you encounter OOM errors, try using `--vae-use-slicing` and `--vae-use-tiling` to reduce memory usage.
examples/offline_inference/text_to_image/text_to_image.py
+73 -18
View File
@@ -16,13 +16,26 @@ from vllm_omni.outputs import OmniRequestOutput
from vllm_omni.platforms import current_omni_platform
def is_nextstep_model(model_name: str) -> bool:
"""Check if the model is a NextStep model by reading its config."""
from vllm.transformers_utils.config import get_hf_file_to_dict
try:
cfg = get_hf_file_to_dict("config.json", model_name)
if cfg and cfg.get("model_type") == "nextstep":
return True
except Exception:
pass
return False
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="Generate an image with Qwen-Image.")
parser = argparse.ArgumentParser(description="Generate an image with supported diffusion models.")
parser.add_argument(
"--model",
default="Qwen/Qwen-Image",
help="Diffusion model name or local path. Supported models: "
"Qwen/Qwen-Image, Tongyi-MAI/Z-Image-Turbo, Qwen/Qwen-Image-2512",
"Qwen/Qwen-Image, Tongyi-MAI/Z-Image-Turbo, Qwen/Qwen-Image-2512, stepfun-ai/NextStep-1.1",
)
parser.add_argument("--prompt", default="a cup of coffee on the table", help="Text prompt for image generation.")
parser.add_argument(
@@ -153,16 +166,43 @@ def parse_args() -> argparse.Namespace:
default=1,
help="Number of ranks used for VAE patch/tile parallelism (decode/encode).",
)
# NextStep-1.1 specific arguments
parser.add_argument(
"--guidance-scale-2",
type=float,
default=1.0,
help="Secondary guidance scale (e.g. image-level CFG for NextStep-1.1).",
)
parser.add_argument(
"--timesteps-shift",
type=float,
default=1.0,
help="[NextStep-1.1 only] Timesteps shift parameter for sampling.",
)
parser.add_argument(
"--cfg-schedule",
type=str,
default="constant",
choices=["constant", "linear"],
help="[NextStep-1.1 only] CFG schedule type.",
)
parser.add_argument(
"--use-norm",
action="store_true",
help="[NextStep-1.1 only] Apply layer normalization to sampled tokens.",
)
return parser.parse_args()
def main():
args = parse_args()
generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
use_nextstep = is_nextstep_model(args.model)
# Configure cache based on backend type
cache_config = None
if args.cache_backend == "cache_dit":
cache_backend = args.cache_backend
if cache_backend == "cache_dit":
# cache-dit configuration: Hybrid DBCache + SCM + TaylorSeer
# All parameters marked with [cache-dit only] in DiffusionCacheConfig
cache_config = {
@@ -179,7 +219,7 @@ def main():
"scm_steps_mask_policy": None, # SCM mask policy: None (disabled), "slow", "medium", "fast", "ultra"
"scm_steps_policy": "dynamic", # SCM steps policy: "dynamic" or "static"
}
elif args.cache_backend == "tea_cache":
elif cache_backend == "tea_cache":
# TeaCache configuration
# All parameters marked with [tea_cache only] in DiffusionCacheConfig
cache_config = {
@@ -213,19 +253,24 @@ def main():
elif args.quantization:
quant_kwargs["quantization"] = args.quantization
omni = Omni(
model=args.model,
enable_layerwise_offload=args.enable_layerwise_offload,
vae_use_slicing=args.vae_use_slicing,
vae_use_tiling=args.vae_use_tiling,
cache_backend=args.cache_backend,
cache_config=cache_config,
enable_cache_dit_summary=args.enable_cache_dit_summary,
parallel_config=parallel_config,
enforce_eager=args.enforce_eager,
enable_cpu_offload=args.enable_cpu_offload,
# Initialize Omni with model-specific settings
omni_kwargs = {
"model": args.model,
"enable_layerwise_offload": args.enable_layerwise_offload,
"vae_use_slicing": args.vae_use_slicing,
"vae_use_tiling": args.vae_use_tiling,
"cache_backend": cache_backend,
"cache_config": cache_config,
"enable_cache_dit_summary": args.enable_cache_dit_summary,
"parallel_config": parallel_config,
"enforce_eager": args.enforce_eager,
"enable_cpu_offload": args.enable_cpu_offload,
**quant_kwargs,
)
}
if use_nextstep:
# NextStep-1.1 requires explicit pipeline class
omni_kwargs["model_class_name"] = "NextStep11Pipeline"
omni = Omni(**omni_kwargs)
if profiler_enabled:
print("[Profiler] Starting profiling...")
@@ -236,7 +281,7 @@ def main():
print("Generation Configuration:")
print(f" Model: {args.model}")
print(f" Inference steps: {args.num_inference_steps}")
print(f" Cache backend: {args.cache_backend if args.cache_backend else 'None (no acceleration)'}")
print(f" Cache backend: {cache_backend if cache_backend else 'None (no acceleration)'}")
print(f" Quantization: {args.quantization if args.quantization else 'None (BF16)'}")
if ignored_layers:
print(f" Ignored layers: {ignored_layers}")
@@ -250,6 +295,13 @@ def main():
print(f"{'=' * 60}\n")
generation_start = time.perf_counter()
extra_args = {
"timesteps_shift": args.timesteps_shift,
"cfg_schedule": args.cfg_schedule,
"use_norm": args.use_norm,
}
outputs = omni.generate(
{
"prompt": args.prompt,
@@ -261,10 +313,13 @@ def main():
generator=generator,
true_cfg_scale=args.cfg_scale,
guidance_scale=args.guidance_scale,
guidance_scale_2=args.guidance_scale_2,
num_inference_steps=args.num_inference_steps,
num_outputs_per_prompt=args.num_images_per_prompt,
extra_args=extra_args,
),
)
generation_end = time.perf_counter()
generation_time = generation_end - generation_start
examples/online_serving/qwen3_tts/README.md
+19
View File
@@ -90,6 +90,25 @@ curl -X POST http://localhost:8091/v1/audio/speech \
curl http://localhost:8091/v1/audio/voices
```
## Server Configuration
### TTS-Specific CLI Parameters
| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `--tts-max-instructions-length` | int | 500 | Maximum character length for voice style instructions |
**Example with custom instructions length:**
```bash
vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
--omni \
--port 8000 \
--tts-max-instructions-length 1000
```
This is useful when you need longer voice style descriptions, especially for the VoiceDesign task type.
## API Reference
### Endpoint
tests/conftest.py
+33 -32
View File
@@ -848,41 +848,42 @@ def modify_stage_config(
del config[key]
# Apply updates
for key, value in updates.items():
if key == "stage_args":
if value and isinstance(value, dict):
stage_args = config.get("stage_args", [])
if not stage_args:
raise ValueError("stage_args does not exist in config")
if updates:
for key, value in updates.items():
if key == "stage_args":
if value and isinstance(value, dict):
stage_args = config.get("stage_args", [])
if not stage_args:
raise ValueError("stage_args does not exist in config")
for stage_id, stage_updates in value.items():
# Find stage by ID
target_stage = None
for stage in stage_args:
if stage.get("stage_id") == int(stage_id):
target_stage = stage
break
for stage_id, stage_updates in value.items():
# Find stage by ID
target_stage = None
for stage in stage_args:
if stage.get("stage_id") == int(stage_id):
target_stage = stage
break
if target_stage is None:
available_ids = [s.get("stage_id") for s in stage_args if "stage_id" in s]
raise KeyError(f"Stage ID {stage_id} not found, available: {available_ids}")
if target_stage is None:
available_ids = [s.get("stage_id") for s in stage_args if "stage_id" in s]
raise KeyError(f"Stage ID {stage_id} not found, available: {available_ids}")
# Apply updates to this stage
for path, val in stage_updates.items():
# Check if this is a simple key (not dot-separated)
# Example: 'engine_input_source' vs 'engine_args.max_model_len'
if "." not in path:
# Direct key assignment (e.g., updating a list value)
target_stage[path] = val
else:
# Dot-separated path (e.g., nested dict access)
apply_update(target_stage, path, val)
elif "." in key:
# Apply using dot-separated path
apply_update(config, key, value)
else:
# Direct top-level key
config[key] = value
# Apply updates to this stage
for path, val in stage_updates.items():
# Check if this is a simple key (not dot-separated)
# Example: 'engine_input_source' vs 'engine_args.max_model_len'
if "." not in path:
# Direct key assignment (e.g., updating a list value)
target_stage[path] = val
else:
# Dot-separated path (e.g., nested dict access)
apply_update(target_stage, path, val)
elif "." in key:
# Apply using dot-separated path
apply_update(config, key, value)
else:
# Direct top-level key
config[key] = value
# Save to new file with timestamp
timestamp = int(time.time())
tests/diffusion/models/nextstep_1_1/test_nextstep_cfg_parallel_layout.py
+217
View File
@@ -0,0 +1,217 @@
from types import SimpleNamespace
import pytest
import torch
from PIL import Image
import vllm_omni.diffusion.models.nextstep_1_1.pipeline_nextstep_1_1 as nextstep_pipeline_module
from vllm_omni.diffusion.models.nextstep_1_1.modeling_nextstep_heads import FlowMatchingHead
from vllm_omni.diffusion.models.nextstep_1_1.pipeline_nextstep_1_1 import NextStep11Pipeline
class _DummyImageHead:
def __init__(self, token_dim: int):
self.token_dim = token_dim
self.calls = []
def sample(
self,
c: torch.Tensor,
cfg: float,
cfg_img: float,
cfg_mult: int,
timesteps_shift: float,
num_sampling_steps: int,
noise_repeat: int,
) -> torch.Tensor:
self.calls.append(
{
"batch": c.shape[0],
"cfg": cfg,
"cfg_img": cfg_img,
"cfg_mult": cfg_mult,
"noise_repeat": noise_repeat,
}
)
batch_per_prompt = c.shape[0] // cfg_mult
return torch.ones(batch_per_prompt, self.token_dim, dtype=c.dtype, device=c.device)
class _DummyModel:
def __init__(self, hidden_dim: int, token_dim: int):
self.hidden_dim = hidden_dim
self.image_head = _DummyImageHead(token_dim)
self.forward_batches = []
def image_out_projector(self, c: torch.Tensor) -> torch.Tensor:
return c
def image_in_projector(self, sampled_tokens: torch.Tensor) -> torch.Tensor:
bsz = sampled_tokens.shape[0]
return torch.zeros(bsz, 1, self.hidden_dim, dtype=sampled_tokens.dtype, device=sampled_tokens.device)
def forward_model(self, inputs_embeds: torch.Tensor, attention_mask, past_key_values, use_cache: bool):
del attention_mask, use_cache
self.forward_batches.append(inputs_embeds.shape[0])
return SimpleNamespace(
last_hidden_state=torch.zeros(
inputs_embeds.shape[0],
1,
self.hidden_dim,
dtype=inputs_embeds.dtype,
device=inputs_embeds.device,
),
past_key_values=past_key_values,
)
def _make_minimal_pipeline_for_decoding(hidden_dim: int = 8, token_dim: int = 4):
pipeline = object.__new__(NextStep11Pipeline)
pipeline.config = SimpleNamespace(latent_channels=token_dim, latent_patch_size=1, use_gen_pos_embed=False)
pipeline.model = _DummyModel(hidden_dim=hidden_dim, token_dim=token_dim)
return pipeline
@pytest.mark.parametrize(
("cfg", "cfg_img", "has_image_conditions", "expected_cfg_mult", "expected_cfg_img"),
[
(1.0, 1.0, False, 1, 1.0),
(7.5, 1.0, False, 2, 1.0),
(7.5, 8.0, False, 2, 1.0),
(7.5, 1.5, True, 3, 1.5),
],
)
def test_resolve_cfg_layout(cfg, cfg_img, has_image_conditions, expected_cfg_mult, expected_cfg_img):
cfg_mult, effective_cfg_img = NextStep11Pipeline._resolve_cfg_layout(cfg, cfg_img, has_image_conditions)
assert cfg_mult == expected_cfg_mult
assert effective_cfg_img == expected_cfg_img
def test_build_captions_ignores_image_cfg_without_image_conditions():
pipeline = object.__new__(NextStep11Pipeline)
pipeline._image_str = lambda hw: f"<image:{hw}>"
captions, images, cfg_mult, effective_cfg_img = pipeline._build_captions(
captions=["a prompt"],
images=None,
num_images_per_caption=1,
positive_prompt=None,
negative_prompt="bad quality",
cfg=7.5,
cfg_img=8.0,
)
assert cfg_mult == 2
assert effective_cfg_img == 1.0
assert images is None
assert captions == ["a prompt", "bad quality"]
def test_build_captions_enables_three_way_cfg_when_image_conditions_exist():
pipeline = object.__new__(NextStep11Pipeline)
pipeline._image_str = lambda hw: f"<image:{hw}>"
image = Image.new("RGB", (64, 32))
captions, images, cfg_mult, effective_cfg_img = pipeline._build_captions(
captions=["a prompt"],
images=[image],
num_images_per_caption=1,
positive_prompt=None,
negative_prompt="bad quality",
cfg=7.5,
cfg_img=1.5,
)
assert cfg_mult == 3
assert effective_cfg_img == 1.5
assert len(captions) == 3
assert captions[1].startswith("<image:")
assert captions[2] == "bad quality"
assert len(images) == 2
def test_decoding_non_parallel_uses_cfg_mult_for_sampling_and_duplication(monkeypatch):
pipeline = _make_minimal_pipeline_for_decoding()
monkeypatch.setattr(nextstep_pipeline_module, "get_classifier_free_guidance_world_size", lambda: 1)
c = torch.zeros(2, 1, 8)
attention_mask = torch.ones(2, 3, dtype=torch.long)
tokens = pipeline.decoding(
c=c,
attention_mask=attention_mask,
past_key_values=None,
max_new_len=1,
num_images_per_caption=1,
cfg=7.5,
cfg_img=1.0,
cfg_mult=2,
progress=False,
)
assert tokens.shape == (1, 1, 4)
assert pipeline.model.image_head.calls[0]["cfg_mult"] == 2
assert pipeline.model.forward_batches == [2]
def test_decoding_cfg_parallel_mismatch_falls_back_to_non_parallel(monkeypatch):
pipeline = _make_minimal_pipeline_for_decoding()
monkeypatch.setattr(nextstep_pipeline_module, "get_classifier_free_guidance_world_size", lambda: 2)
def _unexpected(*args, **kwargs):
del args, kwargs
raise AssertionError("CFG rank/group should not be queried on mismatch fallback.")
monkeypatch.setattr(nextstep_pipeline_module, "get_classifier_free_guidance_rank", _unexpected)
monkeypatch.setattr(nextstep_pipeline_module, "get_cfg_group", _unexpected)
c = torch.zeros(3, 1, 8)
attention_mask = torch.ones(3, 3, dtype=torch.long)
tokens = pipeline.decoding(
c=c,
attention_mask=attention_mask,
past_key_values=None,
max_new_len=1,
num_images_per_caption=1,
cfg=7.5,
cfg_img=1.5,
cfg_mult=3,
progress=False,
)
assert tokens.shape == (1, 1, 4)
assert pipeline.model.image_head.calls[0]["cfg_mult"] == 3
assert pipeline.model.forward_batches == [3]
def test_decoding_rejects_incompatible_batch_and_cfg_mult(monkeypatch):
pipeline = _make_minimal_pipeline_for_decoding()
monkeypatch.setattr(nextstep_pipeline_module, "get_classifier_free_guidance_world_size", lambda: 1)
with pytest.raises(ValueError, match="not divisible"):
pipeline.decoding(
c=torch.zeros(5, 1, 8),
attention_mask=torch.ones(5, 3, dtype=torch.long),
past_key_values=None,
max_new_len=1,
num_images_per_caption=1,
cfg=7.5,
cfg_img=1.0,
cfg_mult=2,
progress=False,
)
def test_flow_matching_head_sample_validates_cfg_mult_divisibility():
head = FlowMatchingHead(input_dim=4, cond_dim=8, dim=8, layers=1)
with pytest.raises(ValueError, match="not divisible"):
head.sample(
c=torch.zeros(5, 8),
cfg=2.0,
cfg_img=1.0,
cfg_mult=2,
num_sampling_steps=1,
)
tests/diffusion/test_diffusion_model_runner.py
+161
View File
@@ -0,0 +1,161 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from contextlib import contextmanager
from types import SimpleNamespace
import pytest
import torch
import vllm_omni.diffusion.worker.diffusion_model_runner as model_runner_module
from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
@contextmanager
def _noop_forward_context(*args, **kwargs):
del args, kwargs
yield
class _DummyPipeline:
def __init__(self, output):
self._output = output
self.forward_calls = 0
def forward(self, req):
del req
self.forward_calls += 1
return self._output
def _make_request(skip_cache_refresh: bool = True):
sampling_params = SimpleNamespace(
generator=None,
seed=None,
generator_device=None,
num_inference_steps=4,
)
return SimpleNamespace(
prompts=["a prompt"],
sampling_params=sampling_params,
skip_cache_refresh=skip_cache_refresh,
)
def _make_runner(cache_backend, cache_backend_name: str, enable_cache_dit_summary: bool = True):
runner = object.__new__(DiffusionModelRunner)
runner.vllm_config = object()
runner.device = torch.device("cpu")
runner.pipeline = _DummyPipeline(output=SimpleNamespace(output="ok"))
runner.cache_backend = cache_backend
runner.offload_backend = None
runner.od_config = SimpleNamespace(
cache_backend=cache_backend_name,
enable_cache_dit_summary=enable_cache_dit_summary,
)
runner.kv_transfer_manager = SimpleNamespace(receive_kv_cache=lambda req, target_device: None)
return runner
def test_execute_model_skips_cache_summary_without_active_cache_backend(monkeypatch):
"""Guard cache diagnostics with runtime backend state to avoid stale-config crashes."""
runner = _make_runner(cache_backend=None, cache_backend_name="cache_dit")
req = _make_request(skip_cache_refresh=True)
cache_summary_calls = []
monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
monkeypatch.setattr(
model_runner_module,
"cache_summary",
lambda pipeline, details: cache_summary_calls.append((pipeline, details)),
)
output = DiffusionModelRunner.execute_model(runner, req)
assert output.output == "ok"
assert cache_summary_calls == []
def test_execute_model_emits_cache_summary_with_active_cache_dit_backend(monkeypatch):
class _EnabledCacheBackend:
def is_enabled(self):
return True
runner = _make_runner(cache_backend=_EnabledCacheBackend(), cache_backend_name="cache_dit")
req = _make_request(skip_cache_refresh=True)
cache_summary_calls = []
monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
monkeypatch.setattr(
model_runner_module,
"cache_summary",
lambda pipeline, details: cache_summary_calls.append((pipeline, details)),
)
output = DiffusionModelRunner.execute_model(runner, req)
assert output.output == "ok"
assert cache_summary_calls == [(runner.pipeline, True)]
def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch):
class _DummyLoader:
def __init__(self, load_config):
del load_config
def load_model(self, **kwargs):
del kwargs
return SimpleNamespace(transformer=torch.nn.Identity())
class _DummyMemoryProfiler:
consumed_memory = 0
def __enter__(self):
return self
def __exit__(self, exc_type, exc, tb):
del exc_type, exc, tb
return False
class _DummyCacheBackend:
def __init__(self):
self.enabled = False
def enable(self, pipeline):
del pipeline
self.enabled = True
dummy_cache_backend = _DummyCacheBackend()
runner = object.__new__(DiffusionModelRunner)
runner.vllm_config = object()
runner.device = torch.device("cpu")
runner.pipeline = None
runner.cache_backend = None
runner.offload_backend = None
runner.od_config = SimpleNamespace(
enable_cpu_offload=False,
enable_layerwise_offload=False,
cache_backend="cache_dit",
cache_config={},
model_class_name="NextStep11Pipeline",
enforce_eager=True,
)
monkeypatch.setattr(model_runner_module, "LoadConfig", lambda: object())
monkeypatch.setattr(model_runner_module, "DiffusersPipelineLoader", _DummyLoader)
monkeypatch.setattr(model_runner_module, "DeviceMemoryProfiler", _DummyMemoryProfiler)
monkeypatch.setattr(model_runner_module, "get_offload_backend", lambda od_config, device: None)
monkeypatch.setattr(
model_runner_module, "get_cache_backend", lambda cache_backend, cache_config: dummy_cache_backend
)
DiffusionModelRunner.load_model(runner)
assert runner.cache_backend is None
assert runner.od_config.cache_backend is None
assert dummy_cache_backend.enabled is False
tests/distributed/omni_connectors/test_kv_flow.py
+92
View File
@@ -105,6 +105,98 @@ def test_manager_extraction(kv_config, mock_connector, common_constants):
assert data["layer_blocks"]["key_cache"][0].shape == expected_shape
def test_manager_extraction_tuple_layout(kv_config, mock_connector, common_constants):
"""Test extraction with tuple layout."""
num_layers = common_constants["num_layers"]
block_size = common_constants["block_size"]
num_heads = common_constants["num_heads"]
head_dim = common_constants["head_dim"]
seq_len = common_constants["seq_len"]
req_id = common_constants["req_id"]
num_blocks = 10
kv_caches = []
for _ in range(num_layers):
k_cache = torch.randn(num_blocks, block_size, num_heads, head_dim)
v_cache = torch.randn(num_blocks, block_size, num_heads, head_dim)
kv_caches.append((k_cache, v_cache))
block_ids = [1, 3, 5]
finished_reqs = {req_id: {"block_ids": block_ids, "seq_len": seq_len}}
manager = OmniKVTransferManager(kv_config)
manager._connector = mock_connector
processed = manager.handle_finished_requests_kv_transfer(finished_reqs, kv_caches, block_size, "float32")
assert req_id in processed
full_request_id = f"omni_stage1_to_stage2_kv_cache_{req_id}"
expected_key = f"stage1->stage2:{full_request_id}"
assert expected_key in mock_connector.store
data = mock_connector.store[expected_key]
expected_shape = (seq_len, num_heads, head_dim)
for idx in range(len(kv_caches)):
assert data["layer_blocks"]["key_cache"][idx].shape == expected_shape
assert data["layer_blocks"]["value_cache"][idx].shape == expected_shape
def test_manager_extraction_mismatched_kv_block_counts(kv_config, mock_connector, common_constants):
"""Mismatched key/value block counts should not crash extraction."""
block_size = common_constants["block_size"]
num_heads = common_constants["num_heads"]
head_dim = common_constants["head_dim"]
req_id = common_constants["req_id"]
key_blocks = torch.randn(3, block_size, num_heads, head_dim)
value_blocks = torch.randn(2, block_size, num_heads, head_dim)
kv_caches = [(key_blocks, value_blocks)]
finished_reqs = {req_id: {"block_ids": [0, 1, 2], "seq_len": 32}}
manager = OmniKVTransferManager(kv_config)
manager._connector = mock_connector
processed = manager.handle_finished_requests_kv_transfer(finished_reqs, kv_caches, block_size, "float32")
assert req_id in processed
full_request_id = f"omni_stage1_to_stage2_kv_cache_{req_id}"
expected_key = f"stage1->stage2:{full_request_id}"
assert expected_key in mock_connector.store
data = mock_connector.store[expected_key]
expected_shape = (2 * block_size, num_heads, head_dim)
assert data["layer_blocks"]["key_cache"][0].shape == expected_shape
assert data["layer_blocks"]["value_cache"][0].shape == expected_shape
@pytest.mark.parametrize(
"invalid_case",
["invalid_stacked_shape", "invalid_tuple_length", "non_tensor_entries"],
)
def test_normalize_layer_kv_rejects_invalid_inputs(kv_config, common_constants, invalid_case):
"""_normalize_layer_kv should reject malformed KV representations."""
block_size = common_constants["block_size"]
num_heads = common_constants["num_heads"]
head_dim = common_constants["head_dim"]
req_id = common_constants["req_id"]
if invalid_case == "invalid_stacked_shape":
layer_kv = torch.randn(3, block_size, num_heads, head_dim)
elif invalid_case == "invalid_tuple_length":
layer_kv = (
torch.randn(2, block_size, num_heads, head_dim),
torch.randn(2, block_size, num_heads, head_dim),
torch.randn(2, block_size, num_heads, head_dim),
)
else:
layer_kv = (torch.randn(2, block_size, num_heads, head_dim), "not-a-tensor")
manager = OmniKVTransferManager(kv_config)
normalized = manager._normalize_layer_kv(layer_kv, req_id=req_id, layer_idx=0)
assert normalized is None
def test_manager_reception(kv_config, mock_connector, common_constants):
"""Test reception and injection logic in OmniKVTransferManager."""
num_layers = common_constants["num_layers"]
tests/e2e/online_serving/test_qwen3_omni.py
+1 -1
View File
@@ -114,7 +114,7 @@ def test_mix_to_text_audio_001(omni_server, openai_client) -> None:
"messages": messages,
"stream": True,
"key_words": {
"audio": ["water", "chirping", "crackling"],
"audio": ["water", "chirping", "crackling", "rain"],
"image": ["square", "quadrate"],
},
}
tests/e2e/online_serving/test_qwen3_omni_expansion.py
+3 -1
View File
@@ -22,7 +22,7 @@ from tests.utils import hardware_test
models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"]
AUDIO_KEY = ["water", "chirping", "crackling"]
AUDIO_KEY = ["water", "chirping", "crackling", "rain"]
IMAGE_KEY = ["square", "quadrate"]
VIDEO_KEY = ["sphere", "globe", "circle", "round", "ball"]
@@ -267,6 +267,7 @@ def test_text_image_to_text_audio_001(omni_server, openai_client) -> None:
openai_client.send_request(request_config)
@pytest.mark.skip(reason="There is a known issue with oom error.")
@pytest.mark.advanced_model
@pytest.mark.omni
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
@@ -287,6 +288,7 @@ def test_text_video_to_text_audio_001(omni_server, openai_client) -> None:
openai_client.send_request(request_config, request_num=get_max_batch_size())
@pytest.mark.skip(reason="There is a known issue with shape mismatch error.")
@pytest.mark.advanced_model
@pytest.mark.omni
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
tests/entrypoints/openai_api/test_serving_speech.py
+132 -7
View File
@@ -291,6 +291,7 @@ class TestTTSMethods:
mock_engine_client = MagicMock()
mock_engine_client.errored = False
mock_engine_client.stage_list = None
mock_engine_client.tts_max_instructions_length = None
mock_models = MagicMock()
mock_models.is_base_model.return_value = True
return OmniOpenAIServingSpeech(
@@ -299,16 +300,35 @@ class TestTTSMethods:
request_logger=MagicMock(),
)
def test_is_tts_model(self, speech_server):
"""Test TTS model detection."""
# No stage_list -> False
assert speech_server._is_tts_model() is False
def test_is_tts_detection_no_stage(self, speech_server):
"""Test TTS model detection when no TTS stage exists."""
# Fixture creates server with stage_list = None -> _is_tts should be False
assert speech_server._is_tts is False
assert speech_server._tts_stage is None
# With qwen3_tts stage -> True
def test_is_tts_detection_with_tts_stage(self):
"""Test TTS model detection when TTS stage exists."""
mock_engine_client = MagicMock()
mock_engine_client.errored = False
mock_engine_client.tts_max_instructions_length = None
# Create a TTS stage
mock_stage = MagicMock()
mock_stage.model_stage = "qwen3_tts"
speech_server.engine_client.stage_list = [mock_stage]
assert speech_server._is_tts_model() is True
mock_stage.tts_args = {}
mock_engine_client.stage_list = [mock_stage]
mock_models = MagicMock()
mock_models.is_base_model.return_value = True
server = OmniOpenAIServingSpeech(
engine_client=mock_engine_client,
models=mock_models,
request_logger=MagicMock(),
)
assert server._is_tts is True
assert server._tts_stage is mock_stage
def test_estimate_prompt_len_fallback(self, speech_server):
"""Test prompt length estimation falls back to 2048 when model is unavailable."""
@@ -381,3 +401,108 @@ class TestTTSMethods:
# Verify speakers are normalized to lowercase
assert server.supported_speakers == {"ryan", "vivian", "aiden"}
def test_max_instructions_length_default(self, speech_server):
"""Test default max instructions length (500) when no config provided."""
# Fixture creates server with no CLI override and no TTS stage
assert speech_server._max_instructions_length == 500
def test_max_instructions_length_cli_override(self):
"""Test CLI override (stored in engine_client) takes highest priority."""
mock_engine_client = MagicMock()
mock_engine_client.errored = False
mock_engine_client.stage_list = None
# CLI override is stored in engine_client
mock_engine_client.tts_max_instructions_length = 1000
mock_models = MagicMock()
mock_models.is_base_model.return_value = True
server = OmniOpenAIServingSpeech(
engine_client=mock_engine_client,
models=mock_models,
request_logger=MagicMock(),
)
# Value is cached during __init__
assert server._max_instructions_length == 1000
def test_max_instructions_length_stage_config(self):
"""Test stage config value is used when no CLI override."""
mock_engine_client = MagicMock()
mock_engine_client.errored = False
mock_engine_client.tts_max_instructions_length = None # No CLI override
mock_models = MagicMock()
mock_models.is_base_model.return_value = True
# Mock stage with tts_args
mock_stage = MagicMock()
mock_stage.model_stage = "qwen3_tts"
mock_stage.tts_args = {"max_instructions_length": 750}
mock_engine_client.stage_list = [mock_stage]
server = OmniOpenAIServingSpeech(
engine_client=mock_engine_client,
models=mock_models,
request_logger=MagicMock(),
)
# Value is cached during __init__
assert server._max_instructions_length == 750
def test_max_instructions_length_cli_overrides_stage_config(self):
"""Test CLI override (in engine_client) takes precedence over stage config."""
mock_engine_client = MagicMock()
mock_engine_client.errored = False
# CLI override stored in engine_client
mock_engine_client.tts_max_instructions_length = 2000
mock_models = MagicMock()
mock_models.is_base_model.return_value = True
# Mock stage with tts_args
mock_stage = MagicMock()
mock_stage.model_stage = "qwen3_tts"
mock_stage.tts_args = {"max_instructions_length": 750}
mock_engine_client.stage_list = [mock_stage]
server = OmniOpenAIServingSpeech(
engine_client=mock_engine_client,
models=mock_models,
request_logger=MagicMock(),
)
# CLI value (2000) should override stage config (750)
assert server._max_instructions_length == 2000
def test_validate_instructions_length_uses_cached_value(self):
"""Test instructions length validation uses cached _max_instructions_length."""
mock_engine_client = MagicMock()
mock_engine_client.errored = False
mock_engine_client.stage_list = None
# CLI override with max length of 10 characters
mock_engine_client.tts_max_instructions_length = 10
mock_models = MagicMock()
mock_models.is_base_model.return_value = True
server = OmniOpenAIServingSpeech(
engine_client=mock_engine_client,
models=mock_models,
request_logger=MagicMock(),
)
# Verify cached value
assert server._max_instructions_length == 10
# Instructions within limit should pass
req = OpenAICreateSpeechRequest(
input="Hello",
task_type="VoiceDesign",
instructions="short",
)
assert server._validate_tts_request(req) is None
# Instructions exceeding limit should fail
req = OpenAICreateSpeechRequest(
input="Hello",
task_type="VoiceDesign",
instructions="this is too long",
)
error = server._validate_tts_request(req)
assert error is not None
assert "max 10 characters" in error
tests/perf/tests/test.json
-114
View File
@@ -6,28 +6,6 @@
"stage_config_name": "qwen3_omni.yaml"
},
"benchmark_params": [
{
"dataset_name": "random",
"num_prompts": [
10,
40,
100
],
"request_rate": [
0.1,
0.5,
1
],
"random_input_len": 2500,
"random_output_len": 900,
"ignore_eos": true,
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
"baseline": {
"mean_ttft_ms": 100000,
"mean_audio_ttfp_ms": 100000,
"mean_audio_rtf": 100000
}
},
{
"dataset_name": "random",
"num_prompts": [
@@ -84,41 +62,6 @@
"mean_audio_ttfp_ms": 100000,
"mean_audio_rtf": 100000
}
},
{
"dataset_name": "random-mm",
"num_prompts": [
10,
40,
100
],
"request_rate": [
0.1,
0.5,
1
],
"random_input_len": 100,
"random_output_len": 100,
"random_range_ratio": 0.0,
"ignore_eos": true,
"random_mm_base_items_per_request": 3,
"random_mm_num_mm_items_range_ratio": 0,
"random_mm_limit_mm_per_prompt": {
"image": 1,
"video": 1,
"audio": 1
},
"random_mm_bucket_config": {
"(32, 32, 1)": 0.5,
"(0, 1, 1)": 0.1,
"(32, 32, 2)": 0.4
},
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
"baseline": {
"mean_ttft_ms": 100000,
"mean_audio_ttfp_ms": 100000,
"mean_audio_rtf": 100000
}
}
]
},
@@ -147,28 +90,6 @@
}
},
"benchmark_params": [
{
"dataset_name": "random",
"num_prompts": [
10,
40,
100
],
"request_rate": [
0.1,
0.5,
1
],
"random_input_len": 2500,
"random_output_len": 900,
"ignore_eos": true,
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
"baseline": {
"mean_ttft_ms": 100000,
"mean_audio_ttfp_ms": 100000,
"mean_audio_rtf": 100000
}
},
{
"dataset_name": "random",
"num_prompts": [
@@ -225,41 +146,6 @@
"mean_audio_ttfp_ms": 100000,
"mean_audio_rtf": 100000
}
},
{
"dataset_name": "random-mm",
"num_prompts": [
10,
40,
100
],
"request_rate": [
0.1,
0.5,
1
],
"random_input_len": 100,
"random_output_len": 100,
"random_range_ratio": 0.0,
"ignore_eos": true,
"random_mm_base_items_per_request": 3,
"random_mm_num_mm_items_range_ratio": 0,
"random_mm_limit_mm_per_prompt": {
"image": 1,
"video": 1,
"audio": 1
},
"random_mm_bucket_config": {
"(32, 32, 1)": 0.5,
"(0, 1, 1)": 0.1,
"(32, 32, 2)": 0.4
},
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
"baseline": {
"mean_ttft_ms": 100000,
"mean_audio_ttfp_ms": 100000,
"mean_audio_rtf": 100000
}
}
]
}
vllm_omni/core/sched/omni_generation_scheduler.py
+3
View File
@@ -137,6 +137,9 @@ class OmniGenerationScheduler(VLLMScheduler):
# Pop the finished request from waiting queue and don't schedule it
self.waiting.pop_request()
continue
# Count the number of prefix cached tokens.
if request.num_cached_tokens < 0:
request.num_cached_tokens = request.num_computed_tokens
# async_chunk: wait for the first upstream chunk (don't start with placeholders).
if self.chunk_transfer_adapter is not None and len(request.prompt_token_ids) == 0:
vllm_omni/diffusion/attention/backends/sdpa.py
+43 -12
View File
@@ -1,6 +1,8 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from typing import Literal
import torch
from vllm.logger import init_logger
@@ -13,17 +15,27 @@ from vllm_omni.diffusion.attention.backends.abstract import (
logger = init_logger(__name__)
def _maybe_reshape_attn_mask(query: torch.Tensor, key: torch.Tensor, attn_mask: torch.Tensor | None = None):
SDPAMaskMode = Literal["broadcast_k", "full_qk"]
def _maybe_reshape_attn_mask(
query: torch.Tensor,
key: torch.Tensor,
attn_mask: torch.Tensor | None = None,
mask_mode: SDPAMaskMode = "broadcast_k",
):
"""
Reshape Attention Mask
[batch_size, seq_len_k] -> [batch_size, 1, seq_len_q, seq_len_k]
2D [batch_size, seq_len_k] ->
- broadcast_k: [batch_size, 1, 1, seq_len_k]
- full_qk: [batch_size, 1, seq_len_q, seq_len_k]
"""
# Skip Attention Mask if all values are 1, `None` mask can speedup the computation
if attn_mask is not None and torch.all(attn_mask != 0):
attn_mask = None
# Reshape Attention Mask
# [batch_size, seq_len_k] -> [batch_size, 1, seq_len_q, seq_len_k]
# 2D [batch_size, seq_len_k] mask only.
if (
attn_mask is not None
and attn_mask.ndim == 2
@@ -32,7 +44,14 @@ def _maybe_reshape_attn_mask(query: torch.Tensor, key: torch.Tensor, attn_mask:
):
B, Sq, Skv = attn_mask.shape[0], query.shape[1], key.shape[1]
attn_mask = attn_mask.to(torch.bool)
attn_mask = attn_mask.unsqueeze(1).expand(B, Sq, Skv).unsqueeze(1).contiguous()
if mask_mode == "full_qk":
# NPU path requires explicit [B, 1, Q, K] mask layout.
attn_mask = attn_mask.unsqueeze(1).expand(B, Sq, Skv).unsqueeze(1).contiguous()
elif mask_mode == "broadcast_k":
# CUDA-like backends prefer [B, 1, 1, K] and rely on SDPA broadcast.
attn_mask = attn_mask.unsqueeze(1).unsqueeze(2)
else:
raise ValueError(f"Unsupported SDPA mask mode: {mask_mode}")
return attn_mask
@@ -70,15 +89,21 @@ class SDPAImpl(AttentionImpl):
self.causal = causal
self.softmax_scale = softmax_scale
def forward_cuda(
def _forward_impl(
self,
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
attn_metadata: AttentionMetadata | None = None,
mask_mode: SDPAMaskMode = "broadcast_k",
) -> torch.Tensor:
# Normalize mask before permuting q/k/v.
# _maybe_reshape_attn_mask expects sequence length on dim=1.
attention_mask = None
if attn_metadata:
attention_mask = _maybe_reshape_attn_mask(query, key, attn_metadata.attn_mask, mask_mode=mask_mode)
query, key, value = (x.permute(0, 2, 1, 3) for x in (query, key, value))
attention_mask = attn_metadata.attn_mask if attn_metadata else None
output = torch.nn.functional.scaled_dot_product_attention(
query,
key,
@@ -91,6 +116,15 @@ class SDPAImpl(AttentionImpl):
out = output.permute(0, 2, 1, 3)
return out
def forward_cuda(
self,
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
attn_metadata: AttentionMetadata | None = None,
) -> torch.Tensor:
return self._forward_impl(query, key, value, attn_metadata, mask_mode="broadcast_k")
def forward_xpu(
self,
query: torch.Tensor,
@@ -98,7 +132,7 @@ class SDPAImpl(AttentionImpl):
value: torch.Tensor,
attn_metadata: AttentionMetadata | None = None,
) -> torch.Tensor:
return self.forward_cuda(query, key, value, attn_metadata)
return self._forward_impl(query, key, value, attn_metadata, mask_mode="broadcast_k")
def forward_hip(
self,
@@ -107,7 +141,7 @@ class SDPAImpl(AttentionImpl):
value: torch.Tensor,
attn_metadata: AttentionMetadata | None = None,
) -> torch.Tensor:
return self.forward_cuda(query, key, value, attn_metadata)
return self._forward_impl(query, key, value, attn_metadata, mask_mode="broadcast_k")
def forward_npu(
self,
@@ -116,7 +150,4 @@ class SDPAImpl(AttentionImpl):
value: torch.Tensor,
attn_metadata: AttentionMetadata | None = None,
) -> torch.Tensor:
if attn_metadata:
attention_mask = _maybe_reshape_attn_mask(query, key, attn_metadata.attn_mask)
setattr(attn_metadata, "attn_mask", attention_mask)
return self.forward_cuda(query, key, value, attn_metadata)
return self._forward_impl(query, key, value, attn_metadata, mask_mode="full_qk")
vllm_omni/diffusion/distributed/vae_patch_parallel.py
+5 -5
View File
@@ -456,11 +456,11 @@ def maybe_wrap_vae_decode_with_patch_parallelism(
vae = getattr(pipeline, "vae", None)
if vae is None or not hasattr(vae, "decode"):
return
try:
from diffusers.models.autoencoders import AutoencoderKL
except Exception:
return
if not isinstance(vae, AutoencoderKL):
# NOTE: Use capability checks (not strict diffusers type checks) because
# NextStep's custom VAE is compatible but not a diffusers AutoencoderKL.
# TODO(vae-pp): Replace duck-typing with an explicit VAE compatibility
# validator/protocol and add per-model integration tests.
if not hasattr(vae, "decoder"):
return
if getattr(vae, "_vllm_vae_patch_parallel_installed", False):
vllm_omni/diffusion/models/nextstep_1_1/__init__.py
+12
View File
@@ -0,0 +1,12 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from vllm_omni.diffusion.models.nextstep_1_1.pipeline_nextstep_1_1 import (
NextStep11Pipeline,
get_nextstep11_post_process_func,
)
__all__ = [
"NextStep11Pipeline",
"get_nextstep11_post_process_func",
]
vllm_omni/diffusion/models/nextstep_1_1/modeling_flux_vae.py
+479
View File
@@ -0,0 +1,479 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# Adapted from NextStep-1 (https://github.com/stepfun-ai/NextStep-1)
import json
import os
from dataclasses import dataclass, field
import torch
import torch.nn.functional as F
from diffusers.models.autoencoders.vae import (
DecoderOutput,
DiagonalGaussianDistribution,
)
from diffusers.models.modeling_outputs import AutoencoderKLOutput
from diffusers.utils.torch_utils import randn_tensor
from einops import rearrange
from torch import Tensor, nn
@dataclass
class AutoEncoderParams:
resolution: int = 256
in_channels: int = 3
ch: int = 128
out_ch: int = 3
ch_mult: list[int] = field(default_factory=lambda: [1, 2, 4, 4])
num_res_blocks: int = 2
z_channels: int = 16
scaling_factor: float = 0.3611
shift_factor: float = 0.1159
deterministic: bool = False
encoder_norm: bool = False
psz: int | None = None
def swish(x: Tensor) -> Tensor:
return x * torch.sigmoid(x)
class AttnBlock(nn.Module):
def __init__(self, in_channels: int):
super().__init__()
self.in_channels = in_channels
self.norm = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
self.q = nn.Conv2d(in_channels, in_channels, kernel_size=1)
self.k = nn.Conv2d(in_channels, in_channels, kernel_size=1)
self.v = nn.Conv2d(in_channels, in_channels, kernel_size=1)
self.proj_out = nn.Conv2d(in_channels, in_channels, kernel_size=1)
def attention(self, h_: Tensor) -> Tensor:
h_ = self.norm(h_)
q = self.q(h_)
k = self.k(h_)
v = self.v(h_)
b, c, h, w = q.shape
q = rearrange(q, "b c h w -> b 1 (h w) c").contiguous()
k = rearrange(k, "b c h w -> b 1 (h w) c").contiguous()
v = rearrange(v, "b c h w -> b 1 (h w) c").contiguous()
h_ = nn.functional.scaled_dot_product_attention(q, k, v)
return rearrange(h_, "b 1 (h w) c -> b c h w", h=h, w=w, c=c, b=b)
def forward(self, x: Tensor) -> Tensor:
return x + self.proj_out(self.attention(x))
class ResnetBlock(nn.Module):
def __init__(self, in_channels: int, out_channels: int):
super().__init__()
self.in_channels = in_channels
out_channels = in_channels if out_channels is None else out_channels
self.out_channels = out_channels
self.norm1 = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
self.norm2 = nn.GroupNorm(num_groups=32, num_channels=out_channels, eps=1e-6, affine=True)
self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
if self.in_channels != self.out_channels:
self.nin_shortcut = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
def forward(self, x):
h = x
h = self.norm1(h)
h = swish(h)
h = self.conv1(h)
h = self.norm2(h)
h = swish(h)
h = self.conv2(h)
if self.in_channels != self.out_channels:
x = self.nin_shortcut(x)
return x + h
class Downsample(nn.Module):
def __init__(self, in_channels: int):
super().__init__()
self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=2, padding=0)
def forward(self, x: Tensor):
pad = (0, 1, 0, 1)
x = nn.functional.pad(x, pad, mode="constant", value=0)
x = self.conv(x)
return x
class Upsample(nn.Module):
def __init__(self, in_channels: int):
super().__init__()
self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1)
def forward(self, x: Tensor):
x = nn.functional.interpolate(x, scale_factor=2.0, mode="nearest")
x = self.conv(x)
return x
class Encoder(nn.Module):
def __init__(
self,
resolution: int,
in_channels: int,
ch: int,
ch_mult: list[int],
num_res_blocks: int,
z_channels: int,
):
super().__init__()
self.ch = ch
self.num_resolutions = len(ch_mult)
self.num_res_blocks = num_res_blocks
self.resolution = resolution
self.in_channels = in_channels
# downsampling
self.conv_in = nn.Conv2d(in_channels, self.ch, kernel_size=3, stride=1, padding=1)
curr_res = resolution
in_ch_mult = (1,) + tuple(ch_mult)
self.in_ch_mult = in_ch_mult
self.down = nn.ModuleList()
block_in = self.ch
for i_level in range(self.num_resolutions):
block = nn.ModuleList()
attn = nn.ModuleList()
block_in = ch * in_ch_mult[i_level]
block_out = ch * ch_mult[i_level]
for _ in range(self.num_res_blocks):
block.append(ResnetBlock(in_channels=block_in, out_channels=block_out))
block_in = block_out
down = nn.Module()
down.block = block
down.attn = attn
if i_level != self.num_resolutions - 1:
down.downsample = Downsample(block_in)
curr_res = curr_res // 2
self.down.append(down)
# middle
self.mid = nn.Module()
self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in)
self.mid.attn_1 = AttnBlock(block_in)
self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in)
# end
self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in, eps=1e-6, affine=True)
self.conv_out = nn.Conv2d(block_in, 2 * z_channels, kernel_size=3, stride=1, padding=1)
def forward(self, x: Tensor) -> Tensor:
# downsampling
hs = [self.conv_in(x)]
for i_level in range(self.num_resolutions):
for i_block in range(self.num_res_blocks):
h = self.down[i_level].block[i_block](hs[-1])
if len(self.down[i_level].attn) > 0:
h = self.down[i_level].attn[i_block](h)
hs.append(h)
if i_level != self.num_resolutions - 1:
hs.append(self.down[i_level].downsample(hs[-1]))
# middle
h = hs[-1]
h = self.mid.block_1(h)
h = self.mid.attn_1(h)
h = self.mid.block_2(h)
# end
h = self.norm_out(h)
h = swish(h)
h = self.conv_out(h)
return h
class Decoder(nn.Module):
def __init__(
self,
ch: int,
out_ch: int,
ch_mult: list[int],
num_res_blocks: int,
in_channels: int,
resolution: int,
z_channels: int,
):
super().__init__()
self.ch = ch
self.num_resolutions = len(ch_mult)
self.num_res_blocks = num_res_blocks
self.resolution = resolution
self.in_channels = in_channels
self.ffactor = 2 ** (self.num_resolutions - 1)
# compute in_ch_mult, block_in and curr_res at lowest res
block_in = ch * ch_mult[self.num_resolutions - 1]
curr_res = resolution // 2 ** (self.num_resolutions - 1)
self.z_shape = (1, z_channels, curr_res, curr_res)
# z to block_in
self.conv_in = nn.Conv2d(z_channels, block_in, kernel_size=3, stride=1, padding=1)
# middle
self.mid = nn.Module()
self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in)
self.mid.attn_1 = AttnBlock(block_in)
self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in)
# upsampling
self.up = nn.ModuleList()
for i_level in reversed(range(self.num_resolutions)):
block = nn.ModuleList()
attn = nn.ModuleList()
block_out = ch * ch_mult[i_level]
for _ in range(self.num_res_blocks + 1):
block.append(ResnetBlock(in_channels=block_in, out_channels=block_out))
block_in = block_out
up = nn.Module()
up.block = block
up.attn = attn
if i_level != 0:
up.upsample = Upsample(block_in)
curr_res = curr_res * 2
self.up.insert(0, up)
# end
self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in, eps=1e-6, affine=True)
self.conv_out = nn.Conv2d(block_in, out_ch, kernel_size=3, stride=1, padding=1)
def forward(self, z: Tensor) -> Tensor:
# get dtype for proper tracing
upscale_dtype = next(self.up.parameters()).dtype
# z to block_in
h = self.conv_in(z)
# middle
h = self.mid.block_1(h)
h = self.mid.attn_1(h)
h = self.mid.block_2(h)
# cast to proper dtype
h = h.to(upscale_dtype)
# upsampling
for i_level in reversed(range(self.num_resolutions)):
for i_block in range(self.num_res_blocks + 1):
h = self.up[i_level].block[i_block](h)
if len(self.up[i_level].attn) > 0:
h = self.up[i_level].attn[i_block](h)
if i_level != 0:
h = self.up[i_level].upsample(h)
# end
h = self.norm_out(h)
h = swish(h)
h = self.conv_out(h)
return h
def layer_norm_2d(input: torch.Tensor, eps: float = 1e-6) -> torch.Tensor:
# input.shape = (bsz, c, h, w)
_input = input.permute(0, 2, 3, 1)
_input = F.layer_norm(_input, _input.size()[-1:], None, None, eps)
_input = _input.permute(0, 3, 1, 2)
return _input
class AutoencoderKL(nn.Module):
def __init__(self, params: AutoEncoderParams):
super().__init__()
self.params = params
# Create a config-like object for compatibility
class Config:
def __init__(self, params):
self.latent_channels = params.z_channels
self.block_out_channels = params.ch_mult
self.scaling_factor = params.scaling_factor
self.shift_factor = params.shift_factor
self.out_channels = params.out_ch
self.config = Config(params)
self.encoder = Encoder(
resolution=params.resolution,
in_channels=params.in_channels,
ch=params.ch,
ch_mult=params.ch_mult,
num_res_blocks=params.num_res_blocks,
z_channels=params.z_channels,
)
self.decoder = Decoder(
resolution=params.resolution,
in_channels=params.in_channels,
ch=params.ch,
out_ch=params.out_ch,
ch_mult=params.ch_mult,
num_res_blocks=params.num_res_blocks,
z_channels=params.z_channels,
)
self.encoder_norm = params.encoder_norm
self.psz = params.psz
# Tiling attributes for VAE patch parallelism
self.use_tiling = False
vae_scale_factor = 2 ** (len(params.ch_mult) - 1)
self.tile_latent_min_size = int(params.resolution / vae_scale_factor)
self.tile_overlap_factor = 0.25
self.tile_sample_min_size = params.resolution
self.apply(self._init_weights)
def _init_weights(self, module):
std = 0.02
if isinstance(module, (nn.Conv2d, nn.Linear)):
module.weight.data.normal_(mean=0.0, std=std)
if module.bias is not None:
module.bias.data.zero_()
elif isinstance(module, nn.GroupNorm):
if module.weight is not None:
module.weight.data.fill_(1.0)
if module.bias is not None:
module.bias.data.zero_()
@property
def dtype(self):
return self.encoder.conv_in.weight.dtype
@property
def device(self):
return self.encoder.conv_in.weight.device
def patchify(self, img: torch.Tensor):
"""
img: (bsz, C, H, W)
x: (bsz, patch_size**2 * C, H / patch_size, W / patch_size)
"""
bsz, c, h, w = img.shape
p = self.psz
h_, w_ = h // p, w // p
img = img.reshape(bsz, c, h_, p, w_, p)
img = torch.einsum("nchpwq->ncpqhw", img)
x = img.reshape(bsz, c * p**2, h_, w_)
return x
def unpatchify(self, x: torch.Tensor):
"""
x: (bsz, patch_size**2 * C, H / patch_size, W / patch_size)
img: (bsz, C, H, W)
"""
bsz = x.shape[0]
p = self.psz
c = self.config.latent_channels
h_, w_ = x.shape[2], x.shape[3]
x = x.reshape(bsz, c, p, p, h_, w_)
x = torch.einsum("ncpqhw->nchpwq", x)
img = x.reshape(bsz, c, h_ * p, w_ * p)
return img
def encode(self, x: torch.Tensor, return_dict: bool = True):
moments = self.encoder(x)
mean, logvar = torch.chunk(moments, 2, dim=1)
if self.psz is not None:
mean = self.patchify(mean)
if self.encoder_norm:
mean = layer_norm_2d(mean)
if self.psz is not None:
mean = self.unpatchify(mean)
moments = torch.cat([mean, logvar], dim=1).contiguous()
posterior = DiagonalGaussianDistribution(moments, deterministic=self.params.deterministic)
if not return_dict:
return (posterior,)
return AutoencoderKLOutput(latent_dist=posterior)
def decode(self, z: torch.Tensor, return_dict: bool = True):
dec = self.decoder(z)
if not return_dict:
return (dec,)
return DecoderOutput(sample=dec)
def blend_v(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor:
blend_extent = min(a.shape[2], b.shape[2], blend_extent)
for y in range(blend_extent):
b[:, :, y, :] = a[:, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[:, :, y, :] * (y / blend_extent)
return b
def blend_h(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor:
blend_extent = min(a.shape[3], b.shape[3], blend_extent)
for x in range(blend_extent):
b[:, :, :, x] = a[:, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[:, :, :, x] * (x / blend_extent)
return b
def forward(self, input, sample_posterior=True, noise_strength=0.0):
posterior = self.encode(input).latent_dist
z = posterior.sample() if sample_posterior else posterior.mode()
if noise_strength > 0.0:
p = torch.distributions.Uniform(0, noise_strength)
z = z + p.sample((z.shape[0],)).reshape(-1, 1, 1, 1).to(z.device) * randn_tensor(
z.shape, device=z.device, dtype=z.dtype
)
dec = self.decode(z).sample
return dec, posterior
@classmethod
def from_pretrained(cls, model_path, **kwargs):
config_path = os.path.join(model_path, "config.json")
ckpt_path = os.path.join(model_path, "checkpoint.pt")
if not os.path.isdir(model_path):
raise ValueError(f"Model path does not exist: {model_path}")
if not os.path.isfile(config_path):
raise ValueError(f"Config file not found: {config_path}")
if not os.path.isfile(ckpt_path):
raise ValueError(f"Checkpoint file not found: {ckpt_path}")
state_dict = torch.load(ckpt_path, map_location="cpu", weights_only=True)
with open(config_path) as f:
config: dict = json.load(f)
config.update(kwargs)
# Filter out kwargs that are not in AutoEncoderParams
valid_kwargs = {}
valid_params = {
"resolution",
"in_channels",
"ch",
"out_ch",
"ch_mult",
"num_res_blocks",
"z_channels",
"scaling_factor",
"shift_factor",
"deterministic",
"encoder_norm",
"psz",
}
for key, value in config.items():
if key in valid_params:
valid_kwargs[key] = value
params = AutoEncoderParams(**valid_kwargs)
model = cls(params)
model.load_state_dict(state_dict, strict=False)
return model
vllm_omni/diffusion/models/nextstep_1_1/modeling_nextstep.py
+446
View File
@@ -0,0 +1,446 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# Adapted from NextStep-1.1 (https://huggingface.co/stepfun-ai/NextStep-1.1)
# Original: models/nextstep_model.py — local version with TP-aware layers.
from __future__ import annotations
import json
from collections.abc import Iterable
import numpy as np
import torch
import torch.nn as nn
from transformers.cache_utils import Cache, DynamicCache, StaticCache
from transformers.modeling_attn_mask_utils import AttentionMaskConverter
from transformers.modeling_outputs import BaseModelOutputWithPast
from transformers.models.llama.configuration_llama import LlamaConfig
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
from vllm_omni.diffusion.models.nextstep_1_1.modeling_nextstep_heads import (
FlowMatchingHead,
)
from vllm_omni.diffusion.models.nextstep_1_1.modeling_nextstep_llama import (
LlamaDecoderLayer,
LlamaRMSNorm,
LlamaRotaryEmbedding,
)
# ---------------------------------------------------------------------------
# Positional embedding utilities (inlined from remote utils/model_utils.py)
# ---------------------------------------------------------------------------
def _get_1d_sincos_pos_embed_from_grid(embed_dim: int, pos: np.ndarray) -> np.ndarray:
assert embed_dim % 2 == 0
omega = np.arange(embed_dim // 2, dtype=np.float64)
omega /= embed_dim / 2.0
omega = 1.0 / 10000**omega
pos = pos.reshape(-1)
out = np.einsum("m,d->md", pos, omega)
return np.concatenate([np.sin(out), np.cos(out)], axis=1)
def _get_2d_sincos_pos_embed_from_grid(embed_dim: int, grid: np.ndarray) -> np.ndarray:
assert embed_dim % 2 == 0
emb_h = _get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])
emb_w = _get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])
return np.concatenate([emb_h, emb_w], axis=1)
def get_2d_sincos_pos_embed(embed_dim: int, grid_size: int) -> np.ndarray:
grid_h = np.arange(grid_size, dtype=np.float32)
grid_w = np.arange(grid_size, dtype=np.float32)
grid = np.meshgrid(grid_w, grid_h)
grid = np.stack(grid, axis=0).reshape(2, 1, grid_size, grid_size)
return _get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
# ---------------------------------------------------------------------------
# NextStepConfig — extends LlamaConfig with NextStep-specific fields.
# This mirrors the remote models/config.py.
# ---------------------------------------------------------------------------
class NextStepConfig(LlamaConfig):
model_type = "nextstep"
def __init__(
self,
vae_name_or_path: str | None = None,
latent_size: int = 32,
latent_patch_size: int = 2,
latent_channels: int = 16,
boi: int | None = None,
eoi: int | None = None,
image_placeholder_id: int | None = None,
pad_token_id_added: int | None = None,
lm_loss_weight: float = 0.01,
im_loss_weight: float = 1.0,
fm_head_dim: int = 1536,
fm_head_layers: int = 12,
fm_head_batch_mul: int = 4,
o_attention_bias: bool | None = None,
**kwargs,
):
super().__init__(**kwargs)
self.vae_name_or_path = vae_name_or_path
self.latent_size = latent_size
self.latent_patch_size = latent_patch_size
self.latent_channels = latent_channels
self.boi = boi
self.eoi = eoi
self.image_placeholder_id = image_placeholder_id
self.pad_token_id_added = pad_token_id_added
self.lm_loss_weight = lm_loss_weight
self.im_loss_weight = im_loss_weight
self.fm_head_dim = fm_head_dim
self.fm_head_layers = fm_head_layers
self.fm_head_batch_mul = fm_head_batch_mul
self.o_attention_bias = self.attention_bias if o_attention_bias is None else o_attention_bias
@classmethod
def from_json(cls, path: str) -> NextStepConfig:
with open(path) as f:
data = json.load(f)
# Remove keys that are not constructor parameters (accept extra keys)
# LlamaConfig.__init__ uses **kwargs to absorb extras.
return cls(**data)
# ---------------------------------------------------------------------------
# NextStepModel — main model class
# ---------------------------------------------------------------------------
class NextStepModel(nn.Module):
def __init__(self, config: NextStepConfig):
super().__init__()
self.config = config
self.padding_idx = config.pad_token_id
self.vocab_size = config.vocab_size
self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
self.layers = nn.ModuleList(
[LlamaDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
)
self.norm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
self.rotary_emb = LlamaRotaryEmbedding(config=config)
# lm_head is part of the checkpoint but not used during image generation
self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
# Image projectors.
token_dim = config.latent_channels * config.latent_patch_size**2
self.image_in_projector = nn.Linear(token_dim, config.hidden_size)
self.image_out_projector = nn.Linear(config.hidden_size, config.hidden_size)
# Flow-matching head (no TP — tiny network)
self.image_head = FlowMatchingHead(
input_dim=token_dim,
cond_dim=config.hidden_size,
dim=config.fm_head_dim,
layers=config.fm_head_layers,
)
# Optional generation position embeddings
if getattr(config, "use_gen_pos_embed", False):
self._init_gen_pos_embed()
# ------------------------------------------------------------------
# Generation positional embeddings
# ------------------------------------------------------------------
def _init_gen_pos_embed(self):
self.register_buffer(
"gen_pos_embed",
torch.from_numpy(get_2d_sincos_pos_embed(self.config.hidden_size, self.config.base_image_grid_size))
.float()
.unsqueeze(0),
)
def gen_pos_embed_with_ar(self, h: int, w: int) -> torch.Tensor:
bsz, hw, dim = self.gen_pos_embed.shape
side = int(hw**0.5)
gen_pos_embed = self.gen_pos_embed.reshape(bsz, side, side, dim)
gen_pos_embed = gen_pos_embed[:, :h, :w, :]
return gen_pos_embed.reshape(bsz, -1, dim)
# ------------------------------------------------------------------
# Patchify / Unpatchify
# ------------------------------------------------------------------
def patchify(self, img: torch.Tensor) -> torch.Tensor:
bsz, c, h, w = img.shape
p = self.config.latent_patch_size
h_, w_ = h // p, w // p
img = img.reshape(bsz, c, h_, p, w_, p)
img = torch.einsum("nchpwq->nhwcpq", img)
return img.reshape(bsz, h_ * w_, c * p**2)
def unpatchify(
self,
x: torch.Tensor,
h: int | None = None,
w: int | None = None,
) -> torch.Tensor:
bsz = x.shape[0]
p = self.config.latent_patch_size
c = self.config.latent_channels
if h is None and w is None:
h_ = w_ = int(x.shape[1] ** 0.5)
else:
h_, w_ = h, w
assert h_ * w_ == x.shape[1], f"Invalid sequence length {x.shape[1]}."
x = x.reshape(bsz, h_, w_, c, p, p)
x = torch.einsum("nhwcpq->nchpwq", x)
return x.reshape(bsz, c, h_ * p, w_ * p)
# ------------------------------------------------------------------
# Input embedding preparation
# ------------------------------------------------------------------
def prepare_inputs_embeds(
self,
input_ids: torch.LongTensor,
latents: torch.FloatTensor | None = None,
) -> torch.Tensor:
if latents is None:
return self.embed_tokens(input_ids)
bs, seq_length = input_ids.shape
inputs_embeds = torch.zeros(
(bs, seq_length, self.config.hidden_size),
device=self.embed_tokens.weight.device,
dtype=self.embed_tokens.weight.dtype,
)
im_indices = input_ids == self.config.image_placeholder_id
lm_indices = ~im_indices
if isinstance(latents, list):
tokens = torch.cat([self.patchify(latent) for latent in latents], dim=1)
else:
tokens = self.patchify(latents)
image_embeds = self.image_in_projector(tokens)
image_embeds = image_embeds.view(-1, self.config.hidden_size)
token_embeds = self.embed_tokens(input_ids[lm_indices])
inputs_embeds[im_indices] = image_embeds.to(inputs_embeds.dtype)
inputs_embeds[lm_indices] = token_embeds
return inputs_embeds
# ------------------------------------------------------------------
# Causal mask utilities
# ------------------------------------------------------------------
def _update_causal_mask(
self,
attention_mask: torch.Tensor,
input_tensor: torch.Tensor,
cache_position: torch.Tensor,
past_key_values: Cache,
output_attentions: bool,
) -> torch.Tensor | None:
# For SDPA we build the 4D mask; flash_attention_2 uses None.
attn_impl = getattr(self.config, "_attn_implementation", "sdpa")
if attn_impl == "flash_attention_2":
if attention_mask is not None and (attention_mask == 0.0).any():
return attention_mask
return None
past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
using_static_cache = isinstance(past_key_values, StaticCache)
if attn_impl == "sdpa" and not using_static_cache and not output_attentions:
if AttentionMaskConverter._ignore_causal_mask_sdpa(
attention_mask,
inputs_embeds=input_tensor,
past_key_values_length=past_seen_tokens,
is_training=self.training,
):
return None
dtype, device = input_tensor.dtype, input_tensor.device
sequence_length = input_tensor.shape[1]
if using_static_cache:
target_length = past_key_values.get_max_cache_shape()
else:
target_length = (
attention_mask.shape[-1]
if isinstance(attention_mask, torch.Tensor)
else past_seen_tokens + sequence_length + 1
)
causal_mask = self._prepare_4d_causal_attention_mask_with_cache_position(
attention_mask,
sequence_length=sequence_length,
target_length=target_length,
dtype=dtype,
device=device,
cache_position=cache_position,
batch_size=input_tensor.shape[0],
)
if (
attn_impl == "sdpa"
and attention_mask is not None
and attention_mask.device.type == "cuda"
and not output_attentions
):
min_dtype = torch.finfo(dtype).min
causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype)
return causal_mask
@staticmethod
def _prepare_4d_causal_attention_mask_with_cache_position(
attention_mask: torch.Tensor,
sequence_length: int,
target_length: int,
dtype: torch.dtype,
device: torch.device,
cache_position: torch.Tensor,
batch_size: int,
**kwargs,
) -> torch.Tensor:
if attention_mask is not None and attention_mask.dim() == 4:
return attention_mask
min_dtype = torch.finfo(dtype).min
causal_mask = torch.full(
(sequence_length, target_length),
fill_value=min_dtype,
dtype=dtype,
device=device,
)
if sequence_length != 1:
causal_mask = torch.triu(causal_mask, diagonal=1)
causal_mask *= torch.arange(target_length, device=device) > cache_position.reshape(-1, 1)
causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1)
if attention_mask is not None:
causal_mask = causal_mask.clone()
mask_length = attention_mask.shape[-1]
padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :].to(causal_mask.device)
padding_mask = padding_mask == 0
causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill(padding_mask, min_dtype)
return causal_mask
# ------------------------------------------------------------------
# Forward through decoder layers
# ------------------------------------------------------------------
def forward_model(
self,
inputs_embeds: torch.FloatTensor,
attention_mask: torch.Tensor | None = None,
past_key_values: Cache | list[torch.FloatTensor] | None = None,
use_cache: bool | None = None,
output_attentions: bool | None = None,
output_hidden_states: bool | None = None,
cache_position: torch.LongTensor | None = None,
) -> BaseModelOutputWithPast:
output_attentions = (
output_attentions if output_attentions is not None else getattr(self.config, "output_attentions", False)
)
output_hidden_states = (
output_hidden_states
if output_hidden_states is not None
else getattr(self.config, "output_hidden_states", False)
)
use_cache = use_cache if use_cache is not None else getattr(self.config, "use_cache", True)
if use_cache and past_key_values is None:
past_key_values = DynamicCache()
if cache_position is None:
past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
cache_position = torch.arange(
past_seen_tokens,
past_seen_tokens + inputs_embeds.shape[1],
device=inputs_embeds.device,
)
position_ids = cache_position.unsqueeze(0)
causal_mask = self._update_causal_mask(
attention_mask,
inputs_embeds,
cache_position,
past_key_values,
output_attentions,
)
hidden_states = inputs_embeds
position_embeddings = self.rotary_emb(hidden_states, position_ids)
all_hidden_states = () if output_hidden_states else None
all_self_attns = () if output_attentions else None
for decoder_layer in self.layers:
if output_hidden_states:
all_hidden_states += (hidden_states,)
layer_outputs = decoder_layer(
hidden_states,
attention_mask=causal_mask,
past_key_value=past_key_values,
output_attentions=output_attentions,
use_cache=use_cache,
cache_position=cache_position,
position_embeddings=position_embeddings,
)
hidden_states = layer_outputs[0]
if output_attentions:
all_self_attns += (layer_outputs[1],)
hidden_states = self.norm(hidden_states)
if output_hidden_states:
all_hidden_states += (hidden_states,)
return BaseModelOutputWithPast(
last_hidden_state=hidden_states,
past_key_values=past_key_values if use_cache else None,
hidden_states=all_hidden_states,
attentions=all_self_attns,
)
# ------------------------------------------------------------------
# Weight loading with TP sharding support
# ------------------------------------------------------------------
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
stacked_params_mapping = [
# (param_name_in_model, weight_name_in_checkpoint, shard_id)
(".qkv_proj", ".q_proj", "q"),
(".qkv_proj", ".k_proj", "k"),
(".qkv_proj", ".v_proj", "v"),
(".gate_up_proj", ".gate_proj", 0),
(".gate_up_proj", ".up_proj", 1),
]
params_dict = dict(self.named_parameters())
loaded_params: set[str] = set()
for name, loaded_weight in weights:
for param_name, weight_name, shard_id in stacked_params_mapping:
if weight_name not in name:
continue
name = name.replace(weight_name, param_name)
param = params_dict[name]
weight_loader = param.weight_loader
weight_loader(param, loaded_weight, shard_id)
break
else:
if name not in params_dict:
continue
param = params_dict[name]
weight_loader = getattr(param, "weight_loader", default_weight_loader)
weight_loader(param, loaded_weight)
loaded_params.add(name)
return loaded_params
vllm_omni/diffusion/models/nextstep_1_1/modeling_nextstep_heads.py
+309
View File
@@ -0,0 +1,309 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# Adapted from NextStep-1.1 (https://huggingface.co/stepfun-ai/NextStep-1.1)
# Original: models/heads.py — FlowMatchingHead and components.
# No TP needed: the FM head is tiny (dim=1536, 12 layers).
from __future__ import annotations
import math
import torch
import torch.nn as nn
# ---------------------------------------------------------------------------
# Utilities (inlined from remote utils/model_utils.py)
# ---------------------------------------------------------------------------
def expand_t(t: torch.Tensor, x: torch.Tensor) -> torch.Tensor:
dims = [1] * (len(x.size()) - 1)
return t.view(t.size(0), *dims)
def randn_tensor(
shape: tuple[int, ...],
noise_repeat: int,
device: torch.device,
dtype: torch.dtype = torch.float32,
) -> torch.Tensor:
bsz = shape[0]
if bsz % noise_repeat != 0:
raise ValueError(f"Batch size ({bsz}) must be divisible by noise repeat ({noise_repeat})")
_shape = (noise_repeat,) + shape[1:]
_tensor = torch.randn(_shape, device=device, dtype=dtype).repeat(bsz // noise_repeat, 1)
return _tensor
# ---------------------------------------------------------------------------
# Adaptive LayerNorm modulation
# ---------------------------------------------------------------------------
def modulate(
x: torch.Tensor,
shift: torch.Tensor | None,
scale: torch.Tensor | None = None,
) -> torch.Tensor:
if shift is None:
return x * (1 + scale)
return x * (1 + scale) + shift
# ---------------------------------------------------------------------------
# ResBlock
# ---------------------------------------------------------------------------
class ResBlock(nn.Module):
def __init__(self, channels: int, mlp_ratio: float = 1.0):
super().__init__()
self.channels = channels
self.intermediate_size = int(channels * mlp_ratio)
self.in_ln = nn.LayerNorm(self.channels, eps=1e-6)
self.mlp = nn.Sequential(
nn.Linear(self.channels, self.intermediate_size),
nn.SiLU(),
nn.Linear(self.intermediate_size, self.channels),
)
self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(channels, 3 * channels, bias=True))
def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(y).chunk(3, dim=-1)
h = modulate(self.in_ln(x), shift_mlp, scale_mlp)
h = self.mlp(h)
return x + gate_mlp * h
# ---------------------------------------------------------------------------
# FinalLayer
# ---------------------------------------------------------------------------
class FinalLayer(nn.Module):
def __init__(self, model_channels: int, out_channels: int):
super().__init__()
self.norm_final = nn.LayerNorm(model_channels, elementwise_affine=False, eps=1e-6)
self.linear = nn.Linear(model_channels, out_channels, bias=True)
self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(model_channels, 2 * model_channels, bias=True))
def forward(self, x: torch.Tensor, c: torch.Tensor) -> torch.Tensor:
shift, scale = self.adaLN_modulation(c).chunk(2, dim=-1)
x = modulate(self.norm_final(x), shift, scale)
x = self.linear(x)
return x
# ---------------------------------------------------------------------------
# TimestepEmbedder
# ---------------------------------------------------------------------------
class TimestepEmbedder(nn.Module):
def __init__(self, hidden_size: int, frequency_embedding_size: int = 256):
super().__init__()
self.mlp = nn.Sequential(
nn.Linear(frequency_embedding_size, hidden_size, bias=True),
nn.SiLU(),
nn.Linear(hidden_size, hidden_size, bias=True),
)
self.frequency_embedding_size = frequency_embedding_size
@staticmethod
def timestep_embedding(t: torch.Tensor, dim: int, max_period: float = 10000.0) -> torch.Tensor:
half = dim // 2
freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to(
device=t.device
)
args = t[:, None].float() * freqs[None]
embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
if dim % 2:
embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
return embedding
def forward(self, t: torch.Tensor) -> torch.Tensor:
t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
return self.mlp(t_freq.to(self.mlp[0].weight.dtype))
# ---------------------------------------------------------------------------
# SimpleMLPAdaLN
# ---------------------------------------------------------------------------
class SimpleMLPAdaLN(nn.Module):
def __init__(
self,
input_dim: int,
cond_dim: int,
dim: int = 1536,
layers: int = 12,
mlp_ratio: float = 1.0,
):
super().__init__()
self.input_dim = input_dim
self.cond_dim = cond_dim
self.dim = dim
self.time_embed = TimestepEmbedder(dim)
self.cond_embed = nn.Linear(cond_dim, dim)
self.input_proj = nn.Linear(input_dim, dim)
self.res_blocks = nn.ModuleList([ResBlock(dim, mlp_ratio) for _ in range(layers)])
self.final_layer = FinalLayer(dim, input_dim)
self.initialize_weights()
def initialize_weights(self):
def _basic_init(module):
if isinstance(module, nn.Linear):
torch.nn.init.xavier_uniform_(module.weight)
if module.bias is not None:
nn.init.constant_(module.bias, 0)
self.apply(_basic_init)
nn.init.normal_(self.time_embed.mlp[0].weight, std=0.02)
nn.init.normal_(self.time_embed.mlp[2].weight, std=0.02)
for block in self.res_blocks:
nn.init.constant_(block.adaLN_modulation[-1].weight, 0)
nn.init.constant_(block.adaLN_modulation[-1].bias, 0)
nn.init.constant_(self.final_layer.adaLN_modulation[-1].weight, 0)
nn.init.constant_(self.final_layer.adaLN_modulation[-1].bias, 0)
nn.init.constant_(self.final_layer.linear.weight, 0)
nn.init.constant_(self.final_layer.linear.bias, 0)
def forward(self, x: torch.Tensor, t: torch.Tensor, c: torch.Tensor) -> torch.Tensor:
x = self.input_proj(x)
t = self.time_embed(t)
c = self.cond_embed(c)
y = t + c
for block in self.res_blocks:
x = block(x, y)
return self.final_layer(x, y)
# ---------------------------------------------------------------------------
# FlowMatchingHead
# ---------------------------------------------------------------------------
class FlowMatchingHead(nn.Module):
def __init__(
self,
input_dim: int,
cond_dim: int,
dim: int = 1536,
layers: int = 12,
mlp_ratio: float = 1.0,
):
super().__init__()
self.input_dim = input_dim
self.net = SimpleMLPAdaLN(
input_dim=input_dim,
cond_dim=cond_dim,
dim=dim,
layers=layers,
mlp_ratio=mlp_ratio,
)
@property
def dtype(self):
return self.net.input_proj.weight.dtype
@property
def device(self):
return self.net.input_proj.weight.device
def get_score_from_velocity(
self,
velocity: torch.Tensor,
x: torch.Tensor,
t: torch.Tensor,
) -> torch.Tensor:
t = expand_t(t, x)
alpha_t, d_alpha_t = t, 1
sigma_t, d_sigma_t = 1 - t, -1
mean = x
reverse_alpha_ratio = alpha_t / d_alpha_t
var = sigma_t**2 - reverse_alpha_ratio * d_sigma_t * sigma_t
score = (reverse_alpha_ratio * velocity - mean) / var
return score
def get_velocity_from_cfg(
self,
velocity: torch.Tensor,
cfg: float,
cfg_img: float,
cfg_mult: int,
) -> torch.Tensor:
if cfg_mult == 2:
cond_v, uncond_v = torch.chunk(velocity, 2, dim=0)
velocity = uncond_v + cfg * (cond_v - uncond_v)
elif cfg_mult == 3:
cond_v, uncond_v1, uncond_v2 = torch.chunk(velocity, 3, dim=0)
velocity = uncond_v2 + cfg_img * (uncond_v1 - uncond_v2) + cfg * (cond_v - uncond_v1)
return velocity
@torch.no_grad()
def sample(
self,
c: torch.Tensor,
cfg: float = 1.0,
cfg_img: float = 1.0,
cfg_mult: int | None = None,
timesteps_shift: float = 1.0,
num_sampling_steps: int = 20,
last_step_size: float = 0.0,
noise_repeat: int = 1,
) -> torch.Tensor:
if cfg_mult is None:
cfg_mult = 1
if cfg > 1.0:
cfg_mult += 1
if cfg_img > 1.0:
cfg_mult += 1
if cfg_mult <= 0:
raise ValueError(f"Invalid cfg_mult={cfg_mult}; expected a positive value.")
if c.shape[0] % cfg_mult != 0:
raise ValueError(
f"Invalid CFG layout: condition batch size {c.shape[0]} is not divisible by cfg_mult={cfg_mult}."
)
noise = randn_tensor((c.shape[0] // cfg_mult, self.input_dim), noise_repeat, self.device)
x = noise
xs = []
t0, t1 = 0, 1
timesteps = torch.linspace(t0, t1, num_sampling_steps + 1, device=c.device)[:-1]
timesteps = timesteps / (timesteps_shift - (timesteps_shift - 1) * timesteps)
timesteps = torch.cat([timesteps, torch.ones(1, device=c.device)])
for ti, tj in zip(timesteps[:-1], timesteps[1:]):
dt = tj - ti
combined = torch.cat([x] * cfg_mult, dim=0)
velocity = self.net(combined.to(c.dtype), ti.expand(c.shape[0]).to(c), c)
velocity = velocity.to(torch.float32)
velocity = self.get_velocity_from_cfg(velocity, cfg, cfg_img, cfg_mult)
score = self.get_score_from_velocity(velocity, x, ti.expand(x.shape[0]).to(x))
drift = velocity + (1 - expand_t(ti.expand(x.shape[0]).to(x), x)) * score
w_cur = randn_tensor((c.shape[0] // cfg_mult, self.input_dim), noise_repeat, self.device)
dw = w_cur * torch.sqrt(dt)
mean_x = x + drift * dt
x = mean_x + torch.sqrt(2 * (1 - expand_t(ti.expand(x.shape[0]).to(x), x))) * dw
xs.append(x)
if len(xs) != num_sampling_steps:
raise ValueError(f"Samples ({len(xs)}) does not match the number of steps ({num_sampling_steps})")
return xs[-1].to(c.dtype)
vllm_omni/diffusion/models/nextstep_1_1/modeling_nextstep_llama.py
+285
View File
@@ -0,0 +1,285 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# Adapted from NextStep-1.1 (https://huggingface.co/stepfun-ai/NextStep-1.1)
# Original: models/llama_model.py — made TP-aware for vLLM-Omni.
from __future__ import annotations
import torch
import torch.nn as nn
from transformers import ROPE_INIT_FUNCTIONS
from transformers.cache_utils import Cache
from vllm.model_executor.layers.linear import (
MergedColumnParallelLinear,
QKVParallelLinear,
RowParallelLinear,
)
# ---------------------------------------------------------------------------
# Utilities (inlined from remote utils/model_utils.py)
# ---------------------------------------------------------------------------
def rotate_half(x: torch.Tensor) -> torch.Tensor:
x1 = x[..., : x.shape[-1] // 2]
x2 = x[..., x.shape[-1] // 2 :]
return torch.cat((-x2, x1), dim=-1)
def apply_rotary_pos_emb(
q: torch.Tensor,
k: torch.Tensor,
cos: torch.Tensor,
sin: torch.Tensor,
unsqueeze_dim: int = 1,
) -> tuple[torch.Tensor, torch.Tensor]:
cos = cos.unsqueeze(unsqueeze_dim)
sin = sin.unsqueeze(unsqueeze_dim)
q_embed = (q * cos) + (rotate_half(q) * sin)
k_embed = (k * cos) + (rotate_half(k) * sin)
return q_embed, k_embed
def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
batch, num_kv_heads, slen, head_dim = hidden_states.shape
if n_rep == 1:
return hidden_states
hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_kv_heads, n_rep, slen, head_dim)
return hidden_states.reshape(batch, num_kv_heads * n_rep, slen, head_dim)
# ---------------------------------------------------------------------------
# LlamaRMSNorm — no TP needed
# ---------------------------------------------------------------------------
class LlamaRMSNorm(nn.Module):
def __init__(self, hidden_size: int, eps: float = 1e-6):
super().__init__()
self.weight = nn.Parameter(torch.ones(hidden_size))
self.variance_epsilon = eps
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
input_dtype = hidden_states.dtype
hidden_states = hidden_states.to(torch.float32)
variance = hidden_states.pow(2).mean(-1, keepdim=True)
hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
return self.weight * hidden_states.to(input_dtype)
# ---------------------------------------------------------------------------
# LlamaRotaryEmbedding — no TP needed
# ---------------------------------------------------------------------------
class LlamaRotaryEmbedding(nn.Module):
def __init__(self, config, device=None):
super().__init__()
self.rope_type = "default"
self.config = config
self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
self.register_buffer("inv_freq", inv_freq, persistent=False)
@torch.no_grad()
def forward(self, x: torch.Tensor, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)
position_ids_expanded = position_ids[:, None, :].float()
device_type = x.device.type
device_type = device_type if isinstance(device_type, str) and device_type != "mps" else "cpu"
with torch.autocast(device_type=device_type, enabled=False):
freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
emb = torch.cat((freqs, freqs), dim=-1)
cos = emb.cos() * self.attention_scaling
sin = emb.sin() * self.attention_scaling
return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
# ---------------------------------------------------------------------------
# LlamaAttention — TP-aware (fused QKV + RowParallel o_proj)
# ---------------------------------------------------------------------------
class LlamaAttention(nn.Module):
def __init__(self, config, layer_idx: int):
super().__init__()
self.config = config
self.layer_idx = layer_idx
self.attention_dropout = config.attention_dropout
self.hidden_size = config.hidden_size
self.num_heads = config.num_attention_heads
self.head_dim = getattr(config, "head_dim", self.hidden_size // self.num_heads)
self.num_key_value_heads = config.num_key_value_heads
self.num_key_value_groups = self.num_heads // self.num_key_value_heads
self.is_causal = True
# TP-aware: fused QKV projection
self.qkv_proj = QKVParallelLinear(
hidden_size=self.hidden_size,
head_size=self.head_dim,
total_num_heads=self.num_heads,
total_num_kv_heads=self.num_key_value_heads,
bias=config.attention_bias,
)
# TP-aware: row-parallel output projection
self.o_proj = RowParallelLinear(
self.num_heads * self.head_dim,
self.hidden_size,
bias=getattr(config, "o_attention_bias", config.attention_bias),
)
def forward(
self,
hidden_states: torch.Tensor,
attention_mask: torch.Tensor | None = None,
past_key_value: Cache | None = None,
output_attentions: bool = False,
use_cache: bool = False,
cache_position: torch.LongTensor | None = None,
position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
**kwargs,
) -> tuple[torch.Tensor, torch.Tensor | None, Cache | None]:
bsz, q_len, _ = hidden_states.size()
# Fused QKV projection
qkv, _ = self.qkv_proj(hidden_states)
# Split into Q, K, V — sizes account for TP sharding
qkv = qkv.view(bsz, q_len, -1, self.head_dim)
# QKVParallelLinear output: [q_heads_local, k_heads_local, v_heads_local]
num_local_heads = self.qkv_proj.num_heads
num_local_kv_heads = self.qkv_proj.num_kv_heads
split_sizes = [num_local_heads, num_local_kv_heads, num_local_kv_heads]
query_states, key_states, value_states = qkv.split(split_sizes, dim=2)
query_states = query_states.transpose(1, 2)
key_states = key_states.transpose(1, 2)
value_states = value_states.transpose(1, 2)
cos, sin = position_embeddings
query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
if past_key_value is not None:
cache_kwargs = {
"sin": sin,
"cos": cos,
"cache_position": cache_position,
}
key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
num_local_kv_groups = num_local_heads // num_local_kv_heads
key_states = repeat_kv(key_states, num_local_kv_groups)
value_states = repeat_kv(value_states, num_local_kv_groups)
causal_mask = attention_mask
if attention_mask is not None:
causal_mask = causal_mask[:, :, :, : key_states.shape[-2]]
if query_states.device.type == "cuda" and causal_mask is not None:
query_states = query_states.contiguous()
key_states = key_states.contiguous()
value_states = value_states.contiguous()
is_causal = True if causal_mask is None and q_len > 1 else False
attn_output = torch.nn.functional.scaled_dot_product_attention(
query_states,
key_states,
value_states,
attn_mask=causal_mask,
dropout_p=self.attention_dropout if self.training else 0.0,
is_causal=is_causal,
)
attn_output = attn_output.transpose(1, 2).contiguous()
attn_output = attn_output.view(bsz, q_len, -1)
attn_output, _ = self.o_proj(attn_output)
return attn_output, None, past_key_value
# ---------------------------------------------------------------------------
# LlamaMLP — TP-aware (fused gate_up + RowParallel down)
# ---------------------------------------------------------------------------
class LlamaMLP(nn.Module):
def __init__(self, config):
super().__init__()
self.hidden_size = config.hidden_size
self.intermediate_size = config.intermediate_size
# TP-aware: fused gate + up projection
self.gate_up_proj = MergedColumnParallelLinear(
self.hidden_size,
[self.intermediate_size] * 2,
bias=config.mlp_bias,
)
# TP-aware: row-parallel down projection
self.down_proj = RowParallelLinear(
self.intermediate_size,
self.hidden_size,
bias=config.mlp_bias,
)
self.act_fn = nn.SiLU()
def forward(self, x: torch.Tensor) -> torch.Tensor:
gate_up, _ = self.gate_up_proj(x)
gate, up = gate_up.chunk(2, dim=-1)
down, _ = self.down_proj(self.act_fn(gate) * up)
return down
# ---------------------------------------------------------------------------
# LlamaDecoderLayer — uses TP-aware Attention + MLP
# ---------------------------------------------------------------------------
class LlamaDecoderLayer(nn.Module):
def __init__(self, config, layer_idx: int):
super().__init__()
self.hidden_size = config.hidden_size
self.self_attn = LlamaAttention(config=config, layer_idx=layer_idx)
self.mlp = LlamaMLP(config)
self.input_layernorm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
self.post_attention_layernorm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
def forward(
self,
hidden_states: torch.Tensor,
attention_mask: torch.Tensor | None = None,
past_key_value: Cache | None = None,
output_attentions: bool = False,
use_cache: bool = False,
cache_position: torch.LongTensor | None = None,
position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
**kwargs,
) -> tuple:
residual = hidden_states
hidden_states = self.input_layernorm(hidden_states)
hidden_states, self_attn_weights, present_key_value = self.self_attn(
hidden_states=hidden_states,
attention_mask=attention_mask,
past_key_value=past_key_value,
output_attentions=output_attentions,
use_cache=use_cache,
cache_position=cache_position,
position_embeddings=position_embeddings,
**kwargs,
)
hidden_states = residual + hidden_states
residual = hidden_states
hidden_states = self.post_attention_layernorm(hidden_states)
hidden_states = self.mlp(hidden_states)
hidden_states = residual + hidden_states
outputs = (hidden_states,)
if output_attentions:
outputs += (self_attn_weights,)
if use_cache:
outputs += (present_key_value,)
return outputs
vllm_omni/diffusion/models/nextstep_1_1/pipeline_nextstep_1_1.py
+710
View File
@@ -0,0 +1,710 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# Adapted from NextStep-1.1 (https://huggingface.co/stepfun-ai/NextStep-1.1)
import os
import re
from collections.abc import Iterable
from typing import Literal
import numpy as np
import torch
import torch.nn.functional as F
import torchvision.transforms as transforms
from diffusers.image_processor import VaeImageProcessor
from PIL import Image
from torch import nn
from tqdm.auto import tqdm
from transformers import AutoTokenizer, PreTrainedTokenizer
from transformers.cache_utils import StaticCache
from vllm.logger import init_logger
from vllm.model_executor.models.utils import AutoWeightsLoader
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
from vllm_omni.diffusion.distributed.parallel_state import (
get_cfg_group,
get_classifier_free_guidance_rank,
get_classifier_free_guidance_world_size,
)
from vllm_omni.diffusion.distributed.utils import get_local_device
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
from vllm_omni.diffusion.models.nextstep_1_1.modeling_flux_vae import AutoencoderKL
from vllm_omni.diffusion.models.nextstep_1_1.modeling_nextstep import (
NextStepConfig,
NextStepModel,
)
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.model_executor.model_loader.weight_utils import (
download_weights_from_hf_specific,
)
logger = init_logger(__name__)
def layer_norm(input: torch.Tensor, normalized_shape: torch.Size, eps: float = 1e-6) -> torch.Tensor:
return F.layer_norm(input, normalized_shape, None, None, eps)
def set_seed(seed: int):
"""Set random seed for reproducibility."""
torch.manual_seed(seed)
if torch.cuda.is_available():
torch.cuda.manual_seed_all(seed)
def to_pil(image: torch.Tensor, mode: str = "11") -> Image.Image:
"""Convert tensor to PIL Image."""
if mode == "11":
# Assuming image is in [-1, 1] range
image = (image + 1) / 2
image = image.clamp(0, 1)
image = image.permute(1, 2, 0).cpu().numpy()
image = (image * 255).astype(np.uint8)
return Image.fromarray(image)
#################################################################################
# Sine/Cosine Positional Embedding Functions #
#################################################################################
def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False, extra_tokens=0, pe_interpolation=1.0):
"""
grid_size: int of the grid height and width
return:
pos_embed: [grid_size*grid_size, embed_dim] or [1+grid_size*grid_size, embed_dim]
"""
grid_h = np.arange(grid_size, dtype=np.float32) / pe_interpolation
grid_w = np.arange(grid_size, dtype=np.float32) / pe_interpolation
grid = np.meshgrid(grid_w, grid_h)
grid = np.stack(grid, axis=0)
grid = grid.reshape([2, 1, grid_size, grid_size])
pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
if cls_token and extra_tokens > 0:
pos_embed = np.concatenate([np.zeros([extra_tokens, embed_dim]), pos_embed], axis=0)
return pos_embed
def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):
assert embed_dim % 2 == 0
emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])
emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])
emb = np.concatenate([emb_h, emb_w], axis=1)
return emb
def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
assert embed_dim % 2 == 0
omega = np.arange(embed_dim // 2, dtype=np.float64)
omega /= embed_dim / 2.0
omega = 1.0 / 10000**omega
pos = pos.reshape(-1)
out = np.einsum("m,d->md", pos, omega)
emb_sin = np.sin(out)
emb_cos = np.cos(out)
emb = np.concatenate([emb_sin, emb_cos], axis=1)
return emb
def hw2str(h: int, w: int) -> str:
"""Convert height and width to string format."""
return f"{h}*{w}"
DEFAULT_IMAGE_AREA_TOKEN = "<|image_area|>"
def get_nextstep11_post_process_func(od_config: OmniDiffusionConfig):
"""Return post-processing function for NextStep-1.1 pipeline outputs."""
vae_scale_factor = 8 # Default for NextStep VAE
image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor)
def post_process_func(images: torch.Tensor):
return image_processor.postprocess(images)
return post_process_func
class NextStep11Pipeline(nn.Module):
"""
NextStep-1.1 Pipeline for text-to-image generation.
This pipeline implements the autoregressive flow-based image generation
model from StepFun. It uses an LLM backbone with a flow matching head
to generate images autoregressively.
"""
def __init__(
self,
*,
od_config: OmniDiffusionConfig,
prefix: str = "",
):
super().__init__()
self.od_config = od_config
self._execution_device = get_local_device()
model_path = od_config.model
local_files_only = os.path.exists(model_path)
if not local_files_only:
model_path = download_weights_from_hf_specific(model_path, None, ["*"])
# Load tokenizer (still uses trust_remote_code for tokenizer only)
self.tokenizer: PreTrainedTokenizer = AutoTokenizer.from_pretrained(
model_path,
local_files_only=True,
model_max_length=512,
padding_side="left",
use_fast=True,
trust_remote_code=True,
)
self.tokenizer.add_eos_token = False
# Load model from local TP-aware code (weights loaded later via load_weights)
config = NextStepConfig.from_json(os.path.join(model_path, "config.json"))
self.model = NextStepModel(config)
self.model.eval()
# Load config
self.config = self.model.config
# Load VAE
vae_path = getattr(self.config, "vae_name_or_path", None)
if vae_path is None:
vae_path = os.path.join(model_path, "vae")
elif not os.path.isabs(vae_path):
# Resolve relative vae_name_or_path (e.g. "vae/") against model dir
vae_path = os.path.join(model_path, vae_path)
if os.path.exists(vae_path):
self.vae = AutoencoderKL.from_pretrained(vae_path)
else:
# Try loading from model directory
vae_checkpoint = os.path.join(model_path, "vae", "checkpoint.pt")
vae_config = os.path.join(model_path, "vae", "config.json")
if os.path.exists(vae_checkpoint) and os.path.exists(vae_config):
self.vae = AutoencoderKL.from_pretrained(os.path.join(model_path, "vae"))
else:
raise ValueError(f"Could not find VAE at {vae_path}")
self.vae.eval()
# Calculate down factor
vae_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
latent_patch_size = getattr(self.config, "latent_patch_size", 2)
self.down_factor = vae_factor * latent_patch_size
# Get VAE parameters
self.shift_factor = getattr(self.vae.config, "shift_factor", 0.0)
self.scaling_factor = getattr(self.vae.config, "scaling_factor", 1.0)
# Get special token IDs from config
self.boi = getattr(self.config, "boi", None)
self.eoi = getattr(self.config, "eoi", None)
self.image_placeholder_id = getattr(self.config, "image_placeholder_id", None)
# Image processing
self.pil2tensor = transforms.Compose(
[
transforms.ToTensor(),
transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]),
]
)
self.__device = self._execution_device
self.__dtype = od_config.dtype
# Weight sources: model weights from safetensors, prefixed with "model."
# so AutoWeightsLoader dispatches to self.model.load_weights()
self.weights_sources = [
DiffusersPipelineLoader.ComponentSource(
model_or_path=model_path,
subfolder=None,
revision=None,
prefix="model.",
fall_back_to_pt=True,
allow_patterns_overrides=["model-*.safetensors", "model.safetensors"],
)
]
@property
def device(self):
return self.__device
@property
def dtype(self):
return self.__dtype
def to(self, device=None, dtype=None):
if device is not None:
self.__device = device
if dtype is not None:
self.__dtype = dtype
self.model.to(self.__device, dtype=self.__dtype)
self.vae.to(self.__device, dtype=self.__dtype)
return self
def _image_str(self, hw: tuple[int, int] = (256, 256)):
"""Generate image token string for given height/width."""
latent_hw = (hw[0] // self.down_factor, hw[1] // self.down_factor)
image_ids = [self.boi] + [self.image_placeholder_id] * (latent_hw[0] * latent_hw[1]) + [self.eoi]
image_str = DEFAULT_IMAGE_AREA_TOKEN + hw2str(*latent_hw) + self.tokenizer.decode(image_ids)
return image_str
def _check_input(
self, captions: str | list[str], images: Image.Image | list[Image.Image] | None
) -> tuple[list[str], list[Image.Image] | None]:
"""Validate and process input captions and images."""
if not isinstance(captions, list):
captions = [captions]
if images is not None:
if not isinstance(images, list):
images = [images]
# Validate image count matches <image> tokens in captions
image_token_count = 0
for caption in captions:
num_image_token = len(re.findall(r"<image>", caption))
if num_image_token != 1:
raise ValueError(
f"Caption must contain exactly one <image> token. "
f"Found {num_image_token} in: {caption[:100]}..."
)
image_token_count += num_image_token
if image_token_count != len(images):
raise ValueError(
f"Number of images ({len(images)}) does not match number of image tokens ({image_token_count})."
)
hws = [(image.size[1], image.size[0]) for image in images]
# Replace <image> tokens with corresponding image_str
processed_captions = []
image_idx = 0
for caption in captions:
processed_caption = caption
num_image_tokens = processed_caption.count("<image>")
for _ in range(num_image_tokens):
processed_caption = processed_caption.replace("<image>", self._image_str(hws[image_idx]), 1)
image_idx += 1
processed_captions.append(processed_caption)
captions = processed_captions
return captions, images
@staticmethod
def _resolve_cfg_layout(cfg: float, cfg_img: float, has_image_conditions: bool) -> tuple[int, float]:
"""Resolve the active CFG branch layout for the current request."""
use_text_cfg = cfg > 1.0
# Image CFG branch is only meaningful when the request has an image condition.
use_img_cfg = use_text_cfg and has_image_conditions and cfg_img != 1.0
cfg_mult = 1 + int(use_text_cfg) + int(use_img_cfg)
effective_cfg_img = cfg_img if use_img_cfg else 1.0
return cfg_mult, effective_cfg_img
def _build_captions(
self,
captions: str | list[str],
images: list[Image.Image] | None = None,
num_images_per_caption: int = 1,
positive_prompt: str | None = None,
negative_prompt: str | None = None,
cfg: float = 1.0,
cfg_img: float = 1.0,
) -> tuple[list[str], list[Image.Image] | None, int, float]:
"""Build captions with CFG support."""
if not isinstance(captions, list):
captions = [captions]
captions = [caption for caption in captions for _ in range(num_images_per_caption)]
if images is not None:
images = [image for image in images for _ in range(num_images_per_caption)]
# Add positive prompt
if positive_prompt is not None and positive_prompt != "":
captions = [f"{caption} {positive_prompt}" for caption in captions]
cfg_mult, effective_cfg_img = self._resolve_cfg_layout(cfg, cfg_img, images is not None)
# Add negative prompt for CFG
if negative_prompt is None:
negative_prompt = ""
num_samples = len(captions)
if cfg_mult == 3:
w, h = images[0].size
captions = captions + [self._image_str((h, w)) + negative_prompt] * num_samples
images = images + images
captions = captions + [negative_prompt] * num_samples
elif cfg_mult == 2:
captions = captions + [negative_prompt] * num_samples
return captions, images, cfg_mult, effective_cfg_img
def _add_prefix_ids(self, hw: tuple[int, int], input_ids: torch.Tensor, attention_mask: torch.Tensor):
"""Add prefix IDs for image generation."""
prefix_str = DEFAULT_IMAGE_AREA_TOKEN + hw2str(hw[0] // self.down_factor, hw[1] // self.down_factor)
prefix_output = self.tokenizer(prefix_str, truncation=False, add_special_tokens=True, return_tensors="pt")
prefix_input_ids = prefix_output.input_ids.to(input_ids.device, dtype=input_ids.dtype)
prefix_attention_mask = prefix_output.attention_mask.to(attention_mask.device, dtype=attention_mask.dtype)
# Remove bos token
if self.tokenizer.bos_token is not None:
prefix_input_ids = prefix_input_ids[:, 1:]
prefix_attention_mask = prefix_attention_mask[:, 1:]
# Add boi token
prefix_input_ids = torch.cat(
[
prefix_input_ids,
prefix_input_ids.new_tensor([self.boi]).unsqueeze(0),
],
dim=1,
)
prefix_attention_mask = torch.cat(
[
prefix_attention_mask,
prefix_attention_mask.new_ones((prefix_attention_mask.shape[0], 1)),
],
dim=1,
)
bsz = input_ids.shape[0]
input_ids = torch.cat([input_ids, prefix_input_ids.expand(bsz, -1)], dim=1)
attention_mask = torch.cat([attention_mask, prefix_attention_mask.expand(bsz, -1)], dim=1)
return input_ids, attention_mask
@torch.no_grad()
def decoding(
self,
c: torch.Tensor,
attention_mask: torch.Tensor,
past_key_values,
max_new_len: int,
num_images_per_caption: int,
use_norm: bool = False,
cfg: float = 1.0,
cfg_img: float = 1.0,
cfg_mult: int = 1,
cfg_schedule: Literal["linear", "constant"] = "constant",
timesteps_shift: float = 1.0,
num_sampling_steps: int = 20,
progress: bool = True,
hw: tuple[int, int] = (256, 256),
):
"""Autoregressive image token decoding with optional CFG-Parallel."""
if cfg_mult <= 0:
raise ValueError(f"Invalid cfg_mult={cfg_mult}; expected a positive value.")
full_bsz = c.shape[0]
if full_bsz % cfg_mult != 0:
raise ValueError(
f"Invalid CFG layout: batch size {full_bsz} is not divisible by active CFG multiplier {cfg_mult}."
)
# CFG-Parallel: each rank handles one portion of the CFG batch
cfg_world_size = get_classifier_free_guidance_world_size()
cfg_parallel = cfg_world_size > 1 and cfg_mult > 1
if cfg_parallel and cfg_world_size != cfg_mult:
logger.warning(
"CFG parallel world size (%d) does not match the number of active "
"CFG branches (%d); falling back to non-parallel CFG.",
cfg_world_size,
cfg_mult,
)
cfg_parallel = False
if cfg_parallel:
cfg_rank = get_classifier_free_guidance_rank()
cfg_group = get_cfg_group()
# Split batch: rank 0 gets positive, rank 1 gets first uncond, etc.
batch_per_rank = full_bsz // cfg_mult
start = cfg_rank * batch_per_rank
end = start + batch_per_rank
c = c[start:end]
attention_mask = attention_mask[start:end]
# Slice the StaticCache to keep only this rank's portion.
# We rebuild a new StaticCache for this rank's batch size and copy
# the relevant slices from the pre-filled cache.
if isinstance(past_key_values, StaticCache):
old_cache = past_key_values
new_cache = StaticCache(
config=self.config,
max_cache_len=old_cache.get_max_cache_shape(),
)
for layer_idx, layer in enumerate(old_cache.layers):
if not layer.is_initialized:
continue
# Force-initialize the new layer with sliced tensors
new_layer = new_cache.layers[layer_idx]
sliced_keys = layer.keys[start:end]
new_layer.lazy_initialization(sliced_keys)
new_layer.keys.copy_(sliced_keys)
new_layer.values.copy_(layer.values[start:end])
# Preserve cumulative_length for sliding window layers
# (StaticSlidingWindowLayer tracks seq position via this
# counter, not by counting non-zero entries)
if hasattr(layer, "cumulative_length"):
new_layer.cumulative_length = layer.cumulative_length
past_key_values = new_cache
else:
cfg_rank = 0
cfg_group = None
token_dim = self.config.latent_channels * self.config.latent_patch_size**2
indices = list(range(max_new_len))
indices = tqdm(indices, desc="Generating") if progress else indices
tokens = None
for step in indices:
# CFG schedule
if cfg_schedule == "linear":
tokens_len = 0 if tokens is None else tokens.shape[1]
cfg_iter = 1 + (cfg - 1) * (max_new_len - tokens_len) / max_new_len
cfg_img_iter = 1 + (cfg_img - 1) * (max_new_len - tokens_len) / max_new_len
elif cfg_schedule == "constant":
cfg_iter = cfg
cfg_img_iter = cfg_img
else:
raise NotImplementedError(f"Unknown cfg_schedule: {cfg_schedule}")
if cfg_parallel:
# Each rank projects its own portion
c_proj = self.model.image_out_projector(c)
# Gather projected context from all CFG ranks
c_gathered = cfg_group.all_gather(c_proj, separate_tensors=True)
c_full = torch.cat(c_gathered, dim=0) # (full_bsz, 1, hidden)
# Rank 0 runs the FM head sampling (it needs full CFG batch)
if cfg_rank == 0:
token_sampled = self.model.image_head.sample(
c=c_full.squeeze(1),
cfg=cfg_iter,
cfg_img=cfg_img_iter,
cfg_mult=cfg_mult,
timesteps_shift=timesteps_shift,
num_sampling_steps=num_sampling_steps,
noise_repeat=num_images_per_caption,
)
else:
token_sampled = torch.empty(
batch_per_rank,
token_dim,
device=c.device,
dtype=c.dtype,
)
# Broadcast sampled token from rank 0 to all ranks
token_sampled = token_sampled.contiguous()
cfg_group.broadcast(token_sampled, src=0)
else:
c_proj = self.model.image_out_projector(c)
token_sampled = self.model.image_head.sample(
c=c_proj.squeeze(1),
cfg=cfg_iter,
cfg_img=cfg_img_iter,
cfg_mult=cfg_mult,
timesteps_shift=timesteps_shift,
num_sampling_steps=num_sampling_steps,
noise_repeat=num_images_per_caption,
)
if use_norm:
token_sampled = layer_norm(token_sampled, normalized_shape=token_sampled.size()[1:])
if tokens is not None:
tokens = torch.cat([tokens, token_sampled.unsqueeze(1)], dim=1)
else:
tokens = token_sampled.unsqueeze(1)
# Prepare input embeds for next LLM step
cur_inputs_embeds = self.model.image_in_projector(tokens[:, -1:])
if not cfg_parallel and cfg_mult > 1:
# Non-parallel: duplicate embeds for all active CFG branches.
cur_inputs_embeds = torch.cat([cur_inputs_embeds] * cfg_mult, dim=0)
# In CFG-parallel mode, each rank already has its portion — no duplication needed
attention_mask = torch.cat(
[attention_mask, attention_mask.new_ones((attention_mask.shape[0], 1))],
dim=-1,
)
outputs = self.model.forward_model(
inputs_embeds=cur_inputs_embeds,
attention_mask=attention_mask,
past_key_values=past_key_values,
use_cache=True,
)
past_key_values = outputs.past_key_values
c = outputs.last_hidden_state[:, -1:]
if getattr(self.config, "use_gen_pos_embed", False):
c = c + self.model.gen_pos_embed_with_ar(hw[0], hw[1])[:, step + 1 : step + 2, :]
return tokens
@torch.no_grad()
def forward(
self,
req: OmniDiffusionRequest,
prompt: str | list[str] | None = None,
height: int | None = None,
width: int | None = None,
num_inference_steps: int = 28,
guidance_scale: float = 7.5,
negative_prompt: str | list[str] | None = None,
num_images_per_prompt: int = 1,
generator: torch.Generator | None = None,
seed: int | None = None,
**kwargs,
) -> DiffusionOutput:
"""
Generate images from text prompts.
Args:
req: OmniDiffusionRequest containing generation parameters
prompt: Text prompt(s) for generation
height: Output image height
width: Output image width
num_inference_steps: Number of sampling steps (default 28 for NextStep-1.1)
guidance_scale: CFG scale
negative_prompt: Negative prompt for CFG
num_images_per_prompt: Number of images per prompt
generator: Random generator for reproducibility
seed: Random seed
Returns:
DiffusionOutput containing generated images
"""
# Extract parameters from request
# req.prompts is a list of str or dict; req.sampling_params holds all generation params
first_prompt = req.prompts[0] if req.prompts else None
if first_prompt is not None:
if isinstance(first_prompt, str):
prompt = first_prompt
else:
prompt = first_prompt.get("prompt") or prompt
negative_prompt = first_prompt.get("negative_prompt", negative_prompt)
height = req.sampling_params.height or height or 512
width = req.sampling_params.width or width or 512
num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps
if req.sampling_params.guidance_scale_provided:
guidance_scale = req.sampling_params.guidance_scale
num_images_per_prompt = (
req.sampling_params.num_outputs_per_prompt
if req.sampling_params.num_outputs_per_prompt > 0
else num_images_per_prompt
)
seed = req.sampling_params.seed if req.sampling_params.seed is not None else seed
# NextStep-specific parameters from request extra
cfg_img = (
req.sampling_params.guidance_scale_2
if req.sampling_params.guidance_scale_2 is not None
else req.sampling_params.extra_args.get("cfg_img", 1.0)
)
timesteps_shift = req.sampling_params.extra_args.get("timesteps_shift", 1.0)
use_norm = req.sampling_params.extra_args.get("use_norm", False)
cfg_schedule = req.sampling_params.extra_args.get("cfg_schedule", "constant")
positive_prompt = req.sampling_params.extra_args.get("positive_prompt", None)
# Set seed for reproducibility (use generator if provided, else fall back to seed)
if generator is None and seed is not None:
set_seed(seed)
elif generator is not None:
torch.manual_seed(generator.initial_seed())
# Prepare hw tuple
hw = (height, width)
# Check and process inputs (no image inputs for t2i)
captions, images = self._check_input(prompt, None)
# Build captions with CFG
captions, images, cfg_mult, effective_cfg_img = self._build_captions(
captions,
images,
num_images_per_prompt,
positive_prompt,
negative_prompt,
guidance_scale,
cfg_img,
)
# No input images for text-to-image
latents = None
# Add BOS token to captions before tokenizing
captions = [
self.tokenizer.bos_token + caption if self.tokenizer.bos_token is not None else caption
for caption in captions
]
# Tokenize captions and add prefix ids
output = self.tokenizer(
captions,
padding="longest",
truncation=False,
add_special_tokens=True,
return_tensors="pt",
)
input_ids = output.input_ids.to(self.device)
attention_mask = output.attention_mask.to(self.device)
input_ids, attention_mask = self._add_prefix_ids(hw, input_ids, attention_mask)
# LLM prefill
max_new_len = (hw[0] // self.down_factor) * (hw[1] // self.down_factor)
max_cache_len = input_ids.shape[1] + max_new_len
past_key_values = StaticCache(
config=self.config,
max_cache_len=max_cache_len,
)
inputs_embeds = self.model.prepare_inputs_embeds(input_ids, latents)
outputs = self.model.forward_model(
inputs_embeds=inputs_embeds,
attention_mask=attention_mask,
past_key_values=past_key_values,
use_cache=True,
)
past_key_values = outputs.past_key_values
c = outputs.last_hidden_state[:, -1:]
if getattr(self.config, "use_gen_pos_embed", False):
c = c + self.model.gen_pos_embed_with_ar(height, width)[:, 0:1, :]
# Decoding
tokens = self.decoding(
c=c,
attention_mask=attention_mask,
past_key_values=past_key_values,
max_new_len=max_new_len,
num_images_per_caption=num_images_per_prompt,
use_norm=use_norm,
cfg=guidance_scale,
cfg_img=effective_cfg_img,
cfg_mult=cfg_mult,
cfg_schedule=cfg_schedule,
timesteps_shift=timesteps_shift,
num_sampling_steps=num_inference_steps,
progress=True,
hw=hw,
)
# Unpatchify
latents = self.model.unpatchify(tokens)
latents = (latents / self.scaling_factor) + self.shift_factor
# Decode latents
sampled_images = self.vae.decode(latents.to(self.vae.dtype)).sample
sampled_images = sampled_images.detach().cpu().to(torch.float32)
return DiffusionOutput(output=sampled_images)
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
"""Load model weights."""
loader = AutoWeightsLoader(self)
return loader.load_weights(weights)
vllm_omni/diffusion/registry.py
+12
View File
@@ -95,6 +95,11 @@ _DIFFUSION_MODELS = {
"pipeline_flux2_klein",
"Flux2KleinPipeline",
),
"NextStep11Pipeline": (
"nextstep_1_1",
"pipeline_nextstep_1_1",
"NextStep11Pipeline",
),
"FluxPipeline": (
"flux",
"pipeline_flux",
@@ -116,6 +121,12 @@ DiffusionModelRegistry = _ModelRegistry(
_VAE_PATCH_PARALLEL_ALLOWLIST = {
# Only enable for models we have validated end-to-end.
"ZImagePipeline",
"NextStep11Pipeline",
}
_NO_CACHE_ACCELERATION = {
# Pipelines that do not support cache acceleration (cache_dit / tea_cache).
"NextStep11Pipeline",
}
@@ -276,6 +287,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = {
"LongCatImageEditPipeline": "get_longcat_image_post_process_func",
"StableDiffusion3Pipeline": "get_sd3_image_post_process_func",
"Flux2KleinPipeline": "get_flux2_klein_post_process_func",
"NextStep11Pipeline": "get_nextstep11_post_process_func",
"FluxPipeline": "get_flux_post_process_func",
}
vllm_omni/diffusion/worker/diffusion_model_runner.py
+17 -2
View File
@@ -27,6 +27,7 @@ from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
from vllm_omni.diffusion.forward_context import set_forward_context
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
from vllm_omni.diffusion.offloader import get_offload_backend
from vllm_omni.diffusion.registry import _NO_CACHE_ACCELERATION
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.distributed.omni_connectors.kv_transfer_manager import OmniKVTransferManager
from vllm_omni.platforms import current_omni_platform
@@ -148,7 +149,16 @@ class DiffusionModelRunner:
self.cache_backend = get_cache_backend(self.od_config.cache_backend, self.od_config.cache_config)
if self.cache_backend is not None:
self.cache_backend.enable(self.pipeline)
if self.od_config.model_class_name in _NO_CACHE_ACCELERATION:
logger.warning(
"Cache backend '%s' is not supported for %s; disabling cache acceleration.",
self.od_config.cache_backend,
self.od_config.model_class_name,
)
self.cache_backend = None
self.od_config.cache_backend = None
else:
self.cache_backend.enable(self.pipeline)
logger.info("Model runner: Initialization complete.")
@@ -196,7 +206,12 @@ class DiffusionModelRunner:
output = self.pipeline.forward(req)
# NOTE:
if self.od_config.cache_backend == "cache_dit" and self.od_config.enable_cache_dit_summary:
if (
self.cache_backend is not None
and self.cache_backend.is_enabled()
and self.od_config.cache_backend == "cache_dit"
and self.od_config.enable_cache_dit_summary
):
cache_summary(self.pipeline, details=True)
return output
vllm_omni/distributed/omni_connectors/kv_transfer_manager.py
+82 -17
View File
@@ -15,6 +15,8 @@ from .utils.config import ConnectorSpec
logger = init_logger(__name__)
LayerKV = torch.Tensor | tuple[torch.Tensor, torch.Tensor]
@dataclass
class OmniKVCacheConfig:
@@ -157,7 +159,7 @@ class OmniKVTransferManager:
def handle_finished_requests_kv_transfer(
self,
finished_reqs: dict[str, dict[str, Any]],
kv_caches: list[torch.Tensor],
kv_caches: list[LayerKV],
block_size: int,
cache_dtype: str,
request_id_resolver: Callable[[str], str] | None = None,
@@ -169,7 +171,7 @@ class OmniKVTransferManager:
Args:
finished_reqs: Dict mapping request_id to {block_ids, seq_len}
kv_caches: List of KV cache tensors per layer
kv_caches: List of KV cache (tensor or tuple) per layer
block_size: Size of each cache block
cache_dtype: Data type of the cache
request_id_resolver: Optional function to resolve global request ID
@@ -219,7 +221,7 @@ class OmniKVTransferManager:
req_id: str,
block_ids: list[int],
seq_len: int,
kv_caches: list[torch.Tensor],
kv_caches: list[LayerKV],
block_size: int,
cache_dtype: str,
) -> KVCacheTransferData | None:
@@ -229,10 +231,13 @@ class OmniKVTransferManager:
req_id: Request identifier
block_ids: List of block IDs to extract
seq_len: Sequence length
kv_caches: List of KV cache tensors per layer
kv_caches: List of KV cache (tensor or tuple) per layer
block_size: Size of each cache block
cache_dtype: Data type of the cache
Note: If key/value block counts differ, extraction uses only the overlapping
block range. Extra key/value blocks are ignored, so returned KV may be partial.
Returns:
KVCacheTransferData if extraction successful, None otherwise
"""
@@ -240,25 +245,37 @@ class OmniKVTransferManager:
key_cache: list[torch.Tensor | None] = [None] * num_layers
value_cache: list[torch.Tensor | None] = [None] * num_layers
for layer_idx, kv_tensor in enumerate(kv_caches):
# Validate block IDs - shape: [2, num_blocks, block_size, n_heads, head_dim]
max_block = kv_tensor.shape[1] - 1
for layer_idx, layer_kv in enumerate(kv_caches):
kv_pair = self._normalize_layer_kv(layer_kv, req_id=req_id, layer_idx=layer_idx)
if kv_pair is None:
continue
key_blocks, value_blocks = kv_pair
if key_blocks.shape[0] != value_blocks.shape[0]:
logger.warning(
f"Layer {layer_idx} for request {req_id} has mismatched KV block counts: "
f"key={key_blocks.shape[0]}, value={value_blocks.shape[0]}; using shared range"
)
# Validate block IDs - shape: [num_blocks, block_size, n_heads, head_dim]
max_block = min(key_blocks.shape[0], value_blocks.shape[0]) - 1
valid_ids = [bid for bid in block_ids if 0 <= bid <= max_block]
if not valid_ids:
continue
# Extract and reshape: [2, n_blocks, block_size, n_heads, head_dim]
# -> [2, seq_len, n_heads, head_dim]
selected = kv_tensor[:, valid_ids] # [2, n_valid, block_size, n_heads, head_dim]
n_kv, n_blks, blk_sz, n_heads, d_head = selected.shape
flat = selected.reshape(n_kv, n_blks * blk_sz, n_heads, d_head)
if seq_len < flat.shape[1]:
flat = flat[:, :seq_len]
# Extract and reshape: [n_blocks, block_size, n_heads, head_dim]
# -> [seq_len, n_heads, head_dim]
selected_k = key_blocks[valid_ids]
selected_v = value_blocks[valid_ids]
flat_k = selected_k.flatten(0, 1)
flat_v = selected_v.flatten(0, 1)
if seq_len < flat_k.shape[0]:
flat_k = flat_k[:seq_len]
flat_v = flat_v[:seq_len]
# Move to CPU
flat_cpu = flat.detach().cpu().contiguous()
key_cache[layer_idx] = flat_cpu[0]
value_cache[layer_idx] = flat_cpu[1]
key_cache[layer_idx] = flat_k.detach().cpu().contiguous()
value_cache[layer_idx] = flat_v.detach().cpu().contiguous()
if not any(k is not None for k in key_cache):
return None
@@ -275,6 +292,54 @@ class OmniKVTransferManager:
},
)
def _normalize_layer_kv(
self,
layer_kv: LayerKV,
req_id: str,
layer_idx: int,
) -> tuple[torch.Tensor, torch.Tensor] | None:
"""Normalize one layer KV cache to a `(key_blocks, value_blocks)` tuple.
Args:
layer_kv: The raw KV cache (tensor or tuple) for the layer
req_id: Request ID for logging
layer_idx: Layer index for logging
Returns:
Tuple of (key_blocks, value_blocks) if valid, None otherwise
"""
if isinstance(layer_kv, torch.Tensor):
if layer_kv.ndim < 3 or layer_kv.shape[0] != 2:
logger.warning(
f"Layer {layer_idx} for request {req_id} has invalid stacked KV shape: "
f"expected [2, blocks, block_size, ...], got {tuple(layer_kv.shape)}"
)
return None
key_blocks = layer_kv[0]
value_blocks = layer_kv[1]
elif isinstance(layer_kv, tuple):
if len(layer_kv) != 2:
logger.warning(
f"Layer {layer_idx} for request {req_id} has KV pair length {len(layer_kv)} (expected 2)"
)
return None
key_blocks, value_blocks = layer_kv
if not isinstance(key_blocks, torch.Tensor) or not isinstance(value_blocks, torch.Tensor):
logger.warning(f"Layer {layer_idx} for request {req_id} has non-tensor KV pair entries")
return None
else:
logger.warning(f"Layer {layer_idx} for request {req_id} has unsupported KV type {type(layer_kv).__name__}")
return None
# ensure key/value blocks are at least 2D for block indexing
if key_blocks.ndim < 2 or value_blocks.ndim < 2:
logger.warning(
f"Layer {layer_idx} for request {req_id} has invalid KV block shape: "
f"got key={tuple(key_blocks.shape)} value={tuple(value_blocks.shape)}"
)
return None
return key_blocks, value_blocks
def _transfer_kv_cache(self, kv_data: KVCacheTransferData, transfer_req_id: str) -> None:
"""Transfer KV cache data to downstream stage via OmniConnector.
vllm_omni/entrypoints/async_omni_diffusion.py
+17 -5
View File
@@ -80,24 +80,36 @@ class AsyncOmniDiffusion:
if engine_input_source is not None:
self.od_config.omni_kv_config.setdefault("engine_input_source", engine_input_source)
# Diffusers-style models expose `model_index.json` with `_class_name`.
# Non-diffusers models (e.g. Bagel, NextStep) only have `config.json`,
# so we fall back to reading that and mapping model_type manually.
try:
config_dict = get_hf_file_to_dict("model_index.json", od_config.model)
od_config.model_class_name = config_dict.get("_class_name", None)
od_config.update_multimodal_support()
if config_dict is not None:
od_config.model_class_name = config_dict.get("_class_name", None)
od_config.update_multimodal_support()
tf_config_dict = get_hf_file_to_dict("transformer/config.json", od_config.model)
od_config.tf_model_config = TransformerConfig.from_dict(tf_config_dict)
except (AttributeError, OSError, ValueError):
tf_config_dict = get_hf_file_to_dict("transformer/config.json", od_config.model)
od_config.tf_model_config = TransformerConfig.from_dict(tf_config_dict)
else:
raise FileNotFoundError("model_index.json not found")
except (AttributeError, OSError, ValueError, FileNotFoundError):
cfg = get_hf_file_to_dict("config.json", od_config.model)
if cfg is None:
raise ValueError(f"Could not find config.json or model_index.json for model {od_config.model}")
model_type = cfg.get("model_type")
architectures = cfg.get("architectures") or []
# Bagel/NextStep models don't have a model_index.json, so we set the pipeline class name manually
if model_type == "bagel" or "BagelForConditionalGeneration" in architectures:
od_config.model_class_name = "BagelPipeline"
od_config.tf_model_config = TransformerConfig()
od_config.update_multimodal_support()
elif model_type == "nextstep":
if od_config.model_class_name is None:
od_config.model_class_name = "NextStep11Pipeline"
od_config.tf_model_config = TransformerConfig()
od_config.update_multimodal_support()
elif architectures and len(architectures) == 1:
od_config.model_class_name = architectures[0]
else:
vllm_omni/entrypoints/cli/serve.py
+8
View File
@@ -290,6 +290,14 @@ class OmniServeCommand(CLISubcommand):
type=int,
help="The max size of generate image (height * width).",
)
# TTS-specific parameters
omni_config_group.add_argument(
"--tts-max-instructions-length",
type=int,
default=None,
help="Maximum length for TTS voice style instructions (overrides stage config, default: 500).",
)
return serve_parser
vllm_omni/entrypoints/omni.py
+3
View File
@@ -257,6 +257,9 @@ class OmniBase:
stage_configs_path = kwargs.get("stage_configs_path", None)
# TTS-specific CLI overrides
self.tts_max_instructions_length: int | None = kwargs.get("tts_max_instructions_length", None)
# Load stage configurations from YAML
config_path, stage_configs = load_and_resolve_stage_configs(
model,
vllm_omni/entrypoints/omni_diffusion.py
+19 -9
View File
@@ -55,31 +55,41 @@ class OmniDiffusion:
self.od_config.omni_kv_config.setdefault("engine_input_source", engine_input_source)
# Diffusers-style models expose `model_index.json` with `_class_name`.
# Bagel models (and other non-diffusers) typically expose `config.json`.
# Non-diffusers models (e.g. Bagel, NextStep) only have `config.json`,
# so we fall back to reading that and mapping model_type manually.
try:
config_dict = get_hf_file_to_dict(
"model_index.json",
od_config.model,
)
od_config.model_class_name = config_dict.get("_class_name", None)
od_config.update_multimodal_support()
if config_dict is not None:
od_config.model_class_name = config_dict.get("_class_name", None)
od_config.update_multimodal_support()
tf_config_dict = get_hf_file_to_dict(
"transformer/config.json",
od_config.model,
)
od_config.tf_model_config = TransformerConfig.from_dict(tf_config_dict)
except (AttributeError, OSError, ValueError):
tf_config_dict = get_hf_file_to_dict(
"transformer/config.json",
od_config.model,
)
od_config.tf_model_config = TransformerConfig.from_dict(tf_config_dict)
else:
raise FileNotFoundError("model_index.json not found")
except (AttributeError, OSError, ValueError, FileNotFoundError):
cfg = get_hf_file_to_dict("config.json", od_config.model)
if cfg is None:
raise ValueError(f"Could not find config.json or model_index.json for model {od_config.model}")
model_type = cfg.get("model_type")
architectures = cfg.get("architectures") or []
# Bagel/NextStep models don't have a model_index.json, so we set the pipeline class name manually
if model_type == "bagel" or "BagelForConditionalGeneration" in architectures:
od_config.model_class_name = "BagelPipeline"
od_config.tf_model_config = TransformerConfig()
od_config.update_multimodal_support()
elif model_type == "nextstep":
if od_config.model_class_name is None:
od_config.model_class_name = "NextStep11Pipeline"
od_config.tf_model_config = TransformerConfig()
od_config.update_multimodal_support()
elif architectures and len(architectures) == 1:
od_config.model_class_name = architectures[0]
else:
vllm_omni/entrypoints/omni_stage.py
+1
View File
@@ -281,6 +281,7 @@ class OmniStage:
self.final_output = getattr(stage_config, "final_output", False)
self.final_output_type = getattr(stage_config, "final_output_type", None)
self.tts_args = _to_dict(getattr(stage_config, "tts_args", {}))
default_sampling_params = getattr(stage_config, "default_sampling_params", {})
# For LLM stage, this can directly be a SamplingParams-compatible dict;
# For diffusion stage, this only serves as default values for diffusion kwargs.
vllm_omni/entrypoints/openai/serving_chat.py
+7 -5
View File
@@ -1698,12 +1698,14 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
# OMNI: Access multimodal_output from CompletionOutput (outputs[0]), not from RequestOutput
# Reference: examples/offline_inference/qwen3_omni/end2end.py line 421
audio_data = final_res.outputs[0].multimodal_output.get("audio")
if stream:
audio_tensor = audio_data[-1].float().detach().cpu().numpy()
if isinstance(audio_data, list):
if stream:
audio_tensor = audio_data[-1]
else:
audio_tensor = torch.cat(audio_data, dim=-1)
else:
if isinstance(audio_data, list):
audio_data = torch.cat(audio_data, dim=-1)
audio_tensor = audio_data.float().detach().cpu().numpy()
audio_tensor = audio_data
audio_tensor = audio_tensor.float().detach().cpu().numpy()
# Ensure audio is 1D (flatten if needed)
if audio_tensor.ndim > 1:
vllm_omni/entrypoints/openai/serving_speech.py
+41 -4
View File
@@ -61,11 +61,48 @@ _TTS_MAX_NEW_TOKENS_MAX = 4096
class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
# Find and cache the TTS stage (if any) during initialization
self._tts_stage = self._find_tts_stage()
self._is_tts = self._tts_stage is not None
# Cache TTS configuration values (computed once, reused per request)
self._max_instructions_length = self._compute_max_instructions_length()
# Load supported speakers
self.supported_speakers = self._load_supported_speakers()
logger.info(f"Loaded {len(self.supported_speakers)} supported speakers: {sorted(self.supported_speakers)}")
self._tts_tokenizer = None
def _find_tts_stage(self):
"""Find and return the TTS stage from the stage list, or None if not found."""
stage_list = getattr(self.engine_client, "stage_list", None)
if stage_list is None:
return None
for stage in stage_list:
if getattr(stage, "model_stage", None) in _TTS_MODEL_STAGES:
return stage
return None
def _compute_max_instructions_length(self) -> int:
"""Compute max instructions length with precedence: CLI > stage config > default.
Called once during initialization; result is cached in self._max_instructions_length.
"""
# 1. CLI override takes highest priority (stored in engine_client)
cli_override = getattr(self.engine_client, "tts_max_instructions_length", None)
if cli_override is not None:
return cli_override
# 2. Try to get from TTS stage config
if self._tts_stage is not None:
tts_args = getattr(self._tts_stage, "tts_args", {})
if "max_instructions_length" in tts_args:
return tts_args["max_instructions_length"]
# 3. Default fallback
return _TTS_MAX_INSTRUCTIONS_LENGTH
def _load_supported_speakers(self) -> set[str]:
"""Load supported speakers (case-insensitive) from the model configuration."""
try:
@@ -164,9 +201,9 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
if task_type == "VoiceDesign" and not request.instructions:
return "VoiceDesign task requires 'instructions' to describe the voice"
# Validate instructions length
if request.instructions and len(request.instructions) > _TTS_MAX_INSTRUCTIONS_LENGTH:
return f"Instructions too long (max {_TTS_MAX_INSTRUCTIONS_LENGTH} characters)"
# Validate instructions length (using cached value from initialization)
if request.instructions and len(request.instructions) > self._max_instructions_length:
return f"Instructions too long (max {self._max_instructions_length} characters)"
# Validate max_new_tokens range
if request.max_new_tokens is not None:
@@ -300,7 +337,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
request_id = f"speech-{random_uuid()}"
try:
if self._is_tts_model():
if self._is_tts:
# Validate TTS parameters
validation_error = self._validate_tts_request(request)
if validation_error:
vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py
+4 -2
View File
@@ -17,8 +17,6 @@ from vllm.logger import init_logger
from vllm.model_executor.models.interfaces import SupportsMRoPE, SupportsMultiModal, SupportsPP
from vllm.model_executor.models.qwen2_5_omni_thinker import (
Qwen2_5OmniConditionalGenerationMixin,
Qwen2_5OmniThinkerDummyInputsBuilder,
Qwen2_5OmniThinkerMultiModalProcessor,
Qwen2_5OmniThinkerProcessingInfo,
)
from vllm.model_executor.models.utils import init_vllm_registered_model, maybe_prefix
@@ -37,6 +35,10 @@ from vllm.v1.sample.sampler import Sampler
from vllm_omni.model_executor.custom_process_mixin import CustomProcessMixin
from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific
from vllm_omni.model_executor.models.output_templates import OmniOutput
from vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker import (
Qwen2_5OmniThinkerDummyInputsBuilder,
Qwen2_5OmniThinkerMultiModalProcessor,
)
from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights, split_list_into_ranges
from vllm_omni.platforms import current_omni_platform
vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_talker.py
+5 -3
View File
@@ -11,8 +11,6 @@ from vllm.config import VllmConfig
from vllm.logger import init_logger
from vllm.model_executor.models.interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP
from vllm.model_executor.models.qwen2_5_omni_thinker import (
Qwen2_5OmniThinkerDummyInputsBuilder,
Qwen2_5OmniThinkerMultiModalProcessor,
Qwen2_5OmniThinkerProcessingInfo,
)
from vllm.model_executor.models.qwen2_5_vl import Qwen2_5_VisionTransformer
@@ -28,7 +26,11 @@ from vllm.v1.outputs import SamplerOutput
from vllm.v1.sample.metadata import SamplingMetadata
from vllm.v1.sample.sampler import Sampler
from vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker import Qwen2_5OmniConditionalGenerationMixin
from vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker import (
Qwen2_5OmniConditionalGenerationMixin,
Qwen2_5OmniThinkerDummyInputsBuilder,
Qwen2_5OmniThinkerMultiModalProcessor,
)
@MULTIMODAL_REGISTRY.register_processor(
vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py
+73 -2
View File
@@ -1,6 +1,6 @@
"""Thin Omni wrapper: reuse upstream Qwen2.5-Omni thinker (v0.14) with minimal overrides."""
from collections.abc import Iterable
from collections.abc import Iterable, Mapping
from typing import Any
import torch
@@ -25,7 +25,6 @@ from vllm.model_executor.models.module_mapping import MultiModelKeys
from vllm.model_executor.models.qwen2_5_omni_thinker import (
Qwen2_5OmniAudioFeatureInputs,
Qwen2_5OmniThinkerDummyInputsBuilder,
Qwen2_5OmniThinkerMultiModalProcessor,
Qwen2_5OmniThinkerProcessingInfo,
check_interleaved_audio_video,
merge_interleaved_embeddings,
@@ -33,6 +32,9 @@ from vllm.model_executor.models.qwen2_5_omni_thinker import (
from vllm.model_executor.models.qwen2_5_omni_thinker import (
Qwen2_5OmniConditionalGenerationMixin as Qwen2_5OmniConditionalGenerationMixinBase,
)
from vllm.model_executor.models.qwen2_5_omni_thinker import (
Qwen2_5OmniThinkerMultiModalProcessor as _Qwen2_5OmniThinkerMultiModalProcessorBase,
)
from vllm.model_executor.models.qwen2_5_vl import (
Qwen2_5_VisionTransformer,
Qwen2_5_VLImageEmbeddingInputs,
@@ -54,6 +56,12 @@ from vllm.model_executor.models.vision import get_llm_pos_ids_for_vision
from vllm.multimodal import MULTIMODAL_REGISTRY
from vllm.multimodal.inputs import (
MultiModalFeatureSpec,
MultiModalKwargsItems,
)
from vllm.multimodal.parse import MultiModalDataItems
from vllm.multimodal.processing.processor import (
MultiModalPromptUpdates,
PlaceholderFeaturesInfo,
)
from vllm.sequence import IntermediateTensors
@@ -64,6 +72,69 @@ except (ImportError, ModuleNotFoundError):
logger = init_logger(__name__)
class Qwen2_5OmniThinkerMultiModalProcessor(
_Qwen2_5OmniThinkerMultiModalProcessorBase,
):
"""Override to fix use_audio_in_video detection when mm cache returns None."""
def _maybe_apply_prompt_updates(
self,
mm_items: MultiModalDataItems,
prompt_ids: list[int],
mm_kwargs: MultiModalKwargsItems,
mm_prompt_updates: MultiModalPromptUpdates,
is_update_applied: bool,
) -> tuple[list[int], Mapping[str, list[PlaceholderFeaturesInfo]]]:
mm_item_counts = mm_items.get_all_counts()
self._validate_mm_kwargs(mm_kwargs, mm_item_counts)
self._validate_mm_updates(mm_prompt_updates, mm_item_counts)
use_audio_in_video = False
if "video" in mm_kwargs:
non_none_items = [item for item in mm_kwargs["video"] if item is not None]
if non_none_items:
for item in non_none_items:
if item.get("use_audio_in_video"):
uaiv_tensor = item["use_audio_in_video"].data
if uaiv_tensor.numel() > 0:
use_audio_in_video = bool(uaiv_tensor.item())
break
elif "audio" in mm_prompt_updates:
tokenizer = self.info.get_tokenizer()
audio_pad_id = tokenizer.convert_tokens_to_ids("<|audio_pad|>")
use_audio_in_video = audio_pad_id not in prompt_ids
if is_update_applied:
mm_placeholders = self._find_mm_placeholders(
prompt_ids,
mm_prompt_updates,
)
self._validate_mm_placeholders(
mm_placeholders,
mm_item_counts,
)
else:
if use_audio_in_video and "audio" in mm_prompt_updates:
filtered_updates = {k: v for k, v in mm_prompt_updates.items() if k != "audio"}
prompt_ids, mm_placeholders = self._apply_prompt_updates(
prompt_ids,
filtered_updates,
)
mm_placeholders = self._derive_audio_from_video_placeholders(mm_placeholders, mm_prompt_updates)
else:
prompt_ids, mm_placeholders = self._apply_prompt_updates(
prompt_ids,
mm_prompt_updates,
)
self._validate_mm_placeholders(
mm_placeholders,
mm_item_counts,
)
return prompt_ids, mm_placeholders
class Qwen2_5OmniConditionalGenerationMixin(Qwen2_5OmniConditionalGenerationMixinBase):
def _parse_and_validate_audio_input(self, **kwargs: object) -> Qwen2_5OmniAudioFeatureInputs | None:
input_audio_features = kwargs.pop("input_audio_features", None)
vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py
+7 -36
View File
@@ -20,9 +20,6 @@ from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding
from vllm.model_executor.models.interfaces import SupportsMRoPE, SupportsMultiModal, SupportsPP
from vllm.model_executor.models.qwen3_omni_moe_thinker import (
Qwen3OmniMoeConditionalGenerationMixin,
Qwen3OmniMoeThinkerDummyInputsBuilder,
Qwen3OmniMoeThinkerMultiModalProcessor,
Qwen3OmniMoeThinkerProcessingInfo,
)
from vllm.model_executor.models.utils import init_vllm_registered_model, maybe_prefix
from vllm.multimodal import MULTIMODAL_REGISTRY
@@ -34,8 +31,12 @@ from vllm.v1.sample.sampler import Sampler
from vllm_omni.model_executor.custom_process_mixin import CustomProcessMixin
from vllm_omni.model_executor.models.output_templates import OmniOutput
from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker import (
Qwen3OmniMoeThinkerDummyInputsBuilder,
Qwen3OmniMoeThinkerMultiModalProcessor,
Qwen3OmniMoeThinkerProcessingInfo,
)
from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights, safe_tensor_reshape
from vllm_omni.platforms import current_omni_platform
# Special token IDs for Qwen3 Omni MoE
# Reference: https://huggingface.co/Qwen/Qwen3-Omni-30B-A3B-Instruct/blob/main/tokenizer_config.json
@@ -269,27 +270,6 @@ class Qwen3OmniMoeForConditionalGeneration(
# ========== Stage 1: Thinker ==========
if self.model_stage == "thinker":
thinker_dev = self._module_device(self.thinker)
if current_omni_platform.is_npu():
# Normalize to batched inputs if needed
_added_batch_dim = False
if input_ids is not None and input_ids.ndim == 1:
input_ids = input_ids.unsqueeze(0)
_added_batch_dim = True
if positions is not None and positions.ndim == 1:
positions = positions.unsqueeze(0)
_added_batch_dim = True
if inputs_embeds is not None and inputs_embeds.ndim == 2:
inputs_embeds = inputs_embeds.unsqueeze(0)
_added_batch_dim = True
# Handle None input_ids
if input_ids is None:
input_ids = torch.zeros(
inputs_embeds.shape[1],
dtype=torch.long,
device=thinker_dev,
).unsqueeze(0)
_added_batch_dim = True
# Move to thinker device
if input_ids is not None and input_ids.device != thinker_dev:
@@ -308,22 +288,13 @@ class Qwen3OmniMoeForConditionalGeneration(
"capture_layer_indices": [0, int(accept_layer)],
"return_hidden_states": True,
}
if current_omni_platform.is_npu():
# TODO: remove this hack when NPU supports batched inputs properly
thinker_input_ids = input_ids[0] if input_ids is not None and _added_batch_dim else input_ids
thinker_inputs_embeds = (
inputs_embeds[0] if inputs_embeds is not None and _added_batch_dim else inputs_embeds
)
else:
thinker_input_ids = input_ids
thinker_inputs_embeds = inputs_embeds
# Run thinker
text_hidden_states, captured_layer_dict = self.thinker(
input_ids=thinker_input_ids,
input_ids=input_ids,
positions=positions,
intermediate_tensors=intermediate_tensors,
inputs_embeds=thinker_inputs_embeds,
inputs_embeds=inputs_embeds,
**capture_kwargs,
**kwargs,
)
vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py
+1 -1
View File
@@ -99,8 +99,8 @@ class Qwen3OmniMoeTalkerForConditionalGeneration(
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
super().__init__()
talker_config: Qwen3OmniMoeTalkerConfig = vllm_config.model_config.hf_config
# rope_scaling already contains rope_theta, just assign it to rope_parameters
talker_config.text_config.rope_parameters = talker_config.text_config.rope_scaling
talker_config.text_config.rope_parameters["rope_theta"] = talker_config.text_config.rope_theta
self.quant_config = vllm_config.quant_config
self.prefix = prefix
self.vllm_config = vllm_config
vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py
+11 -5
View File
@@ -343,11 +343,17 @@ class Qwen3OmniMoeThinkerMultiModalProcessor(
use_audio_in_video = False
if "video" in mm_kwargs:
for item in mm_kwargs["video"]:
if item and item["use_audio_in_video"].data:
use_audio_in_video = True
else:
use_audio_in_video = False
non_none_items = [item for item in mm_kwargs["video"] if item is not None]
if non_none_items:
# Normal case: at least one non-cached item, read flag directly
use_audio_in_video = any(item["use_audio_in_video"].data for item in non_none_items)
elif "audio" in mm_prompt_updates:
# All video items are from cache (None); infer from prompt:
# use_audio_in_video=True means the prompt has no <|audio_pad|>
# placeholder (audio is embedded in video tokens instead)
tokenizer = self.info.get_tokenizer()
audio_pad_id = tokenizer.convert_tokens_to_ids("<|audio_pad|>")
use_audio_in_video = audio_pad_id not in prompt_ids
# normal case with `use_audio_in_video=False`
if is_update_applied:
vllm_omni/model_executor/models/qwen3_tts/qwen3_tts.py
+15 -3
View File
@@ -38,6 +38,18 @@ from .processing_qwen3_tts import Qwen3TTSProcessor
logger = init_logger(__name__)
_TASK_TYPE_CANONICAL: dict[str, str] = {
"customvoice": "CustomVoice",
"voicedesign": "VoiceDesign",
"base": "Base",
}
def _normalize_task_type(raw: str) -> str:
"""Normalize task type string to its canonical PascalCase form."""
return _TASK_TYPE_CANONICAL.get(raw.lower(), raw)
AudioLike = (
str # wav path, URL, base64
| np.ndarray # waveform (requires sr)
@@ -81,7 +93,7 @@ class Qwen3TTSModelForGeneration(nn.Module):
torch_dtype=torch.bfloat16,
**attn_kwargs,
)
self.task_type = model_path.split("-")[-1].split("/")[0]
self.task_type = _normalize_task_type(model_path.split("-")[-1].split("/")[0])
# Mark that this model produces multimodal outputs
self.have_multimodal_outputs = True
@@ -116,8 +128,8 @@ class Qwen3TTSModelForGeneration(nn.Module):
if isinstance(runtime_additional_information, list) and len(runtime_additional_information) > 0:
runtime_additional_information = runtime_additional_information[0]
text = runtime_additional_information.pop("text", [""])[0]
# Extract task_type from kwargs, default to "instruct"
task_type = runtime_additional_information.pop("task_type", [self.task_type])[0]
# Extract task_type from kwargs, default to self.task_type
task_type = _normalize_task_type(runtime_additional_information.pop("task_type", [self.task_type])[0])
speaker = runtime_additional_information.pop("speaker", ["uncle_fu"])[0]
language = runtime_additional_information.pop("language", ["Auto"])[0]
instruct = runtime_additional_information.pop("instruct", [""])[0]
vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/configuration_qwen3_tts_tokenizer_v1.py
-4
View File
@@ -214,8 +214,6 @@ class Qwen3TTSTokenizerV1EncoderConfig(PretrainedConfig):
output_dim (`int`, *optional*, defaults to 3584):
Output feature dimension produced by the encoder head (before/after projection, implementation-dependent).
grad_checkpointing (`bool`, *optional*, defaults to `False`):
Whether to enable gradient checkpointing to reduce memory usage during training.
enable_mp (`bool`, *optional*, defaults to `False`):
Whether to enable model parallel features (implementation-dependent).
audio_sequence_parallel (`bool`, *optional*, defaults to `False`):
@@ -246,7 +244,6 @@ class Qwen3TTSTokenizerV1EncoderConfig(PretrainedConfig):
n_layer=32,
n_window=100,
output_dim=3584,
grad_checkpointing=False,
enable_mp=False,
audio_sequence_parallel=False,
audio_vq_type="GRVQ",
@@ -265,7 +262,6 @@ class Qwen3TTSTokenizerV1EncoderConfig(PretrainedConfig):
self.n_layer = n_layer
self.n_window = n_window
self.output_dim = output_dim
self.grad_checkpointing = grad_checkpointing
self.enable_mp = enable_mp
self.audio_sequence_parallel = audio_sequence_parallel
self.audio_vq_type = audio_vq_type
vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/modeling_qwen3_tts_tokenizer_v1.py
-2
View File
@@ -1297,8 +1297,6 @@ class Qwen3TTSTokenizerV1Encoder(Qwen3TTSTokenizerV1EncoderPreTrainedModel):
n_layer=config.n_layer,
n_window=config.n_window,
output_dim=config.output_dim,
grad_checkpointing=config.grad_checkpointing,
enable_mp=config.enable_mp,
audio_sequence_parallel=config.audio_sequence_parallel,
audio_vq_type=config.audio_vq_type,
audio_vq_layers=config.audio_vq_layers,
vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/vq/speech_vq.py
-4
View File
@@ -196,8 +196,6 @@ class WhisperEncoderVQ(WhisperEncoder):
n_layer: int,
n_window: int = 1500,
output_dim: int = 512,
grad_checkpointing: bool = False,
enable_mp: bool = False,
audio_sequence_parallel: bool = False,
audio_vq_layers: int = -1,
audio_vq_type: str = "NULL",
@@ -219,8 +217,6 @@ class WhisperEncoderVQ(WhisperEncoder):
n_layer,
n_window,
output_dim,
grad_checkpointing,
enable_mp,
audio_sequence_parallel,
)
vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/vq/whisper_encoder.py
+7 -46
View File
@@ -23,19 +23,7 @@ import torch
import torch.nn.functional as F
from torch import Tensor, nn
try:
from flash_attn.flash_attn_interface import flash_attn_varlen_func as flash_attn_varlen_func
except ImportError:
try:
from flash_attn.flash_attn_interface import flash_attn_unpadded_func as flash_attn_varlen_func
except ImportError:
print(
"\n********\nWarning: flash-attn is not installed. "
"Will only run the manual PyTorch version. "
"Please install flash-attn for faster inference.\n********\n "
)
flash_attn_varlen_func = None
from vllm_omni.diffusion.attention.backends.utils.fa import HAS_FLASH_ATTN, flash_attn_varlen_func
N_FFT = 400
HOP_LENGTH = 160
@@ -152,7 +140,7 @@ class Linear(nn.Linear):
class MultiHeadAttention(nn.Module):
def __init__(self, n_state: int, n_head: int):
def __init__(self, n_state: int, n_head: int, use_flash_attention: bool = True):
super().__init__()
self.n_head = n_head
self.query = Linear(n_state, n_state)
@@ -160,7 +148,7 @@ class MultiHeadAttention(nn.Module):
self.value = Linear(n_state, n_state)
self.out = Linear(n_state, n_state)
self.use_flash_attention = True
self.use_flash_attention = use_flash_attention and HAS_FLASH_ATTN
def forward(
self,
@@ -171,15 +159,8 @@ class MultiHeadAttention(nn.Module):
k = self.key(x)
v = self.value(x)
if self.use_flash_attention:
if flash_attn_varlen_func is None:
x = self.qkv_attention_manual(q, k, v, cu_seqlens=cu_seqlens)
else:
if q.dtype not in [torch.float16, torch.bfloat16]:
x = self.qkv_attention_manual(q, k, v, cu_seqlens=cu_seqlens)
self.use_flash_attention = False
else:
x = self.qkv_flash_attention(q, k, v, cu_seqlens=cu_seqlens)
if self.use_flash_attention and q.dtype in [torch.float16, torch.bfloat16]:
x = self.qkv_flash_attention(q, k, v, cu_seqlens=cu_seqlens)
else:
x = self.qkv_attention_manual(q, k, v, cu_seqlens=cu_seqlens)
@@ -249,7 +230,7 @@ class MultiHeadAttention(nn.Module):
class ResidualAttentionBlock(nn.Module):
def __init__(self, n_state: int, n_head: int, enable_mp: bool = False, sequence_parallel: bool = False):
def __init__(self, n_state: int, n_head: int, sequence_parallel: bool = False):
super().__init__()
n_mlp = n_state * 4
self.attn_ln = nn.LayerNorm(n_state)
@@ -274,8 +255,6 @@ class WhisperEncoder(nn.Module):
n_layer: int,
n_window: int = 1500,
output_dim: int = 512,
grad_checkpointing: bool = False,
enable_mp: bool = False,
audio_sequence_parallel: bool = False,
):
super().__init__()
@@ -286,10 +265,7 @@ class WhisperEncoder(nn.Module):
self.n_mels = n_mels
self.blocks = nn.ModuleList(
[
ResidualAttentionBlock(n_state, n_head, enable_mp=enable_mp, sequence_parallel=audio_sequence_parallel)
for _ in range(n_layer)
]
[ResidualAttentionBlock(n_state, n_head, sequence_parallel=audio_sequence_parallel) for _ in range(n_layer)]
)
self.ln_post = nn.LayerNorm(n_state)
self.avg_pooler = nn.AvgPool1d(2, stride=2)
@@ -299,8 +275,6 @@ class WhisperEncoder(nn.Module):
self.audio_bos_eos_token = nn.Embedding(2, output_dim)
self.output_dim = output_dim
self.grad_checkpointing = grad_checkpointing
self.enable_mp = enable_mp
self.n_head = n_head
self.n_state = n_state
self.n_window = n_window
@@ -309,13 +283,6 @@ class WhisperEncoder(nn.Module):
self.tp_world_size = 1
self.set_audio_sync()
def set_audio_sync(self):
for name, param in self.named_parameters():
if not name.startswith("blocks"):
setattr(param, "audio_sync", True)
def forward(
self, x_list: list[Tensor], audio_mellens: list[int], audio_aftercnnlens: list[int], audio_seqlens: list[int]
):
@@ -377,9 +344,3 @@ class WhisperEncoder(nn.Module):
output[end_ids] = self.audio_bos_eos_token.weight[1].to(x.dtype)
output[audio_tokens_mask] = x
return output
def lock(self, layers: int):
self.conv1.requires_grad_(False)
self.conv2.requires_grad_(False)
for i in range(min(layers, len(self.blocks))):
self.blocks[i].requires_grad_(False)
vllm_omni/model_executor/stage_configs/qwen3_tts.yaml
+2
View File
@@ -65,6 +65,8 @@ stage_args:
# Distributed connector configuration
input_connectors:
from_stage_0: connector_of_shared_memory
tts_args:
max_instructions_length: 500
default_sampling_params:
temperature: 0.0
top_p: 1.0
vllm_omni/platforms/npu/stage_configs/qwen3_omni_moe.yaml
+6 -6
View File
@@ -7,7 +7,7 @@
stage_args:
- stage_id: 0
runtime:
devices: "0,1,2,3"
devices: "0,1"
max_batch_size: 1
engine_args:
model_stage: thinker
@@ -21,7 +21,7 @@ stage_args:
distributed_executor_backend: "mp"
enable_prefix_caching: false
hf_config_name: thinker_config
tensor_parallel_size: 4
tensor_parallel_size: 2
final_output: true
final_output_type: text
is_comprehension: true
@@ -36,14 +36,14 @@ stage_args:
- stage_id: 1
runtime:
devices: "4"
devices: "2"
max_batch_size: 1
engine_args:
model_stage: talker
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
gpu_memory_utilization: 0.2
gpu_memory_utilization: 0.6
enforce_eager: true # haven't supported talker ACL graph on NPU
trust_remote_code: true
engine_output_type: latent # Output codec codes for code2wav
@@ -66,7 +66,7 @@ stage_args:
- stage_id: 2
runtime:
devices: "0"
devices: "2"
max_batch_size: 1
engine_args:
model_stage: code2wav
@@ -78,7 +78,7 @@ stage_args:
async_scheduling: false
enable_prefix_caching: false
engine_output_type: audio # Final output: audio waveform
gpu_memory_utilization: 0.1
gpu_memory_utilization: 0.3
distributed_executor_backend: "mp"
max_num_batched_tokens: 1000000
hf_config_name: thinker_config
vllm_omni/platforms/npu/stage_configs/qwen3_tts.yaml
+2
View File
@@ -59,6 +59,8 @@ stage_args:
final_output_type: audio
input_connectors:
from_stage_0: connector_of_shared_memory
tts_args:
max_instructions_length: 500
default_sampling_params:
temperature: 0.0
top_p: 1.0
vllm_omni/platforms/npu/worker/npu_ar_model_runner.py
+254 -180
View File
@@ -3,16 +3,16 @@
from __future__ import annotations
from copy import copy
from copy import copy, deepcopy
from typing import Any, NamedTuple
import numpy as np
import torch
from vllm.compilation.cuda_graph import CUDAGraphStat
from vllm.config import CUDAGraphMode
from vllm.distributed.ec_transfer import get_ec_transfer, has_ec_transfer
from vllm.distributed.kv_transfer import get_kv_transfer_group, has_kv_transfer_group
from vllm.distributed.parallel_state import get_pp_group, get_tp_group
from vllm.forward_context import get_forward_context
from vllm.logger import logger
from vllm.sequence import IntermediateTensors
from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
@@ -24,16 +24,16 @@ from vllm.v1.outputs import (
)
from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
from vllm.v1.structured_output.utils import apply_grammar_bitmask
from vllm.v1.worker.gpu_model_runner import AsyncGPUModelRunnerOutput
from vllm.v1.worker.kv_connector_model_runner_mixin import KVConnectorOutput
from vllm.v1.utils import record_function_or_nullcontext
from vllm.v1.worker.gpu_model_runner import AsyncGPUModelRunnerOutput, PerLayerAttnMetadata
from vllm.v1.worker.ubatch_utils import maybe_create_ubatch_slices
from vllm_ascend.ascend_forward_context import set_ascend_forward_context
from vllm_ascend.attention.attention_v1 import AscendAttentionState
from vllm_ascend.attention.utils import AscendCommonAttentionMetadata
# yapf conflicts with isort for this block
# yapf: disable
from vllm_ascend.compilation.acl_graph import update_full_graph_params
from vllm_ascend.ops.rotary_embedding import update_cos_sin
from vllm_ascend.utils import ProfileExecuteDuration
from vllm_ascend.utils import enable_sp
from vllm_omni.distributed.omni_connectors.kv_transfer_manager import OmniKVTransferManager
from vllm_omni.outputs import OmniModelRunnerOutput
@@ -47,13 +47,14 @@ class ExecuteModelState(NamedTuple):
scheduler_output: SchedulerOutput
logits: torch.Tensor
spec_decode_metadata: SpecDecodeMetadata | None
spec_decode_common_attn_metadata: AscendCommonAttentionMetadata | None
hidden_states: torch.Tensor
sample_hidden_states: torch.Tensor
aux_hidden_states: list[torch.Tensor] | None
kv_connector_output: KVConnectorOutput | None
attn_metadata: dict[str, Any]
attn_metadata: PerLayerAttnMetadata
positions: torch.Tensor
ec_connector_output: ECConnectorOutput | None
cudagraph_stats: CUDAGraphStat | None
multimodal_outputs: Any
class NPUARModelRunner(OmniNPUModelRunner):
@@ -86,10 +87,9 @@ class NPUARModelRunner(OmniNPUModelRunner):
self,
scheduler_output: SchedulerOutput,
intermediate_tensors: IntermediateTensors | None = None,
) -> OmniModelRunnerOutput | AsyncModelRunnerOutput | IntermediateTensors | None:
) -> OmniModelRunnerOutput | IntermediateTensors | None:
if self.execute_model_state is not None:
raise RuntimeError("State error: sample_tokens() must be called "
"after execute_model() returns None.")
raise RuntimeError("State error: sample_tokens() must be called after execute_model() returns None.")
# -------------------------------------- Omni-new -------------------------------------------------
# [Omni] Handle KV transfer BEFORE updating states (which removes finished requests)
@@ -101,110 +101,202 @@ class NPUARModelRunner(OmniNPUModelRunner):
request_id_resolver=self._resolve_global_request_id,
)
# -------------------------------------- Omni-new -------------------------------------------------
# self._draft_token_ids is None when `input_fits_in_drafter=False`
# and there is no draft tokens scheduled. so it need to update the
# spec_decoding info in scheduler_output with async_scheduling.
# use deepcopy to avoid the modification has influence on the
# scheduler_output in engine core process.
# TODO(Ronald1995): deepcopy is expensive when there is a large
# number of requests, optimize it later.
if (
self.use_async_scheduling and self.num_spec_tokens and self._draft_token_ids is None # type: ignore[has-type]
):
scheduler_output = deepcopy(scheduler_output)
num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
with record_function_or_nullcontext("prepare input"):
with self.synchronize_input_prep():
# -------------------------------------- Omni-new -------------------------------------------------
self._update_states(scheduler_output)
# ------------------------------------------------------------------------------------------------
# Update persistent batch states.
with ProfileExecuteDuration().capture_async("prepare input"):
# -------------------------------------- Omni-new -------------------------------------------------
self._update_states(scheduler_output)
self._decode_and_store_request_payloads(scheduler_output)
# ------------------------------------------------------------------------------------------------
if has_ec_transfer() and get_ec_transfer().is_producer:
with self.maybe_get_ec_connector_output(
if has_ec_transfer() and get_ec_transfer().is_producer:
with self.maybe_get_ec_connector_output(
scheduler_output,
encoder_cache=self.encoder_cache,
) as ec_connector_output:
self._execute_mm_encoder(scheduler_output)
return make_empty_encoder_model_runner_output(
scheduler_output)
) as ec_connector_output:
self._execute_mm_encoder(scheduler_output)
return make_empty_encoder_model_runner_output(scheduler_output)
if not scheduler_output.total_num_scheduled_tokens:
if not has_kv_transfer_group():
logger.debug(
"skip this step for we receive the data from remote disaggregate prefill node"
if not num_scheduled_tokens:
if (
self.parallel_config.distributed_executor_backend == "external_launcher"
and self.parallel_config.data_parallel_size > 1
):
# this is a corner case when both external launcher
# and DP are enabled, num_scheduled_tokens could be
# 0, and has_unfinished_requests in the outer loop
# returns True. before returning early here we call
# dummy run to ensure coordinate_batch_across_dp
# is called into to avoid out of sync issues.
self._dummy_run(1)
if not has_kv_transfer_group():
# Return empty ModelRunnerOutput if no work to do.
return EMPTY_MODEL_RUNNER_OUTPUT
return self.kv_connector_no_forward(scheduler_output, self.vllm_config)
if self.cache_config.kv_sharing_fast_prefill:
assert not self.num_prompt_logprobs, (
"--kv-sharing-fast-prefill produces incorrect "
"logprobs for prompt tokens, tokens, please disable "
"it when the requests need prompt logprobs"
)
# Return empty ModelRunnerOutput if there's no work to do.
return EMPTY_MODEL_RUNNER_OUTPUT
return self.kv_connector_no_forward(scheduler_output,
self.vllm_config)
num_reqs = self.input_batch.num_reqs
req_ids = self.input_batch.req_ids
tokens = [scheduler_output.num_scheduled_tokens[i] for i in req_ids]
num_scheduled_tokens_np = np.array(tokens, dtype=np.int32)
max_num_scheduled_tokens = int(num_scheduled_tokens_np.max())
if self.dynamic_eplb:
self.eplb_updator.forward_before()
(
logits_indices,
spec_decode_metadata,
) = self._prepare_inputs(
scheduler_output,
num_scheduled_tokens_np,
)
(attn_metadata, num_scheduled_tokens_np, num_input_tokens,
num_tokens_across_dp, logits_indices, spec_decode_metadata,
max_query_len) = self._prepare_inputs(scheduler_output)
num_tokens_unpadded = scheduler_output.total_num_scheduled_tokens
if self.pcp_size > 1:
num_tokens_unpadded = self.pcp_manager.total_num_sampled_tokens_pcp
cascade_attn_prefix_lens = None
# Disable cascade attention when using microbatching (DBO)
if self.cascade_attn_enabled and not self.parallel_config.enable_dbo:
# Pre-compute cascade attention prefix lengths
cascade_attn_prefix_lens = self._compute_cascade_attn_prefix_lens(
num_scheduled_tokens_np,
self.input_batch.num_computed_tokens_cpu[:num_reqs],
scheduler_output.num_common_prefix_blocks,
)
(input_ids, inputs_embeds, positions, intermediate_tensors,
model_kwargs, ec_connector_output) = self._preprocess(scheduler_output,
num_input_tokens,
intermediate_tensors)
(
cudagraph_mode,
batch_desc,
should_ubatch,
num_tokens_across_dp,
cudagraph_stats,
) = self._determine_batch_execution_and_padding(
num_tokens=num_tokens_unpadded,
num_reqs=num_reqs,
num_scheduled_tokens_np=num_scheduled_tokens_np,
max_num_scheduled_tokens=max_num_scheduled_tokens,
use_cascade_attn=cascade_attn_prefix_lens is not None,
num_encoder_reqs=len(scheduler_output.scheduled_encoder_inputs),
)
logger.debug(
"Running batch with cudagraph_mode: %s, batch_descriptor: %s, "
"should_ubatch: %s, num_tokens_across_dp: %s",
cudagraph_mode,
batch_desc,
should_ubatch,
num_tokens_across_dp,
)
num_tokens_padded = batch_desc.num_tokens
num_reqs_padded = batch_desc.num_reqs if batch_desc.num_reqs is not None else num_reqs
ubatch_slices, ubatch_slices_padded = maybe_create_ubatch_slices(
should_ubatch,
num_scheduled_tokens_np,
num_tokens_padded,
num_reqs_padded,
self.parallel_config.num_ubatches,
)
pad_attn = cudagraph_mode == CUDAGraphMode.FULL
use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0
ubatch_slices_attn = ubatch_slices_padded if pad_attn else ubatch_slices
if (
cudagraph_mode == CUDAGraphMode.FULL
or (enable_sp() and not self.model_config.use_mla)
and self.pcp_size == 1 # TODO(lxs): fix this
):
# Currently, Graph Mode and SP will both pad num_tokens,
# Another possible condition is num_tokens_padded != num_tokens_unpadded
# but this scope is way too big and the consequences are unpredictable
num_reqs_padded = self._pad_query_start_loc_for_fia(num_tokens_padded, num_reqs_padded, num_reqs)
(attn_metadata, spec_decode_common_attn_metadata) = self._build_attention_metadata(
num_tokens=num_tokens_unpadded,
num_tokens_padded=num_tokens_padded,
num_reqs=num_reqs,
num_reqs_padded=num_reqs_padded,
max_query_len=max_num_scheduled_tokens,
ubatch_slices=ubatch_slices_attn,
logits_indices=logits_indices,
use_spec_decode=use_spec_decode,
num_scheduled_tokens=scheduler_output.num_scheduled_tokens,
num_scheduled_tokens_np=num_scheduled_tokens_np,
cascade_attn_prefix_lens=cascade_attn_prefix_lens,
)
(
input_ids,
inputs_embeds,
positions,
intermediate_tensors,
model_kwargs,
ec_connector_output,
) = self._preprocess(scheduler_output, num_tokens_padded, intermediate_tensors)
# update global cos, sin
update_cos_sin(positions)
if self.dynamic_eplb:
self.eplb_updator.take_update_info_from_eplb_process()
# Set cudagraph mode to none if calc_kv_scales is true.
# KV scales calculation involves dynamic operations that are incompatible
# with CUDA graph capture.
if self.calculate_kv_scales: # type: ignore[has-type]
cudagraph_mode = CUDAGraphMode.NONE
# Mark KV scales as calculated after the first forward pass
self.calculate_kv_scales = False # type: ignore[has-type]
# prevent debugger is None
if self.debugger is not None:
dbg_cfg = getattr(self.debugger, "config", None)
dump_level = str(
getattr(dbg_cfg, "level",
"L1")).upper() if dbg_cfg is not None else "L1"
dump_level = str(getattr(dbg_cfg, "level", "L1")).upper() if dbg_cfg is not None else "L1"
if dump_level in ("L0", "MIX"):
self.debugger.start(model=self.model)
else:
self.debugger.start()
uniform_decode = (max_query_len == self.uniform_decode_query_len) and (
scheduler_output.total_num_scheduled_tokens
== self.input_batch.num_reqs * max_query_len)
has_lora = len(self.input_batch.lora_id_to_lora_request) > 0
aclgraph_runtime_mode, batch_descriptor = \
self.cudagraph_dispatcher.dispatch(
num_tokens=num_input_tokens,
uniform_decode=uniform_decode,
has_lora=has_lora
)
if self.ascend_config.enable_async_exponential:
self.sampler.do_async_exponential(
b_s=logits_indices.shape[0],
head_dim=self.model_config.get_vocab_size(),
generators=self.input_batch.sampling_metadata.generators)
generators=self.input_batch.sampling_metadata.generators,
)
# Encoder-decoder models can only compile the pure decode steps where no
# encoder inputs are present. Use eager for the first pass.
num_encoder_reqs = len(scheduler_output.scheduled_encoder_inputs)
has_encoder_input = self.model_config.is_encoder_decoder and num_encoder_reqs > 0
# Run forward pass
with ProfileExecuteDuration().capture_async("forward"):
with set_ascend_forward_context(
attn_metadata,
self.vllm_config,
num_tokens=num_input_tokens,
num_tokens_across_dp=num_tokens_across_dp,
aclgraph_runtime_mode=aclgraph_runtime_mode,
batch_descriptor=batch_descriptor,
num_actual_tokens=scheduler_output.
total_num_scheduled_tokens,
model_instance=self.model):
self.maybe_setup_kv_connector(scheduler_output)
hidden_states = self._generate_process_reqs_hidden_states(
num_input_tokens, input_ids, positions,
intermediate_tensors, inputs_embeds, model_kwargs)
self.maybe_wait_for_kv_save()
finished_sending, finished_recving = self.get_finished_kv_transfer(
scheduler_output)
aux_hidden_states = None
if self.use_aux_hidden_state_outputs:
hidden_states, aux_hidden_states = hidden_states
kv_connector_output = KVConnectorOutput(
finished_sending=finished_sending,
finished_recving=finished_recving)
finished_sending = None
finished_recving = None
with ProfileExecuteDuration().capture_async("post process"):
with (
record_function_or_nullcontext("forward"),
set_ascend_forward_context(
attn_metadata,
self.vllm_config,
num_tokens=num_tokens_padded,
num_tokens_across_dp=num_tokens_across_dp,
aclgraph_runtime_mode=cudagraph_mode,
batch_descriptor=batch_desc,
num_actual_tokens=scheduler_output.total_num_scheduled_tokens,
model_instance=self.model,
skip_compiled=has_encoder_input,
),
self.maybe_get_kv_connector_output(scheduler_output) as kv_connector_output,
):
hidden_states = self._model_forward(
num_tokens_padded, input_ids, positions, intermediate_tensors, inputs_embeds, **model_kwargs
)
with record_function_or_nullcontext("post process"):
# -------------------------------------- Omni-new -------------------------------------------------
hidden_states, multimodal_outputs = self.extract_multimodal_outputs(hidden_states)
@@ -218,36 +310,36 @@ class NPUARModelRunner(OmniNPUModelRunner):
else:
logger.debug("[AR] execute_model: multimodal_outputs is None")
# -------------------------------------- Omni-new -------------------------------------------------
# Broadcast PP output for external_launcher (torchrun)
# to make sure we are synced across pp ranks
# TODO: Support overlapping mirco-batches
# https://github.com/vllm-project/vllm/issues/18019
broadcast_pp_output = \
self.parallel_config.distributed_executor_backend \
== "external_launcher" and len(get_pp_group().ranks) > 0
if not get_pp_group().is_last_rank:
# For mid-pipeline stages, return the hidden states.
if not broadcast_pp_output:
if self.pcp_size > 1:
# NOTE we must `slice` hidden_states because pcp_allgather_restore_idx
# ignores the padding from CUDA Graph.
hidden_states = self.pcp_manager.get_restore_hidden_states(hidden_states)
aux_hidden_states = None
if self.use_aux_hidden_state_outputs:
hidden_states, aux_hidden_states = hidden_states
if not self.broadcast_pp_output:
# Common case.
if not get_pp_group().is_last_rank:
# Return the intermediate tensors.
assert isinstance(hidden_states, IntermediateTensors)
hidden_states.kv_connector_output = kv_connector_output
self.kv_connector_output = kv_connector_output
if self.debugger is not None:
self.debugger.stop()
self.debugger.step()
return hidden_states
assert isinstance(hidden_states, IntermediateTensors)
get_pp_group().send_tensor_dict(
hidden_states.tensors, all_gather_group=get_tp_group())
logits = None
else:
if self.input_batch.pooling_params:
pool_output = self._pool(
hidden_states,
scheduler_output.total_num_scheduled_tokens,
num_scheduled_tokens_np, kv_connector_output)
if self.is_pooling_model:
# Return the pooling output.
output = self._pool(
hidden_states, num_scheduled_tokens, num_scheduled_tokens_np, kv_connector_output
)
output.kv_connector_output = kv_connector_output
if self.debugger is not None:
self.debugger.stop()
self.debugger.step()
return pool_output
return output
sample_hidden_states = hidden_states[logits_indices]
# -------------------------------------- Omni-new -------------------------------------------------
# Try with sampling_metadata first; fall back to without for models that don't support it
@@ -258,37 +350,56 @@ class NPUARModelRunner(OmniNPUModelRunner):
except TypeError:
logits = self.model.compute_logits(sample_hidden_states)
# -------------------------------------- Omni-new -------------------------------------------------
if broadcast_pp_output:
model_output_broadcast_data = {
"logits": logits.contiguous(),
} if logits is not None else {}
model_output_broadcast_data = get_pp_group(
).broadcast_tensor_dict(model_output_broadcast_data,
src=len(get_pp_group().ranks) - 1)
assert model_output_broadcast_data is not None
logits = model_output_broadcast_data["logits"]
else:
# Rare case.
assert not self.is_pooling_model
if not get_pp_group().is_last_rank:
sample_hidden_states = hidden_states[logits_indices]
get_pp_group().send_tensor_dict(hidden_states.tensors, all_gather_group=get_tp_group())
logits = None
else:
sample_hidden_states = hidden_states[logits_indices]
# -------------------------------------- Omni-new -------------------------------------------------
# Try with sampling_metadata first; fall back to without for models that don't support it
try:
logits = self.model.compute_logits(
sample_hidden_states, sampling_metadata=self.input_batch.sampling_metadata
)
except TypeError:
logits = self.model.compute_logits(sample_hidden_states)
# -------------------------------------- Omni-new -------------------------------------------------
model_output_broadcast_data: dict[str, Any] = {}
if logits is not None:
model_output_broadcast_data["logits"] = logits.contiguous()
broadcasted = get_pp_group().broadcast_tensor_dict(
model_output_broadcast_data, src=len(get_pp_group().ranks) - 1
)
assert broadcasted is not None
logits = broadcasted["logits"]
# Apply structured output bitmasks if present
self.execute_model_state = ExecuteModelState(
scheduler_output,
logits,
spec_decode_metadata,
spec_decode_common_attn_metadata,
hidden_states,
sample_hidden_states,
aux_hidden_states,
kv_connector_output,
attn_metadata,
positions,
ec_connector_output,
multimodal_outputs, # Omni-new
cudagraph_stats,
multimodal_outputs,
)
self.kv_connector_output = kv_connector_output
return None
@torch.inference_mode()
def sample_tokens(
self,
grammar_output: GrammarOutput | None = None,
self, grammar_output: GrammarOutput | None
) -> OmniModelRunnerOutput | AsyncModelRunnerOutput | IntermediateTensors:
kv_connector_output = self.kv_connector_output
self.kv_connector_output = None
@@ -298,6 +409,7 @@ class NPUARModelRunner(OmniNPUModelRunner):
self.kv_extracted_req_ids = None
# -------------------------------------- Omni-new -------------------------------------------------
if self.execute_model_state is None:
# Nothing to do (PP non-final rank case), output isn't used.
if not kv_connector_output:
@@ -316,13 +428,14 @@ class NPUARModelRunner(OmniNPUModelRunner):
scheduler_output,
logits,
spec_decode_metadata,
spec_decode_common_attn_metadata,
hidden_states,
sample_hidden_states,
aux_hidden_states,
kv_connector_output,
attn_metadata,
positions,
ec_connector_output,
cudagraph_stats,
multimodal_outputs,
) = self.execute_model_state
# Clear ephemeral state.
@@ -334,26 +447,26 @@ class NPUARModelRunner(OmniNPUModelRunner):
# the apply_grammar_bitmask uses torch.compile to optimize this,ascend does not support it now
logits_dtype = logits.dtype
logits = logits.to("cpu").float()
apply_grammar_bitmask(scheduler_output, grammar_output,
self.input_batch, logits)
apply_grammar_bitmask(scheduler_output, grammar_output, self.input_batch, logits)
logits = logits.to(self.device).to(logits_dtype)
with ProfileExecuteDuration().capture_async("Sample"):
with record_function_or_nullcontext("sample_token"):
sampler_output = self._sample(logits, spec_decode_metadata)
def propose_draft_token_ids(sampled_token_ids):
assert self.spec_decode_common_attn_metadata is not None
assert spec_decode_common_attn_metadata is not None
self._draft_token_ids = self.propose_draft_token_ids(
sampled_token_ids,
self.input_batch.sampling_metadata,
scheduler_output,
spec_decode_metadata,
spec_decode_common_attn_metadata,
positions,
scheduler_output.total_num_scheduled_tokens,
hidden_states,
attn_metadata,
aux_hidden_states,
sample_hidden_states
sample_hidden_states,
)
self._copy_draft_token_ids_to_cpu(scheduler_output)
@@ -373,11 +486,13 @@ class NPUARModelRunner(OmniNPUModelRunner):
spec_decode_metadata,
)
with ProfileExecuteDuration().capture_async("Draft"):
with record_function_or_nullcontext("draft_token"):
if self.speculative_config:
use_padded_batch_for_eagle = self.speculative_config and \
self.speculative_config.use_eagle() and \
not self.speculative_config.disable_padded_drafter_batch
use_padded_batch_for_eagle = (
self.speculative_config
and self.speculative_config.use_eagle()
and not self.speculative_config.disable_padded_drafter_batch
)
if use_padded_batch_for_eagle:
# EAGLE speculative decoding can use the GPU sampled tokens
# as inputs, and does not need to wait for bookkeeping to finish.
@@ -389,7 +504,6 @@ class NPUARModelRunner(OmniNPUModelRunner):
if has_kv_transfer_group():
get_kv_transfer_group().clear_connector_metadata()
# -------------------------------------- Omni-new -------------------------------------------------
hidden_states_cpu = hidden_states.detach().to("cpu").contiguous()
num_scheduled_tokens_np = getattr(self, "_omni_num_scheduled_tokens_np", None)
@@ -448,28 +562,15 @@ class NPUARModelRunner(OmniNPUModelRunner):
model_runner_output.kv_extracted_req_ids = kv_extracted_req_ids
# -------------------------------------- Omni-new -------------------------------------------------
durations = ProfileExecuteDuration().pop_captured_sync()
if durations:
dr_str = [
f"[{tag}]:{duration:.2f}ms"
for tag, duration in durations.items()
]
captured_name = "Decode" if self.attn_state == AscendAttentionState.DecodeOnly else "Prefill"
logger.info("Profile execute duration [%s]:%s", captured_name,
" ".join(dr_str))
if self.dynamic_eplb:
self.eplb_updator.forward_end()
if not self.use_async_scheduling:
if self.debugger is not None:
assert self.debugger is not None
self.debugger.stop()
self.debugger.step()
return model_runner_output
if self.debugger is not None:
assert self.debugger is not None
self.debugger.stop()
self.debugger.step()
if not self.use_async_scheduling:
return model_runner_output
return AsyncGPUModelRunnerOutput(
model_runner_output=model_runner_output,
sampled_token_ids=sampler_output.sampled_token_ids,
@@ -479,33 +580,6 @@ class NPUARModelRunner(OmniNPUModelRunner):
vocab_size=self.input_batch.vocab_size,
)
def _generate_process_reqs_hidden_states(self, num_input_tokens,
input_ids, positions,
intermediate_tensors,
inputs_embeds, model_kwargs):
assert self.model is not None
hidden_states = self._model_forward(input_ids=input_ids,
positions=positions,
intermediate_tensors=intermediate_tensors,
inputs_embeds=inputs_embeds,
**model_kwargs)
forward_context = get_forward_context()
if forward_context.cudagraph_runtime_mode == CUDAGraphMode.FULL \
and not self.use_sparse:
if self.vllm_config.model_config.use_mla:
if self.pcp_size * self.dcp_size > 1:
update_full_graph_params(self.attn_backend, self.update_stream, forward_context,
num_input_tokens, self.vllm_config,
self.vllm_config.speculative_config)
if get_forward_context().sp_enabled and not isinstance(
hidden_states, IntermediateTensors):
hidden_states = self._all_gather_hidden_states_and_aux(
hidden_states)
return hidden_states if self.pcp_size == 1 else self.pcp_manager.get_restore_hidden_states(
hidden_states)
def _resolve_global_request_id(self, req_id: str) -> str:
"""Resolve global request ID from request state."""
req_state = self.requests.get(req_id)
vllm_omni/platforms/npu/worker/npu_generation_model_runner.py
+429 -209
View File
@@ -4,8 +4,7 @@
from __future__ import annotations
import gc
import math
from copy import copy
from copy import copy, deepcopy
import numpy as np
import torch
@@ -19,12 +18,16 @@ from vllm.sequence import IntermediateTensors
from vllm.utils.math_utils import cdiv
from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT, AsyncModelRunnerOutput, make_empty_encoder_model_runner_output
from vllm.v1.worker.gpu_model_runner import AsyncGPUModelRunnerOutput
from vllm.v1.worker.kv_connector_model_runner_mixin import KVConnectorOutput
from vllm_ascend.ascend_forward_context import MoECommType, get_mc2_tokens_capacity, set_ascend_forward_context
from vllm.v1.utils import record_function_or_nullcontext
from vllm.v1.worker.gpu_model_runner import AsyncGPUModelRunnerOutput, PerLayerAttnMetadata
from vllm.v1.worker.ubatch_utils import maybe_create_ubatch_slices
from vllm.v1.worker.utils import sanity_check_mm_encoder_outputs
from vllm_ascend.ascend_forward_context import set_ascend_forward_context
from vllm_ascend.attention.attention_v1 import AscendAttentionState
from vllm_ascend.attention.utils import using_paged_attention
from vllm_ascend.ops.rotary_embedding import update_cos_sin
from vllm_ascend.platform import NPUPlatform
from vllm_ascend.utils import ProfileExecuteDuration, enable_sp, lmhead_tp_enable
from vllm_ascend.utils import enable_sp, lmhead_tp_enable
from vllm_ascend.worker.model_runner_v1 import SEQ_LEN_WITH_MAX_PA_WORKSPACE
from vllm_omni.outputs import OmniModelRunnerOutput
from vllm_omni.platforms.npu.worker.npu_ar_model_runner import ExecuteModelState
@@ -35,15 +38,31 @@ class NPUGenerationModelRunner(OmniNPUModelRunner):
"""Generation model runner for vLLM-omni on NPU (non-autoregressive)."""
def _update_request_states(self, scheduler_output: SchedulerOutput):
# remove requests
for req_id in scheduler_output.finished_req_ids:
self.input_batch.remove_request(req_id)
scheduled_req_ids = scheduler_output.num_scheduled_tokens.keys()
cached_req_ids = self.input_batch.req_id_to_index.keys()
resumed_req_ids = scheduler_output.scheduled_cached_reqs.resumed_req_ids
unscheduled_req_ids = cached_req_ids - (scheduled_req_ids - resumed_req_ids)
for req_id in unscheduled_req_ids:
self.input_batch.remove_request(req_id)
cached_reqs = scheduler_output.scheduled_cached_reqs
for _, req_id in enumerate(cached_reqs.req_ids):
req_states = []
for req_id in cached_reqs.req_ids:
req_state = self.requests.get(req_id)
assert req_state is not None
req_state.prompt_token_ids = cached_reqs.prompt_token_ids.get(req_id)
self.input_batch.remove_request(req_id)
req_states.append(req_state)
# Remove the request from the current input batch only if it is still present.
if req_id in self.input_batch.req_id_to_index:
self.input_batch.remove_request(req_id)
for req_state in req_states:
# update the request state in self.input_batch
self.input_batch.add_request(req_state)
self._init_mrope_positions(req_state)
# Only relevant for models using M-RoPE (e.g, Qwen2-VL)
if self.uses_mrope:
self._init_mrope_positions(req_state)
@torch.inference_mode()
def execute_model(
@@ -53,51 +72,169 @@ class NPUGenerationModelRunner(OmniNPUModelRunner):
) -> OmniModelRunnerOutput | IntermediateTensors:
if self.execute_model_state is not None:
raise RuntimeError("State error: sample_tokens() must be called after execute_model() returns None.")
with ProfileExecuteDuration().capture_async("prepare input"):
# self._draft_token_ids is None when `input_fits_in_drafter=False`
# and there is no draft tokens scheduled. so it need to update the
# spec_decoding info in scheduler_output with async_scheduling.
# use deepcopy to avoid the modification has influence on the
# scheduler_output in engine core process.
# TODO(Ronald1995): deepcopy is expensive when there is a large
# number of requests, optimize it later.
if (
self.use_async_scheduling and self.num_spec_tokens and self._draft_token_ids is None # type: ignore[has-type]
):
scheduler_output = deepcopy(scheduler_output)
num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
with record_function_or_nullcontext("prepare input"):
# -------------------------------------- Omni-new -------------------------------------------------
if self.model_config.async_chunk:
self._update_request_states(scheduler_output)
# -------------------------------------- Omni-new -------------------------------------------------
self._update_states(scheduler_output)
if has_ec_transfer() and get_ec_transfer().is_producer:
with self.maybe_get_ec_connector_output(
with self.synchronize_input_prep():
# Update persistent batch states.
self._update_states(scheduler_output)
if has_ec_transfer() and get_ec_transfer().is_producer:
with self.maybe_get_ec_connector_output(
scheduler_output,
encoder_cache=self.encoder_cache,
) as ec_connector_output:
self._execute_mm_encoder(scheduler_output)
return make_empty_encoder_model_runner_output(scheduler_output)
if not num_scheduled_tokens:
if (
self.parallel_config.distributed_executor_backend == "external_launcher"
and self.parallel_config.data_parallel_size > 1
):
# this is a corner case when both external launcher
# and DP are enabled, num_scheduled_tokens could be
# 0, and has_unfinished_requests in the outer loop
# returns True. before returning early here we call
# dummy run to ensure coordinate_batch_across_dp
# is called into to avoid out of sync issues.
self._dummy_run(1)
if not has_kv_transfer_group():
# Return empty ModelRunnerOutput if no work to do.
return EMPTY_MODEL_RUNNER_OUTPUT
return self.kv_connector_no_forward(scheduler_output, self.vllm_config)
if self.cache_config.kv_sharing_fast_prefill:
assert not self.num_prompt_logprobs, (
"--kv-sharing-fast-prefill produces incorrect "
"logprobs for prompt tokens, tokens, please disable "
"it when the requests need prompt logprobs"
)
num_reqs = self.input_batch.num_reqs
req_ids = self.input_batch.req_ids
tokens = [scheduler_output.num_scheduled_tokens[i] for i in req_ids]
num_scheduled_tokens_np = np.array(tokens, dtype=np.int32)
max_num_scheduled_tokens = int(num_scheduled_tokens_np.max())
(
logits_indices,
spec_decode_metadata,
) = self._prepare_inputs(
scheduler_output,
encoder_cache=self.encoder_cache,
) as ec_connector_output:
self._execute_mm_encoder(scheduler_output)
return make_empty_encoder_model_runner_output(scheduler_output)
num_scheduled_tokens_np,
)
if not scheduler_output.total_num_scheduled_tokens:
if not has_kv_transfer_group():
logger.debug("skip this step for we receive the data from remote disaggregate prefill node")
# Return empty ModelRunnerOutput if there's no work to do.
return EMPTY_MODEL_RUNNER_OUTPUT
return self.kv_connector_no_forward(scheduler_output, self.vllm_config)
num_tokens_unpadded = scheduler_output.total_num_scheduled_tokens
if self.pcp_size > 1:
num_tokens_unpadded = self.pcp_manager.total_num_sampled_tokens_pcp
cascade_attn_prefix_lens = None
# Disable cascade attention when using microbatching (DBO)
if self.cascade_attn_enabled and not self.parallel_config.enable_dbo:
# Pre-compute cascade attention prefix lengths
cascade_attn_prefix_lens = self._compute_cascade_attn_prefix_lens(
num_scheduled_tokens_np,
self.input_batch.num_computed_tokens_cpu[:num_reqs],
scheduler_output.num_common_prefix_blocks,
)
if self.dynamic_eplb:
self.eplb_updator.forward_before()
(
cudagraph_mode,
batch_desc,
should_ubatch,
num_tokens_across_dp,
cudagraph_stats,
) = self._determine_batch_execution_and_padding(
num_tokens=num_tokens_unpadded,
num_reqs=num_reqs,
num_scheduled_tokens_np=num_scheduled_tokens_np,
max_num_scheduled_tokens=max_num_scheduled_tokens,
use_cascade_attn=cascade_attn_prefix_lens is not None,
num_encoder_reqs=len(scheduler_output.scheduled_encoder_inputs),
)
logger.debug(
"Running batch with cudagraph_mode: %s, batch_descriptor: %s, "
"should_ubatch: %s, num_tokens_across_dp: %s",
cudagraph_mode,
batch_desc,
should_ubatch,
num_tokens_across_dp,
)
num_tokens_padded = batch_desc.num_tokens
num_reqs_padded = batch_desc.num_reqs if batch_desc.num_reqs is not None else num_reqs
ubatch_slices, ubatch_slices_padded = maybe_create_ubatch_slices(
should_ubatch,
num_scheduled_tokens_np,
num_tokens_padded,
num_reqs_padded,
self.parallel_config.num_ubatches,
)
pad_attn = cudagraph_mode == CUDAGraphMode.FULL
use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0
ubatch_slices_attn = ubatch_slices_padded if pad_attn else ubatch_slices
if (
cudagraph_mode == CUDAGraphMode.FULL
or (enable_sp() and not self.model_config.use_mla)
and self.pcp_size == 1 # TODO(lxs): fix this
):
# Currently, Graph Mode and SP will both pad num_tokens,
# Another possible condition is num_tokens_padded != num_tokens_unpadded
# but this scope is way too big and the consequences are unpredictable
num_reqs_padded = self._pad_query_start_loc_for_fia(num_tokens_padded, num_reqs_padded, num_reqs)
(attn_metadata, spec_decode_common_attn_metadata) = self._build_attention_metadata(
num_tokens=num_tokens_unpadded,
num_tokens_padded=num_tokens_padded,
num_reqs=num_reqs,
num_reqs_padded=num_reqs_padded,
max_query_len=max_num_scheduled_tokens,
ubatch_slices=ubatch_slices_attn,
logits_indices=logits_indices,
use_spec_decode=use_spec_decode,
num_scheduled_tokens=scheduler_output.num_scheduled_tokens,
num_scheduled_tokens_np=num_scheduled_tokens_np,
cascade_attn_prefix_lens=cascade_attn_prefix_lens,
)
(
attn_metadata,
num_scheduled_tokens_np,
num_input_tokens,
num_tokens_across_dp,
logits_indices,
spec_decode_metadata,
max_query_len,
) = self._prepare_inputs(scheduler_output)
input_ids,
inputs_embeds,
positions,
intermediate_tensors,
model_kwargs,
ec_connector_output,
) = self._preprocess(scheduler_output, num_tokens_padded, intermediate_tensors)
(input_ids, inputs_embeds, positions, intermediate_tensors, model_kwargs, ec_connector_output) = (
self._preprocess(scheduler_output, num_input_tokens, intermediate_tensors)
)
# [Omni] Pass token counts per request for code2wav output slicing
model_kwargs["seq_token_counts"] = tokens
# update global cos, sin
update_cos_sin(positions)
if self.dynamic_eplb:
self.eplb_updator.take_update_info_from_eplb_process()
# Set cudagraph mode to none if calc_kv_scales is true.
# KV scales calculation involves dynamic operations that are incompatible
# with CUDA graph capture.
if self.calculate_kv_scales: # type: ignore[has-type]
cudagraph_mode = CUDAGraphMode.NONE
# Mark KV scales as calculated after the first forward pass
self.calculate_kv_scales = False # type: ignore[has-type]
# prevent debugger is None
if self.debugger is not None:
dbg_cfg = getattr(self.debugger, "config", None)
@@ -106,15 +243,6 @@ class NPUGenerationModelRunner(OmniNPUModelRunner):
self.debugger.start(model=self.model)
else:
self.debugger.start()
uniform_decode = (max_query_len == self.uniform_decode_query_len) and (
scheduler_output.total_num_scheduled_tokens == self.input_batch.num_reqs * max_query_len
)
has_lora = len(self.input_batch.lora_id_to_lora_request) > 0
aclgraph_runtime_mode, batch_descriptor = self.cudagraph_dispatcher.dispatch(
num_tokens=num_input_tokens, uniform_decode=uniform_decode, has_lora=has_lora
)
if self.ascend_config.enable_async_exponential:
self.sampler.do_async_exponential(
b_s=logits_indices.shape[0],
@@ -122,63 +250,60 @@ class NPUGenerationModelRunner(OmniNPUModelRunner):
generators=self.input_batch.sampling_metadata.generators,
)
# Encoder-decoder models can only compile the pure decode steps where no
# encoder inputs are present. Use eager for the first pass.
num_encoder_reqs = len(scheduler_output.scheduled_encoder_inputs)
has_encoder_input = self.model_config.is_encoder_decoder and num_encoder_reqs > 0
# Run forward pass
with ProfileExecuteDuration().capture_async("forward"):
with set_ascend_forward_context(
with (
record_function_or_nullcontext("forward"),
set_ascend_forward_context(
attn_metadata,
self.vllm_config,
num_tokens=num_input_tokens,
num_tokens=num_tokens_padded,
num_tokens_across_dp=num_tokens_across_dp,
aclgraph_runtime_mode=aclgraph_runtime_mode,
batch_descriptor=batch_descriptor,
aclgraph_runtime_mode=cudagraph_mode,
batch_descriptor=batch_desc,
num_actual_tokens=scheduler_output.total_num_scheduled_tokens,
model_instance=self.model,
):
self.maybe_setup_kv_connector(scheduler_output)
# -------------------------------------- Omni-new -------------------------------------------------
outputs = self._run_generation_model(
input_ids=input_ids,
positions=positions,
intermediate_tensors=intermediate_tensors,
inputs_embeds=inputs_embeds,
model_kwargs=model_kwargs,
logits_indices=logits_indices,
)
# -------------------------------------- Omni-new -------------------------------------------------
self.maybe_wait_for_kv_save()
finished_sending, finished_recving = self.get_finished_kv_transfer(scheduler_output)
aux_hidden_states = None
if self.use_aux_hidden_state_outputs:
outputs, aux_hidden_states = outputs
kv_connector_output = KVConnectorOutput(finished_sending=finished_sending, finished_recving=finished_recving)
finished_sending = None
finished_recving = None
_, multimodal_outputs = self.extract_multimodal_outputs(outputs)
# Apply structured output bitmasks if present
self.execute_model_state = ExecuteModelState(
scheduler_output,
None,
spec_decode_metadata,
outputs,
None,
aux_hidden_states,
kv_connector_output,
attn_metadata,
positions,
ec_connector_output,
multimodal_outputs,
)
self.kv_connector_output = kv_connector_output
skip_compiled=has_encoder_input,
),
self.maybe_get_kv_connector_output(scheduler_output) as kv_connector_output,
):
# -------------------------------------- Omni-new -------------------------------------------------
outputs = self._run_generation_model(
num_tokens_padded=num_tokens_padded,
input_ids=input_ids,
positions=positions,
intermediate_tensors=intermediate_tensors,
inputs_embeds=inputs_embeds,
model_kwargs=model_kwargs,
logits_indices=logits_indices,
)
_, multimodal_outputs = self.extract_multimodal_outputs(outputs)
# -------------------------------------- Omni-new -------------------------------------------------
# Apply structured output bitmasks if present
self.execute_model_state = ExecuteModelState(
scheduler_output,
None,
spec_decode_metadata,
spec_decode_common_attn_metadata,
None,
None,
None,
attn_metadata,
positions,
ec_connector_output,
cudagraph_stats,
multimodal_outputs, # Omni-new: pass multimodal_outputs to ExecuteModelState
)
self.kv_connector_output = kv_connector_output
return None
@torch.inference_mode()
def sample_tokens(
self,
grammar_output: GrammarOutput | None = None,
self, grammar_output: GrammarOutput | None
) -> OmniModelRunnerOutput | AsyncModelRunnerOutput | IntermediateTensors:
kv_connector_output = self.kv_connector_output
self.kv_connector_output = None
@@ -201,13 +326,14 @@ class NPUGenerationModelRunner(OmniNPUModelRunner):
scheduler_output,
logits,
spec_decode_metadata,
spec_decode_common_attn_metadata,
hidden_states,
sample_hidden_states,
aux_hidden_states,
kv_connector_output,
attn_metadata,
positions,
ec_connector_output,
cudagraph_stats,
multimodal_outputs,
) = self.execute_model_state
# Clear ephemeral state.
@@ -231,11 +357,22 @@ class NPUGenerationModelRunner(OmniNPUModelRunner):
{"model_outputs": out.detach().to("cpu").contiguous() if out is not None else None}
)
elif isinstance(multimodal_outputs, dict):
mm_payload = {}
for key, out in multimodal_outputs.items():
if out is not None and isinstance(out, torch.Tensor):
mm_payload[key] = out.detach().to("cpu").contiguous()
pooler_output.append(mm_payload)
num_reqs = self.input_batch.num_reqs
for i in range(num_reqs):
mm_payload = {}
for key, out in multimodal_outputs.items():
if isinstance(out, list):
if len(out) != num_reqs:
raise ValueError(
f"Multimodal output list for key '{key}' has length {len(out)} "
f"but expected {num_reqs} (one entry per request)."
)
mm_payload[key] = out[i].detach().to("cpu").contiguous()
elif isinstance(out, torch.Tensor):
mm_payload[key] = out.detach().to("cpu").contiguous()
else:
logger.warning(f"Unsupported multimodal output type for key '{key}': {type(out)}")
pooler_output.append(mm_payload)
else:
raise RuntimeError("Unsupported diffusion output type")
# [Omni] Copy req_id mappings to avoid async scheduling mutation.
@@ -253,20 +390,25 @@ class NPUGenerationModelRunner(OmniNPUModelRunner):
ec_connector_output=ec_connector_output if self.supports_mm_inputs else None,
)
# -------------------------------------- Omni-new -------------------------------------------------
if self.debugger is not None:
self.debugger.stop()
self.debugger.step()
if not self.use_async_scheduling:
return output
return AsyncGPUModelRunnerOutput(
model_runner_output=output,
sampled_token_ids=torch.tensor([], device=self.device),
logprobs_tensors=None,
invalid_req_indices=[],
async_output_copy_stream=self.async_output_copy_stream,
vocab_size=self.input_batch.vocab_size,
logprobs_tensors=None,
)
def _run_generation_model(
self,
*,
num_tokens_padded: int,
input_ids: torch.Tensor | None,
positions: torch.Tensor | None,
intermediate_tensors: IntermediateTensors | None,
@@ -285,6 +427,7 @@ class NPUGenerationModelRunner(OmniNPUModelRunner):
"""
# Keep inputs identical to AR runner
kwargs = dict(
num_tokens_padded=num_tokens_padded,
input_ids=input_ids,
positions=positions,
intermediate_tensors=intermediate_tensors,
@@ -317,38 +460,16 @@ class NPUGenerationModelRunner(OmniNPUModelRunner):
force_attention: bool = False,
uniform_decode: bool = False,
is_profile: bool = False,
create_mixed_batch: bool = False,
allow_microbatching: bool = True,
skip_eplb: bool = False,
remove_lora: bool = True,
activate_lora: bool = False,
is_graph_capturing: bool = False,
) -> torch.Tensor:
num_active_loras: int = 0,
) -> tuple[torch.Tensor, torch.Tensor]:
# only support eager mode and piecewise graph now
assert cudagraph_runtime_mode is None or cudagraph_runtime_mode in {
CUDAGraphMode.NONE,
CUDAGraphMode.PIECEWISE,
CUDAGraphMode.FULL,
}
# In multi-DP scenarios, there may be situations where all DP groups are executing dummy runs.
# If sequence parallelism is enabled, it is essential to ensure that num_tokens is divisible by tp_size.
if self.use_aclgraph and enable_sp(self.vllm_config):
tp_size = self.vllm_config.parallel_config.tensor_parallel_size
num_tokens = math.ceil(num_tokens / tp_size) * tp_size
# Force dummy run on prefill stage when this node is deemed as kv producer.
if self.is_kv_producer and not self.is_kv_consumer:
with_prefill = True
has_lora = True if self.lora_config and self.compilation_config.cudagraph_specialize_lora else False
_ag_mode, batch_descriptor = self.cudagraph_dispatcher.dispatch(
num_tokens=num_tokens, uniform_decode=uniform_decode, has_lora=has_lora
)
# Padding for DP
(num_tokens, num_tokens_across_dp, with_prefill) = self._sync_metadata_across_dp(
batch_descriptor.num_tokens, with_prefill
)
assert cudagraph_runtime_mode is None or cudagraph_runtime_mode.valid_runtime_modes()
# If cudagraph_mode.decode_mode() == FULL and
# cudagraph_mode.separate_routine(). This means that we are using
# different graphs and/or modes for mixed prefill-decode batches vs.
@@ -363,78 +484,126 @@ class NPUGenerationModelRunner(OmniNPUModelRunner):
# routine of FA2 for pure decode, i.e., Flashdecode + an optimization
# for GQA/MQA.
max_query_len = self.uniform_decode_query_len if uniform_decode else num_tokens
# Set num_scheduled_tokens based on num_tokens and max_num_seqs
# for dummy run with LoRA so that the num_reqs collectively
# has num_tokens in total.
assert num_tokens <= self.scheduler_config.max_num_batched_tokens
max_num_reqs = self.max_num_reqs
if uniform_decode:
num_reqs = cdiv(num_tokens, max_query_len)
max_num_reqs = self.scheduler_config.max_num_seqs
if create_mixed_batch:
raise NotImplementedError("create_mixed_batch is used for warmup deepgemm, vllm-ascend does not need it")
elif uniform_decode:
assert not create_mixed_batch
num_reqs = min(max_num_reqs, cdiv(num_tokens, max_query_len))
num_scheduled_tokens_list = [max_query_len] * num_reqs
if num_tokens % max_query_len != 0:
num_scheduled_tokens_list[-1] = num_tokens % max_query_len
else:
if with_prefill:
num_reqs = num_tokens
else:
num_reqs = (num_tokens + self.decode_token_per_req - 1) // self.decode_token_per_req
num_reqs = min(num_reqs, max_num_reqs)
num_reqs = min(num_tokens, max_num_reqs)
min_tokens_per_req = num_tokens // num_reqs
num_scheduled_tokens_list = [min_tokens_per_req] * num_reqs
num_scheduled_tokens_list[-1] += num_tokens % num_reqs
assert sum(num_scheduled_tokens_list) == num_tokens
assert len(num_scheduled_tokens_list) == num_reqs
num_scheduled_tokens = np.array(num_scheduled_tokens_list, dtype=np.int32)
self.query_lens = torch.from_numpy(num_scheduled_tokens)
num_tokens_unpadded = int(num_scheduled_tokens.sum())
num_sampled_tokens = np.ones(num_reqs, dtype=np.int32)
if not is_profile and self.dynamic_eplb:
self.eplb_updator.forward_before()
if num_tokens != batch_descriptor.num_tokens:
_ag_mode, batch_descriptor = self.cudagraph_dispatcher.dispatch(
num_tokens=num_tokens, uniform_decode=uniform_decode, has_lora=has_lora
_cudagraph_mode, batch_desc, _, num_tokens_across_dp, _ = self._determine_batch_execution_and_padding(
num_tokens=num_tokens_unpadded,
num_reqs=num_reqs,
num_scheduled_tokens_np=num_scheduled_tokens,
max_num_scheduled_tokens=max_query_len,
use_cascade_attn=False,
allow_microbatching=allow_microbatching,
force_eager=is_profile or (cudagraph_runtime_mode == CUDAGraphMode.NONE),
# `force_uniform_decode` is used for cudagraph capture; because for
# capturing mixed prefill-decode batches, we sometimes use
# num_tokens == num_reqs which looks like a uniform decode batch to the
# dispatcher; but we actually want to capture a piecewise cudagraph
force_uniform_decode=uniform_decode,
# `force_has_lora` is used for cudagraph capture; because LoRA is
# activated later in the context manager, but we need to know the
# LoRA state when determining the batch descriptor for capture
force_has_lora=activate_lora,
)
if self.use_cp:
self.pcp_manager.init_batch_info(
num_scheduled_tokens,
num_reqs,
)
num_tokens_padded = batch_descriptor.num_tokens
num_reqs_padded = batch_descriptor.num_reqs if batch_descriptor.num_reqs is not None else num_reqs
if self.speculative_config:
self.pcp_manager.query_lens_pcp_full.cpu[:num_reqs] = torch.from_numpy(num_scheduled_tokens)
self.pcp_manager.query_lens_pcp_full.copy_to_gpu()
if cudagraph_runtime_mode is None:
cudagraph_runtime_mode = _cudagraph_mode
else:
assert cudagraph_runtime_mode == _cudagraph_mode, (
f"Cudagraph runtime mode mismatch in dummy_run. "
f"Expected {_cudagraph_mode}, but got {cudagraph_runtime_mode}."
)
num_tokens_padded = batch_desc.num_tokens
num_reqs_padded = batch_desc.num_reqs if batch_desc.num_reqs is not None else num_reqs
if num_tokens_across_dp is not None and num_tokens_padded != num_tokens:
# pad is needed if the pad of `num_tokens` is triggered inside CudagraphDispatcher
num_tokens_across_dp[:] = num_tokens_padded
num_scheduled_tokens = num_scheduled_tokens.repeat(num_reqs_padded)
# filter out the valid batch descriptor
if cudagraph_runtime_mode is not None:
# we allow forcing NONE when the dispatcher disagrees to support
# warm ups for aclgraph capture
if cudagraph_runtime_mode != CUDAGraphMode.NONE and cudagraph_runtime_mode != _ag_mode:
raise ValueError(
f"Aclgraph runtime mode mismatch at dummy_run. "
f"Expected {_ag_mode}, but got {cudagraph_runtime_mode}."
# vllm-ascend does not support ubatch now
ubatch_slices, ubatch_slices_padded = None, None
attn_metadata: PerLayerAttnMetadata | None = None
# If force_attention is True, we always capture attention. Otherwise,
# it only happens for cudagraph_runtime_mode=FULL.
if force_attention or cudagraph_runtime_mode == CUDAGraphMode.FULL:
if create_mixed_batch:
raise NotImplementedError(
"create_mixed_batch is used for warmup deepgemm, vllm-ascend does not need it"
)
else:
cudagraph_runtime_mode = _ag_mode
self.attn_state = AscendAttentionState.DecodeOnly
if self.speculative_config and self.speculative_config.method == "mtp":
# `AscendAttentionState.SpecDecoding` is only designed for mla
if self.vllm_config.model_config.use_mla:
self.attn_state = AscendAttentionState.SpecDecoding
else:
self.attn_state = AscendAttentionState.ChunkedPrefill
# The reason why we use a fixed seq_len rather than max_query_len is that
# _npu_paged_attention_get_workspace only returns max workspace with specific
# seq_lens. We use this seq_len only when capturing graph, and still use max_query_len
# in inference. This will be removed once npu_fused_infer_attention_score
# outperforms _npu_paged_attention on all cases.
seq_lens = (
SEQ_LEN_WITH_MAX_PA_WORKSPACE
if is_graph_capturing and using_paged_attention(num_tokens, self.vllm_config)
else max_query_len
) # type: ignore[assignment]
self.seq_lens.np[:num_reqs_padded] = seq_lens
self.seq_lens.np[num_reqs_padded:] = 0
self.seq_lens.copy_to_gpu()
# TODO(Mengqing): Set create_mixed_batch to False since it's only used in FI warmup
# and not supported in ASCEND now. We could remove it in the future.
attn_metadata = self._build_dummy_attn_metadata(
False,
num_reqs=num_reqs_padded,
num_tokens=num_tokens_padded,
max_query_len=max_query_len,
aclgraph_runtime_mode=cudagraph_runtime_mode,
force_attention=force_attention,
is_graph_capturing=is_graph_capturing,
num_scheduled_tokens=num_scheduled_tokens,
)
cum_num_tokens, _ = self._get_cumsum_and_arange(num_scheduled_tokens)
self.query_start_loc.np[1 : num_reqs_padded + 1] = cum_num_tokens
self.query_start_loc.copy_to_gpu()
with self.maybe_dummy_run_with_lora(self.lora_config, num_scheduled_tokens, num_sampled_tokens):
num_reqs_padded = self._pad_query_start_loc_for_fia(num_tokens_padded, num_reqs_padded, num_reqs)
pad_attn = cudagraph_runtime_mode == CUDAGraphMode.FULL
attn_metadata, _ = self._build_attention_metadata(
num_tokens=num_tokens_unpadded,
num_tokens_padded=num_tokens_padded,
num_reqs=num_reqs_padded,
max_query_len=max_query_len,
ubatch_slices=ubatch_slices_padded if pad_attn else ubatch_slices,
for_cudagraph_capture=is_graph_capturing,
num_scheduled_tokens_np=num_scheduled_tokens,
)
with self.maybe_dummy_run_with_lora(
self.lora_config,
num_scheduled_tokens,
num_sampled_tokens,
):
# Make sure padding doesn't exceed max_num_tokens
assert num_tokens_padded <= self.max_num_tokens
if self.is_multimodal_model and not self.model_config.is_encoder_decoder:
input_ids = None
inputs_embeds = self.inputs_embeds.gpu[:num_tokens_padded]
elif self.enable_prompt_embeds:
if self.is_multimodal_model and not self.model_config.is_encoder_decoder or self.enable_prompt_embeds:
input_ids = None
inputs_embeds = self.inputs_embeds.gpu[:num_tokens_padded]
else:
@@ -454,20 +623,22 @@ class NPUGenerationModelRunner(OmniNPUModelRunner):
if get_pp_group().is_first_rank:
intermediate_tensors = None
else:
# When PP and flashcomm1 are enabled,
# during dummy_run the estimated space should divide num_tokens by tp_size;
# otherwise, on non-first PP ranks it would effectively perform an extra all-gather,
# leading to incorrect memory estimation and potentially causing OOM.
actual_tokens = num_tokens
# When PP and flashcomm1 are enabled, during dummy_run the estimated space should divide num_tokens by
# tp_size; otherwise, on non-first PP ranks it would effectively perform an extra all-gather, leading
# to incorrect memory estimation and potentially causing OOM.
intermediate_tokens = num_tokens_padded
if enable_sp():
tp_size = get_tensor_model_parallel_world_size()
actual_tokens = num_tokens // tp_size
intermediate_tokens = (num_tokens_padded + tp_size - 1) // tp_size
if self.intermediate_tensors is None:
max_actual_tokens = self.max_num_tokens
if enable_sp():
max_actual_tokens = (self.max_num_tokens + tp_size - 1) // tp_size
self.intermediate_tensors = self.model.make_empty_intermediate_tensors(
batch_size=actual_tokens, dtype=self.dtype, device=self.device
batch_size=max_actual_tokens, dtype=self.dtype, device=self.device
)
intermediate_tensors = IntermediateTensors(
{k: v[:num_tokens_padded] for k, v in self.intermediate_tensors.items()}
{k: v[:intermediate_tokens] for k, v in self.intermediate_tensors.items()}
)
need_dummy_logits = not is_profile and lmhead_tp_enable()
@@ -485,21 +656,32 @@ class NPUGenerationModelRunner(OmniNPUModelRunner):
if hasattr(self.drafter, "model") and hasattr(self.drafter.model, "compute_logits"):
return self.drafter.model.compute_logits(hidden_states[dummy_indices])
model_kwargs = self._init_model_kwargs()
with set_ascend_forward_context(
attn_metadata,
self.vllm_config,
num_tokens=num_tokens_padded,
num_tokens_across_dp=num_tokens_across_dp,
in_profile_run=is_profile,
num_actual_tokens=0,
num_actual_tokens=num_tokens_padded,
aclgraph_runtime_mode=cudagraph_runtime_mode,
batch_descriptor=batch_descriptor,
batch_descriptor=batch_desc,
model_instance=self.model,
):
hidden_states = self._generate_dummy_run_hidden_states(
input_ids, positions, num_tokens_padded, intermediate_tensors, inputs_embeds
outputs = self.model(
input_ids=input_ids,
positions=positions,
intermediate_tensors=intermediate_tensors,
inputs_embeds=inputs_embeds,
**model_kwargs,
)
dummy_compute_logits(hidden_states)
if self.use_aux_hidden_state_outputs:
hidden_states, _ = outputs
else:
hidden_states = outputs
dummy_compute_logits(hidden_states)
if self.drafter:
self.drafter.dummy_run(
@@ -508,7 +690,7 @@ class NPUGenerationModelRunner(OmniNPUModelRunner):
num_reqs=num_reqs_padded,
num_tokens_across_dp=num_tokens_across_dp,
aclgraph_runtime_mode=cudagraph_runtime_mode,
batch_descriptor=batch_descriptor,
batch_descriptor=batch_desc,
dummy_compute_logits=dummy_drafter_compute_logits,
in_graph_capturing=not force_attention,
is_profile=is_profile,
@@ -518,30 +700,68 @@ class NPUGenerationModelRunner(OmniNPUModelRunner):
if self.dynamic_eplb:
self.eplb_updator.take_update_info_from_eplb_process()
self.eplb_updator.forward_end()
# -------------------------------------- Omni-new -------------------------------------------------
hidden_states, _ = self.extract_multimodal_outputs(hidden_states)
# -------------------------------------------------------------------------------------------------
return hidden_states
return hidden_states, None
def profile_run(self) -> None:
# Trigger compilation for general shape.
with self.set_in_profile_run():
hidden_states = self._dummy_run(
self.max_num_tokens // self.pcp_size if self.pcp_size > 1 else self.max_num_tokens, with_prefill=True
)
# MC2 will consume additional NPU memory.
# Therefore, we need to run the MC2 path once here to complete its initialization,
# allowing vLLM to correctly estimate the maximum memory required.
mc2_tokens_capacity = get_mc2_tokens_capacity()
if (
self.max_num_tokens > mc2_tokens_capacity
and self._select_moe_comm_method(mc2_tokens_capacity) == MoECommType.MC2
):
self._dummy_run(mc2_tokens_capacity, with_prefill=True)
# Profile with multimodal encoder & encoder cache.
if self.supports_mm_inputs:
mm_config = self.model_config.multimodal_config
if mm_config is not None and mm_config.skip_mm_profiling:
logger.info("Skipping memory profiling for multimodal encoder and encoder cache.")
else:
mm_budget = self.mm_budget
assert mm_budget is not None
if (encoder_budget := mm_budget.get_encoder_budget()) > 0:
if not mm_budget.mm_max_toks_per_item:
# All modality limits are 0 — embedding-only mode.
# Budget is non-zero for embedding storage, but
# there's no encoder to profile.
logger.info(
"Skipping encoder profiling for embedding-only "
"mode (all modality limits=0 with "
"enable_mm_embeds=True).",
)
else:
# NOTE: Currently model is profiled with a single non-text
# modality with the max possible input tokens even when
# it supports multiple.
dummy_modality = mm_budget.get_modality_with_max_tokens()
max_mm_items_per_batch = mm_budget.mm_max_items_per_batch[dummy_modality]
logger.info(
"Encoder cache will be initialized with a budget of "
"%s tokens, and profiled with %s %s items of the "
"maximum feature size.",
encoder_budget,
max_mm_items_per_batch,
dummy_modality,
)
# Create dummy batch of multimodal inputs.
batched_dummy_mm_inputs = self._get_mm_dummy_batch(
dummy_modality,
max_mm_items_per_batch,
)
# Run multimodal encoder.
dummy_encoder_outputs = self.model.embed_multimodal(**batched_dummy_mm_inputs)
sanity_check_mm_encoder_outputs(
dummy_encoder_outputs,
expected_num_items=max_mm_items_per_batch,
)
for i, output in enumerate(dummy_encoder_outputs):
self.encoder_cache[f"tmp_{i}"] = output
# Add `is_profile` here to pre-allocate communication buffers
hidden_states, _ = self._dummy_run(self.max_num_tokens, is_profile=True)
output = None
NPUPlatform.synchronize()
del hidden_states, output
self._sync_device()
del hidden_states
self.encoder_cache.clear()
gc.collect()
vllm_omni/platforms/npu/worker/npu_model_runner.py
+208 -748
View File
File diff suppressed because it is too large. Load diff
vllm_omni/worker/gpu_model_runner.py
+2 -1
View File
@@ -925,7 +925,8 @@ class OmniGPUModelRunner(GPUModelRunner):
def _update_additional_information(self, scheduler_output: "SchedulerOutput") -> None:
for new_req in scheduler_output.scheduled_new_reqs:
payload_info = getattr(new_req, "additional_information", None)
self._merge_additional_information_update(new_req.req_id, payload_info)
if isinstance(payload_info, dict):
self._merge_additional_information_update(new_req.req_id, payload_info)
if hasattr(scheduler_output.scheduled_cached_reqs, "additional_information"):
cached_infos = getattr(scheduler_output.scheduled_cached_reqs, "additional_information", {})