[LoRA] Gate all_gather on fully_sharded_loras inside _mcp_apply; rewrite regression test (#45715)
Signed-off-by: lcheng <lcheng321@gatech.edu> Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai> Co-authored-by: Jee Jee Li <jeejeelee@inferact.ai> Co-authored-by: Jee Jee Li <pandaleefree@gmail.com>
lcheng committed
Jun 23, 2026 at 22:08 UTC
e51e700470cd04dbb25cb470571882dcbfed2a8a
1 file changed
+5
-5
vllm/lora/layers/column_parallel_linear.py
+5
-5
@@ -22,10 +22,10 @@ from .utils import _fully_sharded_can_replace, _not_fully_sharded_can_replace
22
23
24
def _mcp_apply(x, bias, layer: "ColumnParallelLinearWithLoRA"):
25
- """
26
- For `ColumnParallelLinearWithLoRA` or classes that inherit from
27
- `ColumnParallelLinearWithLoRA`, they share the same `apply` logic.
28
- """
25
+ """Fully-sharded (S-LoRA) apply path for column-parallel LoRA layers."""
26
+ assert layer.lora_config.fully_sharded_loras, (
27
+ "_mcp_apply is only used for fully sharded LoRA"
28
+ )
29
assert (
30
layer.n_slices
31
== len(layer.lora_a_stacked)
@@ -341,7 +341,7 @@ class MergedColumnParallelLinearWithLoRA(ColumnParallelLinearWithLoRA):
341
and base_forward is not merged_forward
342
):
343
return self._apply_base_forward(x)
344
- return _mcp_apply(x, bias, self)
344
+ return super().apply(x, bias)
345
346
@classmethod
347
def can_replace_layer(