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

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>
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>

recipes/Qwen/Qwen3-Omni.md

@@ -73,8 +73,20 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
     "1": {"max_num_seqs": 4},
     "2": {"max_num_seqs": 4}
   }'
+
+# Replay the thinker's image encoder from CUDA graphs (images only)
+vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
+  --stage-overrides '{"0": {"compilation_config": {"cudagraph_mm_encoder": true}}}'
 ```
 
+With the encoder graphs on, the thinker captures one graph per token budget
+(64 to 256 image tokens and up to 4 images per replay by default). A step's
+images replay together when one captured graph holds them all; otherwise they
+run through the eager encoder in one call. Set `encoder_cudagraph_token_budgets`
+and `encoder_cudagraph_max_vision_items_per_batch` in the same
+`compilation_config` to change them; each captured budget adds graph memory on
+the thinker's GPU.
+
 ### Stage-based launch (one stage per process)
 
 Use three terminals (one per stage). Start with the default commands below, then

tests/model_executor/models/qwen3_omni/test_vision_encoder_cudagraph.py

@@ -0,0 +1,394 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""CPU contracts for Qwen3-Omni stage routing and encoder graph buffers."""
+
+from contextlib import nullcontext
+from inspect import getattr_static
+from types import SimpleNamespace
+
+import pytest
+import torch
+from torch import nn
+from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import Qwen3OmniMoeConfig
+from vllm.model_executor.models.interfaces import supports_encoder_cudagraph
+from vllm.v1.attention.backends.registry import AttentionBackendEnum
+from vllm.v1.worker.encoder_cudagraph import BudgetGraphMetadata, EncoderCudaGraphManager
+from vllm.v1.worker.gpu_model_runner import GPUModelRunner
+
+import vllm_omni.model_executor.models.qwen3_omni.qwen3_omni as outer_module
+import vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker as thinker_module
+import vllm_omni.model_executor.models.qwen3_omni.vision_encoder_cudagraph as graph_module
+from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker import (
+    Qwen3Omni_VisionTransformer,
+    Qwen3OmniMoeThinkerForConditionalGeneration,
+)
+from vllm_omni.worker.encoder_cudagraph import SingleReplayEncoderCudaGraphManager
+from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.fixture(autouse=True)
+def _first_pipeline_rank(monkeypatch):
+    monkeypatch.setattr(graph_module, "get_pp_group", lambda: SimpleNamespace(is_first_rank=True))
+
+
+class _PatchEmbed(nn.Module):
+    patch_size = 2
+    temporal_patch_size = 1
+
+    def __init__(self):
+        super().__init__()
+        self.proj = nn.Linear(4, 4, bias=False)
+        self.proj.in_channels = 1
+        with torch.no_grad():
+            self.proj.weight.copy_(torch.eye(4))
+
+    def forward(self, pixels):
+        return self.proj(pixels)
+
+
+class _Block(nn.Module):
+    def forward(self, x, *, cu_seqlens, rotary_pos_emb_cos, rotary_pos_emb_sin, max_seqlen, sequence_lengths):
+        return x + rotary_pos_emb_cos[:, None, :] + rotary_pos_emb_sin[:, None, :]
+
+
+class _Merger(nn.Module):
+    def forward(self, x):
+        return x.reshape(-1, 4, 4).sum(dim=1)
+
+
+def _vision():
+    model = object.__new__(Qwen3Omni_VisionTransformer)
+    nn.Module.__init__(model)
+    model.hidden_size = 4
+    model.spatial_merge_size = 2
+    model.tp_size = 1
+    model.attn_backend = AttentionBackendEnum.FLASH_ATTN
+    model.apply_vit_abs_pos_embed = True
+    model.deepstack_visual_indexes = [0, 1]
+    model.patch_embed = _PatchEmbed()
+    model.blocks = nn.ModuleList([_Block(), _Block()])
+    model.merger = _Merger()
+    model.merger_list = nn.ModuleList([_Merger(), _Merger()])
+
+    def positions(grid):
+        return torch.cat([torch.arange(int(t * h * w)).float().unsqueeze(1).expand(-1, 4) + int(w) for t, h, w in grid])
+
+    model.fast_pos_embed_interpolate = positions
+    model.rot_pos_emb = lambda grid: (positions(grid) / 10, positions(grid) / 20)
+    return model
+
+
+def _config(stage="thinker", *, embeds=False):
+    mm = SimpleNamespace(
+        enable_mm_embeds=embeds,
+        get_limit_per_prompt=lambda modality: 4,
+        mm_encoder_tp_mode="weights",
+        mm_encoder_attn_dtype=None,
+    )
+    hf = Qwen3OmniMoeConfig(tts_bos_token_id=1, tts_eos_token_id=2, tts_pad_token_id=3)
+    cfg = SimpleNamespace(
+        model_config=SimpleNamespace(hf_config=hf, multimodal_config=mm, model_stage=stage, max_model_len=32),
+        scheduler_config=SimpleNamespace(max_num_batched_tokens=32),
+        compilation_config=SimpleNamespace(
+            cudagraph_mm_encoder=True,
+            encoder_cudagraph_token_budgets=[2, 4],
+            encoder_cudagraph_max_vision_items_per_batch=2,
+            encoder_cudagraph_max_frames_per_batch=None,
+        ),
+        parallel_config=SimpleNamespace(tensor_parallel_size=1),
+        device_config=SimpleNamespace(device="cpu"),
+        speculative_config=None,
+    )
+    cfg.with_hf_config = lambda *args, **kwargs: cfg
+    return cfg
+
+
+def _thinker(cfg=None):
+    model = object.__new__(Qwen3OmniMoeThinkerForConditionalGeneration)
+    nn.Module.__init__(model)
+    model.visual = _vision()
+    model.vllm_config = cfg or _config()
+    model.multimodal_config = model.vllm_config.model_config.multimodal_config
+    model.visual_dim = 4
+    model.multiscale_dim = 8
+    model.make_empty_intermediate_tensors = lambda **kwargs: None
+    # The thinker wrapper reads the first pipeline layer to size captures.
+    model.language_model = SimpleNamespace(model=SimpleNamespace(start_layer=0))
+    model._enable_image_encoder_cudagraph()
+    return model
+
+
+def _combined(monkeypatch, stage, *, embeds=False):
+    cfg = _config(stage, embeds=embeds)
+    thinker = _thinker(cfg)
+    monkeypatch.setattr(outer_module, "init_vllm_registered_model", lambda **kwargs: thinker)
+    monkeypatch.setattr(
+        outer_module, "current_omni_platform", SimpleNamespace(supports_talker_mtp_graph_capture=lambda: False)
+    )
+    cls = outer_module.Qwen3OmniMoeForConditionalGeneration
+    monkeypatch.setattr(cls, "_init_special_tokens_embeddings", lambda self: None)
+    monkeypatch.setattr(cls, "_get_talker_suppressed_tokens", lambda self: [])
+    return cls(vllm_config=cfg), cfg
+
+
+def _runner_factory(model, cfg, supports_mm_inputs=True):
+    runner = object.__new__(OmniGPUModelRunner)
+    runner.compilation_config = cfg.compilation_config
+    runner.supports_mm_inputs = supports_mm_inputs
+    runner.model = model
+    runner.vllm_config = cfg
+    runner.device = torch.device("cpu")
+    runner.dtype = torch.float32
+    return runner._create_encoder_cudagraph_manager()
+
+
+@pytest.mark.parametrize("stage,enabled", [("thinker", True), ("talker", False), ("code2wav", False), (None, True)])
+def test_real_wrapper_construction_controls_runner_protocol(monkeypatch, stage, enabled):
+    model, cfg = _combined(monkeypatch, stage)
+    assert supports_encoder_cudagraph(model) is enabled
+    manager = _runner_factory(model, cfg)
+    assert (manager is not None) is enabled
+    if enabled:
+        assert getattr_static(model, "encoder_cudagraph_forward").__self__ is model.thinker
+        assert type(manager) is SingleReplayEncoderCudaGraphManager
+        assert manager.model is model
+        assert manager.config.out_hidden_size == 12
+        assert manager.supports_modality("image")
+        assert not manager.supports_modality("video")
+        assert not manager.supports_modality("audio")
+
+
+def test_runner_keeps_the_upstream_manager_without_the_opt_in(monkeypatch):
+    model, cfg = _combined(monkeypatch, "thinker")
+    model.encoder_cudagraph_single_replay = False
+    assert type(_runner_factory(model, cfg)) is EncoderCudaGraphManager
+
+
+def test_runner_builds_no_manager_without_multimodal_inputs(monkeypatch):
+    model, cfg = _combined(monkeypatch, "thinker")
+    assert _runner_factory(model, cfg, supports_mm_inputs=False) is None
+
+
+def test_thinker_construction_enables_the_image_graph(monkeypatch):
+    cfg = _config()
+    cfg.quant_config = None
+    thinker_config = cfg.model_config.hf_config.thinker_config
+    thinker_config.architectures = ["Qwen3OmniMoeThinkerForConditionalGeneration"]
+    thinker_config.vision_config.deepstack_visual_indexes = [0, 1]
+    thinker_config.vision_config.out_hidden_size = 4
+    thinker_config.text_config.rms_norm_eps = 1e-6
+    thinker_config.text_config.hidden_size = 4
+    language_model = nn.Module()
+    language_model.make_empty_intermediate_tensors = None
+    monkeypatch.setattr(thinker_module, "Qwen3OmniMoeAudioEncoder", lambda *args, **kwargs: nn.Identity())
+    monkeypatch.setattr(thinker_module, "Qwen3Omni_VisionTransformer", lambda **kwargs: _vision())
+    monkeypatch.setattr(thinker_module, "Qwen3MoeLLMForCausalLM", lambda **kwargs: language_model)
+    cls = Qwen3OmniMoeThinkerForConditionalGeneration
+    monkeypatch.setattr(cls, "_mark_tower_model", lambda self, *args: nullcontext())
+    monkeypatch.setattr(cls, "_mark_language_model", lambda self, *args: nullcontext())
+    assert supports_encoder_cudagraph(cls(vllm_config=cfg))
+
+
+@pytest.mark.parametrize("change", ["no_image_limit", "sdpa_backend", "fp8_attention", "later_pipeline_rank", None])
+def test_image_graph_gate(monkeypatch, change):
+    model = _thinker()
+    del model.supports_encoder_cudagraph
+    mm = model.multimodal_config
+    if change == "no_image_limit":
+        mm.get_limit_per_prompt = lambda modality: 0 if modality == "image" else 4
+    elif change == "sdpa_backend":
+        model.visual.attn_backend = AttentionBackendEnum.TORCH_SDPA
+    elif change == "fp8_attention":
+        mm.mm_encoder_attn_dtype = "fp8"
+    elif change == "later_pipeline_rank":
+        monkeypatch.setattr(graph_module, "get_pp_group", lambda: SimpleNamespace(is_first_rank=False))
+    model._enable_image_encoder_cudagraph()
+    assert supports_encoder_cudagraph(model) is (change is None)
+
+
+def test_precomputed_embeddings_keep_existing_encoder_entry(monkeypatch):
+    model, cfg = _combined(monkeypatch, "thinker", embeds=True)
+    assert not supports_encoder_cudagraph(model.thinker)
+    assert _runner_factory(model, cfg) is None
+
+
+def _images(grids, offset=0):
+    patches = sum(t * h * w for t, h, w in grids)
+    return {
+        "image_grid_thw": torch.tensor(grids),
+        "pixel_values": torch.arange(patches * 4).float().reshape(patches, 4) + offset,
+    }
+
+
+def _cpu_manager(model):
+    manager = SingleReplayEncoderCudaGraphManager(model.vllm_config, torch.device("cpu"), torch.float32, model)
+    # Replace only CUDA launch/capture; keep production packing, selectors,
+    # metadata copy, output slicing and ownership behavior.
+    for budget in manager.token_budgets:
+        values = model.prepare_encoder_cudagraph_capture_inputs(
+            budget, manager.max_batch_size, manager.max_frames_per_batch, torch.device("cpu"), torch.float32
+        ).values
+        output = torch.empty_like(model.encoder_cudagraph_forward(values))
+
+        def replay(values=values, output=output):
+            output.copy_(model.encoder_cudagraph_forward(values))
+
+        manager.budget_graphs.setdefault("default", {})[budget] = BudgetGraphMetadata(
+            token_budget=budget,
+            max_batch_size=manager.max_batch_size,
+            max_frames_per_batch=manager.max_frames_per_batch,
+            graph=SimpleNamespace(replay=replay),
+            input_buffers=values,
+            output_buffer=output,
+        )
+    return manager
+
+
+def test_manager_preserves_item_order_deepstack_and_previous_outputs():
+    model = _thinker()
+    manager = _cpu_manager(model)
+    inputs = _images([[1, 2, 4], [1, 2, 2]])  # packing reorders 2/1-token images
+    expected = model.encoder_eager_forward(inputs).split([2, 1])
+    outputs = manager.execute(inputs)
+    assert manager.graph_hits == 2
+    assert [output.shape for output in outputs] == [(2, 12), (1, 12)]
+    for output, reference in zip(outputs, expected):
+        torch.testing.assert_close(output, reference)
+        # The two captured intermediate levels must remain in their own slots.
+        assert not torch.equal(output[:, 4:8], output[:, 8:12])
+    saved = [output.clone() for output in outputs]
+    changed = _images([[1, 4, 2], [1, 2, 4]], offset=1000)
+    actual = manager.execute(changed)
+    for output, reference in zip(actual, model.encoder_eager_forward(changed).split([2, 2])):
+        torch.testing.assert_close(output, reference)
+    for output, reference in zip(outputs, saved):
+        torch.testing.assert_close(output, reference)
+
+
+@pytest.mark.parametrize(
+    "grids",
+    [
+        [[1, 2, 10]],  # one image above the largest budget
+        [[1, 2, 4], [1, 2, 6]],  # 2 + 3 tokens exceed the largest budget
+        [[1, 2, 2]] * 3,  # more images than one replay holds
+    ],
+)
+def test_group_that_does_not_fit_one_replay_goes_back_to_the_runner(grids):
+    model = _thinker()
+    manager = _cpu_manager(model)
+    assert manager.execute(_images(grids)) is None
+    assert manager.graph_hits == 0 and manager.graph_misses == len(grids)
+
+
+def test_empty_shard_is_valid():
+    model = _thinker()
+    manager = _cpu_manager(model)
+    inputs = _images([[1, 2, 10]])
+    empty = model.select_encoder_cudagraph_items(inputs, [])
+    assert empty["pixel_values"].shape == (0, 4)
+    assert manager.execute(empty) == []
+
+
+def test_budget_range_follows_scheduler_and_model_limits():
+    model = _thinker()
+    cfg = _config()
+    assert model.get_encoder_cudagraph_budget_range(cfg) == (32, 32)
+    cfg.scheduler_config.max_num_batched_tokens = 128
+    cfg.model_config.max_model_len = 4096
+    assert model.get_encoder_cudagraph_budget_range(cfg) == (64, 128)
+    # The shipped deploy's 32768-token batch stops at the replay-profitable cap.
+    cfg.scheduler_config.max_num_batched_tokens = 32768
+    cfg.model_config.max_model_len = 65536
+    assert model.get_encoder_cudagraph_budget_range(cfg) == (64, 256)
+
+
+@pytest.mark.parametrize("grid", [[2, 2, 2], [1, 3, 2], [1, 2, 0]])
+def test_multi_frame_or_unmergeable_image_grid_is_rejected(grid):
+    with pytest.raises(ValueError, match="Invalid Qwen3-Omni image grid"):
+        _thinker().get_encoder_cudagraph_item_specs(_images([grid]))
+
+
+def test_capture_has_full_budget_and_replay_keeps_scalar_bound():
+    model = _thinker()
+    values = model.prepare_encoder_cudagraph_capture_inputs(3, 2, 0, torch.device("cpu"), torch.float32).values
+    assert values["pixel_values"].shape == (12, 4)  # non-divisible budget/max-items
+    replay = model.prepare_encoder_cudagraph_replay_buffers(_images([[1, 2, 2]]), 2, 0)
+    assert "max_seqlen" not in replay.values
+    assert values["max_seqlen"].item() == 12
+    model.get_encoder_cudagraph_config().padding_logics["cu_seqlens"](values["cu_seqlens"], replay.values["cu_seqlens"])
+    assert values["cu_seqlens"].tolist() == [0, 4, 4]
+
+
+def test_graph_entry_preserves_vision_module_hooks():
+    model = _thinker()
+    values = model.prepare_encoder_cudagraph_capture_inputs(2, 2, 0, torch.device("cpu"), torch.float32).values
+    calls = []
+    hook = model.visual.register_forward_pre_hook(lambda module, args: calls.append(args[0]))
+    try:
+        model.encoder_cudagraph_forward(values)
+    finally:
+        hook.remove()
+    assert len(calls) == 1 and calls[0] is values["pixel_values"]
+
+
+def _run_encoder_step(monkeypatch, model, manager, batches):
+    import vllm.v1.worker.gpu_model_runner as runner_module
+
+    monkeypatch.setattr(runner_module, "group_and_batch_mm_kwargs", lambda *args, **kwargs: iter(batches))
+    keys = [f"item{index}" for index in range(sum(count for _, count, _ in batches))]
+    cached: dict[str, torch.Tensor] = {}
+    runner = SimpleNamespace(
+        _batch_mm_inputs_from_scheduler=lambda _: (keys, [(m, {}) for m, _, _ in batches], []),
+        observability_config=None,
+        model=model,
+        lora_config=None,
+        device=torch.device("cpu"),
+        is_multimodal_pruning_enabled=False,
+        requires_sequential_video_encoding=False,
+        timed_encoder_operation=lambda *args: nullcontext(),
+        encoder_cudagraph_manager=manager,
+        _cache_encoder_output=lambda key, value, *args: cached.__setitem__(key, value),
+    )
+    scheduler = SimpleNamespace(ec_manager_metadata=None, free_encoder_mm_hashes=[])
+    GPUModelRunner._execute_mm_encoder(runner, scheduler)
+    return [cached[key] for key in keys]
+
+
+def test_actual_runner_dispatch_keeps_video_and_audio_eager(monkeypatch):
+    model = _thinker()
+    manager = _cpu_manager(model)
+    eager_modalities = []
+
+    def embed_multimodal(**kwargs):
+        eager_modalities.append(kwargs["kind"])
+        return [torch.zeros(1, 12)]
+
+    model.embed_multimodal = embed_multimodal
+    image = _images([[1, 2, 2]])
+    expected = model.encoder_eager_forward(image)
+    batches = [("image", 1, image), ("video", 1, {"kind": "video"}), ("audio", 1, {"kind": "audio"})]
+    cached = _run_encoder_step(monkeypatch, model, manager, batches)
+    assert eager_modalities == ["video", "audio"]
+    assert len(cached) == 3
+    torch.testing.assert_close(cached[0], expected)
+    assert manager.graph_hits == 1
+
+
+def test_runner_encodes_a_group_beyond_one_replay_in_one_eager_call(monkeypatch):
+    model = _thinker()
+    manager = _cpu_manager(model)
+    calls = []
+
+    def embed_multimodal(**kwargs):
+        calls.append(kwargs["image_grid_thw"].tolist())
+        return list(model.encoder_eager_forward(kwargs).split([2, 3]))
+
+    model.embed_multimodal = embed_multimodal
+    images = _images([[1, 2, 4], [1, 2, 6]])
+    cached = _run_encoder_step(monkeypatch, model, manager, [("image", 2, images)])
+    assert calls == [[[1, 2, 4], [1, 2, 6]]]
+    for output, reference in zip(cached, model.encoder_eager_forward(images).split([2, 3])):
+        torch.testing.assert_close(output, reference)
+    assert manager.graph_hits == 0 and manager.graph_misses == 2

tests/model_executor/models/qwen3_omni/test_vision_encoder_cudagraph_cuda.py

@@ -0,0 +1,251 @@
+# 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 on CUDA."""
+
+import socket
+from types import SimpleNamespace
+
+import pytest
+import torch
+from torch import nn
+from vllm.v1.attention.backends.registry import AttentionBackendEnum
+
+from tests.helpers.mark import hardware_marks
+
+pytestmark = [
+    pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"),
+    pytest.mark.core_model,
+    *hardware_marks(res={"cuda": "L4"}),
+]
+
+DTYPE = torch.bfloat16
+# Four input patches per output token; budgets are in output tokens.
+BUDGETS = [8, 64]
+
+
+def _tol(reference):
+    # bf16 relative tolerance, with the absolute floor scaled to the output.
+    return {"rtol": 1.6e-2, "atol": 1.6e-2 * reference.abs().max().item()}
+
+
+def _vision_config():
+    return SimpleNamespace(
+        hidden_size=128,
+        num_heads=2,
+        image_size=128,
+        patch_size=8,
+        spatial_merge_size=2,
+        temporal_patch_size=2,
+        in_channels=3,
+        apply_vit_abs_pos_embed=True,
+        deepstack_visual_indexes=[0, 1],
+        intermediate_size=256,
+        hidden_act="gelu_pytorch_tanh",
+        depth=2,
+        out_hidden_size=64,
+    )
+
+
+def _engine_config():
+    mm = SimpleNamespace(
+        enable_mm_embeds=False,
+        get_limit_per_prompt=lambda modality: 4,
+        mm_encoder_tp_mode="weights",
+        mm_encoder_attn_dtype=None,
+    )
+    return SimpleNamespace(
+        model_config=SimpleNamespace(multimodal_config=mm, max_model_len=256),
+        scheduler_config=SimpleNamespace(max_num_batched_tokens=256),
+        compilation_config=SimpleNamespace(
+            cudagraph_mm_encoder=True,
+            encoder_cudagraph_token_budgets=BUDGETS,
+            encoder_cudagraph_max_vision_items_per_batch=2,
+            encoder_cudagraph_max_frames_per_batch=None,
+        ),
+        parallel_config=SimpleNamespace(tensor_parallel_size=1),
+    )
+
+
+@pytest.fixture(scope="module")
+def _parallel_state():
+    from vllm.config import VllmConfig, set_current_vllm_config
+    from vllm.distributed import parallel_state
+
+    if parallel_state.model_parallel_is_initialized():
+        yield
+        return
+    previous_device = torch.accelerator.current_device_index()
+    with socket.socket() as sock:
+        sock.bind(("127.0.0.1", 0))
+        port = sock.getsockname()[1]
+    torch.accelerator.set_device_index(0)
+    parallel_state.init_distributed_environment(
+        world_size=1, rank=0, local_rank=0, distributed_init_method=f"tcp://127.0.0.1:{port}", backend="nccl"
+    )
+    with set_current_vllm_config(VllmConfig()):
+        parallel_state.initialize_model_parallel(tensor_model_parallel_size=1)
+    try:
+        yield
+    finally:
+        parallel_state.destroy_model_parallel()
+        parallel_state.destroy_distributed_environment()
+        torch.accelerator.set_device_index(previous_device)
+
+
+@pytest.fixture(scope="module", params=[AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TRITON_ATTN])
+def encoder(request, _parallel_state):
+    import vllm.model_executor.models.vision as vision_module
+    from vllm.config import VllmConfig, set_current_vllm_config
+    from vllm.config.multimodal import MultiModalConfig
+    from vllm.utils.torch_utils import set_default_torch_dtype
+
+    from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker import (
+        Qwen3Omni_VisionTransformer,
+        Qwen3OmniMoeThinkerForConditionalGeneration,
+    )
+
+    override = MultiModalConfig(mm_encoder_attn_backend=request.param)
+    with (
+        pytest.MonkeyPatch.context() as patch,
+        set_current_vllm_config(VllmConfig()),
+        set_default_torch_dtype(DTYPE),
+        torch.device("cuda"),
+    ):
+        patch.setattr(vision_module, "get_multimodal_config", lambda: override)
+        visual = Qwen3Omni_VisionTransformer(vision_config=_vision_config())
+    assert visual.attn_backend == request.param
+    generator = torch.Generator().manual_seed(0)
+    with torch.no_grad():
+        for name, param in visual.named_parameters():
+            if "norm" in name:
+                param.fill_(1.0 if name.endswith("weight") else 0.0)
+            elif name.startswith("pos_embed"):
+                # Comparable to the patch embedding, so a misplaced row shows.
+                param.copy_(torch.randn(param.shape, generator=generator))
+            elif param.dim() > 1:
+                fan_in = param[0].numel()
+                param.copy_(torch.randn(param.shape, generator=generator) / fan_in**0.5)
+            else:
+                param.copy_(torch.randn(param.shape, generator=generator) * 0.02)
+    model = object.__new__(Qwen3OmniMoeThinkerForConditionalGeneration)
+    nn.Module.__init__(model)
+    model.visual = visual
+    model.vllm_config = _engine_config()
+    model.multimodal_config = model.vllm_config.model_config.multimodal_config
+    model.visual_dim = 64
+    model.multiscale_dim = 128
+    model._enable_image_encoder_cudagraph()
+    return model
+
+
+def _images(grids, seed):
+    generator = torch.Generator().manual_seed(seed)
+    patches = sum(t * h * w for t, h, w in grids)
+    pixels = torch.randn(patches, 3 * 2 * 8 * 8, generator=generator)
+    # The processor keeps image_grid_thw on the CPU; pixels reach the GPU.
+    return {"image_grid_thw": torch.tensor(grids), "pixel_values": pixels.to("cuda", DTYPE)}
+
+
+def _forward(model, pixels, metadata):
+    from vllm.config import VllmConfig
+    from vllm.forward_context import set_forward_context
+
+    # The tower on caller-supplied metadata, inside a forward context like the thinker.
+    with torch.inference_mode(), set_forward_context(None, VllmConfig()):
+        return model.visual.forward_with_encoder_metadata(pixels, metadata)
+
+
+def _eager(model, inputs):
+    from vllm.config import VllmConfig
+
+    # The thinker's own image path, which the runner calls without graphs.
+    engine_config, model.vllm_config = model.vllm_config, VllmConfig()
+    try:
+        with torch.inference_mode():
+            return list(model.embed_multimodal(**inputs))
+    finally:
+        model.vllm_config = engine_config
+
+
+def _assert_detectable(model, inputs, wrong_metadata):
+    # A replay built from wrong metadata must fall outside the tolerance.
+    right = torch.cat(_eager(model, inputs))
+    wrong = _forward(model, inputs["pixel_values"], wrong_metadata)
+    assert not torch.allclose(wrong, right, **_tol(right))
+
+
+def test_tolerance_detects_wrong_encoder_metadata(encoder):
+    visual = encoder.visual
+    two = _images([[1, 8, 8], [1, 4, 8]], seed=2)
+    grids = two["image_grid_thw"].tolist()
+    metadata = visual.prepare_encoder_metadata(grids)
+    shifted = dict(metadata, pos_embeds=metadata["pos_embeds"].roll(1, dims=0))
+    _assert_detectable(encoder, two, shifted)
+    transposed = visual.prepare_encoder_metadata([[1, 8, 8], [1, 8, 4]])
+    rotated = dict(metadata, rotary_pos_emb_cos=transposed["rotary_pos_emb_cos"])
+    rotated["rotary_pos_emb_sin"] = transposed["rotary_pos_emb_sin"]
+    _assert_detectable(encoder, two, rotated)
+    total = int(metadata["cu_seqlens"][-1])
+    merged = dict(metadata, cu_seqlens=metadata["cu_seqlens"][[0, -1]], max_seqlen=torch.tensor(total))
+    merged.pop("sequence_lengths", None)
+    _assert_detectable(encoder, two, merged)
+
+
+def test_replay_metadata_does_not_read_the_device(encoder):
+    # Replay buffers are built from the host grid every step; a device read
+    # here would stall the CPU behind the queued GPU work.
+    torch.accelerator.synchronize()
+    torch.cuda.set_sync_debug_mode("error")
+    try:
+        encoder.visual.prepare_encoder_metadata([[1, 8, 8], [1, 4, 8]])
+    finally:
+        torch.cuda.set_sync_debug_mode("default")
+
+
+def test_captured_graphs_match_eager_and_keep_previous_outputs(encoder):
+    from vllm.compilation.monitor import set_cudagraph_capturing_enabled
+    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
+
+    assert supports_encoder_cudagraph(encoder)
+    manager = SingleReplayEncoderCudaGraphManager(encoder.vllm_config, torch.device("cuda"), DTYPE, encoder)
+    try:
+        # The runner captures inside graph_capture(), which supplies the
+        # non-default capture stream.
+        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)
+        # A small image; two images that exactly fill the smallest budget; one
+        # 256-patch sequence, longer than one attention tile, filling a budget
+        # alone; two images the manager packs into one replay in reverse order.
+        cases = [[[1, 4, 4]], [[1, 2, 8], [1, 4, 4]], [[1, 16, 16]], [[1, 8, 8], [1, 4, 8]]]
+        retained = []
+        for seed, grids in enumerate(cases):
+            inputs = _images(grids, seed)
+            expected = _eager(encoder, inputs)
+            with torch.inference_mode():
+                actual = manager.execute(inputs)
+            assert [tuple(x.shape) for x in actual] == [tuple(x.shape) for x in expected]
+            for output, reference in zip(actual, expected):
+                torch.testing.assert_close(output, reference, **_tol(reference))
+                retained.append((output, output.clone()))
+        stats = manager.get_cumulative_stats()
+        # Hits and misses count images, not replays.
+        assert stats["graph_hits"] == 6 and stats["graph_misses"] == 0, stats
+        for output, saved in retained:
+            torch.testing.assert_close(output, saved, rtol=0, atol=0)
+
+        # Larger than every budget, or more than one replay holds: the 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
+    finally:
+        manager.clear()

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py

@@ -26,6 +26,7 @@ from vllm.model_executor.models.interfaces import (
     SupportsPP,
     SupportsQuant,
     SupportsRealtime,
+    supports_encoder_cudagraph,
 )
 from vllm.model_executor.models.qwen3_asr_realtime import Qwen3ASRRealtimeBuffer
 from vllm.model_executor.models.qwen3_omni_moe_thinker import (
@@ -90,6 +91,24 @@ TALKER_CODEC_THINK_EOS_ID = 4205  # Think mode end
 
 logger = init_logger(__name__)
 
+# The encoder CUDA graph protocol, forwarded from the thinker to the stage
+# wrapper that the runner inspects.
+_THINKER_ENCODER_CUDAGRAPH_MEMBERS = (
+    "supports_encoder_cudagraph",
+    "encoder_cudagraph_single_replay",
+    "get_encoder_cudagraph_config",
+    "get_input_modality",
+    "get_max_frames_per_video",
+    "get_encoder_cudagraph_budget_range",
+    "get_encoder_cudagraph_item_specs",
+    "select_encoder_cudagraph_items",
+    "prepare_encoder_cudagraph_capture_inputs",
+    "prepare_encoder_cudagraph_replay_buffers",
+    "encoder_cudagraph_forward",
+    "encoder_eager_forward",
+    "postprocess_encoder_output",
+)
+
 
 def _cpu_long(values: Any) -> torch.Tensor:
     """Token ids as a host tensor (payload ids arrive as lists or CPU tensors)."""
@@ -221,6 +240,12 @@ class Qwen3OmniMoeForConditionalGeneration(
             self.model = self.thinker
             self.talker = None
             self.code2wav = None
+            if supports_encoder_cudagraph(self.thinker):
+                # Expose the protocol on this stage instance only. Defining
+                # these methods on the shared class also advertises an encoder
+                # for talker/code2wav under runtime structural protocol checks.
+                for name in _THINKER_ENCODER_CUDAGRAPH_MEMBERS:
+                    setattr(self, name, getattr(self.thinker, name))
             self.tts_tokens = torch.tensor(
                 [[self.config.tts_bos_token_id, self.config.tts_eos_token_id, self.config.tts_pad_token_id]],
                 device=self._module_device(self.thinker),

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py

@@ -117,6 +117,7 @@ from vllm.multimodal.processing.processor import (
 from vllm.multimodal.utils import set_mm_embedding_modality
 from vllm.sequence import IntermediateTensors
 from vllm.transformers_utils.processor import cached_processor_from_config
+from vllm.utils.torch_utils import async_tensor_h2d
 
 from vllm_omni.data_entry_keys import OmniPayload
 from vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker import (
@@ -129,6 +130,7 @@ from vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker import (
 from vllm_omni.model_executor.models.qwen3_omni.quantization import (
     Qwen3OmniNestedSupportsQuant,
 )
+from vllm_omni.model_executor.models.qwen3_omni.vision_encoder_cudagraph import Qwen3OmniVisionEncoderCudaGraphMixin
 from vllm_omni.quantization.component_config import (
     PRE_QUANTIZED_METHODS,
     ComponentQuantizationConfig,
@@ -167,79 +169,74 @@ class Qwen3Omni_VisionTransformer(_Qwen3Omni_VisionTransformer):
         super().__init__(*args, **kwargs)
         self.tp_size = get_tensor_model_parallel_world_size()
 
-    def forward(
-        self,
-        x: torch.Tensor,
-        grid_thw,
-    ) -> torch.Tensor:
-        hidden_states = x.to(device=self.device, dtype=self.dtype)
-        hidden_states = self.patch_embed(hidden_states)
-
+    def prepare_encoder_metadata(self, grid_thw) -> dict[str, torch.Tensor]:
+        """Prepare positions and attention metadata outside CUDA graph capture."""
+        if not isinstance(grid_thw, torch.Tensor):
+            grid_thw = torch.as_tensor(grid_thw, dtype=torch.int32)
+        metadata = {}
         if self.apply_vit_abs_pos_embed:
-            pos_embeds = self.fast_pos_embed_interpolate(grid_thw)
-            hidden_states = hidden_states + pos_embeds
+            metadata["pos_embeds"] = self.fast_pos_embed_interpolate(grid_thw)
         rotary_pos_emb_cos, rotary_pos_emb_sin = self.rot_pos_emb(grid_thw)
 
-        if isinstance(grid_thw, torch.Tensor):
-            grid_thw_tensor = grid_thw.to(self.device)
-        else:
-            grid_thw_tensor = torch.as_tensor(grid_thw, dtype=torch.int32, device=self.device)
+        rotary_pos_emb_cos = rotary_pos_emb_cos.to(self.device)
+        rotary_pos_emb_sin = rotary_pos_emb_sin.to(self.device)
 
-        try:
-            cu_seqlens = torch.repeat_interleave(
-                grid_thw_tensor[:, 1] * grid_thw_tensor[:, 2],
-                grid_thw_tensor[:, 0],
-            ).cumsum(
-                dim=0,
-                dtype=grid_thw_tensor.dtype if torch.jit.is_tracing() else torch.int32,
-            )
-            cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0)
-        except RuntimeError:
-            logger.warning(
-                "torch.repeat_interleave not executable, switching to vectorized searchsorted implementation."
-            )
-            repeat_counts = grid_thw_tensor[:, 0]
-            values = grid_thw_tensor[:, 1] * grid_thw_tensor[:, 2]
-            repeat_cumsum = repeat_counts.cumsum(0)
-            total_items = repeat_cumsum[-1].item()
-            indices = torch.searchsorted(
-                repeat_cumsum,
-                torch.arange(total_items, device=grid_thw_tensor.device),
-                right=True,
-            )
-            cu_seqlens = values[indices].cumsum(
-                dim=0,
-                dtype=grid_thw_tensor.dtype if torch.jit.is_tracing() else torch.int32,
-            )
-            cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0)
-
-        hidden_states = hidden_states.unsqueeze(1)
-        rotary_pos_emb_cos = rotary_pos_emb_cos.to(hidden_states.device)
-        rotary_pos_emb_sin = rotary_pos_emb_sin.to(hidden_states.device)
-        max_seqlen = self.compute_attn_mask_seqlen(cu_seqlens)
-
-        grid_thw_np = grid_thw_tensor.cpu().numpy().astype(np.int32)
+        # Sequence boundaries come from the host grid, so neither the eager
+        # path nor a replay reads a device value back.
+        grid_thw_np = grid_thw.cpu().numpy().astype(np.int32)
         cu_seqlens_np = np.repeat(grid_thw_np[:, 1] * grid_thw_np[:, 2], grid_thw_np[:, 0]).cumsum(
             axis=0, dtype=np.int32
         )
         cu_seqlens_np = np.concatenate([np.zeros(1, dtype=np.int32), cu_seqlens_np])
+        cu_seqlens_host = torch.from_numpy(cu_seqlens_np)
+        if self.device.type == "cpu":
+            cu_seqlens = cu_seqlens_host
+        else:
+            cu_seqlens = async_tensor_h2d(cu_seqlens_host, device=self.device)
+        # Attention reads this scalar on the host. Replay keeps the capture
+        # budget's bound instead.
+        max_seqlen = self.compute_attn_mask_seqlen(cu_seqlens_host)
         sequence_lengths = MMEncoderAttention.maybe_compute_seq_lens(
             self.attn_backend,
             cu_seqlens_np,
             self.device,
         )
 
+        metadata.update(
+            rotary_pos_emb_cos=rotary_pos_emb_cos,
+            rotary_pos_emb_sin=rotary_pos_emb_sin,
+            cu_seqlens=cu_seqlens,
+            max_seqlen=max_seqlen,
+        )
+        if sequence_lengths is not None:
+            metadata["sequence_lengths"] = sequence_lengths
+        return metadata
+
+    def forward(
+        self, x: torch.Tensor, grid_thw, *, encoder_metadata: dict[str, torch.Tensor] | None = None
+    ) -> torch.Tensor:
+        if encoder_metadata is None:
+            encoder_metadata = self.prepare_encoder_metadata(grid_thw)
+        return self.forward_with_encoder_metadata(x, encoder_metadata)
+
+    def forward_with_encoder_metadata(self, x: torch.Tensor, metadata: dict[str, torch.Tensor]) -> torch.Tensor:
+        """Run patch embedding, vision blocks and all DeepStack mergers."""
+        hidden_states = self.patch_embed(x.to(device=self.device, dtype=self.dtype))
+        if self.apply_vit_abs_pos_embed:
+            hidden_states = hidden_states + metadata["pos_embeds"]
+        hidden_states = hidden_states.unsqueeze(1)
+
         hidden_states_list = []
         deepstack_visual_indexes = self.deepstack_visual_indexes
 
         for layer_num, blk in enumerate(self.blocks):
             hidden_states = blk(
                 hidden_states,
-                cu_seqlens=cu_seqlens,
-                rotary_pos_emb_cos=rotary_pos_emb_cos,
-                rotary_pos_emb_sin=rotary_pos_emb_sin,
-                max_seqlen=max_seqlen,
-                sequence_lengths=sequence_lengths,
+                cu_seqlens=metadata["cu_seqlens"],
+                rotary_pos_emb_cos=metadata["rotary_pos_emb_cos"],
+                rotary_pos_emb_sin=metadata["rotary_pos_emb_sin"],
+                max_seqlen=metadata["max_seqlen"],
+                sequence_lengths=metadata.get("sequence_lengths"),
             )
 
             if deepstack_visual_indexes is not None and layer_num in deepstack_visual_indexes:
@@ -440,8 +437,6 @@ class Qwen3OmniMoeAudioEncoder(_Qwen3OmniMoeAudioEncoder):
         feature_lens: torch.Tensor,
         aftercnn_lens: torch.Tensor,
     ):
-        import torch.nn.functional as F
-
         chunk_num = torch.ceil(feature_lens / (self.n_window * 2)).long()
 
         chunk_lengths = torch.tensor(
@@ -1024,6 +1019,7 @@ class Qwen3OmniMoeConditionalGenerationMixin(Qwen2_5OmniConditionalGenerationMix
 )
 class Qwen3OmniMoeThinkerForConditionalGeneration(
     nn.Module,
+    Qwen3OmniVisionEncoderCudaGraphMixin,
     SupportsMultiModal,
     SupportsPP,
     SupportsLoRA,
@@ -1159,6 +1155,8 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
                 ]
                 self.deepstack_input_embeds_num_tokens = 0
 
+        self._enable_image_encoder_cudagraph()
+
         with self._mark_language_model(vllm_config):
             lm_vllm_config = vllm_config.with_hf_config(
                 thinker_config.text_config,

vllm_omni/model_executor/models/qwen3_omni/vision_encoder_cudagraph.py

@@ -0,0 +1,182 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Qwen3-Omni image encoder contract for the upstream CUDA graph manager."""
+
+from collections.abc import Hashable
+from typing import Any
+
+import torch
+from vllm.distributed import get_pp_group
+from vllm.model_executor.models.utils import scatter_output_slices
+from vllm.v1.attention.backends.registry import AttentionBackendEnum
+from vllm.v1.worker.encoder_cudagraph_defs import (
+    EncoderCudaGraphCaptureInputs,
+    EncoderCudaGraphConfig,
+    EncoderCudaGraphReplayBuffers,
+    EncoderItemSpec,
+)
+
+
+def _pad_cu_seqlens(dst: torch.Tensor, src: torch.Tensor) -> None:
+    # Unused attention sequences must be empty, not run backwards to zero.
+    dst.copy_(src[-1].expand_as(dst))
+    dst[: src.shape[0]].copy_(src)
+
+
+class Qwen3OmniVisionEncoderCudaGraphMixin:
+    """Image-only protocol methods; video and audio retain embed_multimodal.
+
+    The capability attribute is set on eligible model instances after tower
+    construction. It is deliberately absent from this mixin's class: the
+    runtime protocol checks attribute presence, not the value of a false flag.
+    """
+
+    # The eager tower is launch-bound, so a group that needs several replays
+    # or per-image eager calls runs as one eager call instead; see
+    # SingleReplayEncoderCudaGraphManager.
+    encoder_cudagraph_single_replay = True
+
+    def _enable_image_encoder_cudagraph(self) -> None:
+        # Only the first pipeline rank runs the encoder. Dynamic-scale FP8 ViT
+        # attention advances a host-side amax slot per call, which a graph
+        # would freeze. The pipeline-group query comes last so a thinker built
+        # without a vision tower or distributed state never reaches it.
+        if (
+            not self.multimodal_config.enable_mm_embeds
+            and self.multimodal_config.mm_encoder_attn_dtype is None
+            and self.multimodal_config.get_limit_per_prompt("image") > 0
+            and getattr(self.visual, "attn_backend", None)
+            in {
+                AttentionBackendEnum.FLASH_ATTN,
+                AttentionBackendEnum.ROCM_AITER_FA,
+                AttentionBackendEnum.TRITON_ATTN,
+            }
+            and get_pp_group().is_first_rank
+        ):
+            self.supports_encoder_cudagraph = True
+
+    def get_encoder_cudagraph_config(self) -> EncoderCudaGraphConfig:
+        keys = ["pixel_values", "rotary_pos_emb_cos", "rotary_pos_emb_sin", "cu_seqlens", "max_seqlen"]
+        if self.visual.apply_vit_abs_pos_embed:
+            keys.append("pos_embeds")
+        return EncoderCudaGraphConfig(
+            modalities=["image"],
+            buffer_keys=keys,
+            out_hidden_size=self.visual_dim + self.multiscale_dim,
+            padding_logics={"cu_seqlens": _pad_cu_seqlens},
+        )
+
+    def get_input_modality(self, mm_kwargs: dict[str, Any]) -> str:
+        return "image"
+
+    def get_max_frames_per_video(self) -> int:
+        return 1
+
+    def get_encoder_cudagraph_budget_range(self, vllm_config) -> tuple[int, int]:
+        # Replay pays off while the eager tower is launch-bound: on an A800 a
+        # 196-token image replays 32% faster, while 280-400 tokens replay
+        # 4-10% slower and 560-784 tokens padded to the 1024 budget 12-51%
+        # slower. Every budget also adds captured graph memory. Auto-inferred
+        # budgets therefore stop at 256 tokens; larger images run eagerly
+        # unless encoder_cudagraph_token_budgets asks for more.
+        maximum = min(vllm_config.scheduler_config.max_num_batched_tokens, vllm_config.model_config.max_model_len, 256)
+        return min(64, maximum), maximum
+
+    def _image_grid(self, mm_kwargs: dict[str, Any]) -> list[list[int]]:
+        grid = mm_kwargs["image_grid_thw"]
+        return grid if isinstance(grid, list) else grid.tolist()
+
+    def get_encoder_cudagraph_item_specs(self, mm_kwargs: dict[str, Any]) -> list[EncoderItemSpec]:
+        merge = self.visual.spatial_merge_size
+        result = []
+        for t, h, w in self._image_grid(mm_kwargs):
+            if t != 1 or h <= 0 or w <= 0 or h % merge or w % merge:
+                raise ValueError(f"Invalid Qwen3-Omni image grid: {(t, h, w)}")
+            result.append(EncoderItemSpec(input_size=h * w, output_tokens=(h // merge) * (w // merge)))
+        return result
+
+    def select_encoder_cudagraph_items(self, mm_kwargs: dict[str, Any], indices: list[int]) -> dict[str, Any]:
+        grid = self._image_grid(mm_kwargs)
+        pixels = mm_kwargs["pixel_values"]
+        offsets = [0]
+        for spec in self.get_encoder_cudagraph_item_specs(mm_kwargs):
+            offsets.append(offsets[-1] + spec.input_size)
+        return {
+            "pixel_values": torch.cat([pixels[offsets[i] : offsets[i + 1]] for i in indices])
+            if indices
+            else pixels[:0],
+            "image_grid_thw": [grid[i] for i in indices],
+        }
+
+    def prepare_encoder_cudagraph_capture_inputs(
+        self,
+        token_budget: int,
+        max_batch_size: int,
+        max_frames_per_batch: int,
+        device: torch.device,
+        dtype: torch.dtype,
+        path: str = "default",
+        axis_keys: tuple[Hashable, ...] | None = None,
+    ) -> EncoderCudaGraphCaptureInputs:
+        if path != "default" or axis_keys:
+            raise ValueError("Qwen3-Omni images use one encoder path without capture axes")
+        visual = self.visual
+        merge = visual.spatial_merge_size
+        patches = token_budget * merge**2
+        # Infer the RoPE layout from one real minimal grid. No large dummy
+        # aspect ratio is needed, and no dummy grid indexes past the RoPE cache.
+        small = visual.prepare_encoder_metadata([[1, merge, merge]])
+        values = {
+            "rotary_pos_emb_cos": torch.ones(
+                (patches, *small["rotary_pos_emb_cos"].shape[1:]),
+                device=device,
+                dtype=small["rotary_pos_emb_cos"].dtype,
+            ),
+            "rotary_pos_emb_sin": torch.zeros(
+                (patches, *small["rotary_pos_emb_sin"].shape[1:]),
+                device=device,
+                dtype=small["rotary_pos_emb_sin"].dtype,
+            ),
+            "cu_seqlens": torch.full((max_batch_size + 1,), patches, device=device, dtype=torch.int32),
+            # This host scalar is baked into the captured attention launch.
+            # Any single image in this graph has at most `patches` patches.
+            "max_seqlen": torch.tensor(patches, dtype=torch.int32),
+        }
+        values["cu_seqlens"][0] = 0
+        patch_embed = visual.patch_embed
+        patch_width = patch_embed.proj.in_channels * patch_embed.temporal_patch_size * patch_embed.patch_size**2
+        values["pixel_values"] = torch.zeros((patches, patch_width), device=device, dtype=dtype)
+        if visual.apply_vit_abs_pos_embed:
+            values["pos_embeds"] = torch.zeros((patches, visual.hidden_size), device=device, dtype=visual.dtype)
+        return EncoderCudaGraphCaptureInputs(values=values)
+
+    def prepare_encoder_cudagraph_replay_buffers(
+        self,
+        mm_kwargs: dict[str, Any],
+        max_batch_size: int,
+        max_frames_per_batch: int,
+        path: str = "default",
+    ) -> EncoderCudaGraphReplayBuffers:
+        values = self.visual.prepare_encoder_metadata(self._image_grid(mm_kwargs))
+        # Replaying a smaller request must not lower the scalar bound retained
+        # by the capture. Tensor positions and sequence boundaries do change.
+        values.pop("max_seqlen")
+        values["pixel_values"] = mm_kwargs["pixel_values"]
+        return EncoderCudaGraphReplayBuffers(values=values)
+
+    def encoder_cudagraph_forward(self, inputs: dict[str, torch.Tensor], path: str = "default") -> torch.Tensor:
+        return self.visual(inputs["pixel_values"], None, encoder_metadata=inputs)
+
+    def encoder_eager_forward(self, mm_kwargs: dict[str, Any], path: str = "default") -> torch.Tensor:
+        return self.visual(mm_kwargs["pixel_values"], self._image_grid(mm_kwargs))
+
+    def postprocess_encoder_output(
+        self,
+        outputs: dict[str, torch.Tensor],
+        indices: list[int],
+        per_item_out_tokens: list[int],
+        dest: dict[int, torch.Tensor] | list[torch.Tensor | None],
+        clone: bool = False,
+        batch_mm_kwargs: dict[str, Any] | None = None,
+    ) -> None:
+        scatter_output_slices(outputs["default"], indices, per_item_out_tokens, dest, clone)

vllm_omni/worker/encoder_cudagraph.py

@@ -0,0 +1,32 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Encoder CUDA graph manager that replays a group only when one graph holds it."""
+
+from typing import Any
+
+import torch
+from vllm.v1.worker.encoder_cudagraph import EncoderCudaGraphManager
+
+
+class SingleReplayEncoderCudaGraphManager(EncoderCudaGraphManager):
+    """Replay an encoder group only when one captured graph holds all of it.
+
+    The upstream manager splits a group across several replays and gives each
+    item above the largest budget its own eager call. For a launch-bound vision
+    tower that is usually slower than one eager call over the whole group, so a
+    group that does not fit one graph goes back to the runner, which then
+    encodes it with a single ``embed_multimodal`` call. Models opt in by
+    setting ``encoder_cudagraph_single_replay``.
+    """
+
+    def execute(self, mm_kwargs: dict[str, Any]) -> list[torch.Tensor] | None:  # type: ignore[override]
+        if not self.use_dp:
+            specs = self._get_item_specs(mm_kwargs)
+            fits = len(specs) <= self.max_batch_size and all(
+                sum(spec.get_path_output_tokens(path) for spec in specs) <= max(budgets)
+                for path, budgets in self.path_token_budgets.items()
+            )
+            if not fits:
+                self.graph_misses += len(specs)
+                return None
+        return super().execute(mm_kwargs)

vllm_omni/worker/gpu_model_runner.py

@@ -283,6 +283,25 @@ class OmniGPUModelRunner(PrefixCacheRunnerMixin, GPUModelRunner):
             capture_sizes=capture_sizes,
         )
 
+    @torch.inference_mode()
+    def _create_encoder_cudagraph_manager(self):
+        raw_model = self.get_model()
+        if not getattr(raw_model, "encoder_cudagraph_single_replay", False):
+            return super()._create_encoder_cudagraph_manager()
+        from vllm.model_executor.models.interfaces import supports_encoder_cudagraph
+
+        from vllm_omni.worker.encoder_cudagraph import SingleReplayEncoderCudaGraphManager
+
+        if not (
+            self.compilation_config.cudagraph_mm_encoder
+            and self.supports_mm_inputs
+            and supports_encoder_cudagraph(raw_model)
+        ):
+            return None
+        return SingleReplayEncoderCudaGraphManager(
+            vllm_config=self.vllm_config, device=self.device, dtype=self.dtype, model=raw_model
+        )
+
     def _maybe_attach_attention_metadata_extensions(
         self,
         *,