Repository: vllm-project/vllm
Pull request #16432: [V1][Performance] Implement custom serializaton for MultiModalKwargs [Rebased]
Author: login: p88h
Labels: name: documentation
description: Improvements or additions to documentation
name: performance
description: Performance-related issues
name: ready
description: ONLY add when PR is ready to merge/full CI is needed
name: ci/build
name: v1
name: multi-modality
description: Related to multi-modality (#4194)
Created: 2025-04-10T21:02:18Z | Merged: 2025-04-17T02:28:32Z

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% 



Linked issue #16185: [Bug]: Huge memory overhead with V1 (multiprocessing) when handling several multimodal inputs
### Your current environment

<details>
<summary>The output of `python collect_env.py`</summary>

```text
INFO 04-07 10:56:24 [__init__.py:239] Automatically detected platform cuda.
Collecting environment information...
PyTorch version: 2.6.0+cu124
Is debug build: False
CUDA used to build PyTorch: 12.4
ROCM used to build PyTorch: N/A

OS: Ubuntu 24.04.2 LTS (x86_64)
GCC version: (Ubuntu 13.3.0-6ubuntu2~24.04) 13.3.0
Clang version: 18.1.3 (1ubuntu1)
CMake version: version 3.28.3
Libc version: glibc-2.39

Python version: 3.12.8 | packaged by conda-forge | (main, Dec  5 2024, 14:24:40) [GCC 13.3.0] (64-bit runtime)
Python platform: Linux-6.6.75.1-microsoft-standard-WSL2-x86_64-with-glibc2.39
Is CUDA available: True
CUDA runtime version: 12.0.140
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA GeForce RTX 4080
Nvidia driver version: 572.47
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.9.2
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.9.2
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.9.2
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.9.2
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.9.2
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.9.2
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.9.2
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                         x86_64
CPU op-mode(s):                       32-bit, 64-bit
Address sizes:                        46 bits physical, 48 bits virtual
Byte Order:                           Little Endian
CPU(s):                               24
On-line CPU(s) list:                  0-23
Vendor ID:                            GenuineIntel
Model name:                           13th Gen Intel(R) Core(TM) i7-13700K
CPU family:                           6
Model:                                183
Thread(s) per core:                   2
Core(s) per socket:                   12
Socket(s):                            1
Stepping:                             1
BogoMIPS:                             6835.19
Flags:                                fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ss ht syscall nx pdpe1gb rdtscp lm constant_tsc rep_good nopl xtopology tsc_reliable nonstop_tsc cpuid tsc_known_freq pni pclmulqdq vmx ssse3 fma cx16 sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand hypervisor lahf_lm abm 3dnowprefetch ssbd ibrs ibpb stibp ibrs_enhanced tpr_shadow ept vpid ept_ad fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 xsaves avx_vnni vnmi umip waitpkg gfni vaes vpclmulqdq rdpid movdiri movdir64b fsrm md_clear serialize flush_l1d arch_capabilities
Virtualization:                       VT-x
Hypervisor vendor:                    Microsoft
Virtualization type:                  full
L1d cache:                            576 KiB (12 instances)
L1i cache:                            384 KiB (12 instances)
L2 cache:                             24 MiB (12 instances)
L3 cache:                             30 MiB (1 instance)
NUMA node(s):                         1
NUMA node0 CPU(s):                    0-23
Vulnerability Gather data sampling:   Not affected
Vulnerability Itlb multihit:          Not affected
Vulnerability L1tf:                   Not affected
Vulnerability Mds:                    Not affected
Vulnerability Meltdown:               Not affected
Vulnerability Mmio stale data:        Not affected
Vulnerability Reg file data sampling: Mitigation; Clear Register File
Vulnerability Retbleed:               Mitigation; Enhanced IBRS
Vulnerability Spec rstack overflow:   Not affected
Vulnerability Spec store bypass:      Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1:             Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:             Mitigation; Enhanced / Automatic IBRS; IBPB conditional; RSB filling; PBRSB-eIBRS SW sequence; BHI BHI_DIS_S
Vulnerability Srbds:                  Not affected
Vulnerability Tsx async abort:        Not affected

Versions of relevant libraries:
[pip3] numpy==1.26.4
[pip3] nvidia-cublas-cu12==12.4.5.8
[pip3] nvidia-cuda-cupti-cu12==12.4.127
[pip3] nvidia-cuda-nvrtc-cu12==12.4.127
[pip3] nvidia-cuda-runtime-cu12==12.4.127
[pip3] nvidia-cudnn-cu12==9.1.0.70
[pip3] nvidia-cufft-cu12==11.2.1.3
[pip3] nvidia-curand-cu12==10.3.5.147
[pip3] nvidia-cusolver-cu12==11.6.1.9
[pip3] nvidia-cusparse-cu12==12.3.1.170
[pip3] nvidia-cusparselt-cu12==0.6.2
[pip3] nvidia-ml-py==12.570.86
[pip3] nvidia-nccl-cu12==2.21.5
[pip3] nvidia-nvjitlink-cu12==12.4.127
[pip3] nvidia-nvtx-cu12==12.4.127
[pip3] onnx==1.17.0
[pip3] onnxruntime==1.20.1
[pip3] pytorch-lightning==2.5.0.post0
[pip3] pytorch-metric-learning==2.8.1
[pip3] pyzmq==26.2.1
[pip3] torch==2.6.0
[pip3] torch-audiomentations==0.12.0
[pip3] torch_pitch_shift==1.2.5
[pip3] torchaudio==2.6.0
[pip3] torchmetrics==1.6.1
[pip3] torchvision==0.21.0
[pip3] transformers==4.50.3
[pip3] triton==3.2.0
[conda] numpy                     1.26.4                   pypi_0    pypi
[conda] nvidia-cublas-cu12        12.4.5.8                 pypi_0    pypi
[conda] nvidia-cuda-cupti-cu12    12.4.127                 pypi_0    pypi
[conda] nvidia-cuda-nvrtc-cu12    12.4.127                 pypi_0    pypi
[conda] nvidia-cuda-runtime-cu12  12.4.127                 pypi_0    pypi
[conda] nvidia-cudnn-cu12         9.1.0.70                 pypi_0    pypi
[conda] nvidia-cufft-cu12         11.2.1.3                 pypi_0    pypi
[conda] nvidia-curand-cu12        10.3.5.147               pypi_0    pypi
[conda] nvidia-cusolver-cu12      11.6.1.9                 pypi_0    pypi
[conda] nvidia-cusparse-cu12      12.3.1.170               pypi_0    pypi
[conda] nvidia-cusparselt-cu12    0.6.2                    pypi_0    pypi
[conda] nvidia-ml-py              12.570.86                pypi_0    pypi
[conda] nvidia-nccl-cu12          2.21.5                   pypi_0    pypi
[conda] nvidia-nvjitlink-cu12     12.4.127                 pypi_0    pypi
[conda] nvidia-nvtx-cu12          12.4.127                 pypi_0    pypi
[conda] pytorch-lightning         2.5.0.post0              pypi_0    pypi
[conda] pytorch-metric-learning   2.8.1                    pypi_0    pypi
[conda] pyzmq                     26.2.1                   pypi_0    pypi
[conda] torch                     2.6.0                    pypi_0    pypi
[conda] torch-audiomentations     0.12.0                   pypi_0    pypi
[conda] torch-pitch-shift         1.2.5                    pypi_0    pypi
[conda] torchaudio                2.6.0                    pypi_0    pypi
[conda] torchmetrics              1.6.1                    pypi_0    pypi
[conda] torchvision               0.21.0                   pypi_0    pypi
[conda] transformers              4.50.3                   pypi_0    pypi
[conda] triton                    3.2.0                    pypi_0    pypi
ROCM Version: Could not collect
Neuron SDK Version: N/A
vLLM Version: 0.8.3.dev212+g58e234a75.d20250405
vLLM Build Flags:
CUDA Archs: Not Set; ROCm: Disabled; Neuron: Disabled
GPU Topology:
GPU0    CPU Affinity    NUMA Affinity   GPU NUMA ID
GPU0     X                              N/A

Legend:

  X    = Self
  SYS  = Connection traversing PCIe as well as the SMP interconnect between NUMA nodes (e.g., QPI/UPI)
  NODE = Connection traversing PCIe as well as the interconnect between PCIe Host Bridges within a NUMA node
  PHB  = Connection traversing PCIe as well as a PCIe Host Bridge (typically the CPU)
  PXB  = Connection traversing multiple PCIe bridges (without traversing the PCIe Host Bridge)
  PIX  = Connection traversing at most a single PCIe bridge
  NV#  = Connection traversing a bonded set of # NVLinks

NCCL_CUMEM_ENABLE=0
TORCHINDUCTOR_COMPILE_THREADS=1
CUDA_MODULE_LOADING=LAZY
```

</details>


### 🐛 Describe the bug

This should be reproducible with QWEN VL 2.5 and using `vision_language_multi_image.py` offline inference.
When configured to use several images as input (e.g. just multiply `IMAGE_URLS` 8-16 times), the *CPU* memory usage spikes dramatically. With just around 20 or so images, VLLM will try to consume around `20-30GB` of RAM, with 40 it will get into `50-60GB` range. 

It's interesting that the _number_ of passed elements seem to be a problem,  while the _size_ causes less problems (ie. when merging multiple images into one, without scaling, it's possible to send ~60 images via 4 collages and still fit within 30GB). This happens regardless of whether qwen-vl-utils are installed to resize the images.

At the same time request preprocessing starts to slow down, and while profiling was not super useful due to multiprocessing it did help as the next step was disabling that. With `VLLM_ENABLE_V1_MULTIPROCESSING=0`  the issue completely disappears, processing delays are gone and even with ~100 input files the memory usage is in low GBs. 

I haven't tried profiling MessageQueue yet, but perhaps someone will also run into this, or maybe already ran into?

BTW - the environment above was collected running within WSL, but the OS doesn't seem to be a significant factor here - while WSL does run into memory issues earlier than raw Linux, running on raw Linux exhibits the same behavior. 

### Before submitting a new issue...

- [x] Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the [documentation page](https://docs.vllm.ai/en/latest/), which can answer lots of frequently asked questions.
Linked issue comments: user: login: DarkLight1337
created_at: 2025-04-07T12:31:28Z
body: Can you see if `--disable-mm-preprocessor-cache` alleviates the issue?
author_association: MEMBER
user: login: sfc-gh-spasko
created_at: 2025-04-07T12:55:55Z
body: `--disable-mm-preprocessor-cache` doesn't seem to affect this.
author_association: NONE
user: login: DarkLight1337
created_at: 2025-04-07T12:59:24Z
body: cc @ywang96 @robertgshaw2-redhat maybe it's related to transferring the multimodal inputs between processes?
author_association: MEMBER
user: login: njhill
created_at: 2025-04-07T15:46:07Z
body: > cc [@ywang96](https://github.com/ywang96) [@robertgshaw2-redhat](https://github.com/robertgshaw2-redhat) maybe it's related to transferring the multimodal inputs between processes?

We have some in-progress improvements for this aspect.
author_association: MEMBER
user: login: DarkLight1337
created_at: 2025-04-08T15:42:08Z
body: Can you try out https://github.com/vllm-project/vllm/pull/16273 and see if it improves the memory overhead?
author_association: MEMBER
user: login: p88h
created_at: 2025-04-08T16:50:00Z
body: It's *much* better ... but it's unfortunately still pretty much unusable. 

With muliprocessing, handling ~64 1Mpix images  used to exceed 52GB and OOM at this point (on a 64GB system...)
After the patch, and with mm_preprocessor_cache enabled, it now fits within 37GB and runs correctly. 
At 32 images, before the patch it peaks at 37GB, and with the patch (& preproc cache) it's closer to 18GB.

So while that's likely a 2x improvement, it's still very slow in the preparation phase and the memory use is still very significant.
Without multiprocessing the same workload requires only <4GB. 

~16GB memory load to process ~128MB worth of data seems pretty steep. 
author_association: CONTRIBUTOR
user: login: DarkLight1337
created_at: 2025-04-08T16:51:49Z
body: How much memory does vLLM use if you disable preprocessor cache?
author_association: MEMBER
user: login: p88h
created_at: 2025-04-08T16:55:37Z
body: If the preprocessor cache is disabled, the patch has no effect. 
Without the patch, on 32 images,. it consumes ~4GB regardless of cache setting in single-process mode, and ~37GB in multi-process mode. 
With the patch, on 32 images, it consumes 37/18 GB in multi-process mode depending on the cache setting, and ~4GB in single-process mode. 

There may be some difference at that ~4GB level since i was mostly looking at the multi-GB consumption. 
author_association: CONTRIBUTOR
user: login: DarkLight1337
created_at: 2025-04-08T16:57:38Z
body: Quite surprising that disabling the cache actually increases the memory overhead by that much. I guess @njhill 's WIP should address that. Thanks for reporting the results!
author_association: MEMBER
user: login: p88h
created_at: 2025-04-08T16:58:44Z
body: Well, the patch depends on the items to be in the cache, no ? So I would assume if the cache disabled, it will revert to the pre-patch behavior, basically. 
author_association: CONTRIBUTOR
user: login: DarkLight1337
created_at: 2025-04-08T16:59:28Z
body: Yes, you're correct. So the main problem is still about how the data is being transferred in multiprocessing mode.
author_association: MEMBER
user: login: p88h
created_at: 2025-04-08T18:12:29Z
body: Oh, okay, it seems that serialization itself is not really that big of the problem, but the actual issue was in _encoding_ the multimodal args properly.

This POC seems to fix the problem:
https://github.com/vllm-project/vllm/pull/16279


author_association: CONTRIBUTOR
user: login: njhill
created_at: 2025-04-08T18:27:46Z
body: Thanks @p88h, aren't encoding and serialization essentially the same thing? :)

This is similar to what I was planning - it just needs to be a bit more general since AFAIK `NestedTensors` may not always be a single tensor and it has a `_items_by_modality` field. And combining with https://github.com/vllm-project/vllm/pull/13790 to be even more efficient.

Oh and also to do this in a way that avoids using pickle since we're planning to disable the use of pickle by default (that was actually a secondary motivation). cc @russellb 
author_association: MEMBER
user: login: p88h
created_at: 2025-04-08T18:44:08Z
body: In this case they really aren't the same - just as in the previous similar workaround for bare torch.Tensor, serializing it via pickle uses print format, not bytes.  I agree something more generic that can just do this in zero-copy fashion (and without pickles) would be preferable, I will have a look at your PR, updated mine to handle NestedTensors properly in case it helps.
author_association: CONTRIBUTOR
user: login: p88h
created_at: 2025-04-10T10:27:57Z
body: I completed a simple benchmark to quantify the problems along the impact of PRs affecting this. Results below;


All times Best of 2 runs each, memory averaged
Time = elapsed realtime via 'time'. 
Memory = max RSS reported via 'time'

Note 32 images has 50% duplication, 64 images has 75% duplication (the image files are repeated 2/4 times). Each image is about 1Mpix, and generates ~1000 tokens. 

Model used: Qwen/Qwen2.5-VL-7B-Instruct-AWQ. Output limited to 1024 tokens. 

Times marked with * mean the model hallucinates and actually hits the limit.
Times marked with ! mean the model truncated the response.
This is less relevant for 64 image case, as that requires about these many tokens.

```
patches and parameters applied             | 16 images  | 32 images  | 64 images   |
-------------------------------------------+------------+------------+-------------+
baseline, cache on, single                 | 5.2G 33s   | 5.3G 53s*  | 5.8G 69s*   |
baseline, cache off, single                | 5.2G 34s   | 5.3G 54s*  | 5.4G 68s*   |
baseline, cache on, multi                  | 11.3G 44s  | 39.5G 92s* | OOM         |
baseline, cache off, multi                 | 11.3G 43s  | 39.5G 93s* | OOM         |
#16273, cache on, multi                    | 10.2G 45s  | 18.6G 72s* | 35.5G 107s* |
#16273, cache off, multi                   | 11.4G 46s  | 39.5G 94s* | OOM         |
#16273 + #13790, cache on, multi           | 10.2G 46s  | 18.6G 73s* | 35.5G 108s* |
#16273 + #13790, cache off, multi          | 11.4G 44s  | 39.5G 96s* | OOM         |
#16273 + #13790 + #16279, cache on, multi  | 5.4G 34s   | 9.9G 53s*  | 18.6G 71s*  |
#16273 + #13790 + #16279, cache off, multi | 6.0G 33s   | 20.2G 53s* | OOM         |
```

Note that disabling cache now effectively disables the effect of #16273 when multi-processing. All of these changes have no impact on single-process mode. #16279 is built on top of #13790 now so can't be deployed separately.

with combined communication improvements, this is now _usable_ but there is still something generating a ton of memory use that it does not in the single-threaded version. 

author_association: CONTRIBUTOR
user: login: DarkLight1337
created_at: 2025-04-10T10:34:03Z
body: Just to be sure, what are you referring to by "baseline" here? Is it main branch?
author_association: MEMBER
user: login: p88h
created_at: 2025-04-10T10:49:31Z
body: @DarkLight1337 baseline = main branch @24f6b9a71397539a3d02c801963220b0e9a2aef9 (=yesterday, before #16273 was applied)
author_association: CONTRIBUTOR
user: login: njhill
created_at: 2025-04-10T23:48:32Z
body: Thanks @p88h, this is great! It's expected that https://github.com/vllm-project/vllm/pull/13790 by itself doesn't change much but I expect that the current combination of https://github.com/vllm-project/vllm/pull/13790 with your PR should give a significant reduction over your original PR without it. I understand that to test that you'd have to revert to an older version of the PR though. An easy way to measure the "zero copy" benefits would be to disable that by changing this line in `serial_utils.py`:
```python
        if not obj.shape or obj.nbytes < INLINE_BUF_SIZE_THRESHOLD:
```
to
```python
        if True:
```
author_association: MEMBER
user: login: p88h
created_at: 2025-04-11T06:17:05Z
body: Interesting... 

With inline mode, I do get warnings about read-only Tensors. 
Can either ignore them or copy memory out in this case. 

But that's not the interesting bit. It seems that the overall memory usage is significantly lower with inline mode, and performance is back to the single-thread levels, at least for the one test case I've ran so far. 

I'll add some benchmark results to the PR and figure out a way to make this work efficiently by default. 
author_association: CONTRIBUTOR
user: login: p88h
created_at: 2025-04-11T11:08:01Z
body: I added benchmark  results in the PR. 
I've uploaded the benchmark I wrote for this here: https://github.com/p88h/fake-vqa

author_association: CONTRIBUTOR

Diff:
diff --git a/tests/v1/test_serial_utils.py b/tests/v1/test_serial_utils.py
index bc0e0cbd85e1..e58d3c403c19 100644
--- a/tests/v1/test_serial_utils.py
+++ b/tests/v1/test_serial_utils.py
@@ -1,10 +1,16 @@
 # SPDX-License-Identifier: Apache-2.0
 from collections import UserDict
 from dataclasses import dataclass
+from typing import Optional
 
+import msgspec
 import numpy as np
 import torch
 
+from vllm.multimodal.inputs import (MultiModalBatchedField,
+                                    MultiModalFieldElem, MultiModalKwargs,
+                                    MultiModalKwargsItem,
+                                    MultiModalSharedField, NestedTensors)
 from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder
 
 
@@ -50,7 +56,7 @@ def test_encode_decode():
         large_non_contig_tensor=torch.rand(1024, 512)[:, 10:20],
     )
 
-    encoder = MsgpackEncoder()
+    encoder = MsgpackEncoder(size_threshold=256)
     decoder = MsgpackDecoder(MyType)
 
     encoded = encoder.encode(obj)
@@ -78,6 +84,97 @@ def test_encode_decode():
     assert_equal(decoded2, obj)
 
 
+class MyRequest(msgspec.Struct):
+    mm: Optional[list[MultiModalKwargs]]
+
+
+def test_multimodal_kwargs():
+    d = {
+        "foo":
+        torch.zeros(20000, dtype=torch.float16),
+        "bar": [torch.zeros(i * 1000, dtype=torch.int8) for i in range(3)],
+        "baz": [
+            torch.rand((256), dtype=torch.float16),
+            [
+                torch.rand((1, 12), dtype=torch.float32),
+                torch.rand((3, 5, 7), dtype=torch.float64),
+            ], [torch.rand((4, 4), dtype=torch.float16)]
+        ],
+    }
+
+    # pack mm kwargs into a mock request so that it can be decoded properly
+    req = MyRequest(mm=[MultiModalKwargs(d)])
+
+    encoder = MsgpackEncoder()
+    decoder = MsgpackDecoder(MyRequest)
+
+    encoded = encoder.encode(req)
+
+    assert len(encoded) == 6
+
+    total_len = sum(memoryview(x).cast("B").nbytes for x in encoded)
+
+    # expected total encoding length, should be 44536, +-20 for minor changes
+    assert total_len >= 44516 and total_len <= 44556
+    decoded: MultiModalKwargs = decoder.decode(encoded).mm[0]
+    assert all(nested_equal(d[k], decoded[k]) for k in d)
+
+
+def test_multimodal_items_by_modality():
+    e1 = MultiModalFieldElem("audio", "a0", torch.zeros(1000,
+                                                        dtype=torch.int16),
+                             MultiModalBatchedField())
+    e2 = MultiModalFieldElem(
+        "video",
+        "v0",
+        [torch.zeros(1000, dtype=torch.int8) for _ in range(4)],
+        MultiModalBatchedField(),
+    )
+    e3 = MultiModalFieldElem("image", "i0", torch.zeros(1000,
+                                                        dtype=torch.int32),
+                             MultiModalSharedField(4))
+    e4 = MultiModalFieldElem("image", "i1", torch.zeros(1000,
+                                                        dtype=torch.int32),
+                             MultiModalBatchedField())
+    audio = MultiModalKwargsItem.from_elems([e1])
+    video = MultiModalKwargsItem.from_elems([e2])
+    image = MultiModalKwargsItem.from_elems([e3, e4])
+    mm = MultiModalKwargs.from_items([audio, video, image])
+
+    # pack mm kwargs into a mock request so that it can be decoded properly
+    req = MyRequest([mm])
+
+    encoder = MsgpackEncoder()
+    decoder = MsgpackDecoder(MyRequest)
+
+    encoded = encoder.encode(req)
+
+    assert len(encoded) == 8
+
+    total_len = sum(memoryview(x).cast("B").nbytes for x in encoded)
+
+    # expected total encoding length, should be 14255, +-20 for minor changes
+    assert total_len >= 14235 and total_len <= 14275
+    decoded: MultiModalKwargs = decoder.decode(encoded).mm[0]
+
+    # check all modalities were recovered and do some basic sanity checks
+    assert len(decoded.modalities) == 3
+    images = decoded.get_items("image")
+    assert len(images) == 1
+    assert len(images[0].items()) == 2
+    assert list(images[0].keys()) == ["i0", "i1"]
+
+    # check the tensor contents and layout in the main dict
+    assert all(nested_equal(mm[k], decoded[k]) for k in mm)
+
+
+def nested_equal(a: NestedTensors, b: NestedTensors):
+    if isinstance(a, torch.Tensor):
+        return torch.equal(a, b)
+    else:
+        return all(nested_equal(x, y) for x, y in zip(a, b))
+
+
 def assert_equal(obj1: MyType, obj2: MyType):
     assert torch.equal(obj1.tensor1, obj2.tensor1)
     assert obj1.a_string == obj2.a_string
diff --git a/vllm/envs.py b/vllm/envs.py
index f80bf878f79c..d32968c3d173 100644
--- a/vllm/envs.py
+++ b/vllm/envs.py
@@ -107,6 +107,7 @@
     VLLM_TPU_BUCKET_PADDING_GAP: int = 0
     VLLM_USE_DEEP_GEMM: bool = False
     VLLM_XGRAMMAR_CACHE_MB: int = 0
+    VLLM_MSGPACK_ZERO_COPY_THRESHOLD: int = 256
 
 
 def get_default_cache_root():
@@ -704,6 +705,16 @@ def maybe_convert_int(value: Optional[str]) -> Optional[int]:
     # It can be changed with this variable if needed for some reason.
     "VLLM_XGRAMMAR_CACHE_MB":
     lambda: int(os.getenv("VLLM_XGRAMMAR_CACHE_MB", "512")),
+
+    # Control the threshold for msgspec to use 'zero copy' for
+    # serialization/deserialization of tensors. Tensors below
+    # this limit will be encoded into the msgpack buffer, and
+    # tensors above will instead be sent via a separate message.
+    # While the sending side still actually copies the tensor
+    # in all cases, on the receiving side, tensors above this
+    # limit will actually be zero-copy decoded.
+    "VLLM_MSGPACK_ZERO_COPY_THRESHOLD":
+    lambda: int(os.getenv("VLLM_MSGPACK_ZERO_COPY_THRESHOLD", "256")),
 }
 
 # end-env-vars-definition
diff --git a/vllm/v1/serial_utils.py b/vllm/v1/serial_utils.py
index 3af6793fde74..4f7987ee46a6 100644
--- a/vllm/v1/serial_utils.py
+++ b/vllm/v1/serial_utils.py
@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 
+import dataclasses
 import pickle
 from collections.abc import Sequence
 from inspect import isclass
@@ -12,12 +13,26 @@
 import zmq
 from msgspec import msgpack
 
+from vllm import envs
+from vllm.multimodal.inputs import (BaseMultiModalField,
+                                    MultiModalBatchedField,
+                                    MultiModalFieldConfig, MultiModalFieldElem,
+                                    MultiModalFlatField, MultiModalKwargs,
+                                    MultiModalKwargsItem,
+                                    MultiModalSharedField, NestedTensors)
+
 CUSTOM_TYPE_PICKLE = 1
 CUSTOM_TYPE_CLOUDPICKLE = 2
 CUSTOM_TYPE_RAW_VIEW = 3
 
-# TODO calibrate this size
-MIN_NOCOPY_BUF_SIZE = 512
+# MultiModalField class serialization type map.
+# These need to list all possible field types and match them
+# to factory methods in `MultiModalFieldConfig`.
+MMF_CLASS_TO_FACTORY: dict[type[BaseMultiModalField], str] = {
+    MultiModalFlatField: "flat",
+    MultiModalSharedField: "shared",
+    MultiModalBatchedField: "batched",
+}
 
 bytestr = Union[bytes, bytearray, memoryview, zmq.Frame]
 
@@ -27,14 +42,20 @@ class MsgpackEncoder:
 
     Note that unlike vanilla `msgspec` Encoders, this interface is generally
     not thread-safe when encoding tensors / numpy arrays.
+
+    By default, arrays below 256B are serialized inline Larger will get sent 
+    via dedicated messages. Note that this is a per-tensor limit.
     """
 
-    def __init__(self):
+    def __init__(self, size_threshold: Optional[int] = None):
+        if size_threshold is None:
+            size_threshold = envs.VLLM_MSGPACK_ZERO_COPY_THRESHOLD
         self.encoder = msgpack.Encoder(enc_hook=self.enc_hook)
         # This is used as a local stash of buffers that we can then access from
         # our custom `msgspec` hook, `enc_hook`. We don't have a way to
         # pass custom data to the hook otherwise.
         self.aux_buffers: Optional[list[bytestr]] = None
+        self.size_threshold = size_threshold
 
     def encode(self, obj: Any) -> Sequence[bytestr]:
         try:
@@ -65,6 +86,25 @@ def enc_hook(self, obj: Any) -> Any:
         if isinstance(obj, np.ndarray) and obj.dtype.kind not in ('O', 'V'):
             return self._encode_ndarray(obj)
 
+        if isinstance(obj, MultiModalKwargs):
+            mm: MultiModalKwargs = obj
+            if not mm.modalities:
+                # just return the main dict if there are no modalities.
+                return dict(mm)
+
+            # ignore the main dict, it will be re-indexed.
+            # Encode a list of MultiModalKwargsItems as plain dicts
+            # + special handling for .field.
+            # Any tensors *not* indexed by modality will be ignored.
+            return [[{
+                "modality": elem.modality,
+                "key": elem.key,
+                "data": self._encode_nested_tensors(elem.data),
+                "field": self._encode_mm_field(elem.field),
+            } for elem in item.values()]
+                    for itemlist in mm._items_by_modality.values()
+                    for item in itemlist]
+
         if isinstance(obj, FunctionType):
             # `pickle` is generally faster than cloudpickle, but can have
             # problems serializing methods.
@@ -77,8 +117,9 @@ def _encode_ndarray(
         self, obj: np.ndarray
     ) -> tuple[str, tuple[int, ...], Union[int, memoryview]]:
         assert self.aux_buffers is not None
+        # If the array is non-contiguous, we need to copy it first
         arr_data = obj.data if obj.data.c_contiguous else obj.tobytes()
-        if not obj.shape or obj.nbytes < MIN_NOCOPY_BUF_SIZE:
+        if not obj.shape or obj.nbytes < self.size_threshold:
             # Encode small arrays and scalars inline. Using this extension type
             # ensures we can avoid copying when decoding.
             data = msgpack.Ext(CUSTOM_TYPE_RAW_VIEW, arr_data)
@@ -92,6 +133,26 @@ def _encode_ndarray(
         # backing buffers that we've stashed in `aux_buffers`.
         return obj.dtype.str, obj.shape, data
 
+    def _encode_nested_tensors(self, nt: NestedTensors) -> Any:
+        if isinstance(nt, torch.Tensor):
+            return self._encode_ndarray(nt.numpy())
+        if isinstance(nt, (int, float)):
+            # Although it violates NestedTensors type, MultiModalKwargs
+            # values are sometimes floats.
+            return nt
+        return [self._encode_nested_tensors(x) for x in nt]
+
+    def _encode_mm_field(self, field: BaseMultiModalField):
+        # Figure out the factory name for the field type.
+        name = MMF_CLASS_TO_FACTORY.get(field.__class__)
+        if not name:
+            raise TypeError(f"Unsupported field type: {field.__class__}")
+        # We just need to copy all of the field values in order
+        # which will be then used to reconstruct the field.
+        field_values = (getattr(field, f.name)
+                        for f in dataclasses.fields(field))
+        return name, *field_values
+
 
 class MsgpackDecoder:
     """Decoder with custom torch tensor and numpy array serialization.
@@ -126,13 +187,50 @@ def dec_hook(self, t: type, obj: Any) -> Any:
                 return self._decode_ndarray(obj)
             if issubclass(t, torch.Tensor):
                 return torch.from_numpy(self._decode_ndarray(obj))
+            if issubclass(t, MultiModalKwargs):
+                if isinstance(obj, list):
+                    return MultiModalKwargs.from_items(
+                        self._decode_mm_items(obj))
+                return MultiModalKwargs({
+                    k: self._decode_nested_tensors(v)
+                    for k, v in obj.items()
+                })
         return obj
 
     def _decode_ndarray(self, arr: Any) -> np.ndarray:
         dtype, shape, data = arr
-        buffer = self.aux_buffers[data] if isinstance(data, int) else data
+        # Copy from inline representation, otherwise Torch is unhappy since
+        # the returned memory is non-writeable.
+        buffer = self.aux_buffers[data] if isinstance(data, int) \
+            else bytearray(data)
         return np.ndarray(buffer=buffer, dtype=np.dtype(dtype), shape=shape)
 
+    def _decode_mm_items(self, obj: list) -> list[MultiModalKwargsItem]:
+        decoded_items = []
+        for item in obj:
+            elems = []
+            for v in item:
+                v["data"] = self._decode_nested_tensors(v["data"])
+                # Reconstruct the field processor using MultiModalFieldConfig
+                factory_meth_name, *field_args = v["field"]
+                factory_meth = getattr(MultiModalFieldConfig,
+                                       factory_meth_name)
+                v["field"] = factory_meth(None, *field_args).field
+                elems.append(MultiModalFieldElem(**v))
+            decoded_items.append(MultiModalKwargsItem.from_elems(elems))
+        return decoded_items
+
+    def _decode_nested_tensors(self, obj: Any) -> NestedTensors:
+        if isinstance(obj, (int, float)):
+            # Although it violates NestedTensors type, MultiModalKwargs
+            # values are sometimes floats.
+            return obj
+        if not isinstance(obj, list):
+            raise TypeError(f"Unexpected NestedTensors contents: {type(obj)}")
+        if obj and isinstance(obj[0], str):
+            return torch.from_numpy(self._decode_ndarray(obj))
+        return [self._decode_nested_tensors(x) for x in obj]
+
     def ext_hook(self, code: int, data: memoryview) -> Any:
         if code == CUSTOM_TYPE_RAW_VIEW:
             return data


Reviews:
user: login: njhill
body: 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 
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-10T23:36:28Z
user: login: DarkLight1337
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-11T03:52:50Z
user: login: p88h
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-11T11:29:24Z
user: login: p88h
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-11T11:35:35Z
user: login: p88h
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-11T11:37:01Z
user: login: njhill
body: Thanks for all of the investigation/benchmarking and work on this @p88h!
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-14T22:20:57Z
user: login: p88h
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-15T06:57:01Z
user: login: p88h
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-15T07:32:31Z
user: login: njhill
body: Thanks @p88h, I did another pass ... I haven't looked at your isolated test case yet, will do so when I get a chance.
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-15T18:00:41Z
user: login: p88h
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-15T18:58:58Z
user: login: p88h
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-15T18:59:40Z
user: login: p88h
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-15T19:01:54Z
user: login: njhill
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-15T19:13:30Z
user: login: p88h
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-15T19:17:09Z
user: login: p88h
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-15T19:18:50Z
user: login: p88h
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-15T19:33:44Z
user: login: p88h
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-15T19:48:45Z
user: login: njhill
body: Thanks @p88h I think this is pretty close now that mystery is solved.
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-15T20:58:29Z
user: login: p88h
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-15T21:57:19Z
user: login: njhill
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-16T03:43:42Z
user: login: p88h
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-16T05:19:23Z
user: login: p88h
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-16T05:27:32Z
user: login: njhill
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-16T14:54:49Z
user: login: p88h
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-16T15:00:04Z
user: login: njhill
body: Thanks for the great work @p88h!
state: APPROVED
author_association: MEMBER
submitted_at: 2025-04-16T15:39:40Z

Review comments:
diff_hunk: @@ -64,6 +69,26 @@ def enc_hook(self, obj: Any) -> Any:
         if isinstance(obj, np.ndarray) and obj.dtype.kind not in ('O', 'V'):
             return self._encode_ndarray(obj)
 
+        if isinstance(obj, MultiModalKwargs):
+            mm: MultiModalKwargs = obj
+            if mm.modalities:
+                # ignore the main dict, it will be re-indexed.
+                # pass a list of MultiModalKwargsItem, then see below
+                # Any tensors *not* indexed by modality will be ignored.
+                return [mm.get_items(m) for m in mm.modalities]
path: vllm/v1/serial_utils.py
user: login: njhill
body: ```suggestion
                return mm._items_by_modality.values()
```
created_at: 2025-04-10T22:58:22Z
author_association: MEMBER
diff_hunk: @@ -64,6 +69,26 @@ def enc_hook(self, obj: Any) -> Any:
         if isinstance(obj, np.ndarray) and obj.dtype.kind not in ('O', 'V'):
             return self._encode_ndarray(obj)
 
+        if isinstance(obj, MultiModalKwargs):
+            mm: MultiModalKwargs = obj
+            if mm.modalities:
+                # ignore the main dict, it will be re-indexed.
+                # pass a list of MultiModalKwargsItem, then see below
+                # Any tensors *not* indexed by modality will be ignored.
+                return [mm.get_items(m) for m in mm.modalities]
+            # just return the main dict if there are no modalities
+            return {k: v for k, v in obj.items()}
path: vllm/v1/serial_utils.py
user: login: njhill
body: ```suggestion
            return dict(mm)
```
created_at: 2025-04-10T22:59:16Z
author_association: MEMBER
diff_hunk: @@ -130,6 +161,24 @@ def _decode_ndarray(self, arr: Any) -> np.ndarray:
         buffer = self.aux_buffers[data] if isinstance(data, int) else data
         return np.ndarray(buffer=buffer, dtype=np.dtype(dtype), shape=shape)
 
+    def _decode_items(self, obj: list) -> list[MultiModalKwargsItem]:
+        all = []
+        for item in chain.from_iterable(obj):
+            elems = []
+            for v in item.values():
+                v['data'] = self._decode_nested(v['data'])
+                v['field'] = pickle.loads(v['field'])
+                elems.append(MultiModalFieldElem(**v))
+            all.append(MultiModalKwargsItem.from_elems(elems))
+        return all
+
+    def _decode_nested(self, obj: Any) -> NestedTensors:
+        if isinstance(obj, list) and isinstance(obj[0], str):
+            return torch.from_numpy(self._decode_ndarray(obj))
+        if isinstance(obj, list):
+            return [self._decode_nested(x) for x in obj]
+        raise TypeError(f"Unexpected NestedArray contents: {obj}")
path: vllm/v1/serial_utils.py
user: login: njhill
body: ```suggestion
    def _decode_nested_tensors(self, obj: Any) -> NestedTensors:
        if not isinstance(obj, list):
            raise TypeError(f"Unexpected NestedTensors contents: {type(obj)}")
        if obj and isinstance(obj[0], str):
            return torch.from_numpy(self._decode_ndarray(obj))
        return [self._decode_nested_tensors(x) for x in obj]
```

`type(obj)` because it's the type that's unexpected, and the obj itself could be very large
created_at: 2025-04-10T23:10:14Z
author_association: MEMBER
diff_hunk: @@ -64,6 +69,26 @@ def enc_hook(self, obj: Any) -> Any:
         if isinstance(obj, np.ndarray) and obj.dtype.kind not in ('O', 'V'):
             return self._encode_ndarray(obj)
 
+        if isinstance(obj, MultiModalKwargs):
+            mm: MultiModalKwargs = obj
+            if mm.modalities:
+                # ignore the main dict, it will be re-indexed.
+                # pass a list of MultiModalKwargsItem, then see below
+                # Any tensors *not* indexed by modality will be ignored.
+                return [mm.get_items(m) for m in mm.modalities]
path: vllm/v1/serial_utils.py
user: login: njhill
body: I'm not sure that it makes sense to "re-enter" here. May as well encode the `MultiModalKwargsItem`s before returning.
created_at: 2025-04-10T23:16:48Z
author_association: MEMBER
diff_hunk: @@ -64,6 +69,26 @@ def enc_hook(self, obj: Any) -> Any:
         if isinstance(obj, np.ndarray) and obj.dtype.kind not in ('O', 'V'):
             return self._encode_ndarray(obj)
 
+        if isinstance(obj, MultiModalKwargs):
+            mm: MultiModalKwargs = obj
+            if mm.modalities:
+                # ignore the main dict, it will be re-indexed.
+                # pass a list of MultiModalKwargsItem, then see below
+                # Any tensors *not* indexed by modality will be ignored.
+                return [mm.get_items(m) for m in mm.modalities]
+            # just return the main dict if there are no modalities
+            return {k: v for k, v in obj.items()}
+
+        if isinstance(obj, MultiModalKwargsItem):
+            # Encode as plain dictionary + special handling for '.field'
+            rd = {}
+            for k, v in obj.items():
path: vllm/v1/serial_utils.py
user: login: njhill
body: I don't think we need to encode the keys because they are reconstructed by `MulitModalKwargsItems.from_elems()`.
created_at: 2025-04-10T23:18:05Z
author_association: MEMBER
diff_hunk: @@ -64,6 +69,26 @@ def enc_hook(self, obj: Any) -> Any:
         if isinstance(obj, np.ndarray) and obj.dtype.kind not in ('O', 'V'):
             return self._encode_ndarray(obj)
 
+        if isinstance(obj, MultiModalKwargs):
+            mm: MultiModalKwargs = obj
+            if mm.modalities:
+                # ignore the main dict, it will be re-indexed.
+                # pass a list of MultiModalKwargsItem, then see below
+                # Any tensors *not* indexed by modality will be ignored.
+                return [mm.get_items(m) for m in mm.modalities]
+            # just return the main dict if there are no modalities
+            return {k: v for k, v in obj.items()}
+
+        if isinstance(obj, MultiModalKwargsItem):
+            # Encode as plain dictionary + special handling for '.field'
+            rd = {}
+            for k, v in obj.items():
+                vv = asdict(v)
+                vv['field'] = pickle.dumps(v.field,
+                                           protocol=pickle.HIGHEST_PROTOCOL)
path: vllm/v1/serial_utils.py
user: login: njhill
body: We want to avoid using pickle here ... I think it shouldn't be too hard to have custom handling?
created_at: 2025-04-10T23:22:43Z
author_association: MEMBER
diff_hunk: @@ -121,6 +146,12 @@ def dec_hook(self, t: type, obj: Any) -> Any:
         if isclass(t):
             if issubclass(t, np.ndarray):
                 return self._decode_ndarray(obj)
+            if issubclass(t, MultiModalKwargs) and isinstance(obj, dict):
+                return MultiModalKwargs(
+                    {k: self._decode_nested(obj[k])
+                     for k in obj})
+            if issubclass(t, MultiModalKwargs) and isinstance(obj, list):
+                return MultiModalKwargs.from_items(self._decode_items(obj))
path: vllm/v1/serial_utils.py
user: login: njhill
body: ```suggestion
            if issubclass(t, MultiModalKwargs) 
                if isinstance(obj, list):
                    return MultiModalKwargs.from_items(self._decode_items(obj))
                return MultiModalKwargs(
                    {k: self._decode_nested(v)
                     for k in obj.items()})
```

And move below the Tensor case since this is more specific?
created_at: 2025-04-10T23:26:43Z
author_association: MEMBER
diff_hunk: @@ -130,6 +161,24 @@ def _decode_ndarray(self, arr: Any) -> np.ndarray:
         buffer = self.aux_buffers[data] if isinstance(data, int) else data
         return np.ndarray(buffer=buffer, dtype=np.dtype(dtype), shape=shape)
 
+    def _decode_items(self, obj: list) -> list[MultiModalKwargsItem]:
path: vllm/v1/serial_utils.py
user: login: njhill
body: Maybe rename to something less generic like `_decode_mm_items`?
created_at: 2025-04-10T23:27:37Z
author_association: MEMBER
diff_hunk: @@ -70,6 +75,99 @@ def test_encode_decode():
     assert_equal(decoded2, obj)
 
 
+class MyRequest(msgspec.Struct):
+    mm: Optional[list[MultiModalKwargs]]
+
+
+def test_multimodal_kwargs():
+    d = {
+        "foo":
+        torch.zeros(1000, dtype=torch.float16),
+        "bar": [torch.zeros(i * 1000, dtype=torch.int8) for i in range(3)],
+        "baz": [
+            torch.rand((256), dtype=torch.float16),
+            [
+                torch.rand((1, 12), dtype=torch.float32),
+                torch.rand((3, 5, 7), dtype=torch.float64),
+            ], [torch.rand((4, 4), dtype=torch.float16)]
+        ],
+    }
+
+    # pack mm kwargs into a mock request so that it can be decoded properly
+    req = MyRequest(mm=[MultiModalKwargs(d)])
+
+    encoder = MsgpackEncoder()
+    decoder = MsgpackDecoder(MyRequest)
+
+    encoded = encoder.encode(req)
+
+    # 8 total tensors + top level buffer
+    assert len(encoded) == 6
+
+    total_len = sum([len(x) for x in encoded])
+
+    # expected total encoding length, should be 4440, +-20 for minor changes
+    assert total_len >= 4420 and total_len <= 4460
+    decoded: MultiModalKwargs = decoder.decode(encoded).mm[0]
+    assert all(nested_equal(d[k], decoded[k]) for k in d)
+
+
+def test_multimodal_items_by_modality():
+    e1 = MultiModalFieldElem("audio", "a0", torch.zeros(1000,
+                                                        dtype=torch.int16),
+                             MultiModalBatchedField())
+    e2 = MultiModalFieldElem(
+        "video",
+        "v0",
+        [torch.zeros(1000, dtype=torch.int8) for _ in range(4)],
+        MultiModalBatchedField(),
+    )
+    e3 = MultiModalFieldElem("image", "i0", torch.zeros(1000,
+                                                        dtype=torch.int32),
+                             MultiModalBatchedField())
+    e4 = MultiModalFieldElem("image", "i1", torch.zeros(1000,
+                                                        dtype=torch.int32),
+                             MultiModalBatchedField())
+    audio = MultiModalKwargsItem.from_elems([e1])
+    video = MultiModalKwargsItem.from_elems([e2])
+    image = MultiModalKwargsItem.from_elems([e3, e4])
+    mm = MultiModalKwargs.from_items([audio, video, image])
+
+    # pack mm kwargs into a mock request so that it can be decoded properly
+    req = MyRequest([mm])
+
+    encoder = MsgpackEncoder()
+    decoder = MsgpackDecoder(MyRequest)
+
+    encoded = encoder.encode(req)
+
+    # 5 total tensors + top level buffer
+    assert len(encoded) == 8
+
+    total_len = sum([len(x) for x in encoded])
+
+    # expected total encoding length, should be 7507, +-20 for minor changes
+    assert total_len >= 7487 and total_len <= 7527
+    decoded: MultiModalKwargs = decoder.decode(encoded).mm[0]
+
+    # check all modalities were recovered and do some basic sanity checks
+    assert len(decoded.modalities) == 3
+    images = decoded.get_items("image")
+    assert len(images) == 1
+    assert len(images[0].items()) == 2
+    assert list(images[0].keys()) == ["i0", "i1"]
+
+    # check the tensor contents and layout in the main dict
+    assert all(nested_equal(mm[k], decoded[k]) for k in mm)
+
+
+def nested_equal(a: NestedTensors, b: NestedTensors):
+    if isinstance(a, torch.Tensor):
+        return torch.equal(a, b)
+    else:
+        return all([nested_equal(x, y) for (x, y) in zip(a, b)])
path: tests/v1/test_serial_utils.py
user: login: njhill
body: nit:
```suggestion
        return all(nested_equal(x, y) for x, y in zip(a, b))
```
created_at: 2025-04-10T23:30:10Z
author_association: MEMBER
diff_hunk: @@ -70,6 +75,99 @@ def test_encode_decode():
     assert_equal(decoded2, obj)
 
 
+class MyRequest(msgspec.Struct):
+    mm: Optional[list[MultiModalKwargs]]
+
+
+def test_multimodal_kwargs():
+    d = {
+        "foo":
+        torch.zeros(1000, dtype=torch.float16),
+        "bar": [torch.zeros(i * 1000, dtype=torch.int8) for i in range(3)],
+        "baz": [
+            torch.rand((256), dtype=torch.float16),
+            [
+                torch.rand((1, 12), dtype=torch.float32),
+                torch.rand((3, 5, 7), dtype=torch.float64),
+            ], [torch.rand((4, 4), dtype=torch.float16)]
+        ],
+    }
+
+    # pack mm kwargs into a mock request so that it can be decoded properly
+    req = MyRequest(mm=[MultiModalKwargs(d)])
+
+    encoder = MsgpackEncoder()
+    decoder = MsgpackDecoder(MyRequest)
+
+    encoded = encoder.encode(req)
+
+    # 8 total tensors + top level buffer
+    assert len(encoded) == 6
+
+    total_len = sum([len(x) for x in encoded])
path: tests/v1/test_serial_utils.py
user: login: DarkLight1337
body: ```suggestion
    total_len = sum(len(x) for x in encoded)
```

Nit. Same for the next test 
created_at: 2025-04-11T03:48:30Z
author_association: MEMBER
diff_hunk: @@ -64,6 +69,26 @@ def enc_hook(self, obj: Any) -> Any:
         if isinstance(obj, np.ndarray) and obj.dtype.kind not in ('O', 'V'):
             return self._encode_ndarray(obj)
 
+        if isinstance(obj, MultiModalKwargs):
+            mm: MultiModalKwargs = obj
+            if mm.modalities:
+                # ignore the main dict, it will be re-indexed.
+                # pass a list of MultiModalKwargsItem, then see below
+                # Any tensors *not* indexed by modality will be ignored.
+                return [mm.get_items(m) for m in mm.modalities]
+            # just return the main dict if there are no modalities
+            return {k: v for k, v in obj.items()}
+
+        if isinstance(obj, MultiModalKwargsItem):
+            # Encode as plain dictionary + special handling for '.field'
+            rd = {}
+            for k, v in obj.items():
path: vllm/v1/serial_utils.py
user: login: p88h
body: right, that should be safe to drop.
created_at: 2025-04-11T11:29:24Z
author_association: CONTRIBUTOR
diff_hunk: @@ -64,6 +69,26 @@ def enc_hook(self, obj: Any) -> Any:
         if isinstance(obj, np.ndarray) and obj.dtype.kind not in ('O', 'V'):
             return self._encode_ndarray(obj)
 
+        if isinstance(obj, MultiModalKwargs):
+            mm: MultiModalKwargs = obj
+            if mm.modalities:
+                # ignore the main dict, it will be re-indexed.
+                # pass a list of MultiModalKwargsItem, then see below
+                # Any tensors *not* indexed by modality will be ignored.
+                return [mm.get_items(m) for m in mm.modalities]
+            # just return the main dict if there are no modalities
+            return {k: v for k, v in obj.items()}
+
+        if isinstance(obj, MultiModalKwargsItem):
+            # Encode as plain dictionary + special handling for '.field'
+            rd = {}
+            for k, v in obj.items():
+                vv = asdict(v)
+                vv['field'] = pickle.dumps(v.field,
+                                           protocol=pickle.HIGHEST_PROTOCOL)
path: vllm/v1/serial_utils.py
user: login: p88h
body: Sure, I guess... there is a few of `BaseMultiModalField` variants, so we'd likely need to represent them here and re-create correctly. Catching a case when a new implementation is added will be tricky though. 

created_at: 2025-04-11T11:35:35Z
author_association: CONTRIBUTOR
diff_hunk: @@ -130,6 +161,24 @@ def _decode_ndarray(self, arr: Any) -> np.ndarray:
         buffer = self.aux_buffers[data] if isinstance(data, int) else data
         return np.ndarray(buffer=buffer, dtype=np.dtype(dtype), shape=shape)
 
+    def _decode_items(self, obj: list) -> list[MultiModalKwargsItem]:
path: vllm/v1/serial_utils.py
user: login: p88h
body: ok
created_at: 2025-04-11T11:37:01Z
author_association: CONTRIBUTOR
diff_hunk: @@ -27,29 +30,32 @@ class MsgpackEncoder:
 
     Note that unlike vanilla `msgspec` Encoders, this interface is generally
     not thread-safe when encoding tensors / numpy arrays.
+
+    By default, arrays below 256B are serialized inline.
+    Larger will get sent via dedicated messages. 
+    Note that this is a per-tensor limit.
+
+    Sending multiple large messages via zeromq saturates memory very quickly.
+    See: https://github.com/vllm-project/vllm/issues/16185
     """
 
-    def __init__(self):
+    def __init__(self, size_threshold=256):
         self.encoder = msgpack.Encoder(enc_hook=self.enc_hook)
         # This is used as a local stash of buffers that we can then access from
         # our custom `msgspec` hook, `enc_hook`. We don't have a way to
         # pass custom data to the hook otherwise.
         self.aux_buffers: Optional[list[bytestr]] = None
+        self.size_threshold = size_threshold
 
     def encode(self, obj: Any) -> Sequence[bytestr]:
+        return self.encode_into(obj, bytearray())
+
+    def encode_into(self, obj: Any, buf: bytearray) -> Sequence[bytestr]:
path: vllm/v1/serial_utils.py
user: login: njhill
body: what's the reason for this change? I can see it's a bit less code but might be better to let msgspec handle the allocation in this case
created_at: 2025-04-14T21:29:16Z
author_association: MEMBER
diff_hunk: @@ -27,29 +30,32 @@ class MsgpackEncoder:
 
     Note that unlike vanilla `msgspec` Encoders, this interface is generally
     not thread-safe when encoding tensors / numpy arrays.
+
+    By default, arrays below 256B are serialized inline.
+    Larger will get sent via dedicated messages. 
+    Note that this is a per-tensor limit.
+
+    Sending multiple large messages via zeromq saturates memory very quickly.
+    See: https://github.com/vllm-project/vllm/issues/16185
     """
 
-    def __init__(self):
+    def __init__(self, size_threshold=256):
path: vllm/v1/serial_utils.py
user: login: njhill
body: If we want this to be tuneable I think it would be better to do via an env var, see https://github.com/vllm-project/vllm/blob/main/vllm/envs.py.
created_at: 2025-04-14T21:31:07Z
author_association: MEMBER
diff_hunk: @@ -65,6 +71,25 @@ def enc_hook(self, obj: Any) -> Any:
         if isinstance(obj, np.ndarray) and obj.dtype.kind not in ('O', 'V'):
             return self._encode_ndarray(obj)
 
+        if isinstance(obj, MultiModalKwargs):
+            mm: MultiModalKwargs = obj
+            if mm.modalities:
+                # ignore the main dict, it will be re-indexed.
+                # pass a list of MultiModalKwargsItem, then see below
+                # Any tensors *not* indexed by modality will be ignored.
+                return list(mm._items_by_modality.values())
+            # just return the main dict if there are no modalities
+            return dict(mm)
+
+        if isinstance(obj, MultiModalKwargsItem):
+            ret = []
+            for elem in obj.values():
+                # Encode as plain dictionary + special handling for .field
+                d = asdict(elem)
+                d["field"] = elem.field.field_type()
path: vllm/v1/serial_utils.py
user: login: njhill
body: Rather than requiring changes to those types, could we instead just encode as e.g. `(field.__class__.__name__, asdict(field))`? Then when decoding we would check the classname matches one of the three valid names.
created_at: 2025-04-14T21:44:47Z
author_association: MEMBER
diff_hunk: @@ -77,8 +102,13 @@ def _encode_ndarray(
         self, obj: np.ndarray
     ) -> tuple[str, tuple[int, ...], Union[int, memoryview]]:
         assert self.aux_buffers is not None
-        arr_data = obj.data if obj.data.c_contiguous else obj.tobytes()
-        if not obj.shape or obj.nbytes < MIN_NOCOPY_BUF_SIZE:
+        # Either copy the memoryview directly or flatten the array to bytes.
+        # Sending memoryviews is theoretically faster, but in this particular
+        # case, it triggers some unnecessary copies anyway.
+        # With this, the tensors can still be zero-copy read.
+        arr_data = obj.data.tobytes() if obj.data.c_contiguous \
+            else obj.tobytes()
path: vllm/v1/serial_utils.py
user: login: njhill
body: Do you know if there's a difference between `arr.data.tobytes()` and `arr.tobytes()`?

If we are always copying here, what's the value of keeping it as a separate buffer? It might as well be written into the main buffer (i.e. the inline path). Or is it that this has the effect of keeping it split into smaller messages and this somehow avoids the memory blowup problem?

It would still be really good to understand what the issue is with zmq pulling directly from a memoryview ...
created_at: 2025-04-14T21:53:53Z
author_association: MEMBER
diff_hunk: @@ -126,13 +156,45 @@ def dec_hook(self, t: type, obj: Any) -> Any:
                 return self._decode_ndarray(obj)
             if issubclass(t, torch.Tensor):
                 return torch.from_numpy(self._decode_ndarray(obj))
+            if issubclass(t, MultiModalKwargs):
+                if isinstance(obj, list):
+                    return MultiModalKwargs.from_items(
+                        self._decode_mm_items(obj))
+                return MultiModalKwargs({
+                    k: self._decode_nested_tensors(v)
+                    for k, v in obj.items()
+                })
         return obj
 
     def _decode_ndarray(self, arr: Any) -> np.ndarray:
         dtype, shape, data = arr
-        buffer = self.aux_buffers[data] if isinstance(data, int) else data
+        # Copy from inline representation, otherwise Torch is unhappy since
+        # the returned memory is non-writeable.
+        buffer = self.aux_buffers[data] if isinstance(data, int) \
+            else bytearray(data)
path: vllm/v1/serial_utils.py
user: login: njhill
body: I really don't think that this is needed. Even if there's a warning I don't think it poses any problem. I can't think of any reason this data would be modified in-place, and even if it is, I don't think it will actually cause any problems. We could look into other ways to avoid the warning if needed.
created_at: 2025-04-14T22:04:50Z
author_association: MEMBER
diff_hunk: @@ -27,29 +30,32 @@ class MsgpackEncoder:
 
     Note that unlike vanilla `msgspec` Encoders, this interface is generally
     not thread-safe when encoding tensors / numpy arrays.
+
+    By default, arrays below 256B are serialized inline.
+    Larger will get sent via dedicated messages. 
+    Note that this is a per-tensor limit.
+
+    Sending multiple large messages via zeromq saturates memory very quickly.
+    See: https://github.com/vllm-project/vllm/issues/16185
     """
 
-    def __init__(self):
+    def __init__(self, size_threshold=256):
         self.encoder = msgpack.Encoder(enc_hook=self.enc_hook)
         # This is used as a local stash of buffers that we can then access from
         # our custom `msgspec` hook, `enc_hook`. We don't have a way to
         # pass custom data to the hook otherwise.
         self.aux_buffers: Optional[list[bytestr]] = None
+        self.size_threshold = size_threshold
 
     def encode(self, obj: Any) -> Sequence[bytestr]:
+        return self.encode_into(obj, bytearray())
+
+    def encode_into(self, obj: Any, buf: bytearray) -> Sequence[bytestr]:
path: vllm/v1/serial_utils.py
user: login: p88h
body: Ah, what I was going for but didn't finish yet was making management of this work buffer automatic (and get rid of the encode_into altogether perhaps) so that the bytearray is kept  at the MsgSpecEncoder level. 
Encode() on it's own does basically the same as what this is doing right now - allocates a new buffer every time.
created_at: 2025-04-15T05:56:41Z
author_association: CONTRIBUTOR
diff_hunk: @@ -65,6 +71,25 @@ def enc_hook(self, obj: Any) -> Any:
         if isinstance(obj, np.ndarray) and obj.dtype.kind not in ('O', 'V'):
             return self._encode_ndarray(obj)
 
+        if isinstance(obj, MultiModalKwargs):
+            mm: MultiModalKwargs = obj
+            if mm.modalities:
+                # ignore the main dict, it will be re-indexed.
+                # pass a list of MultiModalKwargsItem, then see below
+                # Any tensors *not* indexed by modality will be ignored.
+                return list(mm._items_by_modality.values())
+            # just return the main dict if there are no modalities
+            return dict(mm)
+
+        if isinstance(obj, MultiModalKwargsItem):
+            ret = []
+            for elem in obj.values():
+                # Encode as plain dictionary + special handling for .field
+                d = asdict(elem)
+                d["field"] = elem.field.field_type()
path: vllm/v1/serial_utils.py
user: login: p88h
body: Ok
created_at: 2025-04-15T05:57:13Z
author_association: CONTRIBUTOR
diff_hunk: @@ -77,8 +102,13 @@ def _encode_ndarray(
         self, obj: np.ndarray
     ) -> tuple[str, tuple[int, ...], Union[int, memoryview]]:
         assert self.aux_buffers is not None
-        arr_data = obj.data if obj.data.c_contiguous else obj.tobytes()
-        if not obj.shape or obj.nbytes < MIN_NOCOPY_BUF_SIZE:
+        # Either copy the memoryview directly or flatten the array to bytes.
+        # Sending memoryviews is theoretically faster, but in this particular
+        # case, it triggers some unnecessary copies anyway.
+        # With this, the tensors can still be zero-copy read.
+        arr_data = obj.data.tobytes() if obj.data.c_contiguous \
+            else obj.tobytes()
path: vllm/v1/serial_utils.py
user: login: p88h
body: Writing to main buffer was a bit more expensive in testing (even without copying out on the decode path). 
The threshold = 256MB benchmark tests this. Memory usage is similar (there is no big difference between keeping a list of byte buffers and one larger buffer), but performance wise there was a noticeable difference. 

I kept two different paths for tobytes() for benchmarking convenience mostly, I don't think it's necessary, but I'll test it as well. 
created_at: 2025-04-15T06:15:41Z
author_association: CONTRIBUTOR
diff_hunk: @@ -126,13 +156,45 @@ def dec_hook(self, t: type, obj: Any) -> Any:
                 return self._decode_ndarray(obj)
             if issubclass(t, torch.Tensor):
                 return torch.from_numpy(self._decode_ndarray(obj))
+            if issubclass(t, MultiModalKwargs):
+                if isinstance(obj, list):
+                    return MultiModalKwargs.from_items(
+                        self._decode_mm_items(obj))
+                return MultiModalKwargs({
+                    k: self._decode_nested_tensors(v)
+                    for k, v in obj.items()
+                })
         return obj
 
     def _decode_ndarray(self, arr: Any) -> np.ndarray:
         dtype, shape, data = arr
-        buffer = self.aux_buffers[data] if isinstance(data, int) else data
+        # Copy from inline representation, otherwise Torch is unhappy since
+        # the returned memory is non-writeable.
+        buffer = self.aux_buffers[data] if isinstance(data, int) \
+            else bytearray(data)
path: vllm/v1/serial_utils.py
user: login: p88h
body: I found the original bug that inspired this warning - https://github.com/pytorch/pytorch/issues/5442
You are probably right that this should not affect VLLM use cases in terms of causing any errors. 

However, it will affect the ability to free up the whole msgspec buffer received from zmq - the whole array will have to stay alive as long as there are tensors pointing into it. Given this copy is only performed for very small tensors now, the impact of copying is probably negligible, but the impact of copying these is unknown - the Tensors will end up in the cache on the receiving side. 

created_at: 2025-04-15T06:43:10Z
author_association: CONTRIBUTOR
diff_hunk: @@ -27,29 +30,32 @@ class MsgpackEncoder:
 
     Note that unlike vanilla `msgspec` Encoders, this interface is generally
     not thread-safe when encoding tensors / numpy arrays.
+
+    By default, arrays below 256B are serialized inline.
+    Larger will get sent via dedicated messages. 
+    Note that this is a per-tensor limit.
+
+    Sending multiple large messages via zeromq saturates memory very quickly.
+    See: https://github.com/vllm-project/vllm/issues/16185
     """
 
-    def __init__(self):
+    def __init__(self, size_threshold=256):
         self.encoder = msgpack.Encoder(enc_hook=self.enc_hook)
         # This is used as a local stash of buffers that we can then access from
         # our custom `msgspec` hook, `enc_hook`. We don't have a way to
         # pass custom data to the hook otherwise.
         self.aux_buffers: Optional[list[bytestr]] = None
+        self.size_threshold = size_threshold
 
     def encode(self, obj: Any) -> Sequence[bytestr]:
+        return self.encode_into(obj, bytearray())
+
+    def encode_into(self, obj: Any, buf: bytearray) -> Sequence[bytestr]:
path: vllm/v1/serial_utils.py
user: login: p88h
body: I added an internal msg_buffer that follows this path and a TODO to finish up by changing the usage pattern in core.py
created_at: 2025-04-15T06:46:38Z
author_association: CONTRIBUTOR
diff_hunk: @@ -27,29 +30,32 @@ class MsgpackEncoder:
 
     Note that unlike vanilla `msgspec` Encoders, this interface is generally
     not thread-safe when encoding tensors / numpy arrays.
+
+    By default, arrays below 256B are serialized inline.
+    Larger will get sent via dedicated messages. 
+    Note that this is a per-tensor limit.
+
+    Sending multiple large messages via zeromq saturates memory very quickly.
+    See: https://github.com/vllm-project/vllm/issues/16185
     """
 
-    def __init__(self):
+    def __init__(self, size_threshold=256):
path: vllm/v1/serial_utils.py
user: login: p88h
body: I added this mostly for the unit test convenience, but sure, can also add an env var.
created_at: 2025-04-15T06:57:01Z
author_association: CONTRIBUTOR
diff_hunk: @@ -77,8 +102,13 @@ def _encode_ndarray(
         self, obj: np.ndarray
     ) -> tuple[str, tuple[int, ...], Union[int, memoryview]]:
         assert self.aux_buffers is not None
-        arr_data = obj.data if obj.data.c_contiguous else obj.tobytes()
-        if not obj.shape or obj.nbytes < MIN_NOCOPY_BUF_SIZE:
+        # Either copy the memoryview directly or flatten the array to bytes.
+        # Sending memoryviews is theoretically faster, but in this particular
+        # case, it triggers some unnecessary copies anyway.
+        # With this, the tensors can still be zero-copy read.
+        arr_data = obj.data.tobytes() if obj.data.c_contiguous \
+            else obj.tobytes()
path: vllm/v1/serial_utils.py
user: login: p88h
body: ok, the tests are identical so I'll just simplify this.
created_at: 2025-04-15T07:32:31Z
author_association: CONTRIBUTOR
diff_hunk: @@ -65,6 +77,24 @@ def enc_hook(self, obj: Any) -> Any:
         if isinstance(obj, np.ndarray) and obj.dtype.kind not in ('O', 'V'):
             return self._encode_ndarray(obj)
 
+        if isinstance(obj, MultiModalKwargs):
+            mm: MultiModalKwargs = obj
+            if mm.modalities:
+                # ignore the main dict, it will be re-indexed.
+                # pass a list of MultiModalKwargsItem, then see below
+                # Any tensors *not* indexed by modality will be ignored.
+                return list(mm._items_by_modality.values())
+            # just return the main dict if there are no modalities
+            return dict(mm)
+
+        if isinstance(obj, MultiModalKwargsItem):
+            ret = []
+            for elem in obj.values():
+                # Encode as plain dictionary + special handling for .field
+                ret.append(
+                    asdict(elem) | {"field": self._encode_field(elem.field)})
+            return ret
path: vllm/v1/serial_utils.py
user: login: njhill
body: ```suggestion
            return [
                (elem.modality, elem.key, elem.data, self._encode_field(elem.field))
                for elem in obj.values()
            ]
```

I think `asdict` may be the source of the problems :skull:  ... it deep-copies everything, including the tensor `data` here...
created_at: 2025-04-15T17:36:33Z
author_association: MEMBER
diff_hunk: @@ -92,6 +126,13 @@ def _encode_ndarray(
         # backing buffers that we've stashed in `aux_buffers`.
         return obj.dtype.str, obj.shape, data
 
+    def _encode_field(self, field: BaseMultiModalField):
+        # Encode the field as a dictionary + special handling for .field
+        d = asdict(field)
path: vllm/v1/serial_utils.py
user: login: njhill
body: Let's avoid `asdict` per above.

Can use something like this as workaround:
```python
tuple(getattr(obj, f.name) for f in dataclasses.fields(field))
```
created_at: 2025-04-15T17:39:43Z
author_association: MEMBER
diff_hunk: @@ -92,6 +126,13 @@ def _encode_ndarray(
         # backing buffers that we've stashed in `aux_buffers`.
         return obj.dtype.str, obj.shape, data
 
+    def _encode_field(self, field: BaseMultiModalField):
+        # Encode the field as a dictionary + special handling for .field
+        d = asdict(field)
+        # Strip first 10 characters and last 5 characters from the class name
+        # to get the field type name that matches the factory function name.
+        return (field.__class__.__name__[10:-5].lower(), *d.values())
path: vllm/v1/serial_utils.py
user: login: njhill
body: It feels a bit precarious to depend implicitly on the specific wording of the classnames .. I know it's more verbose but what about a lookup like:
```python
MMF_FACTORIES = { MultiModalFlatField: "flat", ... }

# ...

fac = MMF_FACTORIES.get(field.__class__)
if not fac:
    raise RuntimeError(f"unrecognized: {field.__class__}")
```

created_at: 2025-04-15T17:49:21Z
author_association: MEMBER
diff_hunk: @@ -27,29 +30,32 @@ class MsgpackEncoder:
 
     Note that unlike vanilla `msgspec` Encoders, this interface is generally
     not thread-safe when encoding tensors / numpy arrays.
+
+    By default, arrays below 256B are serialized inline.
+    Larger will get sent via dedicated messages. 
+    Note that this is a per-tensor limit.
+
+    Sending multiple large messages via zeromq saturates memory very quickly.
+    See: https://github.com/vllm-project/vllm/issues/16185
     """
 
-    def __init__(self):
+    def __init__(self, size_threshold=256):
         self.encoder = msgpack.Encoder(enc_hook=self.enc_hook)
         # This is used as a local stash of buffers that we can then access from
         # our custom `msgspec` hook, `enc_hook`. We don't have a way to
         # pass custom data to the hook otherwise.
         self.aux_buffers: Optional[list[bytestr]] = None
+        self.size_threshold = size_threshold
 
     def encode(self, obj: Any) -> Sequence[bytestr]:
+        return self.encode_into(obj, bytearray())
+
+    def encode_into(self, obj: Any, buf: bytearray) -> Sequence[bytestr]:
path: vllm/v1/serial_utils.py
user: login: njhill
body: But we want `encode_into`, it was included here intentionally. It's used to avoid reallocating a buffer every time when sending back responses from the engine.
created_at: 2025-04-15T17:53:41Z
author_association: MEMBER
diff_hunk: @@ -77,8 +102,13 @@ def _encode_ndarray(
         self, obj: np.ndarray
     ) -> tuple[str, tuple[int, ...], Union[int, memoryview]]:
         assert self.aux_buffers is not None
-        arr_data = obj.data if obj.data.c_contiguous else obj.tobytes()
-        if not obj.shape or obj.nbytes < MIN_NOCOPY_BUF_SIZE:
+        # Either copy the memoryview directly or flatten the array to bytes.
+        # Sending memoryviews is theoretically faster, but in this particular
+        # case, it triggers some unnecessary copies anyway.
+        # With this, the tensors can still be zero-copy read.
+        arr_data = obj.data.tobytes() if obj.data.c_contiguous \
+            else obj.tobytes()
path: vllm/v1/serial_utils.py
user: login: njhill
body: My guess about how writing to the main buffer would be more expensive is that it means growing it incrementally, resulting in multiple re-sizes that involve a re-allocation and copy of the whole thing.

We don't need to have intermediate `tobytes()` for the non-contiguous in-line case right?
created_at: 2025-04-15T17:59:25Z
author_association: MEMBER
diff_hunk: @@ -65,6 +77,24 @@ def enc_hook(self, obj: Any) -> Any:
         if isinstance(obj, np.ndarray) and obj.dtype.kind not in ('O', 'V'):
             return self._encode_ndarray(obj)
 
+        if isinstance(obj, MultiModalKwargs):
+            mm: MultiModalKwargs = obj
+            if mm.modalities:
+                # ignore the main dict, it will be re-indexed.
+                # pass a list of MultiModalKwargsItem, then see below
+                # Any tensors *not* indexed by modality will be ignored.
+                return list(mm._items_by_modality.values())
+            # just return the main dict if there are no modalities
+            return dict(mm)
+
+        if isinstance(obj, MultiModalKwargsItem):
+            ret = []
+            for elem in obj.values():
+                # Encode as plain dictionary + special handling for .field
+                ret.append(
+                    asdict(elem) | {"field": self._encode_field(elem.field)})
+            return ret
path: vllm/v1/serial_utils.py
user: login: p88h
body: yeah, just noticed this as well - already patched that. Now testing whether this solves the zmq issue as well, but it seems it does !!! 
created_at: 2025-04-15T18:58:58Z
author_association: CONTRIBUTOR
diff_hunk: @@ -92,6 +126,13 @@ def _encode_ndarray(
         # backing buffers that we've stashed in `aux_buffers`.
         return obj.dtype.str, obj.shape, data
 
+    def _encode_field(self, field: BaseMultiModalField):
+        # Encode the field as a dictionary + special handling for .field
+        d = asdict(field)
path: vllm/v1/serial_utils.py
user: login: p88h
body: this asdict is likely much cheaper, but yeah, i'll switch this too.
created_at: 2025-04-15T18:59:40Z
author_association: CONTRIBUTOR
diff_hunk: @@ -92,6 +126,13 @@ def _encode_ndarray(
         # backing buffers that we've stashed in `aux_buffers`.
         return obj.dtype.str, obj.shape, data
 
+    def _encode_field(self, field: BaseMultiModalField):
+        # Encode the field as a dictionary + special handling for .field
+        d = asdict(field)
+        # Strip first 10 characters and last 5 characters from the class name
+        # to get the field type name that matches the factory function name.
+        return (field.__class__.__name__[10:-5].lower(), *d.values())
path: vllm/v1/serial_utils.py
user: login: p88h
body: I was going between this approach and that - overall any change to the class naming will be caught by tests, and both of these are not 'immediately' visible when doing any changes to the MMFieldConfigs. 

That was why I ended up with an explicit method originally - to ensure whoever touches that code, pays attention to serialization. But I guess referencing all classes here at least will pop up in references, so maybe that's better than string manipulation
created_at: 2025-04-15T19:01:54Z
author_association: CONTRIBUTOR
diff_hunk: @@ -65,6 +77,24 @@ def enc_hook(self, obj: Any) -> Any:
         if isinstance(obj, np.ndarray) and obj.dtype.kind not in ('O', 'V'):
             return self._encode_ndarray(obj)
 
+        if isinstance(obj, MultiModalKwargs):
+            mm: MultiModalKwargs = obj
+            if mm.modalities:
+                # ignore the main dict, it will be re-indexed.
+                # pass a list of MultiModalKwargsItem, then see below
+                # Any tensors *not* indexed by modality will be ignored.
+                return list(mm._items_by_modality.values())
+            # just return the main dict if there are no modalities
+            return dict(mm)
+
+        if isinstance(obj, MultiModalKwargsItem):
+            ret = []
+            for elem in obj.values():
+                # Encode as plain dictionary + special handling for .field
+                ret.append(
+                    asdict(elem) | {"field": self._encode_field(elem.field)})
+            return ret
path: vllm/v1/serial_utils.py
user: login: njhill
body: > solves the zmq issue as well,

You mean the apparent issue with using memoryviews directly?
created_at: 2025-04-15T19:13:30Z
author_association: MEMBER
diff_hunk: @@ -77,8 +102,13 @@ def _encode_ndarray(
         self, obj: np.ndarray
     ) -> tuple[str, tuple[int, ...], Union[int, memoryview]]:
         assert self.aux_buffers is not None
-        arr_data = obj.data if obj.data.c_contiguous else obj.tobytes()
-        if not obj.shape or obj.nbytes < MIN_NOCOPY_BUF_SIZE:
+        # Either copy the memoryview directly or flatten the array to bytes.
+        # Sending memoryviews is theoretically faster, but in this particular
+        # case, it triggers some unnecessary copies anyway.
+        # With this, the tensors can still be zero-copy read.
+        arr_data = obj.data.tobytes() if obj.data.c_contiguous \
+            else obj.tobytes()
path: vllm/v1/serial_utils.py
user: login: p88h
body: Well writing to the main buffer is yet another memory copy as well..
But now it  should be simpler and yes - moved this below to the aux buffers case
created_at: 2025-04-15T19:17:09Z
author_association: CONTRIBUTOR
diff_hunk: @@ -27,29 +30,32 @@ class MsgpackEncoder:
 
     Note that unlike vanilla `msgspec` Encoders, this interface is generally
     not thread-safe when encoding tensors / numpy arrays.
+
+    By default, arrays below 256B are serialized inline.
+    Larger will get sent via dedicated messages. 
+    Note that this is a per-tensor limit.
+
+    Sending multiple large messages via zeromq saturates memory very quickly.
+    See: https://github.com/vllm-project/vllm/issues/16185
     """
 
-    def __init__(self):
+    def __init__(self, size_threshold=256):
         self.encoder = msgpack.Encoder(enc_hook=self.enc_hook)
         # This is used as a local stash of buffers that we can then access from
         # our custom `msgspec` hook, `enc_hook`. We don't have a way to
         # pass custom data to the hook otherwise.
         self.aux_buffers: Optional[list[bytestr]] = None
+        self.size_threshold = size_threshold
 
     def encode(self, obj: Any) -> Sequence[bytestr]:
+        return self.encode_into(obj, bytearray())
+
+    def encode_into(self, obj: Any, buf: bytearray) -> Sequence[bytestr]:
path: vllm/v1/serial_utils.py
user: login: p88h
body: Right, but that can be useful in both directions. The engine also keeps a MsgpackEncoder around, so that can keep its encoding buffer automatically, rather than requiring encode_into()
created_at: 2025-04-15T19:18:49Z
author_association: CONTRIBUTOR
diff_hunk: @@ -65,6 +77,24 @@ def enc_hook(self, obj: Any) -> Any:
         if isinstance(obj, np.ndarray) and obj.dtype.kind not in ('O', 'V'):
             return self._encode_ndarray(obj)
 
+        if isinstance(obj, MultiModalKwargs):
+            mm: MultiModalKwargs = obj
+            if mm.modalities:
+                # ignore the main dict, it will be re-indexed.
+                # pass a list of MultiModalKwargsItem, then see below
+                # Any tensors *not* indexed by modality will be ignored.
+                return list(mm._items_by_modality.values())
+            # just return the main dict if there are no modalities
+            return dict(mm)
+
+        if isinstance(obj, MultiModalKwargsItem):
+            ret = []
+            for elem in obj.values():
+                # Encode as plain dictionary + special handling for .field
+                ret.append(
+                    asdict(elem) | {"field": self._encode_field(elem.field)})
+            return ret
path: vllm/v1/serial_utils.py
user: login: p88h
body: I'm still a bit at a loss as to why when this copying was followed by yet another copy into bytes, it resulted in _less_ memory usage when compared to copying and then using memoryviews. After fixing this, the good news is that performance improves a lot on longer inputs (same in all tests except last one, which improved by a whole 1 second) - and the memory usage is the same. I would assume there would be some drop.... but guess this is enough.
created_at: 2025-04-15T19:33:44Z
author_association: CONTRIBUTOR
diff_hunk: @@ -77,8 +102,13 @@ def _encode_ndarray(
         self, obj: np.ndarray
     ) -> tuple[str, tuple[int, ...], Union[int, memoryview]]:
         assert self.aux_buffers is not None
-        arr_data = obj.data if obj.data.c_contiguous else obj.tobytes()
-        if not obj.shape or obj.nbytes < MIN_NOCOPY_BUF_SIZE:
+        # Either copy the memoryview directly or flatten the array to bytes.
+        # Sending memoryviews is theoretically faster, but in this particular
+        # case, it triggers some unnecessary copies anyway.
+        # With this, the tensors can still be zero-copy read.
+        arr_data = obj.data.tobytes() if obj.data.c_contiguous \
+            else obj.tobytes()
path: vllm/v1/serial_utils.py
user: login: p88h
body: Actually, no. Msgspec needs the buffer to be contiguous in order to return it via the Ext handling. 

created_at: 2025-04-15T19:48:45Z
author_association: CONTRIBUTOR
diff_hunk: @@ -65,6 +84,28 @@ def enc_hook(self, obj: Any) -> Any:
         if isinstance(obj, np.ndarray) and obj.dtype.kind not in ('O', 'V'):
             return self._encode_ndarray(obj)
 
+        if isinstance(obj, MultiModalKwargs):
+            mm: MultiModalKwargs = obj
+            if mm.modalities:
+                # ignore the main dict, it will be re-indexed.
+                # pass a list of MultiModalKwargsItem, then see below
+                # Any tensors *not* indexed by modality will be ignored.
+                return list(mm._items_by_modality.values())
+            # just return the main dict if there are no modalities
+            return dict(mm)
+
+        if isinstance(obj, MultiModalKwargsItem):
+            ret = []
+            for elem in obj.values():
+                # Encode as plain dictionary + special handling for .field
+                ret.append({
+                    "modality": elem.modality,
+                    "key": elem.key,
+                    "data": self._encode_nested_tensors(elem.data),
+                    "field": self._encode_field(elem.field),
+                })
+            return ret
path: vllm/v1/serial_utils.py
user: login: njhill
body: Suggest to combine this to avoid re-entrance, and means the chain isn't needed when decoding:

```suggestion
    if isinstance(obj, MultiModalKwargs):
        mm: MultiModalKwargs = obj
        if not mm.modalities:
            # just return the main dict if there are no modalities.
            return dict(mm)
        
        # ignore the main dict, it will be re-indexed.
        # Encode a list of MultiModalKwargsItems as plain dicts
        # + special handling for .field.
        # Any tensors *not* indexed by modality will be ignored.
        return [{
            "modality": elem.modality,
            "key": elem.key,
            "data": self._encode_nested_tensors(elem.data),
            "field": self._encode_field(elem.field),
        } for item in mm._items_by_modality.values() for elem in item]
```
created_at: 2025-04-15T20:46:45Z
author_association: MEMBER
diff_hunk: @@ -27,32 +43,35 @@ class MsgpackEncoder:
 
     Note that unlike vanilla `msgspec` Encoders, this interface is generally
     not thread-safe when encoding tensors / numpy arrays.
+
+    By default, arrays below 256B are serialized inline.
+    Larger will get sent via dedicated messages. 
+    Note that this is a per-tensor limit.
+
+    Sending multiple large messages via zeromq saturates memory very quickly.
+    See: https://github.com/vllm-project/vllm/issues/16185
     """
 
-    def __init__(self):
+    def __init__(self, size_threshold=None):
+        if (size_threshold is None):
+            size_threshold = envs.VLLM_MSGPACK_ZERO_COPY_THRESHOLD
         self.encoder = msgpack.Encoder(enc_hook=self.enc_hook)
         # This is used as a local stash of buffers that we can then access from
         # our custom `msgspec` hook, `enc_hook`. We don't have a way to
         # pass custom data to the hook otherwise.
+        self.msg_buffer = bytearray()
path: vllm/v1/serial_utils.py
user: login: njhill
body: I'm still not sure about this. It implicitly assumes the returned bytearray is finished with by the caller before the method is called again.
created_at: 2025-04-15T20:49:20Z
author_association: MEMBER
diff_hunk: @@ -27,32 +43,35 @@ class MsgpackEncoder:
 
     Note that unlike vanilla `msgspec` Encoders, this interface is generally
     not thread-safe when encoding tensors / numpy arrays.
+
+    By default, arrays below 256B are serialized inline.
+    Larger will get sent via dedicated messages. 
+    Note that this is a per-tensor limit.
+
+    Sending multiple large messages via zeromq saturates memory very quickly.
+    See: https://github.com/vllm-project/vllm/issues/16185
     """
 
-    def __init__(self):
+    def __init__(self, size_threshold=None):
+        if (size_threshold is None):
path: vllm/v1/serial_utils.py
user: login: njhill
body: ```suggestion
        if size_threshold is None:
```
created_at: 2025-04-15T20:49:33Z
author_association: MEMBER
diff_hunk: @@ -27,32 +43,35 @@ class MsgpackEncoder:
 
     Note that unlike vanilla `msgspec` Encoders, this interface is generally
     not thread-safe when encoding tensors / numpy arrays.
+
+    By default, arrays below 256B are serialized inline.
+    Larger will get sent via dedicated messages. 
+    Note that this is a per-tensor limit.
+
+    Sending multiple large messages via zeromq saturates memory very quickly.
+    See: https://github.com/vllm-project/vllm/issues/16185
     """
 
-    def __init__(self):
+    def __init__(self, size_threshold=None):
path: vllm/v1/serial_utils.py
user: login: njhill
body: ```suggestion
    def __init__(self, size_threshold: Optional[int] = None):
```
created_at: 2025-04-15T20:50:57Z
author_association: MEMBER
diff_hunk: @@ -27,32 +43,35 @@ class MsgpackEncoder:
 
     Note that unlike vanilla `msgspec` Encoders, this interface is generally
     not thread-safe when encoding tensors / numpy arrays.
+
+    By default, arrays below 256B are serialized inline.
+    Larger will get sent via dedicated messages. 
+    Note that this is a per-tensor limit.
+
+    Sending multiple large messages via zeromq saturates memory very quickly.
+    See: https://github.com/vllm-project/vllm/issues/16185
     """
 
-    def __init__(self):
+    def __init__(self, size_threshold=None):
+        if (size_threshold is None):
+            size_threshold = envs.VLLM_MSGPACK_ZERO_COPY_THRESHOLD
         self.encoder = msgpack.Encoder(enc_hook=self.enc_hook)
         # This is used as a local stash of buffers that we can then access from
         # our custom `msgspec` hook, `enc_hook`. We don't have a way to
         # pass custom data to the hook otherwise.
+        self.msg_buffer = bytearray()
path: vllm/v1/serial_utils.py
user: login: p88h
body: ok, i see - you're right. The code in DP mode does not follow the simple encode&send pattern but rather keeps one message around. _Probably_ it would be more efficient to switch that around since that just keeps a tiny initialization message, but let's settle with a TODO for now - it wasn't the point of this change to fix this
created_at: 2025-04-15T21:57:18Z
author_association: CONTRIBUTOR
diff_hunk: @@ -126,13 +188,45 @@ def dec_hook(self, t: type, obj: Any) -> Any:
                 return self._decode_ndarray(obj)
             if issubclass(t, torch.Tensor):
                 return torch.from_numpy(self._decode_ndarray(obj))
+            if issubclass(t, MultiModalKwargs):
+                if isinstance(obj, list):
+                    return MultiModalKwargs.from_items(
+                        self._decode_mm_items(obj))
+                return MultiModalKwargs({
+                    k: self._decode_nested_tensors(v)
+                    for k, v in obj.items()
+                })
         return obj
 
     def _decode_ndarray(self, arr: Any) -> np.ndarray:
         dtype, shape, data = arr
-        buffer = self.aux_buffers[data] if isinstance(data, int) else data
+        # Copy from inline representation, otherwise Torch is unhappy since
+        # the returned memory is non-writeable.
+        buffer = self.aux_buffers[data] if isinstance(data, int) \
+            else bytearray(data)
         return np.ndarray(buffer=buffer, dtype=np.dtype(dtype), shape=shape)
 
+    def _decode_mm_items(self, obj: list) -> list[MultiModalKwargsItem]:
+        all = []
+        for item in obj:
+            elems = []
+            for v in item:
+                v["data"] = self._decode_nested_tensors(v["data"])
+                # Reconstruct the field processor using MultiModalFieldConfig
+                field = v["field"]
+                ctor = getattr(MultiModalFieldConfig, field[0])
+                v["field"] = ctor(None, *field[1:]).field
path: vllm/v1/serial_utils.py
user: login: njhill
body: small sugggestion
```suggestion
                factory_meth_name, *field_args = v["field"]
                factory_meth = getattr(MultiModalFieldConfig, factory_meth_name)
                v["field"] = factory_meth(None, *field_args).field
```
created_at: 2025-04-16T02:13:23Z
author_association: MEMBER
diff_hunk: @@ -92,6 +138,22 @@ def _encode_ndarray(
         # backing buffers that we've stashed in `aux_buffers`.
         return obj.dtype.str, obj.shape, data
 
+    def _encode_nested_tensors(self, obj: Any) -> NestedTensors:
+        if isinstance(obj, torch.Tensor):
+            return self._encode_ndarray(obj.numpy())
+        return [self._encode_nested_tensors(x) for x in obj]
+
+    def _encode_field(self, field: BaseMultiModalField):
path: vllm/v1/serial_utils.py
user: login: njhill
body: Rename to `_encode_mm_field` to sound less generic?
created_at: 2025-04-16T03:14:05Z
author_association: MEMBER
diff_hunk: @@ -48,10 +73,11 @@ def encode(self, obj: Any) -> Sequence[bytestr]:
         finally:
             self.aux_buffers = None
 
+    # TODO: would be nice to make this automatic
     def encode_into(self, obj: Any, buf: bytearray) -> Sequence[bytestr]:
         try:
             self.aux_buffers = [buf]
-            bufs = self.aux_buffers
+            bufs = [buf]
path: vllm/v1/serial_utils.py
user: login: njhill
body: Revert these changes?
created_at: 2025-04-16T03:15:00Z
author_association: MEMBER
diff_hunk: @@ -27,14 +42,24 @@ class MsgpackEncoder:
 
     Note that unlike vanilla `msgspec` Encoders, this interface is generally
     not thread-safe when encoding tensors / numpy arrays.
+
+    By default, arrays below 256B are serialized inline.
+    Larger will get sent via dedicated messages. 
+    Note that this is a per-tensor limit.
+
+    Sending multiple large messages via zeromq saturates memory very quickly.
+    See: https://github.com/vllm-project/vllm/issues/16185
path: vllm/v1/serial_utils.py
user: login: njhill
body: Should this comment still be here?
created_at: 2025-04-16T03:18:29Z
author_association: MEMBER
diff_hunk: @@ -27,32 +43,35 @@ class MsgpackEncoder:
 
     Note that unlike vanilla `msgspec` Encoders, this interface is generally
     not thread-safe when encoding tensors / numpy arrays.
+
+    By default, arrays below 256B are serialized inline.
+    Larger will get sent via dedicated messages. 
+    Note that this is a per-tensor limit.
+
+    Sending multiple large messages via zeromq saturates memory very quickly.
+    See: https://github.com/vllm-project/vllm/issues/16185
     """
 
-    def __init__(self):
+    def __init__(self, size_threshold=None):
+        if (size_threshold is None):
+            size_threshold = envs.VLLM_MSGPACK_ZERO_COPY_THRESHOLD
         self.encoder = msgpack.Encoder(enc_hook=self.enc_hook)
         # This is used as a local stash of buffers that we can then access from
         # our custom `msgspec` hook, `enc_hook`. We don't have a way to
         # pass custom data to the hook otherwise.
+        self.msg_buffer = bytearray()
path: vllm/v1/serial_utils.py
user: login: njhill
body: I wasn't really thinking of that usage specifically, more that if a buffer is reused it should be owned/provided by the caller.
created_at: 2025-04-16T03:26:29Z
author_association: MEMBER
diff_hunk: @@ -126,13 +188,45 @@ def dec_hook(self, t: type, obj: Any) -> Any:
                 return self._decode_ndarray(obj)
             if issubclass(t, torch.Tensor):
                 return torch.from_numpy(self._decode_ndarray(obj))
+            if issubclass(t, MultiModalKwargs):
+                if isinstance(obj, list):
+                    return MultiModalKwargs.from_items(
+                        self._decode_mm_items(obj))
+                return MultiModalKwargs({
+                    k: self._decode_nested_tensors(v)
+                    for k, v in obj.items()
+                })
         return obj
 
     def _decode_ndarray(self, arr: Any) -> np.ndarray:
         dtype, shape, data = arr
-        buffer = self.aux_buffers[data] if isinstance(data, int) else data
+        # Copy from inline representation, otherwise Torch is unhappy since
+        # the returned memory is non-writeable.
+        buffer = self.aux_buffers[data] if isinstance(data, int) \
+            else bytearray(data)
         return np.ndarray(buffer=buffer, dtype=np.dtype(dtype), shape=shape)
 
+    def _decode_mm_items(self, obj: list) -> list[MultiModalKwargsItem]:
+        all = []
+        for item in obj:
+            elems = []
+            for v in item:
+                v["data"] = self._decode_nested_tensors(v["data"])
+                # Reconstruct the field processor using MultiModalFieldConfig
+                field = v["field"]
+                ctor = getattr(MultiModalFieldConfig, field[0])
+                v["field"] = ctor(None, *field[1:]).field
path: vllm/v1/serial_utils.py
user: login: njhill
body: Also don't need to change it but I was wondering why not call the class constructor directly rather than creating via a temporary `MulitModalFieldConfig` obj?
created_at: 2025-04-16T03:42:09Z
author_association: MEMBER
diff_hunk: @@ -48,10 +73,11 @@ def encode(self, obj: Any) -> Sequence[bytestr]:
         finally:
             self.aux_buffers = None
 
+    # TODO: would be nice to make this automatic
     def encode_into(self, obj: Any, buf: bytearray) -> Sequence[bytestr]:
         try:
             self.aux_buffers = [buf]
-            bufs = self.aux_buffers
+            bufs = [buf]
path: vllm/v1/serial_utils.py
user: login: p88h
body: oops.
created_at: 2025-04-16T05:19:23Z
author_association: CONTRIBUTOR
diff_hunk: @@ -126,13 +188,45 @@ def dec_hook(self, t: type, obj: Any) -> Any:
                 return self._decode_ndarray(obj)
             if issubclass(t, torch.Tensor):
                 return torch.from_numpy(self._decode_ndarray(obj))
+            if issubclass(t, MultiModalKwargs):
+                if isinstance(obj, list):
+                    return MultiModalKwargs.from_items(
+                        self._decode_mm_items(obj))
+                return MultiModalKwargs({
+                    k: self._decode_nested_tensors(v)
+                    for k, v in obj.items()
+                })
         return obj
 
     def _decode_ndarray(self, arr: Any) -> np.ndarray:
         dtype, shape, data = arr
-        buffer = self.aux_buffers[data] if isinstance(data, int) else data
+        # Copy from inline representation, otherwise Torch is unhappy since
+        # the returned memory is non-writeable.
+        buffer = self.aux_buffers[data] if isinstance(data, int) \
+            else bytearray(data)
         return np.ndarray(buffer=buffer, dtype=np.dtype(dtype), shape=shape)
 
+    def _decode_mm_items(self, obj: list) -> list[MultiModalKwargsItem]:
+        all = []
+        for item in obj:
+            elems = []
+            for v in item:
+                v["data"] = self._decode_nested_tensors(v["data"])
+                # Reconstruct the field processor using MultiModalFieldConfig
+                field = v["field"]
+                ctor = getattr(MultiModalFieldConfig, field[0])
+                v["field"] = ctor(None, *field[1:]).field
path: vllm/v1/serial_utils.py
user: login: p88h
body: Mostly to avoid a switch statement, I guess. That would still end up calling the factory though.  
Alternative would belooking up the contructor name in globals, but that would also require serializing the field names. 
And the factory was just already there.... Also It's not a temporary MMFieldConfig, these are statics.

created_at: 2025-04-16T05:27:31Z
author_association: CONTRIBUTOR
diff_hunk: @@ -48,6 +69,7 @@ def encode(self, obj: Any) -> Sequence[bytestr]:
         finally:
             self.aux_buffers = None
 
+    # TODO: would be nice to make this automatic
path: vllm/v1/serial_utils.py
user: login: njhill
body: @p88h could you remove this comment too? I'm still not sure I understand it and it's unrelated to the PR anyhow.
created_at: 2025-04-16T14:54:48Z
author_association: MEMBER
diff_hunk: @@ -48,6 +69,7 @@ def encode(self, obj: Any) -> Sequence[bytestr]:
         finally:
             self.aux_buffers = None
 
+    # TODO: would be nice to make this automatic
path: vllm/v1/serial_utils.py
user: login: p88h
body: Sure, done
created_at: 2025-04-16T15:00:03Z
author_association: CONTRIBUTOR

PR comments:
user: login: github-actions[bot]
created_at: 2025-04-10T21:02:28Z
body: 👋 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.

🚀
author_association: NONE
user: login: ywang96
created_at: 2025-04-10T21:47:35Z
body: @p88h This is amazing! Have you tried running some benchmarks to see the throughput performance impact of this PR?
author_association: MEMBER
user: login: p88h
created_at: 2025-04-10T23:00:54Z
body: @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
author_association: CONTRIBUTOR
user: login: njhill
created_at: 2025-04-11T01:16:49Z
body: > 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: https://github.com/vllm-project/vllm/pull/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.
author_association: MEMBER
user: login: p88h
created_at: 2025-04-11T11:06:15Z
body: 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
```
author_association: CONTRIBUTOR
user: login: njhill
created_at: 2025-04-11T15:21:45Z
body: Thanks @p88h, the benchmark scripts look great!

> Interestingly it performs _much_ better with zero-copy disabled

In your table, you have the first two rows with disabled/enabled, and second two enabled/disabled... just thought to check that doesn't mean that the numbers are mixed up. It appears that without this PR, zero-copy helps, but in conjunction with this PR, it hurts (memory use especially)?

~If it is indeed the case that it hurts the performance, I have an idea what the reason _might_ be and have opened another PR to address that - it doesn't make sense to follow the zero-copy path for non-contiguous tensors since they'll incur an additional copy anyhow: https://github.com/vllm-project/vllm/pull/16492~ edit: still checking this


author_association: MEMBER
user: login: p88h
created_at: 2025-04-11T18:00:48Z
body: Regarding the table, yes, it's sorted weird, but the contents are correct. What matters most is the last row, anyways.

Overall I think this is some problem with either zmq itself or the way we are using it, which somehow ends up allocating ~500MB per async message used to send files. Which naturally becomes problematic when sending multiple messages, zero-copy or not.

That number is eerily similar to the buffer size specified in make_zmq_socket. But that sets the RCV/SNDBUF sizes, which control the _kernel_ level  socket buffers. HWM pools were set to 0 (unbound) . Unfortunately even setting the SNDHWM to some ridiculously low value doesn't seem to affect memory consumption much. 

It's all really suspect, since we're not even sending that much data through those sockets - so something, somewhere, definitely tries to allocate a very large buffer for every sent message. Still not sure what/where/why. Limiting the _size_ of the sent messages seems to help as well, as apparently that hits some other issues. 

For now, in this PR, the value of zero-copy-threshold is set to an absurd high value which makes it work reasonably well (but effectively turns off zero copy... and has to additionally copy memory on receive to make Tensors happy). 

On the positive side, the performance is significantly improved on large inputs, and zero-copy does not seem to actually improve it even when it does fit into memory, and when it doesn't, performance suffers.  On the small inputs, it seems to be just as good as it was. I can try replacing ZMQ with something different just to check whether that helps, but that's a topic for a followup PR. 

author_association: CONTRIBUTOR
user: login: njhill
created_at: 2025-04-11T19:33:16Z
body: Thanks @p88h. I have a few questions...

- This strange per-message usage is happening on the client/encode side?
- When you say per message do you mean per call to send_multipart in the zero-copy case? Or per call to send? Since send_multipart just makes multiple send calls.
- Presumably when you say per async message you mean when a bunch of different inputs are sent in a single call to `LLM.generate`?
- It looks like you're doing all of the benchmarking via `LLM.generate/chat` rather than via the openai API server, is that correct?

> zero-copy does not seem to actually improve it even when it does fit into memory, and when it doesn't, performance suffers.

What are the fits-into-memory and doesn't-fit-into-memory cases referring to here? And you're saying that the use of zero-copy still makes the memory use much higher and the performance worse, is that correct? I'm still a bit confused about the reason for that.

> has to additionally copy memory on receive to make Tensors happy

I have another PR which should avoid the extra copy on receive by ensuring a memoryview is used when decoding: https://github.com/vllm-project/vllm/pull/16492. I don't think it matters really if the backing data for the tensors is read-only. 


author_association: MEMBER
user: login: p88h
created_at: 2025-04-11T21:08:27Z
body: > Thanks @p88h. I have a few questions...
> 
> * This strange per-message usage is happening on the client/encode side?

The memory usage is on the sender / encode size. 
The number of messages sent grows primarily with the number of multimodal elements.
The elements are being sent sequentially, though - In the test I'm always handling 32 images in each phase, but they are split differently. So the memory usage in the '4 images' column is actually 8 prompts with 4 images each, 40 messages total.
And so, on for '32 images' case it's 1 prompt with 32 images, and 33 messages total.

That's what triggers the memory use -> when multiple messages are being sent as a multipart request. I suppose they are sent quickly enough that the receiving end doesn't have a chance to read them, but the amount of memory used is really disproportionate to the amount of data sent. 


> * When you say per message do you mean per call to send_multipart in the zero-copy case? Or per call to send? Since send_multipart just makes multiple send calls.

Per message in a largest send_multipart. 
Multiple send_multipart calls don't seem to trigger this. 

> * Presumably when you say per async message you mean when a bunch of different inputs are sent in a single call to `LLM.generate`?

See above - sorry for confusion, yes, there are multiple async prompts _as well_, but the problem is triggered depending on number of messages on a single prompt. 

> * It looks like you're doing all of the benchmarking via `LLM.generate/chat` rather than via the openai API server, is that correct?

Yes, I'm using the .generate path / offline-inference style. 
 
> > zero-copy does not seem to actually improve it even when it does fit into memory, and when it doesn't, performance suffers.
> 
> What are the fits-into-memory and doesn't-fit-into-memory cases referring to here? And you're saying that the use of zero-copy still makes the memory use much higher and the performance worse, is that correct? I'm still a bit confused about the reason for that.

I mean that as long as there is enough physical memory free, the performance seems to be OK, even though it does consume a lot of RAM. That could be an indication this memory is just _allocated_ and not necessarily used. Once the system starts running low on memory, though, the performance tanks. 

> > has to additionally copy memory on receive to make Tensors happy
> 
> I have another PR which should avoid the extra copy on receive by ensuring a memoryview is used when decoding: #16492. I don't think it matters really if the backing data for the tensors is read-only.

Well, Torch likely complains about this for some reason. If the tensors are processed later / stacked or sliced, then that may actually become an issue. 
author_association: CONTRIBUTOR
user: login: p88h
created_at: 2025-04-13T11:17:11Z
body: Spent some more time investigating this - it seems that no combination of zmq setup affects the problem. 

However, I found the more specific trigger : It's not just that zmq behaves weird with multiple messages, but, specifically when those messages are sent as memoryview. If I convert that to bytes (which creates a copy, but only on one side, so it's still better than packing it with the rest of the message) the problem disappears. 

Now, zmq does have some special handling of bytes - specifically, when a Frame is constructed on a bytes object, then that object is used as a backing buffer for the frame. If you provide a memoryview, it can create a copy of the buffer as frame._bytes, anytime bytes property is accessed, and one of the things that accesses it is the equality operator. Not sure _where_ that is triggered, but it's the only clue I have now. 

Now, copying _once_ still seems a better idea than wrapping into a msgpack buffer, and then potentially having to copy on the receiving side anyway, *and* we have to copy non-contiguous tensors anyway, so perhaps this is a good enough compromise for now. Performance-wise it's better or equal to the previous attempt - consumes 5.9GB, but the worst-case times dropped to 13.7 / 35.9 so even faster than previously _with_ zero-copy  (that used 22GB). 
author_association: CONTRIBUTOR
user: login: p88h
created_at: 2025-04-13T13:07:55Z
body: Benchmarks of the current approach:
```
config / prompts x images     | 8 x 4        | 4 x 8       | 2 x 16       | 1 x 32        |
------------------------------+--------------+-------------+--------------+---------------+
zero-copy read and write      | 3.4 GB 10.7s | 4.8 GB 9.2s | 8.5 GB 10.5s | 22.0 GB 14.6s |
zero-copy threshold = 256MB   | 3.3 GB 10.7s | 4.0 GB 9.3s | 4.8 GB 10.7s | 5.9 GB  13.8s |
zero-copy read, copy on write | 3.3 GB 10.7s | 4.0 GB 9.2s | 4.8 GB 10.5s | 5.9 GB  13.6s |
```
Each group run has 32 images total, sent in a differnt layout.
Each images is 1Mpix. Each processed image being sent is 16MB.

This now seems always better than 'full' zero-copy version, and on-par memory-wise but faster than the non-zero-copy version with custom serialization. 

author_association: CONTRIBUTOR
user: login: njhill
created_at: 2025-04-14T17:17:54Z
body: Thanks @p88h! I am afk most of today but will review/respond towards eod hopefully.
author_association: MEMBER
user: login: p88h
created_at: 2025-04-15T08:01:55Z
body: Re: zmq & memoryview I've extracted the communication pattern into a separate test ... and that does *not* trigger the issue -- I tried making the comm pattern as close as possible to the VLLM codebase, and it's still just working happily -- with memoryview variant performing just a tiny amount better than the copying one (the small magnitude of the change is interesting - it seems that the copied buffer pages are actually passed directly via kernel space, probably in COW mode. This is demonstrated by the RSS of the sender + RSS of the receiver being significantly higher than the cumulative RSS use of the whole process, the only case that's true is when they share memory)

However, that doesn't seem to hold true for memoryview pages - they must get copied over on the receiving end. So one way or another _something_ is copying this memory; meaning the current approach is likely the best possible for now at least. Using shared memory directly could potentially alleviate that, but that would not scale to DP over TCP cases. 

Isolated zmq test : https://gist.github.com/p88h/81659b458acc244d7117d2f84f05614c

author_association: CONTRIBUTOR
user: login: p88h
created_at: 2025-04-15T20:50:38Z
body: Current results. Switched to async processing in the benchmark to improve throughput a bit,
so had to re-run this for consistency (also, the asdict() removal helps the in-line case as well)

```
config / prompts*imgs | 8 x 4        | 4 x 8        | 2 x 16       | 1 x 32        |
----------------------+--------------+--------------+--------------+---------------+
single threaded       | 
main (baseline)       | 3.6 GB 12.1s | 5.2 GB 13.3s | 9.9 GB 24.5s | DNC / OOM     |
#16432 + in-line only | 3.4 GB 11.2s | 4.1 GB 9.3s  | 5.0 GB 10.7s | 6.1 GB  12.2s |
#16432 + zero-copy    | 3.3 GB 10.7s | 4.0 GB 9.1s  | 4.8 GB 10.0s | 5.9 GB  12.0s |
```
author_association: CONTRIBUTOR
user: login: njhill
created_at: 2025-04-16T22:25:28Z
body: 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 https://github.com/vllm-project/vllm/pull/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.
author_association: MEMBER
user: login: njhill
created_at: 2025-04-16T23:28:45Z
body: 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.
author_association: MEMBER
user: login: p88h
created_at: 2025-04-17T06:33:56Z
body: Thank you ! I was about to go back to debugging this morning ;) 

author_association: CONTRIBUTOR