[V1] Add V1 support of Qwen2-VL

[ywang96]

Continued from #11668 for supporting Qwen2-VL on V1.

Co-authored-by: @imkero who has done the great work on the V1 MRoPE implementation.

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.
Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can do one of these:

• Add ready label to the PR
• Enable auto-merge.

🚀

[ywang96]

@baifanxxx @Zhiy-Zhang I have verified this PR works end-to-end on V1 with all following commands, so please try it out and let me know if there's any issue

VLLM_USE_V1=1 python3 examples/offline_inference/vision_language.py --model-type qwen2_vl
VLLM_USE_V1=1 python3 examples/offline_inference/vision_language.py --model-type qwen2_vl --modality video
VLLM_USE_V1=1 python3 examples/offline_inference/vision_language_multi_image.py --model-type qwen2_vl 

[WoosukKwon]

LGTM wrt the V1 model runner.

[WoosukKwon]

A dumb question: Why don't create the tensors with the permuted shape from the first place?

[ywang96]

The comment above explains it:

# a permuted mrope_positions tensor satisfying the following
#   properties to allow `torch.compile` work properly:
# - shape: (3, <variable>)
# - stride: (1, <variable>)

but creating the tensors with (3, self.max_num_tokens) directly will result in a stride of (self.max_num_tokens, 1). Alternatively, we can use torch.as_strided but I think using permute is cleaner.

@imkero @youkaichao feel free to comment since I'm not a torch.compile expert at all 😄.

[youkaichao]

why do we need stride: (1, <variable>) ? I think stride: (<variable>, 1) should work, with a normal (self.max_num_tokens, 3) shape tensor.

[imkero]

Here is the reason why I impl MRoPE in this way. @youkaichao @ywang96 @WoosukKwon

I'm trying to give a clear explanation below, and feel free to ask me if there's still any confusion.

TL;DR

1. we should pass positions=self.mrope_positions[:, :total_num_scheduled_tokens] in the Qwen2-VL's forward pass (just like common rope's positions=self.positions[:total_num_scheduled_tokens])
2. however, self.mrope_positions[:, :total_num_scheduled_tokens] is not contiguous, and torch.compiled code complain about this. (causing an assertion error)
3. i discovered that if we exchange dim 0 and dim 1, then we should do self.mrope_positions[:total_num_scheduled_tokens, :] and we will get a contiguous tensor view to make it work.
4. it's possible to make a logically shape(3, dynamic) and physically shape(dynamic, 3) tensor via:

self.mrope_positions = torch.zeros((self.max_num_tokens, 3),
                                    dtype=torch.int64,
                                    device=self.device)
self.mrope_positions = self.mrope_positions.permute((1, 0))

5. this approach is compatible with Qwen2-VL's impl, and it works with torch.compile

Reason in detail

These lines of code organize the memory layout of mrope_positons, to avoid the following conflict:

1. Qwen2-VL's model code expects positions in shape (3, seq_len) (dynamic at dim 1)

2. torch.compiled model.forward func expects positions arg to have following attributes:

   • size: (3, shape[1]) = (3, total_num_scheduled_tokens)
   • stride: (shape[1], 1) = (total_num_scheduled_tokens, 1)

   a.k.a. contiguous

3. CUDA graphs use persistant buffers (including self.mrope_positions), that means:

   • self.mrope_positions is initialized with shape (3, max_num_tokens)
   • model.forward func's positions arg is given by self.mrope_positions[:, :total_num_scheduled_tokens]

4. And actually positions arg has following attributes:

   • size: (3, total_num_scheduled_tokens)
   • stride: (max_num_tokens, 1)
   • ↑ notice the difference here, view created by self.mrope_positions[:, :total_num_scheduled_tokens] is not contiguous now (if total_num_scheduled_tokens < max_num_tokens)

   This conflict with 2.

Approach in detail

We can avoid this conflict by organizing mrope_positons's memory layout to be F-contiguous (column-major) instead of C-contiguous (row-major) via:

self.mrope_positions = torch.zeros((self.max_num_tokens, 3),
                                    dtype=torch.int64,
                                    device=self.device)
self.mrope_positions = self.mrope_positions.permute((1, 0))

In this approch, torch compile's assertion of positions will become:

• size: (3, shape[1]) = (3, total_num_scheduled_tokens)
• stride: (1, shape[0]) = (1, 3)

This can always be satisfied because shape[0] of positions is a model-level constant (= 3), it will never change across different seq len

Illustrate C-Contiguous

mrope_dim_num = 3 # a constant across all seq

def run(total_num_scheduled_tokens, max_num_tokens):
    positions_c = torch.Tensor([
        list(range(max_num_tokens)),
    ] * mrope_dim_num)

    positions_c = positions_c[:, :total_num_scheduled_tokens]

    print("total_num_scheduled_tokens", total_num_scheduled_tokens)
    print("max_num_tokens", max_num_tokens)
    print("tensor", positions_c)
    print("shape:", positions_c.shape) # (3, total_num_scheduled_tokens)
    print("stride:", positions_c.stride()) # (max_num_tokens, 1)
    print()

run(total_num_scheduled_tokens=2, max_num_tokens=4)
run(total_num_scheduled_tokens=3, max_num_tokens=4)

it prints

total_num_scheduled_tokens 2
max_num_tokens 4
tensor tensor([[0., 1.],
        [0., 1.],
        [0., 1.]])
shape: torch.Size([3, 2])
stride: (4, 1)

total_num_scheduled_tokens 3
max_num_tokens 4
tensor tensor([[0., 1., 2.],
        [0., 1., 2.],
        [0., 1., 2.]])
shape: torch.Size([3, 3])
stride: (4, 1)

physically these two tensors would be (* means unused)

0   1   *   *   0   1   *   *   0   1   *   *

0   1   2   *   0   1   2   *   0   1   2   *

Illustrate F-Contiguous (this PR's approach)

mrope_dim_num = 3 # a constant across all seq

def run(total_num_scheduled_tokens, max_num_tokens):
    positions_f = torch.Tensor([
        [i] * mrope_dim_num for i in range(max_num_tokens)
    ])
    positions_f = positions_f.permute((1, 0))

    positions_f = positions_f[:, :total_num_scheduled_tokens]

    print("total_num_scheduled_tokens", total_num_scheduled_tokens)
    print("max_num_tokens", max_num_tokens)
    print("tensor", positions_f)
    print("shape:", positions_f.shape) # (3, total_num_scheduled_tokens)
    print("stride:", positions_f.stride()) # (max_num_tokens, 1)
    print()

run(total_num_scheduled_tokens=2, max_num_tokens=4)
run(total_num_scheduled_tokens=3, max_num_tokens=4)

it prints

total_num_scheduled_tokens 2
max_num_tokens 4
tensor tensor([[0., 1.],
        [0., 1.],
        [0., 1.]])
shape: torch.Size([3, 2])
stride: (1, 3)

total_num_scheduled_tokens 3
max_num_tokens 4
tensor tensor([[0., 1., 2.],
        [0., 1., 2.],
        [0., 1., 2.]])
shape: torch.Size([3, 3])
stride: (1, 3)

physically these two tensors would be (* means unused)

0   0   0   1   1   1   *   *   *   *   *   *

0   0   0   1   1   1   2   2   2   *   *   *

↑ As total_num_scheduled_tokens grows, positions_f is still contiguous, and torch.compile is happy to work with this contiguous tensor!

[imkero]

Maybe we should point out the stride of mrope_positions is a constant in that comment, and add a link to the discussion here

# a permuted (column-major) mrope_positions tensor satisfying the following
#   properties to allow `torch.compile` work properly:
# - shape: (3, <variable>)
# - stride: (1, 3)
# see https://github.com/vllm-project/vllm/pull/12128#discussion_r1920908088 for detail.

[imkero]

torch.compiled code here for reference:

A won't work approach

source code:

self.mrope_positions = torch.zeros((3, self.max_num_tokens),
                                    dtype=torch.int64,
                                    device=self.device)

compiled code:

def call(args):
    arg0_1, arg1_1, arg2_1, arg3_1, arg4_1, arg5_1, arg6_1 = args
    args.clear()
    s0 = arg1_1
    assert_size_stride(arg0_1, (s0, 3584), (3584, 1))
    assert_size_stride(arg2_1, (3584, ), (1, ))
    assert_size_stride(arg3_1, (4608, 3584), (3584, 1))
    assert_size_stride(arg4_1, (4608, ), (1, ))
    assert_size_stride(arg5_1, (32768, 128), (128, 1))
    assert_size_stride(arg6_1, (3, s0), (s0, 1)) # exception thrown here

exception thrown here:

[rank0]: Traceback (most recent call last):
...
[rank0]:   File "/data/home/imkero/workspace/git/vllm-project/vllm/vllm/v1/worker/gpu_worker.py", line 134, in determine_num_available_blocks
[rank0]:     self.model_runner.profile_run()
[rank0]:   File "/data/home/imkero/workspace/git/vllm-project/vllm/vllm/v1/worker/gpu_model_runner.py", line 836, in profile_run
[rank0]:     hidden_states = self._dummy_run(self.model, self.max_num_tokens,
[rank0]:                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
...
[rank0]:   File "/data/home/imkero/workspace/git/vllm-project/vllm/vllm/model_executor/models/qwen2_vl.py", line 1205, in forward
[rank0]:     hidden_states = self.language_model.model(
[rank0]:                     ^^^^^^^^^^^^^^^^^^^^^^^^^^
...
[rank0]:   File "/data/home/imkero/.cache/vllm/torch_compile_cache/b06b53d306/rank_0/inductor_cache/il/ciloq6dlbeaiigv2l2eqyzyv7qihhat45i5xy4fjrhxyvuj3s6sb.py", line 319, in call
[rank0]:     assert_size_stride(arg6_1, (3, s0), (s0, 1))
[rank0]: AssertionError: expected size 3==3, stride 512==504 at dim=0

This PR's approach

source code:

self.mrope_positions = torch.zeros((self.max_num_tokens, 3),
                                    dtype=torch.int64,
                                    device=self.device)
self.mrope_positions = self.mrope_positions.permute((1, 0))

compiled code:

def call(args):
    arg0_1, arg1_1, arg2_1, arg3_1, arg4_1, arg5_1, arg6_1 = args
    args.clear()
    s0 = arg1_1
    assert_size_stride(arg0_1, (s0, 3584), (3584, 1))
    assert_size_stride(arg2_1, (3584, ), (1, ))
    assert_size_stride(arg3_1, (4608, 3584), (3584, 1))
    assert_size_stride(arg4_1, (4608, ), (1, ))
    assert_size_stride(arg5_1, (32768, 128), (128, 1))
    assert_size_stride(arg6_1, (3, s0), (1, 3))

[youkaichao]

thanks for the explanation! I think it might be a torch.compile 's bug, it should handle non-contiguous tensors as well. Let's take this work-around first, and investigate later.

[imkero]

thanks for the explanation! I think it might be a torch.compile 's bug, it should handle non-contiguous tensors as well. Let's take this work-around first, and investigate later.

@youkaichao Here I found another workaround (maybe a direct approach?) that would make torch.compile "realize" that mrope_positions tensor is not always contiguous, and then works with mrope_positions' original shape (3, seq_len)

Approach

self.mrope_positions = torch.zeros((3, self.max_num_tokens + 1),
                                   dtype=torch.int64,
                                   device=self.device)

how it works:

• extend the size of mrope_positions in dim 1 (plus 1 on it)
• mrope_positions will be non-contiguous always in the upcoming usage (self.mrope_positions[:, :total_num_scheduled_tokens]) since total_num_scheduled_tokens is always less than max_num_tokens + 1
  • and torch.compile will notice it while compiling

Compiled code

def call(args):
    arg0_1, arg1_1, arg2_1, arg3_1, arg4_1, arg5_1, arg6_1, arg7_1 = args
    args.clear()
    s0 = arg1_1
    s2 = arg7_1
    assert_size_stride(arg0_1, (s0, 3584), (3584, 1))
    assert_size_stride(arg2_1, (3584, ), (1, ))
    assert_size_stride(arg3_1, (4608, 3584), (3584, 1))
    assert_size_stride(arg4_1, (4608, ), (1, ))
    assert_size_stride(arg5_1, (32768, 128), (128, 1))
    assert_size_stride(arg6_1, (3, s0), (s2, 1))

we can discover that:

• in the compiled code, the stride_of_dim_0 (s2) becomes independent of size_of_dim_1 (s0), which means non-contiguous tensor supported.

[youkaichao]

this is amazing! cc @ywang96 if we can make the change, it will make the code look simpler.

[ywang96]

this is amazing! cc @ywang96 if we can make the change, it will make the code look simpler.

Yep - changed in #12352

[Isotr0py]

LGTM now!

[ywang96]

Looks like the tensor splitting breaks the embeddings as input interface, so I think we will need to move the split out of the encoder forward pass.

[DarkLight1337]

@imkero it would be great if you have time to help look into the failing embedding tests (since you originally created them)

[DarkLight1337]

Nvm I think I found the issue, solving it now.