diff --git a/python/sglang/srt/lora/lora.py b/python/sglang/srt/lora/lora.py index 245c15e3bcb2..704fba9d3a3a 100644 --- a/python/sglang/srt/lora/lora.py +++ b/python/sglang/srt/lora/lora.py @@ -208,6 +208,8 @@ def _normalize_weights(self): 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) + # Stack in_proj_b + in_proj_a → in_proj_ba for GDN layers + self._normalize_in_proj_ba(layer.weights) weight_names = list(layer.weights.keys()) self.normalize_gate_up_proj(weight_names, layer.weights) weight_names = list(layer.weights.keys()) @@ -347,6 +349,25 @@ def _normalize_in_proj_qkvz(self, weights: Dict[str, torch.Tensor]): weights.pop(k_name) weights.pop(v_name) weights.pop(z_name) + elif "in_proj_qkv." in weight_name: + # 2-way split (Megatron-Bridge adapter export): in_proj_qkv + # covers the q|k|v rows and in_proj_z the z rows, with one + # shared lora_A. The stacked buffer expects one A block per + # slice, so repeat the qkv A block 3x (q, k, v) before + # appending the z block; B rows concatenate directly. + z_name = weight_name.replace("in_proj_qkv.", "in_proj_z.") + if z_name not in weights: + continue + qkvz_name = weight_name.replace("in_proj_qkv.", "in_proj_qkvz.") + cat_dim = weights[weight_name].dim() - 2 + qkv_w = weights[weight_name] + if "lora_A" in weight_name: + repeat_dims = [1] * qkv_w.dim() + repeat_dims[cat_dim] = 3 + qkv_w = qkv_w.repeat(*repeat_dims) + weights[qkvz_name] = torch.cat((qkv_w, weights[z_name]), cat_dim) + weights.pop(weight_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). @@ -356,6 +377,45 @@ def _normalize_in_proj_qkvz(self, weights: Dict[str, torch.Tensor]): weights[weight_name] = weights[weight_name].repeat(*repeat_dims) # else (in_proj_qkvz lora_B, or unrelated): no-op. + def _normalize_in_proj_ba(self, weights: Dict[str, torch.Tensor]): + """Normalize in_proj_ba weights for GDN (GatedDeltaNet) layers like + Qwen3.5. + + Two adapter formats are handled: + + 1. Split: ``in_proj_b + in_proj_a`` (HF checkpoint naming, also the + Megatron-Bridge adapter export) are present as separate weights → + concatenate them into ``in_proj_ba`` (B rows b|a; A blocks b, a). + + 2. Already-merged: the adapter has a single ``in_proj_ba`` weight + (PEFT trained against SGLang's fused Linear). The stacked buffer + expects two per-slice ``A`` blocks, so repeat ``lora_A`` 2x along + the rank dim. ``lora_B`` is already full-output-dim and matches + the buffer directly. + """ + for weight_name in list(weights.keys()): + # NB: match with the trailing dot so the merged "in_proj_ba." + # names don't take the split branch. + if "in_proj_b." in weight_name: + a_name = weight_name.replace("in_proj_b.", "in_proj_a.") + if a_name not in weights: + continue + ba_name = weight_name.replace("in_proj_b.", "in_proj_ba.") + cat_dim = weights[weight_name].dim() - 2 + weights[ba_name] = torch.cat( + (weights[weight_name], weights[a_name]), cat_dim + ) + weights.pop(weight_name) + weights.pop(a_name) + elif "in_proj_ba" in weight_name and "lora_A" in weight_name: + # Already-merged adapter: replicate the shared A across the 2 + # stacked slots the buffer expects (b, a). + ndim = weights[weight_name].dim() + repeat_dims = [1] * ndim + repeat_dims[ndim - 2] = 2 + weights[weight_name] = weights[weight_name].repeat(*repeat_dims) + # else (in_proj_ba lora_B, or unrelated): no-op. + def normalize_gate_up_proj( self, weight_names: List[str], weights: Dict[str, torch.Tensor] ): diff --git a/python/sglang/srt/lora/utils.py b/python/sglang/srt/lora/utils.py index e8301d5b416c..befa86b96288 100644 --- a/python/sglang/srt/lora/utils.py +++ b/python/sglang/srt/lora/utils.py @@ -253,6 +253,8 @@ def get_normalized_target_modules( "v_proj": "qkv_proj", "gate_proj": "gate_up_proj", "up_proj": "gate_up_proj", + "in_proj_b": "in_proj_ba", + "in_proj_a": "in_proj_ba", "out_proj": "out_proj", "embed_tokens": "embed_tokens", "vocab_emb": "embed_tokens", @@ -291,6 +293,7 @@ def get_stacked_multiply( stacked_rank = { "qkv_proj": 3, "in_proj_qkvz": 4, # GDN packed input projection + "in_proj_ba": 2, # GDN packed b/a input projection "gate_up_proj": 2, "gate_up_proj_moe": 2, "in_proj": 2, @@ -342,6 +345,7 @@ def get_target_module_name(full_module_name: str, target_modules: Set[str]) -> s "out_proj", "in_proj", "in_proj_qkvz", + "in_proj_ba", "up_proj", "gate_up_proj", "down_proj", diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py index 7157eb23e969..0017590f7625 100644 --- a/python/sglang/srt/models/qwen3_5.py +++ b/python/sglang/srt/models/qwen3_5.py @@ -1071,6 +1071,7 @@ class Qwen3_5ForCausalLM(nn.Module): "o_proj", "out_proj", "in_proj_qkvz", + "in_proj_ba", "gate_up_proj", "down_proj", "lm_head", @@ -1096,6 +1097,9 @@ def get_hidden_dim(self, module_name: str, layer_idx: int): 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 == "in_proj_ba": + # b + a projections: one scalar per linear-attention value head each + return config.hidden_size, config.linear_num_value_heads * 2 elif module_name == "gate_up_proj": # MoE: shared expert uses shared_expert_intermediate_size # Dense: regular MLP uses intermediate_size