diff --git a/python/sglang/srt/configs/laguna.py b/python/sglang/srt/configs/laguna.py index 45f9ad8a2..5c05de584 100644 --- a/python/sglang/srt/configs/laguna.py +++ b/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 - # default above always plants one at index 0; let .index() raise if a - # caller passes an all-sliding layer_types — silent fallback would wire - # the SWA head count into a "full" attribute and corrupt downstream sizes. - full_idx = self.layer_types.index("full_attention") + # DFlash draft configs can be all-SWA. In that case there is no full + # layer geometry to expose, so use layer 0 for the default attention + # fields and keep per-layer Q-head geometry explicit. + full_idx = ( + 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 diff --git a/python/sglang/srt/models/dflash.py b/python/sglang/srt/models/dflash.py index a8b2aeefd..5f27b7fb5 100644 --- a/python/sglang/srt/models/dflash.py +++ b/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] diff --git a/python/sglang/srt/models/laguna.py b/python/sglang/srt/models/laguna.py index 691f835ae..b2f6b0e63 100644 --- a/python/sglang/srt/models/laguna.py +++ b/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"): - 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" + gating = normalize_gating(gating) + self.gating = gating != "disabled" + self.gate_per_head = 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 diff --git a/python/sglang/srt/speculative/dflash_worker_v2.py b/python/sglang/srt/speculative/dflash_worker_v2.py index 2553a2927..51f3fadeb 100644 --- a/python/sglang/srt/speculative/dflash_worker_v2.py +++ b/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)