[Bugfix] Fix broadcasting logic for multi_modal_kwargs

[DarkLight1337]

Currently multi_modal_kwargs are placed on GPU before distributed broadcast. Since the broadcast operation only allows flat dictionaries of tensors, custom logic is required to handle nested tensors. Currently, it only handles singly-nested dictionaries and fails to consider nested lists of tensors. It is complicated and even inefficient to extend this to arbitrarily nested structures such as interleaved nested dictionaries and lists.

This PR addresses the root issue by keeping the tensors in CPU before the broadcast operation, only moving them to GPU right before they are inputted to the model. This lets us broadcast them as opaque Python objects, removing the need to handle nested structures.

FIX #5905 (comment)
FIX #6946

cc @xwjiang2010 @youkaichao

[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 consists a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of default ones by unblocking the steps in your fast-check build on Buildkite UI.

Once the PR is approved and ready to go, please make sure to run full CI as it is required to merge (or just use auto-merge).

To run full CI, you can do one of these:

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

🚀

[DarkLight1337]

Closing as superseded by #6925.

Reopening to address the comment from @youkaichao .

[youkaichao]

thanks for the change! will take a look tmrw.

[xwjiang2010]

Hmm, does this have performance impact? Like given a CPU tensor, how do you compare the following

• CPU-side broadcasting followed by multiple CPU-to-GPU transfers
• Single CPU-to-GPU transfer followed by NCCL GPU-to-GPU broadcast

I feel like on a single node with high bandwidth NCCL, option 2 may be faster?

[youkaichao]

LGTM, thanks for the refactor to make our life easier!

[youkaichao]

@xwjiang2010 the tensor here to broadcast is small, and NCCL does not have a win here. It only adds code complexity.

[xwjiang2010]

@xwjiang2010 the tensor here to broadcast is small, and NCCL does not have a win here. It only adds code complexity.

yeah while I am happy that I don't have to think of a way of making nested list of nested dict of nested list work,
but just want to make sure that we all on the same page that these tensors are small.

For example, are they small when there are dozens of queries in a batch? Or multiple images per query? Or when we are sending in embeddings (in the future)? Or when factors like these are compounding together?

[youkaichao]

In the future we will only broadcast CPU data, without sending any GPU data via broadcast. GPU data should live in the process that produces it.

[xwjiang2010]

And this also applies to things like input_ids right?

[youkaichao]

Yes, in the future input_ids live in worker process, and we just need to append new generated tokens, without re-constructing all the input_ids at every iteration.