diff --git a/python/sglang/srt/layers/quantization/base_config.py b/python/sglang/srt/layers/quantization/base_config.py index 48511d09f..5cac5d420 100644 --- a/python/sglang/srt/layers/quantization/base_config.py +++ b/python/sglang/srt/layers/quantization/base_config.py @@ -10,6 +10,7 @@ from torch import nn if TYPE_CHECKING: from sglang.srt.layers.moe.moe_runner import MoeRunnerConfig + from sglang.srt.layers.moe.moe_runner.triton import TritonMoeQuantInfo from sglang.srt.layers.moe.token_dispatcher import CombineInput, DispatchOutput from sglang.srt.models.utils import WeightsMapper @@ -106,6 +107,19 @@ class FusedMoEMethodBase(QuantizeMethodBase): ) -> CombineInput: raise NotImplementedError + def get_triton_quant_info(self, layer: torch.nn.Module) -> "TritonMoeQuantInfo": + """Return a ``TritonMoeQuantInfo`` describing the quantisation state + stored on *layer*. + + The LoRA MoE runner calls this so that ``invoke_fused_moe_kernel`` + receives the correct flags / scales / block-shape for the base + weights. Each quantisation method must override this with the + same construction it already uses inside ``apply()``. + """ + raise NotImplementedError( + f"{type(self).__name__} must implement get_triton_quant_info()" + ) + class QuantizationConfig(ABC): """Base class for quantization configs.""" diff --git a/python/sglang/srt/layers/quantization/blockwise_int8.py b/python/sglang/srt/layers/quantization/blockwise_int8.py index 60d4e3929..fec99da3c 100644 --- a/python/sglang/srt/layers/quantization/blockwise_int8.py +++ b/python/sglang/srt/layers/quantization/blockwise_int8.py @@ -360,13 +360,8 @@ class BlockInt8MoEMethod(FusedMoEMethodBase): self.moe_runner_config = moe_runner_config self.runner = MoeRunner(MoeRunnerBackend.TRITON, moe_runner_config) - def apply( - self, - layer: torch.nn.Module, - dispatch_output: StandardDispatchOutput, - ) -> CombineInput: - - quant_info = TritonMoeQuantInfo( + def get_triton_quant_info(self, layer: torch.nn.Module) -> TritonMoeQuantInfo: + return TritonMoeQuantInfo( w13_weight=layer.w13_weight, w2_weight=layer.w2_weight, use_int8_w8a8=True, @@ -377,4 +372,12 @@ class BlockInt8MoEMethod(FusedMoEMethodBase): block_shape=self.quant_config.weight_block_size, ) + def apply( + self, + layer: torch.nn.Module, + dispatch_output: StandardDispatchOutput, + ) -> CombineInput: + + quant_info = self.get_triton_quant_info(layer) + return self.runner.run(dispatch_output, quant_info) diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index 45ac15649..50b330b52 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/python/sglang/srt/layers/quantization/fp8.py @@ -1487,6 +1487,26 @@ class Fp8MoEMethod(FusedMoEMethodBase): # TODO(cwan): refactor other backends pass + def get_triton_quant_info(self, layer: torch.nn.Module) -> TritonMoeQuantInfo: + return TritonMoeQuantInfo( + w13_weight=layer.w13_weight, + w2_weight=layer.w2_weight, + b13=getattr(layer, "w13_weight_bias", None), + b2=getattr(layer, "w2_weight_bias", None), + use_fp8_w8a8=True, + w13_scale=( + layer.w13_weight_scale_inv + if self.block_quant + else layer.w13_weight_scale + ), + w2_scale=( + layer.w2_weight_scale_inv if self.block_quant else layer.w2_weight_scale + ), + a13_scale=layer.w13_input_scale, + a2_scale=layer.w2_input_scale, + block_shape=self.quant_config.weight_block_size, + ) + def apply( self, layer: torch.nn.Module, @@ -1663,26 +1683,7 @@ class Fp8MoEMethod(FusedMoEMethodBase): ), ) elif self.runner.runner_backend.is_triton(): - quant_info = TritonMoeQuantInfo( - w13_weight=layer.w13_weight, - w2_weight=layer.w2_weight, - b13=getattr(layer, "w13_weight_bias", None), - b2=getattr(layer, "w2_weight_bias", None), - use_fp8_w8a8=True, - w13_scale=( - layer.w13_weight_scale_inv - if self.block_quant - else layer.w13_weight_scale - ), - w2_scale=( - layer.w2_weight_scale_inv - if self.block_quant - else layer.w2_weight_scale - ), - a13_scale=layer.w13_input_scale, - a2_scale=layer.w2_input_scale, - block_shape=self.quant_config.weight_block_size, - ) + quant_info = self.get_triton_quant_info(layer) else: raise NotImplementedError( "Unsupported runner backend: %s" % self.runner.runner_backend diff --git a/python/sglang/srt/layers/quantization/moe_wna16.py b/python/sglang/srt/layers/quantization/moe_wna16.py index ce6397dd1..2da526402 100644 --- a/python/sglang/srt/layers/quantization/moe_wna16.py +++ b/python/sglang/srt/layers/quantization/moe_wna16.py @@ -364,19 +364,10 @@ class MoeWNA16Method(FusedMoEMethodBase): self.moe_runner_config = moe_runner_config self.runner = MoeRunner(MoeRunnerBackend.TRITON, moe_runner_config) - def apply( - self, - layer: torch.nn.Module, - dispatch_output: StandardDispatchOutput, - ) -> CombineInput: - assert ( - self.moe_runner_config.activation == "silu" - ), "Only SiLU activation is supported." - + def get_triton_quant_info(self, layer: torch.nn.Module) -> TritonMoeQuantInfo: weight_bits = self.quant_config.weight_bits has_zp = self.quant_config.has_zp - - quant_info = TritonMoeQuantInfo( + return TritonMoeQuantInfo( w13_weight=layer.w13_qweight, w2_weight=layer.w2_qweight, use_int4_w4a16=weight_bits == 4, @@ -387,6 +378,17 @@ class MoeWNA16Method(FusedMoEMethodBase): w2_zp=layer.w2_qzeros if has_zp else None, block_shape=[0, layer.group_size], ) + + def apply( + self, + layer: torch.nn.Module, + dispatch_output: StandardDispatchOutput, + ) -> CombineInput: + assert ( + self.moe_runner_config.activation == "silu" + ), "Only SiLU activation is supported." + + quant_info = self.get_triton_quant_info(layer) return self.runner.run(dispatch_output, quant_info) @staticmethod diff --git a/python/sglang/srt/layers/quantization/unquant.py b/python/sglang/srt/layers/quantization/unquant.py index 5fce65159..7339640c4 100644 --- a/python/sglang/srt/layers/quantization/unquant.py +++ b/python/sglang/srt/layers/quantization/unquant.py @@ -554,6 +554,14 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, MultiPlatformOp): ) return StandardCombineInput(hidden_states=output) + def get_triton_quant_info(self, layer: torch.nn.Module) -> TritonMoeQuantInfo: + return TritonMoeQuantInfo( + w13_weight=layer.w13_weight, + w2_weight=layer.w2_weight, + b13=getattr(layer, "w13_weight_bias", None), + b2=getattr(layer, "w2_weight_bias", None), + ) + def forward_xpu( self, layer: torch.nn.Module, @@ -594,12 +602,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, MultiPlatformOp): ), f"activation = {moe_runner_config.activation} is not supported \ for Triton PATH, please set ENV SGLANG_USE_SGL_XPU=1." - quant_info = TritonMoeQuantInfo( - w13_weight=layer.w13_weight, - w2_weight=layer.w2_weight, - b13=getattr(layer, "w13_weight_bias", None), - b2=getattr(layer, "w2_weight_bias", None), - ) + quant_info = self.get_triton_quant_info(layer) return self.runner.run(dispatch_output, quant_info) def forward_npu( diff --git a/python/sglang/srt/layers/quantization/w8a8_fp8.py b/python/sglang/srt/layers/quantization/w8a8_fp8.py index 808e3e822..aeea826cd 100644 --- a/python/sglang/srt/layers/quantization/w8a8_fp8.py +++ b/python/sglang/srt/layers/quantization/w8a8_fp8.py @@ -286,13 +286,8 @@ class W8A8FP8MoEMethod(FusedMoEMethodBase): self.moe_runner_config = moe_runner_config self.runner = MoeRunner(MoeRunnerBackend.TRITON, moe_runner_config) - def apply( - self, - layer: torch.nn.Module, - dispatch_output: StandardDispatchOutput, - ) -> CombineInput: - - quant_info = TritonMoeQuantInfo( + def get_triton_quant_info(self, layer: torch.nn.Module) -> TritonMoeQuantInfo: + return TritonMoeQuantInfo( w13_weight=layer.w13_weight, w2_weight=layer.w2_weight, use_fp8_w8a8=True, @@ -302,4 +297,12 @@ class W8A8FP8MoEMethod(FusedMoEMethodBase): a13_scale=layer.w13_input_scale, a2_scale=layer.w2_input_scale, ) + + def apply( + self, + layer: torch.nn.Module, + dispatch_output: StandardDispatchOutput, + ) -> CombineInput: + + quant_info = self.get_triton_quant_info(layer) return self.runner.run(dispatch_output, quant_info) diff --git a/python/sglang/srt/layers/quantization/w8a8_int8.py b/python/sglang/srt/layers/quantization/w8a8_int8.py index fb12bc4fe..d7bfd888b 100644 --- a/python/sglang/srt/layers/quantization/w8a8_int8.py +++ b/python/sglang/srt/layers/quantization/w8a8_int8.py @@ -331,6 +331,18 @@ class W8A8Int8MoEMethod(FusedMoEMethodBase): self.moe_runner_config = moe_runner_config self.runner = MoeRunner(MoeRunnerBackend.TRITON, moe_runner_config) + def get_triton_quant_info(self, layer: torch.nn.Module) -> TritonMoeQuantInfo: + return TritonMoeQuantInfo( + w13_weight=layer.w13_weight, + w2_weight=layer.w2_weight, + use_int8_w8a8=True, + per_channel_quant=True, + w13_scale=layer.w13_weight_scale, + w2_scale=layer.w2_weight_scale, + a13_scale=layer.w13_input_scale, + a2_scale=layer.w2_input_scale, + ) + def apply( self, layer: torch.nn.Module, @@ -365,14 +377,5 @@ class W8A8Int8MoEMethod(FusedMoEMethodBase): ) return StandardCombineInput(hidden_states=output) - quant_info = TritonMoeQuantInfo( - w13_weight=layer.w13_weight, - w2_weight=layer.w2_weight, - use_int8_w8a8=True, - per_channel_quant=True, - w13_scale=layer.w13_weight_scale, - w2_scale=layer.w2_weight_scale, - a13_scale=layer.w13_input_scale, - a2_scale=layer.w2_input_scale, - ) + quant_info = self.get_triton_quant_info(layer) return self.runner.run(dispatch_output, quant_info) diff --git a/python/sglang/srt/lora/backend/chunked_backend.py b/python/sglang/srt/lora/backend/chunked_backend.py index 2e0415b87..75163cb9a 100644 --- a/python/sglang/srt/lora/backend/chunked_backend.py +++ b/python/sglang/srt/lora/backend/chunked_backend.py @@ -66,6 +66,7 @@ class ChunkedSgmvLoRABackend(BaseLoRABackend): x: torch.Tensor, weights: torch.Tensor, pruned_batch_info: LoRABatchInfo = None, + stack_num: int = 1, *args, **kwargs, ) -> torch.Tensor: @@ -76,7 +77,7 @@ class ChunkedSgmvLoRABackend(BaseLoRABackend): x=x, weights=weights, batch_info=batch_info, - num_slices=1, + num_slices=stack_num, ) def run_lora_b_sgemm( diff --git a/python/sglang/srt/lora/backend/torch_backend.py b/python/sglang/srt/lora/backend/torch_backend.py index 64337acce..11bb1b31f 100644 --- a/python/sglang/srt/lora/backend/torch_backend.py +++ b/python/sglang/srt/lora/backend/torch_backend.py @@ -39,7 +39,12 @@ class TorchNativeLoRABackend(BaseLoRABackend): super().__init__(max_loras_per_batch, device) def run_lora_a_sgemm( - self, x: torch.Tensor, weights: torch.Tensor, *args, **kwargs + self, + x: torch.Tensor, + weights: torch.Tensor, + stack_num: int = 1, + *args, + **kwargs, ) -> torch.Tensor: output_tensor = sgemm_lora_a_fwd( inputs=x, @@ -48,7 +53,7 @@ class TorchNativeLoRABackend(BaseLoRABackend): seg_len_tensor=self.batch_info.seg_lens_cpu, lora_ranks=self.batch_info.lora_ranks_cpu, scaling_tensor=self.batch_info.scalings_cpu, - num_slices=1, + num_slices=stack_num, ) return output_tensor diff --git a/python/sglang/srt/lora/backend/triton_backend.py b/python/sglang/srt/lora/backend/triton_backend.py index 3ea54bb49..de15d4e16 100644 --- a/python/sglang/srt/lora/backend/triton_backend.py +++ b/python/sglang/srt/lora/backend/triton_backend.py @@ -53,13 +53,14 @@ class TritonLoRABackend(BaseLoRABackend): x: torch.Tensor, weights: torch.Tensor, pruned_batch_info: LoRABatchInfo = None, + stack_num: int = 1, *args, **kwargs, ) -> torch.Tensor: batch_info = ( pruned_batch_info if pruned_batch_info is not None else self.batch_info ) - return sgemm_lora_a_fwd(x, weights, batch_info) + return sgemm_lora_a_fwd(x, weights, batch_info, stack_num=stack_num) def run_lora_b_sgemm( self, diff --git a/python/sglang/srt/lora/layers.py b/python/sglang/srt/lora/layers.py index 412ff5824..4a327f1e4 100644 --- a/python/sglang/srt/lora/layers.py +++ b/python/sglang/srt/lora/layers.py @@ -14,6 +14,7 @@ from sglang.srt.layers.linear import ( ColumnParallelLinear, MergedColumnParallelLinear, QKVParallelLinear, + ReplicatedLinear, RowParallelLinear, ) from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE @@ -687,6 +688,94 @@ class RowParallelLinearWithLoRA(BaseLayerWithLoRA): return B +class ReplicatedLinearWithLoRA(BaseLayerWithLoRA): + """LoRA wrapper for ReplicatedLinear (no TP sharding). + + Used for DeepSeek MLA's fused_qkv_a_proj_with_mqa, which fuses + q_a_proj and kv_a_proj_with_mqa into a single replicated linear. + The two sub-projections have unequal output dimensions, so LoRA B + is applied via two separate sgemm calls, one per partition. + """ + + first_output_dim: int = 0 + + def __init__( + self, + base_layer: ReplicatedLinear, + lora_backend: BaseLoRABackend, + ) -> None: + super().__init__(base_layer, lora_backend) + self.output_size = base_layer.output_size + + def set_lora_info(self, A_buffer: torch.Tensor, B_buffer: torch.Tensor): + self.set_lora = True + self.A_buffer = A_buffer + self.B_buffer = B_buffer + first = self.first_output_dim + if first > 0 and first < B_buffer.shape[-2]: + self.B_first = B_buffer[:, :first, :].contiguous() + self.B_second = B_buffer[:, first:, :].contiguous() + output_size = B_buffer.shape[-2] + self.first_offset = torch.tensor( + [0, first], dtype=torch.int32, device=B_buffer.device + ) + self.second_offset = torch.tensor( + [0, output_size - first], dtype=torch.int32, device=B_buffer.device + ) + else: + self.B_first = None + self.B_second = None + self.output_offset = torch.tensor( + [0, self.output_size], + dtype=torch.int32, + device=B_buffer.device, + ) + + def apply_lora(self, base_output: torch.Tensor, x: torch.Tensor) -> torch.Tensor: + if self.B_first is not None: + rank = self.B_buffer.shape[-1] + lora_a_output = self.lora_backend.run_lora_a_sgemm( + x, self.A_buffer, stack_num=2 + ) + first_out = base_output[:, : self.first_output_dim] + second_out = base_output[:, self.first_output_dim :] + self.lora_backend.run_lora_b_sgemm( + x=lora_a_output[:, :rank].contiguous(), + weights=self.B_first, + output_offset=self.first_offset, + base_output=first_out, + ) + self.lora_backend.run_lora_b_sgemm( + x=lora_a_output[:, rank:].contiguous(), + weights=self.B_second, + output_offset=self.second_offset, + base_output=second_out, + ) + return base_output + else: + lora_a_output = self.lora_backend.run_lora_a_sgemm(x, self.A_buffer) + return self.lora_backend.run_lora_b_sgemm( + x=lora_a_output, + weights=self.B_buffer, + output_offset=self.output_offset, + base_output=base_output, + ) + + def forward(self, x: torch.Tensor): + bias = self.base_layer.bias if not self.base_layer.skip_bias_add else None + output = self.base_layer.quant_method.apply(self.base_layer, x, bias) + if self.set_lora: + output = self.apply_lora(output, x) + output_bias = self.base_layer.bias if self.base_layer.skip_bias_add else None + return output, output_bias + + def slice_lora_a_weights(self, A: torch.Tensor, tp_rank: int): + return A + + def slice_lora_b_weights(self, B: torch.Tensor, tp_rank: int): + return B + + class FusedMoEWithLoRA(BaseLayerWithLoRA): """ Wrapper around FusedMoE that integrates LoRA into the MoE computation. @@ -721,7 +810,6 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA): # initialize triton_lora moe runner for batches with lora enabled from sglang.srt.layers.moe.moe_runner.runner import MoeRunner - from sglang.srt.layers.moe.moe_runner.triton import TritonMoeQuantInfo self._lora_runner = MoeRunner( base_layer.quant_method.runner.runner_backend, @@ -730,12 +818,7 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA): ) # Pre-compute quant info for efficiency (weights don't change during inference) - self._quant_info = TritonMoeQuantInfo( - w13_weight=base_layer.w13_weight, - w2_weight=base_layer.w2_weight, - b13=getattr(base_layer, "w13_weight_bias", None), - b2=getattr(base_layer, "w2_weight_bias", None), - ) + self._quant_info = base_layer.quant_method.get_triton_quant_info(base_layer) def set_lora_info( self, @@ -941,6 +1024,7 @@ def get_lora_layer( FusedMoE: FusedMoEWithLoRA, ParallelLMHead: ParallelLMHeadWithLoRA, VocabParallelEmbedding: VocabParallelEmbeddingWithLoRA, + ReplicatedLinear: ReplicatedLinearWithLoRA, QKVParallelLinear: QKVParallelLinearWithLoRA, MergedColumnParallelLinear: MergedColumnParallelLinearWithLoRA, ColumnParallelLinear: ColumnParallelLinearWithLoRA, diff --git a/python/sglang/srt/lora/lora.py b/python/sglang/srt/lora/lora.py index fbae373cb..082d1fb0a 100644 --- a/python/sglang/srt/lora/lora.py +++ b/python/sglang/srt/lora/lora.py @@ -133,13 +133,14 @@ class LoRAAdapter(nn.Module): ) def _normalize_weights(self): - # normalize kv_proj and gate_up_proj for layer in self.layers: weight_names = list(layer.weights.keys()) self.normalize_qkv_proj(weight_names, layer.weights) self._rename_expert_w_to_proj(layer.weights) weight_names = list(layer.weights.keys()) self.normalize_gate_up_proj(weight_names, layer.weights) + weight_names = list(layer.weights.keys()) + self.normalize_fused_qkv_a_proj(weight_names, layer.weights) def normalize_qkv_proj( self, weight_names: List[str], weights: Dict[str, torch.Tensor] @@ -242,6 +243,33 @@ class LoRAAdapter(nn.Module): weights[gate_up_name] = weights[gate_up_name].repeat(*repeat_dims) # else: no-op as LoRA B weight is already stacked. + def normalize_fused_qkv_a_proj( + self, weight_names: List[str], weights: Dict[str, torch.Tensor] + ): + """Fuse separate q_a_proj and kv_a_proj_with_mqa LoRA weights into + a single fused_qkv_a_proj_with_mqa entry (concat along dim 0 for + both A and B), matching the DeepSeek MLA fused projection layout.""" + for weight_name in weight_names: + if "q_a_proj" not in weight_name: + continue + if "fused_qkv_a_proj_with_mqa" in weight_name: + continue + + q_a_name = weight_name + kv_a_name = weight_name.replace("q_a_proj", "kv_a_proj_with_mqa") + fused_name = weight_name.replace("q_a_proj", "fused_qkv_a_proj_with_mqa") + + kv_a_weight = ( + weights[kv_a_name] + if kv_a_name in weights + else torch.zeros_like(weights[q_a_name]) + ) + + weights[fused_name] = torch.cat((weights[q_a_name], kv_a_weight), dim=0) + weights.pop(q_a_name) + if kv_a_name in weights: + weights.pop(kv_a_name) + def pin_weights_in_cpu(self): for layer in self.layers: for name, weight in layer.weights.items(): diff --git a/python/sglang/srt/lora/lora_manager.py b/python/sglang/srt/lora/lora_manager.py index c8f338cdc..b14de3f4b 100644 --- a/python/sglang/srt/lora/lora_manager.py +++ b/python/sglang/srt/lora/lora_manager.py @@ -763,6 +763,25 @@ class LoRAManager: self.lm_head_module = lora_module continue + # Handle DeepSeek MLA fused projection: set the boundary + # between q_a and kv_a output partitions so the LoRA layer + # can apply separate B projections for each. + if ( + "fused_qkv_a_proj_with_mqa" in self.target_modules + and module_name.endswith("fused_qkv_a_proj_with_mqa") + ): + from sglang.srt.lora.layers import ReplicatedLinearWithLoRA + + layer_id = get_layer_id(module_name) + if layer_id is None: + continue + lora_module = self.set_lora_module(module_name, module) + if isinstance(lora_module, ReplicatedLinearWithLoRA): + q_lora_rank = getattr(self.base_hf_config, "q_lora_rank", None) or 0 + lora_module.first_output_dim = q_lora_rank + self.lora_modules[layer_id][module_name] = lora_module + continue + # The module should be converted if it is included in target_names if module_name.split(".")[-1] in self.target_modules: layer_id = get_layer_id(module_name) diff --git a/python/sglang/srt/lora/mem_pool.py b/python/sglang/srt/lora/mem_pool.py index 22ef4d6ad..6c923afc5 100644 --- a/python/sglang/srt/lora/mem_pool.py +++ b/python/sglang/srt/lora/mem_pool.py @@ -12,6 +12,7 @@ from sglang.srt.lora.lora_config import LoRAConfig from sglang.srt.lora.lora_registry import LoRARef from sglang.srt.lora.utils import ( EMBEDDING_NAMES, + REPLICATED_LINEAR_LORA_NAMES, ROW_PARALLELISM_LINEAR_LORA_NAMES, LoRAType, get_hidden_dim, @@ -188,7 +189,11 @@ class LoRAMemoryPool: module_name, self.base_hf_config, base_model, layer_idx ) c = get_stacked_multiply(module_name) - if self.tp_size > 1 and module_name in ROW_PARALLELISM_LINEAR_LORA_NAMES: + if ( + self.tp_size > 1 + and module_name in ROW_PARALLELISM_LINEAR_LORA_NAMES + and module_name not in REPLICATED_LINEAR_LORA_NAMES + ): input_dim = divide(input_dim, self.tp_size) if self.is_moe_module(module_name): @@ -240,7 +245,11 @@ class LoRAMemoryPool: _, output_dim = get_hidden_dim( module_name, self.base_hf_config, base_model, layer_idx ) - if self.tp_size > 1 and module_name not in ROW_PARALLELISM_LINEAR_LORA_NAMES: + if ( + self.tp_size > 1 + and module_name not in ROW_PARALLELISM_LINEAR_LORA_NAMES + and module_name not in REPLICATED_LINEAR_LORA_NAMES + ): output_dim = divide(output_dim, self.tp_size) # Check if MoE module and return appropriate shape @@ -619,7 +628,12 @@ class LoRAMemoryPool: if name in ["gate_up_proj_moe", "down_proj_moe"]: if self.experts_shared_outer_loras and name == "gate_up_proj_moe": - if isinstance(weights, torch.Tensor) and weights.dim() == 3: + if weights is None: + buffer_view = target_buffer[ + buffer_id, 0, : lora_rank * c, : + ] + load_lora_weight_tensor(buffer_view, None) + elif isinstance(weights, torch.Tensor) and weights.dim() == 3: if weights.shape[0] != 1: raise ValueError( f"experts_shared_outer_loras is enabled but " @@ -669,7 +683,10 @@ class LoRAMemoryPool: if name in ["gate_up_proj_moe", "down_proj_moe"]: if self.experts_shared_outer_loras and name == "down_proj_moe": - if isinstance(weights, torch.Tensor) and weights.dim() == 3: + if weights is None: + buffer_view = target_buffer[buffer_id, 0, :, :lora_rank] + load_lora_weight_tensor(buffer_view, None) + elif isinstance(weights, torch.Tensor) and weights.dim() == 3: if weights.shape[0] != 1: raise ValueError( f"experts_shared_outer_loras is enabled but " diff --git a/python/sglang/srt/lora/utils.py b/python/sglang/srt/lora/utils.py index 4ed9691fa..2ef127053 100644 --- a/python/sglang/srt/lora/utils.py +++ b/python/sglang/srt/lora/utils.py @@ -83,14 +83,47 @@ def get_hidden_dim( config.num_attention_heads + config.num_key_value_heads * 2 ) elif module_name == "o_proj": + o_head_dim = getattr(config, "v_head_dim", None) or head_dim return ( - head_dim * config.num_attention_heads, + o_head_dim * config.num_attention_heads, config.hidden_size, ) elif module_name == "gate_up_proj": - return config.hidden_size, config.intermediate_size * 2 + inter = config.intermediate_size + first_k = getattr(config, "first_k_dense_replace", None) + moe_freq = getattr(config, "moe_layer_freq", 1) + if ( + first_k is not None + and layer_idx >= first_k + and layer_idx % moe_freq == 0 + ): + moe_inter = getattr(config, "moe_intermediate_size", None) + n_shared = getattr(config, "n_shared_experts", None) + if moe_inter is not None and n_shared is not None: + inter = moe_inter * n_shared + return config.hidden_size, inter * 2 elif module_name == "down_proj": - return config.intermediate_size, config.hidden_size + inter = config.intermediate_size + first_k = getattr(config, "first_k_dense_replace", None) + moe_freq = getattr(config, "moe_layer_freq", 1) + if ( + first_k is not None + and layer_idx >= first_k + and layer_idx % moe_freq == 0 + ): + moe_inter = getattr(config, "moe_intermediate_size", None) + n_shared = getattr(config, "n_shared_experts", None) + if moe_inter is not None and n_shared is not None: + inter = moe_inter * n_shared + return inter, config.hidden_size + elif module_name == "fused_qkv_a_proj_with_mqa": + q_lora_rank = getattr(config, "q_lora_rank", None) or 0 + kv_lora_rank = config.kv_lora_rank + qk_rope_head_dim = config.qk_rope_head_dim + return ( + config.hidden_size, + q_lora_rank + kv_lora_rank + qk_rope_head_dim, + ) elif module_name == "gate_up_proj_moe": moe_inter = ( getattr(config, "moe_intermediate_size", None) @@ -151,6 +184,8 @@ def get_normalized_target_modules( "lm_head": "lm_head", "output": "lm_head", "unembed_tokens": "lm_head", + "q_a_proj": "fused_qkv_a_proj_with_mqa", + "kv_a_proj_with_mqa": "fused_qkv_a_proj_with_mqa", } result = set() @@ -169,6 +204,7 @@ def get_stacked_multiply(module_name: str) -> int: "qkv_proj": 3, "gate_up_proj": 2, "gate_up_proj_moe": 2, + "fused_qkv_a_proj_with_mqa": 2, } return stacked_rank[module_name] if module_name in stacked_rank else 1 @@ -190,6 +226,7 @@ def get_target_module_name(full_module_name: str, target_modules: Set[str]) -> s EMBEDDING_NAMES = ["embed_tokens", "lm_head"] ROW_PARALLELISM_LINEAR_LORA_NAMES = ["o_proj", "down_proj", "down_proj_moe"] +REPLICATED_LINEAR_LORA_NAMES = ["fused_qkv_a_proj_with_mqa"] # Normalized module names that the LoRA system fully supports # (i.e. get_hidden_dim, init_buffers, and init_lora_modules can handle them). @@ -201,6 +238,7 @@ _KNOWN_LORA_TARGET_MODULES = frozenset( "down_proj", "embed_tokens", "lm_head", + "fused_qkv_a_proj_with_mqa", } ) diff --git a/test/registered/lora/test_lora_deepseek_v3_base_logprob_diff.py b/test/registered/lora/test_lora_deepseek_v3_base_logprob_diff.py new file mode 100644 index 000000000..8ae8dbd15 --- /dev/null +++ b/test/registered/lora/test_lora_deepseek_v3_base_logprob_diff.py @@ -0,0 +1,156 @@ +# Copyright 2023-2025 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +""" +Regression test for DeepSeek-V3.1-Base MLA LoRA logprob accuracy. + +Compares SGLang LoRA logprobs against reference training logprobs from a +pre-computed dataset. The LoRA adapter and reference data are downloaded from: +https://huggingface.co/datasets/yushengsu/lora-diff-DeepSeek-V3.1-Base + +Usage: + python -m unittest test_lora_deepseek_v3_base_logprob_diff +""" + +import multiprocessing as mp +import os +import unittest + +import torch +from huggingface_hub import snapshot_download + +import sglang as sgl +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci( + est_time=300, + suite="nightly-8-gpu-b200", +) + +BASE_MODEL = "deepseek-ai/DeepSeek-V3.1-Base" +LORA_HF_REPO = "yushengsu/lora-diff-DeepSeek-V3.1-Base" +LORA_BACKEND = "triton" +MAX_LORA_RANK = 32 +TP_SIZE = 8 +MOE_RUNNER_BACKEND = "triton" +EXPERTS_SHARED_OUTER_LORAS = True +PREFILL_ATTENTION_BACKEND = "fa4" +DECODE_ATTENTION_BACKEND = "flashinfer" + +KL_THRESHOLD = 6e-3 + + +def kl_v2(a, b): + a = torch.tensor(a) if not torch.is_tensor(a) else a + b = torch.tensor(b) if not torch.is_tensor(b) else b + return (((a - b) ** 2) * 0.5).mean().item() + + +def get_prompt_logprobs(engine, input_ids, lora_path): + if isinstance(input_ids, torch.Tensor): + input_ids = [input_ids.tolist()] + elif not isinstance(input_ids[0], list): + input_ids = [input_ids] + out = engine.generate( + input_ids=input_ids, + sampling_params={"max_new_tokens": 0, "temperature": 0.0}, + return_logprob=True, + logprob_start_len=0, + lora_path=lora_path, + ) + if isinstance(out, list): + out = out[0] + return [logprob for logprob, _, _ in out["meta_info"]["input_token_logprobs"]][1:] + + +class TestLoRADeepSeekV3BaseLogprobDiff(CustomTestCase): + + def test_lora_deepseek_v3_base_logprob_accuracy(self): + adapter_path = snapshot_download( + LORA_HF_REPO, + repo_type="dataset", + ) + + engine = sgl.Engine( + model_path=BASE_MODEL, + tp_size=TP_SIZE, + enable_lora=True, + max_lora_rank=MAX_LORA_RANK, + lora_paths={"my_lora": adapter_path}, + lora_backend=LORA_BACKEND, + attention_backend="flashinfer", + moe_runner_backend=MOE_RUNNER_BACKEND, + experts_shared_outer_loras=EXPERTS_SHARED_OUTER_LORAS, + prefill_attention_backend=PREFILL_ATTENTION_BACKEND, + decode_attention_backend=DECODE_ATTENTION_BACKEND, + disable_shared_experts_fusion=True, + ) + + try: + cdata = torch.load( + os.path.join(adapter_path, "compare_sample_train_data.pt"), + weights_only=False, + ) + + base_logprobs = get_prompt_logprobs(engine, cdata["tokens"], lora_path=None) + logprobs = get_prompt_logprobs(engine, cdata["tokens"], lora_path="my_lora") + + base_t = torch.tensor(base_logprobs) + lora_t = torch.tensor(logprobs) + diff = (base_t - lora_t).abs() + print( + f"[VERIFY] base vs lora: mean_diff={diff.mean().item():.6f}, " + f"max_diff={diff.max().item():.6f}, " + f"identical={torch.equal(base_t, lora_t)}" + ) + + self.assertFalse( + torch.equal(base_t, lora_t), + "LoRA logprobs should differ from base model logprobs", + ) + + kl_sglang_trainer = kl_v2(cdata["training_logprobs"], logprobs) + kl_orig_trainer = kl_v2( + cdata["training_logprobs"], cdata["sampling_logprobs"] + ) + kl_sglang_orig = kl_v2(logprobs, cdata["sampling_logprobs"]) + + print(f"KL(orig_sampler, trainer) = {kl_orig_trainer:.6e}") + print(f"KL(sglang, trainer) = {kl_sglang_trainer:.6e}") + print(f"KL(sglang, orig_sampler) = {kl_sglang_orig:.6e}") + + self.assertLessEqual( + kl_sglang_trainer, + KL_THRESHOLD, + f"KL(sglang, trainer) = {kl_sglang_trainer:.6e} exceeds " + f"threshold {KL_THRESHOLD}", + ) + + finally: + engine.shutdown() + + +if __name__ == "__main__": + try: + mp.set_start_method("spawn") + except RuntimeError: + pass + + try: + unittest.main(warnings="ignore", verbosity=2) + finally: + if torch.cuda.is_available(): + torch.cuda.empty_cache() + torch.cuda.synchronize()