@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / a250f1bd5f

Fork 0

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 4, 20:00 UTC

[Bugfix] LoRA for DeepSeek V3.2 (#35077)

Signed-off-by: Hollow Man <hollowman@opensuse.org> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Jee Jee Li <pandaleefree@gmail.com>

ℍ𝕠𝕝𝕝𝕠𝕨 𝕄𝕒𝕟 committed Apr 22, 2026 at 14:33 UTC a250f1bd5fd8ac7a7d97bc8b5b2082417b4564d2

13 files changed +614 -37

tests/kernels/moe/test_moe_layer.py

+7 -12

@@ -1620,18 +1620,13 @@ def _parallel_worker(

1620 else:

1621 print("F", end="")

1622 finally:

1623 - # Note: for some reason DeepEP buffers don't seem to be

1624 - # entirely reusable on B200. In order to work around this

1625 - # we clear the all2all manager's cache after each testpoint.

1626 - cap = current_platform.get_device_capability()

1627 - if (

1628 - cap is not None

1629 - and cap.major == 10

1630 - and (

1631 - test_config.backend == "deepep_low_latency"

1632 - or test_config.backend == "deepep_high_throughput"

1633 - )

1634 - ):

1623 + # DeepEP managers are not reliably reusable across many subtests in

1624 + # a single worker process. Tear them down after each DeepEP case so

1625 + # later subtests do not inherit stale communication state.

1626 + if test_config.backend in {

1627 + "deepep_low_latency",

1628 + "deepep_high_throughput",

1629 + }:

1630 torch.accelerator.synchronize()

1631 all2all_manager = get_ep_group().device_communicator.all2all_manager

1632 if all2all_manager is not None:

tests/lora/test_layers.py

+270 -2

@@ -44,6 +44,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (

44 VocabParallelEmbedding,

45 get_masked_input_and_mask,

46 )

47 +from vllm.model_executor.models.deepseek_v2 import DeepSeekV2FusedQkvAProjLinear

48 from vllm.platforms import current_platform

49 from vllm.utils.torch_utils import set_random_seed

@@ -1422,7 +1423,107 @@ def test_variable_slice_lora_class_selection(default_vllm_config, dist_init):

1423 f"for 2 packed modules, got {type(selected_layer_merged).__name__}"

1424 )

1425

1425 - # Case 5: Plain ColumnParallelLinear (not merged) - common in many models

1426 + fully_sharded_tp_lora_config = LoRAConfig(

1427 + max_loras=8,

1428 + max_lora_rank=16,

1429 + lora_dtype=torch.float16,

1430 + fully_sharded_loras=True,

1431 + )

1432 + fully_sharded_tp_layer = MergedColumnParallelLinear(

1433 + 4096, [2048, 2048], bias=False, params_dtype=torch.float16

1434 + )

1435 + fully_sharded_tp_layer.tp_size = 2

1436 +

1437 + assert not MergedColumnParallelLinearWithLoRA.can_replace_layer(

1438 + source_layer=fully_sharded_tp_layer,

1439 + lora_config=fully_sharded_tp_lora_config,

1440 + packed_modules_list=packed_modules_two,

1441 + ), "Generic merged wrapper should reject fully sharded TP layers"

1442 +

1443 + assert MergedColumnParallelLinearWithShardedLoRA.can_replace_layer(

1444 + source_layer=fully_sharded_tp_layer,

1445 + lora_config=fully_sharded_tp_lora_config,

1446 + packed_modules_list=packed_modules_two,

1447 + ), "Sharded merged wrapper should remain eligible for fully sharded TP layers"

1448 +

1449 + selected_fully_sharded_tp_layer = from_layer(

1450 + fully_sharded_tp_layer,

1451 + max_loras=8,

1452 + lora_config=fully_sharded_tp_lora_config,

1453 + packed_modules_list=packed_modules_two,

1454 + )

1455 + assert isinstance(

1456 + selected_fully_sharded_tp_layer,

1457 + MergedColumnParallelLinearWithShardedLoRA,

1458 + ), (

1459 + "from_layer should select MergedColumnParallelLinearWithShardedLoRA "

1460 + "for fully sharded TP merged layers, got "

1461 + f"{type(selected_fully_sharded_tp_layer).__name__}"

1462 + )

1463 +

1464 + # Case 5: DeepSeek's fused_qkv_a_proj should reuse the generic merged

1465 + # wrapper while preserving its custom base forward path.

1466 + deepseek_fused_layer = DeepSeekV2FusedQkvAProjLinear(

1467 + 4096, [2048, 2048], prefix="model.layers.0.self_attn.fused_qkv_a_proj"

1468 + )

1469 + selected_deepseek_layer = from_layer(

1470 + deepseek_fused_layer,

1471 + max_loras=8,

1472 + lora_config=lora_config,

1473 + packed_modules_list=packed_modules_two,

1474 + )

1475 + assert isinstance(selected_deepseek_layer, MergedColumnParallelLinearWithLoRA), (

1476 + "from_layer should select MergedColumnParallelLinearWithLoRA "

1477 + f"for DeepSeek fused_qkv_a_proj, got {type(selected_deepseek_layer).__name__}"

1478 + )

1479 +

1480 + fully_sharded_lora_config = LoRAConfig(

1481 + max_loras=8,

1482 + max_lora_rank=16,

1483 + lora_dtype=torch.float16,

1484 + fully_sharded_loras=True,

1485 + )

1486 + selected_fully_sharded_deepseek_layer = from_layer(

1487 + deepseek_fused_layer,

1488 + max_loras=8,

1489 + lora_config=fully_sharded_lora_config,

1490 + packed_modules_list=packed_modules_two,

1491 + )

1492 + assert isinstance(

1493 + selected_fully_sharded_deepseek_layer,

1494 + MergedColumnParallelLinearWithLoRA,

1495 + ), (

1496 + "from_layer should keep using MergedColumnParallelLinearWithLoRA "

1497 + "for fused_qkv_a_proj when the base layer is effectively unsharded, got "

1498 + f"{type(selected_fully_sharded_deepseek_layer).__name__}"

1499 + )

1500 +

1501 + # Case 6: Generic subclass of MergedColumnParallelLinear with 2 packed

1502 + # modules should still use the generic merged wrapper.

1503 + class CustomMergedColumnParallelLinear(MergedColumnParallelLinear):

1504 + pass

1505 +

1506 + custom_merged_layer = CustomMergedColumnParallelLinear(

1507 + 4096, [2048, 2048], bias=False, params_dtype=torch.float16

1508 + )

1509 + assert MergedColumnParallelLinearWithLoRA.can_replace_layer(

1510 + source_layer=custom_merged_layer,

1511 + lora_config=lora_config,

1512 + packed_modules_list=packed_modules_two,

1513 + ), "MergedColumnParallelLinearWithLoRA should handle subclasses"

1514 +

1515 + selected_custom_layer = from_layer(

1516 + custom_merged_layer,

1517 + max_loras=8,

1518 + lora_config=lora_config,

1519 + packed_modules_list=packed_modules_two,

1520 + )

1521 + assert isinstance(selected_custom_layer, MergedColumnParallelLinearWithLoRA), (

1522 + f"from_layer should select MergedColumnParallelLinearWithLoRA "

1523 + f"for subclassed merged layers, got {type(selected_custom_layer).__name__}"

1524 + )

1525 +

1526 + # Case 7: Plain ColumnParallelLinear (not merged) - common in many models

1527 # -> ColumnParallelLinearWithLoRA should be selected

1528 plain_column_parallel = ColumnParallelLinear(

1529 4096, 4096, bias=False, params_dtype=torch.float16

@@ -1455,7 +1556,7 @@ def test_variable_slice_lora_class_selection(default_vllm_config, dist_init):

1556 f"for plain ColumnParallelLinear, got {type(selected_plain).__name__}"

1557 )

1558

1458 - # Case 6: MergedColumnParallelLinear with exactly 2 output sizes

1559 + # Case 8: MergedColumnParallelLinear with exactly 2 output sizes

1560 # and empty packed_modules_list

1561 # -> ColumnParallelLinearWithLoRA should NOT match (packed_modules_list != 1)

1562 # -> MergedColumnParallelLinearVariableSliceWithLoRA should NOT match (< 3 slices)

@@ -1473,3 +1574,170 @@ def test_variable_slice_lora_class_selection(default_vllm_config, dist_init):

1574 "MergedColumnParallelLinearVariableSliceWithLoRA "

1575 "should NOT handle 2 slices even with empty packed_modules_list"

1576 )

1577 +

1578 +

1579 +@pytest.mark.parametrize(

1580 + "wrapper_cls",

1581 + [ColumnParallelLinearWithLoRA, ColumnParallelLinearWithShardedLoRA],

1582 +)

1583 +def test_get_and_maybe_dequant_weights_accepts_lora_wrappers(dist_init, wrapper_cls):

1584 + from vllm.model_executor.layers.quantization.utils.quant_utils import (

1585 + get_and_maybe_dequant_weights,

1586 + )

1587 +

1588 + linear = ColumnParallelLinear(4096, 4096, bias=False, params_dtype=torch.float16)

1589 + lora_linear = wrapper_cls(linear)

1590 +

1591 + # Should work with LoRA wrappers and return [out, in] weights.

1592 + dequant_weight = get_and_maybe_dequant_weights(lora_linear, out_dtype=torch.float16)

1593 + assert dequant_weight.shape == linear.weight.shape

1594 +

1595 +

1596 +@torch.inference_mode()

1597 +@pytest.mark.parametrize("device", DEVICES)

1598 +@pytest.mark.parametrize("stage", STAGES)

1599 +@pytest.mark.parametrize("fully_sharded", [False, True])

1600 +def test_deepseek_fused_qkv_a_proj_lora_preserves_base_forward(

1601 + default_vllm_config, dist_init, device, stage, fully_sharded

1602 +):

1603 + if current_platform.is_cuda_alike():

1604 + torch.accelerator.set_device_index(device)

1605 +

1606 + torch.set_default_device(device)

1607 + dtype = torch.float16 if current_platform.is_cuda_alike() else torch.float32

1608 + max_loras = 8

1609 + lora_config = LoRAConfig(

1610 + max_loras=max_loras,

1611 + max_lora_rank=8,

1612 + lora_dtype=dtype,

1613 + fully_sharded_loras=fully_sharded,

1614 + )

1615 + punica_wrapper = get_punica_wrapper(8192, 256, device, lora_config=lora_config)

1616 + assert check_punica_wrapper(punica_wrapper)

1617 +

1618 + class OffsetDeepSeekFusedQkvAProjLinear(DeepSeekV2FusedQkvAProjLinear):

1619 + def forward(self, input_):

1620 + output, output_bias = super().forward(input_)

1621 + return output + 1, output_bias

1622 +

1623 + layer = OffsetDeepSeekFusedQkvAProjLinear(

1624 + 32, [16, 16], prefix="model.layers.0.self_attn.fused_qkv_a_proj"

1625 + )

1626 + layer.weight.data = torch.rand_like(layer.weight.data, dtype=dtype)

1627 +

1628 + lora_layer = MergedColumnParallelLinearWithLoRA(layer)

1629 + lora_layer.create_lora_weights(max_loras, lora_config)

1630 + lora_layer.set_mapping(punica_wrapper)

1631 +

1632 + id_to_index = get_random_id_to_index(1, max_loras, log=False)

1633 + active_slot = next(i for i, lora_id in enumerate(id_to_index) if lora_id == 1)

1634 + lora_a = [

1635 + torch.rand(8, 32, dtype=dtype, device=device),

1636 + torch.rand(8, 32, dtype=dtype, device=device),

1637 + ]

1638 + lora_b = [

1639 + torch.rand(16, 8, dtype=dtype, device=device),

1640 + torch.rand(16, 8, dtype=dtype, device=device),

1641 + ]

1642 + lora_layer.set_lora(active_slot, lora_a=lora_a, lora_b=lora_b)

1643 +

1644 + inputs, index_mapping, prompt_mapping = create_random_inputs(

1645 + active_lora_ids=[1],

1646 + num_inputs=4,

1647 + input_size=(1, 32),

1648 + input_range=(0, 1),

1649 + input_type=dtype,

1650 + device=device,

1651 + )

1652 + lora_mapping = LoRAMapping(index_mapping, prompt_mapping, is_prefill=stage)

1653 + punica_wrapper.update_metadata(lora_mapping, id_to_index, max_loras, 512)

1654 +

1655 + lora_result = lora_layer(torch.cat(inputs))[0]

1656 +

1657 + expected_results = []

1658 + for input_ in inputs:

1659 + result = layer(input_)[0]

1660 + result[:, :16] += input_ @ lora_a[0].T @ lora_b[0].T

1661 + result[:, 16:] += input_ @ lora_a[1].T @ lora_b[1].T

1662 + expected_results.append(result)

1663 +

1664 + rtol, atol = TOLERANCES[lora_result.dtype]

1665 + torch.testing.assert_close(

1666 + lora_result, torch.cat(expected_results), rtol=rtol, atol=atol

1667 + )

1668 +

1669 + merged_layer = OffsetDeepSeekFusedQkvAProjLinear(

1670 + 32, [16, 16], prefix="model.layers.0.self_attn.fused_qkv_a_proj"

1671 + )

1672 + merged_layer.weight.data = layer.weight.data.clone()

1673 + merged_layer.weight.data[:16].add_(lora_b[0] @ lora_a[0])

1674 + merged_layer.weight.data[16:].add_(lora_b[1] @ lora_a[1])

1675 + merged_result = merged_layer(torch.cat(inputs))[0]

1676 +

1677 + torch.testing.assert_close(lora_result, merged_result, rtol=rtol, atol=atol)

1678 +

1679 +

1680 +@torch.inference_mode()

1681 +@pytest.mark.parametrize("device", DEVICES)

1682 +@pytest.mark.parametrize("stage", STAGES)

1683 +def test_replicated_lora_preserves_base_forward_for_subclasses(

1684 + default_vllm_config, dist_init, device, stage

1685 +):

1686 + if current_platform.is_cuda_alike():

1687 + torch.accelerator.set_device_index(device)

1688 +

1689 + torch.set_default_device(device)

1690 + dtype = torch.float16 if current_platform.is_cuda_alike() else torch.float32

1691 + max_loras = 8

1692 + lora_config = LoRAConfig(max_loras=max_loras, max_lora_rank=8, lora_dtype=dtype)

1693 + punica_wrapper = get_punica_wrapper(8192, 256, device, lora_config=lora_config)

1694 + assert check_punica_wrapper(punica_wrapper)

1695 +

1696 + class OffsetReplicatedLinear(ReplicatedLinear):

1697 + def forward(self, input_):

1698 + output, output_bias = super().forward(input_)

1699 + return output + 1, output_bias

1700 +

1701 + layer = OffsetReplicatedLinear(32, 16, bias=False, params_dtype=dtype)

1702 + layer.weight.data = torch.rand_like(layer.weight.data, dtype=dtype)

1703 +

1704 + lora_layer = ReplicatedLinearWithLoRA(layer)

1705 + lora_layer.create_lora_weights(max_loras, lora_config)

1706 + lora_layer.set_mapping(punica_wrapper)

1707 +

1708 + id_to_index = get_random_id_to_index(1, max_loras, log=False)

1709 + active_slot = next(i for i, lora_id in enumerate(id_to_index) if lora_id == 1)

1710 + lora_a = torch.rand(8, 32, dtype=dtype, device=device)

1711 + lora_b = torch.rand(16, 8, dtype=dtype, device=device)

1712 + lora_layer.set_lora(active_slot, lora_a=lora_a, lora_b=lora_b)

1713 +

1714 + inputs, index_mapping, prompt_mapping = create_random_inputs(

1715 + active_lora_ids=[1],

1716 + num_inputs=4,

1717 + input_size=(1, 32),

1718 + input_range=(0, 1),

1719 + input_type=dtype,

1720 + device=device,

1721 + )

1722 + lora_mapping = LoRAMapping(index_mapping, prompt_mapping, is_prefill=stage)

1723 + punica_wrapper.update_metadata(lora_mapping, id_to_index, max_loras, 512)

1724 +

1725 + lora_result = lora_layer(torch.cat(inputs))[0]

1726 +

1727 + expected_results = []

1728 + for input_ in inputs:

1729 + result = layer(input_)[0]

1730 + result += input_ @ lora_a.T @ lora_b.T

1731 + expected_results.append(result)

1732 +

1733 + rtol, atol = TOLERANCES[lora_result.dtype]

1734 + torch.testing.assert_close(

1735 + lora_result, torch.cat(expected_results), rtol=rtol, atol=atol

1736 + )

1737 +

1738 + merged_layer = OffsetReplicatedLinear(32, 16, bias=False, params_dtype=dtype)

1739 + merged_layer.weight.data = layer.weight.data.clone()

1740 + merged_layer.weight.data.add_(lora_b @ lora_a)

1741 + merged_result = merged_layer(torch.cat(inputs))[0]

1742 +

1743 + torch.testing.assert_close(lora_result, merged_result, rtol=rtol, atol=atol)

tests/lora/test_lora_manager.py

+150

@@ -13,6 +13,7 @@ from vllm.config.lora import LoRAConfig

13 from vllm.lora.layers import (

14 ColumnParallelLinearWithLoRA,

15 MergedColumnParallelLinearWithLoRA,

16 + ReplicatedLinearWithLoRA,

17 RowParallelLinearWithLoRA,

18 )

19 from vllm.lora.lora_model import LoRAModel

@@ -26,6 +27,7 @@ from vllm.lora.model_manager import (

27 from vllm.lora.peft_helper import PEFTHelper

28 from vllm.lora.request import LoRARequest

29 from vllm.lora.worker_manager import LRUCacheWorkerLoRAManager, WorkerLoRAManager

30 +from vllm.model_executor.layers.fused_moe import GateLinear

31 from vllm.platforms import current_platform

33 from .utils import create_peft_lora

@@ -132,6 +134,135 @@ def test_replace_submodules(default_vllm_config, dist_init, dummy_model):

134 assert isinstance(model.get_submodule("layer1.dense2"), RowParallelLinearWithLoRA)

135

136

137 +def test_wrap_replicated_linear_subclasses(default_vllm_config, dist_init, dummy_model):

138 + from vllm.model_executor.layers.linear import ReplicatedLinear

139 +

140 + class CustomReplicatedLinear(ReplicatedLinear):

141 + pass

142 +

143 + model = dummy_model

144 + model.add_module("custom_gate", CustomReplicatedLinear(10, 10, bias=False))

145 +

146 + manager = LoRAModelManager(

147 + model,

148 + 1,

149 + 1,

150 + 1,

151 + LoRAConfig(

152 + max_lora_rank=8, max_cpu_loras=8, max_loras=8, lora_dtype=DEFAULT_DTYPE

153 + ),

154 + torch.device(DEVICES[0]),

155 + )

156 +

157 + assert isinstance(

158 + manager.model.get_submodule("custom_gate"), ReplicatedLinearWithLoRA

159 + )

160 +

161 +

162 +def test_wrap_gate_linear(default_vllm_config, dist_init, dummy_model):

163 + model = dummy_model

164 + model.add_module("router_gate", GateLinear(10, 4, bias=False))

165 +

166 + manager = LoRAModelManager(

167 + model,

168 + 1,

169 + 1,

170 + 1,

171 + LoRAConfig(

172 + max_lora_rank=8, max_cpu_loras=8, max_loras=8, lora_dtype=DEFAULT_DTYPE

173 + ),

174 + torch.device(DEVICES[0]),

175 + )

176 +

177 + assert isinstance(

178 + manager.model.get_submodule("router_gate"), ReplicatedLinearWithLoRA

179 + )

180 +

181 +

182 +def test_skip_unsupported_matched_modules(default_vllm_config, dist_init, dummy_model):

183 + class UnsupportedContainer(nn.Module):

184 + def __init__(self):

185 + super().__init__()

186 + # This name matches a supported target suffix ("dense1"),

187 + # but nn.Linear is not currently a LoRA-wrappable layer type.

188 + self.dense1 = nn.Linear(10, 10, bias=False)

189 +

190 + model = dummy_model

191 + model.add_module("unsupported", UnsupportedContainer())

192 +

193 + manager = LoRAModelManager(

194 + model,

195 + 1,

196 + 1,

197 + 1,

198 + LoRAConfig(

199 + max_lora_rank=8, max_cpu_loras=8, max_loras=8, lora_dtype=DEFAULT_DTYPE

200 + ),

201 + torch.device(DEVICES[0]),

202 + )

203 +

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

205 + assert isinstance(manager.model.get_submodule("unsupported.dense1"), nn.Linear)

206 + assert "unsupported.dense1" not in manager.modules

207 +

208 +

209 +def test_target_modules_fail_closed_on_unsupported_matched_modules(

210 + default_vllm_config, dist_init, dummy_model

211 +):

212 + class UnsupportedContainer(nn.Module):

213 + def __init__(self):

214 + super().__init__()

215 + self.dense1 = nn.Linear(10, 10, bias=False)

216 +

217 + model = dummy_model

218 + model.add_module("unsupported", UnsupportedContainer())

219 +

220 + with pytest.raises(ValueError, match="unsupported.dense1"):

221 + LoRAModelManager(

222 + model,

223 + 1,

224 + 1,

225 + 1,

226 + LoRAConfig(

227 + max_lora_rank=8,

228 + max_cpu_loras=8,

229 + max_loras=8,

230 + lora_dtype=DEFAULT_DTYPE,

231 + target_modules=["dense1"],

232 + ),

233 + torch.device(DEVICES[0]),

234 + )

235 +

236 +

237 +def test_get_dummy_lora_warmup_rank_for_fully_sharded_moe():

238 + manager = LoRAModelManager.__new__(LoRAModelManager)

239 + manager.lora_config = LoRAConfig(

240 + max_lora_rank=64,

241 + max_cpu_loras=1,

242 + max_loras=1,

243 + lora_dtype=DEFAULT_DTYPE,

244 + fully_sharded_loras=True,

245 + )

246 +

247 + class DummyModule:

248 + def __init__(self, tp_size: int, fully_sharded: bool):

249 + self.tp_size = tp_size

250 + self.fully_sharded = fully_sharded

251 +

252 + manager.modules = {

253 + "model.layers.0.self_attn.q_proj": DummyModule(

254 + tp_size=32,

255 + fully_sharded=True,

256 + ),

257 + "model.layers.0.mlp.experts": DummyModule(

258 + tp_size=32,

259 + fully_sharded=True,

260 + ),

261 + }

262 +

263 + assert manager.get_dummy_lora_warmup_rank(8) == 32

264 +

265 +

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

267 def test_lora_model_manager(default_vllm_config, dist_init, dummy_model, device):

268 model = dummy_model

@@ -795,6 +926,25 @@ def test_target_modules_none_uses_all(

926 )

927

928

929 +@pytest.mark.parametrize("device", DEVICES)

930 +def test_target_modules_match_packed_runtime_modules(

931 + default_vllm_config, dist_init, dummy_model_gate_up, device

932 +):

933 + """Packed runtime modules should be selected by their adapter-visible names."""

934 + _test_target_modules(

935 + dummy_model_gate_up,

936 + ["gate_proj"],

937 + device,

938 + expected_lora=[("gate_up_proj", MergedColumnParallelLinearWithLoRA)],

939 + expected_no_lora=[

940 + ("dense1", ColumnParallelLinearWithLoRA),

941 + ("dense2", RowParallelLinearWithLoRA),

942 + ("layer1.dense1", ColumnParallelLinearWithLoRA),

943 + ("layer1.dense2", RowParallelLinearWithLoRA),

944 + ],

945 + )

946 +

947 +

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

949 def test_load_adapter_warns_on_unsupported_modules(

950 default_vllm_config, dist_init, dummy_model_gate_up, device, tmp_path

tests/lora/test_lora_utils.py

+21

@@ -58,3 +58,24 @@ class TestIsInTargetModules:

59 def test_exact_name_no_match(self):

60 assert not is_in_target_modules("dense3", ["dense1", "dense2"])

61 +

62 + def test_packed_parent_matches_child_target_modules(self):

63 + assert is_in_target_modules(

64 + "model.layers.0.mlp.gate_up_proj",

65 + ["gate_proj", "up_proj"],

66 + {"gate_up_proj": ["gate_proj", "up_proj"]},

67 + )

68 +

69 + def test_packed_child_matches_parent_target_modules(self):

70 + assert is_in_target_modules(

71 + "model.layers.0.mlp.gate_proj",

72 + ["gate_up_proj"],

73 + {"gate_up_proj": ["gate_proj", "up_proj"]},

74 + )

75 +

76 + def test_fused_parent_matches_child_target_modules(self):

77 + assert is_in_target_modules(

78 + "model.layers.0.self_attn.fused_qkv_a_proj",

79 + ["q_a_proj", "kv_a_proj_with_mqa"],

80 + {"fused_qkv_a_proj": ["q_a_proj", "kv_a_proj_with_mqa"]},

81 + )

vllm/lora/layers/base_linear.py

@@ -203,7 +203,16 @@ class BaseLinearLayerWithLoRA(BaseLayerWithLoRA):

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

204 ) -> torch.Tensor:

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

206 + return self._apply_lora_to_output(x, output)

207

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

209 + base_output = self.base_layer(x)

210 + output = base_output[0] if isinstance(base_output, tuple) else base_output

211 + return self._apply_lora_to_output(x, output)

212 +

213 + def _apply_lora_to_output(

214 + self, x: torch.Tensor, output: torch.Tensor

215 + ) -> torch.Tensor:

216 original_shape = output.shape if output.ndim == 3 else None

217

218 # In transformers backend, x and output have extra batch dimension like

vllm/lora/layers/column_parallel_linear.py

+42 -10

@@ -40,11 +40,19 @@ def _mcp_apply(x, bias, layer: "ColumnParallelLinearWithLoRA"):

41 # Since communication is needed, the buffer is directly initialized as a

42 # tensor rather than a tuple of tensor.

43 - buffers = torch.zeros(

44 - (layer.n_slices, x.shape[0], layer.lora_a_stacked[0].shape[2]),

43 + local_lora_rank = layer.lora_a_stacked[0].shape[2]

44 + buffer_shape = (layer.n_slices, x.shape[0], local_lora_rank)

45 + # Under torch.compile, the local-rank-1 fully-sharded path can otherwise

46 + # get lowered to a reinterpret view with a non-canonical layout. The

47 + # Triton shrink op mutates this buffer in place and expects the standard

48 + # contiguous [slice, token, rank] stride contract.

49 + buffers = torch.empty_strided(

50 + buffer_shape,

51 + (x.shape[0] * local_lora_rank, local_lora_rank, 1),

52 dtype=torch.float32,

53 device=x.device,

54 )

55 + buffers.zero_()

57 shrunk_buffers: torch.Tensor | None = layer.punica_wrapper.add_shrink(

58 buffers, x, layer.lora_a_stacked, 1.0

@@ -86,7 +94,7 @@ class ColumnParallelLinearWithLoRA(BaseLinearLayerWithLoRA):

94 # The base_layer type is ColumnParallelLinear or

95 # MergedColumnParallelLinear, their weight sharding logic is

96 # inconsistent when TP is greater than 1.

89 - self.is_merged_col_linear = type(base_layer) is MergedColumnParallelLinear

97 + self.is_merged_col_linear = isinstance(base_layer, MergedColumnParallelLinear)

98 self.output_size = self.base_layer.output_size_per_partition

99 # There is only one LoRA layer

100 self.n_slices = 1

@@ -158,7 +166,7 @@ class ColumnParallelLinearWithLoRA(BaseLinearLayerWithLoRA):

166 ) -> bool:

167 if type(source_layer) is maybe_get_oot_by_class(ColumnParallelLinear):

168 return True

161 - if type(source_layer) is maybe_get_oot_by_class(MergedColumnParallelLinear):

169 + if isinstance(source_layer, maybe_get_oot_by_class(MergedColumnParallelLinear)):

170 if len(packed_modules_list) != 1:

171 return False

172 # Exclude layers with 3+ output sizes - those are handled by

@@ -275,19 +283,41 @@ class MergedColumnParallelLinearWithLoRA(ColumnParallelLinearWithLoRA):

283 index, 0, : lora_b_i.shape[0], : lora_b_i.shape[1]

284 ].copy_(lora_b_i, non_blocking=True)

285

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

287 + merged_cls = maybe_get_oot_by_class(MergedColumnParallelLinear)

288 + # Effectively unsharded subclasses can safely reuse their custom

289 + # forward() implementation before applying the LoRA delta.

290 + if (

291 + self.tp_size == 1

292 + and type(self.base_layer) is not merged_cls

293 + and type(self.base_layer).forward is not merged_cls.forward

294 + ):

295 + return self._apply_base_forward(x)

296 + return _mcp_apply(x, bias, self)

297 +

298 @classmethod

279 - @_not_fully_sharded_can_replace

299 def can_replace_layer(

300 cls,

301 source_layer: nn.Module,

302 lora_config: LoRAConfig,

303 packed_modules_list: list,

304 model_config: PretrainedConfig | None = None,

305 + decorate: bool = True,

306 ) -> bool:

287 - return (

288 - type(source_layer) is MergedColumnParallelLinear

289 - and len(packed_modules_list) == 2

290 - )

307 + merged_cls = maybe_get_oot_by_class(MergedColumnParallelLinear)

308 + if not isinstance(source_layer, merged_cls) or len(packed_modules_list) != 2:

309 + return False

310 +

311 + tp_size = getattr(source_layer, "tp_size", 1)

312 + if type(source_layer) is merged_cls:

313 + if not decorate:

314 + return True

315 + return not lora_config.fully_sharded_loras or tp_size == 1

316 +

317 + # Only support effectively unsharded subclasses here. Sharded

318 + # subclasses may have custom communication semantics that the generic

319 + # merged-column LoRA path does not know how to preserve.

320 + return tp_size == 1

321

322

323 class QKVParallelLinearWithLoRA(ColumnParallelLinearWithLoRA):

@@ -607,7 +637,9 @@ class MergedColumnParallelLinearVariableSliceWithLoRA(

637 ) -> bool:

638 # Support MergedColumnParallelLinear with 3 or more slices

639 # (2 slices are handled by MergedColumnParallelLinearWithLoRA)

610 - if type(source_layer) is not maybe_get_oot_by_class(MergedColumnParallelLinear):

640 + if not isinstance(

641 + source_layer, maybe_get_oot_by_class(MergedColumnParallelLinear)

642 + ):

643 return False

644

645 # If packed_modules_list has 3+ items, use this class

vllm/lora/layers/replicated_linear.py

+7 -1

@@ -46,6 +46,12 @@ class ReplicatedLinearWithLoRA(BaseLinearLayerWithLoRA):

47 return output, output_bias

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

50 + # ReplicatedLinear subclasses such as GateLinear override forward() to

51 + # dispatch custom kernels and/or adjust the output dtype. Apply LoRA on

52 + # top of the actual base-layer output instead of bypassing that path.

53 + return self._apply_base_forward(x)

54 +

55 # ReplicatedLinear should always be replaced, regardless of the fully

56 # sharded LoRAs setting, because it is, by definition, copied per GPU.

57 @classmethod

@@ -56,7 +62,7 @@ class ReplicatedLinearWithLoRA(BaseLinearLayerWithLoRA):

62 packed_modules_list: list,

63 model_config: PretrainedConfig | None = None,

64 ) -> bool:

59 - return type(source_layer) is maybe_get_oot_by_class(ReplicatedLinear)

65 + return isinstance(source_layer, maybe_get_oot_by_class(ReplicatedLinear))

67 def slice_lora_a(

68 self, lora_a: torch.Tensor | list[torch.Tensor | None]

vllm/lora/model_manager.py

+52 -7

@@ -437,12 +437,21 @@ class LoRAModelManager:

437 ),

438 )

439

440 - # In some models, especially multimodal ones, layers with the same

441 - # name may have different types, such as nn.Linear and

442 - # ReplicatedLinear. The nn.Linear layers cannot be replaced with

443 - # LoRA layers, leading to assertion error. The following check

444 - # aims to prevent this error

445 - if self.supports_mm and not isinstance(new_module, BaseLayerWithLoRA):

440 + # Some matched modules can be unsupported by LoRA wrappers

441 + # (e.g. subclasses with specialized forward behavior).

442 + if not isinstance(new_module, BaseLayerWithLoRA):

443 + error_msg = (

444 + "LoRA target module "

445 + f"{module_name} ({type(module).__name__}) matched the "

446 + "deployment configuration but could not be wrapped by any "

447 + "LoRA layer implementation."

448 + )

449 + if self.lora_config.target_modules is not None:

450 + raise ValueError(

451 + f"{error_msg} target_modules="

452 + f"{sorted(self.lora_config.target_modules)}"

453 + )

454 + logger.warning_once("%s It will be ignored.", error_msg)

455 continue

456 self.register_module(module_name, new_module)

457

@@ -578,6 +587,38 @@ class LoRAModelManager:

587 model.loras[module_name] = lora

588 return model

589

590 + def get_dummy_lora_warmup_rank(self, default_rank: int) -> int:

591 + """Return a dummy LoRA rank compatible with wrapped modules.

592 +

593 + Dummy LoRAs keep warmup memory low by using a small rank. Fully

594 + sharded MoE wrappers additionally require the dummy rank to be divisible

595 + by tensor parallel size because they shard W13 along the rank axis.

596 + """

597 + if not self.lora_config.fully_sharded_loras:

598 + return default_rank

599 +

600 + required_multiple = 1

601 + for module in self.modules.values():

602 + if not getattr(module, "fully_sharded", False):

603 + continue

604 + required_multiple = math.lcm(required_multiple, module.tp_size)

605 +

606 + if required_multiple == 1 or default_rank % required_multiple == 0:

607 + return default_rank

608 +

609 + adjusted_rank = (

610 + (default_rank + required_multiple - 1) // required_multiple

611 + ) * required_multiple

612 + if adjusted_rank > self.lora_config.max_lora_rank:

613 + raise ValueError(

614 + "Unable to choose a dummy LoRA warmup rank compatible with "

615 + "fully sharded MoE modules: "

616 + f"default_rank={default_rank}, "

617 + f"required_multiple={required_multiple}, "

618 + f"max_lora_rank={self.lora_config.max_lora_rank}"

619 + )

620 + return adjusted_rank

621 +

622 def _match_target_modules(self, module_name: str) -> bool:

623 """Check if a module should have LoRA applied.

624

@@ -594,7 +635,11 @@ class LoRAModelManager:

635 """

636 if not is_supported_lora_module(module_name, self.supported_lora_modules):

637 return False

597 - return is_in_target_modules(module_name, self.lora_config.target_modules)

638 + return is_in_target_modules(

639 + module_name,

640 + self.lora_config.target_modules,

641 + self.packed_modules_mapping,

642 + )

643

644 def _get_punica_wrapper(self, module_name: str) -> PunicaWrapperBase | None:

645 """

vllm/lora/utils.py

+25 -3

@@ -73,7 +73,9 @@ def get_lora_id():

73 return _GLOBAL_LORA_ID

76 -_all_lora_classes: set[type[BaseLayerWithLoRA]] = {

76 +# Order matters here: more specific wrappers must be checked before generic

77 +# merged/column-parallel wrappers in from_layer().

78 +_all_lora_classes: tuple[type[BaseLayerWithLoRA], ...] = (

79 VocabParallelEmbeddingWithLoRA,

80 ColumnParallelLinearWithLoRA,

81 MergedColumnParallelLinearWithLoRA,

@@ -90,7 +92,7 @@ _all_lora_classes: set[type[BaseLayerWithLoRA]] = {

92 RowParallelLinearWithShardedLoRA,

93 FusedMoEWithLoRA,

94 FusedMoE3DWithLoRA,

93 -}

95 +)

98 def is_moe_model(model: nn.Module) -> bool:

@@ -258,6 +260,7 @@ def is_supported_lora_module(

260 def is_in_target_modules(

261 module_name: str,

262 target_modules: list[str] | None,

263 + packed_modules_mapping: dict[str, list[str]] | None = None,

264 ) -> bool:

265 """Check if a module passes the deployment-time target_modules filter.

266

@@ -268,14 +271,33 @@ def is_in_target_modules(

271 module_name: Full dot-separated module name.

272 target_modules: Optional deployment-time restriction list from

273 LoRAConfig.target_modules.

274 + packed_modules_mapping: Optional model-defined mapping from packed

275 + runtime module names to their adapter-visible submodule names

276 + (e.g. ``{"gate_up_proj": ["gate_proj", "up_proj"]}``).

277

278 Returns:

279 True if the module passes the filter, False otherwise.

280 """

281 if target_modules is None:

282 return True

283 + target_module_set = set(target_modules)

284 module_suffix = module_name.split(".")[-1]

278 - return module_suffix in set(target_modules)

285 + if module_suffix in target_module_set or module_name in target_module_set:

286 + return True

287 +

288 + if not packed_modules_mapping:

289 + return False

290 +

291 + # Runtime packed parent matched by deployment-time child targets.

292 + packed_children = packed_modules_mapping.get(module_suffix)

293 + if packed_children and any(child in target_module_set for child in packed_children):

294 + return True

295 +

296 + # Adapter-visible packed child matched by deployment-time parent target.

297 + return any(

298 + module_suffix in children and packed_parent in target_module_set

299 + for packed_parent, children in packed_modules_mapping.items()

300 + )

301

302

303 def get_adapter_absolute_path(lora_path: str) -> str:

vllm/lora/worker_manager.py

+8 -1

@@ -160,7 +160,11 @@ class WorkerLoRAManager:

160 lora_request.lora_path,

161 ", ".join(sorted(expected_lora_modules_lst)),

162 )

163 - elif not is_in_target_modules(module_name, target_modules):

163 + elif not is_in_target_modules(

164 + module_name,

165 + target_modules,

166 + packed_modules_mapping,

167 + ):

168 logger.warning_once(

169 "LoRA module '%s' in adapter '%s' is not in the "

170 "deployment-time target_modules restriction [%s]."

@@ -197,6 +201,9 @@ class WorkerLoRAManager:

201 self._cached_dummy_lora = dummy_lora

202 return self._adapter_manager.add_adapter(dummy_lora)

203

204 + def get_dummy_lora_warmup_rank(self, default_rank: int) -> int:

205 + return self._adapter_manager.get_dummy_lora_warmup_rank(default_rank)

206 +

207 def pin_adapter(self, adapter_id: int) -> bool:

208 return self._adapter_manager.pin_adapter(adapter_id)

209

vllm/model_executor/layers/fused_moe/oracle/unquantized.py

+13

@@ -214,6 +214,19 @@ def select_unquantized_moe_backend(

214 return backend, k_cls

215 raise ValueError(_make_log_unsupported(backend, reason))

216

217 + # LoRA needs Triton's unfused activation/reduction hooks. Selecting the

218 + # backend here ensures weights stay in a LoRA-compatible layout instead of

219 + # being permuted for a backend like FlashInfer or AITER during load.

220 + if moe_config.is_lora_enabled:

221 + backend = UnquantizedMoeBackend.TRITON

222 + if activation_format == mk.FusedMoEActivationFormat.BatchedExperts:

223 + backend = UnquantizedMoeBackend.BATCHED_TRITON

224 + return _return_or_raise(

225 + backend,

226 + moe_config,

227 + activation_format,

228 + )

229 +

230 runner_backend = moe_config.moe_backend

231 if runner_backend != "auto":

232 requested_backend = map_unquantized_backend(runner_backend)

vllm/model_executor/layers/quantization/utils/quant_utils.py

@@ -356,6 +356,12 @@ def get_and_maybe_dequant_weights(

356 from vllm.model_executor.layers.linear import UnquantizedLinearMethod

357 from vllm.model_executor.layers.quantization.fp8 import Fp8LinearMethod

358

359 + # LoRA linear wrappers store quantization metadata on `base_layer`.

360 + # Unwrap here so callers can pass either a raw linear layer or its LoRA

361 + # wrapper without special-casing.

362 + while hasattr(layer, "base_layer") and hasattr(layer.base_layer, "quant_method"):

363 + layer = layer.base_layer

364 +

365 weight = get_attribute_fallback(layer, ["weight", "qweight", "weight_packed"])

366

367 # Unquantized layer: just return base weights

vllm/v1/worker/lora_model_runner_mixin.py

+4 -1

@@ -101,9 +101,12 @@ class LoRAModelRunnerMixin:

101 assert self.lora_manager is not None, "LoRA is not enabled"

102

103 num_loras = lora_config.max_loras

104 - lora_warmup_rank = (

104 + lora_warmup_rank: int = (

105 lora_config.max_lora_rank if lora_config.max_lora_rank < 8 else 8

106 )

107 + lora_warmup_rank = self.lora_manager.get_dummy_lora_warmup_rank(

108 + lora_warmup_rank

109 + )

110 # Make dummy lora requests

111 lora_requests: set[LoRARequest] = {

112 LoRARequest(