[Model] Laguna: support per-element output gating (#28400)

Co-authored-by: Jimmy Shong <69131491+Jiminator@users.noreply.github.com>
Co-authored-by: Jimmy Shong <jimmy.shong@radixark.ai>

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

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,11 +295,15 @@ class LagunaAttention(nn.Module):
             prefix=add_prefix("o_proj", prefix),
         )
 
-        # Per-head softplus gate (`gating=True` in HF). Shard like Q so the
+        if self.gating:
-        # local output dim matches `num_heads`.
+            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,
-            self.total_num_heads,
+                g_proj_dim,
                 bias=False,
                 gather_output=False,
                 quant_config=None,
@@ -299,6 +311,8 @@ class LagunaAttention(nn.Module):
                 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)
 
+        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]