Implement efficient serialization of MultiModalKwargs

In addition to serializing base Tensors, this now allows to pass
Tensors embedded in MultiModalKwargs correctly.

Handles both V0 and V1 style args.

Improves memory usage with large multimodal payloads by a further
50% (but still not on par with single-threaded behavior).

Author: p88h

tests/v1/test_serial_utils.py

@@ -1,10 +1,15 @@
 # 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, NestedTensors)
 from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder
 
 
@@ -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)])
+
+
 def assert_equal(obj1: MyType, obj2: MyType):
     assert torch.equal(obj1.tensor1, obj2.tensor1)
     assert obj1.a_string == obj2.a_string

vllm/v1/serial_utils.py

@@ -2,7 +2,9 @@
 
 import pickle
 from collections.abc import Sequence
+from dataclasses import asdict
 from inspect import isclass
+from itertools import chain
 from types import FunctionType
 from typing import Any, Optional, Union
 
@@ -12,6 +14,9 @@
 import zmq
 from msgspec import msgpack
 
+from vllm.multimodal.inputs import (MultiModalFieldElem, MultiModalKwargs,
+                                    MultiModalKwargsItem, NestedTensors)
+
 CUSTOM_TYPE_PICKLE = 1
 CUSTOM_TYPE_CLOUDPICKLE = 2
 
@@ -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)
+                rd[k] = vv
+            return rd
+
         if isinstance(obj, FunctionType):
             # `pickle` is generally faster than cloudpickle, but can have
             # problems serializing methods.
@@ -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))
             if issubclass(t, torch.Tensor):
                 return torch.from_numpy(self._decode_ndarray(obj))
         return obj
@@ -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}")
+
     def ext_hook(self, code: int, data: memoryview) -> Any:
         if code == CUSTOM_TYPE_PICKLE:
             return pickle.loads(data)