[Bugfix][Multimodal] Fix reference counting of the SHM processor cache (#59695)

Signed-off-by: tenotek <50898317+hdq66666@users.noreply.github.com>
Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>

tests/distributed/test_shm_storage.py

@@ -291,6 +291,93 @@ class TestSingleWriterShmObjectStorage(unittest.TestCase):
             reader.touch("key", address, monotonic_id, signature)
 
 
+class TestShmObjectStorageReferenceCounting(unittest.TestCase):
+    """One writer and four readers, as with tensor-parallel workers."""
+
+    def setUp(self):
+        ring_buffer = SingleWriterShmRingBuffer(
+            data_buffer_size=1024 * 100, create=True
+        )
+        self.writer = SingleWriterShmObjectStorage(
+            max_object_size=1024 * 10,
+            n_readers=4,
+            ring_buffer=ring_buffer,
+            serde_class=MsgpackSerde,
+        )
+        self.addCleanup(self.writer.close)
+        handle = self.writer.handle()
+        handle.reader_lock = Lock()
+        self.readers = [
+            SingleWriterShmObjectStorage.create_from_handle(handle) for _ in range(4)
+        ]
+        for reader in self.readers:
+            self.addCleanup(reader.close)
+
+    def _reference(self, key: str) -> tuple[int, int, list[int], str]:
+        """Take one writer reference to the key, as one request item does."""
+        if self.writer.is_cached(key):
+            address, monotonic_id = self.writer.get_cached(key)
+        else:
+            address, monotonic_id = self.writer.put(key, "value")
+        return address, monotonic_id, self.writer.get_signature(key), key
+
+    @staticmethod
+    def _touch(reader: SingleWriterShmObjectStorage, reference) -> None:
+        address, monotonic_id, signature, key = reference
+        reader.touch(key, address, monotonic_id, signature)
+
+    def _is_evictable(self, key: str) -> bool:
+        self.writer.free_unused()
+        return not self.writer.is_cached(key)
+
+    def test_item_referenced_again_before_every_reader_fetched_it(self):
+        """A second request must not leave the item unevictable forever."""
+        first = self._reference("image")
+        self.writer.touch("image")
+        second = self._reference("image")
+        self.writer.release_touches()
+
+        # Every reader starts on the first request; the ones that finish it
+        # go on to the second while the others are still fetching.
+        for reader in self.readers:
+            self._touch(reader, first)
+        for reader in self.readers:
+            reader.get(*first)
+            self._touch(reader, second)
+            reader.get(*second)
+
+        self.assertTrue(self._is_evictable("image"))
+
+    def test_new_item_twice_in_one_request_is_not_evicted_early(self):
+        """The item must stay until the last reader has fetched it."""
+        self.writer.touch("image")
+        references = [self._reference("image"), self._reference("image")]
+        self.writer.release_touches()
+
+        for reader in self.readers[:-1]:
+            for reference in references:
+                self._touch(reader, reference)
+            for reference in references:
+                reader.get(*reference)
+        for reference in references:
+            self._touch(self.readers[-1], reference)
+        self.assertFalse(self._is_evictable("image"))
+
+        for reference in references:
+            self.readers[-1].get(*reference)
+        self.assertTrue(self._is_evictable("image"))
+
+    def test_touch_protects_item_until_released(self):
+        reference = self._reference("image")
+        for reader in self.readers:
+            reader.get(*reference)
+
+        self.writer.touch("image")
+        self.assertFalse(self._is_evictable("image"))
+        self.writer.release_touches()
+        self.assertTrue(self._is_evictable("image"))
+
+
 # Reader process function
 def reader_process(process_id, storage_handle, items_to_read):
     """Reader process that connects to existing shared memory and reads data."""

vllm/distributed/device_communicators/shm_object_storage.py

@@ -527,6 +527,9 @@ class SingleWriterShmObjectStorage:
             self.id_index: dict[int, str] = {}
             # Writer flag to track in-use status: monotonic_id -> count
             self.writer_flag: dict[int, int] = {}
+            # Items touch() keeps from being freed until they are looked up
+            # or release_touches() is called
+            self._touched: set[int] = set()
         else:
             if reader_lock is None:
                 raise ValueError("Lock must be provided for readers.")
@@ -540,6 +543,7 @@ class SingleWriterShmObjectStorage:
             self.key_index.clear()
             self.id_index.clear()
             self.writer_flag.clear()
+            self._touched.clear()
             logger.debug("Object storage cleared and reinitialized.")
 
     def copy_to_buffer(
@@ -664,8 +668,13 @@ class SingleWriterShmObjectStorage:
         """Get the cached object by key if it exists."""
         address, monotonic_id = self.key_index[key]
         self.increment_writer_flag(monotonic_id)
+        self._touched.discard(monotonic_id)
         return address, monotonic_id
 
+    def release_touches(self) -> None:
+        """Stop protecting touched items that were not looked up."""
+        self._touched.clear()
+
     def get_signature(self, key: str) -> list[int]:
         """Sign the handle of a cached object so readers can verify it."""
         address, monotonic_id = self.key_index[key]
@@ -756,10 +765,11 @@ class SingleWriterShmObjectStorage:
         monotonic_id: int = 0,
         signature: list[int] | None = None,
     ) -> None:
-        """Touch an existing cached item to update its eviction status.
+        """Touch an existing cached item.
 
-        For writers (ShmObjectStoreSenderCache): Increment writer_flag
-        For readers (ShmObjectStoreReceiverCache): Increment reader_count
+        For writers (ShmObjectStoreSenderCache): Protect the item from
+        eviction until get_cached or release_touches
+        For readers (ShmObjectStoreReceiverCache): Validate the handle
 
         Args:
             key: String key of the object to touch
@@ -771,25 +781,12 @@ class SingleWriterShmObjectStorage:
         if self._reader_lock is None:
             if key not in self.key_index:
                 return None
-            address, monotonic_id = self.key_index[key]
-            # Writer side: increment writer_flag to raise eviction threshold
-            self.increment_writer_flag(monotonic_id)
+            self._touched.add(self.key_index[key][1])
         else:
             # Reject forged handles before dereferencing the supplied address.
             self.verify_signature(key, address, monotonic_id, signature)
-            with (
-                self._reader_lock,
-                self.ring_buffer.access_buf(address) as (data_view, buf_metadata),
-            ):
+            with self.ring_buffer.access_buf(address) as (_, buf_metadata):
                 self._validate_monotonic_id(address, monotonic_id, buf_metadata)
-                reader_count = self.ring_buffer.byte2int(data_view[: self.flag_bytes])
-
-                # NOTE(Long):
-                # Avoid increasing flag on newly added item (sync with sender)
-                # Since when a new item is added
-                # pre-touch has no effect on writer side
-                if reader_count >= self.n_readers:
-                    self.increment_reader_flag(data_view[: self.flag_bytes])
 
     def close(self) -> None:
         """Close the shared memory."""
@@ -823,6 +820,8 @@ class SingleWriterShmObjectStorage:
         """Default is_free function that checks if the first 4 bytes are zero.
         This indicates that the buffer is free.
         """
+        if id in self._touched:
+            return False
         reader_count = int.from_bytes(buf[0:4], "little", signed=True)
         writer_count = self.writer_flag[id]
         return reader_count >= writer_count * self.n_readers

vllm/multimodal/cache/base.py

@@ -377,6 +377,10 @@ class BaseMultiModalProcessorCache(
         """
         raise NotImplementedError
 
+    def release_sender_touches(self) -> None:
+        """Release the items that were touched but not updated afterwards."""
+        return None
+
     @abstractmethod
     def make_stats(self, *, delta: bool = False) -> CacheInfo:
         """Get (and reset) the multi-modal cache stats.

vllm/multimodal/cache/shm.py

@@ -170,10 +170,13 @@ class ShmObjectStoreSenderCache(BaseMultiModalProcessorCache):
 
     @override
     def touch_sender_cache_item(self, mm_hash: str) -> None:
-        """Touch the item in shared memory cache to prevent eviction.
-        Increments writer_flag on sender side."""
+        """Touch the item in shared memory cache to prevent eviction."""
         self._shm_cache.touch(mm_hash)
 
+    @override
+    def release_sender_touches(self) -> None:
+        self._shm_cache.release_touches()
+
     @override
     def validate_input_item(
         self,
@@ -305,8 +308,7 @@ class ShmObjectStoreReceiverCache(BaseMultiModalReceiverCache):
         mm_hash: str,
         mm_item: MultiModalKwargsItem | None = None,
     ) -> None:
-        """Touch the item in shared memory cache to prevent eviction.
-        Increments reader_count on receiver side."""
+        """Validate the item's handle in shared memory cache."""
         assert mm_item is not None
         if (handle := _get_shm_handle(mm_item)) is not None:
             address, monotonic_id, signature = handle

vllm/multimodal/processing/processor.py

@@ -1534,13 +1534,16 @@ class BaseMultiModalProcessor(ABC, Generic[_I]):
         )
 
         with timing_ctx.record("merge_mm_kwargs"):
-            mm_kwargs, mm_prompt_updates = self._merge_mm_kwargs(
-                cache,
-                mm_hashes=mm_hashes,
-                mm_is_cached=mm_is_cached,
-                mm_missing_kwargs=mm_missing_kwargs,
-                mm_missing_prompt_updates=mm_missing_prompt_updates,
-            )
+            try:
+                mm_kwargs, mm_prompt_updates = self._merge_mm_kwargs(
+                    cache,
+                    mm_hashes=mm_hashes,
+                    mm_is_cached=mm_is_cached,
+                    mm_missing_kwargs=mm_missing_kwargs,
+                    mm_missing_prompt_updates=mm_missing_prompt_updates,
+                )
+            finally:
+                cache.release_sender_touches()
 
         return MultiModalProcessingResult(
             prompt_ids=self._postprocess_prompt(inputs.prompt),