diff --git a/.buildkite/test-nightly.yml b/.buildkite/test-nightly.yml index 1f1506870..b70e73c48 100644 --- a/.buildkite/test-nightly.yml +++ b/.buildkite/test-nightly.yml @@ -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" diff --git a/.buildkite/test-ready.yml b/.buildkite/test-ready.yml index e0cb91b08..d7595e26c 100644 --- a/.buildkite/test-ready.yml +++ b/.buildkite/test-ready.yml @@ -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 diff --git a/docker/Dockerfile.ci b/docker/Dockerfile.ci index 059c0eb89..ac6adf8a0 100644 --- a/docker/Dockerfile.ci +++ b/docker/Dockerfile.ci @@ -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// -# 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 [] diff --git a/docker/Dockerfile.npu b/docker/Dockerfile.npu index c555372d9..1c65d43b0 100644 --- a/docker/Dockerfile.npu +++ b/docker/Dockerfile.npu @@ -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 diff --git a/docker/Dockerfile.npu.a3 b/docker/Dockerfile.npu.a3 index 413ed88c3..33b5d3c4e 100644 --- a/docker/Dockerfile.npu.a3 +++ b/docker/Dockerfile.npu.a3 @@ -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 diff --git a/docs/configuration/stage_configs.md b/docs/configuration/stage_configs.md index 1f1a18902..0e671bebf 100644 --- a/docs/configuration/stage_configs.md +++ b/docs/configuration/stage_configs.md @@ -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. diff --git a/docs/getting_started/installation/npu.md b/docs/getting_started/installation/npu.md index 5ee2c0ee4..0d538f25d 100644 --- a/docs/getting_started/installation/npu.md +++ b/docs/getting_started/installation/npu.md @@ -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" diff --git a/docs/getting_started/installation/npu/npu.inc.md b/docs/getting_started/installation/npu/npu.inc.md index c473820fe..50e18697b 100644 --- a/docs/getting_started/installation/npu/npu.inc.md +++ b/docs/getting_started/installation/npu/npu.inc.md @@ -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 \ diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 0fbbe6eb7..34a16a313 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -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 diff --git a/docs/user_guide/diffusion_acceleration.md b/docs/user_guide/diffusion_acceleration.md index 1a2e0a7d2..4a7ba65ee 100644 --- a/docs/user_guide/diffusion_acceleration.md +++ b/docs/user_guide/diffusion_acceleration.md @@ -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 diff --git a/examples/offline_inference/text_to_image/README.md b/examples/offline_inference/text_to_image/README.md index 9c5ceb927..2bfae8472 100644 --- a/examples/offline_inference/text_to_image/README.md +++ b/examples/offline_inference/text_to_image/README.md @@ -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. diff --git a/examples/offline_inference/text_to_image/text_to_image.py b/examples/offline_inference/text_to_image/text_to_image.py index db87abf00..04f66863c 100644 --- a/examples/offline_inference/text_to_image/text_to_image.py +++ b/examples/offline_inference/text_to_image/text_to_image.py @@ -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 diff --git a/examples/online_serving/qwen3_tts/README.md b/examples/online_serving/qwen3_tts/README.md index e1ad3edc2..dd4619cc4 100644 --- a/examples/online_serving/qwen3_tts/README.md +++ b/examples/online_serving/qwen3_tts/README.md @@ -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 diff --git a/tests/conftest.py b/tests/conftest.py index 08d26a6d1..c3ff10246 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -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()) diff --git a/tests/diffusion/models/nextstep_1_1/test_nextstep_cfg_parallel_layout.py b/tests/diffusion/models/nextstep_1_1/test_nextstep_cfg_parallel_layout.py new file mode 100644 index 000000000..9caf667be --- /dev/null +++ b/tests/diffusion/models/nextstep_1_1/test_nextstep_cfg_parallel_layout.py @@ -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"" + + 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 = 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("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"] diff --git a/tests/e2e/online_serving/test_qwen3_omni.py b/tests/e2e/online_serving/test_qwen3_omni.py index d8ea61550..74eaa9c6b 100644 --- a/tests/e2e/online_serving/test_qwen3_omni.py +++ b/tests/e2e/online_serving/test_qwen3_omni.py @@ -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"], }, } diff --git a/tests/e2e/online_serving/test_qwen3_omni_expansion.py b/tests/e2e/online_serving/test_qwen3_omni_expansion.py index 84cc85ba1..c4731ffc7 100644 --- a/tests/e2e/online_serving/test_qwen3_omni_expansion.py +++ b/tests/e2e/online_serving/test_qwen3_omni_expansion.py @@ -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) diff --git a/tests/entrypoints/openai_api/test_serving_speech.py b/tests/entrypoints/openai_api/test_serving_speech.py index e55ff6812..11d676839 100644 --- a/tests/entrypoints/openai_api/test_serving_speech.py +++ b/tests/entrypoints/openai_api/test_serving_speech.py @@ -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 diff --git a/tests/perf/tests/test.json b/tests/perf/tests/test.json index da965a518..32bbf1b5a 100644 --- a/tests/perf/tests/test.json +++ b/tests/perf/tests/test.json @@ -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 - } } ] } diff --git a/vllm_omni/core/sched/omni_generation_scheduler.py b/vllm_omni/core/sched/omni_generation_scheduler.py index 071620de9..96e7757f8 100644 --- a/vllm_omni/core/sched/omni_generation_scheduler.py +++ b/vllm_omni/core/sched/omni_generation_scheduler.py @@ -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: diff --git a/vllm_omni/diffusion/attention/backends/sdpa.py b/vllm_omni/diffusion/attention/backends/sdpa.py index cfe460e3d..3585689dd 100644 --- a/vllm_omni/diffusion/attention/backends/sdpa.py +++ b/vllm_omni/diffusion/attention/backends/sdpa.py @@ -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") diff --git a/vllm_omni/diffusion/distributed/vae_patch_parallel.py b/vllm_omni/diffusion/distributed/vae_patch_parallel.py index 1a3edde3c..c5a63ba62 100644 --- a/vllm_omni/diffusion/distributed/vae_patch_parallel.py +++ b/vllm_omni/diffusion/distributed/vae_patch_parallel.py @@ -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): diff --git a/vllm_omni/diffusion/models/nextstep_1_1/__init__.py b/vllm_omni/diffusion/models/nextstep_1_1/__init__.py new file mode 100644 index 000000000..1dfaf51c8 --- /dev/null +++ b/vllm_omni/diffusion/models/nextstep_1_1/__init__.py @@ -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", +] diff --git a/vllm_omni/diffusion/models/nextstep_1_1/modeling_flux_vae.py b/vllm_omni/diffusion/models/nextstep_1_1/modeling_flux_vae.py new file mode 100644 index 000000000..8c30d4b13 --- /dev/null +++ b/vllm_omni/diffusion/models/nextstep_1_1/modeling_flux_vae.py @@ -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 diff --git a/vllm_omni/diffusion/models/nextstep_1_1/modeling_nextstep.py b/vllm_omni/diffusion/models/nextstep_1_1/modeling_nextstep.py new file mode 100644 index 000000000..ded307926 --- /dev/null +++ b/vllm_omni/diffusion/models/nextstep_1_1/modeling_nextstep.py @@ -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 diff --git a/vllm_omni/diffusion/models/nextstep_1_1/modeling_nextstep_heads.py b/vllm_omni/diffusion/models/nextstep_1_1/modeling_nextstep_heads.py new file mode 100644 index 000000000..7768ac765 --- /dev/null +++ b/vllm_omni/diffusion/models/nextstep_1_1/modeling_nextstep_heads.py @@ -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) diff --git a/vllm_omni/diffusion/models/nextstep_1_1/modeling_nextstep_llama.py b/vllm_omni/diffusion/models/nextstep_1_1/modeling_nextstep_llama.py new file mode 100644 index 000000000..7b367b6ff --- /dev/null +++ b/vllm_omni/diffusion/models/nextstep_1_1/modeling_nextstep_llama.py @@ -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 diff --git a/vllm_omni/diffusion/models/nextstep_1_1/pipeline_nextstep_1_1.py b/vllm_omni/diffusion/models/nextstep_1_1/pipeline_nextstep_1_1.py new file mode 100644 index 000000000..4a20c28c5 --- /dev/null +++ b/vllm_omni/diffusion/models/nextstep_1_1/pipeline_nextstep_1_1.py @@ -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 tokens in captions + image_token_count = 0 + for caption in captions: + num_image_token = len(re.findall(r"", caption)) + if num_image_token != 1: + raise ValueError( + f"Caption must contain exactly one 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 tokens with corresponding image_str + processed_captions = [] + image_idx = 0 + for caption in captions: + processed_caption = caption + num_image_tokens = processed_caption.count("") + + for _ in range(num_image_tokens): + processed_caption = processed_caption.replace("", 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) diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index 173c576a2..f862dd0b5 100644 --- a/vllm_omni/diffusion/registry.py +++ b/vllm_omni/diffusion/registry.py @@ -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", } diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py index 134a2ae39..5fee65645 100644 --- a/vllm_omni/diffusion/worker/diffusion_model_runner.py +++ b/vllm_omni/diffusion/worker/diffusion_model_runner.py @@ -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 diff --git a/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py b/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py index 82c06fdaf..7c7d90ca1 100644 --- a/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py +++ b/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py @@ -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. diff --git a/vllm_omni/entrypoints/async_omni_diffusion.py b/vllm_omni/entrypoints/async_omni_diffusion.py index 11a3c07e1..072daee30 100644 --- a/vllm_omni/entrypoints/async_omni_diffusion.py +++ b/vllm_omni/entrypoints/async_omni_diffusion.py @@ -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: diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index 73bb03805..74a4ff7f0 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -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 diff --git a/vllm_omni/entrypoints/omni.py b/vllm_omni/entrypoints/omni.py index 94409b3e6..dc18978f9 100644 --- a/vllm_omni/entrypoints/omni.py +++ b/vllm_omni/entrypoints/omni.py @@ -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, diff --git a/vllm_omni/entrypoints/omni_diffusion.py b/vllm_omni/entrypoints/omni_diffusion.py index 9f6dde15b..e0b8911d4 100644 --- a/vllm_omni/entrypoints/omni_diffusion.py +++ b/vllm_omni/entrypoints/omni_diffusion.py @@ -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: diff --git a/vllm_omni/entrypoints/omni_stage.py b/vllm_omni/entrypoints/omni_stage.py index 89eda4f16..f245df84c 100644 --- a/vllm_omni/entrypoints/omni_stage.py +++ b/vllm_omni/entrypoints/omni_stage.py @@ -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. diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index 6874b4202..7278552cc 100644 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/vllm_omni/entrypoints/openai/serving_chat.py @@ -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: diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index 201be69da..98f79f87f 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -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: diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py index 936a18427..39baa5573 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py +++ b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py @@ -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 diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_talker.py b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_talker.py index 927bc5525..988d8c236 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_talker.py +++ b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_talker.py @@ -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( diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py index 8ee516376..a4c8dcab5 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py +++ b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py @@ -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) diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py index 58eb1ff22..4bd9bd3ba 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py @@ -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, ) diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py index ac7d8120d..400696f18 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py @@ -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 diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py index 6e1b4b9b4..5c6f71d5d 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py @@ -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: diff --git a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts.py b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts.py index 664d2a295..418322672 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts.py +++ b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts.py @@ -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] diff --git a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/configuration_qwen3_tts_tokenizer_v1.py b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/configuration_qwen3_tts_tokenizer_v1.py index 74272f936..74775591d 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/configuration_qwen3_tts_tokenizer_v1.py +++ b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/configuration_qwen3_tts_tokenizer_v1.py @@ -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 diff --git a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/modeling_qwen3_tts_tokenizer_v1.py b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/modeling_qwen3_tts_tokenizer_v1.py index bceafc98e..cf4622fac 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/modeling_qwen3_tts_tokenizer_v1.py +++ b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/modeling_qwen3_tts_tokenizer_v1.py @@ -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, diff --git a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/vq/speech_vq.py b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/vq/speech_vq.py index 805feb81f..de2c69702 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/vq/speech_vq.py +++ b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/vq/speech_vq.py @@ -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, ) diff --git a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/vq/whisper_encoder.py b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/vq/whisper_encoder.py index c1c3c987a..e3bd6e1c3 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/vq/whisper_encoder.py +++ b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/vq/whisper_encoder.py @@ -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) diff --git a/vllm_omni/model_executor/stage_configs/qwen3_tts.yaml b/vllm_omni/model_executor/stage_configs/qwen3_tts.yaml index 1f29f0796..90053852a 100644 --- a/vllm_omni/model_executor/stage_configs/qwen3_tts.yaml +++ b/vllm_omni/model_executor/stage_configs/qwen3_tts.yaml @@ -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 diff --git a/vllm_omni/platforms/npu/stage_configs/qwen3_omni_moe.yaml b/vllm_omni/platforms/npu/stage_configs/qwen3_omni_moe.yaml index f99ed22e8..706bee5dd 100644 --- a/vllm_omni/platforms/npu/stage_configs/qwen3_omni_moe.yaml +++ b/vllm_omni/platforms/npu/stage_configs/qwen3_omni_moe.yaml @@ -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 diff --git a/vllm_omni/platforms/npu/stage_configs/qwen3_tts.yaml b/vllm_omni/platforms/npu/stage_configs/qwen3_tts.yaml index fbfbf10a4..6e37ec018 100644 --- a/vllm_omni/platforms/npu/stage_configs/qwen3_tts.yaml +++ b/vllm_omni/platforms/npu/stage_configs/qwen3_tts.yaml @@ -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 diff --git a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py index fbf9e632f..b0b4ea09d 100644 --- a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py @@ -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) diff --git a/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py b/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py index d263fb0d3..1fb1ac07a 100644 --- a/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py @@ -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() diff --git a/vllm_omni/platforms/npu/worker/npu_model_runner.py b/vllm_omni/platforms/npu/worker/npu_model_runner.py index 3aa68431b..14ec559dd 100644 --- a/vllm_omni/platforms/npu/worker/npu_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_model_runner.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import math -from typing import TYPE_CHECKING, Any, cast +from typing import Any import numpy as np import torch @@ -10,41 +9,30 @@ from vllm.config import CUDAGraphMode from vllm.distributed import get_tensor_model_parallel_world_size from vllm.distributed.parallel_state import get_pp_group from vllm.logger import init_logger -from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding -from vllm.model_executor.models.interfaces import supports_mrope -from vllm.model_executor.models.interfaces_base import VllmModelForPooling -from vllm.sampling_params import SamplingType from vllm.sequence import IntermediateTensors from vllm.utils.math_utils import cdiv -from vllm.v1.worker.gpu_input_batch import CachedRequestState -from vllm_ascend.ascend_forward_context import set_ascend_forward_context -from vllm_ascend.compilation.acl_graph import ACLGraphWrapper +from vllm.v1.worker.gpu_model_runner import PerLayerAttnMetadata +from vllm_ascend.ascend_forward_context import get_forward_context, set_ascend_forward_context +from vllm_ascend.attention.attention_v1 import AscendAttentionState +from vllm_ascend.attention.utils import using_paged_attention +from vllm_ascend.compilation.acl_graph import ACLGraphWrapper, update_full_graph_params from vllm_ascend.ops.rotary_embedding import update_cos_sin from vllm_ascend.utils import enable_sp, lmhead_tp_enable -from vllm_ascend.worker.model_runner_v1 import NPUModelRunner +from vllm_ascend.worker.model_runner_v1 import SEQ_LEN_WITH_MAX_PA_WORKSPACE, NPUModelRunner from vllm_omni.model_executor.models.output_templates import OmniOutput from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner -if TYPE_CHECKING: - from vllm.v1.core.sched.output import SchedulerOutput - logger = init_logger(__name__) -class OmniNPUModelRunner(NPUModelRunner): - """ - Base class for NPU model runners with multimodality support. - """ - - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - self._omni_per_req_additional_information: dict[str, dict] | None = None - self._omni_num_scheduled_tokens_np: np.ndarray | None = None - self._omni_last_model_output: object | None = None - +class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner): def load_model(self, *args, **kwargs) -> None: - super().load_model(*args, **kwargs) + NPUModelRunner.load_model(self, *args, **kwargs) + # Initialize enable_sp cache to avoid get_current_vllm_config() error + # in _pad_for_sequence_parallelism during execute_model. + # This is a workaround for vllm-ascend not passing vllm_config to enable_sp(). + enable_sp(self.vllm_config) # TODO move this model specific logic to a separate class if hasattr(self.model, "talker_mtp") and self.model.talker is not None: self.talker_mtp = self.model.talker_mtp @@ -63,308 +51,6 @@ class OmniNPUModelRunner(NPUModelRunner): self.last_talker_hidden = self._make_buffer(max_batch_size, hidden_size, dtype=self.dtype, numpy=False) self.text_step = self._make_buffer(max_batch_size, hidden_size, dtype=self.dtype, numpy=False) - def _init_mrope_positions(self, req_state: CachedRequestState): - image_grid_thw = [] - video_grid_thw = [] - second_per_grid_ts = [] - audio_feature_lengths = [] - use_audio_in_video = False - for mm_feature in req_state.mm_features: - mm_item = mm_feature.data - if mm_item is None: - continue - mm_input = mm_item.get_data() - if (t := mm_input.get("image_grid_thw")) is not None: - image_grid_thw.append(t.tolist()) - if (t := mm_input.get("video_grid_thw")) is not None: - video_grid_thw.append(t.tolist()) - if (t := mm_input.get("second_per_grid_ts")) is not None: - second_per_grid_ts.append(t) - if (t := mm_input.get("audio_feature_lengths")) is not None: - audio_feature_lengths.append(t) - # Check for use_audio_in_video - use_audio_in_video_value = mm_input.get("use_audio_in_video") - if use_audio_in_video_value is not None: - use_audio_in_video = bool(use_audio_in_video_value.item()) - - if supports_mrope(self.get_model()): - req_state.mrope_positions, req_state.mrope_position_delta = self.model.get_mrope_input_positions( - req_state.prompt_token_ids, - mm_features=req_state.mm_features, - hf_config=self.model_config.hf_config, - image_grid_thw=image_grid_thw, - video_grid_thw=video_grid_thw, - second_per_grid_ts=second_per_grid_ts, - audio_feature_lengths=audio_feature_lengths, - use_audio_in_video=use_audio_in_video, - ) - else: - req_state.mrope_positions, req_state.mrope_position_delta = MRotaryEmbedding.get_input_positions_tensor( - req_state.prompt_token_ids, - hf_config=self.model_config.hf_config, - image_grid_thw=image_grid_thw, - video_grid_thw=video_grid_thw, - second_per_grid_ts=second_per_grid_ts, - audio_feature_lengths=audio_feature_lengths, - use_audio_in_video=use_audio_in_video, - ) - - def _update_states(self, scheduler_output: "SchedulerOutput") -> None: - """Update the cached states and the persistent batch with the scheduler - output. - - The updated states are used by the `_prepare_inputs` function to create - the input GPU tensors for the model. - - The SamplingMetadata is updated and copied to the GPU if there is a - new/resumed/paused/finished request in the batch. - """ - # Remove finished requests from the cached states. - for req_id in scheduler_output.finished_req_ids: - self.requests.pop(req_id, None) - self.num_prompt_logprobs.pop(req_id, None) - # Remove the finished requests from the persistent batch. - # NOTE(woosuk): There could be an edge case where finished_req_ids and - # scheduled_req_ids overlap. This happens when a request is aborted and - # then resubmitted with the same ID. In this case, we treat them as two - # distinct requests - clearing the cached states for the first request - # and handling the second as a new request. - for req_id in scheduler_output.finished_req_ids: - self.input_batch.remove_request(req_id) - - # Free the cached encoder outputs. - for mm_hash in scheduler_output.free_encoder_mm_hashes: - self.encoder_cache.pop(mm_hash, None) - - # Remove the unscheduled requests from the persistent batch. - # NOTE(woosuk): The unscheduled requests are either preempted requests - # or running requests that are not scheduled in this step. We remove - # them from the persistent batch but keep their cached states since - # they will be scheduled again sometime in the future. - 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 - # NOTE(zhuohan): cached_req_ids and resumed_req_ids are usually disjoint, - # so `(scheduled_req_ids - resumed_req_ids) == scheduled_req_ids` holds - # apart from the forced-preemption case in reset_prefix_cache. And in - # that case we include the resumed_req_ids in the unscheduled set so - # that they get cleared from the persistent batch before being re-scheduled - # in the normal resumed request path. - unscheduled_req_ids = cached_req_ids - (scheduled_req_ids - resumed_req_ids) - # NOTE(woosuk): The persistent batch optimization assumes that - # consecutive batches contain mostly the same requests. If batches - # have low request overlap (e.g., alternating between two distinct - # sets of requests), this optimization becomes very inefficient. - for req_id in unscheduled_req_ids: - self.input_batch.remove_request(req_id) - - reqs_to_add: list[CachedRequestState] = [] - # Add new requests to the cached states. - for new_req_data in scheduler_output.scheduled_new_reqs: - req_id = new_req_data.req_id - sampling_params = new_req_data.sampling_params - pooling_params = new_req_data.pooling_params - - if sampling_params and sampling_params.sampling_type == SamplingType.RANDOM_SEED: - generator = torch.Generator(device=self.device) - generator.manual_seed(sampling_params.seed) - else: - generator = None - - if self.is_pooling_model: - assert pooling_params is not None - task = pooling_params.task - assert task is not None, "You did not set `task` in the API" - - model = cast(VllmModelForPooling, self.get_model()) - to_update = model.pooler.get_pooling_updates(task) - to_update.apply(pooling_params) - - req_state = CachedRequestState( - req_id=req_id, - prompt_token_ids=new_req_data.prompt_token_ids, - prompt_embeds=new_req_data.prompt_embeds, - mm_features=new_req_data.mm_features, - sampling_params=sampling_params, - pooling_params=pooling_params, - generator=generator, - block_ids=new_req_data.block_ids, - num_computed_tokens=new_req_data.num_computed_tokens, - output_token_ids=[], - lora_request=new_req_data.lora_request, - ) - self.requests[req_id] = req_state - - # Decode and store omni-specific payloads on CPU. - pe_cpu = OmniGPUModelRunner._resolve_prompt_embeds_cpu(new_req_data.prompt_embeds) - if pe_cpu is not None: - setattr(self.requests[req_id], "prompt_embeds_cpu", pe_cpu) - - info_dict = OmniGPUModelRunner._resolve_additional_information( - getattr(new_req_data, "additional_information", None) - ) - if info_dict: - setattr(self.requests[req_id], "additional_information_cpu", info_dict) - - if sampling_params and sampling_params.prompt_logprobs is not None: - self.num_prompt_logprobs[req_id] = ( - self.input_batch.vocab_size - if sampling_params.prompt_logprobs == -1 - else sampling_params.prompt_logprobs - ) - # Only relevant for models using M-RoPE (e.g, Qwen2-VL) - if self.uses_mrope: - self._init_mrope_positions(req_state) - - # Only relevant for models using XD-RoPE (e.g, HunYuan-VL) - if self.uses_xdrope_dim > 0: - self._init_xdrope_positions(req_state) - - reqs_to_add.append(self.requests[req_id]) - - # Update the states of the running/resumed requests. - is_last_rank = get_pp_group().is_last_rank - req_data = scheduler_output.scheduled_cached_reqs - scheduled_spec_tokens = scheduler_output.scheduled_spec_decode_tokens - - # Wait until valid_sampled_tokens_count is copied to cpu, - # then use it to update actual num_computed_tokens of each request. - valid_sampled_token_count = self._get_valid_sampled_token_count() - - for i, req_id in enumerate(req_data.req_ids): - req_state = self.requests[req_id] - num_computed_tokens = req_data.num_computed_tokens[i] - new_block_ids = req_data.new_block_ids[i] - resumed_from_preemption = req_id in req_data.resumed_req_ids - num_output_tokens = req_data.num_output_tokens[i] - req_index = self.input_batch.req_id_to_index.get(req_id) - - if req_state.prev_num_draft_len and self.use_async_scheduling: - # prev_num_draft_len is used in async scheduling mode with - # spec decode. it indicates if need to update num_computed_tokens - # of the request. for example: - # fist step: num_computed_tokens = 0, spec_tokens = [], - # prev_num_draft_len = 0. - # second step: num_computed_tokens = 100(prompt length), - # spec_tokens = [a,b], prev_num_draft_len = 0. - # third step: num_computed_tokens = 100 + 2, spec_tokens = [c,d], - # prev_num_draft_len = 2. - # num_computed_tokens in first step and second step does't contain - # the spec tokens length, but in third step it contains the - # spec tokens length. we only need to update num_computed_tokens - # when prev_num_draft_len > 0. - if req_index is None: - req_state.prev_num_draft_len = 0 - else: - assert self.input_batch.prev_req_id_to_index is not None - prev_req_index = self.input_batch.prev_req_id_to_index[req_id] - num_accepted = valid_sampled_token_count[prev_req_index] - 1 - num_rejected = req_state.prev_num_draft_len - num_accepted - num_computed_tokens -= num_rejected - req_state.output_token_ids.extend([-1] * num_accepted) - - # Update the cached states. - req_state.num_computed_tokens = num_computed_tokens - - if not is_last_rank: - # When using PP, the scheduler sends the sampled tokens back, - # because there's no direct communication between the first- - # stage worker and the last-stage worker. - new_token_ids = req_data.new_token_ids[i] - # Add the sampled token(s) from the previous step (if any). - # This doesn't include "unverified" tokens like spec tokens. - num_new_tokens = num_computed_tokens + len(new_token_ids) - req_state.num_tokens - if num_new_tokens == 1: - # Avoid slicing list in most common case. - req_state.output_token_ids.append(new_token_ids[-1]) - elif num_new_tokens > 0: - req_state.output_token_ids.extend(new_token_ids[-num_new_tokens:]) - elif num_output_tokens < len(req_state.output_token_ids): - # Some output tokens were discarded due to a sync-KV-load - # failure. Align the cached state. - del req_state.output_token_ids[num_output_tokens:] - if req_index is not None: - end_idx = self.input_batch.num_prompt_tokens[req_index] + num_output_tokens - self.input_batch.num_tokens_no_spec[req_index] = end_idx - - # Update the block IDs. - if not resumed_from_preemption: - if new_block_ids is not None: - # Append the new blocks to the existing block IDs. - for block_ids, new_ids in zip(req_state.block_ids, new_block_ids): - block_ids.extend(new_ids) - else: - assert req_index is None - assert new_block_ids is not None - # The request is resumed from preemption. - # Replace the existing block IDs with the new ones. - req_state.block_ids = new_block_ids - - req_index = self.input_batch.req_id_to_index.get(req_id) - if req_index is None: - # The request is not in the persistent batch. - # The request was either preempted and resumed later, or was not - # scheduled in the previous step and needs to be added again. - - if self.use_async_scheduling and num_output_tokens > 0: - # We must recover the output token ids for resumed requests in the - # async scheduling case, so that correct input_ids are obtained. - resumed_token_ids = req_data.all_token_ids[req_id] - req_state.output_token_ids = resumed_token_ids[-num_output_tokens:] - - reqs_to_add.append(req_state) - continue - - # Update the persistent batch. - self.input_batch.num_computed_tokens_cpu[req_index] = num_computed_tokens - if new_block_ids is not None: - self.input_batch.block_table.append_row(new_block_ids, req_index) - - # For the last rank, we don't need to update the token_ids_cpu - # because the sampled tokens are already cached. - if not is_last_rank: - # Add new_token_ids to token_ids_cpu. - start_token_index = num_computed_tokens - end_token_index = num_computed_tokens + len(new_token_ids) - self.input_batch.token_ids_cpu[req_index, start_token_index:end_token_index] = new_token_ids - self.input_batch.num_tokens_no_spec[req_index] = end_token_index - - # Add spec_token_ids to token_ids_cpu. - self.input_batch.update_req_spec_token_ids(req_state, scheduled_spec_tokens) - - # Add the new or resumed requests to the persistent batch. - # The smaller empty indices are filled first. - for request in reqs_to_add: - self.input_batch.add_request(request) - self.input_batch.update_req_spec_token_ids(request, scheduled_spec_tokens) - - # Condense the batched states if there are gaps left by removed requests - self.input_batch.condense() - # Allow attention backend to reorder the batch, potentially - self._may_reorder_batch(scheduler_output) - # Refresh batch metadata with any pending updates. - self.input_batch.refresh_metadata() - - @torch.inference_mode() - def extract_multimodal_outputs(self, hidden_states: torch.Tensor | list[torch.Tensor] | OmniOutput) -> dict: - if ( - hasattr(self.model, "have_multimodal_outputs") - and self.model.have_multimodal_outputs - and isinstance(hidden_states, OmniOutput) - ): - text_hidden_states = hidden_states.text_hidden_states - multimodal_outputs = hidden_states.multimodal_outputs - - elif isinstance(hidden_states, torch.Tensor): - text_hidden_states = hidden_states - multimodal_outputs = {} - elif isinstance(hidden_states, list) or isinstance(hidden_states, tuple): - text_hidden_states = hidden_states[0] - multimodal_outputs = {} - else: - raise ValueError(f"Invalid hidden states type: {type(hidden_states)}") - return text_hidden_states, multimodal_outputs - @torch.inference_mode() def _dummy_run( self, @@ -374,38 +60,16 @@ class OmniNPUModelRunner(NPUModelRunner): 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. @@ -420,78 +84,126 @@ class OmniNPUModelRunner(NPUModelRunner): # 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: @@ -511,20 +223,22 @@ class OmniNPUModelRunner(NPUModelRunner): 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() @@ -548,28 +262,41 @@ class OmniNPUModelRunner(NPUModelRunner): 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, ): + # ---------------------------------------Omni-new---------------------------------------------- if getattr(self.model, "talker", None) is not None and hasattr(self.model, "talker_mtp"): num_tokens_padded_talker_mtp = num_tokens_padded if num_tokens_padded_talker_mtp == self.max_num_tokens: num_tokens_padded_talker_mtp = self.talker_mtp_input_ids.gpu.shape[0] - hidden_states = self.talker_mtp( + outputs = self.talker_mtp( self.talker_mtp_input_ids.gpu[:num_tokens_padded_talker_mtp], self.talker_mtp_inputs_embeds.gpu[:num_tokens_padded_talker_mtp], self.last_talker_hidden.gpu[:num_tokens_padded_talker_mtp], self.text_step.gpu[:num_tokens_padded_talker_mtp], ) self.compilation_config.cache_dir = None - hidden_states = self._generate_dummy_run_hidden_states( - input_ids, positions, num_tokens_padded, intermediate_tensors, inputs_embeds + # ---------------------------------------Omni-new---------------------------------------------- + # NOTE: Directly call self.model() instead of self._model_forward() to match + # GPU behavior. _model_forward contains Omni-specific logic (make_omni_output) + # that requires valid runtime_additional_information, which is empty during dummy run. + outputs = self.model( + input_ids=input_ids, + positions=positions, + intermediate_tensors=intermediate_tensors, + inputs_embeds=inputs_embeds, ) - dummy_compute_logits(hidden_states) - - hidden_states, _ = self.extract_multimodal_outputs(hidden_states) + if self.use_aux_hidden_state_outputs: + hidden_states, _ = outputs + else: + hidden_states = outputs + # ---------------------------------------Omni-new---------------------------------------------- + hidden_states, multimodal_outputs = self.extract_multimodal_outputs(hidden_states) + # ---------------------------------------Omni-new---------------------------------------------- + dummy_compute_logits(hidden_states) if self.drafter: self.drafter.dummy_run( @@ -578,7 +305,7 @@ class OmniNPUModelRunner(NPUModelRunner): 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, @@ -590,292 +317,74 @@ class OmniNPUModelRunner(NPUModelRunner): self.eplb_updator.forward_end() return hidden_states, hidden_states - def _decode_and_store_request_payloads( + def _model_forward( self, - scheduler_output: "SchedulerOutput", - ) -> None: - """Decode per-request prompt_embeds and additional_information for - newly scheduled requests and store them on CPU in the request state. - """ - new_reqs = getattr(scheduler_output, "scheduled_new_reqs", []) - for nr in new_reqs: - req_id = getattr(nr, "req_id", None) or getattr(nr, "request_id", None) - if req_id is None or req_id not in self.requests: - continue - pe_cpu = OmniGPUModelRunner._resolve_prompt_embeds_cpu(getattr(nr, "prompt_embeds", None)) - if pe_cpu is not None: - setattr(self.requests[req_id], "prompt_embeds_cpu", pe_cpu) - info_dict = OmniGPUModelRunner._resolve_additional_information(getattr(nr, "additional_information", None)) - if info_dict: - setattr(self.requests[req_id], "additional_information_cpu", info_dict) - - def _gather_runtime_additional_information(self) -> list[dict]: - """Gather per-request additional_information stored in request state in batch order.""" - per_req_runtime_info = [] - for req_id in self.input_batch.req_ids: - req_state = self.requests.get(req_id) - info = getattr(req_state, "additional_information_cpu", None) if req_state is not None else None - if info and isinstance(info, dict): - per_req_runtime_info.append(info) - if "thinker_reply_part_per_request" in info: - q = info["thinker_reply_part_per_request"] - if hasattr(q, "shape"): - logger.debug(f"[OMNI] req={req_id} has thinker_reply_part_per_request queue shape: {q.shape}") - else: - per_req_runtime_info.append({}) - return per_req_runtime_info - - def _compute_request_token_spans(self, num_scheduled_tokens_np) -> list[tuple[int, int]]: - """Compute (start, end) token spans for each request within the flattened step sequence.""" - req_token_spans: list[tuple[int, int]] = [] - for req_index in range(len(self.input_batch.req_ids)): - start_offset = int(self.query_start_loc.cpu[req_index]) - sched_tokens = int(num_scheduled_tokens_np[req_index]) - req_token_spans.append((start_offset, start_offset + sched_tokens)) - return req_token_spans - - def _build_model_kwargs_extra(self) -> dict: - """Build extra keyword arguments passed to the model for this step, including: - - runtime_additional_information: per-request additional information stored in request state - """ - model_kwargs_extra: dict[str, object] = {} - try: - model_kwargs_extra["runtime_additional_information"] = self._gather_runtime_additional_information() - except Exception as e: - logger.error(f"[OMNI DEBUG] Error building model_kwargs_extra: {e}") - import traceback - - traceback.print_exc() - return model_kwargs_extra - - def _process_additional_information_updates( - self, - hidden_states: torch.Tensor, - multimodal_outputs: object, - num_scheduled_tokens_np: np.ndarray, - scheduler_output: "SchedulerOutput", - ) -> None: - """Process model-provided per-request additional_information updates and merge into request state.""" - try: - # execute the custom postprocess function - # TODO(Peiqi): do we have a more elegant way to do this? - if hasattr(self.model, "has_postprocess") and self.model.has_postprocess: - for req_index, req_id in enumerate(self.input_batch.req_ids): - req_state = self.requests.get(req_id) - req_infos = ( - getattr(req_state, "additional_information_cpu", None) if req_state is not None else None - ) - start_offset = int(self.query_start_loc.cpu[req_index]) - sched_tokens = int(num_scheduled_tokens_np[req_index]) - s, e = start_offset, start_offset + sched_tokens - # only consider to store data into update dict. - hidden_states_slice = hidden_states[s:e] - update_dict = self.model.postprocess(hidden_states_slice, **req_infos) - self._merge_additional_information_update(req_id, update_dict) - except Exception as e: - logger.error( - f"Error merging for requests:{self.input_batch.req_ids} " - f"additional information update: {e}, with the multimodal_outputs " - f"as {multimodal_outputs}" - ) - import traceback - - traceback.print_exc() - - def _collect_additional_information_for_prefill( - self, - num_scheduled_tokens_np: np.ndarray, - ) -> dict[str, dict]: - """Overlay per-request prompt_embeds for the prefill portion and collect - additional_information slices for this step. Returns a map req_id -> dict.""" - for req_index, req_id in enumerate(self.input_batch.req_ids): - req_state = self.requests[req_id] - pe_cpu = getattr(req_state, "prompt_embeds_cpu", None) - num_computed_tokens = int(self.input_batch.num_computed_tokens_cpu[req_index]) - prompt_len = len(req_state.prompt_token_ids) - prompt_remaining = max(0, prompt_len - num_computed_tokens) - sched_tokens = int(num_scheduled_tokens_np[req_index]) - overlay_len = min(sched_tokens, prompt_remaining) - if overlay_len <= 0: - continue - if overlay_len > 0 and pe_cpu is not None: - src = pe_cpu[num_computed_tokens : num_computed_tokens + overlay_len].to( - dtype=self.dtype, device=self.device, non_blocking=True - ) - start_offset = int(self.query_start_loc.cpu[req_index]) - self.inputs_embeds[start_offset : start_offset + overlay_len].copy_(src) - - 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 hasattr(scheduler_output.scheduled_cached_reqs, "additional_information"): - cached_infos = getattr(scheduler_output.scheduled_cached_reqs, "additional_information", {}) - if isinstance(cached_infos, dict): - for req_id, req_infos in cached_infos.items(): - self._merge_additional_information_update(req_id, req_infos) - - def _preprocess( - self, - scheduler_output: "SchedulerOutput", - num_input_tokens: int, + num_tokens_padded: int, + input_ids: torch.Tensor | None = None, + positions: torch.Tensor | None = None, intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + **model_kwargs: dict[str, Any], ): - """Align with v0.14.0 preprocess and omni's additional information handling.""" - num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens - is_first_rank = get_pp_group().is_first_rank - is_encoder_decoder = self.model_config.is_encoder_decoder + """Override to combine NPUModelRunner's signature with OmniGPUModelRunner's logic. - # _prepare_inputs may reorder the batch, so we must gather multi - # modal outputs after that to ensure the correct order - ec_connector_output = None + This method: + 1. Accepts num_tokens_padded as required by NPUModelRunner + 2. Injects omni-specific kwargs (runtime_additional_information) + 3. Caches model output for multimodal results + 4. Handles NPU-specific post-forward logic (graph params update, SP all-gather) + """ + # Omni-specific: build and inject extra model kwargs + model_kwargs_extra = self._build_model_kwargs_extra() - if self.supports_mm_inputs and is_first_rank and not is_encoder_decoder: - # Run the multimodal encoder if any. - with self.maybe_get_ec_connector_output( - scheduler_output, - encoder_cache=self.encoder_cache, - ) as ec_connector_output: - self._execute_mm_encoder(scheduler_output) - mm_embeds, is_mm_embed = self._gather_mm_embeddings(scheduler_output) + runtime_info = model_kwargs_extra.get("runtime_additional_information", []) + if runtime_info: + for i, info in enumerate(runtime_info): + if info: + logger.debug(f"[OMNI] req[{i}] runtime_additional_information keys: {list(info.keys())}") - # NOTE(woosuk): To unify token ids and soft tokens (vision - # embeddings), we always use embeddings (rather than token ids) - # as input to the multimodal model, even when the input is text. - inputs_embeds_scheduled = self.model.embed_input_ids( - self.input_ids.gpu[:num_scheduled_tokens], - multimodal_embeddings=mm_embeds, - is_multimodal=is_mm_embed, + # Call the model forward (same as NPUModelRunner) + assert self.model is not None + model_output = self.model( + input_ids=input_ids, + positions=positions, + intermediate_tensors=intermediate_tensors, + inputs_embeds=inputs_embeds, + **model_kwargs, + **model_kwargs_extra, + ) + + # Omni-specific: wrap output if needed + if not isinstance(model_output, OmniOutput) and hasattr(self.model, "make_omni_output"): + model_output = self.model.make_omni_output(model_output, **model_kwargs_extra) + + # Omni-specific: cache model output for later sample_tokens + self._omni_last_model_output = model_output + + # NPU-specific: update full graph params + forward_context = get_forward_context() + assert forward_context is not None + if ( + forward_context.cudagraph_runtime_mode == CUDAGraphMode.FULL + and not forward_context.capturing + and not self.use_sparse + ): + assert positions is not None + update_full_graph_params( + self.attn_backend, + self.update_stream, + forward_context, + num_tokens_padded, + self.vllm_config, + self.speculative_config, + positions.shape[0], ) - # TODO(woosuk): Avoid the copy. Optimize. - self.inputs_embeds.gpu[:num_scheduled_tokens].copy_(inputs_embeds_scheduled) + # NPU-specific: all-gather for sequence parallelism + if get_forward_context().sp_enabled and not isinstance(model_output, IntermediateTensors): + model_output = self._all_gather_hidden_states_and_aux(model_output) - input_ids, inputs_embeds = self._prepare_mm_inputs(num_input_tokens) - model_kwargs = { - **self._init_model_kwargs(), - **self._extract_mm_kwargs(scheduler_output), - } - elif self.enable_prompt_embeds and is_first_rank: - # Get the input embeddings for the tokens that are not input embeds, - # then put them into the appropriate positions. - # TODO(qthequartermasterman): Since even when prompt embeds are - # enabled, (a) not all requests will use prompt embeds, and (b) - # after the initial prompt is processed, the rest of the generated - # tokens will be token ids, it is not desirable to have the - # embedding layer outside of the CUDA graph all the time. The v0 - # engine avoids this by "double compiling" the CUDA graph, once - # with input_ids and again with inputs_embeds, for all num_tokens. - # If a batch only has token ids, then including the embedding layer - # in the CUDA graph will be more performant (like in the else case - # below). - token_ids_idx = self.is_token_ids.gpu[:num_scheduled_tokens].nonzero(as_tuple=False).squeeze(1) - # Some tokens ids may need to become embeds - if token_ids_idx.numel() > 0: - token_ids = self.input_ids.gpu[token_ids_idx] - tokens_to_embeds = self.model.embed_input_ids(input_ids=token_ids) - self.inputs_embeds.gpu[token_ids_idx] = tokens_to_embeds - - inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens] - model_kwargs = self._init_model_kwargs() - input_ids = self.input_ids.gpu[:num_input_tokens] - else: - # For text-only models, we use token ids as input. - # While it is possible to use embeddings as input just like the - # multimodal models, it is not desirable for performance since - # then the embedding layer is not included in the CUDA graph. - input_ids = self.input_ids.gpu[:num_input_tokens] - inputs_embeds = None - model_kwargs = self._init_model_kwargs() - - if self.uses_mrope: - positions = self.mrope_positions.gpu[:, :num_input_tokens] - elif self.uses_xdrope_dim > 0: - positions = self.xdrope_positions.gpu[:, :num_input_tokens] - else: - positions = self.positions.gpu[:num_input_tokens] - - if is_first_rank: - intermediate_tensors = None - else: - assert intermediate_tensors is not None - intermediate_tensors = self.sync_and_slice_intermediate_tensors( - num_input_tokens, intermediate_tensors, True - ) - - if is_encoder_decoder and scheduler_output.scheduled_encoder_inputs: - # Run the encoder, just like we do with other multimodal inputs. - # For an encoder-decoder model, our processing here is a bit - # simpler, because the outputs are just passed to the decoder. - # We are not doing any prompt replacement. We also will only - # ever have a single encoder input. - encoder_outputs = self._execute_mm_encoder(scheduler_output) - model_kwargs.update({"encoder_outputs": encoder_outputs}) - - req_ids = self.input_batch.req_ids - num_scheduled_tokens_np = np.array( - [scheduler_output.num_scheduled_tokens[rid] for rid in req_ids], - dtype=np.int32, - ) - self._omni_num_scheduled_tokens_np = num_scheduled_tokens_np - - # Note: only prefill need collect additional_information for now. - # Decode don't need per_req_additional_information anymore. - if inputs_embeds is not None: - # Prefill: overlay prompt_embeds and collect additional_information - self._collect_additional_information_for_prefill(num_scheduled_tokens_np) - - if hasattr(self.model, "has_preprocess") and self.model.has_preprocess: - # Overlay custom prompt_embeds per request for the prompt portion; - # collect additional_information (tensor/list) for prefill portion only - decode_req_ids = [] - if self.vllm_config.model_config.async_chunk: - self._update_additional_information(scheduler_output) - for req_index, req_id in enumerate(self.input_batch.req_ids): - # Try to get additional_information from multiple sources - req_state = self.requests.get(req_id) - req_infos = getattr(req_state, "additional_information_cpu", None) if req_state is not None else None - start_offset = int(self.query_start_loc.cpu[req_index]) - sched_tokens = int(num_scheduled_tokens_np[req_index]) - s, e = start_offset, start_offset + sched_tokens - span_len = int(e) - int(s) - - # call the custom process function - req_input_ids, req_embeds, update_dict = self.model.preprocess( - input_ids=input_ids[s:e], input_embeds=inputs_embeds[s:e], **req_infos - ) - - if hasattr(self.model, "talker_mtp") and span_len == 1: - last_talker_hidden, text_step = update_dict.pop("mtp_inputs") - decode_slice = slice(len(decode_req_ids), len(decode_req_ids) + 1) - self.talker_mtp_input_ids.gpu[decode_slice].copy_(req_input_ids) - self.talker_mtp_inputs_embeds.gpu[decode_slice].copy_(req_embeds) - self.last_talker_hidden.gpu[decode_slice].copy_(last_talker_hidden) - self.text_step.gpu[decode_slice].copy_(text_step) - decode_req_ids.append(req_id) - - # TODO(Peiqi): the merge stage could move out from the critical path - self._merge_additional_information_update(req_id, update_dict) - - # update the inputs_embeds and input_ids - seg_len = min(span_len, req_embeds.shape[0]) - inputs_embeds[s : s + seg_len] = req_embeds[:seg_len] - if isinstance(req_input_ids, torch.Tensor) and req_input_ids.numel() == seg_len: - input_ids[s : s + seg_len] = req_input_ids - - # run talker mtp decode - if hasattr(self.model, "talker_mtp"): - self._talker_mtp_forward(decode_req_ids, inputs_embeds) - - return ( - input_ids, - inputs_embeds, - positions, - intermediate_tensors, - model_kwargs, - ec_connector_output, - ) + return model_output def _talker_mtp_forward(self, decode_req_ids: list[str], inputs_embeds: torch.Tensor) -> None: decode_batch_size = len(decode_req_ids) @@ -888,10 +397,11 @@ class OmniNPUModelRunner(NPUModelRunner): max_num_scheduled_tokens=1, use_cascade_attn=False, ) - req_input_ids = self.talker_mtp_input_ids.gpu[:decode_batch_size] - req_embeds = self.talker_mtp_inputs_embeds.gpu[:decode_batch_size] - last_talker_hidden = self.last_talker_hidden.gpu[:decode_batch_size] - text_step = self.text_step.gpu[:decode_batch_size] + num_tokens_padded = batch_desc.num_tokens + req_input_ids = self.talker_mtp_input_ids.gpu[:num_tokens_padded] + req_embeds = self.talker_mtp_inputs_embeds.gpu[:num_tokens_padded] + last_talker_hidden = self.last_talker_hidden.gpu[:num_tokens_padded] + text_step = self.text_step.gpu[:num_tokens_padded] with set_ascend_forward_context( None, self.vllm_config, aclgraph_runtime_mode=_cudagraph_mode, batch_descriptor=batch_desc ): @@ -904,53 +414,3 @@ class OmniNPUModelRunner(NPUModelRunner): inputs_embeds[start_offset : start_offset + 1] = req_embeds[idx : idx + 1] update_dict = {"code_predictor_codes": code_predictor_codes_cpu[idx : idx + 1]} self._merge_additional_information_update(req_id, update_dict) - - def _model_forward( - self, - input_ids: torch.Tensor | None = None, - positions: torch.Tensor | None = None, - intermediate_tensors: IntermediateTensors | None = None, - inputs_embeds: torch.Tensor | None = None, - **model_kwargs: dict[str, Any], - ): - """Inject omni-specific kwargs into forward and cache model output""" - model_kwargs_extra = self._build_model_kwargs_extra() - - runtime_info = model_kwargs_extra.get("runtime_additional_information", []) - if runtime_info: - for i, info in enumerate(runtime_info): - if info: - logger.debug(f"[OMNI] req[{i}] runtime_additional_information keys: {list(info.keys())}") - - model_output = super()._model_forward( - input_ids=input_ids, - positions=positions, - intermediate_tensors=intermediate_tensors, - inputs_embeds=inputs_embeds, - **model_kwargs, - **model_kwargs_extra, - ) - if not isinstance(model_output, OmniOutput) and hasattr(self.model, "make_omni_output"): - model_output = self.model.make_omni_output(model_output, **model_kwargs_extra) - # Cache model output so later sample_tokens can consume multimodal results. - self._omni_last_model_output = model_output - return model_output - - def _merge_additional_information_update(self, req_id: str, upd: dict) -> None: - req_state = self.requests.get(req_id) - if req_state is None: - return - existing = getattr(req_state, "additional_information_cpu", {}) - if not isinstance(existing, dict): - existing = {} - merged = dict(existing) - for k, v in upd.items(): - if isinstance(v, torch.Tensor): - merged[k] = v.detach().to("cpu").contiguous() - elif isinstance(v, list): - merged[k] = [ - (item.detach().to("cpu").contiguous() if isinstance(item, torch.Tensor) else item) for item in v - ] - else: - merged[k] = v - setattr(req_state, "additional_information_cpu", merged) diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index 76cd63654..5d23892c4 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -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", {})