[BugFix][LoRA] Refresh punica metadata when LoRA slots are reassigned under an unchanged mapping (#47725)
Signed-off-by: AmeenP <ameenp360@gmail.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Ameen Patel committed
Jul 7, 2026 at 08:53 UTC
392d1b4d2eaced844026196b2eea3b284a0e2554
2 files changed
+56
-1
tests/lora/test_lora_manager.py
+49
@@ -651,6 +651,55 @@ def test_lru_lora_model_manager(default_vllm_config, dist_init, dummy_model, dev
651
assert manager.device == device
652
653
654
+@pytest.mark.parametrize("device", DEVICES)
655
+def test_set_adapter_mapping_refreshes_after_slot_reassignment(
656
+ default_vllm_config, dist_init, dummy_model, device
657
+):
658
+ # An out-of-band add_lora() can LRU-evict and reassign GPU slots while the
659
+ # running batch -- and therefore its LoRAMapping -- is unchanged. The
660
+ # punica metadata must still be re-derived, otherwise in-flight requests
661
+ # are routed to the evicted layout and decode with the wrong adapter.
662
+ model = dummy_model
663
+ model_lora1 = create_lora(1, model, ["dense1", "dense2", "lm_head"], device=device)
664
+ model_lora2 = create_lora(2, model, ["dense1", "dense2", "lm_head"], device=device)
665
+ model_lora3 = create_lora(3, model, ["dense1", "dense2", "lm_head"], device=device)
666
+ manager = LRUCacheLoRAModelManager(
667
+ model,
668
+ 2,
669
+ 2,
670
+ 2,
671
+ LoRAConfig(
672
+ max_lora_rank=8, max_cpu_loras=3, max_loras=2, lora_dtype=DEFAULT_DTYPE
673
+ ),
674
+ device=device,
675
+ vllm_config=default_vllm_config,
676
+ )
677
+ punica_wrapper = manager.punica_wrapper_mapping[DEFAULT_LANGUAGE_WRAPPER_KEY]
678
+
679
+ assert manager.add_adapter(model_lora1)
680
+ assert manager.activate_adapter(1)
681
+ assert manager.add_adapter(model_lora2)
682
+ assert manager.activate_adapter(2)
683
+ assert manager.lora_index_to_id == [1, 2]
684
+
685
+ # Two in-flight requests, one token each, on adapters 1 and 2.
686
+ manager.set_adapter_mapping(LoRAMapping((1, 2), (1, 2)))
687
+ assert punica_wrapper.token_lora_indices.tolist() == [0, 1]
688
+
689
+ # Out-of-band add_lora() with both slots held by the running batch:
690
+ # activating 3 evicts 1; re-activating the batch's adapters lands them
691
+ # in swapped slots while the batch itself is unchanged.
692
+ assert manager.add_adapter(model_lora3)
693
+ assert manager.activate_adapter(3)
694
+ assert manager.activate_adapter(1)
695
+ assert manager.activate_adapter(2)
696
+ assert manager.lora_index_to_id == [2, 1]
697
+
698
+ # Identical mapping, but the metadata must follow the new slot layout.
699
+ manager.set_adapter_mapping(LoRAMapping((1, 2), (1, 2)))
700
+ assert punica_wrapper.token_lora_indices.tolist() == [1, 0]
701
+
702
+
703
@pytest.mark.parametrize("device", DEVICES)
704
def test_lru_cache_worker_adapter_manager(dist_init, dummy_model, device, tmp_path):
705
lora_config = LoRAConfig(
vllm/lora/model_manager.py
+7
-1
@@ -117,6 +117,7 @@ class LoRAModelManager:
117
self.packed_modules: dict[str, list[str]] = {}
118
self.modules: dict[str, BaseLayerWithLoRA] = {}
119
self._last_mapping: LoRAMapping | None = None
120
+ self._last_slot_layout: tuple[int | None, ...] | None = None
121
is_moe = is_moe_model(self.model)
122
self._is_moe = is_moe
123
@@ -1147,9 +1148,14 @@ class LoRAModelManager:
1148
return True
1149
1150
def set_adapter_mapping(self, mapping: LoRAMapping) -> None:
1150
- if self._last_mapping != mapping:
1151
+ # The punica metadata derives from the slot layout as well as the
1152
+ # mapping: an out-of-band add_lora() can LRU-evict and reassign slots
1153
+ # while the running batch, and thus the mapping, is unchanged.
1154
+ slot_layout = tuple(self.lora_index_to_id)
1155
+ if self._last_mapping != mapping or self._last_slot_layout != slot_layout:
1156
self._set_adapter_mapping(mapping)
1157
self._last_mapping = mapping
1158
+ self._last_slot_layout = slot_layout
1159
1160
def remove_adapter(self, adapter_id: int) -> bool:
1161
self.deactivate_adapter(adapter_id)