Add Laguna XS.2.1 DFlash support to SGLang (#29446)

Co-authored-by: Jimmy Shong <69131491+Jiminator@users.noreply.github.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

python/sglang/srt/configs/laguna.py

@@ -10,7 +10,7 @@
 
 from __future__ import annotations
 
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Literal, Optional
 
 from transformers.configuration_utils import PretrainedConfig
 from transformers.utils import logging
@@ -23,6 +23,20 @@ def _first_not_none(*candidates: Any) -> Any:
     return next((c for c in candidates if c is not None), None)
 
 
+def normalize_gating(value: Any) -> Literal["per-head", "per-element", "disabled"]:
+    if value in (True, "per-head"):
+        return "per-head"
+    if value == "per-element":
+        return "per-element"
+    if value in (False, None, "disabled"):
+        return "disabled"
+    raise ValueError(
+        "gating must be one of True, False, None, "
+        '"per-head", "per-element", or "disabled"; '
+        f"got {value!r}."
+    )
+
+
 def _to_sglang_rope_scaling(rope_params: Dict[str, Any]) -> Optional[Dict[str, Any]]:
     """HF per-layer rope dict → SGLang `get_rope` `rope_scaling`. None means plain RoPE."""
     if not rope_params:
@@ -121,7 +135,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.gating = normalize_gating(gating)
         self.sliding_window = sliding_window
 
         self.num_experts = num_experts
@@ -149,14 +163,23 @@ class LagunaConfig(PretrainedConfig):
             if (num_attention_heads_per_layer)
             else [num_attention_heads] * num_hidden_layers
         )
+        if len(self.num_attention_heads_per_layer) != num_hidden_layers:
+            raise ValueError(
+                "num_attention_heads_per_layer must have one entry per layer: "
+                f"expected num_hidden_layers={num_hidden_layers}, "
+                f"got {len(self.num_attention_heads_per_layer)}."
+            )
 
         # SGLang's hybrid-SWA core reads `swa_*` KV/head_dim from hf_text_config.
         # Per-layer Q-head count is read directly from num_attention_heads_per_layer.
-        # Pure-SWA models would have no full_attention layer, but the synthesized
+        # DFlash draft configs can be all-SWA. In that case there is no full
-        # default above always plants one at index 0; let .index() raise if a
+        # layer geometry to expose, so use layer 0 for the default attention
-        # caller passes an all-sliding layer_types — silent fallback would wire
+        # fields and keep per-layer Q-head geometry explicit.
-        # the SWA head count into a "full" attribute and corrupt downstream sizes.
+        full_idx = (
-        full_idx = self.layer_types.index("full_attention")
+            self.layer_types.index("full_attention")
+            if "full_attention" in self.layer_types
+            else 0
+        )
         self.num_attention_heads = self.num_attention_heads_per_layer[full_idx]
         self.swa_num_key_value_heads = num_key_value_heads
         self.swa_head_dim = head_dim
@@ -164,8 +187,9 @@ class LagunaConfig(PretrainedConfig):
 
         # Released checkpoint nests rope_parameters under layer-type keys.
         rp = rope_parameters if isinstance(rope_parameters, dict) else {}
-        full_rp = rp.get("full_attention") or {}
+        has_full_attention = "full_attention" in self.layer_types
         swa_rp = rp.get("sliding_attention") or {}
+        full_rp = rp.get("full_attention") or (swa_rp if not has_full_attention else {})
 
         # transformers v5 aliases `rope_scaling` ↔ `rope_parameters` on
         # PretrainedConfig — writing one clobbers the other. Keep the nested

python/sglang/srt/models/dflash.py

@@ -12,9 +12,11 @@ import torch
 import torch.nn.functional as F
 from torch import nn
 
+from sglang.srt.configs.laguna import normalize_gating
 from sglang.srt.layers.activation import SiluAndMul
 from sglang.srt.layers.layernorm import RMSNorm
 from sglang.srt.layers.linear import (
+    ColumnParallelLinear,
     MergedColumnParallelLinear,
     QKVParallelLinear,
     RowParallelLinear,
@@ -67,7 +69,7 @@ def _get_dflash_layer_attention_params(
 
 
 class DFlashAttention(nn.Module):
-    def __init__(self, config, layer_id: int) -> None:
+    def __init__(self, config, layer_id: int, quant_config=None) -> None:
         super().__init__()
         hidden_size = int(config.hidden_size)
         tp_size = int(get_parallel().tp_size)
@@ -109,12 +111,14 @@ class DFlashAttention(nn.Module):
             total_num_heads=self.total_num_heads,
             total_num_kv_heads=self.total_num_kv_heads,
             bias=attention_bias,
+            quant_config=quant_config,
             prefix="qkv_proj",
         )
         self.o_proj = RowParallelLinear(
             self.total_num_heads * head_dim,
             hidden_size,
             bias=attention_bias,
+            quant_config=quant_config,
             prefix="o_proj",
         )
 
@@ -186,9 +190,15 @@ class DFlashAttention(nn.Module):
             q, k = apply_qk_norm(q, k, self.q_norm, self.k_norm, self.head_dim)
             q, k = self.rotary_emb(positions, q, k)
         attn_output = self.attn(q, k, v, forward_batch)
+        attn_output = self.apply_attention_output(attn_output, hidden_states)
         output, _ = self.o_proj(attn_output)
         return output
 
+    def apply_attention_output(
+        self, attn_output: torch.Tensor, hidden_states: torch.Tensor
+    ) -> torch.Tensor:
+        return attn_output
+
     def kv_proj_only(
         self, hidden_states: torch.Tensor
     ) -> Tuple[torch.Tensor, torch.Tensor]:
@@ -265,15 +275,19 @@ class DFlashMLP(nn.Module):
 
 
 class DFlashDecoderLayer(nn.Module):
-    def __init__(self, config, layer_id: int) -> None:
+    attention_cls = DFlashAttention
+
+    def __init__(self, config, layer_id: int, quant_config=None) -> None:
         super().__init__()
         hidden_size = int(config.hidden_size)
         rms_norm_eps = float(getattr(config, "rms_norm_eps", 1e-6))
 
         self.input_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps)
-        self.self_attn = DFlashAttention(config=config, layer_id=layer_id)
+        self.self_attn = self.attention_cls(
+            config=config, layer_id=layer_id, quant_config=quant_config
+        )
         self.post_attention_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps)
-        self.mlp = DFlashMLP(config=config)
+        self.mlp = DFlashMLP(config=config, quant_config=quant_config)
 
     def forward(
         self,
@@ -314,6 +328,9 @@ class DFlashDraftModel(nn.Module):
       - `norm.weight` for final normalization
     """
 
+    decoder_layer_cls = DFlashDecoderLayer
+    supports_fused_context_kv = True
+
     def __init__(self, config, quant_config=None, prefix: str = "") -> None:
         super().__init__()
         self.config = config
@@ -323,7 +340,12 @@ class DFlashDraftModel(nn.Module):
         rms_norm_eps = float(getattr(config, "rms_norm_eps", 1e-6))
 
         self.layers = nn.ModuleList(
-            [DFlashDecoderLayer(config=config, layer_id=i) for i in range(num_layers)]
+            [
+                self.decoder_layer_cls(
+                    config=config, layer_id=i, quant_config=quant_config
+                )
+                for i in range(num_layers)
+            ]
         )
         self.norm = RMSNorm(hidden_size, eps=rms_norm_eps)
 
@@ -352,6 +374,11 @@ class DFlashDraftModel(nn.Module):
     def get_attention_sliding_window_size(self) -> Optional[int]:
         return get_dflash_attention_sliding_window_size(self.config)
 
+    def prepare_context_hidden_for_kv(
+        self, layer: DFlashDecoderLayer, ctx_hidden: torch.Tensor
+    ) -> torch.Tensor:
+        return ctx_hidden
+
     def project_target_hidden(self, target_hidden: torch.Tensor) -> torch.Tensor:
         """Project concatenated target-layer hidden states into draft hidden_size."""
         expected = int(self.fc.in_features)
@@ -456,4 +483,97 @@ class DFlashDraftModel(nn.Module):
                 weight_loader(param, loaded_weight)
 
 
-EntryClass = DFlashDraftModel
+class DFlashLagunaAttention(DFlashAttention):
+    """Laguna DFlash attention with the trained Laguna softplus gate."""
+
+    def __init__(self, config, layer_id: int, quant_config=None) -> None:
+        super().__init__(config=config, layer_id=layer_id, quant_config=quant_config)
+        hidden_size = int(config.hidden_size)
+        total_num_heads = self.total_num_heads
+        gating = normalize_gating(getattr(config, "gating", True))
+        self.gating = gating
+        self.gate_per_head = gating == "per-head"
+        if self.gating == "disabled":
+            self.g_proj = None
+        else:
+            g_out = (
+                total_num_heads
+                if self.gate_per_head
+                else total_num_heads * self.head_dim
+            )
+            self.g_proj = ColumnParallelLinear(
+                hidden_size,
+                g_out,
+                bias=False,
+                quant_config=quant_config,
+                prefix="g_proj",
+            )
+
+    def apply_attention_output(
+        self, attn_output: torch.Tensor, hidden_states: torch.Tensor
+    ) -> torch.Tensor:
+        if self.g_proj is None:
+            return attn_output
+
+        gate, _ = self.g_proj(hidden_states)
+        gate = F.softplus(gate.float()).to(attn_output.dtype)
+        if self.gate_per_head:
+            attn_shape = attn_output.shape
+            return (
+                attn_output.view(*attn_shape[:-1], self.num_heads, self.head_dim)
+                * gate.unsqueeze(-1)
+            ).view(attn_shape)
+        else:
+            return attn_output * gate
+
+
+class DFlashLagunaDecoderLayer(DFlashDecoderLayer):
+    attention_cls = DFlashLagunaAttention
+
+
+class DFlashLagunaForCausalLM(DFlashDraftModel):
+    """Laguna DFlash draft model matching the exported Speculators checkpoint."""
+
+    decoder_layer_cls = DFlashLagunaDecoderLayer
+    supports_fused_context_kv = False
+
+    def __init__(self, config, quant_config=None, prefix: str = "") -> None:
+        super().__init__(config=config, quant_config=quant_config, prefix=prefix)
+        rms_norm_eps = float(getattr(config, "rms_norm_eps", 1e-6))
+        hidden_size = int(config.hidden_size)
+        self.aux_hidden_norms = nn.ModuleList(
+            [
+                RMSNorm(hidden_size, eps=rms_norm_eps)
+                for _ in range(self.num_context_features)
+            ]
+        )
+
+    def prepare_context_hidden_for_kv(
+        self, layer: DFlashLagunaDecoderLayer, ctx_hidden: torch.Tensor
+    ) -> torch.Tensor:
+        return layer.input_layernorm(ctx_hidden)
+
+    def project_target_hidden(self, target_hidden: torch.Tensor) -> torch.Tensor:
+        expected = int(self.fc.in_features)
+        if target_hidden.ndim != 2 or int(target_hidden.shape[-1]) != expected:
+            raise ValueError(
+                "Laguna DFLASH target_hidden feature dim mismatch. "
+                f"Expected shape [N, {expected}] "
+                f"(num_context_features={self.num_context_features}, hidden_size={int(self.config.hidden_size)}), "
+                f"but got shape={tuple(target_hidden.shape)}."
+            )
+
+        num_slices = int(self.num_context_features)
+        slice_size = int(target_hidden.shape[-1]) // num_slices
+        slices = target_hidden.view(target_hidden.shape[0], num_slices, slice_size)
+        compute_dtype = self.fc.weight.dtype
+        if slices.dtype != compute_dtype:
+            slices = slices.to(compute_dtype)
+        normed = torch.empty_like(slices)
+        for i, norm in enumerate(self.aux_hidden_norms):
+            normed[:, i, :] = norm(slices[:, i, :])
+        fused = normed.reshape(target_hidden.shape[0], -1)
+        return self.hidden_norm(self.fc(fused))
+
+
+EntryClass = [DFlashDraftModel, DFlashLagunaForCausalLM]

python/sglang/srt/models/laguna.py

@@ -10,17 +10,18 @@ from __future__ import annotations
 
 import logging
 from collections.abc import Iterable
-from typing import Any, Dict, Optional, Tuple, Union
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 import torch
 import torch.nn.functional as F
 from torch import nn
 
-from sglang.srt.configs.laguna import LagunaConfig
+from sglang.srt.configs.laguna import LagunaConfig, normalize_gating
 from sglang.srt.distributed import (
     get_pp_group,
     tensor_model_parallel_all_reduce,
 )
+from sglang.srt.environ import envs
 from sglang.srt.layers.activation import SiluAndMul
 from sglang.srt.layers.communicator import (
     LayerCommunicator,
@@ -68,6 +69,8 @@ class LagunaMLP(nn.Module):
         quant_config: Optional[QuantizationConfig] = None,
         reduce_results: bool = True,
         prefix: str = "",
+        tp_rank: Optional[int] = None,
+        tp_size: Optional[int] = None,
     ) -> None:
         super().__init__()
         if hidden_act != "silu":
@@ -80,6 +83,8 @@ class LagunaMLP(nn.Module):
             bias=False,
             quant_config=quant_config,
             prefix=add_prefix("gate_up_proj", prefix),
+            tp_rank=tp_rank,
+            tp_size=tp_size,
         )
         self.down_proj = RowParallelLinear(
             intermediate_size,
@@ -88,6 +93,8 @@ class LagunaMLP(nn.Module):
             quant_config=quant_config,
             reduce_results=reduce_results,
             prefix=add_prefix("down_proj", prefix),
+            tp_rank=tp_rank,
+            tp_size=tp_size,
         )
         self.act_fn = SiluAndMul()
 
@@ -177,6 +184,11 @@ class LagunaMoE(nn.Module):
 
         # HF safetensors key is singular `shared_expert.…`; mirror so the
         # default loader picks it up without remapping.
+        # SGLANG_SHARED_EXPERT_TP1 replicates the shared expert instead of
+        # TP-sharding it, for checkpoints whose shared-expert quant scales are
+        # not divisible by the global TP size (e.g. block-FP8 [128,128] with
+        # shared_expert_intermediate_size=512 at TP=8 → 64-per-rank shards).
+        self._shared_expert_tp1 = envs.SGLANG_SHARED_EXPERT_TP1.get()
         self.shared_expert = LagunaMLP(
             hidden_size=config.hidden_size,
             intermediate_size=config.shared_expert_intermediate_size,
@@ -184,6 +196,7 @@ class LagunaMoE(nn.Module):
             quant_config=quant_config,
             reduce_results=False,
             prefix=add_prefix("shared_expert", prefix),
+            **(dict(tp_rank=0, tp_size=1) if self._shared_expert_tp1 else {}),
         )
 
     def get_moe_weights(self):
@@ -212,7 +225,13 @@ class LagunaMoE(nn.Module):
         # so scale routed manually before adding the unscaled shared expert.
         if self.routed_scaling_factor != 1.0:
             routed_out = routed_out * self.routed_scaling_factor
-        final = routed_out + shared_out
+        # A TP1 (replicated) shared expert already holds the full result on
+        # every rank, so it must be added after the all-reduce — adding before
+        # would sum it once per TP rank.
+        if self._shared_expert_tp1:
+            final = routed_out
+        else:
+            final = routed_out + shared_out
 
         if self.tp_size > 1 and not should_skip_post_experts_all_reduce(
             is_tp_path=True,
@@ -220,6 +239,8 @@ class LagunaMoE(nn.Module):
             should_allreduce_fusion=should_allreduce_fusion,
         ):
             final = tensor_model_parallel_all_reduce(final)
+        if self._shared_expert_tp1:
+            final = final + shared_out
         return final
 
 
@@ -247,13 +268,9 @@ 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"):
+        gating = normalize_gating(gating)
-            raise ValueError(
+        self.gating = gating != "disabled"
-                f"Unsupported gating value {gating!r}; expected one of "
+        self.gate_per_head = gating == "per-head"
-                '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_parallel().attn_tp_rank
         attn_tp_size = get_parallel().attn_tp_size
@@ -553,6 +570,7 @@ class LagunaModel(nn.Module):
             self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
         else:
             self.norm = PPMissingLayer(return_tuple=True)
+        self.layers_to_capture: List[int] = []
 
     def get_input_embeddings(self) -> nn.Embedding:
         return self.embed_tokens
@@ -576,7 +594,12 @@ class LagunaModel(nn.Module):
             hidden_states = pp_proxy_tensors["hidden_states"]
             residual = pp_proxy_tensors["residual"]
 
+        aux_hidden_states = []
         for i in range(self.start_layer, self.end_layer):
+            if i in self.layers_to_capture:
+                aux_hidden_states.append(
+                    hidden_states + residual if residual is not None else hidden_states
+                )
             layer = self.layers[i]
             hidden_states, residual = layer(
                 positions, hidden_states, forward_batch, residual
@@ -588,11 +611,17 @@ class LagunaModel(nn.Module):
             )
 
         if hidden_states.shape[0] != 0:
+            if self.end_layer in self.layers_to_capture:
+                aux_hidden_states.append(
+                    hidden_states + residual if residual is not None else hidden_states
+                )
             if residual is None:
                 hidden_states = self.norm(hidden_states)
             else:
                 hidden_states, _ = self.norm(hidden_states, residual)
-        return hidden_states
+        if len(aux_hidden_states) == 0:
+            return hidden_states
+        return hidden_states, aux_hidden_states
 
 
 class LagunaForCausalLM(nn.Module):
@@ -625,6 +654,7 @@ class LagunaForCausalLM(nn.Module):
         else:
             self.lm_head = PPMissingLayer()
         self.logits_processor = LogitsProcessor(config)
+        self.capture_aux_hidden_states = False
 
         # Only walk this rank's local layers — out-of-range entries can be PPMissingLayer.
         self._routed_experts_weights_of_layer = LazyValue(
@@ -663,9 +693,12 @@ class LagunaForCausalLM(nn.Module):
             input_embeds,
             pp_proxy_tensors=pp_proxy_tensors,
         )
+        aux_hidden_states = None
+        if self.capture_aux_hidden_states:
+            hidden_states, aux_hidden_states = hidden_states
         if self.pp_group.is_last_rank:
             return self.logits_processor(
-                input_ids, hidden_states, self.lm_head, forward_batch
+                input_ids, hidden_states, self.lm_head, forward_batch, aux_hidden_states
             )
         return hidden_states
 
@@ -806,5 +839,18 @@ class LagunaForCausalLM(nn.Module):
         torch.cuda.empty_cache()
         torch.cuda.synchronize()
 
+    def set_dflash_layers_to_capture(self, layer_ids: List[int]):
+        if not self.pp_group.is_last_rank:
+            return
+        if layer_ids is None:
+            raise ValueError(
+                "DFLASH requires explicit layer_ids for aux hidden capture."
+            )
+
+        self.capture_aux_hidden_states = True
+        # SGLang captures "before layer i". To capture the hidden state after
+        # target layer `k` (HF-style), capture before layer `k + 1`.
+        self.model.layers_to_capture = [val + 1 for val in layer_ids]
+
 
 EntryClass = LagunaForCausalLM

python/sglang/srt/speculative/dflash_worker_v2.py

@@ -392,6 +392,18 @@ class DFlashWorkerV2(BaseSpecWorker):
 
             if len(layers) == 0:
                 fused_disable_reason = "no layers found"
+            elif not getattr(self.draft_model, "supports_fused_context_kv", True):
+                fused_disable_reason = "draft model does not support fused context KV"
+
+            if fused_disable_reason is not None:
+                if self.tp_rank == 0:
+                    logger.info(
+                        "DFLASH fused KV materialization disabled: %s",
+                        fused_disable_reason,
+                    )
+                self._use_fused_kv_materialize = False
+                self._fused_kv_helper = None
+                return
 
             for layer_idx, layer in enumerate(layers):
                 attn = layer.self_attn
@@ -1008,7 +1020,10 @@ class DFlashWorkerV2(BaseSpecWorker):
 
                 for layer in self.draft_model.layers:
                     attn = layer.self_attn
-                    k, v = attn.kv_proj_only(ctx_hidden)
+                    layer_ctx_hidden = self.draft_model.prepare_context_hidden_for_kv(
+                        layer, ctx_hidden
+                    )
+                    k, v = attn.kv_proj_only(layer_ctx_hidden)
                     k = attn.apply_k_norm(k)
                     k = attn.apply_k_rope(positions, k)
                     k = k.view(-1, attn.num_kv_heads, attn.head_dim)
@@ -1055,10 +1070,13 @@ class DFlashWorkerV2(BaseSpecWorker):
     ) -> None:
         for layer in self.draft_model.layers:
             attn = layer.self_attn
+            layer_ctx_hidden = self.draft_model.prepare_context_hidden_for_kv(
+                layer, ctx_hidden
+            )
             if _is_npu:
-                _, k, v = attn.forward_prepare_npu(ctx_positions, ctx_hidden)
+                _, k, v = attn.forward_prepare_npu(ctx_positions, layer_ctx_hidden)
             else:
-                k, v = attn.kv_proj_only(ctx_hidden)
+                k, v = attn.kv_proj_only(layer_ctx_hidden)
                 k = attn.apply_k_norm(k)
                 k = attn.apply_k_rope(ctx_positions, k)
             k = k.view(-1, attn.num_kv_heads, attn.head_dim)