Skip to content

[Perf][Qwen3-Omni] Capture the thinker image encoder with CUDA graphs (part of #7452) - #8215

Merged
amy-why-3459 merged 4 commits into
vllm-project:mainfrom
MrlixiangWE:qwen3-omni-image-encoder-cg
Oct 7, 2026
Merged

amy-why-3459 merged 4 commits into
vllm-project:mainfrom
MrlixiangWE:qwen3-omni-image-encoder-cg

Conversation

@MrlixiangWE

@MrlixiangWE MrlixiangWE commented Sep 27, 2026 •

Copy link
Copy Markdown
Contributor

Purpose

Part of #7452: capture and replay the Qwen3-Omni thinker's image encoder, including DeepStack outputs, through vLLM's encoder CUDA graph protocol. Enable it with cudagraph_mm_encoder on stage 0; it defaults to off.

Propose

  1. Implement SupportsEncoderCudaGraph for images and expose it through the thinker stage of the three-stage wrapper. Audio and video retain eager encoding.
  2. Prepare position embeddings, RoPE and attention metadata from the host image grid before replay. Capture only tensor operations, with full-size token-budget buffers and repeated terminal offsets in unused cu_seqlens entries.
  3. Enable the protocol on the first pipeline rank when images are allowed, precomputed embeddings are disabled, and the encoder uses supported non-FP8 FLASH_ATTN, ROCM_AITER_FA or TRITON_ATTN attention.
  4. Install SingleReplayEncoderCudaGraphManager for models requesting it. Replay a group only when all its images fit one captured graph; otherwise send the whole group through one eager encoding call.
  5. Cap automatically inferred budgets at 256 image tokens. Explicit budgets remain configurable. This bounds padded replay work and startup capture memory for the shipped large token batch.
  6. Document the stage-0 option in the Qwen3-Omni recipe. Merge current main wrapper/runner behavior and use real nested Qwen configuration in wrapper contract tests, including speculative/MTP fields.

Test Plan

vLLM-Omni Commit: 23a6ab2, merged with main b65f97b.

Environment: vLLM 0.30.0, PyTorch 2.13.0+cu130; NVIDIA A800 80GB.

python -m pytest -q -o addopts= --strict-markers --strict-config \
  tests/model_executor/models/qwen3_omni/test_vision_encoder_cudagraph.py \
  -m 'core_model and cpu'
python -m pytest -q -o addopts= --strict-markers --strict-config \
  tests/model_executor/models/qwen3_omni/test_vision_encoder_cudagraph_cuda.py \
  -m 'core_model and cuda' --run-level core_model
python -m pytest -q -o addopts= --strict-markers --strict-config \
  tests/model_executor -m 'core_model and cpu'
python -m pytest -s -v -o addopts= --strict-markers --strict-config \
  tests/e2e/online_serving/test_qwen3_omni.py \
  -m core_model --run-level core_model
SKIP=actionlint pre-commit run --files <changed files>

CUDA protocol tests use a small real vision transformer with BF16 weights and native graph capture on FLASH_ATTN and TRITON_ATTN. They cover replay versus eager outputs, multiple images, retained outputs and metadata errors, plus a synchronization guard on host metadata preparation.

Additional online validation runs the existing test_thinker_prefix_caching_async_chunk_shared_image_prefix case with stage 0's compilation_config.cudagraph_mm_encoder=true. The online tests use the CI deploy with dummy weights for the full Qwen3-Omni architecture. The encoder-enabled case sends two streamed text requests sharing an image and system prefix but using different user text, and retains the original response and block-aligned partial-prefix-cache assertions. The stage-0 process captures three encoder graphs at 64/128/256 tokens before serving the requests.

The component benchmark loads the full image tower's real Qwen3-Omni-30B-A3B-Instruct weights. It compares this commit's eager and single-replay paths on 1x A800 80GB, BF16, FLASH_ATTN, automatically inferred budgets 64/128/256, up to 4 images per replay, one image per call, 10 warm-ups and 30 measured CUDA-event samples per size. These are image-tower timings.

Test Result

Check Candidate main b65f97b
Encoder CPU contracts 26 passed —
Native CUDA protocol tests 6 passed —
Complete model-executor CPU selector 3618 passed, 4 skipped 3592 passed, 4 skipped
Three-stage online ready selector before audio test-dependency repair 10 passed, 14 skipped, 1 fixture import failure 10 passed, 14 skipped, 1 fixture import failure
Original mixed online case after audio test-dependency repair 1 passed, 2 skipped, 22 deselected 1 passed, 2 skipped, 22 deselected
Encoder enabled, shared-image-prefix online case 1 passed, 24 deselected —
Changed-file hooks, including mypy Passed —

Both initial online failures occur while generating test audio, before the API request, because pyttsx3 is absent from the local test environment. The focused rerun provides pyttsx3 and espeak-ng, keeps the original test assertions, and exercises the five mixed text/audio/video/image requests. The skipped online cases target ROCm and NPU.

Current component performance, median of 30 samples:

Image tokens Eager Single-replay path Dispatch Change
196 16.160 ms 9.370 ms Graph hit -42.02%
280 16.171 ms 16.041 ms Eager fallback -0.80%
400 16.341 ms 16.408 ms Eager fallback +0.41%
560 20.973 ms 21.018 ms Eager fallback +0.21%
784 32.351 ms 32.373 ms Eager fallback +0.07%

This replaces the earlier performance tables measured on development prototypes with measurements tied to the current production source. Serving first-token latency, TP=2 output comparisons and stage memory figures from those prototypes are excluded from these current results.

Implement vLLM's encoder CUDA graph protocol for the Qwen3-Omni thinker's
image tower, including the DeepStack outputs. It is enabled by setting
cudagraph_mm_encoder on the thinker stage; video and audio keep the eager
path. The vision metadata is built from the host grid, so preparing a replay
reads nothing back from the device.

A step's images replay only when one captured graph holds them all;
otherwise the group goes back to the runner for one eager call, which for
this launch-bound tower is faster than several replays or per-image eager
calls in every group measured except five or more small images. The
OmniGPUModelRunner installs that manager for models that set
encoder_cudagraph_single_replay. Auto-inferred budgets stop at 256 image
tokens, beyond which replay no longer beats the eager tower.

Co-authored-by: zijianc2 <157244773+zijianc2@users.noreply.github.com>
Co-authored-by: Yifan Tian <221685243+Sendoh-code@users.noreply.github.com>
Signed-off-by: MrlixiangWE <mrdanaer@gmail.com>
@vllm-omni-review-bot

Copy link
Copy Markdown

This PR appears to belong to: docs/design/module/ar_runtime.md.

Module owners: @Gaohan123 @tzhouam @fake0fan

Routing: @Gaohan123 via module of the changed files, CODEOWNERS; @tzhouam via module of the changed files, CODEOWNERS; @fake0fan via module of the changed files

@MrlixiangWE, please review your own changes and leave a short self-review comment describing what you checked. PRs without author self-review may not be assigned a reviewer.

Please take a look when you have a chance. If you would like an automated review, mention @vllm-omni-review-bot in a comment.

@hsliuustc0106 hsliuustc0106 added high priority high priority issue, needs to be done asap enhancement New feature or request labels Sep 28, 2026

@0z5a 0z5a left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Supplemental cross-GPU-architecture validation on RTX 5090 (Blackwell), complementing the PR's A800 results. I tested the unchanged head with BF16 and FLASH_ATTN. The focused CPU suite passed 26/26, the CUDA graph suite passed 6/6, and two runner integration cases passed for graph replay and whole-group eager fallback. I also instantiated the full Qwen3-Omni vision-tower dimensions (depth 27, hidden size 1152) with random weights, captured budgets [64, 128, 256], and compared eager and graph outputs. All tested output groups passed the existing BF16 closeness check.

Image input Eager median Graph median Speedup
One 49-token image 17.17 ms 5.79 ms 2.97×
Two 49-token images 19.49 ms 6.58 ms 2.96×
One 98-token image 17.29 ms 6.19 ms 2.79×

The measured graph capture reduced free GPU memory by 186 MiB in this configuration. These runs use the real tower architecture with random weights; they supplement cross-architecture coverage and do not represent a full checkpoint serving benchmark. I found no actionable issue in the tested path.

LGTM

@vllm-omni-review-bot

vllm-omni-review-bot commented Sep 28, 2026 •

Copy link
Copy Markdown

Omni ReviewBot: resolved

The CI failure noted on 6a6dc516842a no longer applies.

@MrlixiangWE MrlixiangWE reopened this Sep 28, 2026
@vllm-omni-review-bot

Copy link
Copy Markdown
Omni ReviewBot routing record

Assigned Strict under experiment vllm-omni-strict-5050-20260829.

@vllm-omni-review-bot vllm-omni-review-bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Omni ReviewBot review

PR description

This PR opts the Qwen3-Omni thinker image tower into vLLM’s encoder CUDA-graph path, off by default behind cudagraph_mm_encoder. The vision forward is split so position/attention metadata is built on the host before capture or replay, and a SingleReplayEncoderCudaGraphManager replays a step only when one captured budget holds every image; otherwise the runner runs one eager embed_multimodal call. Stage wrappers expose the protocol only on the thinker instance, auto budgets stop at 256 image tokens, and the recipe documents how to enable it.

Change flow

flowchart TD
  A["[EXISTING] Thinker image encode<br/>embed_multimodal / runner"]:::existing
  B["[CHANGED] Vision forward split<br/>prepare_encoder_metadata"]:::changed
  C["[NEW] Qwen3OmniVisionEncoderCudaGraphMixin<br/>SupportsEncoderCudaGraph"]:::new
  D["[CHANGED] OmniGPUModelRunner<br/>_create_encoder_cudagraph_manager"]:::changed
  E["[NEW] SingleReplayEncoderCudaGraphManager<br/>one-graph or whole-group eager"]:::new
  F["[EXISTING] Cached encoder outputs<br/>optional TTFT win"]:::existing
  A --> B
  B --> C
  C --> D
  D --> E
  E --> F
  classDef existing fill:#e5e7eb,stroke:#6b7280,color:#111827
  classDef changed fill:#fef3c7,stroke:#d97706,color:#451a03,stroke-width:2px
  classDef new fill:#dcfce7,stroke:#16a34a,color:#052e16,stroke-width:2px
  classDef removed fill:#fee2e2,stroke:#dc2626,color:#450a0a,stroke-width:2px
Loading

CI at 6a6dc516842a (2026-09-29T20:02:21.939691+00:00): verification incomplete; required-check status is unknown.

See inline comments below.


🤖 This review was generated by InferMatrix Copilot, an open-source repo-maintenance agent for PR review, CI debugging and issue triage. Try it on your own repo, and ⭐ star it if it helped!

metadata["sequence_lengths"] = sequence_lengths
return metadata

def forward(

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P1] Path-gated Qwen3-Omni e2e guard not reported in Test Result

.buildkite/cuda/test-ready.yml step Omni · Qwen3-Omni Test declares source_file_dependencies: omni_qwen3_omni_function, which includes vllm_omni/model_executor/models/qwen3_omni/ (see .buildkite/common/ci_source_file_dependencies.yml). That step runs pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m 'core_model' --run-level 'core_model'. The PR Test Plan/Result cover tests/model_executor CPU/CUDA selectors matching Simple · Model Executor and the model_executor portion of Engine&Model Executor, plus manual serving timings, but do not report that e2e selector and do not state a deliberate skip. The always-on vision forward split in this file is what default (graph-off) e2e would exercise; please run that command or explicitly document the uncovered guard.

@vllm-omni-review-bot

Copy link
Copy Markdown

Omni ReviewBot: finding feedback

[p1] Path-gated Qwen3-Omni e2e guard not reported in Test Result — vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py:215

.buildkite/cuda/test-ready.yml step Omni · Qwen3-Omni Test declares source_file_dependencies: omni_qwen3_omni_function, which includes vllm_omni/model_executor/models/qwen3_omni/ (see .buildkite/common/ci_source_file_dependencies.yml). That step runs pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m 'core_model' --run-level 'core_model'. The PR Test Plan/Result cover tests/model_executor CPU/CUDA selectors matching Simple · Model Executor and the model_executor portion of Engine&Model Executor, plus manual serving timings, but do not report that e2e selector and do not state a deliberate skip. The always-on vision forward split in this file is what default (graph-off) e2e would exercise; please run that command or explicitly document the uncovered guard.

This finding appears in the bot's COMMENT review, but GitHub could not place it as an inline diff comment. If you are the PR author and disagree, react 👎 to this comment. The disagreement will be shown to the maintainer; it does not approve or merge the PR.

Preserve current wrapper and runner behavior while carrying the image
encoder graph protocol. Use real nested Qwen HF configuration in the
wrapper contract tests, including the current speculative/MTP fields.

Co-authored-by: zijianc2 <157244773+zijianc2@users.noreply.github.com>
Signed-off-by: MrlixiangWE <mrdanaer@gmail.com>

Copy link
Copy Markdown
Contributor Author

@hsliuustc0106 Ready for review.

…raph gate

vllm-project#8461 constructs the thinker without a pipeline group and with a stand-in
vision tower. Check the tower's attention backend before asking for the
pipeline rank, so a thinker without a real tower never reaches the
distributed query. The encoder-graph test config gains the device_config
that vllm-project#8461 now reads.

Signed-off-by: MrlixiangWE <mrdanaer@gmail.com>
Co-authored-by: zijianc2 <157244773+zijianc2@users.noreply.github.com>
Co-authored-by: Yifan Tian <221685243+Sendoh-code@users.noreply.github.com>

@amy-why-3459 amy-why-3459 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Reviewed head: 0ce25b37821b825ddab25b3d0659ff54c28b6ec9.

I checked the thinker-only protocol exposure, eligibility gate, host metadata preparation, padded sequence boundaries, DeepStack output ordering/ownership, and runner dispatch for graph hits, whole-group eager fallback, and audio/video. I found no confirmed functional defect in the inspected and tested paths.

Independent focused verification on this exact head:

python -m pytest -q -o addopts= --strict-markers --strict-config \
  tests/model_executor/models/qwen3_omni/test_vision_encoder_cudagraph.py \
  -m 'core_model and cpu'
CUDA_VISIBLE_DEVICES=2 python -m pytest -q -o addopts= --strict-markers --strict-config \
  tests/model_executor/models/qwen3_omni/test_vision_encoder_cudagraph_cuda.py \
  -m 'core_model and cuda' --run-level core_model

Results: 26 CPU tests passed; 6 CUDA tests passed, including native FLASH_ATTN/TRITON_ATTN capture and replay. Environment: NVIDIA L20X, Python 3.12, vLLM 0.31.0, PyTorch 2.13.0+cu130, Transformers 5.14.1, pytest 9.1.1. These are focused checks, not a full-suite run or full-checkpoint serving validation; my vLLM version also differs from the author's reported 0.30.0 environment.

Validation requests before merge:

  1. Provide complete repository UT evidence on the current head. The PR description identifies 23a6ab2 / base b65f97b; those results do not establish coverage for this merged head. The currently visible head checks cover pre-commit, package builds, DCO and docs, not the complete UT suite. In particular, tests/model_executor -m 'core_model and cpu' is only one of the four CPU Simple Test groups in the current CUDA ready pipeline. Please also provide results for the Diffusion, Engine&Entrypoints, and Other CPU groups, the accelerator UT jobs/shards defined by the repository CI, and required backend-specific UT jobs from the supported CUDA/AMD/NPU matrix. Use the pipeline's actual selectors, shard settings and exclusions rather than treating all tests/ as unit tests. Attach exact commands, environment/dependency versions, passed/failed/skipped/deselected/collection-error counts, and CI links or complete logs tied to this SHA. Explain skips and unavailable jobs; obtain an explicit maintainer disposition for any remaining gap. If failures are pre-existing, attach a same-environment base/head comparison. Existing complete CI on this SHA would satisfy this request without another rerun. Please also refresh the path-gated Qwen3-Omni online ready-selector evidence for this head; the earlier bot concern has been addressed in the description for the older revision and should not be presented as a new code defect.

  2. Extend graph coverage for the enabled configurations. The new CUDA fixture initializes TP=1 and tests FLASH_ATTN/TRITON_ATTN, while the capability gate also enables ROCM_AITER_FA and permits TP>1. Please add or attach TP=2 weights-mode eager/replay comparisons covering multiple images, changing replay budgets, all DeepStack slices, and whole-group eager fallback; and provide native ROCM_AITER_FA capture/replay evidence on supported AMD hardware. This is a coverage request, not a claim that those configurations are broken.

The focused path looks sound, but merge readiness is not established until the full-suite evidence and remaining validation gaps are resolved or explicitly dispositioned by a maintainer.

@amy-why-3459 amy-why-3459 added ready label to trigger buildkite CI merge-test label to trigger buildkite merge test CI labels Oct 6, 2026
@amy-why-3459

Copy link
Copy Markdown
Collaborator

Could you please add a Daily-Omni benchmark at concurrency 10 comparing the Qwen3-Omni image encoder CUDA graph enabled versus disabled?

Please use the Stage-0 switch introduced by this PR, compilation_config.cudagraph_mm_encoder=true/false, and keep the checkpoint, hardware, workload, request seeds, and all other deployment settings identical between runs.

Please report end-to-end TTFT and first-audio latency (p50/p95), throughput, failed requests, and GPU memory usage. Include encoder graph hit/miss or eager-fallback statistics and the input image-token distribution so we can understand how often the default capture budgets are exercised.

Please attach the tested head SHA, exact commands/configurations, warmup and measurement settings, and result/log links.

@MrlixiangWE

MrlixiangWE commented Oct 6, 2026 •

Copy link
Copy Markdown
Contributor Author

Provide complete repository UT evidence on the current head.

The ready and merge pipelines both ran on this head: Buildkite build 16785 on 6124bf0, 35 of 35 jobs passed with no failed test.

Job Hardware Result
Simple · Model Executor L4 3828 passed, 2 skipped, 1431 deselected
Simple · Diffusion L4 7558 passed, 48 skipped, 1244 deselected
Simple · Engine&Entrypoints L4 3826 passed, 6 skipped, 133 deselected, 1 xfailed
Simple · Other L4 4980 passed, 34 skipped, 1163 deselected
Diffusion · FlashAttention / Offloader / Cache / Batch L4 8 / 5 / 8 / 6 passed
Diffusion · Model / Other L4 190 passed, 25 skipped / 45 passed, 3 skipped
Tiny Model Tests · Diffusion (base) L4 39 passed, 6 skipped
Engine&Model Executor, including the 6 encoder-graph CUDA cases L4 1359 passed, 14 skipped
Distributed (ready / merge) L4 / H100 25 / 3 passed
Entrypoints (ready) L4 1 + 2 passed
Entrypoint with H100 / L4 / L4 2-GPU (merge) H100 / L4 15 / 1 + 21 / 1 passed
Diffusion · AR Paged Attention GPU / LoRA (merge) L4 16 / 1 passed
Omni · Qwen3-Omni (ready / merge) H100 11 passed, 14 skipped / 15 passed
Omni · Qwen3-Omni Async-Chunk SHM Leak (merge) H100 1 passed
TTS · Qwen3-TTS CustomVoice / Base L4 4 and 6 / 6 passed

The merge pipeline's Simple groups report the same counts as the ready ones. The skips are platform conditions: the 14 in the Qwen3-Omni job are its ROCm and NPU variants, and the rest need NPU, Hopper, more GPUs or optional packages such as msgpack_numpy.

Pipeline on 6124bf0 Failure
NPU 8513 every job stops at import: the image's vLLM has no MultiModalProcessingResult, which qwen2_5_omni_thinker.py imports on main
AMD 13384 1 moss_tts and 7 qwen3_tts failures (__nv_rintf dropped on ROCm), test_z_image timeout
Intel 9679 2 test_qwen2_5_omni_expansion cases fail at engine start

These pipelines fail the same way on other ready PRs from the same day, for example #7947, #8462 and #8339. Could you confirm them as pre-existing for this PR?

Extend graph coverage for the enabled configurations. ... TP=2 weights-mode eager/replay comparisons covering multiple images, changing replay budgets, all DeepStack slices, and whole-group eager fallback

TP=2 on 2x A800 80GB PCIe at 6124bf0, vLLM 0.31.0, PyTorch 2.13.0+cu130. The test below captures the image graphs under TP=2 weights mode with budgets of 8 and 64 tokens. On both ranks it compares every replay with the eager tower, separately for the merger and the two DeepStack mergers, and checks that the two ranks return identical outputs. Consecutive groups alternate between the two budgets: one image, two images, two images filling budget 8, and one 256-patch image. An oversized image and a group needing two replays then go back to one eager call, and a replay after the fallback still matches.

Per group on FLASH_ATTN, in the order the test runs them; TRITON_ATTN gives the same maxima. Grids are (t, h, w) in patches, four patches per image token. The tolerance is the test's bf16 bound, rtol=1.6e-2 with atol 1.6e-2 of the largest eager value.

Group Image tokens Dispatch Max abs error vs eager: merger / DeepStack 0 / DeepStack 1 atol Rank 0 vs 1
(1, 4, 4) 4 replay, budget 8 7.3e-4 / 9.8e-4 / 9.8e-4 1.8e-3 - 2.1e-3 identical
(1, 8, 8) + (1, 4, 8) 16 + 8 replay, budget 64 0 / 0 / 4.9e-4 2.1e-3 - 2.4e-3 identical
(1, 2, 8) + (1, 4, 4) 4 + 4 replay, budget 8 0 / 0 / 0 2.0e-3 - 2.1e-3 identical
(1, 16, 16) 64 replay, budget 64 0 / 0 / 0 2.2e-3 - 2.6e-3 identical
(1, 16, 32) 128 eager fallback, above every budget - - -
(1, 16, 16) + (1, 4, 4) 64 + 4 eager fallback, needs two replays - - -
(1, 8, 8) + (1, 4, 8) 16 + 8 replay, budget 64 2.4e-4 / 2.4e-4 / 0 2.1e-3 - 2.3e-3 identical

The manager counts 8 graph hits and 3 misses, one per image.

Shifting pos_embeds by one row in the replay buffers fails it at the first group (merger of [[1, 4, 4]]). The test reuses the single-GPU file's tower builder, factored out as _build_encoder(backend, tp_size). The ready pipeline has no two-GPU lane for tests/model_executor, so in CI it would skip; I can add it to this PR if you want it in the tree.

tests/model_executor/models/qwen3_omni/test_vision_encoder_cudagraph_tp_cuda.py
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Captured Qwen3-Omni image encoder graphs replay the eager tower under TP=2."""

import socket

import pytest
import torch
from vllm.v1.attention.backends.registry import AttentionBackendEnum

from tests.helpers.distributed import start_collective_workers
from tests.helpers.mark import hardware_marks
from tests.model_executor.models.qwen3_omni.test_vision_encoder_cudagraph_cuda import (
    BUDGETS,
    _build_encoder,
    _eager,
    _images,
    _tol,
)

pytestmark = [
    pytest.mark.skipif(
        not torch.cuda.is_available() or torch.accelerator.device_count() < 2, reason="Two CUDA devices required"
    ),
    pytest.mark.core_model,
    *hardware_marks(res={"cuda": "L4"}, num_cards=2),
]

# Output columns: the main merger, then one merger per DeepStack index.
SLICES = {"merger": slice(0, 64), "deepstack_0": slice(64, 128), "deepstack_1": slice(128, 192)}


def _assert_matches_eager_on_both_ranks(encoder, manager, grids, seed):
    from vllm.distributed import tensor_model_parallel_all_gather

    inputs = _images(grids, seed)
    expected = _eager(encoder, inputs)
    with torch.inference_mode():
        actual = manager.execute(inputs)
    assert actual is not None, grids
    assert [tuple(x.shape) for x in actual] == [tuple(x.shape) for x in expected]
    for output, reference in zip(actual, expected):
        for name, columns in SLICES.items():
            torch.testing.assert_close(
                output[:, columns],
                reference[:, columns],
                **_tol(reference[:, columns]),
                msg=lambda detail, name=name: f"{name} of {grids}: {detail}",
            )
        # The row-parallel merger all-reduces, so both ranks hold the same rows.
        both = tensor_model_parallel_all_gather(output.unsqueeze(0).contiguous(), dim=0)
        torch.testing.assert_close(both[0], both[1], rtol=0, atol=0)


def _worker(rank, port, backend):
    from vllm.compilation.monitor import set_cudagraph_capturing_enabled
    from vllm.config import VllmConfig, set_current_vllm_config
    from vllm.distributed import parallel_state
    from vllm.distributed.parallel_state import graph_capture
    from vllm.model_executor.models.interfaces import supports_encoder_cudagraph

    from vllm_omni.worker.encoder_cudagraph import SingleReplayEncoderCudaGraphManager

    torch.accelerator.set_device_index(rank)
    # The captured graphs hold NCCL all-reduces; custom all-reduce needs
    # peer-to-peer links that two-card test hosts may not have.
    parallel_state.set_custom_all_reduce(False)
    parallel_state.init_distributed_environment(
        world_size=2, rank=rank, local_rank=rank, distributed_init_method=f"tcp://127.0.0.1:{port}", backend="nccl"
    )
    try:
        with set_current_vllm_config(VllmConfig()):
            parallel_state.initialize_model_parallel(tensor_model_parallel_size=2)
        encoder = _build_encoder(backend, tp_size=2)
        assert encoder.visual.tp_size == 2
        assert supports_encoder_cudagraph(encoder)
        manager = SingleReplayEncoderCudaGraphManager(
            encoder.vllm_config, torch.device("cuda"), encoder.visual.dtype, encoder
        )
        try:
            with torch.inference_mode(), graph_capture(device=torch.device("cuda")):
                set_cudagraph_capturing_enabled(True)
                try:
                    manager.capture(graph_pool=torch.cuda.graph_pool_handle())
                finally:
                    set_cudagraph_capturing_enabled(False)
            assert manager.get_cumulative_stats()["num_budgets"] == len(BUDGETS)
            # Consecutive groups alternate between the 8- and 64-token graphs:
            # one image, two images, two images filling budget 8, one long image.
            cases = [[[1, 4, 4]], [[1, 8, 8], [1, 4, 8]], [[1, 2, 8], [1, 4, 4]], [[1, 16, 16]]]
            for seed, grids in enumerate(cases):
                _assert_matches_eager_on_both_ranks(encoder, manager, grids, seed)
            stats = manager.get_cumulative_stats()
            assert stats["graph_hits"] == 6 and stats["graph_misses"] == 0, stats

            # Larger than every budget, or more than one replay: the whole group
            # goes back to the runner's single eager call.
            with torch.inference_mode():
                assert manager.execute(_images([[1, 16, 32]], seed=9)) is None
                assert manager.execute(_images([[1, 16, 16], [1, 4, 4]], seed=10)) is None
            assert manager.get_cumulative_stats()["graph_misses"] == 3
            # Graphs still replay correctly after a fallback.
            _assert_matches_eager_on_both_ranks(encoder, manager, [[1, 8, 8], [1, 4, 8]], seed=11)
        finally:
            manager.clear()
    finally:
        parallel_state.destroy_model_parallel()
        parallel_state.destroy_distributed_environment()


@pytest.mark.parametrize("backend", [AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TRITON_ATTN])
def test_tp2_captured_graphs_match_eager_on_every_rank(backend):
    with socket.socket() as sock:
        sock.bind(("127.0.0.1", 0))
        port = sock.getsockname()[1]
    start_collective_workers(_worker, args=(port, backend), nprocs=2, device_kind="cuda")

@MrlixiangWE

Copy link
Copy Markdown
Contributor Author

Could you please add a Daily-Omni benchmark at concurrency 10 comparing the Qwen3-Omni image encoder CUDA graph enabled versus disabled?

Head 6124bf0, vLLM 0.31.0, PyTorch 2.13.0+cu130, Qwen3-Omni-30B-A3B-Instruct on 2x A800 80GB PCIe with the default deploy config qwen3_omni_moe.yaml (stage 0 on GPU 0, stages 1 and 2 on GPU 1). The two arms differ only in stage 0's cudagraph_mm_encoder. Each run starts a fresh server, sends warm-up requests, then measures at concurrency 10 with text and audio output; runs go off, on, on, off.

Daily-Omni, 1197 prompts, 10 warm-ups Graph Req/s TTFT p50 / p95 ms First audio p50 / p95 ms Failed Correct
run 1 off 0.836 8661 / 13790 11686 / 17297 0 867
run 2 on 0.839 8790 / 14035 11843 / 17639 0 867
run 3 on 0.832 8693 / 13724 12057 / 17345 0 867
run 4 off 0.838 8708 / 13845 11877 / 17372 0 867

The two arms are equal within the spread between repeated runs. The Daily-Omni loader sends each clip as one video_url plus its audio, so the image path never executes; the dispatch table below has the counts. The same prompts sent as 1 fps frames (--daily-omni-pack-mode minicpm-interleave) do reach the image graphs:

Daily-Omni frames, 200 prompts, 5 warm-ups Graph Req/s TTFT p50 / p95 ms First audio p50 / p95 ms Failed Correct
run 1 off 0.480 14535 / 20906 20573 / 26930 0 142
run 2 on 0.561 12469 / 17794 17838 / 23505 0 142
run 3 on 0.563 12492 / 18048 17804 / 23772 0 142
run 4 off 0.483 14657 / 20830 20410 / 26852 1 141
change, mean of two runs +16.72% -14.49% / -14.12% -13.03% / -12.10%

The same pairs with text-only output, where the talker stays idle, give +17.60% req/s and -14.89% TTFT p50.

One 448x448 image per request (random-mm), 300 prompts, 10 warm-ups Req/s TTFT p50 ms First audio p50 ms
graph off 1.429 / 1.457 158 / 155 457 / 454
graph on, default budgets 1.446 / 1.447 144 / 143 440 / 439
graph on, budgets 256/512/1024, 8 images per replay 1.465 / 1.459 141 / 140 438 / 429

Image encoder dispatch, counted in one graph-on pass per workload:

Workload Images Image tokens Replayed Eager, whole group
Daily-Omni 0 (1194 video and 1193 audio encoder calls) - - -
Daily-Omni frames 8714 220 for 91.4%, 66-200 for the rest 8647, one image per group 67 in 20 groups
Single image, default budgets 100 196 82 18 in 4 groups
Single image, budgets 256/512/1024 100 196 91 9 in 1 group
GPU memory, MiB Graph off Graph on
GPU 0 after startup, default budgets 72,709 72,791
GPU 0 after startup, budgets 256/512/1024 72,709 72,969
GPU 0 peak, Daily-Omni 77,385 - 77,445 77,351 - 77,401
GPU 0 peak, Daily-Omni frames 76,505 - 76,655 77,809
GPU 1 after startup 53,945 53,945

At concurrency 10 the scheduler sometimes puts several single images into one step, and those groups exceed the 256-token default cap. Raising the budgets replays half of them, for a further -2.09% TTFT p50 and 178 MiB more graph memory. The single-image p95 values vary more between repeated runs than between arms; they are in the attached results with everything else. Peaks are sampled once per second. The failed requests in the frame runs, here and in the text-only pairs (graph off 1 and 1, graph on 1 and 0), are all prompt 161, rejected in stage 0 with Multi-modal cache miss.

pr8215-daily-omni-results.zip has each run's bench.json, the exact server and benchmark commands, GPU memory samples and the encoder call records. The Daily-Omni benchmark command, local paths shortened:

vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
  --stage-overrides '{"0": {"compilation_config": {"cudagraph_mm_encoder": true}}}'   # false for the other arm
vllm bench serve --omni --backend openai-chat-omni --endpoint /v1/chat/completions \
  --model Qwen/Qwen3-Omni-30B-A3B-Instruct --dataset-name daily-omni \
  --daily-omni-qa-json qa.json --daily-omni-video-dir Videos --daily-omni-inline-local-video \
  --daily-omni-input-mode all --num-prompts 1197 --max-concurrency 10 --num-warmups 10 --seed 0 \
  --percentile-metrics ttft,e2el,audio_ttfp --metric-percentiles 50,95 \
  --extra-body '{"modalities":["text","audio"]}'

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

enhancement New feature or request high priority high priority issue, needs to be done asap merge-test label to trigger buildkite merge test CI ready label to trigger buildkite CI

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants