[Lora] Lora quat info re-factor and support deepseekv3 mla lora (#22323)

python/sglang/srt/layers/quantization/base_config.py

@@ -10,6 +10,7 @@ from torch import nn
 
 if TYPE_CHECKING:
     from sglang.srt.layers.moe.moe_runner import MoeRunnerConfig
+    from sglang.srt.layers.moe.moe_runner.triton import TritonMoeQuantInfo
     from sglang.srt.layers.moe.token_dispatcher import CombineInput, DispatchOutput
     from sglang.srt.models.utils import WeightsMapper
 
@@ -106,6 +107,19 @@ class FusedMoEMethodBase(QuantizeMethodBase):
     ) -> CombineInput:
         raise NotImplementedError
 
+    def get_triton_quant_info(self, layer: torch.nn.Module) -> "TritonMoeQuantInfo":
+        """Return a ``TritonMoeQuantInfo`` describing the quantisation state
+        stored on *layer*.
+
+        The LoRA MoE runner calls this so that ``invoke_fused_moe_kernel``
+        receives the correct flags / scales / block-shape for the base
+        weights.  Each quantisation method must override this with the
+        same construction it already uses inside ``apply()``.
+        """
+        raise NotImplementedError(
+            f"{type(self).__name__} must implement get_triton_quant_info()"
+        )
+
 
 class QuantizationConfig(ABC):
     """Base class for quantization configs."""

python/sglang/srt/layers/quantization/blockwise_int8.py

@@ -360,13 +360,8 @@ class BlockInt8MoEMethod(FusedMoEMethodBase):
         self.moe_runner_config = moe_runner_config
         self.runner = MoeRunner(MoeRunnerBackend.TRITON, moe_runner_config)
 
-    def apply(
+    def get_triton_quant_info(self, layer: torch.nn.Module) -> TritonMoeQuantInfo:
-        self,
+        return TritonMoeQuantInfo(
-        layer: torch.nn.Module,
-        dispatch_output: StandardDispatchOutput,
-    ) -> CombineInput:
-
-        quant_info = TritonMoeQuantInfo(
             w13_weight=layer.w13_weight,
             w2_weight=layer.w2_weight,
             use_int8_w8a8=True,
@@ -377,4 +372,12 @@ class BlockInt8MoEMethod(FusedMoEMethodBase):
             block_shape=self.quant_config.weight_block_size,
         )
 
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        dispatch_output: StandardDispatchOutput,
+    ) -> CombineInput:
+
+        quant_info = self.get_triton_quant_info(layer)
+
         return self.runner.run(dispatch_output, quant_info)

python/sglang/srt/layers/quantization/fp8.py

@@ -1487,6 +1487,26 @@ class Fp8MoEMethod(FusedMoEMethodBase):
             # TODO(cwan): refactor other backends
             pass
 
+    def get_triton_quant_info(self, layer: torch.nn.Module) -> TritonMoeQuantInfo:
+        return TritonMoeQuantInfo(
+            w13_weight=layer.w13_weight,
+            w2_weight=layer.w2_weight,
+            b13=getattr(layer, "w13_weight_bias", None),
+            b2=getattr(layer, "w2_weight_bias", None),
+            use_fp8_w8a8=True,
+            w13_scale=(
+                layer.w13_weight_scale_inv
+                if self.block_quant
+                else layer.w13_weight_scale
+            ),
+            w2_scale=(
+                layer.w2_weight_scale_inv if self.block_quant else layer.w2_weight_scale
+            ),
+            a13_scale=layer.w13_input_scale,
+            a2_scale=layer.w2_input_scale,
+            block_shape=self.quant_config.weight_block_size,
+        )
+
     def apply(
         self,
         layer: torch.nn.Module,
@@ -1663,26 +1683,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
                 ),
             )
         elif self.runner.runner_backend.is_triton():
-            quant_info = TritonMoeQuantInfo(
+            quant_info = self.get_triton_quant_info(layer)
-                w13_weight=layer.w13_weight,
-                w2_weight=layer.w2_weight,
-                b13=getattr(layer, "w13_weight_bias", None),
-                b2=getattr(layer, "w2_weight_bias", None),
-                use_fp8_w8a8=True,
-                w13_scale=(
-                    layer.w13_weight_scale_inv
-                    if self.block_quant
-                    else layer.w13_weight_scale
-                ),
-                w2_scale=(
-                    layer.w2_weight_scale_inv
-                    if self.block_quant
-                    else layer.w2_weight_scale
-                ),
-                a13_scale=layer.w13_input_scale,
-                a2_scale=layer.w2_input_scale,
-                block_shape=self.quant_config.weight_block_size,
-            )
         else:
             raise NotImplementedError(
                 "Unsupported runner backend: %s" % self.runner.runner_backend

python/sglang/srt/layers/quantization/moe_wna16.py

@@ -364,19 +364,10 @@ class MoeWNA16Method(FusedMoEMethodBase):
         self.moe_runner_config = moe_runner_config
         self.runner = MoeRunner(MoeRunnerBackend.TRITON, moe_runner_config)
 
-    def apply(
+    def get_triton_quant_info(self, layer: torch.nn.Module) -> TritonMoeQuantInfo:
-        self,
-        layer: torch.nn.Module,
-        dispatch_output: StandardDispatchOutput,
-    ) -> CombineInput:
-        assert (
-            self.moe_runner_config.activation == "silu"
-        ), "Only SiLU activation is supported."
-
         weight_bits = self.quant_config.weight_bits
         has_zp = self.quant_config.has_zp
-
+        return TritonMoeQuantInfo(
-        quant_info = TritonMoeQuantInfo(
             w13_weight=layer.w13_qweight,
             w2_weight=layer.w2_qweight,
             use_int4_w4a16=weight_bits == 4,
@@ -387,6 +378,17 @@ class MoeWNA16Method(FusedMoEMethodBase):
             w2_zp=layer.w2_qzeros if has_zp else None,
             block_shape=[0, layer.group_size],
         )
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        dispatch_output: StandardDispatchOutput,
+    ) -> CombineInput:
+        assert (
+            self.moe_runner_config.activation == "silu"
+        ), "Only SiLU activation is supported."
+
+        quant_info = self.get_triton_quant_info(layer)
         return self.runner.run(dispatch_output, quant_info)
 
     @staticmethod

python/sglang/srt/layers/quantization/unquant.py

@@ -554,6 +554,14 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, MultiPlatformOp):
             )
             return StandardCombineInput(hidden_states=output)
 
+    def get_triton_quant_info(self, layer: torch.nn.Module) -> TritonMoeQuantInfo:
+        return TritonMoeQuantInfo(
+            w13_weight=layer.w13_weight,
+            w2_weight=layer.w2_weight,
+            b13=getattr(layer, "w13_weight_bias", None),
+            b2=getattr(layer, "w2_weight_bias", None),
+        )
+
     def forward_xpu(
         self,
         layer: torch.nn.Module,
@@ -594,12 +602,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, MultiPlatformOp):
             ), f"activation = {moe_runner_config.activation} is not supported \
             for Triton PATH, please set ENV SGLANG_USE_SGL_XPU=1."
 
-            quant_info = TritonMoeQuantInfo(
+            quant_info = self.get_triton_quant_info(layer)
-                w13_weight=layer.w13_weight,
-                w2_weight=layer.w2_weight,
-                b13=getattr(layer, "w13_weight_bias", None),
-                b2=getattr(layer, "w2_weight_bias", None),
-            )
             return self.runner.run(dispatch_output, quant_info)
 
     def forward_npu(

python/sglang/srt/layers/quantization/w8a8_fp8.py

@@ -286,13 +286,8 @@ class W8A8FP8MoEMethod(FusedMoEMethodBase):
         self.moe_runner_config = moe_runner_config
         self.runner = MoeRunner(MoeRunnerBackend.TRITON, moe_runner_config)
 
-    def apply(
+    def get_triton_quant_info(self, layer: torch.nn.Module) -> TritonMoeQuantInfo:
-        self,
+        return TritonMoeQuantInfo(
-        layer: torch.nn.Module,
-        dispatch_output: StandardDispatchOutput,
-    ) -> CombineInput:
-
-        quant_info = TritonMoeQuantInfo(
             w13_weight=layer.w13_weight,
             w2_weight=layer.w2_weight,
             use_fp8_w8a8=True,
@@ -302,4 +297,12 @@ class W8A8FP8MoEMethod(FusedMoEMethodBase):
             a13_scale=layer.w13_input_scale,
             a2_scale=layer.w2_input_scale,
         )
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        dispatch_output: StandardDispatchOutput,
+    ) -> CombineInput:
+
+        quant_info = self.get_triton_quant_info(layer)
         return self.runner.run(dispatch_output, quant_info)

python/sglang/srt/layers/quantization/w8a8_int8.py

@@ -331,6 +331,18 @@ class W8A8Int8MoEMethod(FusedMoEMethodBase):
         self.moe_runner_config = moe_runner_config
         self.runner = MoeRunner(MoeRunnerBackend.TRITON, moe_runner_config)
 
+    def get_triton_quant_info(self, layer: torch.nn.Module) -> TritonMoeQuantInfo:
+        return TritonMoeQuantInfo(
+            w13_weight=layer.w13_weight,
+            w2_weight=layer.w2_weight,
+            use_int8_w8a8=True,
+            per_channel_quant=True,
+            w13_scale=layer.w13_weight_scale,
+            w2_scale=layer.w2_weight_scale,
+            a13_scale=layer.w13_input_scale,
+            a2_scale=layer.w2_input_scale,
+        )
+
     def apply(
         self,
         layer: torch.nn.Module,
@@ -365,14 +377,5 @@ class W8A8Int8MoEMethod(FusedMoEMethodBase):
             )
             return StandardCombineInput(hidden_states=output)
 
-        quant_info = TritonMoeQuantInfo(
+        quant_info = self.get_triton_quant_info(layer)
-            w13_weight=layer.w13_weight,
-            w2_weight=layer.w2_weight,
-            use_int8_w8a8=True,
-            per_channel_quant=True,
-            w13_scale=layer.w13_weight_scale,
-            w2_scale=layer.w2_weight_scale,
-            a13_scale=layer.w13_input_scale,
-            a2_scale=layer.w2_input_scale,
-        )
         return self.runner.run(dispatch_output, quant_info)

python/sglang/srt/lora/backend/chunked_backend.py

@@ -66,6 +66,7 @@ class ChunkedSgmvLoRABackend(BaseLoRABackend):
         x: torch.Tensor,
         weights: torch.Tensor,
         pruned_batch_info: LoRABatchInfo = None,
+        stack_num: int = 1,
         *args,
         **kwargs,
     ) -> torch.Tensor:
@@ -76,7 +77,7 @@ class ChunkedSgmvLoRABackend(BaseLoRABackend):
             x=x,
             weights=weights,
             batch_info=batch_info,
-            num_slices=1,
+            num_slices=stack_num,
         )
 
     def run_lora_b_sgemm(

python/sglang/srt/lora/backend/torch_backend.py

@@ -39,7 +39,12 @@ class TorchNativeLoRABackend(BaseLoRABackend):
         super().__init__(max_loras_per_batch, device)
 
     def run_lora_a_sgemm(
-        self, x: torch.Tensor, weights: torch.Tensor, *args, **kwargs
+        self,
+        x: torch.Tensor,
+        weights: torch.Tensor,
+        stack_num: int = 1,
+        *args,
+        **kwargs,
     ) -> torch.Tensor:
         output_tensor = sgemm_lora_a_fwd(
             inputs=x,
@@ -48,7 +53,7 @@ class TorchNativeLoRABackend(BaseLoRABackend):
             seg_len_tensor=self.batch_info.seg_lens_cpu,
             lora_ranks=self.batch_info.lora_ranks_cpu,
             scaling_tensor=self.batch_info.scalings_cpu,
-            num_slices=1,
+            num_slices=stack_num,
         )
 
         return output_tensor

python/sglang/srt/lora/backend/triton_backend.py

@@ -53,13 +53,14 @@ class TritonLoRABackend(BaseLoRABackend):
         x: torch.Tensor,
         weights: torch.Tensor,
         pruned_batch_info: LoRABatchInfo = None,
+        stack_num: int = 1,
         *args,
         **kwargs,
     ) -> torch.Tensor:
         batch_info = (
             pruned_batch_info if pruned_batch_info is not None else self.batch_info
         )
-        return sgemm_lora_a_fwd(x, weights, batch_info)
+        return sgemm_lora_a_fwd(x, weights, batch_info, stack_num=stack_num)
 
     def run_lora_b_sgemm(
         self,

python/sglang/srt/lora/layers.py

@@ -14,6 +14,7 @@ from sglang.srt.layers.linear import (
     ColumnParallelLinear,
     MergedColumnParallelLinear,
     QKVParallelLinear,
+    ReplicatedLinear,
     RowParallelLinear,
 )
 from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE
@@ -687,6 +688,94 @@ class RowParallelLinearWithLoRA(BaseLayerWithLoRA):
         return B
 
 
+class ReplicatedLinearWithLoRA(BaseLayerWithLoRA):
+    """LoRA wrapper for ReplicatedLinear (no TP sharding).
+
+    Used for DeepSeek MLA's fused_qkv_a_proj_with_mqa, which fuses
+    q_a_proj and kv_a_proj_with_mqa into a single replicated linear.
+    The two sub-projections have unequal output dimensions, so LoRA B
+    is applied via two separate sgemm calls, one per partition.
+    """
+
+    first_output_dim: int = 0
+
+    def __init__(
+        self,
+        base_layer: ReplicatedLinear,
+        lora_backend: BaseLoRABackend,
+    ) -> None:
+        super().__init__(base_layer, lora_backend)
+        self.output_size = base_layer.output_size
+
+    def set_lora_info(self, A_buffer: torch.Tensor, B_buffer: torch.Tensor):
+        self.set_lora = True
+        self.A_buffer = A_buffer
+        self.B_buffer = B_buffer
+        first = self.first_output_dim
+        if first > 0 and first < B_buffer.shape[-2]:
+            self.B_first = B_buffer[:, :first, :].contiguous()
+            self.B_second = B_buffer[:, first:, :].contiguous()
+            output_size = B_buffer.shape[-2]
+            self.first_offset = torch.tensor(
+                [0, first], dtype=torch.int32, device=B_buffer.device
+            )
+            self.second_offset = torch.tensor(
+                [0, output_size - first], dtype=torch.int32, device=B_buffer.device
+            )
+        else:
+            self.B_first = None
+            self.B_second = None
+            self.output_offset = torch.tensor(
+                [0, self.output_size],
+                dtype=torch.int32,
+                device=B_buffer.device,
+            )
+
+    def apply_lora(self, base_output: torch.Tensor, x: torch.Tensor) -> torch.Tensor:
+        if self.B_first is not None:
+            rank = self.B_buffer.shape[-1]
+            lora_a_output = self.lora_backend.run_lora_a_sgemm(
+                x, self.A_buffer, stack_num=2
+            )
+            first_out = base_output[:, : self.first_output_dim]
+            second_out = base_output[:, self.first_output_dim :]
+            self.lora_backend.run_lora_b_sgemm(
+                x=lora_a_output[:, :rank].contiguous(),
+                weights=self.B_first,
+                output_offset=self.first_offset,
+                base_output=first_out,
+            )
+            self.lora_backend.run_lora_b_sgemm(
+                x=lora_a_output[:, rank:].contiguous(),
+                weights=self.B_second,
+                output_offset=self.second_offset,
+                base_output=second_out,
+            )
+            return base_output
+        else:
+            lora_a_output = self.lora_backend.run_lora_a_sgemm(x, self.A_buffer)
+            return self.lora_backend.run_lora_b_sgemm(
+                x=lora_a_output,
+                weights=self.B_buffer,
+                output_offset=self.output_offset,
+                base_output=base_output,
+            )
+
+    def forward(self, x: torch.Tensor):
+        bias = self.base_layer.bias if not self.base_layer.skip_bias_add else None
+        output = self.base_layer.quant_method.apply(self.base_layer, x, bias)
+        if self.set_lora:
+            output = self.apply_lora(output, x)
+        output_bias = self.base_layer.bias if self.base_layer.skip_bias_add else None
+        return output, output_bias
+
+    def slice_lora_a_weights(self, A: torch.Tensor, tp_rank: int):
+        return A
+
+    def slice_lora_b_weights(self, B: torch.Tensor, tp_rank: int):
+        return B
+
+
 class FusedMoEWithLoRA(BaseLayerWithLoRA):
     """
     Wrapper around FusedMoE that integrates LoRA into the MoE computation.
@@ -721,7 +810,6 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
 
         # initialize triton_lora moe runner for batches with lora enabled
         from sglang.srt.layers.moe.moe_runner.runner import MoeRunner
-        from sglang.srt.layers.moe.moe_runner.triton import TritonMoeQuantInfo
 
         self._lora_runner = MoeRunner(
             base_layer.quant_method.runner.runner_backend,
@@ -730,12 +818,7 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
         )
 
         # Pre-compute quant info for efficiency (weights don't change during inference)
-        self._quant_info = TritonMoeQuantInfo(
+        self._quant_info = base_layer.quant_method.get_triton_quant_info(base_layer)
-            w13_weight=base_layer.w13_weight,
-            w2_weight=base_layer.w2_weight,
-            b13=getattr(base_layer, "w13_weight_bias", None),
-            b2=getattr(base_layer, "w2_weight_bias", None),
-        )
 
     def set_lora_info(
         self,
@@ -941,6 +1024,7 @@ def get_lora_layer(
         FusedMoE: FusedMoEWithLoRA,
         ParallelLMHead: ParallelLMHeadWithLoRA,
         VocabParallelEmbedding: VocabParallelEmbeddingWithLoRA,
+        ReplicatedLinear: ReplicatedLinearWithLoRA,
         QKVParallelLinear: QKVParallelLinearWithLoRA,
         MergedColumnParallelLinear: MergedColumnParallelLinearWithLoRA,
         ColumnParallelLinear: ColumnParallelLinearWithLoRA,

python/sglang/srt/lora/lora.py

@@ -133,13 +133,14 @@ class LoRAAdapter(nn.Module):
             )
 
     def _normalize_weights(self):
-        # normalize kv_proj and gate_up_proj
         for layer in self.layers:
             weight_names = list(layer.weights.keys())
             self.normalize_qkv_proj(weight_names, layer.weights)
             self._rename_expert_w_to_proj(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_qkv_proj(
         self, weight_names: List[str], weights: Dict[str, torch.Tensor]
@@ -242,6 +243,33 @@ class LoRAAdapter(nn.Module):
                     weights[gate_up_name] = weights[gate_up_name].repeat(*repeat_dims)
                 # else: no-op as LoRA B weight is already stacked.
 
+    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():

python/sglang/srt/lora/lora_manager.py

@@ -763,6 +763,25 @@ class LoRAManager:
                     self.lm_head_module = lora_module
                     continue
 
+            # Handle DeepSeek MLA fused projection: set the boundary
+            # between q_a and kv_a output partitions so the LoRA layer
+            # can apply separate B projections for each.
+            if (
+                "fused_qkv_a_proj_with_mqa" in self.target_modules
+                and module_name.endswith("fused_qkv_a_proj_with_mqa")
+            ):
+                from sglang.srt.lora.layers import ReplicatedLinearWithLoRA
+
+                layer_id = get_layer_id(module_name)
+                if layer_id is None:
+                    continue
+                lora_module = self.set_lora_module(module_name, module)
+                if isinstance(lora_module, ReplicatedLinearWithLoRA):
+                    q_lora_rank = getattr(self.base_hf_config, "q_lora_rank", None) or 0
+                    lora_module.first_output_dim = q_lora_rank
+                self.lora_modules[layer_id][module_name] = lora_module
+                continue
+
             # The module should be converted if it is included in target_names
             if module_name.split(".")[-1] in self.target_modules:
                 layer_id = get_layer_id(module_name)

python/sglang/srt/lora/mem_pool.py

@@ -12,6 +12,7 @@ from sglang.srt.lora.lora_config import LoRAConfig
 from sglang.srt.lora.lora_registry import LoRARef
 from sglang.srt.lora.utils import (
     EMBEDDING_NAMES,
+    REPLICATED_LINEAR_LORA_NAMES,
     ROW_PARALLELISM_LINEAR_LORA_NAMES,
     LoRAType,
     get_hidden_dim,
@@ -188,7 +189,11 @@ class LoRAMemoryPool:
             module_name, self.base_hf_config, base_model, layer_idx
         )
         c = get_stacked_multiply(module_name)
-        if self.tp_size > 1 and module_name in ROW_PARALLELISM_LINEAR_LORA_NAMES:
+        if (
+            self.tp_size > 1
+            and module_name in ROW_PARALLELISM_LINEAR_LORA_NAMES
+            and module_name not in REPLICATED_LINEAR_LORA_NAMES
+        ):
             input_dim = divide(input_dim, self.tp_size)
 
         if self.is_moe_module(module_name):
@@ -240,7 +245,11 @@ class LoRAMemoryPool:
         _, output_dim = get_hidden_dim(
             module_name, self.base_hf_config, base_model, layer_idx
         )
-        if self.tp_size > 1 and module_name not in ROW_PARALLELISM_LINEAR_LORA_NAMES:
+        if (
+            self.tp_size > 1
+            and module_name not in ROW_PARALLELISM_LINEAR_LORA_NAMES
+            and module_name not in REPLICATED_LINEAR_LORA_NAMES
+        ):
             output_dim = divide(output_dim, self.tp_size)
 
         # Check if MoE module and return appropriate shape
@@ -619,7 +628,12 @@ class LoRAMemoryPool:
 
                 if name in ["gate_up_proj_moe", "down_proj_moe"]:
                     if self.experts_shared_outer_loras and name == "gate_up_proj_moe":
-                        if isinstance(weights, torch.Tensor) and weights.dim() == 3:
+                        if weights is None:
+                            buffer_view = target_buffer[
+                                buffer_id, 0, : lora_rank * c, :
+                            ]
+                            load_lora_weight_tensor(buffer_view, None)
+                        elif isinstance(weights, torch.Tensor) and weights.dim() == 3:
                             if weights.shape[0] != 1:
                                 raise ValueError(
                                     f"experts_shared_outer_loras is enabled but "
@@ -669,7 +683,10 @@ class LoRAMemoryPool:
 
                 if name in ["gate_up_proj_moe", "down_proj_moe"]:
                     if self.experts_shared_outer_loras and name == "down_proj_moe":
-                        if isinstance(weights, torch.Tensor) and weights.dim() == 3:
+                        if weights is None:
+                            buffer_view = target_buffer[buffer_id, 0, :, :lora_rank]
+                            load_lora_weight_tensor(buffer_view, None)
+                        elif isinstance(weights, torch.Tensor) and weights.dim() == 3:
                             if weights.shape[0] != 1:
                                 raise ValueError(
                                     f"experts_shared_outer_loras is enabled but "

python/sglang/srt/lora/utils.py

@@ -83,14 +83,47 @@ def get_hidden_dim(
                 config.num_attention_heads + config.num_key_value_heads * 2
             )
         elif module_name == "o_proj":
+            o_head_dim = getattr(config, "v_head_dim", None) or head_dim
             return (
-                head_dim * config.num_attention_heads,
+                o_head_dim * config.num_attention_heads,
                 config.hidden_size,
             )
         elif module_name == "gate_up_proj":
-            return config.hidden_size, config.intermediate_size * 2
+            inter = config.intermediate_size
+            first_k = getattr(config, "first_k_dense_replace", None)
+            moe_freq = getattr(config, "moe_layer_freq", 1)
+            if (
+                first_k is not None
+                and layer_idx >= first_k
+                and layer_idx % moe_freq == 0
+            ):
+                moe_inter = getattr(config, "moe_intermediate_size", None)
+                n_shared = getattr(config, "n_shared_experts", None)
+                if moe_inter is not None and n_shared is not None:
+                    inter = moe_inter * n_shared
+            return config.hidden_size, inter * 2
         elif module_name == "down_proj":
-            return config.intermediate_size, config.hidden_size
+            inter = config.intermediate_size
+            first_k = getattr(config, "first_k_dense_replace", None)
+            moe_freq = getattr(config, "moe_layer_freq", 1)
+            if (
+                first_k is not None
+                and layer_idx >= first_k
+                and layer_idx % moe_freq == 0
+            ):
+                moe_inter = getattr(config, "moe_intermediate_size", None)
+                n_shared = getattr(config, "n_shared_experts", None)
+                if moe_inter is not None and n_shared is not None:
+                    inter = moe_inter * n_shared
+            return inter, config.hidden_size
+        elif module_name == "fused_qkv_a_proj_with_mqa":
+            q_lora_rank = getattr(config, "q_lora_rank", None) or 0
+            kv_lora_rank = config.kv_lora_rank
+            qk_rope_head_dim = config.qk_rope_head_dim
+            return (
+                config.hidden_size,
+                q_lora_rank + kv_lora_rank + qk_rope_head_dim,
+            )
         elif module_name == "gate_up_proj_moe":
             moe_inter = (
                 getattr(config, "moe_intermediate_size", None)
@@ -151,6 +184,8 @@ def get_normalized_target_modules(
         "lm_head": "lm_head",
         "output": "lm_head",
         "unembed_tokens": "lm_head",
+        "q_a_proj": "fused_qkv_a_proj_with_mqa",
+        "kv_a_proj_with_mqa": "fused_qkv_a_proj_with_mqa",
     }
 
     result = set()
@@ -169,6 +204,7 @@ def get_stacked_multiply(module_name: str) -> int:
         "qkv_proj": 3,
         "gate_up_proj": 2,
         "gate_up_proj_moe": 2,
+        "fused_qkv_a_proj_with_mqa": 2,
     }
     return stacked_rank[module_name] if module_name in stacked_rank else 1
 
@@ -190,6 +226,7 @@ def get_target_module_name(full_module_name: str, target_modules: Set[str]) -> s
 
 EMBEDDING_NAMES = ["embed_tokens", "lm_head"]
 ROW_PARALLELISM_LINEAR_LORA_NAMES = ["o_proj", "down_proj", "down_proj_moe"]
+REPLICATED_LINEAR_LORA_NAMES = ["fused_qkv_a_proj_with_mqa"]
 
 # Normalized module names that the LoRA system fully supports
 # (i.e. get_hidden_dim, init_buffers, and init_lora_modules can handle them).
@@ -201,6 +238,7 @@ _KNOWN_LORA_TARGET_MODULES = frozenset(
         "down_proj",
         "embed_tokens",
         "lm_head",
+        "fused_qkv_a_proj_with_mqa",
     }
 )
 

test/registered/lora/test_lora_deepseek_v3_base_logprob_diff.py

@@ -0,0 +1,156 @@
+# Copyright 2023-2025 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""
+Regression test for DeepSeek-V3.1-Base MLA LoRA logprob accuracy.
+
+Compares SGLang LoRA logprobs against reference training logprobs from a
+pre-computed dataset. The LoRA adapter and reference data are downloaded from:
+https://huggingface.co/datasets/yushengsu/lora-diff-DeepSeek-V3.1-Base
+
+Usage:
+    python -m unittest test_lora_deepseek_v3_base_logprob_diff
+"""
+
+import multiprocessing as mp
+import os
+import unittest
+
+import torch
+from huggingface_hub import snapshot_download
+
+import sglang as sgl
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(
+    est_time=300,
+    suite="nightly-8-gpu-b200",
+)
+
+BASE_MODEL = "deepseek-ai/DeepSeek-V3.1-Base"
+LORA_HF_REPO = "yushengsu/lora-diff-DeepSeek-V3.1-Base"
+LORA_BACKEND = "triton"
+MAX_LORA_RANK = 32
+TP_SIZE = 8
+MOE_RUNNER_BACKEND = "triton"
+EXPERTS_SHARED_OUTER_LORAS = True
+PREFILL_ATTENTION_BACKEND = "fa4"
+DECODE_ATTENTION_BACKEND = "flashinfer"
+
+KL_THRESHOLD = 6e-3
+
+
+def kl_v2(a, b):
+    a = torch.tensor(a) if not torch.is_tensor(a) else a
+    b = torch.tensor(b) if not torch.is_tensor(b) else b
+    return (((a - b) ** 2) * 0.5).mean().item()
+
+
+def get_prompt_logprobs(engine, input_ids, lora_path):
+    if isinstance(input_ids, torch.Tensor):
+        input_ids = [input_ids.tolist()]
+    elif not isinstance(input_ids[0], list):
+        input_ids = [input_ids]
+    out = engine.generate(
+        input_ids=input_ids,
+        sampling_params={"max_new_tokens": 0, "temperature": 0.0},
+        return_logprob=True,
+        logprob_start_len=0,
+        lora_path=lora_path,
+    )
+    if isinstance(out, list):
+        out = out[0]
+    return [logprob for logprob, _, _ in out["meta_info"]["input_token_logprobs"]][1:]
+
+
+class TestLoRADeepSeekV3BaseLogprobDiff(CustomTestCase):
+
+    def test_lora_deepseek_v3_base_logprob_accuracy(self):
+        adapter_path = snapshot_download(
+            LORA_HF_REPO,
+            repo_type="dataset",
+        )
+
+        engine = sgl.Engine(
+            model_path=BASE_MODEL,
+            tp_size=TP_SIZE,
+            enable_lora=True,
+            max_lora_rank=MAX_LORA_RANK,
+            lora_paths={"my_lora": adapter_path},
+            lora_backend=LORA_BACKEND,
+            attention_backend="flashinfer",
+            moe_runner_backend=MOE_RUNNER_BACKEND,
+            experts_shared_outer_loras=EXPERTS_SHARED_OUTER_LORAS,
+            prefill_attention_backend=PREFILL_ATTENTION_BACKEND,
+            decode_attention_backend=DECODE_ATTENTION_BACKEND,
+            disable_shared_experts_fusion=True,
+        )
+
+        try:
+            cdata = torch.load(
+                os.path.join(adapter_path, "compare_sample_train_data.pt"),
+                weights_only=False,
+            )
+
+            base_logprobs = get_prompt_logprobs(engine, cdata["tokens"], lora_path=None)
+            logprobs = get_prompt_logprobs(engine, cdata["tokens"], lora_path="my_lora")
+
+            base_t = torch.tensor(base_logprobs)
+            lora_t = torch.tensor(logprobs)
+            diff = (base_t - lora_t).abs()
+            print(
+                f"[VERIFY] base vs lora: mean_diff={diff.mean().item():.6f}, "
+                f"max_diff={diff.max().item():.6f}, "
+                f"identical={torch.equal(base_t, lora_t)}"
+            )
+
+            self.assertFalse(
+                torch.equal(base_t, lora_t),
+                "LoRA logprobs should differ from base model logprobs",
+            )
+
+            kl_sglang_trainer = kl_v2(cdata["training_logprobs"], logprobs)
+            kl_orig_trainer = kl_v2(
+                cdata["training_logprobs"], cdata["sampling_logprobs"]
+            )
+            kl_sglang_orig = kl_v2(logprobs, cdata["sampling_logprobs"])
+
+            print(f"KL(orig_sampler, trainer) = {kl_orig_trainer:.6e}")
+            print(f"KL(sglang, trainer)       = {kl_sglang_trainer:.6e}")
+            print(f"KL(sglang, orig_sampler)  = {kl_sglang_orig:.6e}")
+
+            self.assertLessEqual(
+                kl_sglang_trainer,
+                KL_THRESHOLD,
+                f"KL(sglang, trainer) = {kl_sglang_trainer:.6e} exceeds "
+                f"threshold {KL_THRESHOLD}",
+            )
+
+        finally:
+            engine.shutdown()
+
+
+if __name__ == "__main__":
+    try:
+        mp.set_start_method("spawn")
+    except RuntimeError:
+        pass
+
+    try:
+        unittest.main(warnings="ignore", verbosity=2)
+    finally:
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+            torch.cuda.synchronize()