@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 443dccd3c9

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

tenotek committed Oct 2, 2026 at 22:14 UTC 443dccd3c9448f586634bfdb7079acd7eab35ead

5 files changed +124 -29

tests/distributed/test_shm_storage.py

+87

@@ -291,6 +291,93 @@ class TestSingleWriterShmObjectStorage(unittest.TestCase):

291 reader.touch("key", address, monotonic_id, signature)

292

293

294 +class TestShmObjectStorageReferenceCounting(unittest.TestCase):

295 + """One writer and four readers, as with tensor-parallel workers."""

296 +

297 + def setUp(self):

298 + ring_buffer = SingleWriterShmRingBuffer(

299 + data_buffer_size=1024 * 100, create=True

300 + )

301 + self.writer = SingleWriterShmObjectStorage(

302 + max_object_size=1024 * 10,

303 + n_readers=4,

304 + ring_buffer=ring_buffer,

305 + serde_class=MsgpackSerde,

306 + )

307 + self.addCleanup(self.writer.close)

308 + handle = self.writer.handle()

309 + handle.reader_lock = Lock()

310 + self.readers = [

311 + SingleWriterShmObjectStorage.create_from_handle(handle) for _ in range(4)

312 + ]

313 + for reader in self.readers:

314 + self.addCleanup(reader.close)

315 +

316 + def _reference(self, key: str) -> tuple[int, int, list[int], str]:

317 + """Take one writer reference to the key, as one request item does."""

318 + if self.writer.is_cached(key):

319 + address, monotonic_id = self.writer.get_cached(key)

320 + else:

321 + address, monotonic_id = self.writer.put(key, "value")

322 + return address, monotonic_id, self.writer.get_signature(key), key

323 +

324 + @staticmethod

325 + def _touch(reader: SingleWriterShmObjectStorage, reference) -> None:

326 + address, monotonic_id, signature, key = reference

327 + reader.touch(key, address, monotonic_id, signature)

328 +

329 + def _is_evictable(self, key: str) -> bool:

330 + self.writer.free_unused()

331 + return not self.writer.is_cached(key)

332 +

333 + def test_item_referenced_again_before_every_reader_fetched_it(self):

334 + """A second request must not leave the item unevictable forever."""

335 + first = self._reference("image")

336 + self.writer.touch("image")

337 + second = self._reference("image")

338 + self.writer.release_touches()

339 +

340 + # Every reader starts on the first request; the ones that finish it

341 + # go on to the second while the others are still fetching.

342 + for reader in self.readers:

343 + self._touch(reader, first)

344 + for reader in self.readers:

345 + reader.get(*first)

346 + self._touch(reader, second)

347 + reader.get(*second)

348 +

349 + self.assertTrue(self._is_evictable("image"))

350 +

351 + def test_new_item_twice_in_one_request_is_not_evicted_early(self):

352 + """The item must stay until the last reader has fetched it."""

353 + self.writer.touch("image")

354 + references = [self._reference("image"), self._reference("image")]

355 + self.writer.release_touches()

356 +

357 + for reader in self.readers[:-1]:

358 + for reference in references:

359 + self._touch(reader, reference)

360 + for reference in references:

361 + reader.get(*reference)

362 + for reference in references:

363 + self._touch(self.readers[-1], reference)

364 + self.assertFalse(self._is_evictable("image"))

365 +

366 + for reference in references:

367 + self.readers[-1].get(*reference)

368 + self.assertTrue(self._is_evictable("image"))

369 +

370 + def test_touch_protects_item_until_released(self):

371 + reference = self._reference("image")

372 + for reader in self.readers:

373 + reader.get(*reference)

374 +

375 + self.writer.touch("image")

376 + self.assertFalse(self._is_evictable("image"))

377 + self.writer.release_touches()

378 + self.assertTrue(self._is_evictable("image"))

379 +

380 +

381 # Reader process function

382 def reader_process(process_id, storage_handle, items_to_read):

383 """Reader process that connects to existing shared memory and reads data."""

vllm/distributed/device_communicators/shm_object_storage.py

+17 -18

@@ -527,6 +527,9 @@ class SingleWriterShmObjectStorage:

527 self.id_index: dict[int, str] = {}

528 # Writer flag to track in-use status: monotonic_id -> count

529 self.writer_flag: dict[int, int] = {}

530 + # Items touch() keeps from being freed until they are looked up

531 + # or release_touches() is called

532 + self._touched: set[int] = set()

533 else:

534 if reader_lock is None:

535 raise ValueError("Lock must be provided for readers.")

@@ -540,6 +543,7 @@ class SingleWriterShmObjectStorage:

543 self.key_index.clear()

544 self.id_index.clear()

545 self.writer_flag.clear()

546 + self._touched.clear()

547 logger.debug("Object storage cleared and reinitialized.")

548

549 def copy_to_buffer(

@@ -664,8 +668,13 @@ class SingleWriterShmObjectStorage:

668 """Get the cached object by key if it exists."""

669 address, monotonic_id = self.key_index[key]

670 self.increment_writer_flag(monotonic_id)

671 + self._touched.discard(monotonic_id)

672 return address, monotonic_id

673

674 + def release_touches(self) -> None:

675 + """Stop protecting touched items that were not looked up."""

676 + self._touched.clear()

677 +

678 def get_signature(self, key: str) -> list[int]:

679 """Sign the handle of a cached object so readers can verify it."""

680 address, monotonic_id = self.key_index[key]

@@ -756,10 +765,11 @@ class SingleWriterShmObjectStorage:

765 monotonic_id: int = 0,

766 signature: list[int] | None = None,

767 ) -> None:

759 - """Touch an existing cached item to update its eviction status.

768 + """Touch an existing cached item.

769

761 - For writers (ShmObjectStoreSenderCache): Increment writer_flag

762 - For readers (ShmObjectStoreReceiverCache): Increment reader_count

770 + For writers (ShmObjectStoreSenderCache): Protect the item from

771 + eviction until get_cached or release_touches

772 + For readers (ShmObjectStoreReceiverCache): Validate the handle

773

774 Args:

775 key: String key of the object to touch

@@ -771,25 +781,12 @@ class SingleWriterShmObjectStorage:

781 if self._reader_lock is None:

782 if key not in self.key_index:

783 return None

774 - address, monotonic_id = self.key_index[key]

775 - # Writer side: increment writer_flag to raise eviction threshold

776 - self.increment_writer_flag(monotonic_id)

784 + self._touched.add(self.key_index[key][1])

785 else:

786 # Reject forged handles before dereferencing the supplied address.

787 self.verify_signature(key, address, monotonic_id, signature)

780 - with (

781 - self._reader_lock,

782 - self.ring_buffer.access_buf(address) as (data_view, buf_metadata),

783 - ):

788 + with self.ring_buffer.access_buf(address) as (_, buf_metadata):

789 self._validate_monotonic_id(address, monotonic_id, buf_metadata)

785 - reader_count = self.ring_buffer.byte2int(data_view[: self.flag_bytes])

786 -

787 - # NOTE(Long):

788 - # Avoid increasing flag on newly added item (sync with sender)

789 - # Since when a new item is added

790 - # pre-touch has no effect on writer side

791 - if reader_count >= self.n_readers:

792 - self.increment_reader_flag(data_view[: self.flag_bytes])

790

791 def close(self) -> None:

792 """Close the shared memory."""

@@ -823,6 +820,8 @@ class SingleWriterShmObjectStorage:

820 """Default is_free function that checks if the first 4 bytes are zero.

821 This indicates that the buffer is free.

822 """

823 + if id in self._touched:

824 + return False

825 reader_count = int.from_bytes(buf[0:4], "little", signed=True)

826 writer_count = self.writer_flag[id]

827 return reader_count >= writer_count * self.n_readers

vllm/multimodal/cache/base.py

@@ -377,6 +377,10 @@ class BaseMultiModalProcessorCache(

377 """

378 raise NotImplementedError

379

380 + def release_sender_touches(self) -> None:

381 + """Release the items that were touched but not updated afterwards."""

382 + return None

383 +

384 @abstractmethod

385 def make_stats(self, *, delta: bool = False) -> CacheInfo:

386 """Get (and reset) the multi-modal cache stats.

vllm/multimodal/cache/shm.py

+6 -4

@@ -170,10 +170,13 @@ class ShmObjectStoreSenderCache(BaseMultiModalProcessorCache):

170

171 @override

172 def touch_sender_cache_item(self, mm_hash: str) -> None:

173 - """Touch the item in shared memory cache to prevent eviction.

174 - Increments writer_flag on sender side."""

173 + """Touch the item in shared memory cache to prevent eviction."""

174 self._shm_cache.touch(mm_hash)

175

176 + @override

177 + def release_sender_touches(self) -> None:

178 + self._shm_cache.release_touches()

179 +

180 @override

181 def validate_input_item(

182 self,

@@ -305,8 +308,7 @@ class ShmObjectStoreReceiverCache(BaseMultiModalReceiverCache):

308 mm_hash: str,

309 mm_item: MultiModalKwargsItem | None = None,

310 ) -> None:

308 - """Touch the item in shared memory cache to prevent eviction.

309 - Increments reader_count on receiver side."""

311 + """Validate the item's handle in shared memory cache."""

312 assert mm_item is not None

313 if (handle := _get_shm_handle(mm_item)) is not None:

314 address, monotonic_id, signature = handle

vllm/multimodal/processing/processor.py

+10 -7

@@ -1534,13 +1534,16 @@ class BaseMultiModalProcessor(ABC, Generic[_I]):

1534 )

1535

1536 with timing_ctx.record("merge_mm_kwargs"):

1537 - mm_kwargs, mm_prompt_updates = self._merge_mm_kwargs(

1538 - cache,

1539 - mm_hashes=mm_hashes,

1540 - mm_is_cached=mm_is_cached,

1541 - mm_missing_kwargs=mm_missing_kwargs,

1542 - mm_missing_prompt_updates=mm_missing_prompt_updates,

1543 - )

1537 + try:

1538 + mm_kwargs, mm_prompt_updates = self._merge_mm_kwargs(

1539 + cache,

1540 + mm_hashes=mm_hashes,

1541 + mm_is_cached=mm_is_cached,

1542 + mm_missing_kwargs=mm_missing_kwargs,

1543 + mm_missing_prompt_updates=mm_missing_prompt_updates,

1544 + )

1545 + finally:

1546 + cache.release_sender_touches()

1547

1548 return MultiModalProcessingResult(

1549 prompt_ids=self._postprocess_prompt(inputs.prompt),