[diffusion] Per-section LoRA adapters on fused linear layers (#34933)

python/sglang/multimodal_gen/configs/models/dits/cosmos3video.py

@@ -34,6 +34,8 @@ def _build_cosmos3_param_names_mapping(gated_mlp: bool = True) -> dict:
     unchanged.
     """
     mapping = {
+        # PEFT saves adapter keys as <module>.lora_{A,B}.default.weight.
+        r"^(.*\.lora_[AB])\.default$": r"\1",
         # Inherited from text pretraining; unused at diffusion inference.
         r"^lm_head\.weight$": "",
         r"^norm\.weight$": "",

python/sglang/multimodal_gen/runtime/pipelines/cosmos3_pipeline.py

@@ -11,6 +11,7 @@ import os
 from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import (
     ComposedPipelineBase,
 )
+from sglang.multimodal_gen.runtime.pipelines_core.lora_pipeline import LoRAPipeline
 from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.cosmos3 import (
     Cosmos3DecodingStage,
     Cosmos3DenoisingStage,
@@ -25,7 +26,7 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 logger = init_logger(__name__)
 
 
-class Cosmos3Pipeline(ComposedPipelineBase):
+class Cosmos3Pipeline(LoRAPipeline, ComposedPipelineBase):
     """Cosmos3 diffusion pipeline shared by T2V, I2V, and T2I.
 
     Text is tokenized and embedded directly inside the transformer; there is

python/sglang/multimodal_gen/runtime/pipelines_core/lora_pipeline.py

@@ -55,6 +55,60 @@ def _swap_peft_swiglu_fc1_lora_b(
     return torch.cat([gate, value], dim=0)
 
 
+def stack_or_compose_fused_lora(
+    a_list: list[torch.Tensor],
+    b_list: list[torch.Tensor],
+    adapter_alpha: int | None,
+) -> tuple[torch.Tensor, torch.Tensor, int | None]:
+    """Equal sections stay 3D-stacked; GQA sections become one 2D pair."""
+    if len({t.shape for t in a_list}) == 1 and len({t.shape for t in b_list}) == 1:
+        return torch.stack(a_list), torch.stack(b_list), None
+    ranks = [a.shape[0] for a in a_list]
+    outs = [b.shape[0] for b in b_list]
+    a_2d = torch.cat(a_list, dim=0)
+    b_2d = b_list[0].new_zeros(sum(outs), sum(ranks))
+    row = col = 0
+    for a, b, rank, out in zip(a_list, b_list, ranks, outs):
+        scale = 1.0 if adapter_alpha is None else adapter_alpha / rank
+        b_2d[row : row + out, col : col + rank] = (
+            b if scale == 1.0 else (b.float() * scale).to(dtype=b.dtype)
+        )
+        row += out
+        col += rank
+    return a_2d, b_2d, a_2d.shape[0]
+
+
+def _store_fused_lora_groups(
+    adapter: dict[str, torch.Tensor],
+    to_merge_params: dict[Hashable, dict[Any, Any]],
+    adapter_alpha: int | None,
+    device: torch.device | str,
+) -> None:
+    """Write deferred fused lora_A/B groups into the adapter dict."""
+    for a_key, a_parts in list(to_merge_params.items()):
+        if not str(a_key).endswith(".lora_A"):
+            continue
+        base = str(a_key)[: -len(".lora_A")]
+        b_key = f"{base}.lora_B"
+        b_parts = to_merge_params.get(b_key)
+        n = max(a_parts) + 1
+        if (
+            b_parts is None
+            or set(a_parts) != set(range(n))
+            or set(b_parts) != set(range(n))
+        ):
+            continue
+        a, b, fused_alpha = stack_or_compose_fused_lora(
+            [a_parts[i] for i in range(n)],
+            [b_parts[i] for i in range(n)],
+            adapter_alpha,
+        )
+        adapter[str(a_key)] = a.to(device)
+        adapter[b_key] = b.to(device)
+        if fused_alpha is not None:
+            adapter[f"{base}.alpha"] = torch.tensor(float(fused_alpha), device=device)
+
+
 class LoRAPipeline(ComposedPipelineBase):
     """
     Pipeline that supports injecting LoRA adapters into the diffusion transformer.
@@ -780,6 +834,9 @@ class LoRAPipeline(ComposedPipelineBase):
             # see param mapping in HunyuanVideoArchConfig
             if merge_index is not None:
                 to_merge_params[target_name][merge_index] = weight
+                # A/B of one fused layer must be laid out together (GQA B cannot stack).
+                if target_name.endswith((".lora_A", ".lora_B")):
+                    continue
                 if len(to_merge_params[target_name]) == num_params_to_merge:
                     sorted_tensors = [
                         to_merge_params[target_name][i]
@@ -798,6 +855,13 @@ class LoRAPipeline(ComposedPipelineBase):
                 )
             self.lora_adapters[lora_nickname][target_name] = weight.to(self.device)
 
+        _store_fused_lora_groups(
+            self.lora_adapters[lora_nickname],
+            to_merge_params,
+            adapter_lora_alpha,
+            self.device,
+        )
+
         self.loaded_adapter_paths[lora_nickname] = lora_path
         self.loaded_adapter_alphas[lora_nickname] = adapter_lora_alpha
         logger.info("Rank %d: loaded LoRA adapter %s", rank, lora_path)

python/sglang/multimodal_gen/runtime/post_training/weights_updater.py

@@ -58,7 +58,10 @@ from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload im
     is_layerwise_offloaded_module,
 )
 from sglang.multimodal_gen.runtime.pipelines.diffusers_pipeline import DiffusersPipeline
-from sglang.multimodal_gen.runtime.pipelines_core.lora_pipeline import LoRAPipeline
+from sglang.multimodal_gen.runtime.pipelines_core.lora_pipeline import (
+    LoRAPipeline,
+    stack_or_compose_fused_lora,
+)
 from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import maybe_download_model
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 from sglang.srt.weight_sync.tensor_bucket import (
@@ -215,11 +218,14 @@ def _build_module_weight_name_mapper(module: torch.nn.Module):
     if not mapping_fns:
         return None
 
-    def map_name(name: str) -> str:
+    def map_name(name: str) -> tuple[str, Any]:
         mapped_name = name
+        merge_index = None
         for mapping_fn in mapping_fns:
-            mapped_name = mapping_fn(mapped_name)[0]
-        return mapped_name
+            mapped_name, index, _ = mapping_fn(mapped_name)
+            if index is not None:
+                merge_index = index
+        return mapped_name, merge_index
 
     return map_name
 
@@ -236,23 +242,24 @@ def _resolve_lora_ipc_layer_dict_key(
     layer_prefix: str,
     layer_dict: dict,
     module: torch.nn.Module,
-) -> tuple[Any | None, str]:
+) -> tuple[Any | None, str, int | None]:
     """Map training-side LoRA layer prefix to lora_layers key (Layer 2)."""
     layer = layer_dict.get(layer_prefix)
     if layer is not None:
-        return layer, layer_prefix
+        return layer, layer_prefix, None
 
     map_name = _build_module_weight_name_mapper(module)
     if map_name is None:
-        return None, layer_prefix
+        return None, layer_prefix, None
 
-    mapped = _strip_param_weight_suffix(map_name(f"{layer_prefix}.weight"))
+    mapped_name, merge_index = map_name(f"{layer_prefix}.weight")
+    mapped = _strip_param_weight_suffix(mapped_name)
     if mapped != layer_prefix:
         layer = layer_dict.get(mapped)
         if layer is not None:
-            return layer, mapped
+            return layer, mapped, merge_index
 
-    return None, layer_prefix
+    return None, layer_prefix, None
 
 
 def _iter_module_weight_updates(
@@ -268,7 +275,7 @@ def _iter_module_weight_updates(
             yield name, loaded_weight
             continue
 
-        mapped_name = map_name(name) if map_name is not None else name
+        mapped_name = map_name(name)[0] if map_name is not None else name
         if mapped_name in model_params:
             yield mapped_name, loaded_weight
             continue
@@ -628,35 +635,72 @@ class WeightsUpdater:
         updated = 0
         skipped = 0
         unknown_layers: list[str] = []
-        with lora_pipeline._temporarily_disable_offload(target=target_module):
-            for layer_name, (lora_a, lora_b) in pairs.items():
-                layer, _resolved_key = _resolve_lora_ipc_layer_dict_key(
-                    layer_name, layer_dict, dit_module
+        # Honor lora_merge_mode: merged and unmerged evaluation differ bitwise.
+        merge_mode = lora_pipeline._resolve_lora_merge_mode(None, None)
+        merge_weights = lora_pipeline._should_merge_lora_for_layers(
+            target_module, layer_dict, merge_mode
+        )
+        plain_pairs: list[tuple[torch.Tensor, torch.Tensor, Any]] = []
+        fused_sections: dict[Any, dict[int, tuple[torch.Tensor, torch.Tensor]]] = {}
+        for layer_name, (lora_a, lora_b) in pairs.items():
+            layer, _resolved_key, merge_index = _resolve_lora_ipc_layer_dict_key(
+                layer_name, layer_dict, dit_module
+            )
+            if layer is None:
+                logger.warning(
+                    "Unknown LoRA layer name %s for target %s; skipping",
+                    layer_name,
+                    target_module,
                 )
-                if layer is None:
-                    logger.warning(
-                        "Unknown LoRA layer name %s for target %s; skipping",
-                        layer_name,
-                        target_module,
-                    )
-                    unknown_layers.append(layer_name)
-                    skipped += 1
-                    continue
-                inferred_rank = int(lora_a.shape[0])
-                alpha = lora_alpha if lora_alpha is not None else inferred_rank
-                if lora_rank is not None and lora_rank != inferred_rank:
-                    logger.warning(
-                        "LoRA rank mismatch for %s: payload=%d request=%d; using payload rank",
-                        layer_name,
-                        inferred_rank,
-                        lora_rank,
-                    )
+                unknown_layers.append(layer_name)
+                skipped += 1
+                continue
+            inferred_rank = int(lora_a.shape[-2])
+            if lora_rank is not None and lora_rank != inferred_rank:
+                logger.warning(
+                    "LoRA rank mismatch for %s: payload=%d request=%d; using payload rank",
+                    layer_name,
+                    inferred_rank,
+                    lora_rank,
+                )
+            if merge_index is None:
+                plain_pairs.append((lora_a, lora_b, layer))
+            else:
+                # Per-section pairs of one fused layer compose into one adapter.
+                fused_sections.setdefault(layer, {})[merge_index] = (lora_a, lora_b)
+
+        with lora_pipeline._temporarily_disable_offload(target=target_module):
+            for lora_a, lora_b, layer in plain_pairs:
+                inferred_rank = int(lora_a.shape[-2])
                 layer.lora_rank = inferred_rank
-                layer.lora_alpha = alpha
+                layer.lora_alpha = (
+                    lora_alpha if lora_alpha is not None else inferred_rank
+                )
                 layer.set_lora_weights(
-                    lora_a, lora_b, merge_weights=True, clear_existing=True
+                    lora_a, lora_b, merge_weights=merge_weights, clear_existing=True
                 )
                 updated += 1
+            for layer, layer_sections in fused_sections.items():
+                indices = sorted(layer_sections)
+                lora_a, lora_b, fused_alpha = stack_or_compose_fused_lora(
+                    [layer_sections[i][0] for i in indices],
+                    [layer_sections[i][1] for i in indices],
+                    lora_alpha,
+                )
+                if fused_alpha is not None:
+                    # Composed pairs fold the scale; alpha == rank keeps it neutral.
+                    layer.lora_rank = fused_alpha
+                    layer.lora_alpha = fused_alpha
+                else:
+                    inferred_rank = int(lora_a.shape[-2])
+                    layer.lora_rank = inferred_rank
+                    layer.lora_alpha = (
+                        lora_alpha if lora_alpha is not None else inferred_rank
+                    )
+                layer.set_lora_weights(
+                    lora_a, lora_b, merge_weights=merge_weights, clear_existing=True
+                )
+                updated += len(layer_sections)
 
         gc.collect()
         torch.cuda.empty_cache()

python/sglang/multimodal_gen/test/unit/test_fused_lora_compose.py

@@ -0,0 +1,373 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Fused-layer LoRA groups: stack-vs-compose decision, loader wiring, IPC resolution."""
+
+from collections import defaultdict
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import torch
+import torch.nn.functional as F
+
+from sglang.multimodal_gen.runtime.layers.lora.linear import (
+    MergedColumnParallelLinearWithLoRA,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.lora_pipeline import (
+    LoRAPipeline,
+    stack_or_compose_fused_lora,
+)
+from sglang.multimodal_gen.runtime.post_training.weights_updater import (
+    _resolve_lora_ipc_layer_dict_key,
+)
+
+_TP_RANK_PATCH = "sglang.multimodal_gen.runtime.layers.lora.linear.get_tp_rank"
+_LOCAL_DEVICE_PATCH = (
+    "sglang.multimodal_gen.runtime.layers.lora.linear.get_local_torch_device"
+)
+
+# GQA-like fused layout: unequal q/k/v sections.
+OUTPUT_SIZES = [8, 2, 2]
+IN_DIM = 4
+
+
+def _make_ab_lists(ranks: list[int]):
+    a_list, b_list = [], []
+    for index, rank in enumerate(ranks):
+        torch.manual_seed(100 + index * 10 + rank)
+        a_list.append(torch.randn(rank, IN_DIM))
+        b_list.append(torch.randn(OUTPUT_SIZES[index], rank))
+    return a_list, b_list
+
+
+def _reference_delta(x, a_list, b_list, adapter_alpha):
+    out = torch.zeros(*x.shape[:-1], sum(OUTPUT_SIZES))
+    row = 0
+    for a, b in zip(a_list, b_list):
+        scale = 1.0 if adapter_alpha is None else adapter_alpha / a.shape[0]
+        out[..., row : row + b.shape[0]] = (x @ a.T @ b.T) * scale
+        row += b.shape[0]
+    return out
+
+
+def test_compose_unequal_sections_matches_reference():
+    a_list, b_list = _make_ab_lists([2, 3, 1])
+    a_2d, b_2d, fused_alpha = stack_or_compose_fused_lora(a_list, b_list, 4)
+    assert a_2d.shape == (6, IN_DIM)
+    assert b_2d.shape == (sum(OUTPUT_SIZES), 6)
+    assert fused_alpha == 6
+
+    x = torch.randn(5, IN_DIM)
+    torch.testing.assert_close(
+        x @ a_2d.T @ b_2d.T,
+        _reference_delta(x, a_list, b_list, 4),
+        rtol=1e-5,
+        atol=1e-5,
+    )
+
+
+def test_stack_kept_for_equal_sections():
+    torch.manual_seed(0)
+    a_list = [torch.randn(2, IN_DIM) for _ in range(2)]
+    b_list = [torch.randn(4, 2) for _ in range(2)]
+    a, b, fused_alpha = stack_or_compose_fused_lora(a_list, b_list, 4)
+    assert a.shape == (2, 2, IN_DIM)
+    assert b.shape == (2, 4, 2)
+    assert fused_alpha is None
+    torch.testing.assert_close(a[0], a_list[0])
+    torch.testing.assert_close(b[1], b_list[1])
+
+
+class _FakeMergedLinear(torch.nn.Module):
+    def __init__(
+        self,
+        output_sizes: list[int],
+        in_dim: int,
+        weight: torch.Tensor | None = None,
+        output_partition_sizes: list[int] | None = None,
+    ):
+        super().__init__()
+        self.output_sizes = output_sizes
+        self.output_partition_sizes = (
+            output_partition_sizes
+            if output_partition_sizes is not None
+            else output_sizes
+        )
+        if weight is None:
+            weight = torch.randn(sum(output_sizes), in_dim)
+        self.weight = torch.nn.Parameter(weight)
+        self.bias = None
+        self.skip_bias_add = False
+        self.gather_output = False
+        self.quant_method = SimpleNamespace(
+            apply=lambda layer, x, bias=None: F.linear(x, layer.weight, bias)
+        )
+
+
+def _make_layer() -> MergedColumnParallelLinearWithLoRA:
+    torch.manual_seed(0)
+    return MergedColumnParallelLinearWithLoRA(_FakeMergedLinear(OUTPUT_SIZES, IN_DIM))
+
+
+def _set_composed(layer, a_list, b_list, adapter_alpha, merge_weights=False):
+    a_2d, b_2d, fused_alpha = stack_or_compose_fused_lora(a_list, b_list, adapter_alpha)
+    layer.lora_rank = fused_alpha
+    layer.lora_alpha = fused_alpha
+    layer.set_lora_weights(a_2d, b_2d, merge_weights=merge_weights)
+
+
+def test_unmerged_forward_applies_composed_pair():
+    layer = _make_layer()
+    base_weight = layer.base_layer.weight.detach().clone()
+    a_list, b_list = _make_ab_lists([2, 3, 1])
+    _set_composed(layer, a_list, b_list, 4)
+
+    assert not layer.merged
+    assert not layer.disable_lora
+
+    x = torch.randn(5, IN_DIM)
+    with patch(_TP_RANK_PATCH, return_value=0):
+        out, _ = layer.forward(x)
+    expected = x @ base_weight.T + _reference_delta(x, a_list, b_list, 4)
+    torch.testing.assert_close(out, expected, rtol=1e-5, atol=1e-5)
+
+
+def test_merge_writes_section_rows_and_unmerge_restores():
+    layer = _make_layer()
+    base_weight = layer.base_layer.weight.detach().clone()
+    a_list, b_list = _make_ab_lists([2, 2, 2])
+    with (
+        patch(_TP_RANK_PATCH, return_value=0),
+        patch(_LOCAL_DEVICE_PATCH, return_value=torch.device("cpu")),
+    ):
+        _set_composed(layer, a_list, b_list, 4, merge_weights=True)
+
+    assert layer.merged
+    merged = layer.base_layer.weight.detach().cpu()
+    row = 0
+    for a, b in zip(a_list, b_list):
+        expected_rows = base_weight[row : row + b.shape[0]] + (4 / a.shape[0]) * (b @ a)
+        torch.testing.assert_close(
+            merged[row : row + b.shape[0]], expected_rows, rtol=1e-5, atol=1e-5
+        )
+        row += b.shape[0]
+
+    with patch(_TP_RANK_PATCH, return_value=0):
+        layer.unmerge_lora_weights()
+    torch.testing.assert_close(layer.base_layer.weight.detach().cpu(), base_weight)
+
+
+def test_composed_pair_shards_correctly_under_mock_tp():
+    tp_size = 2
+    a_list, b_list = _make_ab_lists([2, 3, 1])
+    a_2d, b_2d, fused_alpha = stack_or_compose_fused_lora(a_list, b_list, 4)
+    torch.manual_seed(0)
+    full_weight = torch.randn(sum(OUTPUT_SIZES), IN_DIM)
+    x = torch.randn(5, IN_DIM)
+    expected = x @ full_weight.T + _reference_delta(x, a_list, b_list, 4)
+
+    part_sizes = [size // tp_size for size in OUTPUT_SIZES]
+    rank_outputs = []
+    for tp_rank in range(tp_size):
+        local_rows = []
+        for index, (size, part) in enumerate(zip(OUTPUT_SIZES, part_sizes)):
+            row = sum(OUTPUT_SIZES[:index]) + tp_rank * part
+            local_rows.append(full_weight[row : row + part])
+        layer = MergedColumnParallelLinearWithLoRA(
+            _FakeMergedLinear(
+                OUTPUT_SIZES,
+                IN_DIM,
+                weight=torch.cat(local_rows, dim=0),
+                output_partition_sizes=part_sizes,
+            )
+        )
+        layer.lora_rank = fused_alpha
+        layer.lora_alpha = fused_alpha
+        layer.set_lora_weights(a_2d, b_2d, merge_weights=False)
+        with patch(_TP_RANK_PATCH, return_value=tp_rank):
+            out, _ = layer.forward(x)
+        rank_outputs.append(out)
+
+    # All-gather equivalent: stitch each section's rank slices back together.
+    stitched = []
+    for index, part in enumerate(part_sizes):
+        local_col = sum(part_sizes[:index])
+        stitched.append(
+            torch.cat(
+                [out[..., local_col : local_col + part] for out in rank_outputs],
+                dim=-1,
+            )
+        )
+    torch.testing.assert_close(
+        torch.cat(stitched, dim=-1), expected, rtol=1e-5, atol=1e-5
+    )
+
+
+def test_ipc_key_resolution_returns_fused_merge_index():
+    module = torch.nn.Module()
+    module.param_names_mapping = {
+        r"^attn\.q\.(.*)$": (r"attn.to_qkv.\1", 0, 3),
+        r"^attn\.k\.(.*)$": (r"attn.to_qkv.\1", 1, 3),
+        r"^attn\.v\.(.*)$": (r"attn.to_qkv.\1", 2, 3),
+        r"^proj\.(.*)$": r"out_proj.\1",
+    }
+    sentinel = object()
+    layer_dict = {"attn.to_qkv": sentinel, "out_proj": sentinel}
+
+    layer, key, merge_index = _resolve_lora_ipc_layer_dict_key(
+        "attn.k", layer_dict, module
+    )
+    assert layer is sentinel
+    assert key == "attn.to_qkv"
+    assert merge_index == 1
+
+    layer, key, merge_index = _resolve_lora_ipc_layer_dict_key(
+        "proj", layer_dict, module
+    )
+    assert layer is sentinel
+    assert key == "out_proj"
+    assert merge_index is None
+
+    layer, key, merge_index = _resolve_lora_ipc_layer_dict_key(
+        "attn.to_qkv", layer_dict, module
+    )
+    assert layer is sentinel
+    assert merge_index is None
+
+
+class _TestLoRAPipeline(LoRAPipeline):
+    def create_pipeline_stages(self, server_args):
+        return None
+
+
+_QKV_MAPPING = {
+    r"^attn\.q\.(.*)$": (r"attn.to_qkv.\1", 0, 3),
+    r"^attn\.k\.(.*)$": (r"attn.to_qkv.\1", 1, 3),
+    r"^attn\.v\.(.*)$": (r"attn.to_qkv.\1", 2, 3),
+    r"^mlp\.gate\.(.*)$": (r"mlp.gate_up.\1", 0, 2),
+    r"^mlp\.up\.(.*)$": (r"mlp.gate_up.\1", 1, 2),
+}
+
+
+def _make_loader_pipeline() -> _TestLoRAPipeline:
+    pipeline = object.__new__(_TestLoRAPipeline)
+    pipeline.lora_adapters = defaultdict(dict)
+    pipeline.loaded_adapter_paths = {}
+    pipeline.loaded_adapter_alphas = {}
+    pipeline.device = "cpu"
+    pipeline.modules = {"transformer": torch.nn.Module()}
+    pipeline.server_args = SimpleNamespace(
+        lora_path=None,
+        lora_weight_name=None,
+        pipeline_config=SimpleNamespace(
+            dit_config=SimpleNamespace(
+                arch_config=SimpleNamespace(
+                    param_names_mapping=_QKV_MAPPING,
+                    lora_param_names_mapping={
+                        r"^(.*\.lora_[AB])\.default$": r"\1",
+                    },
+                )
+            )
+        ),
+    )
+    return pipeline
+
+
+def _load_adapter(pipeline, state_dict, lora_alpha=None):
+    loader_mod = "sglang.multimodal_gen.runtime.pipelines_core.lora_pipeline"
+    with (
+        patch(f"{loader_mod}.maybe_download_lora", return_value="/adapter"),
+        patch(f"{loader_mod}.load_file", return_value=state_dict),
+        patch(f"{loader_mod}.dist.is_initialized", return_value=False),
+    ):
+        pipeline.load_lora_adapter("/adapter", "adapter", rank=0, lora_alpha=lora_alpha)
+
+
+def _qkv_state_dict(rank=2):
+    torch.manual_seed(3)
+    state = {}
+    for name, rows in (("q", 8), ("k", 2), ("v", 2)):
+        state[f"attn.{name}.lora_A.default.weight"] = torch.randn(rank, IN_DIM)
+        state[f"attn.{name}.lora_B.default.weight"] = torch.randn(rows, rank)
+    return state
+
+
+def test_loader_composes_unequal_fused_sections():
+    pipeline = _make_loader_pipeline()
+    _load_adapter(pipeline, _qkv_state_dict(), lora_alpha=4)
+
+    adapter = pipeline.lora_adapters["adapter"]
+    assert adapter["attn.to_qkv.lora_A"].shape == (6, IN_DIM)
+    assert adapter["attn.to_qkv.lora_B"].shape == (sum(OUTPUT_SIZES), 6)
+    assert adapter["attn.to_qkv.alpha"].item() == 6.0
+    assert set(adapter) == {
+        "attn.to_qkv.lora_A",
+        "attn.to_qkv.lora_B",
+        "attn.to_qkv.alpha",
+    }
+
+
+def test_loader_stacks_equal_fused_sections():
+    pipeline = _make_loader_pipeline()
+    state_dict = {
+        "mlp.gate.lora_A.weight": torch.randn(2, IN_DIM),
+        "mlp.gate.lora_B.weight": torch.randn(4, 2),
+        "mlp.up.lora_A.weight": torch.randn(2, IN_DIM),
+        "mlp.up.lora_B.weight": torch.randn(4, 2),
+    }
+    _load_adapter(pipeline, state_dict)
+
+    adapter = pipeline.lora_adapters["adapter"]
+    assert adapter["mlp.gate_up.lora_A"].shape == (2, 2, IN_DIM)
+    assert adapter["mlp.gate_up.lora_B"].shape == (2, 4, 2)
+    assert "mlp.gate_up.alpha" not in adapter
+
+
+def test_loader_drops_incomplete_fused_groups():
+    pipeline = _make_loader_pipeline()
+    state_dict = _qkv_state_dict()
+    for key in list(state_dict):
+        if ".k." in key:
+            del state_dict[key]
+    _load_adapter(pipeline, state_dict)
+
+    assert pipeline.lora_adapters["adapter"] == {}
+
+
+def test_apply_composed_adapter_end_to_end():
+    layer = _make_layer()
+    base_weight = layer.base_layer.weight.detach().clone()
+    pipeline = _make_loader_pipeline()
+
+    rank = 2
+    adapter_alpha = 4
+    state_dict = _qkv_state_dict(rank=rank)
+    _load_adapter(pipeline, state_dict, lora_alpha=adapter_alpha)
+
+    strength = 2.0
+    with patch(
+        "sglang.multimodal_gen.runtime.pipelines_core.lora_pipeline.dist.get_rank",
+        return_value=0,
+    ):
+        applied = pipeline._apply_lora_to_layers(
+            {"attn.to_qkv": layer},
+            ["adapter"],
+            ["/adapter"],
+            rank=0,
+            strengths=[strength],
+            merge_weights=False,
+        )
+
+    assert applied == 1
+    a_list = [
+        state_dict[f"attn.{name}.lora_A.default.weight"] for name in ("q", "k", "v")
+    ]
+    b_list = [
+        state_dict[f"attn.{name}.lora_B.default.weight"] for name in ("q", "k", "v")
+    ]
+    x = torch.randn(3, IN_DIM)
+    with patch(_TP_RANK_PATCH, return_value=0):
+        out, _ = layer.forward(x)
+    expected = (
+        x @ base_weight.T
+        + _reference_delta(x, a_list, b_list, adapter_alpha) * strength
+    )
+    torch.testing.assert_close(out, expected, rtol=1e-5, atol=1e-5)