LoRA support for qwen3.5 and nemotron3 (#23594)

Co-authored-by: Yanbin Jiang <jybsuper@gmail.com>

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

@@ -87,16 +87,16 @@ class AscendLoRABackend(BaseLoRABackend):
         output_offset_cpu: torch.Tensor,
         max_qkv_out_dim: int,
         base_output: torch.Tensor = None,
+        n_slices: int = 3,
         *args,
         **kwargs,
     ) -> torch.Tensor:
-        num_slices = 3
         assert isinstance(qkv_lora_b, torch.Tensor)
 
         total_seq_len, _ = x.shape
         _, weight_intermediate_dim, _ = qkv_lora_a.shape
         _, weight_out_dim, _ = qkv_lora_b.shape
-        max_rank = weight_intermediate_dim // num_slices
+        max_rank = weight_intermediate_dim // n_slices
 
         if base_output is None:
             output_tensor = torch.zeros(
@@ -124,7 +124,7 @@ class AscendLoRABackend(BaseLoRABackend):
         )
         lora_a_output *= scaling
 
-        for slice_id in range(num_slices):
+        for slice_id in range(n_slices):
             slice_offset = output_offset_cpu[slice_id]
             slice_offset_next = output_offset_cpu[slice_id + 1]
             slice_size = slice_offset_next - slice_offset

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

@@ -113,20 +113,21 @@ class ChunkedSgmvLoRABackend(BaseLoRABackend):
         output_offset: torch.Tensor,
         max_qkv_out_dim: int,
         base_output: torch.Tensor = None,
+        n_slices: int = 3,
         *args,
         **kwargs,
     ) -> torch.Tensor:
 
         # x: (s, input_dim)
-        # qkv_lora_a: (num_lora, 3 * r, input_dim)
-        # qkv_lora_b: (num_lora, output_dim_q + 2 * output_dim_kv, r)
+        # qkv_lora_a: (num_lora, n_slices * r, input_dim)
+        # qkv_lora_b: (num_lora, total_output_dim, r)
         assert isinstance(qkv_lora_b, torch.Tensor)
 
         lora_a_output = chunked_sgmv_lora_shrink_forward(
             x=x,
             weights=qkv_lora_a,
             batch_info=self.batch_info,
-            num_slices=3,
+            num_slices=n_slices,
         )
         lora_output = chunked_sgmv_lora_expand_forward(
             x=lora_a_output,

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

@@ -91,10 +91,10 @@ class TorchNativeLoRABackend(BaseLoRABackend):
         output_offset_cpu: torch.Tensor,
         max_qkv_out_dim: int,
         base_output: torch.Tensor = None,
+        n_slices: int = 3,
         *args,
         **kwargs,
     ) -> torch.Tensor:
-        num_slices = 3
         lora_a_output = sgemm_lora_a_fwd(
             inputs=x,
             weights=qkv_lora_a,
@@ -102,7 +102,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=num_slices,
+            num_slices=n_slices,
         )
 
         output_tensor = sgemm_lora_b_fwd(

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

@@ -88,17 +88,18 @@ class TritonLoRABackend(BaseLoRABackend):
         output_offset: torch.Tensor,
         max_qkv_out_dim: int,
         base_output: torch.Tensor = None,
+        n_slices: int = 3,
         *args,
         **kwargs,
     ) -> torch.Tensor:
 
         # x: (s, input_dim)
-        # qkv_lora_a: (num_lora, 3 * r, input_dim)
-        # qkv_lora_b: (num_lora, output_dim_q + 2 * output_dim_kv, r)
+        # qkv_lora_a: (num_lora, n_slices * r, input_dim)
+        # qkv_lora_b: (num_lora, total_output_dim, r)
         assert isinstance(qkv_lora_b, torch.Tensor)
 
         sgemm_info = self._sgemm_info()
-        lora_a_output = sgemm_lora_a_fwd(x, qkv_lora_a, sgemm_info, stack_num=3)
+        lora_a_output = sgemm_lora_a_fwd(x, qkv_lora_a, sgemm_info, stack_num=n_slices)
         lora_output = qkv_lora_b_fwd(
             lora_a_output,
             qkv_lora_b,
@@ -106,6 +107,7 @@ class TritonLoRABackend(BaseLoRABackend):
             output_offset,
             max_qkv_out_dim,
             base_output,
+            n_slices=n_slices,
         )
         return lora_output
 

python/sglang/srt/lora/layers.py

@@ -474,6 +474,7 @@ class MergedColumnParallelLinearWithLoRA(ColumnParallelLinearWithLoRA):
         lora_backend: BaseLoRABackend,
     ) -> None:
         super().__init__(base_layer, lora_backend)
+        self.n_slices = len(self.base_layer.output_partition_sizes)
 
     def set_lora_info(
         self,
@@ -481,46 +482,89 @@ class MergedColumnParallelLinearWithLoRA(ColumnParallelLinearWithLoRA):
         B_buffer: torch.Tensor,
     ):
         self.set_lora = True
-        self.A_buffer_gate_up = A_buffer
-        self.B_buffer_gate_up = B_buffer
+        self.A_buffer = A_buffer
+        self.B_buffer = B_buffer
 
-        shard_size = self.base_layer.output_partition_sizes[0]
+        # Build cumulative output offsets from the first `lora_n_slices`
+        # base partitions. `lora_n_slices` may be smaller than self.n_slices
+        # when only a subset of partitions are LoRA'd (e.g. Mamba in_proj
+        # has 5 partitions but stacked_multiply=2), so we can't precompute
+        # these in __init__.
+        lora_n_slices = self._get_lora_n_slices()
+        if lora_n_slices <= 0 or lora_n_slices > self.n_slices:
+            raise ValueError(
+                f"Invalid LoRA slice count {lora_n_slices} for "
+                f"{self.n_slices} base output partitions."
+            )
+        partition_sizes = list(self.base_layer.output_partition_sizes[:lora_n_slices])
+        offsets = [0]
+        for ps in partition_sizes:
+            offsets.append(offsets[-1] + ps)
+        if offsets[-1] != B_buffer.shape[-2]:
+            raise ValueError(
+                f"LoRA B output dim {B_buffer.shape[-2]} does not match "
+                f"base partition prefix dim {offsets[-1]} for {lora_n_slices} slices."
+            )
         self.output_offset = torch.tensor(
-            [
-                0,
-                shard_size,
-                2 * shard_size,
-            ],
+            offsets,
             dtype=torch.int32,
             device=next(self.base_layer.parameters()).device,
         )
+        self.output_offset_cpu = self.output_offset.cpu()
+        self.max_out_dim = max(partition_sizes)
+        self.use_gate_up_lora = (
+            lora_n_slices == 2 and partition_sizes[0] == partition_sizes[1]
+        )
+
+    def _get_lora_n_slices(self) -> int:
+        """Actual number of LoRA slices from the buffer shapes.
+
+        May differ from self.n_slices (base layer partitions) when only a
+        subset of partitions are LoRA'd (e.g. Mamba in_proj has 5 partitions
+        but stacked_multiply=2).
+        """
+        lora_rank = self.B_buffer.shape[-1]
+        if lora_rank == 0:
+            return self.n_slices
+        return self.A_buffer.shape[-2] // lora_rank
 
     def apply_lora(self, base_output: torch.Tensor, x: torch.Tensor) -> torch.Tensor:
-        lora_output = self.lora_backend.run_gate_up_lora(
-            x=x,
-            gate_up_lora_a=self.A_buffer_gate_up,
-            gate_up_lora_b=self.B_buffer_gate_up,
-            output_offset=self.output_offset,
-            base_output=base_output,
-        )
+        lora_n_slices = self._get_lora_n_slices()
+        if lora_n_slices == 2 and self.use_gate_up_lora:
+            lora_output = self.lora_backend.run_gate_up_lora(
+                x=x,
+                gate_up_lora_a=self.A_buffer,
+                gate_up_lora_b=self.B_buffer,
+                output_offset=self.output_offset,
+                base_output=base_output,
+            )
+        else:
+            lora_output = self.lora_backend.run_qkv_lora(
+                x=x,
+                qkv_lora_a=self.A_buffer,
+                qkv_lora_b=self.B_buffer,
+                output_offset=self.output_offset,
+                output_offset_cpu=self.output_offset_cpu,
+                max_qkv_out_dim=self.max_out_dim,
+                base_output=base_output,
+                n_slices=lora_n_slices,
+            )
         return lora_output
 
     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):
-        # Since the outputs for both gate and up are identical, we use a random one.
-        shard_size = self.base_layer.output_partition_sizes[0]
-        gate_size = self.base_layer.output_sizes[0]
-        start_idx = tp_rank * shard_size
-        end_idx = (tp_rank + 1) * shard_size
-        return torch.concat(
-            (
-                B[start_idx:end_idx, :],
-                B[gate_size + start_idx : gate_size + end_idx],
-            ),
-            dim=0,
-        )
+        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 = tp_rank * part_size
+            end_idx = start_idx + part_size
+            slices.append(B[offset + start_idx : offset + end_idx, :])
+            offset += full_size
+        return torch.concat(slices, dim=0)
 
 
 class QKVParallelLinearWithLoRA(ColumnParallelLinearWithLoRA):
@@ -696,8 +740,12 @@ class ReplicatedLinearWithLoRA(BaseLayerWithLoRA):
 
     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.
+    The two sub-projections have unequal output dimensions, so we use
+    the N-component fused kernel (run_qkv_lora) with n_slices=2 to
+    handle the split inside the triton kernel rather than in Python.
+
+    ``first_output_dim`` (set by LoRAManager after construction) marks the
+    boundary between the first and second sub-projection in the output.
     """
 
     first_output_dim: int = 0
@@ -714,55 +762,53 @@ class ReplicatedLinearWithLoRA(BaseLayerWithLoRA):
         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
+        first_dim = self.first_output_dim
+        if first_dim > 0:
+            second_dim = B_buffer.shape[-2] - first_dim
+            self._output_offset = torch.tensor(
+                [0, first_dim, first_dim + second_dim],
+                dtype=torch.int32,
+                device=B_buffer.device,
             )
+            self._output_offset_cpu = self._output_offset.cpu()
+            self._max_out_dim = max(first_dim, second_dim)
         else:
-            self.B_first = None
-            self.B_second = None
-            self.output_offset = torch.tensor(
-                [0, self.output_size],
+            # Single-projection path: csgmv backend requires an explicit
+            # slice_offsets tensor of shape [0, output_dim].
+            self._output_offset = torch.tensor(
+                [0, B_buffer.shape[-2]],
                 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:
+        first_dim = self.first_output_dim
+
+        if first_dim == 0:
+            # Simple single-projection (e.g. fc1_latent_proj, fc2_latent_proj)
             lora_a_output = self.lora_backend.run_lora_a_sgemm(x, self.A_buffer)
-            return self.lora_backend.run_lora_b_sgemm(
+            lora_output = self.lora_backend.run_lora_b_sgemm(
                 x=lora_a_output,
                 weights=self.B_buffer,
-                output_offset=self.output_offset,
+                output_offset=self._output_offset,
                 base_output=base_output,
             )
+            return lora_output
+
+        # Use the fused N-component kernel with n_slices=2 to handle the
+        # split inside the triton kernel, avoiding Python-level splitting
+        # which breaks when adapter rank < max_lora_rank.
+        lora_output = self.lora_backend.run_qkv_lora(
+            x=x,
+            qkv_lora_a=self.A_buffer,
+            qkv_lora_b=self.B_buffer,
+            output_offset=self._output_offset,
+            output_offset_cpu=self._output_offset_cpu,
+            max_qkv_out_dim=self._max_out_dim,
+            base_output=base_output,
+            n_slices=2,
+        )
+        return lora_output
 
     def forward(self, x: torch.Tensor):
         bias = self.base_layer.bias if not self.base_layer.skip_bias_add else None
@@ -1039,6 +1085,17 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
         self, B: torch.Tensor, tp_rank: int, target_module: str
     ) -> torch.Tensor:
         if target_module == "gate_up_proj_moe":
+            # Non-gated MoE (e.g. Nemotron-H): only w1, no w3.
+            # B has shape [intermediate_size, rank] — TP-shard directly.
+            is_gated = self.base_layer.moe_runner_config.is_gated
+            if not is_gated:
+                if self.tp_size > 1:
+                    shard_size = self.intermediate_size_per_partition
+                    start = tp_rank * shard_size
+                    end = start + shard_size
+                    return B[start:end, :]
+                return B
+
             shard_size = self.intermediate_size_per_partition
             start = tp_rank * shard_size
             end = start + shard_size

python/sglang/srt/lora/lora.py

@@ -19,7 +19,8 @@
 # https://github.com/vllm-project/vllm/blob/4abf6336ec65c270343eb895e7b18786e9274176/vllm/lora/layers.py
 
 import logging
-from typing import Dict, List
+import re
+from typing import Dict, List, Optional
 
 import torch
 from torch import nn
@@ -27,11 +28,15 @@ from torch import nn
 from sglang.srt.configs.load_config import LoadConfig
 from sglang.srt.layers.utils import get_layer_id
 from sglang.srt.lora.backend.base_backend import BaseLoRABackend
-from sglang.srt.lora.backend.lora_registry import LORA_SUPPORTED_BACKENDS
 from sglang.srt.lora.lora_config import LoRAConfig
 from sglang.srt.model_loader.loader import DefaultModelLoader
 from sglang.srt.utils.hf_transformers_utils import AutoConfig
 
+# Matches both per-expert keys ("...experts.0.<module>...") and shared-outer
+# keys ("...experts.<module>..."), while excluding "shared_experts." (where the
+# preceding char is "_", not ".").
+_ROUTED_EXPERT_PATTERN = re.compile(r"\.experts\.")
+
 logger = logging.getLogger(__name__)
 
 
@@ -53,6 +58,7 @@ class LoRAAdapter(nn.Module):
         base_hf_config: AutoConfig,
         load_config: LoadConfig,
         lora_backend: BaseLoRABackend,
+        base_model: Optional[torch.nn.Module] = None,
     ):
         super().__init__()
         self.uid: str = uid
@@ -63,6 +69,16 @@ class LoRAAdapter(nn.Module):
         self.lora_backend: BaseLoRABackend = lora_backend
         self.scaling: float = self.config.lora_alpha / self.config.r
 
+        # Bypass nn.Module.__setattr__ so the base model is held as a plain
+        # reference rather than auto-registered as a submodule (which would
+        # leak its parameters into our state_dict / parameters() / .to()).
+        object.__setattr__(self, "base_model", base_model)
+        object.__setattr__(
+            self,
+            "_moe_is_gated_by_layer",
+            self._build_moe_gated_map(base_model) if base_model is not None else {},
+        )
+
         self.layers: List[LoRALayer] = nn.ModuleList(
             [
                 LoRALayer(config, base_hf_config)
@@ -73,6 +89,54 @@ class LoRAAdapter(nn.Module):
         self.embedding_layers: Dict[str, torch.Tensor] = {}
         self.added_tokens_embeddings: Dict[str, torch.Tensor] = {}
 
+    @staticmethod
+    def _build_moe_gated_map(base_model: torch.nn.Module) -> Dict[int, bool]:
+        """Map layer_id -> moe_runner_config.is_gated for FusedMoE base layers.
+
+        Only used by normalize_gate_up_proj to decide whether per-expert
+        gate_proj weights should be zero-padded and stacked (gated → c=2 buffer)
+        or just renamed (non-gated → c=1 buffer via model's get_stacked_multiply
+        override on gate_up_proj_moe).
+
+        Adapters can be loaded both before `init_lora_modules` (initial
+        --lora-paths) and after (dynamic API loads), so the FusedMoE may
+        appear either directly or under a `BaseLayerWithLoRA.base_layer`.
+        """
+        from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE
+
+        gated_map: Dict[int, bool] = {}
+        for name, module in base_model.named_modules():
+            inner = (
+                module
+                if isinstance(module, FusedMoE)
+                else getattr(module, "base_layer", None)
+            )
+            if not isinstance(inner, FusedMoE):
+                continue
+            layer_id = get_layer_id(name)
+            if layer_id is not None:
+                gated_map[layer_id] = bool(inner.moe_runner_config.is_gated)
+        return gated_map
+
+    def _is_non_gated_moe_weight(self, weight_name: str) -> bool:
+        """True iff this adapter weight targets a non-gated MoE expert.
+
+        Such weights flow into the `gate_up_proj_moe` buffer, which the model
+        overrides to stacked_multiply=1 — so the weight must be stored without
+        being stacked with a synthetic up_proj zero-pad.
+
+        Matches both adapter key conventions:
+        - per-expert: ``...experts.0.<module>...`` (one tensor per expert)
+        - shared-outer: ``...experts.<module>...`` (3D tensor with the expert
+          dim baked into the shape)
+        """
+        if not _ROUTED_EXPERT_PATTERN.search(weight_name):
+            return False
+        layer_id = get_layer_id(weight_name)
+        if layer_id is None:
+            return False
+        return self._moe_is_gated_by_layer.get(layer_id) is False
+
     def initialize_weights(self):
         model_path = self.config.path
         loader = DefaultModelLoader(self.load_config)
@@ -137,6 +201,10 @@ class LoRAAdapter(nn.Module):
             weight_names = list(layer.weights.keys())
             self.normalize_qkv_proj(weight_names, layer.weights)
             self._rename_expert_w_to_proj(layer.weights)
+            # Stack gate_proj + x_proj → in_proj for Mamba layers (before gate_up normalization)
+            self._normalize_in_proj(layer.weights)
+            # Stack in_proj_q + in_proj_k + in_proj_v + in_proj_z → in_proj_qkvz for GDN layers
+            self._normalize_in_proj_qkvz(layer.weights)
             weight_names = list(layer.weights.keys())
             self.normalize_gate_up_proj(weight_names, layer.weights)
             weight_names = list(layer.weights.keys())
@@ -212,6 +280,79 @@ class LoRAAdapter(nn.Module):
         for old_name, new_name in renames.items():
             weights[new_name] = weights.pop(old_name)
 
+    def _normalize_in_proj(self, weights: Dict[str, torch.Tensor]):
+        """Stack gate_proj + x_proj → in_proj for Mamba layers.
+
+        Detects Mamba layers by the presence of both gate_proj and x_proj.
+        Must run BEFORE normalize_gate_up_proj to prevent gate_proj from
+        being consumed by the gate+up stacking.
+        """
+        # Find gate_proj weights that have a matching x_proj (Mamba pattern)
+        for weight_name in list(weights.keys()):
+            if "gate_proj" not in weight_name:
+                continue
+            x_name = weight_name.replace("gate_proj", "x_proj")
+            if x_name not in weights:
+                continue
+            # This is a Mamba layer: stack gate_proj + x_proj → in_proj
+            in_proj_name = weight_name.replace("gate_proj", "in_proj")
+            cat_dim = weights[weight_name].dim() - 2
+            weights[in_proj_name] = torch.cat(
+                (weights[weight_name], weights[x_name]), cat_dim
+            )
+            weights.pop(weight_name)
+            weights.pop(x_name)
+
+    def _normalize_in_proj_qkvz(self, weights: Dict[str, torch.Tensor]):
+        """Normalize in_proj_qkvz weights for GDN (GatedDeltaNet) layers like
+        Qwen3.5.
+
+        Two adapter formats are handled:
+
+        1. Split: ``in_proj_q + in_proj_k + in_proj_v + in_proj_z`` are present
+           as separate weights → concatenate them into ``in_proj_qkvz``.
+
+        2. Already-merged: the adapter has a single ``in_proj_qkvz`` weight
+           (PEFT trained against SGLang's fused Linear). The stacked buffer
+           expects four per-slice ``A`` blocks, so repeat ``lora_A`` 4× along
+           the rank dim. ``lora_B`` is already full-output-dim and matches
+           the buffer directly.
+        """
+        for weight_name in list(weights.keys()):
+            if "in_proj_q." in weight_name:
+                k_name = weight_name.replace("in_proj_q", "in_proj_k")
+                v_name = weight_name.replace("in_proj_q", "in_proj_v")
+                z_name = weight_name.replace("in_proj_q", "in_proj_z")
+                if (
+                    k_name not in weights
+                    or v_name not in weights
+                    or z_name not in weights
+                ):
+                    continue
+                qkvz_name = weight_name.replace("in_proj_q", "in_proj_qkvz")
+                cat_dim = weights[weight_name].dim() - 2
+                weights[qkvz_name] = torch.cat(
+                    (
+                        weights[weight_name],
+                        weights[k_name],
+                        weights[v_name],
+                        weights[z_name],
+                    ),
+                    cat_dim,
+                )
+                weights.pop(weight_name)
+                weights.pop(k_name)
+                weights.pop(v_name)
+                weights.pop(z_name)
+            elif "in_proj_qkvz" in weight_name and "lora_A" in weight_name:
+                # Already-merged adapter: replicate the shared A across the 4
+                # stacked slots the buffer expects (q, k, v, z).
+                ndim = weights[weight_name].dim()
+                repeat_dims = [1] * ndim
+                repeat_dims[ndim - 2] = 4
+                weights[weight_name] = weights[weight_name].repeat(*repeat_dims)
+            # else (in_proj_qkvz lora_B, or unrelated): no-op.
+
     def normalize_gate_up_proj(
         self, weight_names: List[str], weights: Dict[str, torch.Tensor]
     ):
@@ -219,20 +360,27 @@ class LoRAAdapter(nn.Module):
             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")
-                if up_name not in weights:
+                # PEFT can ship up_proj in two forms when there's no real
+                # up_proj content: the key may be absent, or present as a
+                # numel-zero placeholder. Treat both as "no up_proj".
+                if up_name not in weights or weights[up_name].numel() == 0:
+                    if self._is_non_gated_moe_weight(weight_name):
+                        # Non-gated MoE expert: the gate_up_proj_moe buffer
+                        # uses stacked_multiply=1 (per model override), so just
+                        # rename without stacking.
+                        weights[gate_up_name] = weights.pop(weight_name)
+                        if up_name in weights:
+                            weights.pop(up_name)
+                        continue
+                    # Gated path: buffer expects stacked [2r, hidden] (c=2);
+                    # synthesize a properly-shaped zero up_proj.
                     weights[up_name] = torch.zeros_like(weights[weight_name])
-                    assert self.lora_backend.name in LORA_SUPPORTED_BACKENDS, (
-                        f"LoRA weight initialization currently only supported for LoRA backends: {', '.join(b for b in LORA_SUPPORTED_BACKENDS)}"
-                        f"Received backend: {self.lora_backend.name}. Please verify your backend configuration "
-                        f"or consider implementing custom initialization logic for other backends."
-                    )
                 cat_dim = weights[weight_name].dim() - 2
                 weights[gate_up_name] = torch.cat(
                     (weights[weight_name], weights[up_name]), cat_dim
                 )
                 weights.pop(weight_name)
-                if up_name in weights:
-                    weights.pop(up_name)
+                weights.pop(up_name)
             elif "gate_up_proj" in weight_name:
                 # If gate_up_proj is already stacked, we normalize it following the SGL convention
                 gate_up_name = weight_name
@@ -242,6 +390,9 @@ class LoRAAdapter(nn.Module):
                     repeat_dims[ndim - 2] = 2
                     weights[gate_up_name] = weights[gate_up_name].repeat(*repeat_dims)
                 # else: no-op as LoRA B weight is already stacked.
+        # Orphan up_proj weights (no matching gate_proj) are kept as-is.
+        # 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]

python/sglang/srt/lora/lora_manager.py

@@ -620,6 +620,7 @@ class LoRAManager:
             self.base_hf_config,
             self.load_config,
             self.lora_backend,
+            base_model=self.base_model,
         )
         lora_adapter.initialize_weights()
 
@@ -641,6 +642,7 @@ class LoRAManager:
             self.base_hf_config,
             self.load_config,
             self.lora_backend,
+            base_model=self.base_model,
         )
         lora_adapter.initialize_weights_from_tensors(tensors)
         self.loras[lora_ref.lora_id] = lora_adapter

python/sglang/srt/lora/lora_moe_runners.py

@@ -465,18 +465,23 @@ def _add_lora_gate_up_delta(
     r = lora_info.max_lora_rank
     gate_up_a = lora_info.gate_up_lora_a_weights
     gate_up_b = lora_info.gate_up_lora_b_weights
-    inter_size = gate_up_b.shape[2] // 2
-    M, top_k, gate_up_dim = intermediate_cache.shape
-    r = lora_info.max_lora_rank
-    gate_up_a = lora_info.gate_up_lora_a_weights
-    gate_up_b = lora_info.gate_up_lora_b_weights
-    inter_size = gate_up_b.shape[2] // 2
 
     if lora_info.experts_shared_outer_loras and not lora_info.lora_use_virtual_experts:
         gate_up_a = gate_up_a.expand(-1, lora_info.num_experts, -1, -1)
-    inter_size = gate_up_b.shape[2] // 2
-    lora_a_stacked = [gate_up_a[:, :, :r, :], gate_up_a[:, :, r : 2 * r, :]]
-    lora_b_stacked = [gate_up_b[:, :, :inter_size, :], gate_up_b[:, :, inter_size:, :]]
+
+    # Detect gated vs non-gated from A buffer rank dimension.
+    # Gated: A has 2*r rows (gate + up). Non-gated: A has 1*r rows (w1 only).
+    is_gated = gate_up_a.shape[2] > r
+    if is_gated:
+        inter_size = gate_up_b.shape[2] // 2
+        lora_a_stacked = [gate_up_a[:, :, :r, :], gate_up_a[:, :, r : 2 * r, :]]
+        lora_b_stacked = [
+            gate_up_b[:, :, :inter_size, :],
+            gate_up_b[:, :, inter_size:, :],
+        ]
+    else:
+        lora_a_stacked = [gate_up_a]
+        lora_b_stacked = [gate_up_b]
 
     if lora_info.lora_use_virtual_experts:
         merged_experts_fused_moe_lora_add(

python/sglang/srt/lora/mem_pool.py

@@ -227,6 +227,16 @@ class LoRAMemoryPool:
             or 1
         )
 
+    @staticmethod
+    def _has_moe_module(base_model: torch.nn.Module) -> bool:
+        # Config-only detection isn't reliable: some dense configs (e.g.
+        # `Qwen3_5TextConfig`) inherit `num_experts > 1` from an MoE parent.
+        # Walk the loaded model for an actual FusedMoE instance before we
+        # commit to allocating 4D per-expert LoRA buffers.
+        from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE
+
+        return any(isinstance(m, FusedMoE) for m in base_model.modules())
+
     def _get_num_local_experts(self, base_model: torch.nn.Module) -> int:
         """Experts owned by this rank. Equals the global count when EP is
         off, the runner keeps global IDs, or the split isn't even (all
@@ -285,7 +295,7 @@ class LoRAMemoryPool:
         input_dim, _ = get_hidden_dim(
             module_name, self.base_hf_config, base_model, layer_idx
         )
-        c = get_stacked_multiply(module_name)
+        c = get_stacked_multiply(module_name, base_model)
         if self.tp_size > 1 and module_name in ROW_PARALLELISM_LINEAR_LORA_NAMES:
             input_dim = divide(input_dim, self.tp_size)
         return (self.max_loras_per_batch, max_lora_dim * c, input_dim)
@@ -307,7 +317,7 @@ class LoRAMemoryPool:
         input_dim, _ = get_hidden_dim(
             module_name, self.base_hf_config, base_model, layer_idx
         )
-        c = get_stacked_multiply(module_name)
+        c = get_stacked_multiply(module_name, base_model)
         # MoE modules shard along `moe_tp_size`, not the outer `tp_size`.
         effective_tp_size = (
             self.moe_tp_size if self.is_moe_module(module_name) else self.tp_size
@@ -411,6 +421,7 @@ class LoRAMemoryPool:
         )
 
     def init_buffers(self, base_model: torch.nn.Module):
+        self.base_model = base_model
         device = next(base_model.parameters()).device
 
         # Cached once so the per-expert load path doesn't re-walk the HF
@@ -429,7 +440,7 @@ class LoRAMemoryPool:
                 hasattr(cfg, "shared_expert_intermediate_size")
                 and cfg.shared_expert_intermediate_size > 0
             ) or (getattr(cfg, "n_shared_experts", 0) or 0) > 0
-            has_moe = self._get_num_experts(base_model) > 1
+            has_moe = self._has_moe_module(base_model)
 
             # Shape functions automatically handle both 3D (standard) and 4D (MoE)
             target_modules = target_modules - set(EMBEDDING_NAMES)
@@ -786,7 +797,7 @@ class LoRAMemoryPool:
                 )
 
             for name, weights in temp_A_buffer.items():
-                c = get_stacked_multiply(name)
+                c = get_stacked_multiply(name, self.base_model)
                 max_r = self.max_lora_rank
                 target_buffer = self.A_buffer[name][layer_id]
 

python/sglang/srt/lora/triton_ops/chunked_sgmv_expand.py

@@ -12,12 +12,16 @@ from sglang.srt.utils import cached_triton_kernel
 @cached_triton_kernel(
     lambda _, kwargs: (kwargs["NUM_SLICES"], kwargs["BLOCK_M"], kwargs["OUTPUT_DIM"])
 )
-@triton.jit(do_not_specialize=["num_segs"])
+@triton.jit(do_not_specialize=["num_segs", "output_stride_0", "output_stride_1"])
 def _chunked_lora_expand_kernel(
     # Pointers to matrices
     x,
     weights,
     output,
+    # Output strides may differ from OUTPUT_DIM when compact LoRA output is
+    # accumulated into a wider base projection.
+    output_stride_0,
+    output_stride_1,
     # Information on sequence lengths and weight id
     seg_indptr,
     weight_indices,
@@ -50,10 +54,8 @@ def _chunked_lora_expand_kernel(
         weights (Tensor): The LoRA B weights for all adapters.
             Shape: (num_lora, output_dim, K).
         output (Tensor): The output tensor where the result is stored.
-            Shape: (s, output_dim).
+            Shape: (s, output_dim) or a wider base output.
     """
-    tl.static_assert(NUM_SLICES <= 3)
-
     x_stride_0: tl.constexpr = NUM_SLICES * MAX_RANK
     x_stride_1: tl.constexpr = 1
 
@@ -61,9 +63,6 @@ def _chunked_lora_expand_kernel(
     w_stride_1: tl.constexpr = MAX_RANK
     w_stride_2: tl.constexpr = 1
 
-    output_stride_0: tl.constexpr = OUTPUT_DIM
-    output_stride_1: tl.constexpr = 1
-
     pid_s = tl.program_id(axis=2)
     if pid_s >= num_segs:
         return
@@ -210,6 +209,8 @@ def chunked_sgmv_lora_expand_forward(
         x=x,
         weights=weights,
         output=output,
+        output_stride_0=output.stride(0),
+        output_stride_1=output.stride(1),
         seg_indptr=batch_info.seg_indptr,
         weight_indices=batch_info.weight_indices,
         lora_ranks=batch_info.lora_ranks,

python/sglang/srt/lora/triton_ops/qkv_lora_b.py

@@ -145,12 +145,13 @@ def qkv_lora_b_fwd(
     output_offset: torch.Tensor,
     max_qkv_out_dim: int,
     base_output: torch.Tensor = None,
+    n_slices: int = 3,
 ) -> torch.Tensor:
 
-    # x: (s, 3 * r)
+    # x: (s, n_slices * r)
     # qkv_lora_b: (num_lora, output_dim_q + 2 * output_dim_kv, r)
     # output_offset = [0, output_dim_q, output_dim_q + output_dim_kv,
-    #                     output_dim_q + 2 * output_dim_kv]
+    #                     output_dim_q + 2 * output_dim_kv]  (length n_slices + 1)
     # max_qkv_out_dim = max(output_dim_q, output_dim_kv)
     # output: (s, output_dim_q + 2 * output_dim_kv)
 
@@ -166,8 +167,8 @@ def qkv_lora_b_fwd(
     input_dim = x.shape[1]
     r = qkv_lora_b.shape[-1]
     output_dim = qkv_lora_b.shape[-2]
-    assert input_dim == 3 * r
-    assert output_offset.shape[0] == 4
+    assert input_dim == n_slices * r
+    assert output_offset.shape[0] == n_slices + 1
 
     BLOCK_S = 16
     BLOCK_R = 16
@@ -176,7 +177,7 @@ def qkv_lora_b_fwd(
     grid_b = (
         triton.cdiv(batch_info.max_len, BLOCK_S)
         * triton.cdiv(max_qkv_out_dim, BLOCK_OUT),
-        3,  # this dimension decides current block computes on q, k or v
+        n_slices,
         batch_info.bs,
     )
 

python/sglang/srt/lora/utils.py

@@ -177,6 +177,7 @@ def get_normalized_target_modules(
         "v_proj": "qkv_proj",
         "gate_proj": "gate_up_proj",
         "up_proj": "gate_up_proj",
+        "out_proj": "out_proj",
         "embed_tokens": "embed_tokens",
         "vocab_emb": "embed_tokens",
         "embeddings": "embed_tokens",
@@ -196,14 +197,21 @@ def get_normalized_target_modules(
     return result
 
 
-def get_stacked_multiply(module_name: str) -> int:
+def get_stacked_multiply(
+    module_name: str, base_model: Optional[torch.nn.Module] = None
+) -> int:
     """
-    Mapping a lora module name to its magnification at output dimension
+    Mapping a lora module name to its magnification at output dimension.
+    Models can override via a get_stacked_multiply(module_name) method.
     """
+    if base_model is not None and hasattr(base_model, "get_stacked_multiply"):
+        return base_model.get_stacked_multiply(module_name)
     stacked_rank = {
         "qkv_proj": 3,
+        "in_proj_qkvz": 4,  # GDN packed input projection
         "gate_up_proj": 2,
         "gate_up_proj_moe": 2,
+        "in_proj": 2,
         "fused_qkv_a_proj_with_mqa": 2,
     }
     return stacked_rank[module_name] if module_name in stacked_rank else 1
@@ -215,18 +223,29 @@ def get_target_module_name(full_module_name: str, target_modules: Set[str]) -> s
 
     If there is a target module name in target_modules that can match full_module_name, return this name
     Else raise ValueError.
+
+    When multiple target modules match (e.g. both "up_proj" and "gate_up_proj"
+    are substrings), the longest match wins to avoid ambiguity.
     """
+    best = None
     for target_module in target_modules:
         if target_module in full_module_name:
-            return target_module
+            if best is None or len(target_module) > len(best):
+                best = target_module
+    if best is not None:
+        return best
     raise ValueError(
         f"Cannot find target module name for {full_module_name} in {target_modules}"
     )
 
 
 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"]
+ROW_PARALLELISM_LINEAR_LORA_NAMES = ["o_proj", "out_proj", "down_proj", "down_proj_moe"]
+REPLICATED_LINEAR_LORA_NAMES = [
+    "fused_qkv_a_proj_with_mqa",
+    "fc1_latent_proj",
+    "fc2_latent_proj",
+]
 
 # Normalized module names that the LoRA system fully supports
 # (i.e. get_hidden_dim, init_buffers, and init_lora_modules can handle them).
@@ -234,8 +253,14 @@ _KNOWN_LORA_TARGET_MODULES = frozenset(
     {
         "qkv_proj",
         "o_proj",
+        "out_proj",
+        "in_proj",
+        "in_proj_qkvz",
+        "up_proj",
         "gate_up_proj",
         "down_proj",
+        "fc1_latent_proj",
+        "fc2_latent_proj",
         "embed_tokens",
         "lm_head",
         "fused_qkv_a_proj_with_mqa",
@@ -271,7 +296,15 @@ def auto_detect_lora_target_modules(model: "torch.nn.Module") -> set:
             raw_names.add(name.split(".")[-1])
 
     normalized = get_normalized_target_modules(raw_names)
-    return normalized & _KNOWN_LORA_TARGET_MODULES
+    result = normalized & _KNOWN_LORA_TARGET_MODULES
+
+    # Allow models to declare additional LoRA-compatible modules that
+    # cannot be auto-discovered or need to bypass normalization
+    # (e.g. Mamba in_proj, non-gated up_proj).
+    if hasattr(model, "supported_lora_modules"):
+        result.update(set(model.supported_lora_modules) & _KNOWN_LORA_TARGET_MODULES)
+
+    return result
 
 
 def get_lm_head_lora_b_shard_size(output_dim: int, shard_indices=None) -> int:

python/sglang/srt/models/nemotron_h.py

@@ -665,6 +665,17 @@ class NemotronHForCausalLM(nn.Module):
     packed_modules_mapping = {
         "qkv_proj": ["q_proj", "k_proj", "v_proj"],
     }
+    supported_lora_modules = [
+        "qkv_proj",
+        "o_proj",
+        "out_proj",
+        "in_proj",
+        "up_proj",
+        "gate_up_proj",
+        "down_proj",
+        "fc1_latent_proj",
+        "fc2_latent_proj",
+    ]
 
     remap_prefix = {"backbone": "model"}
     remap_substr = {
@@ -748,6 +759,100 @@ class NemotronHForCausalLM(nn.Module):
     def get_input_embeddings(self) -> VocabParallelEmbedding:
         return self.model.embed_tokens
 
+    def get_stacked_multiply(self, module_name):
+        """Non-gated MoE uses stacked_multiply=1 for gate_up_proj_moe."""
+        if module_name == "gate_up_proj_moe":
+            return 1  # Non-gated: only w1, no w3
+        # Fall back to defaults for everything else
+        from sglang.srt.lora.utils import get_stacked_multiply
+
+        return get_stacked_multiply(module_name)
+
+    def get_hidden_dim(self, module_name, layer_idx):
+        """Return (input_dim, output_dim) for LoRA buffers, per layer type."""
+        config = self.config
+        layer_type = config.layers_block_type[layer_idx]
+        hidden_size = config.hidden_size
+        head_dim = getattr(
+            config, "head_dim", hidden_size // config.num_attention_heads
+        )
+
+        if module_name == "qkv_proj":
+            return (
+                hidden_size,
+                head_dim
+                * (config.num_attention_heads + config.num_key_value_heads * 2),
+            )
+        elif module_name == "o_proj":
+            return (
+                head_dim * config.num_attention_heads,
+                hidden_size,
+            )
+        elif module_name == "out_proj":
+            # Mamba out_proj: RowParallelLinear from mamba_intermediate to hidden_size
+            mamba_intermediate = config.mamba_num_heads * config.mamba_head_dim
+            return mamba_intermediate, hidden_size
+        elif module_name == "gate_up_proj":
+            if layer_type == "mamba":
+                # Mamba in_proj gate component: output = mamba_num_heads * mamba_head_dim
+                mamba_intermediate = config.mamba_num_heads * config.mamba_head_dim
+                return hidden_size, mamba_intermediate * 2
+            elif layer_type == "moe":
+                # Shared expert: only has up_proj (no gate), but gets stacked
+                shared_inter = (
+                    config.moe_shared_expert_intermediate_size * config.n_shared_experts
+                )
+                return hidden_size, shared_inter * 2
+            else:
+                # MLP layer
+                return hidden_size, config.intermediate_size * 2
+        elif module_name == "up_proj":
+            if layer_type == "moe":
+                shared_inter = (
+                    config.moe_shared_expert_intermediate_size * config.n_shared_experts
+                )
+                return hidden_size, shared_inter
+            else:
+                return hidden_size, config.intermediate_size
+        elif module_name == "down_proj":
+            if layer_type == "moe":
+                shared_inter = (
+                    config.moe_shared_expert_intermediate_size * config.n_shared_experts
+                )
+                return shared_inter, hidden_size
+            else:
+                return config.intermediate_size, hidden_size
+        elif module_name == "in_proj":
+            # Mamba in_proj: gate_proj + x_proj, each mamba_intermediate wide
+            mamba_intermediate = config.mamba_num_heads * config.mamba_head_dim
+            return hidden_size, mamba_intermediate * 2
+        elif module_name == "x_proj":
+            # Mamba x_proj: projects from hidden_size to mamba_intermediate
+            mamba_intermediate = config.mamba_num_heads * config.mamba_head_dim
+            return hidden_size, mamba_intermediate
+        elif module_name == "gate_up_proj_moe":
+            # Non-gated MoE: only w1, no w3. stacked_multiply=1.
+            # For latent MoE, experts operate in moe_latent_size space.
+            moe_hidden = getattr(config, "moe_latent_size", None) or hidden_size
+            return moe_hidden, config.moe_intermediate_size
+        elif module_name == "down_proj_moe":
+            moe_hidden = getattr(config, "moe_latent_size", None) or hidden_size
+            return config.moe_intermediate_size, moe_hidden
+        elif module_name == "fc1_latent_proj":
+            moe_latent = getattr(config, "moe_latent_size", None) or hidden_size
+            return hidden_size, moe_latent
+        elif module_name == "fc2_latent_proj":
+            moe_latent = getattr(config, "moe_latent_size", None) or hidden_size
+            return moe_latent, hidden_size
+        elif module_name == "embed_tokens":
+            return config.vocab_size, hidden_size
+        elif module_name == "lm_head":
+            return hidden_size, config.vocab_size
+        else:
+            raise NotImplementedError(
+                f"get_hidden_dim not implemented for {module_name}"
+            )
+
     @torch.no_grad()
     def forward(
         self,

python/sglang/srt/models/qwen3_5.py

@@ -945,6 +945,65 @@ class Qwen3_5ForCausalLM(nn.Module):
         "in_proj_ba": ["in_proj_b", "in_proj_a"],
     }
 
+    supported_lora_modules = [
+        "qkv_proj",
+        "o_proj",
+        "out_proj",
+        "in_proj_qkvz",
+        "gate_up_proj",
+        "down_proj",
+        "lm_head",
+    ]
+
+    def get_hidden_dim(self, module_name: str, layer_idx: int):
+        config = self.config
+        head_dim = config.head_dim or (config.hidden_size // config.num_attention_heads)
+
+        if module_name == "qkv_proj":
+            attn_output_gate = getattr(config, "attn_output_gate", True)
+            q_heads = config.num_attention_heads * (2 if attn_output_gate else 1)
+            return (
+                config.hidden_size,
+                head_dim * (q_heads + config.num_key_value_heads * 2),
+            )
+        elif module_name == "o_proj":
+            return config.num_attention_heads * head_dim, config.hidden_size
+        elif module_name == "out_proj":
+            value_dim = config.linear_value_head_dim * config.linear_num_value_heads
+            return value_dim, config.hidden_size
+        elif module_name == "in_proj_qkvz":
+            key_dim = config.linear_key_head_dim * config.linear_num_key_heads
+            value_dim = config.linear_value_head_dim * config.linear_num_value_heads
+            return config.hidden_size, key_dim * 2 + value_dim * 2
+        elif module_name == "gate_up_proj":
+            # MoE: shared expert uses shared_expert_intermediate_size
+            # Dense: regular MLP uses intermediate_size
+            is_moe = "moe" in getattr(config, "model_type", "")
+            if is_moe:
+                inter = config.shared_expert_intermediate_size
+            else:
+                inter = config.intermediate_size
+            return config.hidden_size, inter * 2
+        elif module_name == "down_proj":
+            is_moe = "moe" in getattr(config, "model_type", "")
+            if is_moe:
+                inter = config.shared_expert_intermediate_size
+            else:
+                inter = config.intermediate_size
+            return inter, config.hidden_size
+        elif module_name == "gate_up_proj_moe":
+            return config.hidden_size, config.moe_intermediate_size * 2
+        elif module_name == "down_proj_moe":
+            return config.moe_intermediate_size, config.hidden_size
+        elif module_name == "embed_tokens":
+            return config.vocab_size, config.hidden_size
+        elif module_name == "lm_head":
+            return config.hidden_size, config.vocab_size
+        else:
+            raise NotImplementedError(
+                f"get_hidden_dim not implemented for {module_name}"
+            )
+
     def __init__(
         self,
         config: Qwen3_5TextConfig,
@@ -1386,6 +1445,8 @@ class Qwen3_5ForConditionalGeneration(Qwen3VLForConditionalGeneration):
     packed_modules_mapping = Qwen3_5ForCausalLM.packed_modules_mapping
     hf_to_sglang_mapper = None
 
+    supported_lora_modules = Qwen3_5ForCausalLM.supported_lora_modules
+
     def __init__(
         self,
         config: Qwen3_5Config,
@@ -1402,6 +1463,12 @@ class Qwen3_5ForConditionalGeneration(Qwen3VLForConditionalGeneration):
 
         self.deepstack_visual_indexes = self.visual.deepstack_visual_indexes
 
+    def get_hidden_dim(self, module_name: str, layer_idx: int):
+        return self.model.get_hidden_dim(module_name, layer_idx)
+
+    def should_apply_lora(self, module_name: str) -> bool:
+        return module_name.startswith("model.layers.")
+
     @property
     def start_layer(self) -> int:
         return getattr(getattr(self, "model", None), "start_layer", 0)
@@ -1532,6 +1599,8 @@ class Qwen3_5MoeForConditionalGeneration(Qwen3VLForConditionalGeneration):
     packed_modules_mapping = Qwen3_5ForCausalLM.packed_modules_mapping
     hf_to_sglang_mapper = None
 
+    supported_lora_modules = Qwen3_5ForCausalLM.supported_lora_modules
+
     def __init__(
         self,
         config: Qwen3_5MoeConfig,
@@ -1552,6 +1621,13 @@ class Qwen3_5MoeForConditionalGeneration(Qwen3VLForConditionalGeneration):
 
         self.enable_shared_expert_fusion = self.num_fused_shared_experts > 0
 
+    def get_hidden_dim(self, module_name: str, layer_idx: int):
+        return self.model.get_hidden_dim(module_name, layer_idx)
+
+    def should_apply_lora(self, module_name: str) -> bool:
+        # Accept all language model layer modules (attention, linear_attn, mlp).
+        return module_name.startswith("model.layers.")
+
     def _get_num_fused_shared_experts(self):
         if not (
             hasattr(self.model, "layers")

test/registered/lora/test_chunked_sgmv_backend.py

@@ -614,6 +614,110 @@ class TestChunkedSGMV(unittest.TestCase):
                     f"QKV missing k_proj batch_size={batch_size}",
                 )
 
+    def test_4_slice_gdn_qkvz(self):
+        """Test 4-slice shrink+expand operations (GDN in_proj_qkvz)."""
+        num_slices = 4
+        # GDN-style: 4 slices with different sizes [2048, 2048, 4096, 4096]
+        slice_offsets = torch.tensor(
+            [0, 2048, 4096, 8192, 12288], dtype=torch.int32, device=self.device
+        )
+        total_out = 12288
+        max_slice_size = 4096
+
+        for batch_size in [1, 2, 16]:
+            with self.subTest(batch_size=batch_size):
+                # Build batch (reuse the 3-slice helper just for x, batch_info, etc.)
+                x, _, batch_info, seq_lengths, lora_assignments = (
+                    self.create_test_batch(BatchComposition.MIXED, batch_size)
+                )
+
+                # Build 4-slice LoRA weights and stack them manually
+                lora_names = [n for n in self.lora_configs if n != "_NO_LORA_"]
+                max_rank = max(self.lora_configs[n][0] for n in lora_names)
+                stacked_a = torch.zeros(
+                    len(lora_names),
+                    num_slices * max_rank,
+                    self.input_dim,
+                    dtype=self.dtype,
+                    device=self.device,
+                )
+                stacked_b = torch.zeros(
+                    len(lora_names),
+                    total_out,
+                    max_rank,
+                    dtype=self.dtype,
+                    device=self.device,
+                )
+                for i, name in enumerate(lora_names):
+                    rank = self.lora_configs[name][0]
+                    if rank > 0:
+                        stacked_a[i, : num_slices * rank, :] = torch.randn(
+                            num_slices * rank,
+                            self.input_dim,
+                            dtype=self.dtype,
+                            device=self.device,
+                        )
+                        stacked_b[i, :, :rank] = torch.randn(
+                            total_out, rank, dtype=self.dtype, device=self.device
+                        )
+
+                lora_assignments_tensor = torch.tensor(
+                    lora_assignments, dtype=torch.int32, device="cpu"
+                )
+                seq_lengths_tensor = torch.tensor(
+                    seq_lengths, dtype=torch.int32, device="cpu"
+                )
+                lora_ranks_tensor = batch_info.lora_ranks.detach().cpu()
+                scalings_tensor = batch_info.scalings.detach().cpu()
+
+                # Shrink
+                chunked_shrink = chunked_sgmv_lora_shrink_forward(
+                    x, stacked_a, batch_info, num_slices=num_slices
+                )
+                reference_shrink = reference_sgmv_shrink(
+                    x,
+                    stacked_a,
+                    lora_assignments_tensor,
+                    seq_lengths_tensor,
+                    lora_ranks_tensor,
+                    scalings_tensor,
+                    num_slices=num_slices,
+                )
+                self._compare_shrink_outputs(
+                    chunked_shrink,
+                    reference_shrink,
+                    seq_lengths,
+                    lora_assignments,
+                    batch_info,
+                    num_slices=num_slices,
+                    test_name=f"4-slice shrink bs={batch_size}",
+                )
+
+                # Expand
+                chunked_expand = chunked_sgmv_lora_expand_forward(
+                    reference_shrink,
+                    stacked_b,
+                    batch_info,
+                    slice_offsets,
+                    max_slice_size,
+                    base_output=None,
+                )
+                reference_expand = reference_sgmv_expand(
+                    reference_shrink,
+                    stacked_b,
+                    lora_assignments_tensor,
+                    seq_lengths_tensor,
+                    lora_ranks_tensor,
+                    slice_offsets,
+                )
+                torch.testing.assert_close(
+                    chunked_expand,
+                    reference_expand,
+                    rtol=self.RTOL,
+                    atol=self.ATOL,
+                    msg=f"4-slice expand failed bs={batch_size}",
+                )
+
     # === Batch Composition Tests ===
 
     def test_uniform_lora_batch(self):

test/registered/lora/test_lora_nemotron_3_super_120b_a12b_logprob_diff.py

@@ -0,0 +1,154 @@
+# 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 NVIDIA-Nemotron-3-Super-120B-A12B-BF16 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/opherlie/lora-test-case-NVIDIA-Nemotron-3-Super-120B-A12B-BF16
+
+Usage:
+    python -m unittest test_lora_nemotron_3_super_120b_a12b_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="stage-c-test-4-gpu-b200-small",
+)
+
+BASE_MODEL = "nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16"
+LORA_HF_REPO = "opherlie/lora-test-case-NVIDIA-Nemotron-3-Super-120B-A12B-BF16"
+LORA_BACKEND = "triton"
+MAX_LORA_RANK = 64
+TP_SIZE = 4
+MOE_RUNNER_BACKEND = "triton"
+EXPERTS_SHARED_OUTER_LORAS = True
+LORA_USE_VIRTUAL_EXPERTS = True
+DISABLE_SHARED_EXPERTS_FUSION = True
+
+KL_THRESHOLD = 2.5e-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 TestLoRANemotron3Super120B_A12B_LogprobDiff(CustomTestCase):
+
+    def test_lora_nemotron_3_super_120b_a12b_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,
+            moe_runner_backend=MOE_RUNNER_BACKEND,
+            experts_shared_outer_loras=EXPERTS_SHARED_OUTER_LORAS,
+            lora_use_virtual_experts=LORA_USE_VIRTUAL_EXPERTS,
+            disable_shared_experts_fusion=DISABLE_SHARED_EXPERTS_FUSION,
+        )
+
+        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()

test/registered/lora/test_lora_qwen3_5_35b_a3b_logprob_diff.py

@@ -0,0 +1,157 @@
+# 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 Qwen3.5-35B-A3B 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/opherlie/lora-test-case-Qwen3.5-35B-A3B
+
+Usage:
+    python -m unittest test_lora_qwen3_5_35b_a3b_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=160,
+    suite="stage-c-test-4-gpu-b200-small",
+)
+
+BASE_MODEL = "Qwen/Qwen3.5-35B-A3B"
+LORA_HF_REPO = "opherlie/lora-test-case-Qwen3.5-35B-A3B"
+LORA_BACKEND = "triton"
+MAX_LORA_RANK = 64
+TP_SIZE = 4
+MOE_RUNNER_BACKEND = "triton"
+EXPERTS_SHARED_OUTER_LORAS = True
+LORA_USE_VIRTUAL_EXPERTS = True
+DISABLE_SHARED_EXPERTS_FUSION = True
+
+KL_THRESHOLD = 1e-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 TestLoRAQwen3_5_35B_A3B_LogprobDiff(CustomTestCase):
+
+    def test_lora_qwen3_5_35b_a3b_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,
+            moe_runner_backend=MOE_RUNNER_BACKEND,
+            experts_shared_outer_loras=EXPERTS_SHARED_OUTER_LORAS,
+            lora_use_virtual_experts=LORA_USE_VIRTUAL_EXPERTS,
+            disable_shared_experts_fusion=DISABLE_SHARED_EXPERTS_FUSION,
+            # OOM's on logits with the defaults here, as this test-case is longer context and qwen3.5 has larger vocab
+            chunked_prefill_size=8192,
+            mem_fraction_static=0.8,
+        )
+
+        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()

test/registered/lora/test_lora_qwen3_5_4b_logprob_diff.py

@@ -0,0 +1,146 @@
+# 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 Qwen3.5-4B 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/opherlie/lora-test-case-Qwen3.5-4B
+
+Usage:
+    python -m unittest test_lora_qwen3_5_4b_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=90,
+    suite="stage-b-test-1-gpu-large",
+)
+
+BASE_MODEL = "Qwen/Qwen3.5-4B"
+LORA_HF_REPO = "opherlie/lora-test-case-Qwen3.5-4B"
+LORA_BACKEND = "triton"
+MAX_LORA_RANK = 64
+TP_SIZE = 1
+
+KL_THRESHOLD = 4e-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 TestLoRAQwen3_5_4BLogprobDiff(CustomTestCase):
+
+    def test_lora_qwen3_5_4b_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,
+        )
+
+        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()