[V1][Performance] Implement custom serializaton for MultiModalKwargs

[p88h]

FIX #16185

WIP, just handles the basic case of simple str->Tensor dict.

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

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 either: Add ready label to the PR or enable auto-merge.

🚀

[njhill]

Thanks @p88h, this is pretty much what I had planned! It would be great if you could take this on though.

We are planning to disable the use of pickle by default, so it would be good to do this in a way at avoid that. For handling tensors/ndarrays in general, I have a PR #13790 which also eliminates some mem copies, so it would be good to base on top of that (I'll try to get it merged asap, just need to add a unit test).

Unfortunately msgspec doesn't support custom nested types, but I thought to have intermediate types like:

@dataclass
class FlatNestedTensors:
    tensors: list[torch.Tensor]
    structure: list[Union[int, list]]

where structure is a nested list of indexes into the tensors list. And then the tensors should automatically be handled by the zero-copy logic.

[njhill]

@p88h I've added a test to #13790 so we can hopefully merge it shortly.

[p88h]

I added a comment on your PR, I think it can be simplified to sth like this:
https://gist.github.com/p88h/1daec6374c35293f6bced9333d6f2c4c

... if it worked. That inner msgspec serialization is not working as expected.

[p88h]

So when encoding msgspec is actually smart enough to just handle everything recursively.
Unfortunately when decoding not everything is automatic, so some fix-up is necessary - but a lot of simpler now.

Well, except...

_items_by_modality doesn't want to be easily serialized.

For now, I implemented something that works, unfortunately, the MultiModalField part requires a pickle - msgpack doesn't even notice it's a custom type, just encodes the whole MultiModalFieldElem as a dict and just ignores that part.

It also adds additional complexity to serializing MultiModalKwargs, too - since the old and new layout overlap (share tensors), we likely don't want to send that over twice, even with zero copy that seems unnecesary. The approach is now to use either items by modality (and reconstruct UserDict) OR just pass through the dict for V0-style usage (which I think is still in use even in V1)

[njhill]

Thanks @p88h I will review again once #13790 is merged (imminently) and this is rebased on top of that.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @p88h.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[p88h]

@njhill @DarkLight1337 PTAL #16432