diff --git a/python/sglang/srt/layers/cp/cp_decode_attn_tp.py b/python/sglang/srt/layers/cp/cp_decode_attn_tp.py index 7ad40c137cb0..5ba5778a04af 100644 --- a/python/sglang/srt/layers/cp/cp_decode_attn_tp.py +++ b/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)) diff --git a/python/sglang/srt/lora/backend/base_backend.py b/python/sglang/srt/lora/backend/base_backend.py index 5f70a02f5e0a..8dce14956b21 100644 --- a/python/sglang/srt/lora/backend/base_backend.py +++ b/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, diff --git a/python/sglang/srt/lora/backend/chunked_backend.py b/python/sglang/srt/lora/backend/chunked_backend.py index 92a77b8d77f8..6533a2553f46 100644 --- a/python/sglang/srt/lora/backend/chunked_backend.py +++ b/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( diff --git a/python/sglang/srt/lora/backend/triton_backend.py b/python/sglang/srt/lora/backend/triton_backend.py index f8f0e6d12d0a..b8f30ca751e9 100644 --- a/python/sglang/srt/lora/backend/triton_backend.py +++ b/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" diff --git a/python/sglang/srt/lora/layers.py b/python/sglang/srt/lora/layers.py index f92342f34cfd..78c06a00ffb6 100644 --- a/python/sglang/srt/lora/layers.py +++ b/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 diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index 30ee337f9675..605f2a75b8c4 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/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) diff --git a/test/registered/unit/lora/test_deepseek_v4_grouped_wo_a.py b/test/registered/unit/lora/test_deepseek_v4_grouped_wo_a.py new file mode 100644 index 000000000000..7638e42f1809 --- /dev/null +++ b/test/registered/unit/lora/test_deepseek_v4_grouped_wo_a.py @@ -0,0 +1,179 @@ +import pytest +import torch +from torch import nn + +from sglang.srt.layers.cp.cp_decode_attn_tp import CpDecodeAttnTpContext +from sglang.srt.lora.backend.base_backend import BaseLoRABackend +from sglang.srt.lora.layers import ColumnParallelLinearWithLoRA +from sglang.srt.lora.utils import LoRABatchInfo +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=1, suite="base-a-test-cpu") + + +def _batch_info(*, permutation=None): + return LoRABatchInfo( + use_cuda_graph=False, + bs=2, + num_segments=2, + seg_indptr=torch.tensor([0, 1, 3], dtype=torch.int32), + weight_indices=torch.tensor([0, 1], dtype=torch.int32), + lora_ranks=torch.tensor([3, 3], dtype=torch.int32), + scalings=torch.tensor([1.0, 1.0]), + max_len=2, + seg_lens=torch.tensor([1, 2], dtype=torch.int32), + permutation=permutation, + expected_tokens=3, + ) + + +class _TorchSegmentedBackend(BaseLoRABackend): + name = "test" + supports_grouped_sgemm_batch_info = True + + def __init__(self, batch_info): + self.batch_info = batch_info + self._grouped_sgemm_batch_info_cache = {} + self.a_input_shapes = [] + self.b_weight_shapes = [] + self.a_batch_infos = [] + self.b_batch_infos = [] + + def _adapter_per_row(self, batch_info, num_rows): + adapters = torch.empty(num_rows, dtype=torch.long) + permutation = batch_info.permutation + for segment in range(batch_info.num_segments): + start = int(batch_info.seg_indptr[segment]) + end = int(batch_info.seg_indptr[segment + 1]) + physical_rows = ( + torch.arange(start, end) + if permutation is None + else permutation[start:end].long() + ) + adapters[physical_rows] = batch_info.weight_indices[segment].long() + return adapters + + def run_lora_a_sgemm(self, x, weights, pruned_batch_info=None, **_kwargs): + self.a_input_shapes.append(tuple(x.shape)) + batch_info = pruned_batch_info or self.batch_info + self.a_batch_infos.append(batch_info) + adapters = self._adapter_per_row(batch_info, x.shape[0]) + return torch.bmm(x.unsqueeze(1), weights[adapters].transpose(1, 2)).squeeze(1) + + def run_lora_b_sgemm(self, x, weights, pruned_batch_info=None, **_kwargs): + self.b_weight_shapes.append(tuple(weights.shape)) + batch_info = pruned_batch_info or self.batch_info + self.b_batch_infos.append(batch_info) + adapters = self._adapter_per_row(batch_info, x.shape[0]) + output = torch.bmm(x.unsqueeze(1), weights[adapters].transpose(1, 2)).squeeze(1) + return output * batch_info.scalings[adapters].unsqueeze(1) + + +def _layer(backend, lora_a, lora_b): + layer = object.__new__(ColumnParallelLinearWithLoRA) + nn.Module.__init__(layer) + layer.base_layer = nn.Linear(lora_a.shape[-1], lora_b.shape[-2], bias=False) + layer.set_lora = True + layer.A_buffer = lora_a + layer.B_buffer = lora_b + layer.lora_backend = backend + layer.output_offset = torch.tensor([0, lora_b.shape[-2]], dtype=torch.int32) + layer.output_offset_cpu = layer.output_offset + return layer + + +def test_grouped_wo_a_lora_selects_matching_group_diagonal(): + torch.manual_seed(0) + tokens, groups, input_dim, output_dim, rank = 3, 4, 5, 2, 3 + x = torch.randn(tokens, groups, input_dim) + base_output = torch.randn(tokens, groups, output_dim) + lora_a = torch.randn(2, rank, input_dim) + lora_b = torch.randn(2, groups * output_dim, rank) + backend = _TorchSegmentedBackend(_batch_info()) + layer = _layer(backend, lora_a, lora_b) + + output = layer.apply_grouped_lora(base_output, x) + + adapter_per_token = torch.tensor([0, 1, 1]) + delta_weight = torch.bmm(lora_b, lora_a).view(2, groups, output_dim, input_dim) + expected = base_output + torch.einsum( + "tgd,tgrd->tgr", x, delta_weight[adapter_per_token] + ) + torch.testing.assert_close(output, expected) + + +def test_grouped_wo_a_preserves_backend_segments_and_chunk_size(): + info = _batch_info(permutation=torch.tensor([2, 0, 1], dtype=torch.int32)) + backend = _TorchSegmentedBackend(info) + groups = 4 + layer = _layer( + backend, + torch.randn(2, 3, 4), + torch.randn(2, groups * 5, 3), + ) + + layer.apply_grouped_lora( + torch.randn(3, groups, 5), + torch.randn(3, groups, 4), + ) + + assert backend.batch_info is info + assert info.max_len == 2 + torch.testing.assert_close( + info.seg_indptr, torch.tensor([0, 1, 3], dtype=torch.int32) + ) + assert backend.a_input_shapes == [(12, 4)] + assert backend.b_weight_shapes == [(8, 5, 3)] + + a_info = backend.a_batch_infos[0] + b_info = backend.b_batch_infos[0] + assert a_info.max_len == b_info.max_len == 2 + assert a_info.expected_tokens == b_info.expected_tokens == 12 + torch.testing.assert_close( + a_info.seg_indptr, + torch.tensor([0, 1, 3, 4, 6, 7, 9, 10, 12], dtype=torch.int32), + ) + torch.testing.assert_close( + a_info.permutation, + torch.tensor([2, 0, 1, 5, 3, 4, 8, 6, 7, 11, 9, 10], dtype=torch.int32), + ) + torch.testing.assert_close( + a_info.weight_indices, + torch.tensor([0, 1, 0, 1, 0, 1, 0, 1], dtype=torch.int32), + ) + torch.testing.assert_close( + b_info.weight_indices, + torch.tensor([0, 4, 1, 5, 2, 6, 3, 7], dtype=torch.int32), + ) + + +def test_grouped_wo_a_rejects_incompatible_b_layout(): + groups = 2 + backend = _TorchSegmentedBackend(_batch_info()) + layer = _layer( + backend, + torch.randn(1, 3, 4), + torch.randn(1, groups * 5 + 1, 3), + ) + + with pytest.raises(RuntimeError, match="B/output mismatch"): + layer.apply_grouped_lora( + torch.randn(3, groups, 5), + torch.randn(3, groups, 4), + ) + + +def test_cp_decode_tp_rejects_active_lora_and_unwraps_inactive_wrapper(): + backend = _TorchSegmentedBackend(_batch_info()) + layer = _layer( + backend, + torch.randn(2, 3, 4), + torch.randn(2, 10, 3), + ) + context = object.__new__(CpDecodeAttnTpContext) + + with pytest.raises(RuntimeError, match="does not support active LoRA"): + context._unwrap_inactive_lora(layer) + + backend.batch_info = None + assert context._unwrap_inactive_lora(layer) is layer.base_layer