diff --git a/docs_new/docs/supported-models/generative_models.mdx b/docs_new/docs/supported-models/generative_models.mdx index b5258001b..eb83f3e9d 100644 --- a/docs_new/docs/supported-models/generative_models.mdx +++ b/docs_new/docs/supported-models/generative_models.mdx @@ -283,5 +283,10 @@ in the GitHub search bar. sarvamai/sarvam-2 Sarvam's Mixture-of-Experts models. The 105B variant uses MLA (Multi-head Latent Attention) and the 30B variant uses GQA, both with 128 routed experts. + + Laguna XS.2 (poolside) + poolside/Laguna-XS.2 + Poolside's hybrid sliding-window-attention MoE model (256 routed experts, 1 shared expert, sigmoid router) with per-layer-type RoPE (YARN on full-attention layers, default on sliding-attention layers) and a softplus per-head attention gate. + diff --git a/python/sglang/srt/configs/__init__.py b/python/sglang/srt/configs/__init__.py index bbc121eed..35e3193eb 100644 --- a/python/sglang/srt/configs/__init__.py +++ b/python/sglang/srt/configs/__init__.py @@ -15,6 +15,7 @@ from sglang.srt.configs.kimi_k25 import KimiK25Config from sglang.srt.configs.kimi_linear import KimiLinearConfig from sglang.srt.configs.kimi_vl import KimiVLConfig from sglang.srt.configs.kimi_vl_moonvit import MoonViTConfig +from sglang.srt.configs.laguna import LagunaConfig from sglang.srt.configs.lfm2 import Lfm2Config from sglang.srt.configs.lfm2_moe import Lfm2MoeConfig from sglang.srt.configs.lfm2_vl import Lfm2VlConfig @@ -52,6 +53,7 @@ __all__ = [ "Olmo3Config", "KimiLinearConfig", "KimiK25Config", + "LagunaConfig", "Qwen3NextConfig", "Qwen3_5Config", "Qwen3_5MoeConfig", diff --git a/python/sglang/srt/configs/laguna.py b/python/sglang/srt/configs/laguna.py new file mode 100644 index 000000000..ff11f0269 --- /dev/null +++ b/python/sglang/srt/configs/laguna.py @@ -0,0 +1,209 @@ +# 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, 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 _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, + 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.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" if i % 4 == 0 else "sliding_attention" + for i 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 + ) + + # 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") + 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 {} + full_rp = rp.get("full_attention") or {} + swa_rp = rp.get("sliding_attention") or {} + + # 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 + ) diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index 0e62338f6..2529012f7 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -1660,6 +1660,7 @@ def is_hybrid_swa_model(model_architectures: List[str]): "Step3p5MTP", "Gemma4ForCausalLM", "Gemma4ForConditionalGeneration", + "LagunaForCausalLM", } return any(arch in hybrid_swa_archs for arch in model_architectures) @@ -1721,6 +1722,14 @@ def get_hybrid_layer_ids( full_attention_layer_ids = [ i for i, x in enumerate(layer_types) if x == "full_attention" ] + elif "LagunaForCausalLM" in model_architectures: + layer_types = getattr(hf_text_config, "layer_types", []) + swa_attention_layer_ids = [ + i for i, x in enumerate(layer_types) if x == "sliding_attention" + ] + full_attention_layer_ids = [ + i for i, x in enumerate(layer_types) if x == "full_attention" + ] else: swa_attention_layer_ids = None full_attention_layer_ids = None diff --git a/python/sglang/srt/function_call/function_call_parser.py b/python/sglang/srt/function_call/function_call_parser.py index 8123da197..602e93fdd 100644 --- a/python/sglang/srt/function_call/function_call_parser.py +++ b/python/sglang/srt/function_call/function_call_parser.py @@ -30,6 +30,7 @@ from sglang.srt.function_call.llama32_detector import Llama32Detector from sglang.srt.function_call.mimo_detector import MiMoDetector from sglang.srt.function_call.minimax_m2 import MinimaxM2Detector from sglang.srt.function_call.mistral_detector import MistralDetector +from sglang.srt.function_call.poolside_v1_detector import PoolsideV1Detector from sglang.srt.function_call.pythonic_detector import PythonicDetector from sglang.srt.function_call.qwen3_coder_detector import Qwen3CoderDetector from sglang.srt.function_call.qwen25_detector import Qwen25Detector @@ -66,6 +67,7 @@ class FunctionCallParser: "llama3": Llama32Detector, "mimo": MiMoDetector, "mistral": MistralDetector, + "poolside_v1": PoolsideV1Detector, "pythonic": PythonicDetector, "qwen": Qwen25Detector, "qwen25": Qwen25Detector, diff --git a/python/sglang/srt/function_call/poolside_v1_detector.py b/python/sglang/srt/function_call/poolside_v1_detector.py new file mode 100644 index 000000000..4261d9060 --- /dev/null +++ b/python/sglang/srt/function_call/poolside_v1_detector.py @@ -0,0 +1,452 @@ +import ast +import json +import re +from enum import Enum, auto +from typing import Any, List, Optional + +from sglang.srt.entrypoints.openai.protocol import Tool +from sglang.srt.function_call.base_format_detector import BaseFormatDetector +from sglang.srt.function_call.core_types import ( + StreamingParseResult, + StructureInfo, + ToolCallItem, + _GetInfoFunc, +) + + +class _ParseState(Enum): + """5 FSM states for the streaming parser. + + Entry guard: READING_VALUE is reachable only from READING_KEY, so the + "stray before " bug class is structurally + impossible. + + Exit guard: both READING_KEY and READING_VALUE recover on `` + by closing the active call (orphan key dropped if any). READING_VALUE + additionally recovers on `` by replacing the orphan pending key + with the new one. Both guards match the (regex-tightened) non-streaming + path. Without them, malformed inputs would leave the FSM stuck in + READING_VALUE and mis-attribute subsequent values to stale state. + """ + + OUTSIDE = auto() + READING_NAME = auto() + READING_KEY = auto() + READING_VALUE = auto() + DRAINING = auto() + + +class PoolsideV1Detector(BaseFormatDetector): + """ + Detector for poolside Laguna-XS.2 (poolside_v1 series) tool-call wire format. + + Wire format: + {name}\\n + {key}\\n + {val}\\n + ... + + + String values are emitted as raw text; non-strings are JSON-encoded by + the chat template. The parser does schema-based type coercion to round-trip + them: schema type `string` keeps the raw value; other types attempt + `json.loads` and fall back to `ast.literal_eval`, then to the raw string. + """ + + # Wire-format tag tokens — constants, not per-instance. + tool_call_start_token = "" + tool_call_end_token = "" + arg_key_start = "" + arg_key_end = "" + arg_value_start = "" + arg_value_end = "" + + tool_call_regex = re.compile(r"(.*?)", re.DOTALL) + # Key uses [^<]*? to prevent the non-greedy `.*?` from backtracking + # across an `` boundary on malformed inputs like + # `K1K2V` + # — without the `[^<]` constraint, the regex matches the entire orphan + # span as a single key (`K1K2`). Param names never + # contain `<` in practice, so this is safe. The value side keeps `.*?` + # because legitimate values can contain `<` (HTML, paths, etc.); the + # `` boundary is anchored enough. + arg_pair_regex = re.compile( + r"([^<]*?)\s*(.*?)", + re.DOTALL, + ) + + _partial_tag_prefixes = ( + tool_call_start_token, + tool_call_end_token, + arg_key_start, + arg_key_end, + arg_value_start, + arg_value_end, + ) + + def __init__(self): + super().__init__() + self.parsed_pos: int = 0 + self._state: _ParseState = _ParseState.OUTSIDE + self.current_func_name: Optional[str] = None + self.current_pending_key: Optional[str] = None + self.json_started: bool = False + + # ---------- Helpers ---------- + + def _reset_call_state(self) -> None: + """Reset per-call FSM scratch fields. Called when entering a new + and on close.""" + self.current_func_name = None + self.current_pending_key = None + self.json_started = False + + def _consume_arg_key(self, slice_: str) -> bool: + """Consume `K`, set `current_pending_key` to K. + Returns True if consumed, False if `` hasn't arrived yet + (caller should break to wait for more bytes). Shared by READING_KEY + (well-formed: transitions to READING_VALUE) and READING_VALUE + (orphan-key-replace: stays in READING_VALUE).""" + end = slice_.find(self.arg_key_end) + if end == -1: + return False + self.current_pending_key = slice_[len(self.arg_key_start) : end].strip() + self.parsed_pos += end + len(self.arg_key_end) + return True + + def _close_current_call(self, calls: List[ToolCallItem]) -> None: + """Emit the closing `}` (or `{}` for zero-arg) for the active call, + advance past ``, return to OUTSIDE, and reset per-call + state. Called from both READING_KEY (the well-formed close path) and + READING_VALUE (malformed close: `...` + with no value — orphan key is discarded, matching the regex + non-streaming path which drops unmatched ... + pairs).""" + fragment = "}" if self.json_started else "{}" + calls.append( + ToolCallItem( + tool_index=self.current_tool_id, + parameters=fragment, + ) + ) + self.streamed_args_for_tool[self.current_tool_id] += fragment + self.parsed_pos += len(self.tool_call_end_token) + self._state = _ParseState.OUTSIDE + self._reset_call_state() + + def has_tool_call(self, text: str) -> bool: + return self.tool_call_start_token in text + + @staticmethod + def _get_param_schema( + func_name: Optional[str], tools: Optional[List[Tool]] + ) -> dict: + if not tools or not func_name: + return {} + for tool in tools: + try: + if ( + tool.type == "function" + and tool.function.name == func_name + and isinstance(tool.function.parameters, dict) + ): + return tool.function.parameters.get("properties", {}) + except AttributeError: + continue + return {} + + _STRING_TYPES = frozenset({"string", "str", "text", "enum"}) + + @staticmethod + def _convert_param_value(raw: str, schema: dict, key: str) -> Any: + """Coerce a raw arg_value string per schema; fall back to raw on failure. + + Decoder selection by schema type: + - string-like types → identity (raw text) + - no schema entry → json.loads only (conservative; don't + ast-eval untyped values) + - everything else (int, + number, bool, object, …) → json.loads, then ast.literal_eval + + Each decoder result is round-tripped through `json.dumps` before being + returned; non-JSON-serializable values (sets / complex / bytes from + `ast.literal_eval`) are rejected to the next decoder, ultimately + falling through to the raw-string fallback rather than crashing the + streaming JSON emission downstream. + """ + spec = schema.get(key) if isinstance(schema, dict) else None + param_type = str(spec.get("type", "")).lower() if isinstance(spec, dict) else "" + if param_type in PoolsideV1Detector._STRING_TYPES: + return raw + + decoders = (json.loads,) if not param_type else (json.loads, ast.literal_eval) + for decoder in decoders: + try: + result = decoder(raw) + # ast.literal_eval can return non-JSON-serializable values + # (sets, complex numbers); reject so json.dumps downstream + # doesn't choke. + json.dumps(result) + return result + except (ValueError, SyntaxError, TypeError): + continue + return raw + + def _find_name_boundary(self, text: str) -> int: + """Earliest of `\\n`, ``, ``. -1 if none.""" + hits = ( + text.find("\n"), + text.find(self.arg_key_start), + text.find(self.tool_call_end_token), + ) + positive = [h for h in hits if h != -1] + return min(positive) if positive else -1 + + def _is_partial_tag(self, slice_: str) -> bool: + """True if slice_ is a strict prefix of any known tag — i.e. more + bytes might complete it into a real tag.""" + return any( + tag.startswith(slice_) and tag != slice_ + for tag in self._partial_tag_prefixes + ) + + # ---------- Non-streaming ---------- + + def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult: + if self.tool_call_start_token not in text: + return StreamingParseResult(normal_text=text) + + tool_indices = self._get_tool_indices(tools) + first_idx = text.find(self.tool_call_start_token) + normal_text = text[:first_idx] if first_idx > 0 else "" + + calls: List[ToolCallItem] = [] + for body in self.tool_call_regex.findall(text): + # _find_name_boundary searches for `\n` / `` / + # ``, but the regex already stripped ``, + # so a no-arg call without a trailing newline + # (`now`) gives boundary == -1. Treat + # that case as "name == entire body". + boundary = self._find_name_boundary(body) + name = (body if boundary == -1 else body[:boundary]).strip() + if not name or name not in tool_indices: + continue + + schema = self._get_param_schema(name, tools) + args: dict = {} + for raw_key, raw_val in self.arg_pair_regex.findall(body): + key = raw_key.strip() + # Strip at most one wrapping `\n` on each side (template adds + # them around the value); preserve newlines that are part of + # the value itself. + val = raw_val.removeprefix("\n").removesuffix("\n") + args[key] = self._convert_param_value(val, schema, key) + + calls.append( + ToolCallItem( + tool_index=tool_indices[name], + name=name, + parameters=json.dumps(args, ensure_ascii=False), + ) + ) + + return StreamingParseResult(normal_text=normal_text, calls=calls) + + # ---------- Streaming ---------- + + def parse_streaming_increment( + self, new_text: str, tools: List[Tool] + ) -> StreamingParseResult: + self._buffer += new_text + if not self._buffer: + return StreamingParseResult() + + tool_indices = self._get_tool_indices(tools) + calls: List[ToolCallItem] = [] + normal_text_chunks: List[str] = [] + + # No try/except: the FSM's invariants make the prior masked-IndexError + # class unreachable, and TypeError from json.dumps is prevented at the + # source (_convert_param_value round-trips its decoder output). If a + # real bug surfaces, let it surface. + while True: + slice_ = self._buffer[self.parsed_pos :] + if not slice_: + break + state = self._state + + if state is _ParseState.OUTSIDE: + if slice_.startswith(self.tool_call_start_token): + self.parsed_pos += len(self.tool_call_start_token) + self._state = _ParseState.READING_NAME + self._reset_call_state() + continue + if slice_.startswith("<"): + if self._is_partial_tag(slice_): + break # could be a partial + normal_text_chunks.append("<") + self.parsed_pos += 1 + continue + next_lt = slice_.find("<") + segment = slice_ if next_lt == -1 else slice_[:next_lt] + normal_text_chunks.append(segment) + self.parsed_pos += len(segment) + continue + + if state is _ParseState.READING_NAME: + boundary = self._find_name_boundary(slice_) + if boundary == -1: + break # name still incoming + name = slice_[:boundary].strip() + # Consume the name and a single delimiting newline (if + # present). The other boundary types (, + # ) are left for the next state. boundary may + # be 0 for a malformed `...` (no + # name); the state transition below is the loop-progress + # guarantee. + consume = boundary + if boundary < len(slice_) and slice_[boundary : boundary + 1] == "\n": + consume += 1 + self.parsed_pos += consume + + if name and name in tool_indices: + self.current_tool_id += 1 + while len(self.streamed_args_for_tool) <= self.current_tool_id: + self.streamed_args_for_tool.append("") + self.current_func_name = name + # Per-response sequential index — OpenAI clients group + # chunks by tool_index, so the name event and later + # parameter fragments must share this value. + calls.append( + ToolCallItem( + tool_index=self.current_tool_id, + name=name, + parameters="", + ) + ) + self._state = _ParseState.READING_KEY + else: + # Unknown / empty name — drain to with no + # client-visible emission. + self._state = _ParseState.DRAINING + continue + + if state is _ParseState.READING_KEY: + if slice_.startswith(self.tool_call_end_token): + self._close_current_call(calls) + continue + if slice_.startswith(self.arg_key_start): + if not self._consume_arg_key(slice_): + break # incomplete + self._state = _ParseState.READING_VALUE + continue + if slice_.startswith("<"): + if self._is_partial_tag(slice_): + break + # Bare '<' that's not any known tag — discard silently + # (inside a tool call, this is not normal_text). + self.parsed_pos += 1 + continue + # Inter-tag whitespace / newline — discard. + next_lt = slice_.find("<") + self.parsed_pos += len(slice_) if next_lt == -1 else next_lt + continue + + if state is _ParseState.READING_VALUE: + # Recover from a malformed `K` + # (no ) by closing the call here. Without this + # branch the FSM would stay stuck in READING_VALUE and + # mis-attribute the next call's to the orphan + # `current_pending_key`, silently swallowing the next call's + # name. Matches the regex non-streaming path, which drops + # unmatched ... pairs. + if slice_.startswith(self.tool_call_end_token): + self._close_current_call(calls) + continue + # Recover from a malformed `K1K2` + # (no value for K1, model went straight to a new key) by + # replacing the orphan pending_key with the new one. Stays + # in READING_VALUE so the next binds to K2. + # Without this branch the FSM treats the second + # as bare-`<` garbage and the next binds to + # the stale K1 — wrong-argument corruption. + if slice_.startswith(self.arg_key_start): + if not self._consume_arg_key(slice_): + break # incomplete + continue # stay in READING_VALUE: orphan replaced + if slice_.startswith(self.arg_value_start): + end = slice_.find(self.arg_value_end) + if end == -1: + break # incomplete — no partial emission + raw = ( + slice_[len(self.arg_value_start) : end] + .removeprefix("\n") + .removesuffix("\n") + ) + # READING_VALUE is reachable only via READING_KEY + # consuming an ..., so + # current_pending_key is set by construction. + schema = self._get_param_schema(self.current_func_name, tools) + converted = self._convert_param_value( + raw, schema, self.current_pending_key + ) + kv = ( + f"{json.dumps(self.current_pending_key)}: " + f"{json.dumps(converted, ensure_ascii=False)}" + ) + fragment = "{" + kv if not self.json_started else ", " + kv + self.json_started = True + calls.append( + ToolCallItem( + tool_index=self.current_tool_id, + parameters=fragment, + ) + ) + self.streamed_args_for_tool[self.current_tool_id] += fragment + self.current_pending_key = None + self.parsed_pos += end + len(self.arg_value_end) + self._state = _ParseState.READING_KEY + continue + if slice_.startswith("<"): + if self._is_partial_tag(slice_): + break + self.parsed_pos += 1 + continue + next_lt = slice_.find("<") + self.parsed_pos += len(slice_) if next_lt == -1 else next_lt + continue + + if state is _ParseState.DRAINING: + end_idx = slice_.find(self.tool_call_end_token) + if end_idx != -1: + self.parsed_pos += end_idx + len(self.tool_call_end_token) + self._state = _ParseState.OUTSIDE + continue + # Hold back trailing bytes that could be a prefix of + # ; the next chunk extends the tail. + holdback = self._ends_with_partial_token( + slice_, self.tool_call_end_token + ) + self.parsed_pos += len(slice_) - holdback + break + + if self.parsed_pos > 0: + self._buffer = self._buffer[self.parsed_pos :] + self.parsed_pos = 0 + + return StreamingParseResult( + calls=calls, + normal_text="".join(normal_text_chunks), + ) + + # ---------- Constrained generation ---------- + + def supports_structural_tag(self) -> bool: + return False + + def structure_info(self) -> _GetInfoFunc: + return lambda name: StructureInfo( + begin=f"{name}\n", + end="", + trigger="", + ) diff --git a/python/sglang/srt/models/laguna.py b/python/sglang/srt/models/laguna.py new file mode 100644 index 000000000..93cd9b3a3 --- /dev/null +++ b/python/sglang/srt/models/laguna.py @@ -0,0 +1,787 @@ +# 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 +"""Inference-only Laguna (poolside/Laguna-XS.2) model.""" + +from __future__ import annotations + +import logging +from collections.abc import Iterable +from typing import Any, Dict, 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.distributed import ( + get_pp_group, + get_tensor_model_parallel_world_size, + tensor_model_parallel_all_reduce, +) +from sglang.srt.layers.activation import SiluAndMul +from sglang.srt.layers.communicator import ( + LayerCommunicator, + LayerScatterModes, +) +from sglang.srt.layers.dp_attention import ( + get_attention_tp_rank, + get_attention_tp_size, + is_dp_attention_enabled, +) +from sglang.srt.layers.layernorm import RMSNorm +from sglang.srt.layers.linear import ( + ColumnParallelLinear, + MergedColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear, +) +from sglang.srt.layers.logits_processor import LogitsProcessor +from sglang.srt.layers.moe import should_skip_post_experts_all_reduce +from sglang.srt.layers.moe.ep_moe.layer import get_moe_impl_class +from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE +from sglang.srt.layers.moe.topk import TopK +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.layers.radix_attention import RadixAttention +from sglang.srt.layers.rotary_embedding import get_rope +from sglang.srt.layers.utils import PPMissingLayer, get_layer_id +from sglang.srt.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors +from sglang.srt.model_loader.weight_utils import default_weight_loader +from sglang.srt.models.utils import apply_qk_norm +from sglang.srt.server_args import get_global_server_args +from sglang.srt.utils import LazyValue, add_prefix, make_layers + +logger = logging.getLogger(__name__) + + +class LagunaMLP(nn.Module): + def __init__( + self, + hidden_size: int, + intermediate_size: int, + hidden_act: str, + quant_config: Optional[QuantizationConfig] = None, + reduce_results: bool = True, + prefix: str = "", + ) -> None: + super().__init__() + if hidden_act != "silu": + raise ValueError( + f"Unsupported activation: {hidden_act}. Only silu is supported." + ) + self.gate_up_proj = MergedColumnParallelLinear( + hidden_size, + [intermediate_size] * 2, + bias=False, + quant_config=quant_config, + prefix=add_prefix("gate_up_proj", prefix), + ) + self.down_proj = RowParallelLinear( + intermediate_size, + hidden_size, + bias=False, + quant_config=quant_config, + reduce_results=reduce_results, + prefix=add_prefix("down_proj", prefix), + ) + self.act_fn = SiluAndMul() + + def forward( + self, + x: torch.Tensor, + forward_batch: Optional[ForwardBatch] = None, + should_allreduce_fusion: bool = False, + use_reduce_scatter: bool = False, + ) -> torch.Tensor: + gate_up, _ = self.gate_up_proj(x) + x = self.act_fn(gate_up) + # Skip the in-block reduce when LayerCommunicator will fuse it or when + # the next layer expects reduce-scatter — otherwise we'd double-reduce. + x, _ = self.down_proj( + x, + skip_all_reduce=should_allreduce_fusion or use_reduce_scatter, + ) + return x + + +class LagunaMoEGate(nn.Module): + def __init__( + self, + config: LagunaConfig, + prefix: str = "", + ): + super().__init__() + self.weight = nn.Parameter( + torch.empty(config.num_experts, config.hidden_size, dtype=torch.float32) + ) + # Released checkpoint stores this under `mlp.experts.e_score_correction_bias` + # (load_weights remaps it) but every value is 0.0; zero-init keeps us + # correct if a future checkpoint omits the tensor entirely. + self.e_score_correction_bias = nn.Parameter( + torch.zeros(config.num_experts, dtype=torch.float32), + requires_grad=False, + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return F.linear(hidden_states.to(torch.float32), self.weight, None) + + +class LagunaMoE(nn.Module): + def __init__( + self, + config: LagunaConfig, + layer_id: int, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + self.tp_size = get_tensor_model_parallel_world_size() + self.routed_scaling_factor = config.moe_routed_scaling_factor + self.router_logit_softcapping = getattr( + config, "moe_router_logit_softcapping", 0.0 + ) + + if self.tp_size > config.num_experts: + raise ValueError( + f"TP size {self.tp_size} > num_experts {config.num_experts}." + ) + + self.gate = LagunaMoEGate(config, prefix=add_prefix("gate", prefix)) + + self.experts = get_moe_impl_class(quant_config)( + num_experts=config.num_experts + + get_global_server_args().ep_num_redundant_experts, + top_k=config.num_experts_per_tok, + layer_id=layer_id, + hidden_size=config.hidden_size, + intermediate_size=config.moe_intermediate_size, + quant_config=quant_config, + reduce_results=False, + apply_router_weight_on_input=bool(config.moe_apply_router_weight_on_input), + prefix=add_prefix("experts", prefix), + ) + + self.topk = TopK( + top_k=config.num_experts_per_tok, + layer_id=layer_id, + renormalize=True, + use_grouped_topk=False, + scoring_func="sigmoid", + correction_bias=self.gate.e_score_correction_bias, + ) + + # HF safetensors key is singular `shared_expert.…`; mirror so the + # default loader picks it up without remapping. + self.shared_expert = LagunaMLP( + hidden_size=config.hidden_size, + intermediate_size=config.shared_expert_intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + reduce_results=False, + prefix=add_prefix("shared_expert", prefix), + ) + + def get_moe_weights(self): + return [x.data for x in self.experts.parameters()] + + def forward( + self, + hidden_states: torch.Tensor, + forward_batch: Optional[ForwardBatch] = None, + should_allreduce_fusion: bool = False, + use_reduce_scatter: bool = False, + ) -> torch.Tensor: + if hidden_states.shape[0] == 0: + return hidden_states + + shared_out = self.shared_expert(hidden_states) + + router_logits = self.gate(hidden_states) + if self.router_logit_softcapping > 0.0: + cap = self.router_logit_softcapping + router_logits = torch.tanh(router_logits / cap) * cap + topk_output = self.topk(hidden_states, router_logits) + routed_out = self.experts(hidden_states, topk_output) + + # Non-grouped TopK doesn't honor apply_routed_scaling_factor_on_output, + # 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 + + if self.tp_size > 1 and not should_skip_post_experts_all_reduce( + is_tp_path=True, + use_reduce_scatter=use_reduce_scatter, + should_allreduce_fusion=should_allreduce_fusion, + ): + final = tensor_model_parallel_all_reduce(final) + return final + + +class LagunaAttention(nn.Module): + def __init__( + self, + hidden_size: int, + num_heads: int, + num_kv_heads: int, + head_dim: int, + layer_id: int, + rms_norm_eps: float, + rope_theta: float, + rope_scaling: Optional[Dict[str, Any]], + partial_rotary_factor: float, + max_position_embeddings: int, + attention_bias: bool, + sliding_window_size: int, + layer_type: str, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.hidden_size = hidden_size + self.head_dim = head_dim + self.layer_id = layer_id + + attn_tp_rank = get_attention_tp_rank() + attn_tp_size = get_attention_tp_size() + + self.total_num_heads = num_heads + assert self.total_num_heads % attn_tp_size == 0 + self.num_heads = self.total_num_heads // attn_tp_size + self.total_num_kv_heads = num_kv_heads + if self.total_num_kv_heads >= attn_tp_size: + assert self.total_num_kv_heads % attn_tp_size == 0 + else: + assert attn_tp_size % self.total_num_kv_heads == 0 + self.num_kv_heads = max(1, self.total_num_kv_heads // attn_tp_size) + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + self.scaling = self.head_dim**-0.5 + + self.qkv_proj = QKVParallelLinear( + hidden_size, + self.head_dim, + self.total_num_heads, + self.total_num_kv_heads, + bias=attention_bias, + quant_config=quant_config, + tp_rank=attn_tp_rank, + tp_size=attn_tp_size, + prefix=add_prefix("qkv_proj", prefix), + ) + self.o_proj = RowParallelLinear( + self.total_num_heads * self.head_dim, + hidden_size, + bias=attention_bias, + quant_config=quant_config, + tp_rank=attn_tp_rank, + tp_size=attn_tp_size, + reduce_results=False, + 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), + ) + + self.q_norm = RMSNorm(self.head_dim, eps=rms_norm_eps) + self.k_norm = RMSNorm(self.head_dim, eps=rms_norm_eps) + + self.rotary_emb = get_rope( + self.head_dim, + rotary_dim=self.head_dim, + max_position=max_position_embeddings, + base=int(rope_theta), + rope_scaling=rope_scaling, + partial_rotary_factor=partial_rotary_factor, + ) + + assert layer_type in {"sliding_attention", "full_attention"} + use_sliding = layer_type == "sliding_attention" + self.attn = RadixAttention( + self.num_heads, + self.head_dim, + self.scaling, + num_kv_heads=self.num_kv_heads, + layer_id=layer_id, + quant_config=quant_config, + prefix=add_prefix("attn", prefix), + sliding_window_size=sliding_window_size if use_sliding else -1, + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + if hidden_states.shape[0] == 0: + return hidden_states + + qkv, _ = self.qkv_proj(hidden_states) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + + q, k = apply_qk_norm( + q=q, + k=k, + q_norm=self.q_norm, + k_norm=self.k_norm, + head_dim=self.head_dim, + ) + q, k = self.rotary_emb(positions, q, k) + + 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) + + output, _ = self.o_proj(attn_output) + return output + + +class LagunaDecoderLayer(nn.Module): + def __init__( + self, + config: LagunaConfig, + layer_id: int, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.config = config + self.layer_id = layer_id + self.hidden_size = config.hidden_size + + layer_types = config.layer_types + layer_type = layer_types[layer_id] + is_swa = layer_type == "sliding_attention" + + layer_num_heads = config.num_attention_heads_per_layer[layer_id] + + if is_swa: + rope_theta = config.swa_rope_theta + rope_scaling = config.swa_rope_scaling + partial_rotary_factor = config.swa_partial_rotary_factor + else: + rope_theta = config.rope_theta + rope_scaling = config.full_rope_scaling + partial_rotary_factor = config.partial_rotary_factor + + self.self_attn = LagunaAttention( + hidden_size=self.hidden_size, + num_heads=layer_num_heads, + num_kv_heads=config.num_key_value_heads, + head_dim=config.head_dim, + layer_id=layer_id, + rms_norm_eps=config.rms_norm_eps, + rope_theta=rope_theta, + rope_scaling=rope_scaling, + partial_rotary_factor=partial_rotary_factor, + max_position_embeddings=config.max_position_embeddings, + attention_bias=config.attention_bias, + # SGLang's window is exclusive; HF's `sliding_window` is inclusive. + sliding_window_size=config.sliding_window - 1, + layer_type=layer_type, + quant_config=quant_config, + prefix=add_prefix("self_attn", prefix), + ) + + mlp_types = config.mlp_layer_types + self.is_layer_sparse = mlp_types[layer_id] == "sparse" + is_previous_layer_sparse = layer_id > 0 and mlp_types[layer_id - 1] == "sparse" + is_next_layer_sparse = ( + layer_id + 1 < config.num_hidden_layers + and mlp_types[layer_id + 1] == "sparse" + ) + + if self.is_layer_sparse: + self.mlp = LagunaMoE( + config=config, + layer_id=layer_id, + quant_config=quant_config, + prefix=add_prefix("mlp", prefix), + ) + else: + self.mlp = LagunaMLP( + hidden_size=self.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + reduce_results=True, + prefix=add_prefix("mlp", prefix), + ) + + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + + self.layer_scatter_modes = LayerScatterModes.init_new( + layer_id=layer_id, + num_layers=config.num_hidden_layers, + is_layer_sparse=self.is_layer_sparse, + is_previous_layer_sparse=is_previous_layer_sparse, + is_next_layer_sparse=is_next_layer_sparse, + ) + self.layer_communicator = LayerCommunicator( + layer_scatter_modes=self.layer_scatter_modes, + input_layernorm=self.input_layernorm, + post_attention_layernorm=self.post_attention_layernorm, + allow_reduce_scatter=True, + is_last_layer=(layer_id == config.num_hidden_layers - 1), + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + residual: Optional[torch.Tensor], + ) -> Tuple[torch.Tensor, torch.Tensor]: + hidden_states, residual = self.layer_communicator.prepare_attn( + hidden_states, residual, forward_batch + ) + if hidden_states.shape[0] != 0: + hidden_states = self.self_attn( + positions=positions, + hidden_states=hidden_states, + forward_batch=forward_batch, + ) + hidden_states, residual = self.layer_communicator.prepare_mlp( + hidden_states, residual, forward_batch + ) + + should_allreduce_fusion = ( + self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer( + forward_batch + ) + ) + use_reduce_scatter = self.layer_communicator.should_use_reduce_scatter( + forward_batch + ) + + hidden_states = self.mlp( + hidden_states, + forward_batch=forward_batch, + should_allreduce_fusion=should_allreduce_fusion, + use_reduce_scatter=use_reduce_scatter, + ) + + if should_allreduce_fusion: + hidden_states._sglang_needs_allreduce_fusion = True + else: + hidden_states, residual = self.layer_communicator.postprocess_layer( + hidden_states, residual, forward_batch + ) + return hidden_states, residual + + +class LagunaModel(nn.Module): + def __init__( + self, + config: LagunaConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + decoder_layer_type: type = LagunaDecoderLayer, + ) -> None: + super().__init__() + self.config = config + self.padding_idx = getattr(config, "pad_token_id", None) + self.vocab_size = config.vocab_size + self.pp_group = get_pp_group() + + if self.pp_group.is_first_rank: + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + use_attn_tp_group=is_dp_attention_enabled(), + prefix=add_prefix("embed_tokens", prefix), + ) + else: + self.embed_tokens = PPMissingLayer() + + decoder_layer_type = decoder_layer_type or LagunaDecoderLayer + self.layers, self.start_layer, self.end_layer = make_layers( + config.num_hidden_layers, + lambda idx, prefix: decoder_layer_type( + layer_id=idx, + config=config, + quant_config=quant_config, + prefix=prefix, + ), + pp_rank=self.pp_group.rank_in_group, + pp_size=self.pp_group.world_size, + prefix=add_prefix("layers", prefix), + ) + if self.pp_group.is_last_rank: + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + else: + self.norm = PPMissingLayer(return_tuple=True) + + def get_input_embeddings(self) -> nn.Embedding: + return self.embed_tokens + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor = None, + pp_proxy_tensors: Optional[PPProxyTensors] = None, + ) -> Union[torch.Tensor, PPProxyTensors]: + if self.pp_group.is_first_rank: + if input_embeds is None: + hidden_states = self.embed_tokens(input_ids) + else: + hidden_states = input_embeds + residual = None + else: + assert pp_proxy_tensors is not None + hidden_states = pp_proxy_tensors["hidden_states"] + residual = pp_proxy_tensors["residual"] + + for i in range(self.start_layer, self.end_layer): + layer = self.layers[i] + hidden_states, residual = layer( + positions, hidden_states, forward_batch, residual + ) + + if not self.pp_group.is_last_rank: + return PPProxyTensors( + {"hidden_states": hidden_states, "residual": residual} + ) + + if hidden_states.shape[0] != 0: + if residual is None: + hidden_states = self.norm(hidden_states) + else: + hidden_states, _ = self.norm(hidden_states, residual) + return hidden_states + + +class LagunaForCausalLM(nn.Module): + fall_back_to_pt_during_load = False + packed_modules_mapping = { + "qkv_proj": ["q_proj", "k_proj", "v_proj"], + "gate_up_proj": ["gate_proj", "up_proj"], + } + + def __init__( + self, + config: LagunaConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.pp_group = get_pp_group() + self.config = config + self.model = LagunaModel( + config, quant_config=quant_config, prefix=add_prefix("model", prefix) + ) + if self.pp_group.is_last_rank: + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=add_prefix("lm_head", prefix), + use_attn_tp_group=get_global_server_args().enable_dp_lm_head, + ) + else: + self.lm_head = PPMissingLayer() + self.logits_processor = LogitsProcessor(config) + + # Only walk this rank's local layers — out-of-range entries can be PPMissingLayer. + self._routed_experts_weights_of_layer = LazyValue( + lambda: { + layer_id: self.model.layers[layer_id].mlp.get_moe_weights() + for layer_id in range(self.start_layer, self.end_layer) + if isinstance(self.model.layers[layer_id].mlp, LagunaMoE) + } + ) + + @property + def routed_experts_weights_of_layer(self): + return self._routed_experts_weights_of_layer.value + + @property + def start_layer(self): + return self.model.start_layer + + @property + def end_layer(self): + return self.model.end_layer + + @torch.no_grad() + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor = None, + pp_proxy_tensors: Optional[PPProxyTensors] = None, + ) -> torch.Tensor: + hidden_states = self.model( + input_ids, + positions, + forward_batch, + input_embeds, + pp_proxy_tensors=pp_proxy_tensors, + ) + if self.pp_group.is_last_rank: + return self.logits_processor( + input_ids, hidden_states, self.lm_head, forward_batch + ) + return hidden_states + + def get_input_embeddings(self) -> nn.Embedding: + return self.model.embed_tokens + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): + stacked_params_mapping = [ + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ] + + expert_params_mapping = FusedMoE.make_expert_params_mapping( + ckpt_gate_proj_name="gate_proj", + ckpt_down_proj_name="down_proj", + ckpt_up_proj_name="up_proj", + num_experts=self.config.num_experts, + ) + + params_dict = dict(self.named_parameters()) + + # (layer, expert, shard) tuples that hit the per-expert loader, + # cross-checked against `expected` below to fail on dropped weights. + loaded_expert_shards: set[Tuple[int, int, str]] = set() + moe_layer_ids = [ + i + for i, mt in enumerate(self.config.mlp_layer_types) + if mt == "sparse" and self.start_layer <= i < self.end_layer + ] + + for name, loaded_weight in weights: + layer_id = get_layer_id(name) + if layer_id is not None and ( + layer_id < self.start_layer or layer_id >= self.end_layer + ): + continue + + if "rotary_emb.inv_freq" in name: + continue + + if self.config.tie_word_embeddings and "lm_head.weight" in name: + continue + + # HF stores the router correction bias under the experts namespace; + # our parameter lives on the gate. Remap before dispatch. + if name.endswith("mlp.experts.e_score_correction_bias"): + name = name.replace( + "mlp.experts.e_score_correction_bias", + "mlp.gate.e_score_correction_bias", + ) + + # Stacked dense (QKV / gate_up). The `mlp.experts.` guard stops + # `up_proj` substring from false-matching `experts.{i}.up_proj.weight`. + matched_stacked = False + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + if "mlp.experts." in name: + continue + name_mapped = name.replace(weight_name, param_name) + if name_mapped.endswith(".bias") and name_mapped not in params_dict: + continue + if name_mapped not in params_dict: + continue + param = params_dict[name_mapped] + param.weight_loader(param, loaded_weight, shard_id) + matched_stacked = True + break + if matched_stacked: + continue + + matched_expert = False + for param_name, weight_name, expert_id, shard_id in expert_params_mapping: + if weight_name not in name: + continue + name_mapped = name.replace(weight_name, param_name) + if name_mapped not in params_dict: + continue + param = params_dict[name_mapped] + param.weight_loader( + param, + loaded_weight, + name, + shard_id=shard_id, + expert_id=expert_id, + ) + if layer_id is not None: + loaded_expert_shards.add((layer_id, expert_id, shard_id)) + matched_expert = True + break + if matched_expert: + continue + + if name.endswith(".bias") and name not in params_dict: + continue + if name not in params_dict: + logger.warning("Parameter %s not found in params_dict", name) + continue + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + + # If any routed-expert tensor was silently dropped (e.g. a future + # checkpoint renaming `gate_proj`, or a ckpt-vs-mapping shape mismatch), + # fail loud here instead of generating garbage. + expected = { + (layer_id, expert_id, shard_id) + for layer_id in moe_layer_ids + for expert_id in range(self.config.num_experts) + for shard_id in ("w1", "w2", "w3") + } + missing = expected - loaded_expert_shards + if missing: + sample = sorted(missing)[:5] + raise RuntimeError( + f"{len(missing)} routed-expert tensors were not loaded " + f"(sample: {sample}). Expected {len(expected)} (layers={moe_layer_ids}, " + f"num_experts={self.config.num_experts}, shards=3)." + ) + + def get_embed_and_head(self): + return self.model.embed_tokens.weight, self.lm_head.weight + + def set_embed_and_head(self, embed, head): + del self.model.embed_tokens.weight + del self.lm_head.weight + self.model.embed_tokens.weight = embed + self.lm_head.weight = head + torch.cuda.empty_cache() + torch.cuda.synchronize() + + +EntryClass = LagunaForCausalLM diff --git a/python/sglang/srt/parser/reasoning_parser.py b/python/sglang/srt/parser/reasoning_parser.py index b91427eef..6f3cf58f7 100644 --- a/python/sglang/srt/parser/reasoning_parser.py +++ b/python/sglang/srt/parser/reasoning_parser.py @@ -587,6 +587,15 @@ class _MimoDetector(Qwen3Detector): self.reasoning_default = "explicit_enable_thinking" +class _PoolsideV1Detector(Qwen3Detector): + """Poolside v1 (Laguna-XS.2) reuses Qwen3 tokens but the HF chat template + defaults `enable_thinking=False`; reasoning is opt-in via `enable_thinking=True`.""" + + def __init__(self, **kwargs): + super().__init__(**kwargs) + self.reasoning_default = "explicit_enable_thinking" + + class ReasoningParser: """ Parser that handles both streaming and non-streaming scenarios for extracting @@ -608,6 +617,7 @@ class ReasoningParser: "kimi": KimiDetector, "kimi_k2": KimiK2Detector, "mimo": _MimoDetector, + "poolside_v1": _PoolsideV1Detector, "qwen3": Qwen3Detector, "qwen3-thinking": Qwen3Detector, "minimax": Qwen3Detector, diff --git a/python/sglang/srt/utils/hf_transformers/common.py b/python/sglang/srt/utils/hf_transformers/common.py index 9067e6fa6..cd8729798 100644 --- a/python/sglang/srt/utils/hf_transformers/common.py +++ b/python/sglang/srt/utils/hf_transformers/common.py @@ -37,6 +37,7 @@ from sglang.srt.configs import ( KimiK25Config, KimiLinearConfig, KimiVLConfig, + LagunaConfig, LongcatFlashConfig, MultiModalityConfig, NemotronH_Nano_Omni_Reasoning_V3_Config, @@ -79,6 +80,7 @@ _CONFIG_REGISTRY: Dict[str, Type[PretrainedConfig]] = { MultiModalityConfig, KimiVLConfig, InternVLChatConfig, + LagunaConfig, Step3VLConfig, LongcatFlashConfig, Olmo3Config, diff --git a/test/registered/unit/entrypoints/openai/test_serving_chat.py b/test/registered/unit/entrypoints/openai/test_serving_chat.py index 69f3ff4d4..5e4b9f81e 100644 --- a/test/registered/unit/entrypoints/openai/test_serving_chat.py +++ b/test/registered/unit/entrypoints/openai/test_serving_chat.py @@ -1394,6 +1394,52 @@ class ServingChatTestCase(unittest.TestCase): self.assertIsNone(msg.content) self.assertEqual(msg.reasoning_content, "42") + # --- poolside_v1 (Laguna-XS.2) regression tests --- + + def test_poolside_v1_enable_thinking_dispatch(self): + """Laguna chat template defaults `enable_thinking=false`. Parser must + follow that default — must NOT return True via the generic fallback. + After the reasoning-config refactor, this is driven by + `_PoolsideV1Detector.reasoning_default = "explicit_enable_thinking"`.""" + self._setup_fallback("poolside_v1") + req = ChatCompletionRequest( + model="x", messages=[{"role": "user", "content": "hi"}] + ) + cases = [ + (None, False), # no chat_template_kwargs → non-thinking (default) + ({}, False), # empty kwargs → non-thinking + ({"enable_thinking": False}, False), # explicit off + ({"enable_thinking": True}, True), # explicit on + ] + for kwargs, expected in cases: + with self.subTest(kwargs=kwargs): + req.chat_template_kwargs = kwargs + self.assertEqual(self.chat._get_reasoning_from_request(req), expected) + + def test_poolside_v1_does_not_double_prepend_think(self): + """When `enable_thinking=True` for poolside_v1, the HF chat template + already emits `` via add_generation_prompt — server must NOT + append a second ``. After the refactor this is guarded by + `_PoolsideV1Detector.thinks_internally = True` (inherited from Qwen3Detector). + """ + self._setup_fallback("poolside_v1") + req = ChatCompletionRequest( + model="x", + messages=[{"role": "user", "content": "hi"}], + chat_template_kwargs={"enable_thinking": True}, + ) + with patch( + "sglang.srt.entrypoints.openai.serving_chat.generate_chat_conv" + ) as conv_mock: + conv_ins = Mock() + conv_ins.get_prompt.return_value = "BASE_PROMPT" + conv_ins.image_data = conv_ins.audio_data = conv_ins.video_data = None + conv_ins.modalities = [] + conv_ins.stop_str = [] + conv_mock.return_value = conv_ins + result = self.chat._apply_conversation_template(req, is_multimodal=False) + self.assertEqual(result.prompt, "BASE_PROMPT") + class TestProcessToolCallsWithRequiredToolChoice(unittest.TestCase): """Test _process_tool_calls with tool_choice='required' uses model-specific parser.""" diff --git a/test/registered/unit/function_call/test_poolside_v1_detector.py b/test/registered/unit/function_call/test_poolside_v1_detector.py new file mode 100644 index 000000000..2370cf479 --- /dev/null +++ b/test/registered/unit/function_call/test_poolside_v1_detector.py @@ -0,0 +1,483 @@ +"""Unit tests for PoolsideV1Detector — no server, no model loading.""" + +import json + +from sglang.srt.entrypoints.openai.protocol import Function, Tool +from sglang.srt.function_call.function_call_parser import FunctionCallParser +from sglang.srt.function_call.poolside_v1_detector import PoolsideV1Detector +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(1.0, "stage-a-test-cpu") + + +class TestPoolsideV1Detector(CustomTestCase): + def setUp(self): + self.tools = [ + Tool( + type="function", + function=Function( + name="get_weather", + description="Get weather", + parameters={ + "type": "object", + "properties": { + "location": {"type": "string"}, + "count": {"type": "integer"}, + "options": {"type": "object"}, + }, + "required": ["location"], + }, + ), + ), + Tool( + type="function", + function=Function( + name="search", + description="Search", + parameters={ + "type": "object", + "properties": {"query": {"type": "string"}}, + "required": ["query"], + }, + ), + ), + Tool( + type="function", + function=Function( + name="now", + description="Current time", + parameters={"type": "object", "properties": {}}, + ), + ), + ] + self.detector = PoolsideV1Detector() + + # ==================== has_tool_call ==================== + + def test_has_tool_call_true(self): + text = ( + "get_weather\nlocation\n" + "SF\n" + ) + self.assertTrue(self.detector.has_tool_call(text)) + + def test_has_tool_call_false(self): + self.assertFalse(self.detector.has_tool_call("just a sentence.")) + + # ==================== detect_and_parse ==================== + + def test_single_tool_call_string_arg(self): + text = ( + "get_weather\nlocation\n" + "San Francisco\n" + ) + result = self.detector.detect_and_parse(text, self.tools) + self.assertEqual(len(result.calls), 1) + self.assertEqual(result.calls[0].name, "get_weather") + args = json.loads(result.calls[0].parameters) + self.assertEqual(args, {"location": "San Francisco"}) + + def test_single_tool_call_mixed_types(self): + text = ( + "get_weather\n" + "location\nLondon\n" + "count\n3\n" + 'options\n{"verbose": true}\n' + "" + ) + result = self.detector.detect_and_parse(text, self.tools) + self.assertEqual(len(result.calls), 1) + args = json.loads(result.calls[0].parameters) + self.assertEqual(args["location"], "London") + self.assertEqual(args["count"], 3) + self.assertEqual(args["options"], {"verbose": True}) + + def test_multiple_tool_calls(self): + text = ( + "get_weather\nlocation\n" + "NYC\n\n" + "search\nquery\n" + "pizza\n" + ) + result = self.detector.detect_and_parse(text, self.tools) + self.assertEqual(len(result.calls), 2) + self.assertEqual(result.calls[0].name, "get_weather") + self.assertEqual(result.calls[1].name, "search") + self.assertEqual(json.loads(result.calls[1].parameters), {"query": "pizza"}) + + def test_leading_text_extracted_as_normal(self): + text = ( + "Sure, checking now. " + "search\nquery\n" + "tacos\n" + ) + result = self.detector.detect_and_parse(text, self.tools) + self.assertEqual(result.normal_text, "Sure, checking now. ") + self.assertEqual(len(result.calls), 1) + + def test_unknown_tool_dropped(self): + text = ( + "nonexistent_fn\nx\n" + "1\n" + ) + result = self.detector.detect_and_parse(text, self.tools) + self.assertEqual(len(result.calls), 0) + + def test_malformed_value_falls_back_to_string(self): + text = ( + "get_weather\noptions\n" + "not_json\n" + ) + result = self.detector.detect_and_parse(text, self.tools) + self.assertEqual(len(result.calls), 1) + args = json.loads(result.calls[0].parameters) + self.assertEqual(args["options"], "not_json") + + def test_zero_arg_call(self): + text = "now\n" + result = self.detector.detect_and_parse(text, self.tools) + self.assertEqual(len(result.calls), 1) + self.assertEqual(result.calls[0].name, "now") + self.assertEqual(json.loads(result.calls[0].parameters), {}) + + def test_zero_arg_call_no_newline(self): + """`now` (no `\\n` between name and close tag).""" + text = "now" + result = self.detector.detect_and_parse(text, self.tools) + self.assertEqual(len(result.calls), 1) + self.assertEqual(result.calls[0].name, "now") + self.assertEqual(json.loads(result.calls[0].parameters), {}) + + def test_truncated_pre_value_emits_no_calls(self): + """Regression: max-tokens cutoff mid-`` must drop the + in-flight call, matching the old closing-tag-anchored regex behavior. + Without the truncated-call filter in detect_and_parse, streaming-as- + primitive surfaced a tool call with parameters="{}" on this input.""" + text = ( + "get_weather\nlocation\n" "San Fr" + ) + result = self.detector.detect_and_parse(text, self.tools) + self.assertEqual( + len(result.calls), 0, "truncated mid-arg_value must yield 0 calls" + ) + + def test_truncated_post_value_emits_no_calls(self): + """Regression: cutoff after `` but before `` + used to surface a tool call with non-JSON parameters + ('{"location": "SF"' with no closing brace). The filter must drop it.""" + text = ( + "get_weather\nlocation\n" + "SF\n" + ) + result = self.detector.detect_and_parse(text, self.tools) + self.assertEqual( + len(result.calls), + 0, + "truncated after arg_value but before must yield 0 calls", + ) + + def test_set_literal_falls_back_to_raw_string(self): + """Regression: ast.literal_eval('{1,2,3}') returns a set, which + json.dumps cannot serialize. Without the round-trip guard in + _convert_param_value, the parse_streaming_increment loop would + TypeError downstream. The guard rejects sets and falls back to the + raw string (which then matches the underlying schema-string-typed + treatment).""" + tools_with_obj = [ + Tool( + type="function", + function=Function( + name="get_weather", + description="Get weather", + parameters={ + "type": "object", + "properties": {"options": {"type": "object"}}, + }, + ), + ) + ] + detector = PoolsideV1Detector() + text = ( + "get_weather\noptions\n" + "{1, 2, 3}\n" + ) + result = detector.detect_and_parse(text, tools_with_obj) + self.assertEqual(len(result.calls), 1) + args = json.loads(result.calls[0].parameters) + # set literal couldn't round-trip, so it's preserved as the raw + # string (the only sane fallback). + self.assertEqual(args["options"], "{1, 2, 3}") + + def test_truncated_after_complete_call_keeps_complete(self): + """A complete tool_call followed by a truncated second one must keep + the complete one and drop only the truncated tail — matching the old + regex behavior on the same input.""" + text = ( + "get_weather\nlocation\n" + "NYC\n\n" + "search\nq" + ) + result = self.detector.detect_and_parse(text, self.tools) + self.assertEqual(len(result.calls), 1) + self.assertEqual(result.calls[0].name, "get_weather") + self.assertEqual(json.loads(result.calls[0].parameters), {"location": "NYC"}) + + def test_arg_key_without_value_emits_empty_call(self): + """Non-streaming: malformed `K` (no + ``) yields a tool call with empty params — the orphan + `` is dropped because the regex looks for key/value pairs. + Locks in the contract the streaming FSM must match.""" + text = "get_weather\nlocation" + result = self.detector.detect_and_parse(text, self.tools) + self.assertEqual(len(result.calls), 1) + self.assertEqual(result.calls[0].name, "get_weather") + self.assertEqual(json.loads(result.calls[0].parameters), {}) + + def test_streaming_arg_key_without_value_closes_call(self): + """Regression: malformed `K` (no + ``) used to leave the streaming FSM stuck in READING_VALUE + — the bare-`<` discard ate `` byte-by-byte instead of + recognizing it as a close. Worse: a *subsequent* tool call's + `` would mis-attribute its content to the orphan + `current_pending_key`, silently swallowing the second call's name. + Both calls must be emitted with the orphan key dropped.""" + detector = PoolsideV1Detector() + wire = ( + "get_weather\nlocation" + "search\nquery\n" + "tacos\n" + ) + all_calls = [] + for chunk in [wire[i : i + 8] for i in range(0, len(wire), 8)]: + r = detector.parse_streaming_increment(chunk, self.tools) + all_calls.extend(r.calls) + names = [c.name for c in all_calls if c.name] + self.assertEqual( + names, + ["get_weather", "search"], + "second call must not be swallowed when first is malformed", + ) + per_tool: dict = {} + for c in all_calls: + if c.parameters: + per_tool.setdefault(c.tool_index, "") + per_tool[c.tool_index] += c.parameters + # Orphan `location` key dropped — first call has empty params. + self.assertEqual(json.loads(per_tool[0]), {}) + # Second call's value must NOT leak into first call's stale key. + self.assertEqual(json.loads(per_tool[1]), {"query": "tacos"}) + + def test_orphan_key_followed_by_new_key_uses_new_key(self): + """Non-streaming: malformed `K1K2 + V` (model emitted a key, then re-emitted a new + key without a value for the first) yields `{K2: V}` — the orphan K1 + is dropped. Without the `[^<]` constraint in arg_pair_regex, the + non-greedy `.*?` backtracks across the `` boundary and + produces a junk key spanning both tags.""" + text = ( + "get_weather\n" + "location" + "count3" + "\n" + ) + result = self.detector.detect_and_parse(text, self.tools) + self.assertEqual(len(result.calls), 1) + self.assertEqual(result.calls[0].name, "get_weather") + args = json.loads(result.calls[0].parameters) + self.assertEqual( + args, + {"count": 3}, + f"orphan key 'location' should be dropped, count=3 should win, got {args}", + ) + + def test_streaming_orphan_key_followed_by_new_key_uses_new_key(self): + """Regression: streaming on `K1K2 + V` used to mis-attribute V to K1 — the bare-`<` + discard ate the second `` as garbage and the value bound to + the stale `current_pending_key`. With the orphan-key-replace branch + in READING_VALUE, the new key wins and streaming matches the + non-streaming regex path.""" + detector = PoolsideV1Detector() + wire = ( + "get_weather\n" + "location" + "count3" + "\n" + ) + all_calls = [] + for chunk in [wire[i : i + 8] for i in range(0, len(wire), 8)]: + r = detector.parse_streaming_increment(chunk, self.tools) + all_calls.extend(r.calls) + names = [c.name for c in all_calls if c.name] + self.assertEqual(names, ["get_weather"]) + params = "".join(c.parameters for c in all_calls if c.parameters) + self.assertEqual( + json.loads(params), + {"count": 3}, + "orphan key 'location' should be dropped; count=3 must win", + ) + + def test_streaming_malformed_no_name_does_not_hang(self): + """Regression: malformed `...` (no name, no \\n) + used to spin in branch 2 with consume=0. Must drain to .""" + detector = PoolsideV1Detector() + wire = "kv" + result = detector.parse_streaming_increment(wire, self.tools) + self.assertEqual(len(result.calls), 0) + + def test_streaming_arg_tags_without_tool_call_wrapper(self): + """Regression: stray `......` + with no preceding `` used to crash with IndexError on + `streamed_args_for_tool[-1]` (masked by the old broad except). The + FSM's READING_VALUE state is unreachable from OUTSIDE, so this returns + 0 calls without raising — and now that the broad except is gone, any + regression here would propagate as a real test failure.""" + detector = PoolsideV1Detector() + wire = "kv" + result = detector.parse_streaming_increment(wire, self.tools) + self.assertEqual(len(result.calls), 0) + + # ==================== structure_info ==================== + + def test_structure_info(self): + info_func = self.detector.structure_info() + info = info_func("get_weather") + self.assertEqual(info.trigger, "") + self.assertIn("get_weather", info.begin) + self.assertIn("", info.end) + + # ==================== Streaming ==================== + + def test_streaming_single_call_chunked(self): + detector = PoolsideV1Detector() + chunks = [ + "get_weather\n", + "location\nSan Fr", + "ancisco\n", + ] + names, params = self._collect(detector, chunks) + self.assertEqual(names, ["get_weather"]) + self.assertEqual(json.loads(params), {"location": "San Francisco"}) + + def test_streaming_char_by_char_robustness(self): + """Per-arg streaming under one-byte chunks. Values are emitted as a + single `"key": value` fragment when `` arrives; this test + proves the FSM doesn't leak trailing `<` / `get_weather\nlocation\n" + "hello world\n" + ) + chunks = list(wire) + names, params = self._collect(detector, chunks) + self.assertEqual(names, ["get_weather"]) + decoded = json.loads(params) + self.assertEqual(decoded, {"location": "hello world"}) + # And the emitted parameter delta itself contains no stray tag bytes. + self.assertNotIn("<", params) + self.assertNotIn(">", params) + + def test_streaming_index_is_sequential_not_tools_slot(self): + """Regression: streaming emissions must use a per-response sequential + index. If we emit the name with `tools_indices[name]` and the params + with `current_tool_id`, OpenAI clients group chunks by `index` and + split a `search`-only call (slot 1) into two broken calls.""" + detector = PoolsideV1Detector() + # `search` is at tools-list slot 1, NOT 0. + wire = ( + "search\nquery\n" + "tacos\n" + ) + all_calls = [] + for c in list(wire): + r = detector.parse_streaming_increment(c, self.tools) + all_calls.extend(r.calls) + # All chunks for this call must share the same index. + indices = {c.tool_index for c in all_calls} + self.assertEqual( + indices, + {0}, + f"streaming emitted mixed indices {indices}; OpenAI clients would " + "split this into multiple broken calls", + ) + names = [c.name for c in all_calls if c.name] + params = "".join(c.parameters for c in all_calls if c.parameters) + self.assertEqual(names, ["search"]) + self.assertEqual(json.loads(params), {"query": "tacos"}) + + def test_streaming_multiple_calls(self): + detector = PoolsideV1Detector() + wire = ( + "get_weather\nlocation\n" + "NYC\n\n" + "search\nquery\n" + "pizza\n" + ) + all_calls = [] + for chunk in [wire[i : i + 16] for i in range(0, len(wire), 16)]: + r = detector.parse_streaming_increment(chunk, self.tools) + all_calls.extend(r.calls) + names = [c.name for c in all_calls if c.name] + self.assertEqual(names, ["get_weather", "search"]) + # Each tool's argument deltas concatenate to a complete JSON object + per_tool: dict = {} + for c in all_calls: + if c.parameters: + per_tool.setdefault(c.tool_index, "") + per_tool[c.tool_index] += c.parameters + self.assertEqual(json.loads(per_tool[0]), {"location": "NYC"}) + self.assertEqual(json.loads(per_tool[1]), {"query": "pizza"}) + + def test_streaming_zero_arg_call(self): + detector = PoolsideV1Detector() + wire = "now\n" + names, params = self._collect(detector, list(wire)) + self.assertEqual(names, ["now"]) + # Either a single "{}" emission or a sequence whose join parses to {} + self.assertEqual(json.loads(params or "{}"), {}) + + def test_streaming_text_before_tool_call(self): + detector = PoolsideV1Detector() + chunks = [ + "Let me check. ", + "search\nquery\n", + "foo\n", + ] + all_calls = [] + normal = "" + for chunk in chunks: + r = detector.parse_streaming_increment(chunk, self.tools) + all_calls.extend(r.calls) + normal += r.normal_text + self.assertEqual(normal, "Let me check. ") + names = [c.name for c in all_calls if c.name] + self.assertEqual(names, ["search"]) + + # ==================== Registry ==================== + + def test_registered_in_function_call_parser(self): + self.assertIn("poolside_v1", FunctionCallParser.ToolCallParserEnum) + self.assertIs( + FunctionCallParser.ToolCallParserEnum["poolside_v1"], PoolsideV1Detector + ) + + # ==================== Helpers ==================== + + def _collect(self, detector, chunks): + all_calls = [] + for chunk in chunks: + r = detector.parse_streaming_increment(chunk, self.tools) + all_calls.extend(r.calls) + names = [c.name for c in all_calls if c.name] + params = "".join(c.parameters for c in all_calls if c.parameters) + return names, params + + +if __name__ == "__main__": + import unittest + + unittest.main() diff --git a/test/registered/unit/parser/test_reasoning_parser.py b/test/registered/unit/parser/test_reasoning_parser.py index 958168a7d..85fff23aa 100644 --- a/test/registered/unit/parser/test_reasoning_parser.py +++ b/test/registered/unit/parser/test_reasoning_parser.py @@ -1499,5 +1499,21 @@ class TestGptOssDetectorToolCall(CustomTestCase): self.assertIn("done", all_normal) +class TestPoolsideV1Registered(CustomTestCase): + """poolside_v1 (Laguna-XS.2) reuses the Qwen3 `...` envelope. + Request dispatch differs (Mimo-style explicit `enable_thinking=True`, + asserted in test_serving_chat.py), driven by + `reasoning_default = "explicit_enable_thinking"` on the detector.""" + + def test_registered_to_qwen3_subclass(self): + cls = ReasoningParser.DetectorMap["poolside_v1"] + self.assertTrue(issubclass(cls, Qwen3Detector)) + + def test_explicit_enable_thinking_default(self): + rp = ReasoningParser("poolside_v1", stream_reasoning=True) + self.assertEqual(rp.detector.reasoning_default, "explicit_enable_thinking") + self.assertTrue(rp.detector.thinks_internally) + + if __name__ == "__main__": unittest.main()