@tduncan / Nvidia-Optimization-Stack-VLLM / commits / ccd49f6821

[MyPy] Fix mypy for `vllm/lora` (#41722)

Signed-off-by: Martin Hickey <martin.hickey@ie.ibm.com>

Martin Hickey committed Jun 22, 2026 at 15:57 UTC ccd49f6821ee110cc5a2b1aba620a8a1d66c7cbb

15 files changed +140 -50

tests/lora/test_lora_manager.py

+21 -8

@@ -124,6 +124,7 @@ def test_replace_submodules(default_vllm_config, dist_init, dummy_model):

124 max_lora_rank=8, max_cpu_loras=8, max_loras=8, lora_dtype=DEFAULT_DTYPE

125 ),

126 torch.device(DEVICES[0]),

127 + default_vllm_config,

128 )

129 model = manager.model

130 assert isinstance(model.get_submodule("dense1"), ColumnParallelLinearWithLoRA)

@@ -152,6 +153,7 @@ def test_wrap_replicated_linear_subclasses(default_vllm_config, dist_init, dummy

153 max_lora_rank=8, max_cpu_loras=8, max_loras=8, lora_dtype=DEFAULT_DTYPE

154 ),

155 torch.device(DEVICES[0]),

156 + default_vllm_config,

157 )

158

159 assert isinstance(

@@ -172,6 +174,7 @@ def test_wrap_gate_linear(default_vllm_config, dist_init, dummy_model):

174 max_lora_rank=8, max_cpu_loras=8, max_loras=8, lora_dtype=DEFAULT_DTYPE

175 ),

176 torch.device(DEVICES[0]),

177 + default_vllm_config,

178 )

179

180 assert isinstance(

@@ -219,6 +222,7 @@ def test_dedup_shared_module_across_paths(default_vllm_config, dist_init, dummy_

222 max_lora_rank=8, max_cpu_loras=8, max_loras=8, lora_dtype=DEFAULT_DTYPE

223 ),

224 torch.device(DEVICES[0]),

225 + default_vllm_config,

226 )

227

228 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):

267 max_lora_rank=8, max_cpu_loras=8, max_loras=8, lora_dtype=DEFAULT_DTYPE

268 ),

269 torch.device(DEVICES[0]),

270 + default_vllm_config,

271 )

272

273 # 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_

298 max_lora_rank=8, max_cpu_loras=8, max_loras=8, lora_dtype=DEFAULT_DTYPE

299 ),

300 torch.device(DEVICES[0]),

301 + default_vllm_config,

302 )

303

304 # Should not crash and should keep unsupported matched modules unchanged.

@@ -325,6 +331,7 @@ def test_target_modules_fail_closed_on_unsupported_matched_modules(

331 target_modules=["dense1"],

332 ),

333 torch.device(DEVICES[0]),

334 + default_vllm_config,

335 )

336

337

@@ -374,6 +381,7 @@ def test_lora_model_manager(default_vllm_config, dist_init, dummy_model, device)

381 max_lora_rank=8, max_cpu_loras=3, max_loras=2, lora_dtype=DEFAULT_DTYPE

382 ),

383 device=device,

384 + vllm_config=default_vllm_config,

385 )

386 assert all(x is None for x in manager.lora_index_to_id)

387 assert manager.add_adapter(model_lora1)

@@ -442,6 +450,7 @@ def test_lora_lru_cache_model_manager(

450 max_lora_rank=8, max_cpu_loras=3, max_loras=2, lora_dtype=DEFAULT_DTYPE

451 ),

452 device=device,

453 + vllm_config=default_vllm_config,

454 )

455 assert all(x is None for x in manager.lora_index_to_id)

456 assert manager.add_adapter(model_lora1)

@@ -535,6 +544,7 @@ def test_lru_lora_model_manager(default_vllm_config, dist_init, dummy_model, dev

544 max_lora_rank=8, max_cpu_loras=2, max_loras=2, lora_dtype=DEFAULT_DTYPE

545 ),

546 device=device,

547 + vllm_config=default_vllm_config,

548 )

549 assert all(x is None for x in manager.lora_index_to_id)

550

@@ -642,9 +652,7 @@ def test_lru_lora_model_manager(default_vllm_config, dist_init, dummy_model, dev

652

653

654 @pytest.mark.parametrize("device", DEVICES)

645 -def test_lru_cache_worker_adapter_manager(

646 - default_vllm_config, dist_init, dummy_model, device, tmp_path

647 -):

655 +def test_lru_cache_worker_adapter_manager(dist_init, dummy_model, device, tmp_path):

656 lora_config = LoRAConfig(

657 max_lora_rank=8, max_cpu_loras=4, max_loras=4, lora_dtype=DEFAULT_DTYPE

658 )

@@ -670,7 +678,7 @@ def test_lru_cache_worker_adapter_manager(

678 worker_adapter_manager.max_num_seqs = 4

679 worker_adapter_manager.max_num_batched_tokens = 2

680

673 - worker_adapter_manager.create_lora_manager(dummy_model)

681 + worker_adapter_manager.create_lora_manager(dummy_model, vllm_config)

682

683 mapping = LoRAMapping([], [])

684 worker_adapter_manager.set_active_adapters(

@@ -758,9 +766,7 @@ def test_lru_cache_worker_adapter_manager(

766

767

768 @pytest.mark.parametrize("device", DEVICES)

761 -def test_worker_adapter_manager(

762 - default_vllm_config, dist_init, dummy_model_gate_up, device, tmp_path

763 -):

769 +def test_worker_adapter_manager(dist_init, dummy_model_gate_up, device, tmp_path):

770 # Should remove every LoRA not specified in the request.

771 lora_config = LoRAConfig(

772 max_lora_rank=8, max_cpu_loras=4, max_loras=4, lora_dtype=DEFAULT_DTYPE

@@ -774,7 +780,7 @@ def test_worker_adapter_manager(

780

781 worker_adapter_manager = WorkerLoRAManager(vllm_config, device, EMBEDDING_MODULES)

782 worker_adapter_manager.vocab_size = dummy_model_gate_up.unpadded_vocab_size

777 - worker_adapter_manager.create_lora_manager(dummy_model_gate_up)

783 + worker_adapter_manager.create_lora_manager(dummy_model_gate_up, vllm_config)

784

785 dummy_lora_files = f"{tmp_path}/lora_adapter"

786 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

900 max_lora_rank=8, max_cpu_loras=2, max_loras=2, lora_dtype=DEFAULT_DTYPE

901 ),

902 device=device,

903 + vllm_config=default_vllm_config,

904 )

905 model = manager.model

906

@@ -944,6 +951,7 @@ def _test_target_modules(

951 device: str,

952 expected_lora: list[tuple[str, type]],

953 expected_no_lora: list[tuple[str, type]],

954 + vllm_config,

955 ):

956 """Create a LoRAModelManager and assert which modules have LoRA applied."""

957 LoRAModelManager(

@@ -959,6 +967,7 @@ def _test_target_modules(

967 target_modules=target_modules,

968 ),

969 device=device,

970 + vllm_config=vllm_config,

971 )

972 for module_path, lora_cls in expected_lora:

973 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

990 ("dense2", RowParallelLinearWithLoRA),

991 ("layer1.dense2", RowParallelLinearWithLoRA),

992 ],

993 + vllm_config=default_vllm_config,

994 )

995

996

@@ -998,6 +1008,7 @@ def test_target_modules_multiple(default_vllm_config, dist_init, dummy_model, de

1008 ("layer1.dense2", RowParallelLinearWithLoRA),

1009 ],

1010 expected_no_lora=[],

1011 + vllm_config=default_vllm_config,

1012 )

1013

1014

@@ -1017,6 +1028,7 @@ def test_target_modules_none_uses_all(

1028 ("layer1.dense2", RowParallelLinearWithLoRA),

1029 ],

1030 expected_no_lora=[],

1031 + vllm_config=default_vllm_config,

1032 )

1033

1034

@@ -1036,4 +1048,5 @@ def test_target_modules_match_packed_runtime_modules(

1048 ("layer1.dense1", ColumnParallelLinearWithLoRA),

1049 ("layer1.dense2", RowParallelLinearWithLoRA),

1050 ],

1051 + vllm_config=default_vllm_config,

1052 )

tools/pre_commit/mypy.py

-2

@@ -25,8 +25,6 @@ import regex as re

25 # from "skip" to "silent", remove its directory from SEPARATE_GROUPS.

26 SEPARATE_GROUPS = [

27 "tests",

28 - # v0 related

29 - "vllm/lora",

28 ]

30 # TODO(woosuk): Include the code from Megatron and HuggingFace.

vllm/lora/layers/base_linear.py

+11 -2

@@ -17,6 +17,7 @@ from vllm.forward_context import (

17 from vllm.model_executor.layers.linear import (

18 ColumnParallelLinear,

19 LinearBase,

20 + QuantizeMethodBase,

21 ReplicatedLinear,

22 RowParallelLinear,

23 )

@@ -182,6 +183,14 @@ class BaseLinearLayerWithLoRA(BaseLayerWithLoRA):

183 lora_b, non_blocking=True

184 )

185

186 + def _get_quant_method(self) -> QuantizeMethodBase:

187 + quant_method = self.base_layer.quant_method

188 + if quant_method is None:

189 + raise RuntimeError(

190 + f"{type(self.base_layer).__name__} must define quant_method for LoRA."

191 + )

192 + return quant_method

193 +

194 def apply(self, x: torch.Tensor, bias: torch.Tensor | None = None) -> torch.Tensor:

195 # is_forward_context_available for tower modules

196 if self._enable_aux_cuda_stream and is_forward_context_available():

@@ -195,7 +204,7 @@ class BaseLinearLayerWithLoRA(BaseLayerWithLoRA):

204 def _apply_sync(

205 self, x: torch.Tensor, bias: torch.Tensor | None = None

206 ) -> torch.Tensor:

198 - output = self.base_layer.quant_method.apply(self.base_layer, x, bias)

207 + output = self._get_quant_method().apply(self.base_layer, x, bias)

208 return self._apply_lora_to_output(x, output)

209

210 def _apply_base_forward(self, x: torch.Tensor) -> torch.Tensor:

@@ -242,7 +251,7 @@ class BaseLinearLayerWithLoRA(BaseLayerWithLoRA):

251 output_size = sum(self.output_slices)

252

253 def base_fn() -> torch.Tensor:

245 - return self.base_layer.quant_method.apply(self.base_layer, x, bias)

254 + return self._get_quant_method().apply(self.base_layer, x, bias)

255

256 def lora_fn() -> torch.Tensor:

257 # Must be zeros, not empty: _lora_expand_kernel exits early (without

vllm/lora/layers/column_parallel_linear.py

+14 -3

@@ -33,7 +33,7 @@ def _mcp_apply(x, bias, layer: "ColumnParallelLinearWithLoRA"):

33 == len(layer.output_slices)

34 )

36 - output = layer.base_layer.quant_method.apply(layer.base_layer, x, bias)

36 + output = layer._get_quant_method().apply(layer.base_layer, x, bias)

38 x = x.view(-1, x.shape[-1])

39 output, out_orig_shape = output.view(-1, output.shape[-1]), output.shape

@@ -73,6 +73,8 @@ def _mcp_apply(x, bias, layer: "ColumnParallelLinearWithLoRA"):

73 )

75 if not current_platform.can_update_inplace():

76 + if lora_output is None:

77 + raise RuntimeError("LoRA expand must return an output tensor.")

78 output = lora_output

80 output = output.view(*out_orig_shape)

@@ -327,12 +329,16 @@ class MergedColumnParallelLinearWithLoRA(ColumnParallelLinearWithLoRA):

329

330 def apply(self, x: torch.Tensor, bias: torch.Tensor | None = None) -> torch.Tensor:

331 merged_cls = maybe_get_oot_by_class(MergedColumnParallelLinear)

332 + base_forward = getattr(type(self.base_layer), "forward", None)

333 + merged_forward = getattr(merged_cls, "forward", None)

334 # Effectively unsharded subclasses can safely reuse their custom

335 # forward() implementation before applying the LoRA delta.

336 if (

337 self.tp_size == 1

338 and type(self.base_layer) is not merged_cls

335 - and type(self.base_layer).forward is not merged_cls.forward

339 + and base_forward is not None

340 + and merged_forward is not None

341 + and base_forward is not merged_forward

342 ):

343 return self._apply_base_forward(x)

344 return _mcp_apply(x, bias, self)

@@ -482,6 +488,7 @@ class MergedQKVParallelLinearWithLoRA(MergedColumnParallelLinearWithLoRA):

488 lora_config: LoRAConfig,

489 packed_modules_list: list,

490 model_config: PretrainedConfig | None = None,

491 + decorate: bool = True,

492 ) -> bool:

493 return (

494 type(source_layer) is maybe_get_oot_by_class(QKVParallelLinear)

@@ -523,6 +530,7 @@ class ColumnParallelLinearWithShardedLoRA(ColumnParallelLinearWithLoRA):

530 lora_config: LoRAConfig,

531 packed_modules_list: list,

532 model_config: PretrainedConfig | None = None,

533 + decorate: bool = True,

534 ) -> bool:

535 # specifying kwargs so they can be easily accessed in decorator

536 return super().can_replace_layer(

@@ -565,6 +573,7 @@ class MergedColumnParallelLinearWithShardedLoRA(MergedColumnParallelLinearWithLo

573 lora_config: LoRAConfig,

574 packed_modules_list: list,

575 model_config: PretrainedConfig | None = None,

576 + decorate: bool = True,

577 ) -> bool:

578 # specifying kwargs so they can be easily accessed in decorator

579 return super().can_replace_layer(

@@ -650,6 +659,7 @@ class MergedQKVParallelLinearWithShardedLoRA(MergedQKVParallelLinearWithLoRA):

659 lora_config: LoRAConfig,

660 packed_modules_list: list,

661 model_config: PretrainedConfig | None = None,

662 + decorate: bool = True,

663 ) -> bool:

664 # specifying kwargs so they can be easily accessed in decorator

665 return super().can_replace_layer(

@@ -678,6 +688,7 @@ class MergedColumnParallelLinearVariableSliceWithLoRA(

688 lora_config: LoRAConfig,

689 packed_modules_list: list,

690 model_config: PretrainedConfig | None = None,

691 + decorate: bool = True,

692 ) -> bool:

693 # Support MergedColumnParallelLinear with 3 or more slices

694 # (2 slices are handled by MergedColumnParallelLinearWithLoRA)

@@ -727,7 +738,7 @@ class MergedColumnParallelLinearVariableSliceWithLoRA(

738 start_idx = 0

739 for output_size in output_sizes:

740 end_idx = start_idx + output_size

730 - lora_b_list.append(lora_b[start_idx:end_idx, :])

741 + lora_b_list.append(lora_b[start_idx:end_idx])

742 start_idx = end_idx

743 lora_b = lora_b_list

744

vllm/lora/layers/fused_moe.py

+25 -4

@@ -12,10 +12,16 @@ from vllm.lora.layers.base import BaseLayerWithLoRA

12 from vllm.model_executor.custom_op import maybe_get_oot_by_class

13 from vllm.model_executor.layers.fused_moe import MoERunner

14 from vllm.model_executor.layers.fused_moe.experts.lora_context import MoELoRAContext

15 +from vllm.model_executor.layers.fused_moe.experts.lora_experts_mixin import (

16 + LoRAExpertsMixin,

17 +)

18 from vllm.model_executor.layers.fused_moe.fused_moe_modular_method import (

19 FusedMoEModularMethod,

20 )

18 -from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel

21 +from vllm.model_executor.layers.fused_moe.modular_kernel import (

22 + FusedMoEKernel,

23 + FusedMoEKernelModularImpl,

24 +)

25 from vllm.model_executor.layers.fused_moe.prepare_finalize import (

26 MoEPrepareAndFinalizeNoDPEPModular,

27 )

@@ -58,6 +64,13 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):

64 routed_experts._ensure_moe_quant_config_init()

65 if getattr(routed_experts.quant_method, "supports_internal_mk", False):

66 moe_kernel = routed_experts.quant_method.moe_kernel

67 + assert moe_kernel is not None, (

68 + "Fused MoE quant method must provide a moe_kernel."

69 + )

70 + # Don't let the kernel own shared experts so the runner can

71 + # overlap them with routed experts via a separate CUDA stream.

72 + assert isinstance(moe_kernel.impl, FusedMoEKernelModularImpl)

73 + moe_kernel.impl.shared_experts = None

74 else:

75 prepare_finalize = MoEPrepareAndFinalizeNoDPEPModular()

76 moe_kernel = FusedMoEKernel(

@@ -405,7 +418,11 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):

418 def set_mapping(self, punica_wrapper):

419 super().set_mapping(punica_wrapper)

420 lora_context = self._build_lora_context()

408 - self._moe_kernel.fused_experts.set_lora_context(lora_context)

421 + fused_experts = self._moe_kernel.fused_experts

422 + assert isinstance(fused_experts, LoRAExpertsMixin), (

423 + f"{type(fused_experts).__name__} does not support LoRA context setup."

424 + )

425 + fused_experts.set_lora_context(lora_context)

426 prepare_finalize = self._moe_kernel.prepare_finalize

427 if hasattr(prepare_finalize, "set_lora_context"):

428 prepare_finalize.set_lora_context(lora_context)

@@ -482,9 +499,13 @@ class FusedMoE3DWithLoRA(FusedMoEWithLoRA):

499 ) -> None:

500 """Initializes lora matrices."""

501

485 - assert isinstance(model_config, PretrainedConfig)

502 + if model_config is None:

503 + raise ValueError("model_config must be provided for MoE LoRA.")

504 + architectures = model_config.architectures

505 + if not architectures:

506 + raise ValueError("model_config.architectures must be defined for MoE LoRA.")

507 self._verify_ep_fs(lora_config)

487 - self._base_model = model_config.architectures[0]

508 + self._base_model = architectures[0]

509 self.max_loras = lora_config.max_loras

510 self.fully_sharded = lora_config.fully_sharded_loras

511

vllm/lora/layers/row_parallel_linear.py

+1 -1

@@ -116,7 +116,7 @@ class RowParallelLinearWithShardedLoRA(RowParallelLinearWithLoRA):

116 return lora_b

117

118 def apply(self, x: torch.Tensor, bias: torch.Tensor | None = None) -> torch.Tensor:

119 - output = self.base_layer.quant_method.apply(self.base_layer, x, bias)

119 + output = self._get_quant_method().apply(self.base_layer, x, bias)

120

121 x = x.view(-1, x.shape[-1])

122 output, out_orig_shape = output.view(-1, output.shape[-1]), output.shape

vllm/lora/layers/utils.py

+11 -3

@@ -111,19 +111,27 @@ def try_get_optimal_moe_lora_config(

111 # base MoE weight's block-wise quantization, so block_shape is omitted

112 # from the config lookup — the non-quantized branch in get_default_config

113 # ignores it anyway.

114 - config = try_get_optimal_moe_config(w1_shape, w2_shape, top_k, dtype, M).copy()

114 + raw_config = try_get_optimal_moe_config(w1_shape, w2_shape, top_k, dtype, M)

115 + config: dict[str, int | None] = dict(raw_config)

116 if op_type in [

117 "fused_moe_lora_w13_shrink",

118 "fused_moe_lora_w2_shrink",

119 ]:

120 + block_size_n = config.get("BLOCK_SIZE_N")

121 config["BLOCK_SIZE_N"] = min(

120 - config.get("BLOCK_SIZE_N", 64), next_power_of_2(rank)

122 + block_size_n if block_size_n is not None else 64,

123 + next_power_of_2(rank),

124 )

125 elif op_type in [

126 "fused_moe_lora_w13_expand",

127 "fused_moe_lora_w2_expand",

128 ]:

129 + block_size_k = config.get("BLOCK_SIZE_K")

130 config["BLOCK_SIZE_K"] = max(

127 - 16, min(config.get("BLOCK_SIZE_K", 32), next_power_of_2(rank))

131 + 16,

132 + min(

133 + block_size_k if block_size_k is not None else 32,

134 + next_power_of_2(rank),

135 + ),

136 )

137 return config

vllm/lora/lora_model.py

+4 -3

@@ -245,9 +245,10 @@ class LoRAModel:

245 from tensorizer import TensorDeserializer

246

247 tensorizer_config = TensorizerConfig(**tensorizer_config_dict)

248 - lora_tensor_path = os.path.join(

249 - tensorizer_config.tensorizer_dir, "adapter_model.tensors"

250 - )

248 + tensorizer_dir = tensorizer_config.tensorizer_dir

249 + if tensorizer_dir is None:

250 + raise ValueError("tensorizer_dir must be set in tensorizer config.")

251 + lora_tensor_path = os.path.join(tensorizer_dir, "adapter_model.tensors")

252 tensorizer_args = tensorizer_config._construct_tensorizer_args()

253 tensors = TensorDeserializer(

254 lora_tensor_path,

vllm/lora/model_manager.py

+17 -10

@@ -34,6 +34,7 @@ from vllm.lora.utils import (

34 from vllm.model_executor.layers.fused_moe import MoERunner

35 from vllm.model_executor.models import (

36 SupportsLoRA,

37 + SupportsMultiModal,

38 is_pooling_model,

39 supports_multimodal,

40 )

@@ -50,6 +51,12 @@ T = TypeVar("T")

51 DEFAULT_LANGUAGE_WRAPPER_KEY = "language_model"

54 +class SupportsLoRAModel(nn.Module, SupportsLoRA): ...

55 +

56 +

57 +class SupportsLoRAMultiModalModel(SupportsLoRAModel, SupportsMultiModal): ...

58 +

59 +

60 class AdapterLRUCache(LRUCache[int, T]):

61 def __init__(self, capacity: int, deactivate_fn: Callable[[int], object]):

62 super().__init__(capacity)

@@ -66,13 +73,13 @@ class LoRAModelManager:

74 def __init__(

75 self,

69 - model: SupportsLoRA,

76 + model: SupportsLoRAModel,

77 max_num_seqs: int,

78 max_num_batched_tokens: int,

79 vocab_size: int,

80 lora_config: LoRAConfig,

81 device: torch.device,

75 - vllm_config: VllmConfig | None = None,

82 + vllm_config: VllmConfig,

83 ):

84 """Create a LoRAModelManager and adapter for a given model.

@@ -85,7 +92,7 @@ class LoRAModelManager:

92 vocab_size: the vocab size of the model.

93 lora_config: the LoRA configuration.

94 """

88 - self.model: SupportsLoRA = model

95 + self.model: SupportsLoRAModel = model

96 self.supported_lora_modules = get_supported_lora_modules(self.model)

97 assert self.supported_lora_modules, (

98 f"No supported LoRA modules found in {self.model.__class__.__name__}."

@@ -106,7 +113,6 @@ class LoRAModelManager:

113 self.is_pooling_model = is_pooling_model(self.model)

114 self.packed_modules: dict[str, list[str]] = {}

115 self.modules: dict[str, BaseLayerWithLoRA] = {}

109 - # Dict instead of a set for compatibility with LRUCache.

116 self._last_mapping: LoRAMapping | None = None

117 is_moe = is_moe_model(self.model)

118 self._is_moe = is_moe

@@ -272,6 +278,7 @@ class LoRAModelManager:

278

279 @property

280 def capacity(self) -> int:

281 + assert self.lora_config.max_cpu_loras is not None

282 return self.lora_config.max_cpu_loras

283

284 @property

@@ -1156,7 +1163,7 @@ class LoRAModelManager:

1163

1164

1165 class LoRALRUCache(AdapterLRUCache[LoRAModel]):

1159 - def __init__(self, capacity: int, deactivate_lora_fn: Callable[[int], bool]):

1166 + def __init__(self, capacity: int, deactivate_lora_fn: Callable[[int], object]):

1167 super().__init__(capacity, deactivate_lora_fn)

1168

1169

@@ -1165,13 +1172,13 @@ class LRUCacheLoRAModelManager(LoRAModelManager):

1172

1173 def __init__(

1174 self,

1168 - model: nn.Module,

1175 + model: SupportsLoRAModel,

1176 max_num_seqs: int,

1177 max_num_batched_tokens: int,

1178 vocab_size: int,

1179 lora_config: LoRAConfig,

1180 device: torch.device,

1174 - vllm_config: VllmConfig | None = None,

1181 + vllm_config: VllmConfig,

1182 ):

1183 super().__init__(

1184 model,

@@ -1182,10 +1189,10 @@ class LRUCacheLoRAModelManager(LoRAModelManager):

1189 device,

1190 vllm_config,

1191 )

1185 - self._registered_adapters: LoRALRUCache = LoRALRUCache(

1192 + self._registered_adapters: LoRALRUCache = LoRALRUCache( # type: ignore[assignment]

1193 self.capacity, self.deactivate_adapter

1194 )

1188 - self._active_adapters: LoRALRUCache = LoRALRUCache(

1195 + self._active_adapters: LoRALRUCache = LoRALRUCache( # type: ignore[assignment]

1196 self.lora_slots, self._deactivate_adapter

1197 )

1198

@@ -1248,7 +1255,7 @@ class LRUCacheLoRAModelManager(LoRAModelManager):

1255

1256

1257 def create_lora_manager(

1251 - model: nn.Module,

1258 + model: SupportsLoRAModel,

1259 max_num_seqs: int,

1260 max_num_batched_tokens: int,

1261 vocab_size: int,

vllm/lora/peft_helper.py

+5 -3

@@ -91,9 +91,11 @@ class PEFTHelper:

91 tensorizer_args = tensorizer_config._construct_tensorizer_args()

92 from tensorizer.stream_io import open_stream

94 - lora_config_path = os.path.join(

95 - tensorizer_config.tensorizer_dir, "adapter_config.json"

96 - )

94 + tensorizer_dir = tensorizer_config.tensorizer_dir

95 + if tensorizer_dir is None:

96 + raise ValueError("tensorizer_dir must be set in tensorizer config.")

97 +

98 + lora_config_path = os.path.join(tensorizer_dir, "adapter_config.json")

99 with open_stream(

100 lora_config_path, mode="rb", **tensorizer_args.stream_kwargs

101 ) as f:

vllm/lora/utils.py

+16 -5

@@ -173,11 +173,18 @@ def parse_fine_tuned_lora_name(

173 # mapping correctly.

174 if name.startswith("base_model.model."):

175 name = name.replace("base_model.model.", "")

176 - name = weights_mapper._map_name(name) if weights_mapper else name

177 - # recover the prefix `base_model.model.`

178 - name = "base_model.model." + name

176 + if weights_mapper:

177 + mapped_name = weights_mapper._map_name(name)

178 + if mapped_name is None:

179 + raise ValueError("Mapped LoRA weight name cannot be None.")

180 + # recover the prefix `base_model.model.`

181 + name = "base_model.model." + mapped_name

182 else:

180 - name = weights_mapper._map_name(name) if weights_mapper else name

183 + if weights_mapper:

184 + mapped_name = weights_mapper._map_name(name)

185 + if mapped_name is None:

186 + raise ValueError("Mapped LoRA weight name cannot be None.")

187 + name = mapped_name

188

189 # In some situations, we may not start with `base_model.model.`.

190 # If we don't (e.g., ibm-granite/granite-speech-3.3-8b),

@@ -185,7 +192,11 @@ def parse_fine_tuned_lora_name(

192 start_index = 2 if name.startswith("base_model.model.") else 0

193

194 parts = name.split(".")

188 - if parts[-1] == "weight" and (parts[-2] == "lora_A" or parts[-2] == "lora_B"):

195 + if (

196 + parts[-1] == "weight"

197 + and len(parts) >= 2

198 + and (parts[-2] == "lora_A" or parts[-2] == "lora_B")

199 + ):

200 new_name = ".".join(parts[start_index:-2])

201 return new_name, parts[-2] == "lora_A"

202

vllm/lora/worker_manager.py

+11 -3

@@ -7,6 +7,7 @@ from typing import Any, Literal

7 import torch

9 from vllm.config import VllmConfig

10 +from vllm.config.lora import LoRAConfig

11 from vllm.exceptions import LoRAAdapterNotFoundError

12 from vllm.logger import init_logger

13 from vllm.lora.lora_model import LoRAModel

@@ -45,7 +46,10 @@ class WorkerLoRAManager:

46 vllm_config.scheduler_config.max_num_batched_tokens

47 )

48 self.vocab_size = vllm_config.model_config.get_vocab_size()

48 - self.lora_config = vllm_config.lora_config

49 + lora_config = vllm_config.lora_config

50 + if lora_config is None:

51 + raise ValueError("LoRA config must be set for WorkerLoRAManager.")

52 + self.lora_config: LoRAConfig = lora_config

54 # Use get_text_config() in case of multimodal models

55 text_config = vllm_config.model_config.hf_config.get_text_config()

@@ -81,8 +85,10 @@ class WorkerLoRAManager:

85 def create_lora_manager(

86 self,

87 model: torch.nn.Module,

84 - vllm_config: VllmConfig | None = None,

88 + vllm_config: VllmConfig,

89 ) -> Any:

90 + if vllm_config is None:

91 + raise ValueError("vllm_config must be provided to create a LoRA manager.")

92 lora_manager = create_lora_manager(

93 model,

94 max_num_seqs=self.max_num_seqs,

@@ -240,8 +246,10 @@ class LRUCacheWorkerLoRAManager(WorkerLoRAManager):

246 def create_lora_manager(

247 self,

248 model: torch.nn.Module,

243 - vllm_config: VllmConfig | None = None,

249 + vllm_config: VllmConfig,

250 ) -> Any:

251 + if vllm_config is None:

252 + raise ValueError("vllm_config must be provided to create a LoRA manager.")

253 lora_manager = create_lora_manager(

254 model,

255 lora_manager_cls=self._manager_cls,

vllm/model_executor/layers/fused_moe/modular_kernel.py

@@ -1029,6 +1029,7 @@ class FusedMoEKernelModularImpl:

1029 ):

1030 self.prepare_finalize = prepare_finalize

1031 self.fused_experts = fused_experts

1032 + self.shared_experts: SharedExperts | None = None

1033 moe_parallel_config = fused_experts.moe_config.moe_parallel_config

1034 self.moe_parallel_config = moe_parallel_config

1035 self.is_dp_ep = (

vllm/model_executor/models/interfaces.py

@@ -42,6 +42,7 @@ from .interfaces_base import VllmModel

43 if TYPE_CHECKING:

44 from vllm.config import VllmConfig

45 + from vllm.lora.model_manager import LoRAModelManager

46 from vllm.model_executor.models.utils import WeightsMapper

47 from vllm.multimodal.inputs import MultiModalFeatureSpec

48 from vllm.multimodal.registry import _ProcessorFactories

@@ -554,6 +555,7 @@ class SupportsLoRA(Protocol):

555 packed_modules_mapping: dict[str, list[str]] = {}

556 # Module prefixes to skip during LoRA loading (e.g., ["mtp."] for MTP layers)

557 lora_skip_prefixes: ClassVar[list[str]] = []

558 + lora_manager: "LoRAModelManager | None"

559

560

561 # We can't use runtime_checkable with ClassVar for issubclass checks

vllm/models/deepseek_v4/xpu/mtp.py

+1 -3

@@ -24,9 +24,7 @@ from vllm.distributed import (

24 get_tensor_model_parallel_world_size,

25 )

26 from vllm.logger import init_logger

27 -from vllm.model_executor.layers.fused_moe import (

28 - fused_moe_make_expert_params_mapping,

29 -)

27 +from vllm.model_executor.layers.fused_moe import fused_moe_make_expert_params_mapping

28 from vllm.model_executor.layers.layernorm import RMSNorm

29 from vllm.model_executor.layers.linear import ReplicatedLinear

30 from vllm.model_executor.layers.logits_processor import LogitsProcessor