A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[Misc][qwen2_5_vl][torch.compile] Enable `supports_torch_compile` on generic nn.Module and demonstrate speedup on Qwen Vision model (#23207)
Signed-off-by: Lucas Kabela <lucaskabela@meta.com> Signed-off-by: Lucas Kabela <lucasakabela@gmail.com>
Lucas Kabela committed
Oct 28, 2025 at 15:36 UTC
94666612a938380cb643c1555ef9aa68b7ab1e53
6 files changed
+335
-98
tests/compile/test_multimodal_compile.py
new
+36
@@ -0,0 +1,36 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import pytest
4
+
5
+from vllm.compilation.counter import compilation_counter
6
+from vllm.config.compilation import CompilationMode
7
+
8
+
9
+# forked needed to workaround https://github.com/vllm-project/vllm/issues/21073
10
+@pytest.mark.forked
11
+def test_qwen2_5_vl_compilation(vllm_runner, monkeypatch):
12
+ """Test that Qwen2.5-VL vision submodules are compiled.
13
+
14
+ This test verifies that the 3 vision submodules (Qwen2_5_VisionPatchEmbed,
15
+ Qwen2_5_VisionBlock, and Qwen2_5_VisionPatchMerger) are properly tagged
16
+ for compilation by checking that num_models_seen increases by at least 3.
17
+ """
18
+ # Disable multiprocessing so that the counter is in the same process
19
+ monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
20
+
21
+ with (
22
+ # NOTE: Qwen2.5-VL has 35 models in total - the LLM backend
23
+ # Vision Patch Embed, Vision Patch Merger, and then 32 Vision Blocks
24
+ # (one for each layer) - in the future, we should fix vLLM compilation
25
+ # logic to handle this case and only compile the Vision submodules once
26
+ # and reuse the compiled code for all layers
27
+ # See https://github.com/vllm-project/vllm/issues/27590
28
+ compilation_counter.expect(num_models_seen=35),
29
+ vllm_runner(
30
+ "Qwen/Qwen2.5-VL-3B-Instruct",
31
+ max_model_len=2048,
32
+ gpu_memory_utilization=0.7,
33
+ compilation_config={"mode": CompilationMode.VLLM_COMPILE},
34
+ ) as _,
35
+ ):
36
+ pass
vllm/attention/ops/vit_attn_wrappers.py
new
+125
@@ -0,0 +1,125 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""
4
+This file contains ops for ViT attention to be compatible with torch.compile
5
+as there are operations here not supported by torch.compile (for instance,
6
+`to_list` in xformers attn, or `.item()` in flash attention)
7
+
8
+Using these ops and wrapping vision blocks with `torch.compile` can speed up
9
+throughput in vision models by ~5% relative on H100, and improve token
10
+latencies by ~7% (see qwen2_5_vl for example usage)
11
+
12
+To use these ops, you must have a recent version of PyTorch installed (>= 2.4.0)
13
+"""
14
+
15
+import einops
16
+import torch
17
+
18
+from vllm.utils.torch_utils import direct_register_custom_op
19
+
20
+
21
+def xformers_attn_seqlens_wrapper(
22
+ q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, seqlens: torch.Tensor
23
+) -> torch.Tensor:
24
+ from xformers import ops as xops
25
+ from xformers.ops.fmha.attn_bias import BlockDiagonalMask
26
+
27
+ attn_bias = BlockDiagonalMask.from_seqlens(
28
+ q_seqlen=seqlens.tolist(), kv_seqlen=None, device=q.device
29
+ )
30
+ context_layer = xops.memory_efficient_attention_forward(
31
+ q, k, v, attn_bias=attn_bias, p=0, scale=None
32
+ )
33
+ context_layer = einops.rearrange(context_layer, "b s h d -> s b (h d)").contiguous()
34
+ return context_layer
35
+
36
+
37
+def xformers_attn_seqlens_wrapper_fake(
38
+ q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, seqlens: torch.Tensor
39
+) -> torch.Tensor:
40
+ b, s, h, d = q.shape
41
+ return torch.empty((s, b, h * d), dtype=q.dtype, device=q.device)
42
+
43
+
44
+direct_register_custom_op(
45
+ op_name="xformers_attn_seqlens_wrapper",
46
+ op_func=xformers_attn_seqlens_wrapper,
47
+ fake_impl=xformers_attn_seqlens_wrapper_fake,
48
+)
49
+
50
+
51
+def vit_xformers_attn_wrapper(
52
+ q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, seqlens: torch.Tensor
53
+) -> torch.Tensor:
54
+ return torch.ops.vllm.xformers_attn_seqlens_wrapper(q, k, v, seqlens)
55
+
56
+
57
+def flash_attn_maxseqlen_wrapper(
58
+ q: torch.Tensor,
59
+ k: torch.Tensor,
60
+ v: torch.Tensor,
61
+ cu_seqlens: torch.Tensor,
62
+ max_seqlen: torch.Tensor,
63
+ batch_size: int,
64
+ is_rocm_aiter: bool,
65
+ use_upstream_fa: bool,
66
+) -> torch.Tensor:
67
+ if is_rocm_aiter:
68
+ from aiter import flash_attn_varlen_func
69
+ else:
70
+ if use_upstream_fa:
71
+ from flash_attn import flash_attn_varlen_func
72
+ else:
73
+ from vllm.vllm_flash_attn import flash_attn_varlen_func
74
+ q, k, v = (einops.rearrange(x, "b s ... -> (b s) ...") for x in [q, k, v])
75
+ output = flash_attn_varlen_func(
76
+ q,
77
+ k,
78
+ v,
79
+ cu_seqlens_q=cu_seqlens,
80
+ cu_seqlens_k=cu_seqlens,
81
+ max_seqlen_q=max_seqlen.item(),
82
+ max_seqlen_k=max_seqlen.item(),
83
+ dropout_p=0.0,
84
+ causal=False,
85
+ )
86
+ context_layer = einops.rearrange(
87
+ output, "(b s) h d -> s b (h d)", b=batch_size
88
+ ).contiguous()
89
+ return context_layer
90
+
91
+
92
+def flash_attn_maxseqlen_wrapper_fake(
93
+ q: torch.Tensor,
94
+ k: torch.Tensor,
95
+ v: torch.Tensor,
96
+ cu_seqlens: torch.Tensor,
97
+ max_seqlen: torch.Tensor,
98
+ batch_size: int,
99
+ is_rocm_aiter: bool,
100
+ use_upstream_fa: bool,
101
+) -> torch.Tensor:
102
+ b, s, h, d = q.shape
103
+ return torch.empty((s, b, h * d), dtype=q.dtype, device=q.device)
104
+
105
+
106
+direct_register_custom_op(
107
+ op_name="flash_attn_maxseqlen_wrapper",
108
+ op_func=flash_attn_maxseqlen_wrapper,
109
+ fake_impl=flash_attn_maxseqlen_wrapper_fake,
110
+)
111
+
112
+
113
+def vit_flash_attn_wrapper(
114
+ q: torch.Tensor,
115
+ k: torch.Tensor,
116
+ v: torch.Tensor,
117
+ cu_seqlens: torch.Tensor,
118
+ max_seqlen: torch.Tensor,
119
+ batch_size: int,
120
+ is_rocm_aiter: bool,
121
+ use_upstream_fa: bool,
122
+) -> torch.Tensor:
123
+ return torch.ops.vllm.flash_attn_maxseqlen_wrapper(
124
+ q, k, v, cu_seqlens, max_seqlen, batch_size, is_rocm_aiter, use_upstream_fa
125
+ )
vllm/compilation/decorators.py
+56
-4
@@ -18,7 +18,12 @@ from torch._dynamo.symbolic_convert import InliningInstructionTranslator
18
import vllm.envs as envs
19
from vllm.compilation.counter import compilation_counter
20
from vllm.compilation.wrapper import TorchCompileWrapperWithCustomDispatcher
21
-from vllm.config import CompilationMode, VllmConfig, set_current_vllm_config
21
+from vllm.config import (
22
+ CompilationMode,
23
+ VllmConfig,
24
+ get_current_vllm_config,
25
+ set_current_vllm_config,
26
+)
27
from vllm.logger import init_logger
28
from vllm.sequence import IntermediateTensors
29
from vllm.utils.import_utils import resolve_obj_by_qualname
@@ -74,6 +79,21 @@ def support_torch_compile(
79
) -> Callable[[_T], _T]: ...
80
81
82
+@overload
83
+def support_torch_compile(
84
+ *,
85
+ mark_unbacked_dims: dict[str, int | list[int]] | None,
86
+) -> Callable[[_T], _T]: ...
87
+
88
+
89
+@overload
90
+def support_torch_compile(
91
+ *,
92
+ dynamic_arg_dims: dict[str, int | list[int]] | None,
93
+ mark_unbacked_dims: dict[str, int | list[int]] | None,
94
+) -> Callable[[_T], _T]: ...
95
+
96
+
97
@overload
98
def support_torch_compile(cls: _T) -> _T: ...
99
@@ -82,6 +102,7 @@ def support_torch_compile(
102
cls: _T | None = None,
103
*,
104
dynamic_arg_dims: dict[str, int | list[int]] | None = None,
105
+ mark_unbacked_dims: dict[str, int | list[int]] | None = None,
106
enable_if: Callable[[VllmConfig], bool] | None = None,
107
) -> Callable[[_T], _T] | _T:
108
"""
@@ -135,6 +156,11 @@ def support_torch_compile(
156
returns a boolean value indicating whether to compile the model or not.
157
This is useful if you want to compile the model only when certain
158
conditions are met.
159
+
160
+ `mark_unbacked_dims` is a dictionary that maps argument names with a dynamic
161
+ dim to be decorated with `mark_unbacked`. This is useful if we would like to
162
+ enforce that dynamo do not specialize on 0/1 values in the case of dummy input
163
+ such as for vision model compilation
164
"""
165
166
def cls_decorator_helper(cls: _T) -> _T:
@@ -172,7 +198,9 @@ def support_torch_compile(
198
raise ValueError(
199
f"Argument {k} not found in the forward method of {cls}"
200
)
175
- return _support_torch_compile(cls, inferred_dynamic_arg_dims, enable_if)
201
+ return _support_torch_compile(
202
+ cls, inferred_dynamic_arg_dims, mark_unbacked_dims, enable_if
203
+ )
204
205
if cls is not None:
206
# use `support_torch_compile` as a decorator without arguments
@@ -212,6 +240,7 @@ def _verify_source_unchanged(source_info, vllm_config) -> None:
240
def _support_torch_compile(
241
cls: _T,
242
dynamic_arg_dims: dict[str, int | list[int]],
243
+ mark_unbacked_dims: dict[str, int | list[int]] | None = None,
244
enable_if: Callable[[VllmConfig], bool] | None = None,
245
) -> _T:
246
"""
@@ -230,8 +259,22 @@ def _support_torch_compile(
259
260
setattr(cls, IGNORE_COMPILE_KEY, False)
261
233
- def __init__(self, *, vllm_config: VllmConfig, prefix: str = "", **kwargs):
234
- old_init(self, vllm_config=vllm_config, prefix=prefix, **kwargs)
262
+ def __init__(
263
+ self, *, vllm_config: VllmConfig | None = None, prefix: str = "", **kwargs
264
+ ):
265
+ if vllm_config is None:
266
+ vllm_config = get_current_vllm_config()
267
+
268
+ # NOTE: to support multimodal models (such as encoder),
269
+ # we may not have vllm_config so we may need to patch
270
+ # it
271
+ sig = inspect.signature(old_init)
272
+ if "vllm_config" in sig.parameters:
273
+ kwargs["vllm_config"] = vllm_config
274
+ if "prefix" in sig.parameters:
275
+ kwargs["prefix"] = prefix
276
+ old_init(self, **kwargs)
277
+
278
self.vllm_config = vllm_config
279
enable_compile = enable_if is None or enable_if(vllm_config)
280
# for CompilationMode.STOCK_TORCH_COMPILE , the upper level model runner
@@ -344,6 +387,15 @@ def _support_torch_compile(
387
"Unsupported dynamic dimensions"
388
f" {dims} for argument {k} with type {type(arg)}."
389
)
390
+ if mark_unbacked_dims:
391
+ for k, dims in mark_unbacked_dims.items():
392
+ arg = bound_args.arguments.get(k)
393
+ if arg is not None:
394
+ dims = [dims] if isinstance(dims, int) else dims
395
+ if isinstance(arg, torch.Tensor):
396
+ # In case dims is specified with negative indexing
397
+ dims = [arg.ndim + dim if dim < 0 else dim for dim in dims]
398
+ torch._dynamo.decorators.mark_unbacked(arg, dims)
399
# here, it is the starting point of the `torch.compile` process
400
start_monitoring_torch_compile(self.vllm_config)
401
logger.debug("Start compiling function %s", self.original_code_object)
vllm/config/compilation.py
+2
@@ -684,6 +684,8 @@ class CompilationConfig:
684
685
from vllm.compilation.backends import VllmBackend
686
687
+ # TODO[@lucaskabela]: See if we can forward prefix
688
+ # https://github.com/vllm-project/vllm/issues/27045
689
return VllmBackend(vllm_config)
690
691
def post_init_cudagraph_sizes(self) -> None:
vllm/model_executor/models/qwen2_5_omni_thinker.py
+6
-2
@@ -45,6 +45,7 @@ from transformers.models.whisper import WhisperFeatureExtractor
45
46
from vllm.config import VllmConfig
47
from vllm.config.multimodal import BaseDummyOptions
48
+from vllm.forward_context import set_forward_context
49
from vllm.logger import init_logger
50
from vllm.model_executor.models.module_mapping import MultiModelKeys
51
from vllm.model_executor.models.qwen2_5_vl import (
@@ -759,7 +760,8 @@ class Qwen2_5OmniConditionalGenerationMixin:
760
assert grid_thw.ndim == 2
761
762
pixel_values = image_input["pixel_values"].type(self.visual.dtype)
762
- image_embeds = self.visual(pixel_values, grid_thw=grid_thw)
763
+ with set_forward_context(None, self.vllm_config):
764
+ image_embeds = self.visual(pixel_values, grid_thw=grid_thw)
765
# Split concatenated embeddings for each image item.
766
merge_size = self.visual.spatial_merge_size
767
sizes = grid_thw.prod(-1) // merge_size // merge_size
@@ -779,7 +781,8 @@ class Qwen2_5OmniConditionalGenerationMixin:
781
assert grid_thw.ndim == 2
782
783
pixel_values_videos = video_input["pixel_values_videos"].type(self.visual.dtype)
782
- video_embeds = self.visual(pixel_values_videos, grid_thw=grid_thw)
784
+ with set_forward_context(None, self.vllm_config):
785
+ video_embeds = self.visual(pixel_values_videos, grid_thw=grid_thw)
786
# Split concatenated embeddings for each video item.
787
merge_size = self.visual.spatial_merge_size
788
sizes = grid_thw.prod(-1) // merge_size // merge_size
@@ -839,6 +842,7 @@ class Qwen2_5OmniThinkerForConditionalGeneration(
842
843
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
844
super().__init__()
845
+ self.vllm_config = vllm_config
846
thinker_config: Qwen2_5OmniThinkerConfig = (
847
vllm_config.model_config.hf_config.thinker_config
848
)
vllm/model_executor/models/qwen2_5_vl.py
+110
-92
@@ -31,10 +31,10 @@ from collections.abc import Callable, Iterable, Mapping, Sequence
31
from functools import lru_cache, partial
32
from typing import Annotated, Any, Literal, TypeAlias
33
34
+import einops
35
import torch
36
import torch.nn as nn
37
import torch.nn.functional as F
37
-from einops import rearrange
38
from transformers import BatchFeature, PretrainedConfig
39
from transformers.models.qwen2_5_vl import Qwen2_5_VLProcessor
40
from transformers.models.qwen2_5_vl.configuration_qwen2_5_vl import (
@@ -47,9 +47,15 @@ from vllm.attention.layer import (
47
check_upstream_fa_availability,
48
maybe_get_vit_flash_attn_backend,
49
)
50
+from vllm.attention.ops.vit_attn_wrappers import (
51
+ vit_flash_attn_wrapper,
52
+ vit_xformers_attn_wrapper,
53
+)
54
+from vllm.compilation.decorators import support_torch_compile
55
from vllm.config import VllmConfig
56
from vllm.distributed import parallel_state
57
from vllm.distributed import utils as dist_utils
58
+from vllm.forward_context import set_forward_context
59
from vllm.logger import init_logger
60
from vllm.model_executor.layers.activation import get_act_and_mul_fn
61
from vllm.model_executor.layers.layernorm import RMSNorm
@@ -392,8 +398,8 @@ class Qwen2_5_VisionAttention(nn.Module):
398
x: torch.Tensor,
399
cu_seqlens: torch.Tensor,
400
rotary_pos_emb: torch.Tensor,
395
- max_seqlen: int | None = None, # Only used for Flash Attention
396
- seqlens: list[int] | None = None, # Only used for xFormers
401
+ max_seqlen: torch.Tensor, # Only used for Flash Attention
402
+ seqlens: torch.Tensor, # Only used for xFormers
403
) -> torch.Tensor:
404
# [s, b, c] --> [s, b, head * 3 * head_dim]
405
x, _ = self.qkv(x)
@@ -402,7 +408,7 @@ class Qwen2_5_VisionAttention(nn.Module):
408
q, k, v = self.split_qkv(x)
409
batch_size = q.shape[1]
410
405
- q, k, v = (rearrange(x, "s b ... -> b s ...") for x in (q, k, v))
411
+ q, k, v = (einops.rearrange(x, "s b ... -> b s ...") for x in (q, k, v))
412
if rotary_pos_emb is not None:
413
# [2 * b, s, heads, head_dim]
414
qk_concat = torch.cat([q, k], dim=0)
@@ -410,31 +416,18 @@ class Qwen2_5_VisionAttention(nn.Module):
416
q, k = torch.chunk(qk_rotated, 2, dim=0)
417
418
if self.is_flash_attn_backend:
413
- q, k, v = (rearrange(x, "b s ... -> (b s) ...") for x in [q, k, v])
414
-
415
- output = self.flash_attn_varlen_func(
419
+ context_layer = vit_flash_attn_wrapper(
420
q,
421
k,
422
v,
419
- cu_seqlens_q=cu_seqlens,
420
- cu_seqlens_k=cu_seqlens,
421
- max_seqlen_q=max_seqlen,
422
- max_seqlen_k=max_seqlen,
423
- dropout_p=0.0,
424
- causal=False,
423
+ cu_seqlens,
424
+ max_seqlen,
425
+ batch_size,
426
+ self.attn_backend == _Backend.ROCM_AITER_FA,
427
+ self.use_upstream_fa,
428
)
426
-
427
- context_layer = rearrange(
428
- output, "(b s) h d -> s b (h d)", b=batch_size
429
- ).contiguous()
429
elif self.attn_backend == _Backend.TORCH_SDPA:
430
# Execute attention entry by entry for speed & less VRAM.
432
- from vllm.platforms import current_platform
433
-
434
- if current_platform.is_rocm():
435
- q = q.contiguous()
436
- k = k.contiguous()
437
- v = v.contiguous()
431
outputs = []
432
for i in range(1, len(cu_seqlens)):
433
start_idx = cu_seqlens[i - 1]
@@ -443,34 +436,31 @@ class Qwen2_5_VisionAttention(nn.Module):
436
k_i = k[:, start_idx:end_idx]
437
v_i = v[:, start_idx:end_idx]
438
q_i, k_i, v_i = (
446
- rearrange(x, "b s h d -> b h s d") for x in [q_i, k_i, v_i]
439
+ einops.rearrange(x, "b s h d -> b h s d") for x in [q_i, k_i, v_i]
440
)
441
output_i = F.scaled_dot_product_attention(q_i, k_i, v_i, dropout_p=0.0)
449
- output_i = rearrange(output_i, "b h s d -> b s h d ")
442
+ output_i = einops.rearrange(output_i, "b h s d -> b s h d ")
443
outputs.append(output_i)
444
context_layer = torch.cat(outputs, dim=1)
452
- context_layer = rearrange(
445
+ context_layer = einops.rearrange(
446
context_layer, "b s h d -> s b (h d)"
447
).contiguous()
448
elif self.attn_backend == _Backend.XFORMERS:
456
- from xformers import ops as xops
457
- from xformers.ops.fmha.attn_bias import BlockDiagonalMask
458
-
459
- attn_bias = BlockDiagonalMask.from_seqlens(
460
- q_seqlen=seqlens, kv_seqlen=None, device=q.device
461
- )
462
-
463
- context_layer = xops.memory_efficient_attention_forward(
464
- q, k, v, attn_bias=attn_bias, p=0, scale=None
465
- )
466
- context_layer = rearrange(
467
- context_layer, "b s h d -> s b (h d)"
468
- ).contiguous()
449
+ context_layer = vit_xformers_attn_wrapper(q, k, v, seqlens)
450
451
output, _ = self.proj(context_layer)
452
return output
453
454
455
+@support_torch_compile(
456
+ dynamic_arg_dims={
457
+ "x": 0,
458
+ "cu_seqlens": 0,
459
+ "rotary_pos_emb": 0,
460
+ "seqlens": 0,
461
+ },
462
+ mark_unbacked_dims={"seqlens": 0},
463
+)
464
class Qwen2_5_VisionBlock(nn.Module):
465
def __init__(
466
self,
@@ -515,8 +505,8 @@ class Qwen2_5_VisionBlock(nn.Module):
505
x: torch.Tensor,
506
cu_seqlens: torch.Tensor,
507
rotary_pos_emb: torch.Tensor,
518
- max_seqlen: int | None = None, # Only used for Flash Attention
519
- seqlens: list[int] | None = None, # Only used for xFormers
508
+ max_seqlen: torch.Tensor, # Only used for Flash Attention
509
+ seqlens: torch.Tensor, # Only used for xFormers
510
) -> torch.Tensor:
511
x_attn = self.attn(
512
self.norm1(x),
@@ -530,6 +520,11 @@ class Qwen2_5_VisionBlock(nn.Module):
520
return x
521
522
523
+@support_torch_compile(
524
+ dynamic_arg_dims={
525
+ "x": 0,
526
+ }
527
+)
528
class Qwen2_5_VisionPatchEmbed(nn.Module):
529
def __init__(
530
self,
@@ -556,6 +551,11 @@ class Qwen2_5_VisionPatchEmbed(nn.Module):
551
return x
552
553
554
+@support_torch_compile(
555
+ dynamic_arg_dims={
556
+ "x": 0,
557
+ }
558
+)
559
class Qwen2_5_VisionPatchMerger(nn.Module):
560
def __init__(
561
self,
@@ -665,13 +665,18 @@ class Qwen2_5_VisionTransformer(nn.Module):
665
self.spatial_merge_size = vision_config.spatial_merge_size
666
self.fullatt_block_indexes = vision_config.fullatt_block_indexes
667
self.spatial_merge_unit = self.spatial_merge_size**2
668
-
669
- self.patch_embed = Qwen2_5_VisionPatchEmbed(
670
- patch_size=patch_size,
671
- temporal_patch_size=temporal_patch_size,
672
- in_channels=in_channels,
673
- hidden_size=self.hidden_size,
674
- )
668
+ # TODO[@lucaskabela]: Investigate fixing this usage
669
+ # see https://github.com/vllm-project/vllm/issues/27044
670
+ # DO NOT MOVE THIS IMPORT
671
+ from vllm.compilation.backends import set_model_tag
672
+
673
+ with set_model_tag("Qwen2_5_VisionPatchEmbed"):
674
+ self.patch_embed = Qwen2_5_VisionPatchEmbed(
675
+ patch_size=patch_size,
676
+ temporal_patch_size=temporal_patch_size,
677
+ in_channels=in_channels,
678
+ hidden_size=self.hidden_size,
679
+ )
680
681
norm_layer = partial(RMSNorm, eps=norm_eps)
682
head_dim = self.hidden_size // self.num_heads
@@ -701,32 +706,35 @@ class Qwen2_5_VisionTransformer(nn.Module):
706
f"Qwen2.5-VL does not support {self.attn_backend} backend now."
707
)
708
704
- self.blocks = nn.ModuleList(
705
- [
706
- Qwen2_5_VisionBlock(
707
- dim=self.hidden_size,
708
- num_heads=self.num_heads,
709
- mlp_hidden_dim=vision_config.intermediate_size,
710
- act_fn=get_act_and_mul_fn(vision_config.hidden_act),
711
- norm_layer=norm_layer,
712
- quant_config=quant_config,
713
- prefix=f"{prefix}.blocks.{layer_idx}",
714
- use_data_parallel=use_data_parallel,
715
- attn_backend=self.attn_backend,
716
- use_upstream_fa=use_upstream_fa,
717
- )
718
- for layer_idx in range(depth)
719
- ]
720
- )
721
- self.merger = Qwen2_5_VisionPatchMerger(
722
- d_model=vision_config.out_hidden_size,
723
- context_dim=self.hidden_size,
724
- norm_layer=norm_layer,
725
- spatial_merge_size=self.spatial_merge_size,
726
- quant_config=quant_config,
727
- prefix=f"{prefix}.merger",
728
- use_data_parallel=use_data_parallel,
729
- )
709
+ with set_model_tag("Qwen2_5_VisionBlock"):
710
+ self.blocks = nn.ModuleList(
711
+ [
712
+ Qwen2_5_VisionBlock(
713
+ dim=self.hidden_size,
714
+ num_heads=self.num_heads,
715
+ mlp_hidden_dim=vision_config.intermediate_size,
716
+ act_fn=get_act_and_mul_fn(vision_config.hidden_act),
717
+ norm_layer=norm_layer,
718
+ quant_config=quant_config,
719
+ prefix=f"{prefix}.blocks.{layer_idx}",
720
+ use_data_parallel=use_data_parallel,
721
+ attn_backend=self.attn_backend,
722
+ use_upstream_fa=use_upstream_fa,
723
+ )
724
+ for layer_idx in range(depth)
725
+ ]
726
+ )
727
+
728
+ with set_model_tag("Qwen2_5_VisionPatchMerger"):
729
+ self.merger = Qwen2_5_VisionPatchMerger(
730
+ d_model=vision_config.out_hidden_size,
731
+ context_dim=self.hidden_size,
732
+ norm_layer=norm_layer,
733
+ spatial_merge_size=self.spatial_merge_size,
734
+ quant_config=quant_config,
735
+ prefix=f"{prefix}.merger",
736
+ use_data_parallel=use_data_parallel,
737
+ )
738
739
@property
740
def dtype(self) -> torch.dtype:
@@ -827,15 +835,18 @@ class Qwen2_5_VisionTransformer(nn.Module):
835
def compute_attn_mask_seqlen(
836
self,
837
cu_seqlens: torch.Tensor,
830
- ) -> tuple[int | None, list[int] | None]:
831
- max_seqlen, seqlens = None, None
838
+ ) -> tuple[torch.Tensor, torch.Tensor]:
839
+ max_seqlen, seqlens = (
840
+ torch.zeros(1, device=cu_seqlens.device),
841
+ torch.zeros(1, device=cu_seqlens.device),
842
+ )
843
if (
844
self.attn_backend == _Backend.FLASH_ATTN
845
or self.attn_backend == _Backend.ROCM_AITER_FA
846
):
836
- max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()
847
+ max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()
848
elif self.attn_backend == _Backend.XFORMERS:
838
- seqlens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()
849
+ seqlens = cu_seqlens[1:] - cu_seqlens[:-1]
850
return max_seqlen, seqlens
851
852
@staticmethod
@@ -1233,6 +1244,7 @@ class Qwen2_5_VLForConditionalGeneration(
1244
1245
self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data"
1246
self.config = config
1247
+ self.vllm_config = vllm_config
1248
self.multimodal_config = multimodal_config
1249
self.video_pruning_rate = multimodal_config.video_pruning_rate
1250
self.is_multimodal_pruning_enabled = (
@@ -1248,7 +1260,7 @@ class Qwen2_5_VLForConditionalGeneration(
1260
else None
1261
)
1262
self.visual = Qwen2_5_VisionTransformer(
1251
- config.vision_config,
1263
+ vision_config=config.vision_config,
1264
norm_eps=getattr(config, "rms_norm_eps", 1e-6),
1265
quant_config=self.quant_config,
1266
prefix=maybe_prefix(prefix, "visual"),
@@ -1336,13 +1348,13 @@ class Qwen2_5_VLForConditionalGeneration(
1348
image_embeds = image_input["image_embeds"].type(self.visual.dtype)
1349
else:
1350
pixel_values = image_input["pixel_values"]
1339
-
1340
- if self.use_data_parallel:
1341
- return run_dp_sharded_mrope_vision_model(
1342
- self.visual, pixel_values, grid_thw_list, rope_type="rope_3d"
1343
- )
1344
- else:
1345
- image_embeds = self.visual(pixel_values, grid_thw=grid_thw_list)
1351
+ with set_forward_context(None, self.vllm_config):
1352
+ if self.use_data_parallel:
1353
+ return run_dp_sharded_mrope_vision_model(
1354
+ self.visual, pixel_values, grid_thw_list, rope_type="rope_3d"
1355
+ )
1356
+ else:
1357
+ image_embeds = self.visual(pixel_values, grid_thw=grid_thw_list)
1358
1359
# Split concatenated embeddings for each image item.
1360
# Using prod on grid_thw_list instead of grid_thw.prod avoids CUDA sync
@@ -1396,12 +1408,18 @@ class Qwen2_5_VLForConditionalGeneration(
1408
video_embeds = video_input["video_embeds"].type(self.visual.dtype)
1409
else:
1410
pixel_values_videos = video_input["pixel_values_videos"]
1399
- if self.use_data_parallel:
1400
- return run_dp_sharded_mrope_vision_model(
1401
- self.visual, pixel_values_videos, grid_thw_list, rope_type="rope_3d"
1402
- )
1403
- else:
1404
- video_embeds = self.visual(pixel_values_videos, grid_thw=grid_thw_list)
1411
+ with set_forward_context(None, self.vllm_config):
1412
+ if self.use_data_parallel:
1413
+ return run_dp_sharded_mrope_vision_model(
1414
+ self.visual,
1415
+ pixel_values_videos,
1416
+ grid_thw_list,
1417
+ rope_type="rope_3d",
1418
+ )
1419
+ else:
1420
+ video_embeds = self.visual(
1421
+ pixel_values_videos, grid_thw=grid_thw_list
1422
+ )
1423
1424
# Split concatenated embeddings for each video item.
1425
merge_size = self.visual.spatial_merge_size