diff --git a/recipes/Qwen/Qwen3-Omni.md b/recipes/Qwen/Qwen3-Omni.md index e319dddfb..6f30561f7 100644 --- a/recipes/Qwen/Qwen3-Omni.md +++ b/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 diff --git a/tests/model_executor/models/qwen3_omni/test_vision_encoder_cudagraph.py b/tests/model_executor/models/qwen3_omni/test_vision_encoder_cudagraph.py new file mode 100644 index 000000000..b9b73654e --- /dev/null +++ b/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 diff --git a/tests/model_executor/models/qwen3_omni/test_vision_encoder_cudagraph_cuda.py b/tests/model_executor/models/qwen3_omni/test_vision_encoder_cudagraph_cuda.py new file mode 100644 index 000000000..0b8602a70 --- /dev/null +++ b/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() 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 50601d1fa..7b8c6844b 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py +++ b/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), 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 b610e77ff..fd37455af 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 @@ -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, diff --git a/vllm_omni/model_executor/models/qwen3_omni/vision_encoder_cudagraph.py b/vllm_omni/model_executor/models/qwen3_omni/vision_encoder_cudagraph.py new file mode 100644 index 000000000..e2b14d517 --- /dev/null +++ b/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) diff --git a/vllm_omni/worker/encoder_cudagraph.py b/vllm_omni/worker/encoder_cudagraph.py new file mode 100644 index 000000000..6804ade9c --- /dev/null +++ b/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) diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index 931e36d1c..9950e9cd5 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/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, *,