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 2, 07:15 UTC
[Bugfix][Mooncake] Release GPU pin on failed store in MooncakeStoreConnector (#43742)
Signed-off-by: Dao Le <Dao007forever@gmail.com> Co-authored-by: Claude <noreply@anthropic.com>
Dao007forever committed
Jun 1, 2026 at 18:29 UTC
d68f0b220efc2dcadf5e4f24a9bccccf268ff021
2 files changed
+193
-161
tests/v1/kv_connector/unit/test_mooncake_store_worker.py
+29
@@ -461,6 +461,35 @@ def test_store_sending_thread_only_skips_on_no_available_handle():
461
assert store.batch_put_from_multi_buffers.call_count == 2
462
463
464
+def test_store_sending_thread_releases_pin_on_batch_is_exist_failure():
465
+ # `batch_is_exist` raising must still decrement `stored_requests` so the
466
+ # scheduler can drop `delay_free_blocks` and release the pinned GPU blocks.
467
+ store = MagicMock()
468
+ store.batch_is_exist.side_effect = RuntimeError("mooncake down")
469
+ thread = _make_store_sending_thread(store)
470
+
471
+ thread.add_stored_request("req-a")
472
+ with pytest.raises(RuntimeError):
473
+ thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"]))
474
+
475
+ assert thread.stored_requests["req-a"] == 0
476
+ store.batch_put_from_multi_buffers.assert_not_called()
477
+
478
+
479
+def test_store_sending_thread_releases_pin_on_batch_put_failure():
480
+ # `batch_put_from_multi_buffers` raising is logged (not re-raised), and the
481
+ # pin must still be released through the finally block.
482
+ store = MagicMock()
483
+ store.batch_is_exist.return_value = [0, 0]
484
+ store.batch_put_from_multi_buffers.side_effect = RuntimeError("rdma error")
485
+ thread = _make_store_sending_thread(store)
486
+
487
+ thread.add_stored_request("req-a")
488
+ thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"]))
489
+
490
+ assert thread.stored_requests["req-a"] == 0
491
+
492
+
493
def test_store_recving_thread_reports_failed_block_ids():
494
store = MagicMock()
495
store.batch_get_into_multi_buffers.return_value = [256, -5, -7]
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py
+164
-161
@@ -517,187 +517,190 @@ class KVCacheStoreSendingThread(KVTransferThread):
517
if req_id not in self.stored_requests:
518
self.request_queue.task_done()
519
return
520
- if token_len == 0:
521
- self.dec_stored_request(req_id)
522
- self.request_queue.task_done()
523
- return
524
- if self._should_skip_request(req_id):
525
- logger.debug(
526
- "Skipping Mooncake store for request %s while CPU/disk offloading "
527
- "is under pressure",
528
- req_id,
529
- )
530
- self.dec_stored_request(req_id)
531
- self.request_queue.task_done()
532
- return
520
534
- # Within each lcm region only per-spec relevant chunks are loaded
535
- # (e.g., SWA or linear attn), so mask out irrelevant chunks
536
- store_masks = self.coord.store_mask(token_len)
537
- starts: list[int] = []
538
- ends: list[int] = []
539
- keys: list[str] = []
540
- block_hashes: list[BlockHash] = []
541
- group_indices: list[int] = []
542
- for g_idx, db in enumerate(self.token_databases):
543
- mask = store_masks[g_idx]
544
- for chunk_idx, (start, end, key) in enumerate(
545
- db.process_tokens(token_len, req_meta.block_hashes)
546
- ):
547
- if chunk_idx >= len(mask) or not mask[chunk_idx]:
548
- continue
549
- starts.append(start)
550
- ends.append(end)
551
- keys.append(key.to_string())
552
- block_hashes.append(req_meta.block_hashes[chunk_idx])
553
- group_indices.append(g_idx)
554
-
555
- # Apply put_step striding for TP
556
- sl = slice(self.tp_rank % self.put_step, None, self.put_step)
557
- starts = starts[sl]
558
- ends = ends[sl]
559
- keys = keys[sl]
560
- block_hashes = block_hashes[sl]
561
- group_indices = group_indices[sl]
562
-
563
- if not keys:
564
- self.dec_stored_request(req_id)
565
- return
566
-
567
- # Check which blocks already exist (dedup)
568
- save_exists_start = time.perf_counter()
521
+ # Decrement the in-flight counter and signal task_done() in `finally`
522
+ # so the scheduler can release the GPU blocks it pinned for this
523
+ # request (via `delay_free_blocks`) even when the store path raises.
524
try:
570
- exists_states = self.store.batch_is_exist(keys)
571
- except Exception:
525
+ if token_len == 0:
526
+ return
527
+ if self._should_skip_request(req_id):
528
+ logger.debug(
529
+ "Skipping Mooncake store for request %s while CPU/disk "
530
+ "offloading is under pressure",
531
+ req_id,
532
+ )
533
+ return
534
+
535
+ # Within each lcm region only per-spec relevant chunks are loaded
536
+ # (e.g., SWA or linear attn), so mask out irrelevant chunks
537
+ store_masks = self.coord.store_mask(token_len)
538
+ starts: list[int] = []
539
+ ends: list[int] = []
540
+ keys: list[str] = []
541
+ block_hashes: list[BlockHash] = []
542
+ group_indices: list[int] = []
543
+ for g_idx, db in enumerate(self.token_databases):
544
+ mask = store_masks[g_idx]
545
+ for chunk_idx, (start, end, key) in enumerate(
546
+ db.process_tokens(token_len, req_meta.block_hashes)
547
+ ):
548
+ if chunk_idx >= len(mask) or not mask[chunk_idx]:
549
+ continue
550
+ starts.append(start)
551
+ ends.append(end)
552
+ keys.append(key.to_string())
553
+ block_hashes.append(req_meta.block_hashes[chunk_idx])
554
+ group_indices.append(g_idx)
555
+
556
+ # Apply put_step striding for TP
557
+ sl = slice(self.tp_rank % self.put_step, None, self.put_step)
558
+ starts = starts[sl]
559
+ ends = ends[sl]
560
+ keys = keys[sl]
561
+ block_hashes = block_hashes[sl]
562
+ group_indices = group_indices[sl]
563
+
564
+ if not keys:
565
+ return
566
+
567
+ # Check which blocks already exist (dedup)
568
+ save_exists_start = time.perf_counter()
569
+ try:
570
+ exists_states = self.store.batch_is_exist(keys)
571
+ except Exception:
572
+ self._record_operation(
573
+ "save_exists",
574
+ save_exists_start,
575
+ len(keys),
576
+ status="error",
577
+ num_failed_keys=len(keys),
578
+ )
579
+ raise
580
self._record_operation(
581
"save_exists",
582
save_exists_start,
583
len(keys),
576
- status="error",
577
- num_failed_keys=len(keys),
584
)
579
- raise
580
- self._record_operation(
581
- "save_exists",
582
- save_exists_start,
583
- len(keys),
584
- )
585
- missing_indices = [i for i, exists in enumerate(exists_states) if exists != 1]
585
+ missing_indices = [
586
+ i for i, exists in enumerate(exists_states) if exists != 1
587
+ ]
588
587
- if not missing_indices:
588
- self.dec_stored_request(req_id)
589
- return
589
+ if not missing_indices:
590
+ return
591
591
- starts = [starts[i] for i in missing_indices]
592
- ends = [ends[i] for i in missing_indices]
593
- keys = [keys[i] for i in missing_indices]
594
- block_hashes = [block_hashes[i] for i in missing_indices]
595
- group_indices = [group_indices[i] for i in missing_indices]
592
+ starts = [starts[i] for i in missing_indices]
593
+ ends = [ends[i] for i in missing_indices]
594
+ keys = [keys[i] for i in missing_indices]
595
+ block_hashes = [block_hashes[i] for i in missing_indices]
596
+ group_indices = [group_indices[i] for i in missing_indices]
597
597
- logger.debug(
598
- "Storing KV cache for %d blocks (groups=%s) for request %s",
599
- len(keys),
600
- set(group_indices),
601
- req_id,
602
- )
598
+ logger.debug(
599
+ "Storing KV cache for %d blocks (groups=%s) for request %s",
600
+ len(keys),
601
+ set(group_indices),
602
+ req_id,
603
+ )
604
604
- addrs: list[list[int]] = []
605
- sizes: list[list[int]] = []
606
- stored_events: list[BlockStored] = []
607
- # parent_block_hash chains live within a group, not across.
608
- prev_key_per_group: dict[int, Any] = {}
609
- new_block_hashes = [maybe_convert_block_hash(bh) for bh in block_hashes]
605
+ addrs: list[list[int]] = []
606
+ sizes: list[list[int]] = []
607
+ stored_events: list[BlockStored] = []
608
+ # parent_block_hash chains live within a group, not across.
609
+ prev_key_per_group: dict[int, Any] = {}
610
+ new_block_hashes = [maybe_convert_block_hash(bh) for bh in block_hashes]
611
611
- for idx, (s, e, g_idx) in enumerate(
612
- zip(starts, ends, group_indices, strict=True)
613
- ):
614
- db = self.token_databases[g_idx]
615
- addr, size, _ = db.prepare_value(s, e, block_ids_per_group[g_idx])
616
- addrs.append(addr)
617
- sizes.append(size)
618
-
619
- if self.enable_kv_event:
620
- token_ids = (
621
- req_meta.token_ids[s:e] if req_meta.token_ids is not None else None
622
- )
623
- stored_event = BlockStored(
624
- block_hashes=[new_block_hashes[idx]],
625
- parent_block_hash=prev_key_per_group.get(g_idx),
626
- token_ids=token_ids,
627
- block_size=req_meta.original_block_size,
628
- lora_id=None,
629
- medium="cpu",
630
- lora_name=None,
631
- )
632
- stored_events.append(stored_event)
633
- prev_key_per_group[g_idx] = new_block_hashes[idx]
612
+ for idx, (s, e, g_idx) in enumerate(
613
+ zip(starts, ends, group_indices, strict=True)
614
+ ):
615
+ db = self.token_databases[g_idx]
616
+ addr, size, _ = db.prepare_value(s, e, block_ids_per_group[g_idx])
617
+ addrs.append(addr)
618
+ sizes.append(size)
619
+
620
+ if self.enable_kv_event:
621
+ token_ids = (
622
+ req_meta.token_ids[s:e]
623
+ if req_meta.token_ids is not None
624
+ else None
625
+ )
626
+ stored_event = BlockStored(
627
+ block_hashes=[new_block_hashes[idx]],
628
+ parent_block_hash=prev_key_per_group.get(g_idx),
629
+ token_ids=token_ids,
630
+ block_size=req_meta.original_block_size,
631
+ lora_id=None,
632
+ medium="cpu",
633
+ lora_name=None,
634
+ )
635
+ stored_events.append(stored_event)
636
+ prev_key_per_group[g_idx] = new_block_hashes[idx]
637
635
- if current_event is not None:
636
- current_event.synchronize()
638
+ if current_event is not None:
639
+ current_event.synchronize()
640
638
- batch_bytes = _sum_batch_bytes(sizes)
639
- put_start = time.perf_counter()
640
- try:
641
- res = self.store.batch_put_from_multi_buffers(
642
- keys,
643
- addrs,
644
- sizes,
645
- self.replicate_config,
646
- )
647
- failed = [i for i, v in enumerate(res) if v < 0]
648
- self._record_operation(
649
- "save_put",
650
- put_start,
651
- len(keys),
652
- num_bytes=batch_bytes,
653
- status="partial_failure" if failed else "ok",
654
- num_failed_keys=len(failed),
655
- )
656
- if failed:
657
- failed_codes = set(res[i] for i in failed)
658
- logger.warning(
659
- "batch_put failed: %d/%d keys failed "
660
- "(codes=%s, batch_bytes=%d, num_keys=%d), "
661
- "first_key=%s",
662
- len(failed),
663
- len(keys),
664
- failed_codes,
665
- batch_bytes,
641
+ batch_bytes = _sum_batch_bytes(sizes)
642
+ put_start = time.perf_counter()
643
+ try:
644
+ res = self.store.batch_put_from_multi_buffers(
645
+ keys,
646
+ addrs,
647
+ sizes,
648
+ self.replicate_config,
649
+ )
650
+ failed = [i for i, v in enumerate(res) if v < 0]
651
+ self._record_operation(
652
+ "save_put",
653
+ put_start,
654
len(keys),
667
- keys[0] if keys else "N/A",
655
+ num_bytes=batch_bytes,
656
+ status="partial_failure" if failed else "ok",
657
+ num_failed_keys=len(failed),
658
)
669
- if (
670
- MOONCAKE_NO_AVAILABLE_HANDLE in failed_codes
671
- and not self._mark_request_skipped_for_pressure(req_id)
672
- ):
659
+ if failed:
660
+ failed_codes = set(res[i] for i in failed)
661
logger.warning(
674
- "Detected Mooncake CPU/disk offloading pressure "
675
- "(NO_AVAILABLE_HANDLE); skipping future store "
676
- "batches for request %s until a later store "
677
- "batch succeeds",
678
- req_id,
662
+ "batch_put failed: %d/%d keys failed "
663
+ "(codes=%s, batch_bytes=%d, num_keys=%d), "
664
+ "first_key=%s",
665
+ len(failed),
666
+ len(keys),
667
+ failed_codes,
668
+ batch_bytes,
669
+ len(keys),
670
+ keys[0] if keys else "N/A",
671
)
680
- elif self._clear_store_pressure():
681
- logger.info(
682
- "Mooncake CPU/disk offloading pressure cleared after a "
683
- "successful store batch"
672
+ if (
673
+ MOONCAKE_NO_AVAILABLE_HANDLE in failed_codes
674
+ and not self._mark_request_skipped_for_pressure(req_id)
675
+ ):
676
+ logger.warning(
677
+ "Detected Mooncake CPU/disk offloading pressure "
678
+ "(NO_AVAILABLE_HANDLE); skipping future store "
679
+ "batches for request %s until a later store "
680
+ "batch succeeds",
681
+ req_id,
682
+ )
683
+ elif self._clear_store_pressure():
684
+ logger.info(
685
+ "Mooncake CPU/disk offloading pressure cleared after a "
686
+ "successful store batch"
687
+ )
688
+ except Exception as e:
689
+ self._record_operation(
690
+ "save_put",
691
+ put_start,
692
+ len(keys),
693
+ num_bytes=batch_bytes,
694
+ status="error",
695
+ num_failed_keys=len(keys),
696
)
685
- except Exception as e:
686
- self._record_operation(
687
- "save_put",
688
- put_start,
689
- len(keys),
690
- num_bytes=batch_bytes,
691
- status="error",
692
- num_failed_keys=len(keys),
693
- )
694
- logger.error("Failed to put key %s, error: %s", keys, e)
697
+ logger.error("Failed to put key %s, error: %s", keys, e)
698
696
- if self.enable_kv_event and stored_events:
697
- self.update_kv_event(stored_events)
698
-
699
- self.dec_stored_request(req_id)
700
- self.request_queue.task_done()
699
+ if self.enable_kv_event and stored_events:
700
+ self.update_kv_event(stored_events)
701
+ finally:
702
+ self.dec_stored_request(req_id)
703
+ self.request_queue.task_done()
704
705
706
class KVCacheStoreRecvingThread(KVTransferThread):