sgl-project--sglang/python/sglang/srt/configs/laguna.py

# coding=utf-8
# Copyright 2023-2026 SGLang Team
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
"""Laguna (poolside/Laguna-XS.2) model configuration."""

from __future__ import annotations

from typing import Any, Dict, List, Literal, Optional

from transformers.configuration_utils import PretrainedConfig
from transformers.utils import logging

logger = logging.get_logger(__name__)


def _first_not_none(*candidates: Any) -> Any:
    """First non-None candidate. Unlike `a or b`, preserves falsy values."""
    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:
        return None
    rope_type = rope_params.get("rope_type") or rope_params.get("type")
    if rope_type in (None, "default"):
        return None

    out: Dict[str, Any] = {"rope_type": rope_type}
    pass_through = (
        "factor",
        "original_max_position_embeddings",
        "beta_fast",
        "beta_slow",
        "extrapolation_factor",
        "truncate",
        "low_freq_factor",
        "high_freq_factor",
        "mscale",
        "mscale_all_dim",
        "short_factor",
        "long_factor",
        "short_mscale",
        "long_mscale",
    )
    for key in pass_through:
        if key in rope_params:
            out[key] = rope_params[key]
    if "attention_factor" in rope_params:
        # HF spells it attention_factor; SGLang's factory reads attn_factor.
        out["attn_factor"] = rope_params["attention_factor"]
    return out


class LagunaConfig(PretrainedConfig):
    model_type = "laguna"
    keys_to_ignore_at_inference = ["past_key_values"]

    def __init__(
        self,
        vocab_size: int = 100352,
        hidden_size: int = 2048,
        intermediate_size: int = 8192,
        num_hidden_layers: int = 40,
        num_attention_heads: int = 48,
        num_key_value_heads: int = 8,
        head_dim: int = 128,
        hidden_act: str = "silu",
        max_position_embeddings: int = 131072,
        initializer_range: float = 0.02,
        rms_norm_eps: float = 1e-6,
        use_cache: bool = True,
        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,
        num_attention_heads_per_layer: Optional[List[int]] = None,
        num_experts: int = 256,
        num_experts_per_tok: int = 8,
        moe_intermediate_size: int = 512,
        shared_expert_intermediate_size: int = 512,
        moe_routed_scaling_factor: float = 1.0,
        moe_router_logit_softcapping: float = 0.0,
        moe_apply_router_weight_on_input: bool = False,
        # Per-layer-type rope dict; nested under "full_attention" / "sliding_attention".
        rope_parameters: Optional[Dict[str, Any]] = None,
        partial_rotary_factor: Optional[float] = None,
        rope_theta: Optional[float] = None,
        rope_scaling: Optional[Dict[str, Any]] = None,
        bos_token_id: Optional[int] = 2,
        eos_token_id: Optional[Any] = None,
        pad_token_id: Optional[int] = 9,
        **kwargs,
    ):
        super().__init__(
            tie_word_embeddings=tie_word_embeddings,
            bos_token_id=bos_token_id,
            eos_token_id=eos_token_id,
            pad_token_id=pad_token_id,
            **kwargs,
        )

        self.vocab_size = vocab_size
        self.hidden_size = hidden_size
        self.intermediate_size = intermediate_size
        self.num_hidden_layers = num_hidden_layers
        self.num_key_value_heads = num_key_value_heads
        self.head_dim = head_dim
        self.hidden_act = hidden_act
        self.max_position_embeddings = max_position_embeddings
        self.initializer_range = initializer_range
        self.rms_norm_eps = rms_norm_eps
        self.use_cache = use_cache
        self.attention_bias = attention_bias
        self.attention_dropout = attention_dropout
        self.gating = normalize_gating(gating)
        self.sliding_window = sliding_window

        self.num_experts = num_experts
        self.num_experts_per_tok = num_experts_per_tok
        self.moe_intermediate_size = moe_intermediate_size
        self.shared_expert_intermediate_size = shared_expert_intermediate_size
        self.moe_routed_scaling_factor = moe_routed_scaling_factor
        self.moe_router_logit_softcapping = moe_router_logit_softcapping
        self.moe_apply_router_weight_on_input = moe_apply_router_weight_on_input

        # Synthesise per-layer schedules when the caller omits them so the model
        # file can index by layer_id without per-call guards.
        self.layer_types = (
            list(layer_types)
            if layer_types
            else ["full_attention" for _ in range(num_hidden_layers)]
        )
        self.mlp_layer_types = (
            list(mlp_layer_types)
            if mlp_layer_types
            else (["dense"] + ["sparse"] * (num_hidden_layers - 1))
        )
        self.num_attention_heads_per_layer = (
            list(num_attention_heads_per_layer)
            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.
        # 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
        self.swa_v_head_dim = head_dim

        # Released checkpoint nests rope_parameters under layer-type keys.
        rp = rope_parameters if isinstance(rope_parameters, dict) else {}
        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
        # form on those two slots (so HF's reference modeling code can index
        # rope_parameters[layer_type] when invoked via trust_remote_code) and
        # publish our SGLang-shaped flat rope dicts under different names.
        self.rope_parameters = rope_parameters

        self.rope_theta = _first_not_none(
            full_rp.get("rope_theta"), rope_theta, 10000.0
        )
        self.partial_rotary_factor = _first_not_none(
            full_rp.get("partial_rotary_factor"), partial_rotary_factor, 1.0
        )
        self.full_rope_scaling = _first_not_none(
            _to_sglang_rope_scaling(full_rp), rope_scaling
        )

        self.swa_rope_theta = _first_not_none(swa_rp.get("rope_theta"), self.rope_theta)
        self.swa_partial_rotary_factor = _first_not_none(
            swa_rp.get("partial_rotary_factor"), self.partial_rotary_factor
        )
        self.swa_rope_scaling = _to_sglang_rope_scaling(swa_rp)

        # DeepSeek-style aliases consumed by cross-cutting infra outside this
        # model file: `lora/mem_pool.py` and `lora/utils.py` read
        # `n_routed_experts` / `n_shared_experts` / `first_k_dense_replace`,
        # `elastic_ep/expert_backup_*` reads `n_routed_experts`. The
        # hardcoded `n_shared_experts=1` and `norm_topk_prob=True` reflect
        # Laguna's fixed architecture (one shared expert, sigmoid-renormalized
        # top-k routing).
        self.n_routed_experts = num_experts
        self.n_shared_experts = 1
        self.routed_scaling_factor = moe_routed_scaling_factor
        self.norm_topk_prob = True
        self.first_k_dense_replace = (
            self.mlp_layer_types.index("sparse")
            if "sparse" in self.mlp_layer_types
            else num_hidden_layers
        )