diff --git a/python/sglang/srt/lora/backend/ascend_backend.py b/python/sglang/srt/lora/backend/ascend_backend.py index 2cffea189..5ca406ec3 100644 --- a/python/sglang/srt/lora/backend/ascend_backend.py +++ b/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 diff --git a/python/sglang/srt/lora/backend/chunked_backend.py b/python/sglang/srt/lora/backend/chunked_backend.py index 75163cb9a..c35cf6587 100644 --- a/python/sglang/srt/lora/backend/chunked_backend.py +++ b/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, diff --git a/python/sglang/srt/lora/backend/torch_backend.py b/python/sglang/srt/lora/backend/torch_backend.py index 11bb1b31f..13f0b9ba0 100644 --- a/python/sglang/srt/lora/backend/torch_backend.py +++ b/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( diff --git a/python/sglang/srt/lora/backend/triton_backend.py b/python/sglang/srt/lora/backend/triton_backend.py index 4758e640f..c3b8838e9 100644 --- a/python/sglang/srt/lora/backend/triton_backend.py +++ b/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 diff --git a/python/sglang/srt/lora/layers.py b/python/sglang/srt/lora/layers.py index 502aaf23b..b23ba29d2 100644 --- a/python/sglang/srt/lora/layers.py +++ b/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 diff --git a/python/sglang/srt/lora/lora.py b/python/sglang/srt/lora/lora.py index 082d1fb0a..a6570e565 100644 --- a/python/sglang/srt/lora/lora.py +++ b/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....") and shared-outer +# keys ("...experts...."), 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....`` (one tensor per expert) + - shared-outer: ``...experts....`` (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] diff --git a/python/sglang/srt/lora/lora_manager.py b/python/sglang/srt/lora/lora_manager.py index 58d00551a..7b9918cdc 100644 --- a/python/sglang/srt/lora/lora_manager.py +++ b/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 diff --git a/python/sglang/srt/lora/lora_moe_runners.py b/python/sglang/srt/lora/lora_moe_runners.py index 04f8b4867..3a08cb782 100644 --- a/python/sglang/srt/lora/lora_moe_runners.py +++ b/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( diff --git a/python/sglang/srt/lora/mem_pool.py b/python/sglang/srt/lora/mem_pool.py index 9b958d229..22a9faed2 100644 --- a/python/sglang/srt/lora/mem_pool.py +++ b/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] diff --git a/python/sglang/srt/lora/triton_ops/chunked_sgmv_expand.py b/python/sglang/srt/lora/triton_ops/chunked_sgmv_expand.py index 52e3d6038..50830be5a 100644 --- a/python/sglang/srt/lora/triton_ops/chunked_sgmv_expand.py +++ b/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, diff --git a/python/sglang/srt/lora/triton_ops/qkv_lora_b.py b/python/sglang/srt/lora/triton_ops/qkv_lora_b.py index 08ecc40ed..d43f0c64a 100644 --- a/python/sglang/srt/lora/triton_ops/qkv_lora_b.py +++ b/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, ) diff --git a/python/sglang/srt/lora/utils.py b/python/sglang/srt/lora/utils.py index 2ef127053..43b18f917 100644 --- a/python/sglang/srt/lora/utils.py +++ b/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: diff --git a/python/sglang/srt/models/nemotron_h.py b/python/sglang/srt/models/nemotron_h.py index 0426729ce..b0b195750 100644 --- a/python/sglang/srt/models/nemotron_h.py +++ b/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, diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py index e30e02244..8d0457242 100644 --- a/python/sglang/srt/models/qwen3_5.py +++ b/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") diff --git a/test/registered/lora/test_chunked_sgmv_backend.py b/test/registered/lora/test_chunked_sgmv_backend.py index 4073317cd..fe6c8ed96 100644 --- a/test/registered/lora/test_chunked_sgmv_backend.py +++ b/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): diff --git a/test/registered/lora/test_lora_nemotron_3_super_120b_a12b_logprob_diff.py b/test/registered/lora/test_lora_nemotron_3_super_120b_a12b_logprob_diff.py new file mode 100644 index 000000000..1da641680 --- /dev/null +++ b/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() diff --git a/test/registered/lora/test_lora_qwen3_5_35b_a3b_logprob_diff.py b/test/registered/lora/test_lora_qwen3_5_35b_a3b_logprob_diff.py new file mode 100644 index 000000000..f13107f2c --- /dev/null +++ b/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() diff --git a/test/registered/lora/test_lora_qwen3_5_4b_logprob_diff.py b/test/registered/lora/test_lora_qwen3_5_4b_logprob_diff.py new file mode 100644 index 000000000..33115a877 --- /dev/null +++ b/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()