[V1][Performance] Implement custom serializaton for MultiModalKwargs [Rebased]

[p88h]

FIX #16185 (link existing issues this PR will resolve)

This is a rebase of #16279 which had too entangled commits.
Implements additional handling of MultimodalKwargs on top of #13790
Further improves memory usage on top of improvements in #16273 by another 50%

[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.

🚀

[ywang96]

@p88h This is amazing! Have you tried running some benchmarks to see the throughput performance impact of this PR?

[p88h]

@ywang96 I've added a benchmark table to the linked bug #16185

My benchmark focused on memory performance rather than throughput, and only used a single model. It should not really change throughput that much other than in cases that do run into memory issues, though.

I'll try running some throughput checks tomorrow

[njhill]

Thanks @p88h! I think this looks good.

The main thing I think is to add custom serialization for the field. And we'll probably want to add a few more comments since it's tightly coupled with the custom tensor encoding format.

Also, I haven't looked closely at the entire flow, but in the case of MMKs created from items, it might make sense to defer the population of their data (via the "reduce" operations). Since that will be repeated in the receiving process and causes extra cpu and mem overhead since tensors may get stacked etc. It would be nice if there was a way for this to happen lazily but I guess that depends on how the data is later accessed.

cc @ywang96 @DarkLight1337

[njhill]

Also, I haven't looked closely at the entire flow, but in the case of MMKs created from items, it might make sense to defer the population of their data (via the "reduce" operations). Since that will be repeated in the receiving process and causes extra cpu and mem overhead since tensors may get stacked etc. It would be nice if there was a way for this to happen lazily but I guess that depends on how the data is later accessed.

FYI I've opened another PR to help with this: #16440. It should in theory help all of the cases not just the multi-proc case.

It would still be additionally beneficial to postpone doing this reduce operation until after being transferred to the engine though.

[p88h]

I have some experimental data with this PR in place. Interestingly it performs much better with zero-copy disabled

In this new benchmark,` I am feeding gradually increasing document sets to the engine. Turns out custom serialization helps less than expected - I think previously it was augmented by the cache, but now all files are unique so results are a bit different.

The 'mix' performance case measures running all prompts together (15 total, with 128 images total) after they have been initially processed one-by-one, so it's expected that it's performing much better / cached.

config / benchmark case       | 4 images | 8 images | 16 images | 32 images | t.max | t.mix
------------------------------+----------+----------+-----------+-----------+-------+-------
baseline (zero-copy disabled) | 3.55 GB  | 5.11 GB  | 9.96 GB   | 22.54 GB  | 90.4s | 44.1s
baseline (zero-copy enabled)  | 3.50 GB  | 5.01 GB  | 9.87 GB   | 22.56 GB  | 75.3s | 39.4s
#16432 (zero-copy enabled)    | 3.40 GB  | 4.75 GB  | 8.53 GB   | 22.02 GB  | 13.8s | 36.1s
#16432 (zero-copy disabled)   | 3.28 GB  | 3.95 GB  | 4.76 GB   | 5.85 GB   | 14.4s | 36.3s

[njhill]

Thanks for the great work @p88h!

[njhill]

Looks like a CI test is failing - but unfortunately the root cause is obscured (the OOM failure of the subsequent test is a result of improper cleanup after the original failure). This should hopefully be addressed by #11737.

In the meantime I can try running this test locally.

p.s. there's no need to keep rebasing on latest main, this just causes all the tests to start over.

[njhill]

It turns out it was because sometimes MMKwargs can contain non-tensor data (specifically "second_per_grid_ts": [1.0] in this case). So I pushed an update to allow floats and ints too.

[p88h]

Thank you ! I was about to go back to debugging this morning ;)