From 343aeeef39b04e5a884231da279fba2d29696149 Mon Sep 17 00:00:00 2001 From: Joe Rowell Date: Thu, 18 Jun 2026 02:09:04 +0200 Subject: [PATCH] [Model] Laguna: support per-element output gating (#28400) Co-authored-by: Jimmy Shong <69131491+Jiminator@users.noreply.github.com> Co-authored-by: Jimmy Shong --- python/sglang/srt/configs/laguna.py | 2 + python/sglang/srt/models/laguna.py | 59 ++++++++++++++++++++--------- 2 files changed, 44 insertions(+), 17 deletions(-) diff --git a/python/sglang/srt/configs/laguna.py b/python/sglang/srt/configs/laguna.py index ff11f0269..5e172d98f 100644 --- a/python/sglang/srt/configs/laguna.py +++ b/python/sglang/srt/configs/laguna.py @@ -78,6 +78,7 @@ class LagunaConfig(PretrainedConfig): tie_word_embeddings: bool = False, attention_bias: bool = False, attention_dropout: float = 0.0, + gating: bool | str = True, sliding_window: int = 512, layer_types: Optional[List[str]] = None, mlp_layer_types: Optional[List[str]] = None, @@ -120,6 +121,7 @@ class LagunaConfig(PretrainedConfig): self.use_cache = use_cache self.attention_bias = attention_bias self.attention_dropout = attention_dropout + self.gating = "per-head" if gating is True else gating self.sliding_window = sliding_window self.num_experts = num_experts diff --git a/python/sglang/srt/models/laguna.py b/python/sglang/srt/models/laguna.py index 93cd9b3a3..cc8c65894 100644 --- a/python/sglang/srt/models/laguna.py +++ b/python/sglang/srt/models/laguna.py @@ -241,6 +241,7 @@ class LagunaAttention(nn.Module): attention_bias: bool, sliding_window_size: int, layer_type: str, + gating: bool | str = True, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ) -> None: @@ -248,6 +249,13 @@ class LagunaAttention(nn.Module): self.hidden_size = hidden_size self.head_dim = head_dim self.layer_id = layer_id + if gating not in (True, False, None, "per-head", "per-element"): + raise ValueError( + f"Unsupported gating value {gating!r}; expected one of " + 'True, False, None, "per-head", or "per-element".' + ) + self.gating = bool(gating) + self.gate_per_head = gating is True or gating == "per-head" attn_tp_rank = get_attention_tp_rank() attn_tp_size = get_attention_tp_size() @@ -287,18 +295,24 @@ class LagunaAttention(nn.Module): prefix=add_prefix("o_proj", prefix), ) - # Per-head softplus gate (`gating=True` in HF). Shard like Q so the - # local output dim matches `num_heads`. - self.g_proj = ColumnParallelLinear( - hidden_size, - self.total_num_heads, - bias=False, - gather_output=False, - quant_config=None, - tp_rank=attn_tp_rank, - tp_size=attn_tp_size, - prefix=add_prefix("g_proj", prefix), - ) + if self.gating: + g_proj_dim = ( + self.total_num_heads + if self.gate_per_head + else self.total_num_heads * self.head_dim + ) + self.g_proj = ColumnParallelLinear( + hidden_size, + g_proj_dim, + bias=False, + gather_output=False, + quant_config=None, + tp_rank=attn_tp_rank, + tp_size=attn_tp_size, + prefix=add_prefix("g_proj", prefix), + ) + else: + self.g_proj = None self.q_norm = RMSNorm(self.head_dim, eps=rms_norm_eps) self.k_norm = RMSNorm(self.head_dim, eps=rms_norm_eps) @@ -348,11 +362,15 @@ class LagunaAttention(nn.Module): attn_output = self.attn(q, k, v, forward_batch) - gate, _ = self.g_proj(hidden_states) - gate = F.softplus(gate.float()).to(attn_output.dtype) - attn_output = attn_output.view(-1, self.num_heads, self.head_dim) - attn_output = attn_output * gate.view(-1, self.num_heads, 1) - attn_output = attn_output.reshape(-1, self.num_heads * self.head_dim) + if self.gating and self.g_proj is not None: + gate, _ = self.g_proj(hidden_states) + gate = F.softplus(gate.float()).to(attn_output.dtype) + if self.gate_per_head: + attn_output = attn_output.view(-1, self.num_heads, self.head_dim) + attn_output = attn_output * gate.view(-1, self.num_heads, 1) + attn_output = attn_output.reshape(-1, self.num_heads * self.head_dim) + else: + attn_output = attn_output * gate output, _ = self.o_proj(attn_output) return output @@ -401,6 +419,7 @@ class LagunaDecoderLayer(nn.Module): # SGLang's window is exclusive; HF's `sliding_window` is inclusive. sliding_window_size=config.sliding_window - 1, layer_type=layer_type, + gating=config.gating, quant_config=quant_config, prefix=add_prefix("self_attn", prefix), ) @@ -748,6 +767,12 @@ class LagunaForCausalLM(nn.Module): if name.endswith(".bias") and name not in params_dict: continue if name not in params_dict: + if ".g_proj." in name: + raise RuntimeError( + f"Checkpoint provides gate weight {name!r} but the model built no " + "g_proj (gating is disabled in the config). Set gating to True, " + '"per-head", or "per-element" to load this checkpoint.' + ) logger.warning("Parameter %s not found in params_dict", name) continue param = params_dict[name]