From ccd49f6821ee110cc5a2b1aba620a8a1d66c7cbb Mon Sep 17 00:00:00 2001 From: Martin Hickey Date: Mon, 22 Jun 2026 15:57:09 +0100 Subject: [PATCH] [MyPy] Fix mypy for `vllm/lora` (#41722) Signed-off-by: Martin Hickey --- tests/lora/test_lora_manager.py | 29 ++++++++++++++----- tools/pre_commit/mypy.py | 2 -- vllm/lora/layers/base_linear.py | 13 +++++++-- vllm/lora/layers/column_parallel_linear.py | 17 +++++++++-- vllm/lora/layers/fused_moe.py | 29 ++++++++++++++++--- vllm/lora/layers/row_parallel_linear.py | 2 +- vllm/lora/layers/utils.py | 14 +++++++-- vllm/lora/lora_model.py | 7 +++-- vllm/lora/model_manager.py | 27 ++++++++++------- vllm/lora/peft_helper.py | 8 +++-- vllm/lora/utils.py | 21 ++++++++++---- vllm/lora/worker_manager.py | 14 +++++++-- .../layers/fused_moe/modular_kernel.py | 1 + vllm/model_executor/models/interfaces.py | 2 ++ vllm/models/deepseek_v4/xpu/mtp.py | 4 +-- 15 files changed, 140 insertions(+), 50 deletions(-) diff --git a/tests/lora/test_lora_manager.py b/tests/lora/test_lora_manager.py index 49436d66243..80a3b6dd9c6 100644 --- a/tests/lora/test_lora_manager.py +++ b/tests/lora/test_lora_manager.py @@ -124,6 +124,7 @@ def test_replace_submodules(default_vllm_config, dist_init, dummy_model): max_lora_rank=8, max_cpu_loras=8, max_loras=8, lora_dtype=DEFAULT_DTYPE ), torch.device(DEVICES[0]), + default_vllm_config, ) model = manager.model assert isinstance(model.get_submodule("dense1"), ColumnParallelLinearWithLoRA) @@ -152,6 +153,7 @@ def test_wrap_replicated_linear_subclasses(default_vllm_config, dist_init, dummy max_lora_rank=8, max_cpu_loras=8, max_loras=8, lora_dtype=DEFAULT_DTYPE ), torch.device(DEVICES[0]), + default_vllm_config, ) assert isinstance( @@ -172,6 +174,7 @@ def test_wrap_gate_linear(default_vllm_config, dist_init, dummy_model): max_lora_rank=8, max_cpu_loras=8, max_loras=8, lora_dtype=DEFAULT_DTYPE ), torch.device(DEVICES[0]), + default_vllm_config, ) assert isinstance( @@ -219,6 +222,7 @@ def test_dedup_shared_module_across_paths(default_vllm_config, dist_init, dummy_ max_lora_rank=8, max_cpu_loras=8, max_loras=8, lora_dtype=DEFAULT_DTYPE ), torch.device(DEVICES[0]), + default_vllm_config, ) canonical = manager.model.get_submodule("moe.gate") @@ -263,6 +267,7 @@ def test_lm_head_exempt_from_dedup(default_vllm_config, dist_init, dummy_model): max_lora_rank=8, max_cpu_loras=8, max_loras=8, lora_dtype=DEFAULT_DTYPE ), torch.device(DEVICES[0]), + default_vllm_config, ) # lm_head's special handling still ran: logits_processor got wrapped @@ -293,6 +298,7 @@ def test_skip_unsupported_matched_modules(default_vllm_config, dist_init, dummy_ max_lora_rank=8, max_cpu_loras=8, max_loras=8, lora_dtype=DEFAULT_DTYPE ), torch.device(DEVICES[0]), + default_vllm_config, ) # Should not crash and should keep unsupported matched modules unchanged. @@ -325,6 +331,7 @@ def test_target_modules_fail_closed_on_unsupported_matched_modules( target_modules=["dense1"], ), torch.device(DEVICES[0]), + default_vllm_config, ) @@ -374,6 +381,7 @@ def test_lora_model_manager(default_vllm_config, dist_init, dummy_model, device) max_lora_rank=8, max_cpu_loras=3, max_loras=2, lora_dtype=DEFAULT_DTYPE ), device=device, + vllm_config=default_vllm_config, ) assert all(x is None for x in manager.lora_index_to_id) assert manager.add_adapter(model_lora1) @@ -442,6 +450,7 @@ def test_lora_lru_cache_model_manager( max_lora_rank=8, max_cpu_loras=3, max_loras=2, lora_dtype=DEFAULT_DTYPE ), device=device, + vllm_config=default_vllm_config, ) assert all(x is None for x in manager.lora_index_to_id) assert manager.add_adapter(model_lora1) @@ -535,6 +544,7 @@ def test_lru_lora_model_manager(default_vllm_config, dist_init, dummy_model, dev max_lora_rank=8, max_cpu_loras=2, max_loras=2, lora_dtype=DEFAULT_DTYPE ), device=device, + vllm_config=default_vllm_config, ) assert all(x is None for x in manager.lora_index_to_id) @@ -642,9 +652,7 @@ def test_lru_lora_model_manager(default_vllm_config, dist_init, dummy_model, dev @pytest.mark.parametrize("device", DEVICES) -def test_lru_cache_worker_adapter_manager( - default_vllm_config, dist_init, dummy_model, device, tmp_path -): +def test_lru_cache_worker_adapter_manager(dist_init, dummy_model, device, tmp_path): lora_config = LoRAConfig( max_lora_rank=8, max_cpu_loras=4, max_loras=4, lora_dtype=DEFAULT_DTYPE ) @@ -670,7 +678,7 @@ def test_lru_cache_worker_adapter_manager( worker_adapter_manager.max_num_seqs = 4 worker_adapter_manager.max_num_batched_tokens = 2 - worker_adapter_manager.create_lora_manager(dummy_model) + worker_adapter_manager.create_lora_manager(dummy_model, vllm_config) mapping = LoRAMapping([], []) worker_adapter_manager.set_active_adapters( @@ -758,9 +766,7 @@ def test_lru_cache_worker_adapter_manager( @pytest.mark.parametrize("device", DEVICES) -def test_worker_adapter_manager( - default_vllm_config, dist_init, dummy_model_gate_up, device, tmp_path -): +def test_worker_adapter_manager(dist_init, dummy_model_gate_up, device, tmp_path): # Should remove every LoRA not specified in the request. lora_config = LoRAConfig( max_lora_rank=8, max_cpu_loras=4, max_loras=4, lora_dtype=DEFAULT_DTYPE @@ -774,7 +780,7 @@ def test_worker_adapter_manager( worker_adapter_manager = WorkerLoRAManager(vllm_config, device, EMBEDDING_MODULES) worker_adapter_manager.vocab_size = dummy_model_gate_up.unpadded_vocab_size - worker_adapter_manager.create_lora_manager(dummy_model_gate_up) + worker_adapter_manager.create_lora_manager(dummy_model_gate_up, vllm_config) dummy_lora_files = f"{tmp_path}/lora_adapter" os.makedirs(dummy_lora_files, exist_ok=True) @@ -894,6 +900,7 @@ def test_packed_loras(default_vllm_config, dist_init, dummy_model_gate_up, devic max_lora_rank=8, max_cpu_loras=2, max_loras=2, lora_dtype=DEFAULT_DTYPE ), device=device, + vllm_config=default_vllm_config, ) model = manager.model @@ -944,6 +951,7 @@ def _test_target_modules( device: str, expected_lora: list[tuple[str, type]], expected_no_lora: list[tuple[str, type]], + vllm_config, ): """Create a LoRAModelManager and assert which modules have LoRA applied.""" LoRAModelManager( @@ -959,6 +967,7 @@ def _test_target_modules( target_modules=target_modules, ), device=device, + vllm_config=vllm_config, ) for module_path, lora_cls in expected_lora: assert isinstance(model.get_submodule(module_path), lora_cls) @@ -981,6 +990,7 @@ def test_target_modules_config(default_vllm_config, dist_init, dummy_model, devi ("dense2", RowParallelLinearWithLoRA), ("layer1.dense2", RowParallelLinearWithLoRA), ], + vllm_config=default_vllm_config, ) @@ -998,6 +1008,7 @@ def test_target_modules_multiple(default_vllm_config, dist_init, dummy_model, de ("layer1.dense2", RowParallelLinearWithLoRA), ], expected_no_lora=[], + vllm_config=default_vllm_config, ) @@ -1017,6 +1028,7 @@ def test_target_modules_none_uses_all( ("layer1.dense2", RowParallelLinearWithLoRA), ], expected_no_lora=[], + vllm_config=default_vllm_config, ) @@ -1036,4 +1048,5 @@ def test_target_modules_match_packed_runtime_modules( ("layer1.dense1", ColumnParallelLinearWithLoRA), ("layer1.dense2", RowParallelLinearWithLoRA), ], + vllm_config=default_vllm_config, ) diff --git a/tools/pre_commit/mypy.py b/tools/pre_commit/mypy.py index a174208da4c..2b908e39036 100755 --- a/tools/pre_commit/mypy.py +++ b/tools/pre_commit/mypy.py @@ -25,8 +25,6 @@ import regex as re # from "skip" to "silent", remove its directory from SEPARATE_GROUPS. SEPARATE_GROUPS = [ "tests", - # v0 related - "vllm/lora", ] # TODO(woosuk): Include the code from Megatron and HuggingFace. diff --git a/vllm/lora/layers/base_linear.py b/vllm/lora/layers/base_linear.py index cb65cf69504..5c8e829b299 100644 --- a/vllm/lora/layers/base_linear.py +++ b/vllm/lora/layers/base_linear.py @@ -17,6 +17,7 @@ from vllm.forward_context import ( from vllm.model_executor.layers.linear import ( ColumnParallelLinear, LinearBase, + QuantizeMethodBase, ReplicatedLinear, RowParallelLinear, ) @@ -182,6 +183,14 @@ class BaseLinearLayerWithLoRA(BaseLayerWithLoRA): lora_b, non_blocking=True ) + def _get_quant_method(self) -> QuantizeMethodBase: + quant_method = self.base_layer.quant_method + if quant_method is None: + raise RuntimeError( + f"{type(self.base_layer).__name__} must define quant_method for LoRA." + ) + return quant_method + def apply(self, x: torch.Tensor, bias: torch.Tensor | None = None) -> torch.Tensor: # is_forward_context_available for tower modules if self._enable_aux_cuda_stream and is_forward_context_available(): @@ -195,7 +204,7 @@ class BaseLinearLayerWithLoRA(BaseLayerWithLoRA): def _apply_sync( self, x: torch.Tensor, bias: torch.Tensor | None = None ) -> torch.Tensor: - output = self.base_layer.quant_method.apply(self.base_layer, x, bias) + output = self._get_quant_method().apply(self.base_layer, x, bias) return self._apply_lora_to_output(x, output) def _apply_base_forward(self, x: torch.Tensor) -> torch.Tensor: @@ -242,7 +251,7 @@ class BaseLinearLayerWithLoRA(BaseLayerWithLoRA): output_size = sum(self.output_slices) def base_fn() -> torch.Tensor: - return self.base_layer.quant_method.apply(self.base_layer, x, bias) + return self._get_quant_method().apply(self.base_layer, x, bias) def lora_fn() -> torch.Tensor: # Must be zeros, not empty: _lora_expand_kernel exits early (without diff --git a/vllm/lora/layers/column_parallel_linear.py b/vllm/lora/layers/column_parallel_linear.py index 8a86191b891..4df468a2753 100644 --- a/vllm/lora/layers/column_parallel_linear.py +++ b/vllm/lora/layers/column_parallel_linear.py @@ -33,7 +33,7 @@ def _mcp_apply(x, bias, layer: "ColumnParallelLinearWithLoRA"): == len(layer.output_slices) ) - output = layer.base_layer.quant_method.apply(layer.base_layer, x, bias) + output = layer._get_quant_method().apply(layer.base_layer, x, bias) x = x.view(-1, x.shape[-1]) output, out_orig_shape = output.view(-1, output.shape[-1]), output.shape @@ -73,6 +73,8 @@ def _mcp_apply(x, bias, layer: "ColumnParallelLinearWithLoRA"): ) if not current_platform.can_update_inplace(): + if lora_output is None: + raise RuntimeError("LoRA expand must return an output tensor.") output = lora_output output = output.view(*out_orig_shape) @@ -327,12 +329,16 @@ class MergedColumnParallelLinearWithLoRA(ColumnParallelLinearWithLoRA): def apply(self, x: torch.Tensor, bias: torch.Tensor | None = None) -> torch.Tensor: merged_cls = maybe_get_oot_by_class(MergedColumnParallelLinear) + base_forward = getattr(type(self.base_layer), "forward", None) + merged_forward = getattr(merged_cls, "forward", None) # Effectively unsharded subclasses can safely reuse their custom # forward() implementation before applying the LoRA delta. if ( self.tp_size == 1 and type(self.base_layer) is not merged_cls - and type(self.base_layer).forward is not merged_cls.forward + and base_forward is not None + and merged_forward is not None + and base_forward is not merged_forward ): return self._apply_base_forward(x) return _mcp_apply(x, bias, self) @@ -482,6 +488,7 @@ class MergedQKVParallelLinearWithLoRA(MergedColumnParallelLinearWithLoRA): lora_config: LoRAConfig, packed_modules_list: list, model_config: PretrainedConfig | None = None, + decorate: bool = True, ) -> bool: return ( type(source_layer) is maybe_get_oot_by_class(QKVParallelLinear) @@ -523,6 +530,7 @@ class ColumnParallelLinearWithShardedLoRA(ColumnParallelLinearWithLoRA): lora_config: LoRAConfig, packed_modules_list: list, model_config: PretrainedConfig | None = None, + decorate: bool = True, ) -> bool: # specifying kwargs so they can be easily accessed in decorator return super().can_replace_layer( @@ -565,6 +573,7 @@ class MergedColumnParallelLinearWithShardedLoRA(MergedColumnParallelLinearWithLo lora_config: LoRAConfig, packed_modules_list: list, model_config: PretrainedConfig | None = None, + decorate: bool = True, ) -> bool: # specifying kwargs so they can be easily accessed in decorator return super().can_replace_layer( @@ -650,6 +659,7 @@ class MergedQKVParallelLinearWithShardedLoRA(MergedQKVParallelLinearWithLoRA): lora_config: LoRAConfig, packed_modules_list: list, model_config: PretrainedConfig | None = None, + decorate: bool = True, ) -> bool: # specifying kwargs so they can be easily accessed in decorator return super().can_replace_layer( @@ -678,6 +688,7 @@ class MergedColumnParallelLinearVariableSliceWithLoRA( lora_config: LoRAConfig, packed_modules_list: list, model_config: PretrainedConfig | None = None, + decorate: bool = True, ) -> bool: # Support MergedColumnParallelLinear with 3 or more slices # (2 slices are handled by MergedColumnParallelLinearWithLoRA) @@ -727,7 +738,7 @@ class MergedColumnParallelLinearVariableSliceWithLoRA( start_idx = 0 for output_size in output_sizes: end_idx = start_idx + output_size - lora_b_list.append(lora_b[start_idx:end_idx, :]) + lora_b_list.append(lora_b[start_idx:end_idx]) start_idx = end_idx lora_b = lora_b_list diff --git a/vllm/lora/layers/fused_moe.py b/vllm/lora/layers/fused_moe.py index 7b400bc5e97..c3763f5448c 100644 --- a/vllm/lora/layers/fused_moe.py +++ b/vllm/lora/layers/fused_moe.py @@ -12,10 +12,16 @@ from vllm.lora.layers.base import BaseLayerWithLoRA from vllm.model_executor.custom_op import maybe_get_oot_by_class from vllm.model_executor.layers.fused_moe import MoERunner from vllm.model_executor.layers.fused_moe.experts.lora_context import MoELoRAContext +from vllm.model_executor.layers.fused_moe.experts.lora_experts_mixin import ( + LoRAExpertsMixin, +) from vllm.model_executor.layers.fused_moe.fused_moe_modular_method import ( FusedMoEModularMethod, ) -from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel +from vllm.model_executor.layers.fused_moe.modular_kernel import ( + FusedMoEKernel, + FusedMoEKernelModularImpl, +) from vllm.model_executor.layers.fused_moe.prepare_finalize import ( MoEPrepareAndFinalizeNoDPEPModular, ) @@ -58,6 +64,13 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA): routed_experts._ensure_moe_quant_config_init() if getattr(routed_experts.quant_method, "supports_internal_mk", False): moe_kernel = routed_experts.quant_method.moe_kernel + assert moe_kernel is not None, ( + "Fused MoE quant method must provide a moe_kernel." + ) + # Don't let the kernel own shared experts so the runner can + # overlap them with routed experts via a separate CUDA stream. + assert isinstance(moe_kernel.impl, FusedMoEKernelModularImpl) + moe_kernel.impl.shared_experts = None else: prepare_finalize = MoEPrepareAndFinalizeNoDPEPModular() moe_kernel = FusedMoEKernel( @@ -405,7 +418,11 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA): def set_mapping(self, punica_wrapper): super().set_mapping(punica_wrapper) lora_context = self._build_lora_context() - self._moe_kernel.fused_experts.set_lora_context(lora_context) + fused_experts = self._moe_kernel.fused_experts + assert isinstance(fused_experts, LoRAExpertsMixin), ( + f"{type(fused_experts).__name__} does not support LoRA context setup." + ) + fused_experts.set_lora_context(lora_context) prepare_finalize = self._moe_kernel.prepare_finalize if hasattr(prepare_finalize, "set_lora_context"): prepare_finalize.set_lora_context(lora_context) @@ -482,9 +499,13 @@ class FusedMoE3DWithLoRA(FusedMoEWithLoRA): ) -> None: """Initializes lora matrices.""" - assert isinstance(model_config, PretrainedConfig) + if model_config is None: + raise ValueError("model_config must be provided for MoE LoRA.") + architectures = model_config.architectures + if not architectures: + raise ValueError("model_config.architectures must be defined for MoE LoRA.") self._verify_ep_fs(lora_config) - self._base_model = model_config.architectures[0] + self._base_model = architectures[0] self.max_loras = lora_config.max_loras self.fully_sharded = lora_config.fully_sharded_loras diff --git a/vllm/lora/layers/row_parallel_linear.py b/vllm/lora/layers/row_parallel_linear.py index 9460b687f1a..3d3c7f1d279 100644 --- a/vllm/lora/layers/row_parallel_linear.py +++ b/vllm/lora/layers/row_parallel_linear.py @@ -116,7 +116,7 @@ class RowParallelLinearWithShardedLoRA(RowParallelLinearWithLoRA): return lora_b def apply(self, x: torch.Tensor, bias: torch.Tensor | None = None) -> torch.Tensor: - output = self.base_layer.quant_method.apply(self.base_layer, x, bias) + output = self._get_quant_method().apply(self.base_layer, x, bias) x = x.view(-1, x.shape[-1]) output, out_orig_shape = output.view(-1, output.shape[-1]), output.shape diff --git a/vllm/lora/layers/utils.py b/vllm/lora/layers/utils.py index 3662a83acc8..17c21d36f5b 100644 --- a/vllm/lora/layers/utils.py +++ b/vllm/lora/layers/utils.py @@ -111,19 +111,27 @@ def try_get_optimal_moe_lora_config( # base MoE weight's block-wise quantization, so block_shape is omitted # from the config lookup — the non-quantized branch in get_default_config # ignores it anyway. - config = try_get_optimal_moe_config(w1_shape, w2_shape, top_k, dtype, M).copy() + raw_config = try_get_optimal_moe_config(w1_shape, w2_shape, top_k, dtype, M) + config: dict[str, int | None] = dict(raw_config) if op_type in [ "fused_moe_lora_w13_shrink", "fused_moe_lora_w2_shrink", ]: + block_size_n = config.get("BLOCK_SIZE_N") config["BLOCK_SIZE_N"] = min( - config.get("BLOCK_SIZE_N", 64), next_power_of_2(rank) + block_size_n if block_size_n is not None else 64, + next_power_of_2(rank), ) elif op_type in [ "fused_moe_lora_w13_expand", "fused_moe_lora_w2_expand", ]: + block_size_k = config.get("BLOCK_SIZE_K") config["BLOCK_SIZE_K"] = max( - 16, min(config.get("BLOCK_SIZE_K", 32), next_power_of_2(rank)) + 16, + min( + block_size_k if block_size_k is not None else 32, + next_power_of_2(rank), + ), ) return config diff --git a/vllm/lora/lora_model.py b/vllm/lora/lora_model.py index 859ed02f871..01d04963283 100644 --- a/vllm/lora/lora_model.py +++ b/vllm/lora/lora_model.py @@ -245,9 +245,10 @@ class LoRAModel: from tensorizer import TensorDeserializer tensorizer_config = TensorizerConfig(**tensorizer_config_dict) - lora_tensor_path = os.path.join( - tensorizer_config.tensorizer_dir, "adapter_model.tensors" - ) + tensorizer_dir = tensorizer_config.tensorizer_dir + if tensorizer_dir is None: + raise ValueError("tensorizer_dir must be set in tensorizer config.") + lora_tensor_path = os.path.join(tensorizer_dir, "adapter_model.tensors") tensorizer_args = tensorizer_config._construct_tensorizer_args() tensors = TensorDeserializer( lora_tensor_path, diff --git a/vllm/lora/model_manager.py b/vllm/lora/model_manager.py index 39c3bb0ea16..a24a75b8172 100644 --- a/vllm/lora/model_manager.py +++ b/vllm/lora/model_manager.py @@ -34,6 +34,7 @@ from vllm.lora.utils import ( from vllm.model_executor.layers.fused_moe import MoERunner from vllm.model_executor.models import ( SupportsLoRA, + SupportsMultiModal, is_pooling_model, supports_multimodal, ) @@ -50,6 +51,12 @@ T = TypeVar("T") DEFAULT_LANGUAGE_WRAPPER_KEY = "language_model" +class SupportsLoRAModel(nn.Module, SupportsLoRA): ... + + +class SupportsLoRAMultiModalModel(SupportsLoRAModel, SupportsMultiModal): ... + + class AdapterLRUCache(LRUCache[int, T]): def __init__(self, capacity: int, deactivate_fn: Callable[[int], object]): super().__init__(capacity) @@ -66,13 +73,13 @@ class LoRAModelManager: def __init__( self, - model: SupportsLoRA, + model: SupportsLoRAModel, max_num_seqs: int, max_num_batched_tokens: int, vocab_size: int, lora_config: LoRAConfig, device: torch.device, - vllm_config: VllmConfig | None = None, + vllm_config: VllmConfig, ): """Create a LoRAModelManager and adapter for a given model. @@ -85,7 +92,7 @@ class LoRAModelManager: vocab_size: the vocab size of the model. lora_config: the LoRA configuration. """ - self.model: SupportsLoRA = model + self.model: SupportsLoRAModel = model self.supported_lora_modules = get_supported_lora_modules(self.model) assert self.supported_lora_modules, ( f"No supported LoRA modules found in {self.model.__class__.__name__}." @@ -106,7 +113,6 @@ class LoRAModelManager: self.is_pooling_model = is_pooling_model(self.model) self.packed_modules: dict[str, list[str]] = {} self.modules: dict[str, BaseLayerWithLoRA] = {} - # Dict instead of a set for compatibility with LRUCache. self._last_mapping: LoRAMapping | None = None is_moe = is_moe_model(self.model) self._is_moe = is_moe @@ -272,6 +278,7 @@ class LoRAModelManager: @property def capacity(self) -> int: + assert self.lora_config.max_cpu_loras is not None return self.lora_config.max_cpu_loras @property @@ -1156,7 +1163,7 @@ class LoRAModelManager: class LoRALRUCache(AdapterLRUCache[LoRAModel]): - def __init__(self, capacity: int, deactivate_lora_fn: Callable[[int], bool]): + def __init__(self, capacity: int, deactivate_lora_fn: Callable[[int], object]): super().__init__(capacity, deactivate_lora_fn) @@ -1165,13 +1172,13 @@ class LRUCacheLoRAModelManager(LoRAModelManager): def __init__( self, - model: nn.Module, + model: SupportsLoRAModel, max_num_seqs: int, max_num_batched_tokens: int, vocab_size: int, lora_config: LoRAConfig, device: torch.device, - vllm_config: VllmConfig | None = None, + vllm_config: VllmConfig, ): super().__init__( model, @@ -1182,10 +1189,10 @@ class LRUCacheLoRAModelManager(LoRAModelManager): device, vllm_config, ) - self._registered_adapters: LoRALRUCache = LoRALRUCache( + self._registered_adapters: LoRALRUCache = LoRALRUCache( # type: ignore[assignment] self.capacity, self.deactivate_adapter ) - self._active_adapters: LoRALRUCache = LoRALRUCache( + self._active_adapters: LoRALRUCache = LoRALRUCache( # type: ignore[assignment] self.lora_slots, self._deactivate_adapter ) @@ -1248,7 +1255,7 @@ class LRUCacheLoRAModelManager(LoRAModelManager): def create_lora_manager( - model: nn.Module, + model: SupportsLoRAModel, max_num_seqs: int, max_num_batched_tokens: int, vocab_size: int, diff --git a/vllm/lora/peft_helper.py b/vllm/lora/peft_helper.py index 975c3d8fc0a..1443efd4f0c 100644 --- a/vllm/lora/peft_helper.py +++ b/vllm/lora/peft_helper.py @@ -91,9 +91,11 @@ class PEFTHelper: tensorizer_args = tensorizer_config._construct_tensorizer_args() from tensorizer.stream_io import open_stream - lora_config_path = os.path.join( - tensorizer_config.tensorizer_dir, "adapter_config.json" - ) + tensorizer_dir = tensorizer_config.tensorizer_dir + if tensorizer_dir is None: + raise ValueError("tensorizer_dir must be set in tensorizer config.") + + lora_config_path = os.path.join(tensorizer_dir, "adapter_config.json") with open_stream( lora_config_path, mode="rb", **tensorizer_args.stream_kwargs ) as f: diff --git a/vllm/lora/utils.py b/vllm/lora/utils.py index 828aea712d0..a628d70cbad 100644 --- a/vllm/lora/utils.py +++ b/vllm/lora/utils.py @@ -173,11 +173,18 @@ def parse_fine_tuned_lora_name( # mapping correctly. if name.startswith("base_model.model."): name = name.replace("base_model.model.", "") - name = weights_mapper._map_name(name) if weights_mapper else name - # recover the prefix `base_model.model.` - name = "base_model.model." + name + if weights_mapper: + mapped_name = weights_mapper._map_name(name) + if mapped_name is None: + raise ValueError("Mapped LoRA weight name cannot be None.") + # recover the prefix `base_model.model.` + name = "base_model.model." + mapped_name else: - name = weights_mapper._map_name(name) if weights_mapper else name + if weights_mapper: + mapped_name = weights_mapper._map_name(name) + if mapped_name is None: + raise ValueError("Mapped LoRA weight name cannot be None.") + name = mapped_name # In some situations, we may not start with `base_model.model.`. # If we don't (e.g., ibm-granite/granite-speech-3.3-8b), @@ -185,7 +192,11 @@ def parse_fine_tuned_lora_name( start_index = 2 if name.startswith("base_model.model.") else 0 parts = name.split(".") - if parts[-1] == "weight" and (parts[-2] == "lora_A" or parts[-2] == "lora_B"): + if ( + parts[-1] == "weight" + and len(parts) >= 2 + and (parts[-2] == "lora_A" or parts[-2] == "lora_B") + ): new_name = ".".join(parts[start_index:-2]) return new_name, parts[-2] == "lora_A" diff --git a/vllm/lora/worker_manager.py b/vllm/lora/worker_manager.py index 166d5c36ba5..c1aee79bec2 100644 --- a/vllm/lora/worker_manager.py +++ b/vllm/lora/worker_manager.py @@ -7,6 +7,7 @@ from typing import Any, Literal import torch from vllm.config import VllmConfig +from vllm.config.lora import LoRAConfig from vllm.exceptions import LoRAAdapterNotFoundError from vllm.logger import init_logger from vllm.lora.lora_model import LoRAModel @@ -45,7 +46,10 @@ class WorkerLoRAManager: vllm_config.scheduler_config.max_num_batched_tokens ) self.vocab_size = vllm_config.model_config.get_vocab_size() - self.lora_config = vllm_config.lora_config + lora_config = vllm_config.lora_config + if lora_config is None: + raise ValueError("LoRA config must be set for WorkerLoRAManager.") + self.lora_config: LoRAConfig = lora_config # Use get_text_config() in case of multimodal models text_config = vllm_config.model_config.hf_config.get_text_config() @@ -81,8 +85,10 @@ class WorkerLoRAManager: def create_lora_manager( self, model: torch.nn.Module, - vllm_config: VllmConfig | None = None, + vllm_config: VllmConfig, ) -> Any: + if vllm_config is None: + raise ValueError("vllm_config must be provided to create a LoRA manager.") lora_manager = create_lora_manager( model, max_num_seqs=self.max_num_seqs, @@ -240,8 +246,10 @@ class LRUCacheWorkerLoRAManager(WorkerLoRAManager): def create_lora_manager( self, model: torch.nn.Module, - vllm_config: VllmConfig | None = None, + vllm_config: VllmConfig, ) -> Any: + if vllm_config is None: + raise ValueError("vllm_config must be provided to create a LoRA manager.") lora_manager = create_lora_manager( model, lora_manager_cls=self._manager_cls, diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index 0e55e827c20..9f3ac1fd79d 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -1029,6 +1029,7 @@ class FusedMoEKernelModularImpl: ): self.prepare_finalize = prepare_finalize self.fused_experts = fused_experts + self.shared_experts: SharedExperts | None = None moe_parallel_config = fused_experts.moe_config.moe_parallel_config self.moe_parallel_config = moe_parallel_config self.is_dp_ep = ( diff --git a/vllm/model_executor/models/interfaces.py b/vllm/model_executor/models/interfaces.py index f65d7d0b44e..66d1fc6a4e9 100644 --- a/vllm/model_executor/models/interfaces.py +++ b/vllm/model_executor/models/interfaces.py @@ -42,6 +42,7 @@ from .interfaces_base import VllmModel if TYPE_CHECKING: from vllm.config import VllmConfig + from vllm.lora.model_manager import LoRAModelManager from vllm.model_executor.models.utils import WeightsMapper from vllm.multimodal.inputs import MultiModalFeatureSpec from vllm.multimodal.registry import _ProcessorFactories @@ -554,6 +555,7 @@ class SupportsLoRA(Protocol): packed_modules_mapping: dict[str, list[str]] = {} # Module prefixes to skip during LoRA loading (e.g., ["mtp."] for MTP layers) lora_skip_prefixes: ClassVar[list[str]] = [] + lora_manager: "LoRAModelManager | None" # We can't use runtime_checkable with ClassVar for issubclass checks diff --git a/vllm/models/deepseek_v4/xpu/mtp.py b/vllm/models/deepseek_v4/xpu/mtp.py index d4a8d293baf..8baca78b8ba 100644 --- a/vllm/models/deepseek_v4/xpu/mtp.py +++ b/vllm/models/deepseek_v4/xpu/mtp.py @@ -24,9 +24,7 @@ from vllm.distributed import ( get_tensor_model_parallel_world_size, ) from vllm.logger import init_logger -from vllm.model_executor.layers.fused_moe import ( - fused_moe_make_expert_params_mapping, -) +from vllm.model_executor.layers.fused_moe import fused_moe_make_expert_params_mapping from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ReplicatedLinear from vllm.model_executor.layers.logits_processor import LogitsProcessor