[LoRA] Preserve DeepSeek-V4 grouped wo_a semantics

python/sglang/srt/layers/cp/cp_decode_attn_tp.py

@@ -126,6 +126,20 @@ def _restore(self, obj, attr_name: str):
 
     # ==================== Linear helpers ====================
 
+    def _unwrap_inactive_lora(self, linear_instance):
+        """Return the base linear, rejecting active LoRA during decode TP."""
+        from sglang.srt.lora.layers import BaseLayerWithLoRA
+
+        if not isinstance(linear_instance, BaseLayerWithLoRA):
+            return linear_instance
+        if linear_instance.lora_active:
+            raise RuntimeError(
+                "CP decode attention TP does not support active LoRA adapters. "
+                "The replicated LoRA buffers cannot be sliced safely with the "
+                "decode-TP base weights; disable CP decode attention TP or LoRA."
+            )
+        return linear_instance.base_layer
+
     def _get_linear_attrs(self, linear_instance) -> List[Tuple]:
         """Return (obj, attr_name, dim) list for a linear layer."""
         from sglang.srt.layers.linear import ColumnParallelLinear, RowParallelLinear
@@ -171,7 +185,8 @@ def maybe_use_decode_attn_tp(
         row_parallel_decode_flags = []  # (RowParallelLinear, orig_flag) to restore
         orig_tp_q_head_num = None
         try:
-            for linear in modules:
+            for module in modules:
+                linear = self._unwrap_inactive_lora(module)
                 for obj, attr_name, dim in self._get_linear_attrs(linear):
                     self._activate(obj, attr_name, dim)
                     all_attrs.append((obj, attr_name))

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

@@ -1,3 +1,4 @@
+import dataclasses
 from typing import Optional, Tuple, Union
 
 import torch
@@ -26,6 +27,9 @@ class BaseLoRABackend(LoRABackendLmHeadMixing):
     # Supporting backends implement init_prefill_cuda_graph_batch_info() and
     # honor use_prefill_cuda_graph in prepare_lora_batch().
     supports_prefill_cuda_graph: bool = False
+    # Supporting segmented-GEMM backends can apply grouped projections with
+    # group-major routing metadata supplied to both the A and B kernels.
+    supports_grouped_sgemm_batch_info: bool = False
 
     def __init__(self, max_loras_per_batch: int, device: torch.device):
         self.max_loras_per_batch = max_loras_per_batch
@@ -42,6 +46,8 @@ def __init__(self, max_loras_per_batch: int, device: torch.device):
         # Request/token caps for serving a batch from the static metadata.
         self.prefill_cuda_graph_max_bs: int | None = None
         self.prefill_cuda_graph_max_tokens: int | None = None
+        self._grouped_sgemm_batch_info_cache = {}
+        self._grouped_sgemm_capture_state = False
 
     def reset_batch_state(self):
         """Idle-forward counterpart of prepare_lora_batch(): clears all
@@ -51,6 +57,118 @@ def reset_batch_state(self):
         self.lm_head_batch_info = None
         self.lm_head_pass_batch_infos = None
         self._lm_head_pass_idx = None
+        self._reset_grouped_sgemm_batch_info()
+
+    def _reset_grouped_sgemm_batch_info(self):
+        # prepare_lora_batch() calls this once per logical forward. During
+        # CUDA-graph capture, the first grouped layer records the metadata
+        # transforms and later layers reuse their outputs.
+        self._grouped_sgemm_batch_info_cache = {}
+        self._grouped_sgemm_capture_state = False
+
+    def _get_sgemm_batch_info(self) -> LoRABatchInfo:
+        assert self.batch_info is not None, "LoRA batch metadata is not prepared"
+        return self.batch_info
+
+    def get_grouped_sgemm_batch_infos(
+        self, num_groups: int, num_tokens: int
+    ) -> tuple[LoRABatchInfo, LoRABatchInfo]:
+        """Build group-major segmented-GEMM metadata for grouped projections.
+
+        The A pass reuses each token's adapter for every group. The B pass
+        routes to the composite ``adapter * num_groups + group`` weight. Whole
+        token segments are repeated per group, so the backend's original
+        maximum segment length remains unchanged.
+        """
+        if not self.supports_grouped_sgemm_batch_info:
+            raise RuntimeError(
+                f"LoRA backend {self.name!r} does not support grouped "
+                "segmented-GEMM metadata"
+            )
+        if num_groups <= 0 or num_tokens <= 0:
+            raise ValueError(
+                "grouped segmented-GEMM dimensions must be positive, got "
+                f"groups={num_groups}, tokens={num_tokens}"
+            )
+
+        source = self._get_sgemm_batch_info()
+        is_capturing = (
+            source.seg_indptr.is_cuda and torch.cuda.is_current_stream_capturing()
+        )
+        if getattr(self, "_grouped_sgemm_capture_state", False) != is_capturing:
+            # Full CUDA-graph capture invokes the model for eager warmup before
+            # recording it. Rebuild on the eager/capture transition so the
+            # metadata transforms themselves are captured exactly once.
+            self._grouped_sgemm_batch_info_cache = {}
+            self._grouped_sgemm_capture_state = is_capturing
+        key = (id(source), num_groups, num_tokens)
+        cache = getattr(self, "_grouped_sgemm_batch_info_cache", None)
+        if cache is None:
+            cache = self._grouped_sgemm_batch_info_cache = {}
+        if key in cache:
+            return cache[key]
+
+        # CUDA-graph descriptors use fixed-size padded buffers. Include every
+        # allocated slot so captured tensor shapes stay static; padded
+        # zero-length segments remain no-ops.
+        num_segments = (
+            source.weight_indices.shape[0]
+            if source.use_cuda_graph
+            else source.num_segments
+        )
+        if num_segments is None:
+            raise RuntimeError("grouped LoRA batch metadata has no segment count")
+
+        device = source.seg_indptr.device
+        index_dtype = source.weight_indices.dtype
+        group_ids = torch.arange(num_groups, dtype=index_dtype, device=device)
+        token_offsets = group_ids * num_tokens
+        source_starts = source.seg_indptr[:num_segments]
+        segment_starts = (
+            source_starts.unsqueeze(0) + token_offsets.unsqueeze(1)
+        ).reshape(-1)
+        segment_end = source.seg_indptr[:1] + num_tokens * num_groups
+        seg_indptr = torch.cat((segment_starts, segment_end))
+
+        permutation = None
+        if source.permutation is not None:
+            permutation = (
+                source.permutation[:num_tokens].unsqueeze(0)
+                + token_offsets.unsqueeze(1)
+            ).reshape(-1)
+        seg_lens = None
+        if source.seg_lens is not None:
+            seg_lens = source.seg_lens[:num_segments].repeat(num_groups)
+
+        common = {
+            "bs": source.bs * num_groups,
+            "num_segments": num_segments * num_groups,
+            "seg_indptr": seg_indptr,
+            "max_len": source.max_len,
+            "seg_lens": seg_lens,
+            "permutation": permutation,
+            "expected_tokens": num_tokens * num_groups,
+            "req_seg_indptr": None,
+            "req_weight_indices": None,
+            "moe_lora_info": None,
+        }
+        source_weight_indices = source.weight_indices[:num_segments]
+        a_info = dataclasses.replace(
+            source,
+            weight_indices=source_weight_indices.repeat(num_groups),
+            **common,
+        )
+        b_info = dataclasses.replace(
+            source,
+            weight_indices=(
+                source_weight_indices.unsqueeze(0) * num_groups + group_ids.unsqueeze(1)
+            ).reshape(-1),
+            lora_ranks=source.lora_ranks.repeat_interleave(num_groups),
+            scalings=source.scalings.repeat_interleave(num_groups),
+            **common,
+        )
+        cache[key] = (a_info, b_info)
+        return a_info, b_info
 
     def run_lora_a_embedding(
         self,

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

@@ -34,6 +34,7 @@ class ChunkedSgmvLoRABackend(BaseLoRABackend):
 
     name = "csgmv"
     supports_prefill_cuda_graph = True
+    supports_grouped_sgemm_batch_info = True
 
     def __init__(
         self,
@@ -275,6 +276,7 @@ def prepare_lora_batch(
         use_cuda_graph: bool,
         use_prefill_cuda_graph: bool = False,
     ):
+        self._reset_grouped_sgemm_batch_info()
         chunk_size = self._determine_chunk_size(forward_batch)
 
         permutation, weight_indices_reordered = ChunkedSgmvLoRABackend._get_permutation(

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

@@ -26,6 +26,7 @@
 class TritonLoRABackend(BaseLoRABackend):
     name = "triton"
     supports_prefill_cuda_graph = True
+    supports_grouped_sgemm_batch_info = True
 
     def __init__(
         self,
@@ -71,6 +72,9 @@ def _sgemm_info(self, pruned_batch_info=None):
         )
         return self.sgemm_batch_info or self.batch_info
 
+    def _get_sgemm_batch_info(self) -> LoRABatchInfo:
+        return self._sgemm_info()
+
     def run_lora_a_sgemm(
         self,
         x: torch.Tensor,
@@ -273,6 +277,7 @@ def prepare_lora_batch(
         use_cuda_graph: bool,
         use_prefill_cuda_graph: bool = False,
     ):
+        self._reset_grouped_sgemm_batch_info()
         # Use pinned memory to avoid synchronizations during host-to-device transfer
         weight_indices_tensor = torch.tensor(
             weight_indices, dtype=torch.int32, pin_memory=True, device="cpu"

python/sglang/srt/lora/layers.py

@@ -461,6 +461,9 @@ def __init__(
             device="cpu",
             pin_memory=True,
         )
+        self._grouped_output_dim = None
+        self._grouped_output_offset = None
+        self._grouped_output_offset_cpu = None
 
     def set_lora_info(
         self,
@@ -482,6 +485,110 @@ def apply_lora(self, base_output: torch.Tensor, x: torch.Tensor) -> torch.Tensor
         )
         return lora_output
 
+    def _get_grouped_output_offsets(
+        self, group_output_dim: int
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if getattr(self, "_grouped_output_dim", None) != group_output_dim:
+            is_capturing = (
+                self.output_offset.is_cuda and torch.cuda.is_current_stream_capturing()
+            )
+            if is_capturing:
+                raise RuntimeError(
+                    "grouped LoRA output offsets must be initialized during "
+                    "CUDA-graph warmup"
+                )
+            self._grouped_output_dim = group_output_dim
+            self._grouped_output_offset = self.output_offset.new_tensor(
+                [0, group_output_dim]
+            )
+            self._grouped_output_offset_cpu = torch.tensor(
+                [0, group_output_dim],
+                dtype=self.output_offset_cpu.dtype,
+                device="cpu",
+                pin_memory=self.output_offset_cpu.is_pinned(),
+            )
+        return self._grouped_output_offset, self._grouped_output_offset_cpu
+
+    def apply_grouped_lora(
+        self, base_output: torch.Tensor, x: torch.Tensor
+    ) -> torch.Tensor:
+        """Apply a flattened LoRA matrix as independent grouped projections.
+
+        DeepSeek-V4 ``wo_a`` stores ``G`` independent ``[R, D]`` matrices as
+        one column-parallel ``[G * R, D]`` weight. Route the flattened grouped
+        input through matching composite adapter/group B weights so
+        off-diagonal group products are never built.
+        """
+        if not self.lora_active:
+            return base_output
+        if x.ndim != 3 or base_output.ndim != 3:
+            raise RuntimeError(
+                "grouped LoRA expects x=[tokens,groups,input] and "
+                "base_output=[tokens,groups,output]"
+            )
+
+        num_tokens, num_groups, _ = x.shape
+        if base_output.shape[:2] != (num_tokens, num_groups):
+            raise RuntimeError(
+                "grouped LoRA base/input prefix mismatch: "
+                f"x={tuple(x.shape)}, base={tuple(base_output.shape)}"
+            )
+        group_output_dim = base_output.shape[-1]
+        if self.B_buffer.shape[-2] != num_groups * group_output_dim:
+            raise RuntimeError(
+                "grouped LoRA B/output mismatch: "
+                f"B={tuple(self.B_buffer.shape)}, groups={num_groups}, "
+                f"group_output={group_output_dim}"
+            )
+
+        batch_info = self.lora_backend.batch_info
+        if (
+            batch_info.expected_tokens is not None
+            and batch_info.expected_tokens != num_tokens
+        ):
+            raise RuntimeError(
+                "grouped LoRA batch/input token mismatch: "
+                f"metadata={batch_info.expected_tokens}, input={num_tokens}"
+            )
+
+        a_info, b_info = self.lora_backend.get_grouped_sgemm_batch_infos(
+            num_groups, num_tokens
+        )
+        flat_input = x.transpose(0, 1).contiguous().view(num_tokens * num_groups, -1)
+        lora_a_output = self.lora_backend.run_lora_a_sgemm(
+            flat_input,
+            self.A_buffer,
+            pruned_batch_info=a_info,
+        )
+
+        num_adapters, _, lora_rank = self.B_buffer.shape
+        grouped_b = self.B_buffer.view(
+            num_adapters,
+            num_groups,
+            group_output_dim,
+            lora_rank,
+        ).reshape(num_adapters * num_groups, group_output_dim, lora_rank)
+        group_output_offset, group_output_offset_cpu = self._get_grouped_output_offsets(
+            group_output_dim
+        )
+        lora_output = self.lora_backend.run_lora_b_sgemm(
+            x=lora_a_output,
+            weights=grouped_b,
+            output_offset=group_output_offset,
+            output_offset_cpu=group_output_offset_cpu,
+            pruned_batch_info=b_info,
+        )
+        expected_shape = (num_tokens * num_groups, group_output_dim)
+        if tuple(lora_output.shape) != expected_shape:
+            raise RuntimeError(
+                "grouped LoRA kernel output mismatch: "
+                f"got {tuple(lora_output.shape)}, expected {expected_shape}"
+            )
+
+        return base_output + lora_output.view(
+            num_groups, num_tokens, group_output_dim
+        ).transpose(0, 1)
+
     def forward(self, input_: torch.Tensor):
         # duplicate the logic in ColumnParallelLinear
         bias = self.base_layer.bias if not self.base_layer.skip_bias_add else None

python/sglang/srt/models/deepseek_v4.py

@@ -1695,6 +1695,8 @@ def forward(
             )
 
         o = o.view(o.shape[0], self.n_local_groups, -1)
+        wo_a_input = o
+        wo_a_base = getattr(self.wo_a, "base_layer", self.wo_a)
 
         if _FP8_WO_A_GEMM:
             import deep_gemm
@@ -1721,13 +1723,13 @@ def forward(
             deep_gemm.fp8_einsum(
                 "bhr,hdr->bhd",
                 (o_fp8, o_s),
-                (self.wo_a.weight.view(G, R, D), self.wo_a.weight_scale_inv.data),
+                (wo_a_base.weight.view(G, R, D), wo_a_base.weight_scale_inv.data),
                 output,
                 recipe=recipe,
             )
             o = output
         else:
-            wo_a_weight = getattr(self.wo_a, "weight", None)
+            wo_a_weight = getattr(wo_a_base, "weight", None)
             if wo_a_weight is not None:
                 wo_a = wo_a_weight.view(self.n_local_groups, self.o_lora_rank, -1)
                 o = _apply_wo_a_bf16_matmul(
@@ -1736,11 +1738,14 @@ def forward(
             else:
                 o = _apply_gguf_grouped_wo_a(
                     o,
-                    self.wo_a.qweight,
-                    self.wo_a.qweight_type.weight_type,
+                    wo_a_base.qweight,
+                    wo_a_base.qweight_type.weight_type,
                     self.o_lora_rank,
                 )
 
+        if hasattr(self.wo_a, "apply_grouped_lora"):
+            o = self.wo_a.apply_grouped_lora(o, wo_a_input)
+
         o, _ = self.wo_b(o.flatten(1))
         if self.attn_tp_size > 1 and self.attn_tp_size < get_parallel().tp_size:
             o = attn_tp_all_reduce(o)