A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 7, 02:45 UTC
[LMCache MP Patch]: Race Condition + Duplicated Block Ids (#35831)
Samuel Shen committed
Mar 7, 2026 at 13:52 UTC
5d6aae4577590cd6b6a604f9e74c17c5f234271d
1 file changed
+23
-4
vllm/distributed/kv_transfer/kv_connector/v1/lmcache_mp_connector.py
+23
-4
@@ -336,11 +336,21 @@ class LMCacheMPRequestMetadata:
336
start_token_idx = start * vllm_block_size
337
end_token_idx = end * vllm_block_size
338
token_ids = list(tracker.all_token_ids)
339
+
340
+ # Compute how many tokens at the start of the retrieve range
341
+ # overlap with APC-shared blocks. The server must skip writing
342
+ # to these positions to avoid a cross-stream data race: the
343
+ # retrieve writes on the LMCache CUDA stream while concurrent
344
+ # requests may read these APC-shared blocks on the vLLM stream.
345
+ apc_overlap_blocks = tracker.num_vllm_hit_blocks - start
346
+ skip_first_n_tokens = apc_overlap_blocks * vllm_block_size
347
+
348
op = LoadStoreOp(
349
token_ids=token_ids,
350
block_ids=block_ids,
351
start=start_token_idx,
352
end=end_token_idx,
353
+ skip_first_n_tokens=skip_first_n_tokens,
354
)
355
356
ret = LMCacheMPRequestMetadata(
@@ -700,13 +710,22 @@ class LMCacheMPConnector(KVConnectorBase_V1):
710
num_external_tokens (int): the number of tokens that will be
711
loaded from the external KV cache.
712
"""
703
- # NOTE: the `blocks` are NEW BLOCKS allocated for this request.
713
+ # NOTE: `blocks` comes from kv_cache_manager.get_blocks(request_id),
714
+ # which returns ALL blocks for the request (not just newly allocated).
715
+ # This function may be called twice for async-load requests:
716
+ # 1st call: blocks = initial allocation (APC + fresh)
717
+ # 2nd call: blocks = all blocks
718
+ # (initial + newly allocated for remaining tokens)
719
+ # We must only append the NEW blocks beyond what's already tracked
720
+ # to avoid duplication, which would corrupt the store path's block indexing.
721
tracker = self._get_request_tracker(request.request_id)
722
block_ids = reformat_block_ids(blocks.get_block_ids())
723
707
- # No matter we need to retrieve or not, we need to update
708
- # the block ids into the tracker
709
- tracker.append_block_ids(block_ids)
724
+ # Only append blocks beyond what's already tracked
725
+ existing_count = len(tracker.allocated_block_ids)
726
+ new_block_ids = block_ids[existing_count:]
727
+ if new_block_ids:
728
+ tracker.append_block_ids(new_block_ids)
729
730
# Update the state of the tracker
731
condition = tracker.needs_retrieve()