[Fix] lora: Load fused-module members into their own buffer slots

python/sglang/srt/lora/layers.py

@@ -1,4 +1,4 @@
-from typing import Dict, Optional, Union
+from typing import Dict, Optional, Tuple, Union
 
 import torch
 import torch.nn.functional as F
@@ -630,18 +630,32 @@ def slice_lora_a_weights(self, A: torch.Tensor):
         return A
 
     def slice_lora_b_weights(self, B: torch.Tensor):
-        local_tp_rank = self.base_layer.tp_rank
-        partition_sizes = self.base_layer.output_partition_sizes
-        output_sizes = self.base_layer.output_sizes
         slices = []
         offset = 0
-        for full_size, part_size in zip(output_sizes, partition_sizes):
-            start_idx = local_tp_rank * part_size
-            end_idx = start_idx + part_size
-            slices.append(B[offset + start_idx : offset + end_idx, :])
+        for member, full_size in enumerate(self.base_layer.output_sizes):
+            slices.append(
+                self.slice_lora_b_member(member, B[offset : offset + full_size, :])
+            )
             offset += full_size
         return torch.concat(slices, dim=0)
 
+    def lora_b_member_full_rows(self) -> Tuple[int, ...]:
+        """Adapter-side lora_B rows of each member (partition), in slot order."""
+        return tuple(self.base_layer.output_sizes)
+
+    def lora_b_member_local_offsets(self) -> Tuple[int, ...]:
+        """Where each member's rows start in this rank's lora_B."""
+        offsets = [0]
+        for part_size in self.base_layer.output_partition_sizes[:-1]:
+            offsets.append(offsets[-1] + part_size)
+        return tuple(offsets)
+
+    def slice_lora_b_member(self, member: int, B: torch.Tensor) -> torch.Tensor:
+        """This rank's rows of one member's full lora_B."""
+        part_size = self.base_layer.output_partition_sizes[member]
+        start_idx = self.base_layer.tp_rank * part_size
+        return B[start_idx : start_idx + part_size, :]
+
 
 class InklingQKVRLinearWithLoRA(MergedColumnParallelLinearWithLoRA):
     """LoRA wrapper for Inkling's fused q/k/v/r projection.
@@ -654,31 +668,43 @@ class InklingQKVRLinearWithLoRA(MergedColumnParallelLinearWithLoRA):
     """
 
     def slice_lora_b_weights(self, B: torch.Tensor):
+        members = torch.split(B, self.lora_b_member_full_rows(), dim=0)
+        return torch.concat(
+            [
+                self.slice_lora_b_member(i, member_b)
+                for i, member_b in enumerate(members)
+            ],
+            dim=0,
+        )
+
+    def lora_b_member_full_rows(self) -> Tuple[int, ...]:
+        bl = self.base_layer
+        q_size = bl.inkling_head_dim * bl.inkling_num_heads
+        kv_size = bl.inkling_head_dim * bl.inkling_num_kv_heads
+        return q_size, kv_size, kv_size, bl.inkling_d_rel * bl.inkling_num_heads
+
+    def _lora_b_member_local_rows(self) -> Tuple[int, ...]:
         bl = self.base_layer
-        tp_rank = bl.tp_rank
-        hd, nkv, nh, dr, tp = (
-            bl.inkling_head_dim,
-            bl.inkling_num_kv_heads,
-            bl.inkling_num_heads,
-            bl.inkling_d_rel,
-            bl.inkling_tp_size,
+        q_size, kv_size, _, r_size = self.lora_b_member_full_rows()
+        tp = bl.inkling_tp_size
+        kv_local = (
+            bl.inkling_head_dim if tp > bl.inkling_num_kv_heads else kv_size // tp
         )
-        q_size, kv_size, r_size = hd * nh, hd * nkv, dr * nh
-        q_off, k_off, v_off, r_off = 0, q_size, q_size + kv_size, q_size + 2 * kv_size
-        q_per, r_per = q_size // tp, r_size // tp
-        q = B[q_off + tp_rank * q_per : q_off + (tp_rank + 1) * q_per, :]
-        r = B[r_off + tp_rank * r_per : r_off + (tp_rank + 1) * r_per, :]
-        if tp > nkv:
+        return q_size // tp, kv_local, kv_local, r_size // tp
+
+    def lora_b_member_local_offsets(self) -> Tuple[int, ...]:
+        q_local, kv_local, _, _ = self._lora_b_member_local_rows()
+        return 0, q_local, q_local + kv_local, q_local + 2 * kv_local
+
+    def slice_lora_b_member(self, member: int, B: torch.Tensor) -> torch.Tensor:
+        bl = self.base_layer
+        local_rows = self._lora_b_member_local_rows()[member]
+        if member in (1, 2) and bl.inkling_tp_size > bl.inkling_num_kv_heads:
             # Replicate: each rank takes the single kv-head it shares (mirrors #15).
-            replicas = tp // nkv
-            kv_idx = tp_rank // replicas
-            k = B[k_off + kv_idx * hd : k_off + (kv_idx + 1) * hd, :]
-            v = B[v_off + kv_idx * hd : v_off + (kv_idx + 1) * hd, :]
+            shard = bl.tp_rank // (bl.inkling_tp_size // bl.inkling_num_kv_heads)
         else:
-            kv_per = kv_size // tp
-            k = B[k_off + tp_rank * kv_per : k_off + (tp_rank + 1) * kv_per, :]
-            v = B[v_off + tp_rank * kv_per : v_off + (tp_rank + 1) * kv_per, :]
-        return torch.concat([q, k, v, r], dim=0)
+            shard = bl.tp_rank
+        return B[shard * local_rows : (shard + 1) * local_rows, :]
 
 
 class QKVParallelLinearWithLoRA(ColumnParallelLinearWithLoRA):
@@ -741,34 +767,37 @@ def slice_lora_a_weights(self, A: torch.Tensor):
         return A
 
     def slice_lora_b_weights(self, B: torch.Tensor) -> torch.Tensor:
-        base_layer = self.base_layer
-        q_proj_shard_size = base_layer.q_proj_shard_size
-        kv_proj_shard_size = base_layer.kv_proj_shard_size
-        num_kv_head_replicas = base_layer.num_kv_head_replicas
-        local_tp_rank = base_layer.tp_rank
-
-        q_start_idx = q_proj_shard_size * local_tp_rank
-        q_end_idx = q_start_idx + q_proj_shard_size
-
-        kv_shard_id = local_tp_rank // num_kv_head_replicas
-        kv_start_idx = kv_proj_shard_size * kv_shard_id
-        kv_end_idx = kv_start_idx + kv_proj_shard_size
-
-        q_size = base_layer.output_sizes[0]
-        k_size = base_layer.output_sizes[1] // num_kv_head_replicas
-        B_q_shard = B[q_start_idx:q_end_idx, :]
-        B_k_shard = B[q_size + kv_start_idx : q_size + kv_end_idx, :]
-        B_v_shard = B[q_size + k_size + kv_start_idx : q_size + k_size + kv_end_idx, :]
-
+        members = torch.split(B, self.lora_b_member_full_rows(), dim=0)
         return torch.concat(
-            (
-                B_q_shard,
-                B_k_shard,
-                B_v_shard,
-            ),
+            [
+                self.slice_lora_b_member(i, member_b)
+                for i, member_b in enumerate(members)
+            ],
             dim=0,
         )
 
+    def lora_b_member_full_rows(self) -> Tuple[int, ...]:
+        """Adapter-side q/k/v lora_B rows; output_sizes counts replicated kv heads."""
+        base_layer = self.base_layer
+        q_size = base_layer.output_sizes[0]
+        kv_size = base_layer.output_sizes[1] // base_layer.num_kv_head_replicas
+        return q_size, kv_size, kv_size
+
+    def lora_b_member_local_offsets(self) -> Tuple[int, ...]:
+        q_shard = self.base_layer.q_proj_shard_size
+        kv_shard = self.base_layer.kv_proj_shard_size
+        return 0, q_shard, q_shard + kv_shard
+
+    def slice_lora_b_member(self, member: int, B: torch.Tensor) -> torch.Tensor:
+        base_layer = self.base_layer
+        if member == 0:
+            shard_size = base_layer.q_proj_shard_size
+            shard = base_layer.tp_rank
+        else:
+            shard_size = base_layer.kv_proj_shard_size
+            shard = base_layer.tp_rank // base_layer.num_kv_head_replicas
+        return B[shard * shard_size : (shard + 1) * shard_size, :]
+
 
 class RowParallelLinearWithLoRA(BaseLayerWithLoRA):
     def __init__(
@@ -976,6 +1005,15 @@ def slice_lora_a_weights(self, A: torch.Tensor):
     def slice_lora_b_weights(self, B: torch.Tensor):
         return B
 
+    def lora_b_member_full_rows(self) -> Tuple[int, ...]:
+        return self.first_output_dim, self.output_size - self.first_output_dim
+
+    def lora_b_member_local_offsets(self) -> Tuple[int, ...]:
+        return 0, self.first_output_dim
+
+    def slice_lora_b_member(self, member: int, B: torch.Tensor) -> torch.Tensor:
+        return B
+
 
 class FusedMoEWithLoRA(BaseLayerWithLoRA):
     """

python/sglang/srt/lora/lora.py

@@ -37,6 +37,7 @@
 # keys ("...experts.<module>..."), while excluding "shared_experts." (where the
 # preceding char is "_", not ".").
 _ROUTED_EXPERT_PATTERN = re.compile(r"\.experts\.")
+_MOE_EXPERT_PATTERN = re.compile(r"\.(?:experts|shared_experts?)\.")
 
 logger = logging.getLogger(__name__)
 
@@ -216,10 +217,9 @@ def _process_weight(self, name: str, loaded_weight: torch.Tensor):
             )
 
     def _normalize_weights(self):
+        # qkv / qkvr / MLA qkv_a / dense gate_up stay per member: the memory pool writes each
+        # member into its slot of the fused buffer (LORA_FUSED_MEMBERS).
         for layer in self.layers:
-            weight_names = list(layer.weights.keys())
-            self.normalize_qkv_proj(weight_names, layer.weights)
-            self.normalize_inkling_qkvr_proj(weight_names, layer.weights)
             self._normalize_shared_expert_moe(layer.weights)
             self._rename_expert_w_to_proj(layer.weights)
             # Stack gate_proj + x_proj → in_proj for Mamba layers (before gate_up normalization)
@@ -230,107 +230,6 @@ def _normalize_weights(self):
             self._normalize_in_proj_ba(layer.weights)
             weight_names = list(layer.weights.keys())
             self.normalize_gate_up_proj(weight_names, layer.weights)
-            weight_names = list(layer.weights.keys())
-            self.normalize_fused_qkv_a_proj(weight_names, layer.weights)
-
-    def normalize_inkling_qkvr_proj(
-        self, weight_names: List[str], weights: Dict[str, torch.Tensor]
-    ):
-        """Normalize Inkling split attention LoRA keys into the runtime qkvr layer.
-
-        Inkling checkpoints expose separate attention projections
-        (wq_du/wk_dv/wv_dv/wr_du), while the serving model uses a single
-        MergedColumnParallelLinear named qkvr with those slices in that order.
-        """
-        for weight_name in weight_names:
-            if "wq_du" not in weight_name:
-                continue
-
-            q_name = weight_name
-            k_name = weight_name.replace("wq_du", "wk_dv")
-            v_name = weight_name.replace("wq_du", "wv_dv")
-            r_name = weight_name.replace("wq_du", "wr_du")
-            qkvr_name = weight_name.replace("wq_du", "qkvr")
-
-            if any(name not in weights for name in (k_name, v_name, r_name)):
-                continue
-
-            cat_dim = weights[q_name].dim() - 2
-            weights[qkvr_name] = torch.cat(
-                (
-                    weights[q_name],
-                    weights[k_name],
-                    weights[v_name],
-                    weights[r_name],
-                ),
-                cat_dim,
-            )
-            weights.pop(q_name)
-            weights.pop(k_name)
-            weights.pop(v_name)
-            weights.pop(r_name)
-
-        for weight_name in list(weights.keys()):
-            if "qkvr" not in weight_name or "lora_A" not in weight_name:
-                continue
-            if weights[weight_name].shape[-2] != self.config.r:
-                continue
-            repeat_dims = [1] * weights[weight_name].dim()
-            repeat_dims[-2] = 4
-            weights[weight_name] = weights[weight_name].repeat(*repeat_dims)
-
-    def normalize_qkv_proj(
-        self, weight_names: List[str], weights: Dict[str, torch.Tensor]
-    ):
-        # Collect target q/k/v modules. This process is necessary since there might be no lora attached to k_proj
-        target_module = set()
-        for weight_name in weight_names:
-            if "k_proj" in weight_name:
-                target_module.add("k_proj")
-            if "q_proj" in weight_name:
-                target_module.add("q_proj")
-            if "v_proj" in weight_name:
-                target_module.add("v_proj")
-            if "qkv_proj" in weight_name:
-                target_module.add("qkv_proj")
-        if len(target_module) == 0:
-            return
-
-        for weight_name in weight_names:
-            # We assume every lora adaptor should contain lora modules for q_proj
-            if "q_proj" in weight_name:
-                q_name = weight_name
-                k_name = weight_name.replace("q_proj", "k_proj")
-                v_name = weight_name.replace("q_proj", "v_proj")
-                qkv_name = weight_name.replace("q_proj", "qkv_proj")
-
-                # If k_proj doesn't have lora, initialize it to zero
-                k_proj_weight = (
-                    weights[k_name]
-                    if "k_proj" in target_module
-                    else torch.zeros_like(weights[v_name])
-                )
-                weights[qkv_name] = torch.cat(
-                    (
-                        weights[q_name],
-                        k_proj_weight,
-                        weights[v_name],
-                    ),
-                    0,
-                )
-                weights.pop(q_name)
-                if "k_proj" in target_module:
-                    weights.pop(k_name)
-                weights.pop(v_name)
-            elif "qkv_proj" in weight_name:
-                # If qkv_proj is already stacked, we normalize it following the SGL convention.
-                qkv_name = weight_name
-                q_name = weight_name.replace("qkv_proj", "q_proj")
-                k_name = weight_name.replace("qkv_proj", "k_proj")
-                v_name = weight_name.replace("qkv_proj", "v_proj")
-                if "lora_A" in weight_name:
-                    weights[qkv_name] = weights[qkv_name].repeat(3, 1)
-                # else: no-op as LoRA B weight is already stacked.
 
     def _normalize_shared_expert_moe(self, weights: Dict[str, torch.Tensor]):
         """Reshape flat Inkling shared-sink factors to shared-outer 3D form."""
@@ -538,7 +437,10 @@ def _normalize_in_proj_ba(self, weights: Dict[str, torch.Tensor]):
     def normalize_gate_up_proj(
         self, weight_names: List[str], weights: Dict[str, torch.Tensor]
     ):
+        """Stack routed- and shared-expert gate/up factors; dense gate/up load per member."""
         for weight_name in weight_names:
+            if not _MOE_EXPERT_PATTERN.search(weight_name):
+                continue
             if "gate_proj" in weight_name:
                 up_name = weight_name.replace("gate_proj", "up_proj")
                 gate_up_name = weight_name.replace("gate_proj", "gate_up_proj")
@@ -576,33 +478,6 @@ def normalize_gate_up_proj(
         # Models with non-gated MLP/shared-experts declare up_proj in
         # supported_lora_modules so they get their own buffer and wrapping.
 
-    def normalize_fused_qkv_a_proj(
-        self, weight_names: List[str], weights: Dict[str, torch.Tensor]
-    ):
-        """Fuse separate q_a_proj and kv_a_proj_with_mqa LoRA weights into
-        a single fused_qkv_a_proj_with_mqa entry (concat along dim 0 for
-        both A and B), matching the DeepSeek MLA fused projection layout."""
-        for weight_name in weight_names:
-            if "q_a_proj" not in weight_name:
-                continue
-            if "fused_qkv_a_proj_with_mqa" in weight_name:
-                continue
-
-            q_a_name = weight_name
-            kv_a_name = weight_name.replace("q_a_proj", "kv_a_proj_with_mqa")
-            fused_name = weight_name.replace("q_a_proj", "fused_qkv_a_proj_with_mqa")
-
-            kv_a_weight = (
-                weights[kv_a_name]
-                if kv_a_name in weights
-                else torch.zeros_like(weights[q_a_name])
-            )
-
-            weights[fused_name] = torch.cat((weights[q_a_name], kv_a_weight), dim=0)
-            weights.pop(q_a_name)
-            if kv_a_name in weights:
-                weights.pop(kv_a_name)
-
     def pin_weights_in_cpu(self):
         for layer in self.layers:
             for name, weight in layer.weights.items():