From ffd9d3943a0e37e2ce413fe5b613dc7403943364 Mon Sep 17 00:00:00 2001 From: sglang-bot <232288953+sglang-bot@users.noreply.github.com> Date: Mon, 10 Aug 2026 03:01:01 -0700 Subject: [PATCH 01/18] Muse Glimmer native model support Co-Authored-By: Brayden Zhong Co-Authored-By: Jimmy Shong <69131491+Jiminator@users.noreply.github.com> Co-Authored-By: hnyls2002 Co-Authored-By: Alex Nails --- python/sglang/benchmark/utils.py | 4 +- python/sglang/srt/arg_groups/overrides.py | 8 + python/sglang/srt/configs/__init__.py | 6 + python/sglang/srt/configs/model_config.py | 24 +- python/sglang/srt/configs/muse_glimmer.py | 256 ++++ .../srt/configs/muse_glimmer_processing.py | 302 +++++ python/sglang/srt/entrypoints/http_server.py | 1 + .../sglang/srt/entrypoints/openai/protocol.py | 1 + .../srt/entrypoints/openai/serving_chat.py | 41 +- .../entrypoints/openai/serving_responses.py | 29 +- python/sglang/srt/environ.py | 2 + .../srt/function_call/base_format_detector.py | 8 + .../srt/function_call/function_call_parser.py | 13 + .../function_call/muse_glimmer_detector.py | 279 +++++ .../srt/hardware_backend/mlx/model_runner.py | 29 +- .../mlx/models/muse_glimmer_mlx.py | 740 ++++++++++++ .../hardware_backend/mlx/remote_code_gate.py | 112 ++ .../srt/hardware_backend/mlx/tp_worker.py | 1 + python/sglang/srt/managers/scheduler.py | 6 +- .../sglang/srt/managers/tokenizer_manager.py | 7 +- python/sglang/srt/mem_cache/kv_cache_dtype.py | 5 + .../sglang/srt/model_executor/model_runner.py | 1 + .../spec_aux_hidden_state.py | 10 + .../srt/model_executor/runner/base_runner.py | 4 +- .../sglang/srt/model_loader/gguf_name_maps.py | 69 ++ python/sglang/srt/model_loader/loader.py | 6 + python/sglang/srt/models/dflash.py | 31 +- python/sglang/srt/models/muse_glimmer.py | 1027 +++++++++++++++++ .../srt/multimodal/processors/muse_glimmer.py | 61 + python/sglang/srt/parser/reasoning_parser.py | 211 ++++ python/sglang/srt/server_args.py | 87 +- python/sglang/srt/speculative/dflash_utils.py | 8 +- .../srt/speculative/dflash_worker_v2.py | 74 +- .../srt/speculative/draft_worker_common.py | 4 + .../srt/utils/hf_transformers/__init__.py | 2 + .../srt/utils/hf_transformers/common.py | 75 ++ .../srt/utils/hf_transformers/config.py | 20 +- .../srt/utils/hf_transformers/gguf_native.py | 258 +++++ .../srt/utils/hf_transformers/tokenizer.py | 16 +- .../test_muse_glimmer_mlx_correctness.py | 490 ++++++++ ...est_muse_glimmer_dflash_assistant_gsm8k.py | 94 ++ .../openai/test_serving_responses.py | 72 ++ .../openai/test_serving_responses_stream.py | 30 + .../test_muse_glimmer_detector.py | 432 +++++++ .../mlx/test_mlx_remote_code_gate.py | 178 +++ .../mlx/test_muse_glimmer_mlx_model.py | 279 +++++ 46 files changed, 5368 insertions(+), 45 deletions(-) create mode 100644 python/sglang/srt/configs/muse_glimmer.py create mode 100644 python/sglang/srt/configs/muse_glimmer_processing.py create mode 100644 python/sglang/srt/function_call/muse_glimmer_detector.py create mode 100644 python/sglang/srt/hardware_backend/mlx/models/muse_glimmer_mlx.py create mode 100644 python/sglang/srt/hardware_backend/mlx/remote_code_gate.py create mode 100644 python/sglang/srt/model_loader/gguf_name_maps.py create mode 100644 python/sglang/srt/models/muse_glimmer.py create mode 100644 python/sglang/srt/multimodal/processors/muse_glimmer.py create mode 100644 python/sglang/srt/utils/hf_transformers/gguf_native.py create mode 100644 test/registered/mlx/models_e2e/test_muse_glimmer_mlx_correctness.py create mode 100644 test/registered/spec/dflash/test_muse_glimmer_dflash_assistant_gsm8k.py create mode 100644 test/registered/unit/function_call/test_muse_glimmer_detector.py create mode 100644 test/registered/unit/hardware_backend/mlx/test_mlx_remote_code_gate.py create mode 100644 test/registered/unit/hardware_backend/mlx/test_muse_glimmer_mlx_model.py diff --git a/python/sglang/benchmark/utils.py b/python/sglang/benchmark/utils.py index c3cbebf0cc40..cbcf24d92ddc 100644 --- a/python/sglang/benchmark/utils.py +++ b/python/sglang/benchmark/utils.py @@ -48,9 +48,7 @@ def get_tokenizer( pretrained_model_name_or_path is not None and pretrained_model_name_or_path != "" ) - if pretrained_model_name_or_path.endswith( - ".json" - ) or pretrained_model_name_or_path.endswith(".model"): + if pretrained_model_name_or_path.endswith((".json", ".model", ".gguf")): from sglang.srt.utils.hf_transformers_utils import get_tokenizer return get_tokenizer(pretrained_model_name_or_path) diff --git a/python/sglang/srt/arg_groups/overrides.py b/python/sglang/srt/arg_groups/overrides.py index 5e6f5a5791fa..62f9fdb21a4f 100644 --- a/python/sglang/srt/arg_groups/overrides.py +++ b/python/sglang/srt/arg_groups/overrides.py @@ -1930,6 +1930,14 @@ def _deepseek_v4_sm120_moe(view: Any) -> dict: return {} +@_register_for("MuseGlimmerForConditionalGeneration", "MuseGlimmerForCausalLM") +def _muse_glimmer_fp4_gemm_runner_overrides(server_args: Any, hf_config: Any) -> dict: + if is_sm120_supported() and server_args.fp4_gemm_runner_backend == "auto": + logger.info("Use marlin as FP4 GEMM runner backend on SM120 for Muse Glimmer") + return {"fp4_gemm_runner_backend": "marlin"} + return {} + + @register_post_process def _sparse_head_overlap_disable(view: Any) -> dict: diff --git a/python/sglang/srt/configs/__init__.py b/python/sglang/srt/configs/__init__.py index 3b94061c8b86..54b5fa46ce43 100644 --- a/python/sglang/srt/configs/__init__.py +++ b/python/sglang/srt/configs/__init__.py @@ -37,6 +37,10 @@ from sglang.srt.configs.longcat_flash import LongcatFlashConfig from sglang.srt.configs.minicpmv4_6 import MiniCPMV4_6Config, MiniCPMV4_6VisionConfig from sglang.srt.configs.minimax_vl import MiniMaxM3VLConfig +from sglang.srt.configs.muse_glimmer import ( + MuseGlimmerAssistantConfig, + MuseGlimmerConfig, +) from sglang.srt.configs.nano_nemotron_vl import ( NemotronH_Nano_Omni_Reasoning_V3_Config, NemotronH_Nano_VL_V2_Config, @@ -76,6 +80,8 @@ "Step3TextConfig", "Step3VisionEncoderConfig", "Olmo3Config", + "MuseGlimmerConfig", + "MuseGlimmerAssistantConfig", "KimiLinearConfig", "KimiK3Config", "KimiK25Config", diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index 432f74ef0c7c..6c0ac8c53145 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -264,6 +264,7 @@ def __init__( is_multi_layer_eagle: bool = False, encoder_only: bool = False, language_only: bool = False, + language_model_only: bool = False, disable_hybrid_swa_memory: bool = False, model_config_parser: str = "auto", speculative_algorithm: Optional[str] = None, @@ -451,7 +452,8 @@ def __init__( ) # TODO: requires further polishing # Key on the tower, not the attribute: several config classes default - # vision_config to None, which presence alone would read as image-capable. + # vision_config to None, which presence alone would read as image-capable + # (MuseGlimmerConfig's text-only layouts are one such case). self.is_image_understandable_model = ( enable_multimodal and not self.is_lm_only @@ -534,6 +536,7 @@ def __init__( self.hf_config.encoder_only = encoder_only self.hf_config.language_only = language_only + self.hf_config.language_model_only = language_model_only # matryoshka embeddings self.matryoshka_dimensions = getattr( @@ -582,6 +585,7 @@ def from_server_args( override_config_file=override_config_file, is_multi_layer_eagle=server_args.enable_multi_layer_eagle, language_only=server_args.language_only, + language_model_only=server_args.language_model_only, encoder_only=server_args.encoder_only, is_draft_model=is_draft_model, is_draft_quantization_explicit=( @@ -1039,11 +1043,13 @@ def _derive_model_shapes(self): self.num_nextn_predict_layers = getattr( self.hf_text_config, "num_nextn_predict_layers", None ) - self.vocab_size = self.hf_text_config.vocab_size - # GLM-Image is the only model here whose output head predicts vision tokens. - # Use vision_vocab_size for lm_head, LogitsProcessor, and graph-mode logits buffers. - if _hf_arch(self.hf_config) == "GlmImageForConditionalGeneration": - self.vocab_size = self.hf_text_config.vision_vocab_size + # DFlash drafts have no vocab of their own. + if self.is_draft_model and not hasattr(self.hf_text_config, "vocab_size"): + self.vocab_size = None + else: + self.vocab_size = self.hf_text_config.vocab_size + if _hf_arch(self.hf_config) == "GlmImageForConditionalGeneration": + self.vocab_size = self.hf_text_config.vision_vocab_size def _init_mla_scaling(self, rope_scaling: Optional[dict]) -> None: """Base MLA attention scale from the head dims, then the rope mscale.""" @@ -1830,6 +1836,7 @@ def is_generation_model(model_architectures: List[str], is_embedding: bool = Fal "MossVLForConditionalGeneration", "NemotronH_Nano_VL_V2", "NemotronH_Nano_Omni_Reasoning_V3", + "MuseGlimmerForConditionalGeneration", "PixtralForConditionalGeneration", "Qwen2AudioForConditionalGeneration", "Qwen2VLForConditionalGeneration", @@ -1893,6 +1900,7 @@ def is_generation_model(model_architectures: List[str], is_embedding: bool = Fal "InternS2MobiusForConditionalGeneration", "Qwen3_5ForConditionalGeneration", "Qwen3_5MoeForConditionalGeneration", + "MuseGlimmerForConditionalGeneration", ] if external_mm_model_arch := envs.SGLANG_EXTERNAL_MM_MODEL_ARCH.get(): @@ -2036,6 +2044,8 @@ def is_hybrid_swa_model( "Gemma4UnifiedForConditionalGeneration", "LagunaForCausalLM", "MellumForCausalLM", + "MuseGlimmerForCausalLM", + "MuseGlimmerForConditionalGeneration", "InklingForConditionalGeneration", "InklingForConditionalGenerationMTP", "UnlimitedOCRForCausalLM", @@ -2111,6 +2121,8 @@ def get_hybrid_layer_ids( or "Gemma4UnifiedForConditionalGeneration" in model_architectures or "LagunaForCausalLM" in model_architectures or "MellumForCausalLM" in model_architectures + or "MuseGlimmerForCausalLM" in model_architectures + or "MuseGlimmerForConditionalGeneration" in model_architectures ): layer_types = getattr(hf_text_config, "layer_types", []) swa_attention_layer_ids = [ diff --git a/python/sglang/srt/configs/muse_glimmer.py b/python/sglang/srt/configs/muse_glimmer.py new file mode 100644 index 000000000000..ffdddfb76f06 --- /dev/null +++ b/python/sglang/srt/configs/muse_glimmer.py @@ -0,0 +1,256 @@ +# 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 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +import math +from typing import Any, Dict, List, Optional + +from transformers import PretrainedConfig + +from sglang.srt.configs.muse_glimmer_processing import MuseGlimmerProcessor +from sglang.srt.multimodal.customized_mm_processor_utils import ( + register_customized_processor, +) + +_ARCH = "muse-glimmer" + + +class MuseGlimmerAssistantConfig(PretrainedConfig): + + model_type = "muse_glimmer_assistant" + + +class MuseGlimmerVisionConfig(PretrainedConfig): + + model_type = "muse_glimmer_vision" + + def __init__( + self, + hidden_size: int = 1536, + intermediate_size: int = 8960, + num_hidden_layers: int = 50, + num_attention_heads: int = 16, + hidden_act: str = "gelu", + layer_norm_eps: float = 1e-5, + attention_types: Optional[List[str]] = None, + max_position_embeddings: int = 1024, + merge_size: int = 2, + patch_size: int = 14, + patch_temporal: int = 2, + pos_emb_height: int = 32, + pos_emb_width: int = 32, + rope_theta: float = 10000.0, + **kwargs, + ): + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.hidden_act = hidden_act + self.layer_norm_eps = layer_norm_eps + self.attention_types = attention_types + self.max_position_embeddings = max_position_embeddings + self.merge_size = merge_size + self.patch_size = patch_size + self.patch_temporal = patch_temporal + self.pos_emb_height = pos_emb_height + self.pos_emb_width = pos_emb_width + self.rope_theta = rope_theta + super().__init__(**kwargs) + + +@register_customized_processor(MuseGlimmerProcessor) +class MuseGlimmerConfig(PretrainedConfig): + model_type = "muse_glimmer" + sub_configs = {"vision_config": MuseGlimmerVisionConfig} + + def __init__( + self, + vocab_size: int = 202048, + hidden_size: int = 6656, + intermediate_size: int = 19968, + num_hidden_layers: int = 52, + num_attention_heads: int = 32, + num_key_value_heads: int = 2, + head_dim: int = 128, + hidden_act: str = "silu", + max_position_embeddings: int = 16384, + rms_norm_eps: float = 1e-5, + post_norm_eps: float = 1e-8, + rope_theta: float = 500000.0, + sliding_window: int = 2048, + layer_types: Optional[List[str]] = None, + no_rope_layers: Optional[List[int]] = None, + use_qk_norm: bool = True, + use_attn_output_gate: bool = True, + qk_scale_factor: float = 43.7840518911, + rope_is_neox_style: bool = False, + normalize_tok_embeddings: bool = True, + output_multiplier: float = 0.19611613513818404, + output_soft_cap_temp: Optional[float] = 20.0, + tie_word_embeddings: bool = False, + bos_token_id: int = 200000, + eos_token_id: int = 200001, + vision_config: Optional[Dict[str, Any]] = None, + image_token_id: Optional[int] = None, + video_token_id: Optional[int] = None, + out_hidden_size: int = 6144, + projector_hidden_act: str = "gelu", + projector_hidden_size: int = 4096, + **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_attention_heads = num_attention_heads + 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.rms_norm_eps = rms_norm_eps + self.post_norm_eps = post_norm_eps + self.rope_theta = rope_theta + self.sliding_window = sliding_window + self.layer_types = layer_types + self.no_rope_layers = no_rope_layers + self.use_qk_norm = use_qk_norm + self.use_attn_output_gate = use_attn_output_gate + self.qk_scale_factor = qk_scale_factor + self.rope_is_neox_style = rope_is_neox_style + self.normalize_tok_embeddings = normalize_tok_embeddings + self.output_multiplier = output_multiplier + self.output_soft_cap_temp = output_soft_cap_temp + if isinstance(vision_config, dict): + vision_config = MuseGlimmerVisionConfig(**vision_config) + self.vision_config = vision_config + self.image_token_id = image_token_id + self.video_token_id = video_token_id + self.out_hidden_size = out_hidden_size + self.projector_hidden_act = projector_hidden_act + self.projector_hidden_size = projector_hidden_size + super().__init__( + tie_word_embeddings=tie_word_embeddings, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + **kwargs, + ) + + @classmethod + def from_gguf(cls, gguf_path: str) -> "MuseGlimmerConfig": + return cls(**muse_glimmer_config_kwargs_from_gguf(gguf_path)) + + @classmethod + def from_dict(cls, config_dict: Dict[str, Any], **kwargs): + return super().from_dict( + muse_glimmer_config_kwargs_from_hf(config_dict), **kwargs + ) + + +_HF_TEXT_KEYS_TRANSLATED = frozenset( + { + "final_logit_softcapping", + "hidden_activation", + "layer_rope_theta", + "model_type", + "qk_scale_factor", + "rope_parameters", + } +) + +_HF_NESTED_KEYS = ("text_config", "vision_config") + +_HF_VISION_KEYS_TRANSLATED = frozenset({"layer_types", "model_type", "rope_parameters"}) + + +def muse_glimmer_config_kwargs_from_hf(config_dict: Dict[str, Any]) -> Dict[str, Any]: + if "text_config" not in config_dict: + return config_dict + + text = config_dict["text_config"] + kwargs = {k: v for k, v in config_dict.items() if k not in _HF_NESTED_KEYS} + kwargs.update({k: v for k, v in text.items() if k not in _HF_TEXT_KEYS_TRANSLATED}) + kwargs.update( + hidden_act=text["hidden_activation"], + rope_theta=text["rope_parameters"]["rope_theta"], + no_rope_layers=[1 if theta else 0 for theta in text["layer_rope_theta"]], + output_soft_cap_temp=text["final_logit_softcapping"], + qk_scale_factor=text["qk_scale_factor"] * math.sqrt(text["head_dim"]), + rope_is_neox_style=True, + ) + if "vision_config" in config_dict: + kwargs["vision_config"] = muse_glimmer_vision_config_kwargs_from_hf( + config_dict["vision_config"] + ) + return kwargs + + +def muse_glimmer_vision_config_kwargs_from_hf( + vision_config_dict: Dict[str, Any], +) -> Dict[str, Any]: + """Translate the vendor's ``vision_config`` into ``MuseGlimmerVisionConfig`` kwargs.""" + kwargs = { + k: v + for k, v in vision_config_dict.items() + if k not in _HF_VISION_KEYS_TRANSLATED + } + kwargs["rope_theta"] = vision_config_dict["rope_parameters"]["rope_theta"] + kwargs["attention_types"] = vision_config_dict["layer_types"] + return kwargs + + +def muse_glimmer_config_kwargs_from_gguf(gguf_path: str) -> Dict[str, Any]: + from gguf import GGUFReader + + reader = GGUFReader(gguf_path) + meta = {key: field.contents() for key, field in reader.fields.items()} + shapes = {t.name: tuple(int(x) for x in t.shape) for t in reader.tensors} + tensor_names = set(shapes) + + def get(suffix): + return meta[f"{_ARCH}.{suffix}"] + + head_dim = int(get("attention.key_length")) + swa_pattern = [bool(x) for x in get("attention.sliding_window_pattern")] + + return dict( + # token_embd is stored [n_embd, n_vocab] in ggml's reversed order. + vocab_size=shapes["token_embd.weight"][1], + hidden_size=int(get("embedding_length")), + intermediate_size=int(get("feed_forward_length")), + num_hidden_layers=int(get("block_count")), + num_attention_heads=int(get("attention.head_count")), + num_key_value_heads=int(get("attention.head_count_kv")), + head_dim=head_dim, + max_position_embeddings=int(get("context_length")), + rms_norm_eps=float(get("attention.layer_norm_rms_epsilon")), + post_norm_eps=float(get("attention.post_norm_rms_epsilon")), + rope_theta=float(get("rope.freq_base")), + sliding_window=int(get("attention.sliding_window")), + layer_types=[ + "sliding_attention" if s else "full_attention" for s in swa_pattern + ], + no_rope_layers=[1 if s else 0 for s in swa_pattern], + qk_scale_factor=float(get("attention.scale")) * math.sqrt(head_dim), + output_multiplier=float(get("output_multiplier")), + output_soft_cap_temp=float(get("final_logit_softcapping")), + use_qk_norm=any(n.endswith("attn_q_norm.weight") for n in tensor_names), + use_attn_output_gate=any( + n.endswith("attn_output_gate.weight") for n in tensor_names + ), + tie_word_embeddings="output.weight" not in tensor_names, + bos_token_id=int(meta["tokenizer.ggml.bos_token_id"]), + eos_token_id=int(meta["tokenizer.ggml.eos_token_id"]), + architectures=["MuseGlimmerForCausalLM"], + dtype="bfloat16", + ) diff --git a/python/sglang/srt/configs/muse_glimmer_processing.py b/python/sglang/srt/configs/muse_glimmer_processing.py new file mode 100644 index 000000000000..433d35cebbcb --- /dev/null +++ b/python/sglang/srt/configs/muse_glimmer_processing.py @@ -0,0 +1,302 @@ +# 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 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +import itertools +import json +import math +import os +from typing import Optional + +import torch +from transformers import AutoTokenizer +from transformers.image_processing_backends import TorchvisionBackend +from transformers.image_processing_utils import BatchFeature +from transformers.image_transforms import group_images_by_shape, reorder_images +from transformers.image_utils import PILImageResampling, SizeDict +from transformers.processing_utils import ImagesKwargs, MultiModalData, ProcessorMixin +from transformers.utils import TensorType +from transformers.utils.constants import IMAGENET_STANDARD_MEAN, IMAGENET_STANDARD_STD +from transformers.utils.hub import cached_file + +PROCESSOR_CONFIG_NAME = "processor_config.json" + + +def get_aspect_ratio_preserving_size( + height: int, + width: int, + patch_size: int, + max_tokens: int, +) -> tuple[int, int]: + """Patch grid closest to the aspect ratio; returns (height, width) in pixels.""" + ideal_patches_height = height / patch_size + ideal_patches_width = width / patch_size + ratio = ( + ideal_patches_width / ideal_patches_height if ideal_patches_height > 0 else 1.0 + ) + if ideal_patches_height * ideal_patches_width > max_tokens: + ideal_patches_height = (max_tokens / ratio) ** 0.5 + ideal_patches_width = ideal_patches_height * ratio + candidates = list( + set( + itertools.product( + [math.floor(ideal_patches_height), math.ceil(ideal_patches_height)], + [math.floor(ideal_patches_width), math.ceil(ideal_patches_width)], + ) + ) + ) + candidates = [ + (patches_height, patches_width) + for patches_height, patches_width in candidates + if patches_height >= 1 + and patches_width >= 1 + and patches_height * patches_width <= max_tokens + ] + if not candidates: + candidates = [ + (max(1, round(ideal_patches_height)), max(1, round(ideal_patches_width))) + ] + patches_height, patches_width = min( + candidates, key=lambda grid: abs(grid[0] / grid[1] - height / width) + ) + return patches_height * patch_size, patches_width * patch_size + + +class MuseGlimmerImageProcessorKwargs(ImagesKwargs, total=False): + patch_size: int + temporal_patch_size: int + merge_size: int + max_image_tokens: int + + +class MuseGlimmerImageProcessor(TorchvisionBackend): + do_resize = True + resample = PILImageResampling.LANCZOS + size = None + default_to_square = False + do_rescale = True + rescale_factor = 1 / 255 + do_normalize = True + image_mean = IMAGENET_STANDARD_MEAN + image_std = IMAGENET_STANDARD_STD + do_convert_rgb = True + patch_size = 14 + temporal_patch_size = 2 + merge_size = 2 + max_image_tokens = 4096 + valid_kwargs = MuseGlimmerImageProcessorKwargs + model_input_names = ["pixel_values", "image_grid_thw"] + + def _preprocess( + self, + images: list[torch.Tensor], + do_resize: bool, + resample, + do_rescale: bool, + rescale_factor: float, + do_normalize: bool, + image_mean, + image_std, + return_tensors: Optional[TensorType], + patch_size: int, + temporal_patch_size: int, + max_image_tokens: int, + merge_size: int, + disable_grouping: bool = False, + **kwargs, + ) -> BatchFeature: + if resample == PILImageResampling.LANCZOS: + # BICUBIC stands in for LANCZOS, which is CPU-only. + resample = PILImageResampling.BICUBIC + + grouped_images, grouped_images_index = group_images_by_shape( + images, disable_grouping=disable_grouping + ) + resized_images_grouped = {} + for shape, stacked_images in grouped_images.items(): + if do_resize: + height, width = stacked_images.shape[-2:] + resized_height, resized_width = get_aspect_ratio_preserving_size( + height=height, + width=width, + patch_size=patch_size * merge_size, + max_tokens=max_image_tokens, + ) + stacked_images = self.resize( + image=stacked_images, + size=SizeDict(height=resized_height, width=resized_width), + resample=resample, + antialias=True, + ) + resized_images_grouped[shape] = stacked_images + resized_images = reorder_images(resized_images_grouped, grouped_images_index) + + grouped_images, grouped_images_index = group_images_by_shape( + resized_images, disable_grouping=disable_grouping + ) + processed_images_grouped = {} + processed_grids = {} + for shape, stacked_images in grouped_images.items(): + resized_height, resized_width = stacked_images.shape[-2:] + patches = self.rescale_and_normalize( + stacked_images, + do_rescale, + rescale_factor, + do_normalize, + image_mean, + image_std, + ) + if patches.ndim == 4: + patches = patches.unsqueeze(1) + + if patches.shape[1] % temporal_patch_size != 0: + repeats = patches[:, -1:].repeat(1, temporal_patch_size - 1, 1, 1, 1) + patches = torch.cat([patches, repeats], dim=1) + + batch_size, grid_t, channel = patches.shape[:3] + grid_t = grid_t // temporal_patch_size + grid_h, grid_w = resized_height // patch_size, resized_width // patch_size + + patches = patches.view( + batch_size, + grid_t, + temporal_patch_size, + channel, + grid_h, + patch_size, + grid_w, + patch_size, + ) + patches = patches.permute(0, 1, 4, 6, 2, 3, 5, 7) + flatten_patches = patches.reshape( + batch_size, + grid_t * grid_h * grid_w, + temporal_patch_size * channel * patch_size * patch_size, + ) + + processed_images_grouped[shape] = flatten_patches + processed_grids[shape] = [[grid_t, grid_h, grid_w]] * batch_size + + processed_images = reorder_images( + processed_images_grouped, grouped_images_index + ) + processed_grids = reorder_images(processed_grids, grouped_images_index) + pixel_values = torch.cat(processed_images, dim=0) + image_grid_thw = torch.tensor(processed_grids) + + return BatchFeature( + data={"pixel_values": pixel_values, "image_grid_thw": image_grid_thw}, + tensor_type=return_tensors, + ) + + def get_number_of_image_patches(self, height: int, width: int, images_kwargs=None): + """Patch rows a (height, width) image expands to.""" + images_kwargs = images_kwargs or {} + patch_size = images_kwargs.get("patch_size", self.patch_size) + merge_size = images_kwargs.get("merge_size", self.merge_size) + max_image_tokens = images_kwargs.get("max_image_tokens", self.max_image_tokens) + + resized_height, resized_width = get_aspect_ratio_preserving_size( + height=height, + width=width, + patch_size=patch_size * merge_size, + max_tokens=max_image_tokens, + ) + return (resized_height // patch_size) * (resized_width // patch_size) + + def _validate_preprocess_kwargs(self, **kwargs): + kwargs["do_resize"] = False + super()._validate_preprocess_kwargs(**kwargs) + + +class MuseGlimmerProcessor(ProcessorMixin): + """Expands one ``<|patch|>`` placeholder into an image's patch-token run.""" + + def __init__( + self, + image_processor=None, + tokenizer=None, + chat_template=None, + **kwargs, + ): + self.image_token = "<|patch|>" + self.image_token_id = tokenizer.convert_tokens_to_ids(self.image_token) + + super().__init__( + image_processor=image_processor, + tokenizer=tokenizer, + chat_template=chat_template, + ) + + @classmethod + def from_pretrained(cls, pretrained_model_name_or_path, **kwargs): + trust_remote_code = kwargs.pop("trust_remote_code", False) + revision = kwargs.pop("revision", None) + use_fast = kwargs.pop("use_fast", True) + + tokenizer = AutoTokenizer.from_pretrained( + pretrained_model_name_or_path, + trust_remote_code=trust_remote_code, + revision=revision, + use_fast=use_fast, + ) + image_processor = MuseGlimmerImageProcessor( + **_load_image_processor_kwargs(pretrained_model_name_or_path, revision) + ) + return cls( + image_processor=image_processor, + tokenizer=tokenizer, + chat_template=tokenizer.chat_template, + ) + + def replace_image_token(self, image_inputs: dict, image_idx: int, **kwargs) -> str: + merge_length = self.image_processor.merge_size**2 + num_image_tokens = ( + image_inputs["image_grid_thw"][image_idx].prod() // merge_length + ) + return self.image_token * num_image_tokens + + def _get_num_multimodal_tokens(self, image_sizes=None, **kwargs): + """Placeholder counts per image, without running the image processor.""" + vision_data = {} + if image_sizes is not None: + merge_size = self.image_processor.merge_size + num_image_patches = [ + self.image_processor.get_number_of_image_patches(height, width, kwargs) + for height, width in image_sizes + ] + vision_data.update( + num_image_tokens=[ + patches // merge_size**2 for patches in num_image_patches + ], + num_image_patches=num_image_patches, + ) + return MultiModalData(**vision_data) + + +def _load_image_processor_kwargs(model_path: str, revision: Optional[str]) -> dict: + """Read the image_processor block from processor_config.json.""" + local = os.path.join(model_path, PROCESSOR_CONFIG_NAME) + config_file = ( + local + if os.path.isfile(local) + else cached_file(model_path, PROCESSOR_CONFIG_NAME, revision=revision) + ) + with open(config_file) as f: + config = json.load(f) + image_processor_config = config.get("image_processor", {}) + return { + key: value + for key, value in image_processor_config.items() + if key != "image_processor_type" + } diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py index a2f48d9b76c5..5aa6fcca18f5 100644 --- a/python/sglang/srt/entrypoints/http_server.py +++ b/python/sglang/srt/entrypoints/http_server.py @@ -2196,6 +2196,7 @@ def _execute_server_warmup(server_args: ServerArgs): is_vlm = ( bool(model_info.get("has_image_understanding", False)) and not server_args.language_only + and not server_args.language_model_only and not is_mps() ) if model_info["is_generation"]: diff --git a/python/sglang/srt/entrypoints/openai/protocol.py b/python/sglang/srt/entrypoints/openai/protocol.py index ecb458fc603a..a3ea690656b6 100644 --- a/python/sglang/srt/entrypoints/openai/protocol.py +++ b/python/sglang/srt/entrypoints/openai/protocol.py @@ -1937,6 +1937,7 @@ class MessageProcessingResult: tool_call_constraint: Optional[ToolCallConstraint] = None skip_special_tokens: bool = True require_reasoning: bool = False + skip_special_tokens: bool = True class ToolCallProcessingResult(NamedTuple): diff --git a/python/sglang/srt/entrypoints/openai/serving_chat.py b/python/sglang/srt/entrypoints/openai/serving_chat.py index 34dceaa2b4a0..6919dadf67db 100644 --- a/python/sglang/srt/entrypoints/openai/serving_chat.py +++ b/python/sglang/srt/entrypoints/openai/serving_chat.py @@ -727,11 +727,7 @@ async def _generate_stream_content( remaining_logprobs = None # Handle tool calls - if ( - request.tool_choice != "none" - and self._effective_tools(request) - and self.tool_call_parser - ): + if self._tool_call_parsing_active(request): async for chunk in self._process_tool_call_stream( index, delta, @@ -740,6 +736,7 @@ async def _generate_stream_content( request, has_tool_calls, continuous_usage_stats, + flush=finish_reason_type is not None and finish_reason_type != "abort", ): if chunk: yield chunk @@ -802,6 +799,18 @@ async def _generate_stream_content( usage=usage, ) + def _tool_call_parsing_active(self, request: ChatCompletionRequest) -> bool: + """Whether this request's output runs through the tool-call parser. + + The reasoning parser is told the same thing, so channel-framed formats + keep their framing intact exactly when a tool-call parser consumes it. + """ + return bool( + request.tool_choice != "none" + and self._effective_tools(request) + and self.tool_call_parser + ) + def _validate_request(self, request: ChatCompletionRequest) -> Optional[str]: """Validate that the input is valid.""" if not request.messages: @@ -1129,6 +1138,7 @@ def _process_messages( result.tool_call_constraint = tool_call_constraint result.require_reasoning = thinking_mode + result.skip_special_tokens = request.skip_special_tokens return result def _apply_jinja_template( @@ -1797,6 +1807,7 @@ def _build_chat_response( force_reasoning=force_reasoning, request=request, tokenizer=self.tokenizer_manager.tokenizer, + tool_call_parser_active=self._tool_call_parsing_active(request), ) reasoning_text, text = parser.parse_non_stream(text) except Exception as e: @@ -1810,11 +1821,7 @@ def _build_chat_response( # Handle tool calls tool_calls = None effective_tools = self._effective_tools(request) - if ( - request.tool_choice != "none" - and effective_tools - and self.tool_call_parser - ): + if self._tool_call_parsing_active(request): history_tool_calls_cnt = self._get_history_tool_calls_cnt(request) tool_calls, text, finish_reason = self._process_tool_calls( text, @@ -2106,6 +2113,7 @@ def _process_reasoning_stream( is_force_reasoning, request, tokenizer=self.tokenizer_manager.tokenizer, + tool_call_parser_active=self._tool_call_parsing_active(request), ) reasoning_parser = reasoning_parser_dict[index] reasoning_text, normal_text = reasoning_parser.parse_stream_chunk(delta) @@ -2158,6 +2166,8 @@ def _patch_reasoning_skip_special_tokens( request.skip_special_tokens = False elif self.reasoning_parser == "inkling": request.skip_special_tokens = False + elif self.reasoning_parser == "muse": + request.skip_special_tokens = False def wrap_reasoning_history(self, reasoning_text: str) -> str: """Wrap prior-turn reasoning in the detector's own start/end tokens. @@ -2357,8 +2367,13 @@ async def _process_tool_call_stream( request: ChatCompletionRequest, has_tool_calls: Dict[int, bool], continuous_usage_stats: bool = False, + flush: bool = False, ): - """Process tool calls in streaming response""" + """Process tool calls in streaming response. + + With flush=True (the terminal delta), the parser also drains text it + held back waiting for a marker that can no longer arrive. + """ effective_tools = self._effective_tools(request) if index not in parser_dict: is_required = request.tool_choice == "required" or isinstance( @@ -2400,6 +2415,10 @@ async def _process_tool_call_stream( normal_text, calls = result.normal_text, result.calls else: normal_text, calls = parser.parse_stream_chunk(delta) + if flush: + end_text, end_calls = parser.parse_stream_end() + normal_text = (normal_text or "") + end_text + calls = list(calls) + end_calls # Yield normal text if normal_text: diff --git a/python/sglang/srt/entrypoints/openai/serving_responses.py b/python/sglang/srt/entrypoints/openai/serving_responses.py index 5164673596a5..ee00e6a9f81b 100644 --- a/python/sglang/srt/entrypoints/openai/serving_responses.py +++ b/python/sglang/srt/entrypoints/openai/serving_responses.py @@ -394,6 +394,10 @@ async def create_responses( else None ), ) + if processed_messages is not None: + sampling_params["skip_special_tokens"] = ( + processed_messages.skip_special_tokens + ) # _process_messages set skip_special_tokens on a chat_request # we then discard, so re-apply it to the engine sampling dict. @@ -808,6 +812,7 @@ def _make_response_output_items( *, require_reasoning: bool, ): + chat_tools = self._response_tools_to_chat_tools(request) if self.reasoning_parser: reasoning_parser = ReasoningParser( model_type=self.reasoning_parser, @@ -819,6 +824,11 @@ def _make_response_output_items( ), request=request, tokenizer=self.tokenizer_manager.tokenizer, + tool_call_parser_active=bool( + chat_tools + and self.tool_call_parser + and request.tool_choice != "none" + ), ) reasoning_content, content = reasoning_parser.parse_non_stream(final_output) else: @@ -851,7 +861,6 @@ def _make_response_output_items( ) output_items.append(reasoning_item) - chat_tools = self._response_tools_to_chat_tools(request) is_required = request.tool_choice == "required" tool_call_items: list[ResponseFunctionToolCall] = [] parsed_via_native = False @@ -1977,6 +1986,7 @@ def _sanitize_response_dict(d: dict) -> dict: ), request=request, tokenizer=self.tokenizer_manager.tokenizer, + tool_call_parser_active=isinstance(tool_parser, FunctionCallParser), ) current_output_index = -1 @@ -2210,11 +2220,22 @@ def _close_tool_call_state(tool_index: int): stream_offset = len(text) if not delta and finish_reason is None: continue + flush = ( + finish_reason is not None and finish_reason.get("type") != "abort" + ) if reasoning_parser_obj is not None: reasoning_chunk, delta = reasoning_parser_obj.parse_stream_chunk( delta ) + if flush: + end_reasoning, end_normal = ( + reasoning_parser_obj.parse_stream_end() + ) + if end_reasoning: + reasoning_chunk = (reasoning_chunk or "") + end_reasoning + if end_normal: + delta = (delta or "") + end_normal else: reasoning_chunk = None @@ -2278,7 +2299,7 @@ def _close_tool_call_state(tool_index: int): ) ) - if not delta: + if not delta and not flush: continue if isinstance(tool_parser, JsonArrayParser): @@ -2286,6 +2307,10 @@ def _close_tool_call_state(tool_index: int): normal_text, tool_calls = sp.normal_text or "", sp.calls elif tool_parser is not None: normal_text, tool_calls = tool_parser.parse_stream_chunk(delta) + if flush: + end_text, end_calls = tool_parser.parse_stream_end() + normal_text = (normal_text or "") + end_text + tool_calls = list(tool_calls) + end_calls else: normal_text, tool_calls = delta, [] diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 27fdbe107295..4f7f87b93b0a 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -649,6 +649,8 @@ class Envs: # Number of decode steps between periodic mx.clear_cache() calls. # Set to 0 to disable cache clearing entirely. SGLANG_MLX_CLEAR_CACHE_STEPS = EnvInt(256) + # MLX buffer-cache cap in GB. + SGLANG_MLX_CACHE_LIMIT_GB = EnvFloat(None) # NPU SGLANG_NPU_DISABLE_ACL_FORMAT_WEIGHT = EnvBool(False) diff --git a/python/sglang/srt/function_call/base_format_detector.py b/python/sglang/srt/function_call/base_format_detector.py index 4056230067f4..dd848402efb6 100644 --- a/python/sglang/srt/function_call/base_format_detector.py +++ b/python/sglang/srt/function_call/base_format_detector.py @@ -350,6 +350,14 @@ def has_tool_call(self, text: str) -> bool: """ raise NotImplementedError() + def finish(self, tools: List[Tool]) -> StreamingParseResult: + """Called once when the stream ends; flush any buffered state. + + Detectors that hold text back while waiting for a marker that can no + longer arrive (the stream is over) override this to release it. + """ + return StreamingParseResult() + def supports_structural_tag(self) -> bool: """Return True if this detector supports structural tag format.""" return True diff --git a/python/sglang/srt/function_call/function_call_parser.py b/python/sglang/srt/function_call/function_call_parser.py index 669228a8159a..011114b46a1a 100644 --- a/python/sglang/srt/function_call/function_call_parser.py +++ b/python/sglang/srt/function_call/function_call_parser.py @@ -37,6 +37,7 @@ from sglang.srt.function_call.minimax_m2 import MinimaxM2Detector from sglang.srt.function_call.minimax_m3 import MinimaxM3Detector from sglang.srt.function_call.mistral_detector import MistralDetector +from sglang.srt.function_call.muse_glimmer_detector import MuseGlimmerDetector 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 @@ -78,6 +79,7 @@ class FunctionCallParser: "mimo": MiMoDetector, "minicpm5": MiniCPM5Detector, "mistral": MistralDetector, + "muse": MuseGlimmerDetector, "poolside_v1": PoolsideV1Detector, "pythonic": PythonicDetector, "qwen": Qwen25Detector, @@ -175,6 +177,17 @@ def parse_stream_chunk(self, chunk_text: str) -> Tuple[str, list[ToolCallItem]]: return final_normal_text, final_calls + def parse_stream_end(self) -> Tuple[str, list[ToolCallItem]]: + """Flush detector state once the stream ends. + + Text a detector held back waiting for a marker (which can no longer + arrive) is released as normal text; see BaseFormatDetector.finish(). + """ + if not self.tools: + return "", [] + sp_result = self.detector.finish(self.tools) + return sp_result.normal_text, sp_result.calls + def get_legacy_structural_tag( self, at_least_one: bool = False ) -> StructuralTagResponseFormat: diff --git a/python/sglang/srt/function_call/muse_glimmer_detector.py b/python/sglang/srt/function_call/muse_glimmer_detector.py new file mode 100644 index 000000000000..aab5d52640a8 --- /dev/null +++ b/python/sglang/srt/function_call/muse_glimmer_detector.py @@ -0,0 +1,279 @@ +import json +import logging +import re +from typing import Dict, List, Optional, Set + +from sglang.srt.entrypoints.openai.protocol import Tool +from sglang.srt.environ import envs +from sglang.srt.function_call.base_format_detector import BaseFormatDetector +from sglang.srt.function_call.core_types import ( + StreamingParseResult, + StructureInfo, + ToolCallItem, + _GetInfoFunc, +) + +logger = logging.getLogger(__name__) + +# Channel framing, shared with the reasoning-side MuseGlimmerDetector. +MESSAGE = "<|message|>" +EOM = "<|eom|>" +EOT = "<|eot|>" +START = "<|start|>" + +# ATEM payload markers. +FUNCTION_CALLS_OPEN = "" +FUNCTION_CALLS_CLOSE = "" +INVOKE_CLOSE = "" + +_RECIPIENT_RE = re.compile(r"to=([^\s<]+)") +_INVOKE_OPEN_RE = re.compile(r']*?\bname="(?P[^"]+)"[^>]*?>') +_PARAM_RE = re.compile( + r']*?\bname="(?P[^"]+)"[^>]*?>(?P.*?)' + r"", + re.DOTALL, +) + +# Recipients whose bodies are prose, never tool calls. +_NON_TOOL_RECIPIENTS = frozenset({"self", "user"}) + + +def _is_tool_channel(recipient: Optional[str]) -> bool: + """True when this channel routes to a tool.""" + return recipient is not None and recipient not in _NON_TOOL_RECIPIENTS + + +# Longest marker that could straddle a chunk boundary while streaming. +_MAX_MARKER = max(len(m) for m in (MESSAGE, EOM, EOT, START, FUNCTION_CALLS_OPEN)) + + +def _decode_value(raw: str): + try: + return json.loads(raw) + except (json.JSONDecodeError, ValueError): + return raw + + +def _normalize_name(emitted: str, registered: Set[str]) -> str: + """Strip the chat template's doubled namespace.""" + if not registered or emitted in registered: + return emitted + if "." not in emitted: + return emitted + head, _, tail = emitted.partition(".") + if head == tail and head in registered: + return head + leaf = emitted.rsplit(".", 1)[-1] + matches = [n for n in registered if n.rsplit(".", 1)[-1] == leaf] + if len(matches) == 1: + return matches[0] + return emitted + + +def _could_start_header(text: str) -> bool: + """Whether the tail could still grow into a header.""" + stripped = text.lstrip() + if not stripped: + return True + if not (stripped.startswith("to=") or "to=".startswith(stripped[:3])): + return False + if MESSAGE in stripped: + return True + recipient, angle, marker = stripped[3:].partition("<") + if any(c.isspace() for c in recipient): + return False + return not angle or MESSAGE.startswith("<" + marker) + + +class MuseGlimmerDetector(BaseFormatDetector): + """Format detector for Muse Glimmer's ATEM tool-call blocks.""" + + def __init__(self): + super().__init__() + # Streaming channel state. + self._recipient: Optional[str] = None + self._in_body = False + self._at_stream_start = True + # Name of the invoke whose arguments are still arriving, if any. + self._open_invoke: Optional[str] = None + + def has_tool_call(self, text: str) -> bool: + return FUNCTION_CALLS_OPEN in text or " Set[str]: + return {t.function.name for t in tools or [] if t.function and t.function.name} + + def _emit_call( + self, name: str, args: Dict, registered: Set[str] + ) -> Optional[ToolCallItem]: + """Build one ToolCallItem, honoring the unknown-tool policy.""" + name = _normalize_name(name, registered) + if name not in registered: + logger.warning("Model attempted to call undefined function: %s", name) + if not envs.SGLANG_FORWARD_UNKNOWN_TOOLS.get(): + return None + self.current_tool_id += 1 + parameters = json.dumps(args, ensure_ascii=False) + self.prev_tool_call_arr.append({"name": name, "arguments": args}) + self.streamed_args_for_tool.append(parameters) + return ToolCallItem( + tool_index=self.current_tool_id, name=name, parameters=parameters + ) + + def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult: + result = self.parse_streaming_increment(text, tools) + end = self.finish(tools) + return StreamingParseResult( + normal_text=result.normal_text + end.normal_text, + calls=result.calls + end.calls, + ) + + def finish(self, tools: List[Tool]) -> StreamingParseResult: + registered = self._registered_names(tools) + calls: List[ToolCallItem] = [] + normal_parts: List[str] = [] + + if self._buffer: + if self._in_body: + self._consume_body( + self._buffer, + registered, + calls, + normal_parts, + final=True, + ) + else: + # Truncated header: no body ever arrived, keep it as text. + normal_parts.append(self._buffer) + self._buffer = "" + + return StreamingParseResult(normal_text="".join(normal_parts), calls=calls) + + def _held_back(self, text: str) -> int: + """Length of the trailing suffix that could still grow into a marker.""" + markers = (MESSAGE, EOM, EOT, START, FUNCTION_CALLS_OPEN, " StreamingParseResult: + self._buffer += new_text + registered = self._registered_names(tools) + calls: List[ToolCallItem] = [] + normal_parts: List[str] = [] + + while self._buffer: + if not self._in_body: + # Resolve the channel header before anything can be emitted. + if self._at_stream_start and _could_start_header(self._buffer): + pass + else: + ws = len(self._buffer) - len(self._buffer.lstrip()) + head = self._buffer[ws : ws + len(START)] + if not START.startswith(head): + # Unframed prose: no header is coming. + self._in_body = True + self._recipient = None + self._at_stream_start = False + continue + if ws: + normal_parts.append(self._buffer[:ws]) + self._buffer = self._buffer[ws:] + if len(head) < len(START): + break # Partial "<|start|>", wait for the rest. + + idx = self._buffer.find(MESSAGE) + if idx == -1: + break + header = self._buffer[:idx] + m = _RECIPIENT_RE.search(header) + self._recipient = m.group(1) if m else "user" + self._buffer = self._buffer[idx + len(MESSAGE) :] + self._in_body = True + self._at_stream_start = False + continue + + # Inside a body: find the terminator, if it has arrived. + end_at, end_len = -1, 0 + for tok in (EOM, EOT): + i = self._buffer.find(tok) + if i != -1 and (end_at == -1 or i < end_at): + end_at, end_len = i, len(tok) + + if end_at == -1: + keep = self._held_back(self._buffer) + chunk = self._buffer[: len(self._buffer) - keep] + if not chunk: + break + consumed = self._consume_body( + chunk, registered, calls, normal_parts, final=False + ) + if consumed == 0: + break + self._buffer = self._buffer[consumed:] + continue + + self._consume_body( + self._buffer[:end_at], + registered, + calls, + normal_parts, + final=True, + ) + self._buffer = self._buffer[end_at + end_len :] + self._in_body = False + self._recipient = None + self._open_invoke = None + + return StreamingParseResult(normal_text="".join(normal_parts), calls=calls) + + def _consume_body( + self, + chunk: str, + registered: Set[str], + calls: List[ToolCallItem], + normal_parts: List[str], + final: bool, + ) -> int: + if not _is_tool_channel(self._recipient): + normal_parts.append(chunk) + return len(chunk) + + pos = 0 + while pos < len(chunk): + if self._open_invoke is None: + m = _INVOKE_OPEN_RE.search(chunk, pos) + if m is None: + return pos if not final else len(chunk) + self._open_invoke = m.group("name") + pos = m.end() + continue + + close_at = chunk.find(INVOKE_CLOSE, pos) + if close_at == -1: + return pos if not final else len(chunk) + body = chunk[pos:close_at] + args = { + pm.group("key"): _decode_value(pm.group("value")) + for pm in _PARAM_RE.finditer(body) + } + item = self._emit_call(self._open_invoke, args, registered) + if item is not None: + calls.append(item) + self._open_invoke = None + pos = close_at + len(INVOKE_CLOSE) + + return len(chunk) + + def supports_structural_tag(self) -> bool: + return False + + def structure_info(self) -> _GetInfoFunc: + return lambda name: StructureInfo( + begin=f'{FUNCTION_CALLS_OPEN}\n', + end=f"{INVOKE_CLOSE}\n{FUNCTION_CALLS_CLOSE}", + trigger=FUNCTION_CALLS_OPEN, + ) diff --git a/python/sglang/srt/hardware_backend/mlx/model_runner.py b/python/sglang/srt/hardware_backend/mlx/model_runner.py index 3c91ab47a6a7..a9032998b692 100644 --- a/python/sglang/srt/hardware_backend/mlx/model_runner.py +++ b/python/sglang/srt/hardware_backend/mlx/model_runner.py @@ -58,6 +58,10 @@ set_context, uses_sliding_window_attention, ) +from sglang.srt.hardware_backend.mlx.remote_code_gate import ( + ensure_remote_code_allowed, + resolve_model_directory, +) from sglang.srt.hardware_backend.mlx.sampling import ( GREEDY_PARAMS, MlxLazyLogprobs, @@ -166,12 +170,14 @@ def __init__( pool_size: int | None = None, mem_fraction_static: float = 0.8, quantization: str | None = None, + revision: str | None = None, enable_sampling: bool = False, sampling_rng_seed: int = 0, deterministic_seeding: bool = False, ): self.model_path = model_path self.trust_remote_code = trust_remote_code + self.revision = revision self.model = None self.disable_radix_cache = disable_radix_cache self._mem_fraction_static = mem_fraction_static @@ -196,6 +202,19 @@ def __init__( # modules directly. self._quantization: str | None = quantization + # Optionally cap the buffer cache (recycled GPU buffers). MLX never + # returns freed buffers to the OS, so without a cap the process + # footprint ratchets up to the worst transient — which is model + # load/quantization itself, so the cap must be in place before it. + cache_limit_gb = envs.SGLANG_MLX_CACHE_LIMIT_GB.get() + if cache_limit_gb is not None: + if cache_limit_gb < 0: + raise ValueError( + f"SGLANG_MLX_CACHE_LIMIT_GB must be >= 0, got {cache_limit_gb}" + ) + mx.set_cache_limit(int(cache_limit_gb * (1024**3))) + logger.info(f"MLX buffer cache limit set to {cache_limit_gb:.1f} GB") + self._load_model() # Pin MLX allocations to prevent OS paging @@ -480,10 +499,18 @@ def _load_model(self): logger.info(f"Loading MLX model: {self.model_path}") start_time = time.time() + # Resolve the checkpoint directory once and inspect that exact + # directory before mlx-lm can execute any checkpoint-shipped + # model_file; the same directory is then handed to mlx_lm_load + # (identity resolution for local dirs), so the inspected and + # executed snapshots cannot diverge. + model_dir = resolve_model_directory(self.model_path, revision=self.revision) + ensure_remote_code_allowed(model_dir, self.trust_remote_code) + # We need the config dict to pass into quantize_model so it knows tied/embedding # layout. return_config=True is cheap and ignored when no quantization is requested. loaded = mlx_lm_load( - self.model_path, + str(model_dir), tokenizer_config={"trust_remote_code": self.trust_remote_code}, return_config=True, ) diff --git a/python/sglang/srt/hardware_backend/mlx/models/muse_glimmer_mlx.py b/python/sglang/srt/hardware_backend/mlx/models/muse_glimmer_mlx.py new file mode 100644 index 000000000000..31cba44d0f50 --- /dev/null +++ b/python/sglang/srt/hardware_backend/mlx/models/muse_glimmer_mlx.py @@ -0,0 +1,740 @@ +# 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 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Muse Glimmer (dense, text-only) for mlx-lm. + +Loaded via mlx-lm's custom-architecture path: ship this file in the checkpoint +directory as ``mlx_onyx.py``, set ``"model_file": "mlx_onyx.py"`` in +``config.json``. This copy under ``sglang/srt/hardware_backend/mlx/models/`` +is the maintained source; artifacts ship a byte-identical copy. It must stay +importable standalone (mlx / mlx-lm imports only — no sglang imports), +because mlx-lm executes it from the checkpoint directory. + +Ported from the vendor reference implementation (and cross-checked against the +SGLang CUDA port in ``python/sglang/srt/models/muse_glimmer.py``). Deviations from a +llama-style decoder, and how each is mapped: + +* **Sandwich norms.** ``h = h + post_norm(branch(pre_norm(h)))`` — the post + norm applies to the branch output. Pre-norms use ``rms_norm_eps`` (1e-5), + post-norms ``post_norm_eps`` (1e-8). +* **Norm weights are offsets from 1.0.** The four per-layer norms compute + ``rms_norm(x, weight + 1.0)``; ``sanitize`` folds the +1 in at load time so + plain ``nn.RMSNorm`` is exact. The final ``model.norm`` uses its weight + directly and is NOT offset. +* **Non-parametric QK-norm** over ``head_dim`` (no learnable scale), applied + BEFORE RoPE. Exposed as ``q_norm``/``k_norm`` so the SGLang MLX batched + decode wrapper applies them at the same point. +* **Folded attention scale.** The reference multiplies q by + ``qk_scale_factor / sqrt(head_dim)`` after the QK-norm and SDPA then applies + its default ``1/sqrt(head_dim)``. RoPE is orthogonal and softmax scale is + linear in q, so both fold into ``scale = qk_scale_factor / head_dim``. +* **Attention output gate.** ``sigmoid(output_gate_proj(pre_normed_x))`` is + applied elementwise to the attention output before ``o_proj``. The gate + reads the same input as ``q_proj``, so ``sanitize`` fuses it into ``q_proj`` + per-head-interleaved (``[q_head; gate_head]``) — the exact layout the SGLang + ``MLXAttentionWrapper`` gate path splits back out during batched decode. +* **iRoPE.** ``no_rope_layers[i] == 0`` marks NoPE layers (also the + ``full_attention`` layers); they get a ``NoPE`` identity that still + satisfies the wrapper's ``rope(x, offset=...)`` call. RoPE layers use the + interleaved GPT-J convention (``nn.RoPE(traditional=True)``), which the AOT + Metal RoPE kernel does not support — Muse Glimmer always takes the + ``mx.fast.rope`` fallback. +* **Sliding window.** ``layer_types`` marks the non-NoPE layers as + ``sliding_attention`` (window 2048, including the query position — the same + band as HF's ``create_sliding_window_causal_mask``, so no off-by-one). The + container exposes ``layer_types`` + ``sliding_window`` per the gpt-oss + convention that both mlx-lm and the SGLang MLX backend read; windowing is + done by banded masks over full-history KV, never a per-module + ``is_sliding`` flag. +* **Full-history caches.** ``make_cache`` returns a plain ``KVCache`` for + every layer, including sliding ones (unlike gpt-oss's ``RotatingKVCache``). + Banded masks provide the window semantics; keeping full history makes + greedy output exactly reproducible across prefill chunkings and matches how + the SGLang MLX KV pool stores history. +* **Embedding norm** (scaleless RMS) when ``normalize_tok_embeddings``. +* **Logit head.** ``cap * tanh(lm_head(h) * output_multiplier / cap)``, + computed in float32 like the reference (``None`` cap leaves just the + multiplier). + +Checkpoint formats. ``sanitize`` accepts exactly three weight layouts and +rejects everything else with an actionable error: + +* **Raw HF export** (output of the vendor's HF converter): carries + ``output_gate_proj`` and offset-form norm weights. Recognized by the + complete raw key schema; transformed on load. +* **RC multimodal export** (``transformers >= 5.15`` + vendor schema): text weights under ``model.language_model.`` with + HF-canonical names. Normalized to the raw schema first (see the rename + table at ``_RC_SUFFIX_RENAMES`` — the norm renames are POSITIONAL: the + RC ``post_attention_layernorm`` is the post-attn sandwich norm, i.e. the + raw ``post_attn_norm``, while the raw ``post_attention_layernorm`` is the + pre-MLP norm, i.e. the RC ``pre_feedforward_layernorm``), vision tower / + adapter / projection dropped, then transformed like a raw export. Two + converter generations ship this layout: the older one bakes the scaleless + embedding RMS-norm into ``embed_tokens.weight`` and keeps q/k in the native + interleaved layout; the current one ships the raw table and permutes q/k + into the NeoX rotary layout. The RC path reads the current + conventions (``rope_is_neox_style`` pinned True) and leaves + ``normalize_tok_embeddings`` at its default True — the norm is + idempotent on a baked table, so always-on covers both generations, + but an older-generation export served through this path gets the wrong + rope layout (their configs are byte-identical; prefer repackaging). +* **Packaged MLX artifact**: already fused/folded, marked by + ``"onyx_mlx_format": 1`` in ``config.json`` (stamped at packaging time + only, never present on raw HF exports). Passed through untouched. + +Config schemas. ``ModelArgs.from_dict`` accepts the flat schema written at +packaging time and the RC nested schema (``text_config`` present). +The RC schema differs in two conventions beyond field names: +``qk_scale_factor`` is expressed against SDPA's standard ``1/sqrt(head_dim)`` +(flat-schema value = RC value * sqrt(head_dim); both fold to the same +``scale = flat_qk_scale / head_dim``), and NoPE layers are marked by zeros +in ``layer_rope_theta`` rather than ``no_rope_layers``. +""" + +import math +from dataclasses import dataclass +from typing import Any, List, Optional + +import mlx.core as mx +import mlx.nn as nn +from mlx_lm.models.base import ( + BaseModelArgs, + create_attention_mask, + scaled_dot_product_attention, +) +from mlx_lm.models.cache import KVCache + +# Version of the packaged (fused/folded) weight layout this file understands. +ONYX_MLX_FORMAT_VERSION = 1 + +# The four per-layer norms whose checkpoint weight is an offset from 1.0. +# model.norm (MuseGlimmerFinalRMSNorm) is NOT in this list and must not be offset. +_OFFSET_NORM_SUFFIXES = ( + "input_layernorm.weight", + "post_attn_norm.weight", + "post_attention_layernorm.weight", + "post_ffn_norm.weight", +) + +# Text-only port: the vision tower/projector are not built. +_VISION_KEY_MARKERS = ( + "vision_encoder", + "vision_adapter", + "vision_projection", + "vision_tower", + "perception_emb_norm", +) + +# Keys that only appear in a raw HF export, never in a packaged artifact. +# "language_model" catches RC-layout strays (text weights live under +# model.language_model. there). +_RAW_ONLY_KEY_MARKERS = ( + "output_gate_proj", + "rotary_emb", + "language_model", +) + _VISION_KEY_MARKERS + +# RC (transformers >= 5.15 vendor schema) -> raw-schema key renames, applied +# per key after stripping the "model.language_model." prefix. The norm +# renames are positional, not textual: RC's post_attention_layernorm is the +# post-attn sandwich norm (raw post_attn_norm, eps=post_norm_eps) and RC's +# pre_feedforward_layernorm is the pre-MLP norm (raw post_attention_layernorm, +# eps=rms_norm_eps). self_attn.gate_proj is the attention output gate +# (mlp.gate_proj is untouched: the suffixes below carry the self_attn./ +# module context). +_RC_SUFFIX_RENAMES = ( + ("self_attn.gate_proj.weight", "self_attn.output_gate_proj.weight"), + ("post_attention_layernorm.weight", "post_attn_norm.weight"), + ("pre_feedforward_layernorm.weight", "post_attention_layernorm.weight"), + ("post_feedforward_layernorm.weight", "post_ffn_norm.weight"), +) + +_RC_PREFIX = "model.language_model." + + +def flatten_rc_config(config: dict) -> dict: + """Translate the RC nested config schema into this file's flat schema. + + Field mapping plus three convention conversions (see module docstring): + qk_scale_factor gains the sqrt(head_dim) that the RC schema leaves to + SDPA, NoPE layers come from zeros in layer_rope_theta, and the vendor + export permutes q/k into the NeoX rotary layout (``_permute_for_rope``) + so rope_is_neox_style is pinned True -- ``nn.RoPE(traditional=True)`` + on those weights emits garbled text rather than raising. + + normalize_tok_embeddings is left at its default. Older vendor exports baked + the embedding norm into embed_tokens.weight and needed it off; the current + export ships the native table instead. + """ + text = config["text_config"] + + activation = text.get("hidden_activation", "silu") + if activation != "silu": + raise ValueError( + f"RC config has hidden_activation={activation!r}; this port " + "hardcodes silu" + ) + + head_dim = int(text.get("head_dim", 128)) + rope_params = text.get("rope_parameters") or {} + layer_rope_theta = text.get("layer_rope_theta") + + flat = { + "model_type": "onyx", + "hidden_size": text["hidden_size"], + "num_hidden_layers": text["num_hidden_layers"], + "num_attention_heads": text["num_attention_heads"], + "num_key_value_heads": text["num_key_value_heads"], + "head_dim": head_dim, + "intermediate_size": text["intermediate_size"], + "vocab_size": text["vocab_size"], + "rms_norm_eps": text["rms_norm_eps"], + "post_norm_eps": text["post_norm_eps"], + "rope_theta": rope_params.get("rope_theta", text.get("rope_theta", 500_000.0)), + "max_position_embeddings": text["max_position_embeddings"], + "qk_scale_factor": text["qk_scale_factor"] * math.sqrt(head_dim), + "output_multiplier": text["output_multiplier"], + "output_soft_cap_temp": text.get("final_logit_softcapping"), + "rope_is_neox_style": True, + "sliding_window": text["sliding_window"], + } + if "layer_types" in text: + flat["layer_types"] = list(text["layer_types"]) + if layer_rope_theta is not None: + flat["no_rope_layers"] = [0 if not theta else 1 for theta in layer_rope_theta] + return flat + + +@dataclass +class ModelArgs(BaseModelArgs): + model_type: str = "onyx" + hidden_size: int = 6656 + num_hidden_layers: int = 52 + num_attention_heads: int = 32 + num_key_value_heads: int = 2 + head_dim: int = 128 + intermediate_size: int = 19968 + vocab_size: int = 202048 + rms_norm_eps: float = 1e-5 + post_norm_eps: float = 1e-8 + rope_theta: float = 500_000.0 + max_position_embeddings: int = 16384 + use_qk_norm: bool = True + qk_scale_factor: float = 43.7840518911 + use_attn_output_gate: bool = True + output_multiplier: float = 0.19611613513818404 + output_soft_cap_temp: Optional[float] = 20.0 + rope_is_neox_style: bool = False + normalize_tok_embeddings: bool = True + sliding_window: int = 2048 + every_n_layers_nope: int = 4 + no_rope_layers: Optional[List[int]] = None + layer_types: Optional[List[str]] = None + # Set on saved MLX artifacts at packaging time (never on raw HF + # exports); tells sanitize() the weights are already fused/folded. + onyx_mlx_format: Optional[int] = None + + @classmethod + def from_dict(cls, params): + # RC multimodal schema: text fields nested under text_config, with + # convention differences handled by flatten_rc_config. + if "text_config" in params: + params = flatten_rc_config(params) + return super().from_dict(params) + + def __post_init__(self): + # Mirror the vendor config's derivations so a config.json that + # omits the explicit lists still builds the right architecture. + if self.every_n_layers_nope <= 0: + raise ValueError( + f"every_n_layers_nope must be positive, got {self.every_n_layers_nope}" + ) + if self.num_attention_heads % self.num_key_value_heads != 0: + raise ValueError( + f"num_attention_heads ({self.num_attention_heads}) must be a " + f"multiple of num_key_value_heads ({self.num_key_value_heads})" + ) + + derived_no_rope = [ + 0 if (self.num_hidden_layers - i - 1) % self.every_n_layers_nope == 0 else 1 + for i in range(self.num_hidden_layers) + ] + if self.no_rope_layers is None: + self.no_rope_layers = derived_no_rope + else: + if len(self.no_rope_layers) != self.num_hidden_layers: + raise ValueError( + f"no_rope_layers has {len(self.no_rope_layers)} entries but " + f"num_hidden_layers is {self.num_hidden_layers}" + ) + bad_flags = sorted(set(self.no_rope_layers) - {0, 1}) + if bad_flags: + raise ValueError( + f"no_rope_layers contains non-binary entries {bad_flags}; " + "each entry must be 0 (NoPE) or 1 (RoPE)" + ) + + # NoPE layers are the full-attention layers; the rest slide. + derived_layer_types = [ + "full_attention" if rope_flag == 0 else "sliding_attention" + for rope_flag in self.no_rope_layers + ] + if self.layer_types is None: + self.layer_types = derived_layer_types + else: + if len(self.layer_types) != self.num_hidden_layers: + raise ValueError( + f"layer_types has {len(self.layer_types)} entries but " + f"num_hidden_layers is {self.num_hidden_layers}" + ) + bad = sorted( + set(self.layer_types) - {"full_attention", "sliding_attention"} + ) + if bad: + raise ValueError( + f"layer_types contains unknown entries {bad}; expected only " + "'full_attention' or 'sliding_attention'" + ) + if self.layer_types != derived_layer_types: + mismatches = [ + i + for i, (got, want) in enumerate( + zip(self.layer_types, derived_layer_types) + ) + if got != want + ] + raise ValueError( + "layer_types disagrees with no_rope_layers (NoPE layers " + "must be the full_attention layers) at layer indices " + f"{mismatches}" + ) + + if self.onyx_mlx_format is not None and ( + self.onyx_mlx_format != ONYX_MLX_FORMAT_VERSION + ): + raise ValueError( + f"onyx_mlx_format {self.onyx_mlx_format} is not supported by " + f"this model file (expected {ONYX_MLX_FORMAT_VERSION}); " + "regenerate the artifact with a matching packager" + ) + + +class ScalelessRMSNorm(nn.Module): + """RMS norm with no learnable scale (reference MuseGlimmerScalelessRMSNorm).""" + + def __init__(self, dims: int, eps: float): + super().__init__() + self.dims = dims + self.eps = eps + + def __call__(self, x: mx.array) -> mx.array: + return mx.fast.rms_norm(x, None, self.eps) + + +class NoPE(nn.Module): + """Identity standing in for RoPE on NoPE layers. + + Accepts the ``offset`` kwarg so both this file's forward and the SGLang + ``MLXAttentionWrapper`` (which calls ``rope(x, offset=offsets)`` + unconditionally) can treat every layer uniformly. ``dims = 0`` keeps the + AOT Metal RoPE kernel gating disabled for these layers. + """ + + dims = 0 + traditional = True + + def __call__(self, x: mx.array, offset: Any = 0) -> mx.array: + return x + + +class MuseGlimmerAttention(nn.Module): + def __init__(self, args: ModelArgs, layer_idx: int): + super().__init__() + self.num_attention_heads = args.num_attention_heads + self.num_key_value_heads = args.num_key_value_heads + self.head_dim = args.head_dim + self.use_attn_output_gate = args.use_attn_output_gate + + q_dim = args.num_attention_heads * args.head_dim + kv_dim = args.num_key_value_heads * args.head_dim + + # With the output gate, q_proj holds the per-head-interleaved + # [q_head; gate_head] fusion produced by sanitize(): output width + # 2 * q_dim, split back out in the forward pass. + self.q_proj = nn.Linear( + args.hidden_size, + 2 * q_dim if self.use_attn_output_gate else q_dim, + bias=False, + ) + self.k_proj = nn.Linear(args.hidden_size, kv_dim, bias=False) + self.v_proj = nn.Linear(args.hidden_size, kv_dim, bias=False) + self.o_proj = nn.Linear(q_dim, args.hidden_size, bias=False) + + # Bool flag for the forward-pass branch; q_norm/k_norm stay ABSENT + # (not None) when unused — the SGLang batched-decode wrapper + # duck-types them via hasattr. + self.use_qk_norm = args.use_qk_norm + if args.use_qk_norm: + self.q_norm = ScalelessRMSNorm(args.head_dim, args.rms_norm_eps) + self.k_norm = ScalelessRMSNorm(args.head_dim, args.rms_norm_eps) + # Reference: q *= qk_scale_factor / sqrt(head_dim) after the + # QK-norm, then SDPA scales by 1/sqrt(head_dim); folded here. + self.scale = args.qk_scale_factor / args.head_dim + else: + self.scale = args.head_dim**-0.5 + + use_rope = args.no_rope_layers[layer_idx] == 1 + self.rope = ( + nn.RoPE( + args.head_dim, + traditional=not args.rope_is_neox_style, + base=args.rope_theta, + ) + if use_rope + else NoPE() + ) + + def __call__( + self, + x: mx.array, + mask: Optional[mx.array] = None, + cache: Optional[Any] = None, + ) -> mx.array: + B, L, _ = x.shape + H, Hk, D = self.num_attention_heads, self.num_key_value_heads, self.head_dim + + q = self.q_proj(x) + gate = None + if self.use_attn_output_gate: + # Per-head layout [q_head; gate_head]: same split the SGLang MLX + # batched-decode wrapper performs. + q, gate = mx.split(q.reshape(B, L, H, 2 * D), 2, axis=-1) + else: + q = q.reshape(B, L, H, D) + k = self.k_proj(x).reshape(B, L, Hk, D) + v = self.v_proj(x).reshape(B, L, Hk, D) + + # QK-norm BEFORE RoPE, matching the reference. + if self.use_qk_norm: + q = self.q_norm(q) + k = self.k_norm(k) + + q = q.transpose(0, 2, 1, 3) + k = k.transpose(0, 2, 1, 3) + v = v.transpose(0, 2, 1, 3) + + if cache is not None: + q = self.rope(q, offset=cache.offset) + k = self.rope(k, offset=cache.offset) + k, v = cache.update_and_fetch(k, v) + else: + q = self.rope(q) + k = self.rope(k) + + out = scaled_dot_product_attention(q, k, v, cache, scale=self.scale, mask=mask) + + out = out.transpose(0, 2, 1, 3) + if gate is not None: + out = mx.sigmoid(gate) * out + return self.o_proj(out.reshape(B, L, -1)) + + +class MuseGlimmerMLP(nn.Module): + def __init__(self, args: ModelArgs): + super().__init__() + self.gate_proj = nn.Linear(args.hidden_size, args.intermediate_size, bias=False) + self.up_proj = nn.Linear(args.hidden_size, args.intermediate_size, bias=False) + self.down_proj = nn.Linear(args.intermediate_size, args.hidden_size, bias=False) + + def __call__(self, x: mx.array) -> mx.array: + return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x)) + + +class MuseGlimmerDecoderLayer(nn.Module): + def __init__(self, args: ModelArgs, layer_idx: int): + super().__init__() + self.input_layernorm = nn.RMSNorm(args.hidden_size, args.rms_norm_eps) + self.self_attn = MuseGlimmerAttention(args, layer_idx) + self.post_attn_norm = nn.RMSNorm(args.hidden_size, args.post_norm_eps) + self.post_attention_layernorm = nn.RMSNorm(args.hidden_size, args.rms_norm_eps) + self.mlp = MuseGlimmerMLP(args) + self.post_ffn_norm = nn.RMSNorm(args.hidden_size, args.post_norm_eps) + + def __call__( + self, + x: mx.array, + mask: Optional[mx.array] = None, + cache: Optional[Any] = None, + ) -> mx.array: + # Sandwich norms: the post-norm normalizes the branch output before + # the residual add. + x = x + self.post_attn_norm( + self.self_attn(self.input_layernorm(x), mask, cache) + ) + return x + self.post_ffn_norm(self.mlp(self.post_attention_layernorm(x))) + + +class MuseGlimmerModel(nn.Module): + def __init__(self, args: ModelArgs): + super().__init__() + self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size) + self.embed_norm = ( + ScalelessRMSNorm(args.hidden_size, args.rms_norm_eps) + if args.normalize_tok_embeddings + else None + ) + self.layers = [ + MuseGlimmerDecoderLayer(args, i) for i in range(args.num_hidden_layers) + ] + # Reference MuseGlimmerFinalRMSNorm: weight is the scale, not an offset. + self.norm = nn.RMSNorm(args.hidden_size, args.rms_norm_eps) + + # Container-level window declaration (gpt-oss convention), read by + # both this forward and the SGLang MLX backend's + # get_layer_window_sizes(); per-module ``is_sliding`` flags would + # instead trip the backend's uniform-KV-pool check. + self.layer_types = list(args.layer_types) + self.sliding_window = args.sliding_window + + def __call__( + self, + inputs: mx.array, + cache: Optional[Any] = None, + input_embeddings: Optional[mx.array] = None, + ) -> mx.array: + x = ( + input_embeddings + if input_embeddings is not None + else self.embed_tokens(inputs) + ) + if self.embed_norm is not None: + x = self.embed_norm(x) + + if cache is None: + cache = [None] * len(self.layers) + + # One mask per layer type present, anchored to the first cache of + # that type (all caches of a type share the same offset). + masks = {} + for layer_type in ("full_attention", "sliding_attention"): + try: + idx = self.layer_types.index(layer_type) + except ValueError: + continue + window = self.sliding_window if layer_type == "sliding_attention" else None + if window is not None: + masks[layer_type] = create_attention_mask( + x, cache[idx], window_size=window + ) + else: + masks[layer_type] = create_attention_mask(x, cache[idx]) + + for layer, c, layer_type in zip(self.layers, cache, self.layer_types): + x = layer(x, masks[layer_type], c) + + return self.norm(x) + + +class Model(nn.Module): + def __init__(self, args: ModelArgs): + super().__init__() + self.args = args + self.model_type = args.model_type + self.model = MuseGlimmerModel(args) + self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False) + + @property + def layers(self): + return self.model.layers + + def make_cache(self) -> List[Any]: + # Full-history caches for every layer, sliding ones included: banded + # masks provide the window, and full history keeps greedy output + # exactly reproducible across prefill chunkings (a RotatingKVCache + # would diverge once the prompt exceeds the window). + return [KVCache() for _ in range(len(self.model.layers))] + + def __call__( + self, + inputs: mx.array, + cache: Optional[Any] = None, + input_embeddings: Optional[mx.array] = None, + ) -> mx.array: + hidden = self.model(inputs, cache, input_embeddings) + # Reference computes the logit head in float32. + logits = self.lm_head(hidden).astype(mx.float32) + if self.args.output_soft_cap_temp is not None: + cap = self.args.output_soft_cap_temp + logits = cap * mx.tanh(logits * self.args.output_multiplier / cap) + else: + logits = logits * self.args.output_multiplier + return logits + + # ------------------------------------------------------------------ + # Weight loading + # ------------------------------------------------------------------ + + def _expected_raw_keys(self) -> set: + """The complete key schema of a raw HF export (text path only).""" + keys = {"model.embed_tokens.weight", "model.norm.weight", "lm_head.weight"} + for i in range(self.args.num_hidden_layers): + prefix = f"model.layers.{i}." + keys.update( + prefix + suffix + for suffix in ( + "self_attn.q_proj.weight", + "self_attn.k_proj.weight", + "self_attn.v_proj.weight", + "self_attn.o_proj.weight", + "input_layernorm.weight", + "post_attn_norm.weight", + "post_attention_layernorm.weight", + "post_ffn_norm.weight", + "mlp.gate_proj.weight", + "mlp.up_proj.weight", + "mlp.down_proj.weight", + ) + ) + if self.args.use_attn_output_gate: + keys.add(prefix + "self_attn.output_gate_proj.weight") + return keys + + def sanitize(self, weights: dict) -> dict: + if self.args.onyx_mlx_format == ONYX_MLX_FORMAT_VERSION: + # Packaged artifact: weights are already fused/folded. A raw-only + # key here means the marker was stamped on the wrong directory. + stray = sorted( + k + for k in weights + if any(marker in k for marker in _RAW_ONLY_KEY_MARKERS) + ) + if stray: + raise ValueError( + "config.json claims a packaged Muse Glimmer MLX artifact " + f"(onyx_mlx_format={ONYX_MLX_FORMAT_VERSION}) but the " + f"weights contain raw-checkpoint keys {stray[:4]}" + f"{'...' if len(stray) > 4 else ''}; the marker belongs " + "on packaged artifacts only — repackage from the raw HF export" + ) + return weights + + # No marker: a raw HF export, possibly in the RC multimodal layout — + # normalize that to the raw schema first. + if any(k.startswith(_RC_PREFIX) for k in weights): + weights = _normalize_rc_layout(weights) + + text_keys = { + k + for k in weights + if not any(marker in k for marker in _VISION_KEY_MARKERS) + and "rotary_emb" not in k + } + expected = self._expected_raw_keys() + missing = sorted(expected - text_keys) + unexpected = sorted(text_keys - expected) + if missing or unexpected: + hint = "" + gate_missing = all("output_gate_proj" in k for k in missing) and missing + if gate_missing and not unexpected: + hint = ( + " (weights look already fused: if this is a packaged " + 'artifact, its config.json must carry "onyx_mlx_format": ' + f"{ONYX_MLX_FORMAT_VERSION})" + ) + raise ValueError( + "not a complete raw Muse Glimmer HF checkpoint: " + f"{len(missing)} missing keys {missing[:4]}" + f"{'...' if len(missing) > 4 else ''}, " + f"{len(unexpected)} unexpected keys {unexpected[:4]}" + f"{'...' if len(unexpected) > 4 else ''}{hint}" + ) + + H = self.args.num_attention_heads + D = self.args.head_dim + hidden = self.args.hidden_size + + embed_shape = tuple(weights["model.embed_tokens.weight"].shape) + if embed_shape != (self.args.vocab_size, hidden): + raise ValueError( + f"embed_tokens.weight has shape {embed_shape} but config says " + f"(vocab_size, hidden_size) = ({self.args.vocab_size}, {hidden})" + ) + raw_q_shape = tuple(weights["model.layers.0.self_attn.q_proj.weight"].shape) + if raw_q_shape != (H * D, hidden): + raise ValueError( + f"raw q_proj.weight has shape {raw_q_shape}, expected " + f"({H * D}, {hidden}); a width of {2 * H * D} means the gate " + "is already fused — such artifacts must carry " + f'"onyx_mlx_format": {ONYX_MLX_FORMAT_VERSION} in config.json' + ) + + new_weights = {} + for name, w in weights.items(): + # mlx derives RoPE itself; drop cached buffers. + if "rotary_emb" in name: + continue + if any(marker in name for marker in _VISION_KEY_MARKERS): + continue + # Consumed below when its q_proj comes up. + if name.endswith("output_gate_proj.weight"): + continue + + # The reference computes rms_norm(x, weight + 1.0) for these four + # norms; fold the +1 so plain nn.RMSNorm is exact. model.norm + # (MuseGlimmerFinalRMSNorm) is deliberately not offset. + if name.endswith(_OFFSET_NORM_SUFFIXES): + w = w + 1.0 + + if name.endswith("q_proj.weight") and self.args.use_attn_output_gate: + gate_name = name.replace("q_proj.weight", "output_gate_proj.weight") + g = weights[gate_name] + if tuple(g.shape) != (H * D, hidden): + raise ValueError( + f"{gate_name} has shape {tuple(g.shape)}, expected " + f"({H * D}, {hidden})" + ) + # Per-head interleave [q_head; gate_head]: (H*D, hidden) x2 + # -> (H, 2D, hidden) -> (2*H*D, hidden). + w = mx.concatenate( + [w.reshape(H, D, hidden), g.reshape(H, D, hidden)], axis=1 + ).reshape(2 * H * D, hidden) + + new_weights[name] = w + + return new_weights + + +def _normalize_rc_layout(weights: dict) -> dict: + """Rewrite RC multimodal keys to the raw text-only schema. + + Drops the vision tower/adapter/projection, strips the + ``model.language_model.`` prefix, and applies the positional norm/gate + renames from ``_RC_SUFFIX_RENAMES``. Suffix matching happens per key in + one pass, so the post_attention_layernorm name swap cannot cascade. + """ + out = {} + for name, w in weights.items(): + if any(marker in name for marker in _VISION_KEY_MARKERS): + continue + if name.startswith(_RC_PREFIX): + name = "model." + name[len(_RC_PREFIX) :] + for rc_suffix, raw_suffix in _RC_SUFFIX_RENAMES: + if name.endswith(rc_suffix): + name = name[: -len(rc_suffix)] + raw_suffix + break + out[name] = w + return out + + +EntryClass = Model diff --git a/python/sglang/srt/hardware_backend/mlx/remote_code_gate.py b/python/sglang/srt/hardware_backend/mlx/remote_code_gate.py new file mode 100644 index 000000000000..51c7cc1df966 --- /dev/null +++ b/python/sglang/srt/hardware_backend/mlx/remote_code_gate.py @@ -0,0 +1,112 @@ +# 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 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Pre-execution gate for checkpoint-shipped model code on the MLX backend. + +mlx-lm's loader executes ``config.json``'s ``model_file`` unconditionally: +``mlx_lm.utils.load_model`` imports that Python file straight out of the +checkpoint directory, and ``mlx_lm.load()`` exposes no ``trust_remote_code`` +parameter to refuse it. The gate therefore lives on the SGLang side: + +1. Resolve the model path (local directory or HF repo id + revision) to a + local directory exactly once, with mlx-lm's own resolver. +2. Inspect THAT directory's ``config.json``. If it declares ``model_file`` + and the server was not started with ``--trust-remote-code``, refuse + before any checkpoint Python can execute. +3. Hand the same resolved directory to ``mlx_lm.load`` (for which an + existing local directory is a no-op resolution), so the inspected and + executed snapshots cannot diverge. +""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Optional + + +class RemoteCodeGateError(RuntimeError): + """A checkpoint failed the remote-code gate (refusal or bad metadata).""" + + +def resolve_model_directory(model_path: str, revision: Optional[str] = None) -> Path: + """Resolve a model path or HF repo id to a local snapshot directory. + + Uses mlx-lm's resolver so the directory is byte-identical to what a + direct ``mlx_lm.load`` call would consume; existing local paths are + returned as-is (no network access). mlx-lm 0.31.x exposes this as + ``mlx_lm.utils._download`` (formerly ``get_model_path``). + """ + from mlx_lm.utils import _download + + return Path(_download(model_path, revision=revision)) + + +def ensure_remote_code_allowed(model_dir: Path, trust_remote_code: bool) -> None: + """Refuse ``model_file`` checkpoints unless remote code is trusted. + + Must be called with the SAME resolved directory that is subsequently + passed to ``mlx_lm.load``. Raises :class:`RemoteCodeGateError` before + any checkpoint Python executes when the checkpoint declares + ``model_file`` without ``--trust-remote-code``, when its config is + unreadable, or when the ``model_file`` value is malformed. + """ + config_path = model_dir / "config.json" + try: + config = json.loads(config_path.read_text()) + except FileNotFoundError: + raise RemoteCodeGateError( + f"no config.json in resolved model directory {model_dir}; " + "not a loadable MLX checkpoint" + ) from None + except json.JSONDecodeError as e: + raise RemoteCodeGateError( + f"config.json in {model_dir} is not valid JSON ({e}); refusing " + "to load a checkpoint whose metadata cannot be inspected" + ) from None + if not isinstance(config, dict): + raise RemoteCodeGateError( + f"config.json in {model_dir} must contain a JSON object, " + f"found {type(config).__name__}" + ) + + model_file = config.get("model_file") + if model_file is None: + return + + if not isinstance(model_file, str) or not model_file: + raise RemoteCodeGateError( + f"config.json in {model_dir} has a non-string or empty " + f"model_file entry ({model_file!r})" + ) + candidate = Path(model_file) + if candidate.is_absolute() or ".." in candidate.parts: + raise RemoteCodeGateError( + f"model_file {model_file!r} in {model_dir} must be a relative " + "path inside the checkpoint directory (no absolute paths, no " + "'..' traversal)" + ) + if not (model_dir / candidate).is_file(): + raise RemoteCodeGateError( + f"config.json in {model_dir} declares model_file " + f"{model_file!r} but that file does not exist in the " + "checkpoint directory" + ) + + if not trust_remote_code: + raise RemoteCodeGateError( + f"checkpoint {model_dir} ships custom model code " + f"(model_file={model_file!r} in config.json), which mlx-lm " + "would execute at load time. Refusing to load it: restart the " + "server with --trust-remote-code if you trust this checkpoint." + ) diff --git a/python/sglang/srt/hardware_backend/mlx/tp_worker.py b/python/sglang/srt/hardware_backend/mlx/tp_worker.py index 2766e03f5600..8f654c457a7c 100644 --- a/python/sglang/srt/hardware_backend/mlx/tp_worker.py +++ b/python/sglang/srt/hardware_backend/mlx/tp_worker.py @@ -90,6 +90,7 @@ def _init_model_runner(self): disable_radix_cache=get_memory().disable_radix_cache, mem_fraction_static=get_schedule().mem_fraction_static, quantization=get_model().quantization, + revision=get_model().revision, enable_sampling=get_device().mlx_enable_sampling, sampling_rng_seed=get_device().random_seed, deterministic_seeding=( diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index dd035c148a1b..4eed4e8a9511 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -827,7 +827,11 @@ def init_tokenizer(self): # Load multimodal processor for M-RoPE fallback computation. self._mm_processor = None - if self.model_config.is_multimodal and self.processor is not None: + if ( + self.model_config.is_multimodal + and self.processor is not None + and not server_args.language_model_only + ): try: import_processors("sglang.srt.multimodal.processors") self._mm_processor = get_mm_processor( diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index 1983fdd1a802..3e8586fe001c 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -466,7 +466,7 @@ def init_tokenizer_and_processor(self): server_args = self.server_args # Initialize tokenizer and processor - if self.model_config.is_multimodal: + if self.model_config.is_multimodal and not server_args.language_model_only: import_processors("sglang.srt.multimodal.processors") if mm_process_pkg := envs.SGLANG_EXTERNAL_MM_PROCESSOR_PACKAGE.get(): import_processors(mm_process_pkg, overwrite=True) @@ -1023,6 +1023,11 @@ async def _tokenize_one_request( ) contains_mm_input = obj.contains_mm_input() + if contains_mm_input and self.server_args.language_model_only: + raise ValueError( + "Multimodal inputs are not supported when --language-model-only " + "is set; the encoder is not loaded. Restart without the flag." + ) is_mossvl = ( "MossVLForConditionalGeneration" in self.model_config.hf_config.architectures diff --git a/python/sglang/srt/mem_cache/kv_cache_dtype.py b/python/sglang/srt/mem_cache/kv_cache_dtype.py index 2bbdb499aab6..fc71a081f6bc 100644 --- a/python/sglang/srt/mem_cache/kv_cache_dtype.py +++ b/python/sglang/srt/mem_cache/kv_cache_dtype.py @@ -27,8 +27,13 @@ def configure_kv_cache_dtype( is_draft_worker: bool, is_dflash: bool, speculative_draft_attention_backend: str, + speculative_draft_kv_cache_dtype: Optional[str] = None, ) -> tuple[Optional[str], torch.dtype]: resolved_kv_cache_dtype: Optional[str] = None + if is_draft_worker and speculative_draft_kv_cache_dtype is not None: + server_args_kv_cache_dtype = speculative_draft_kv_cache_dtype + if server_args_kv_cache_dtype != "auto": + resolved_kv_cache_dtype = server_args_kv_cache_dtype if server_args_kv_cache_dtype == "auto": quant_config = getattr(model, "quant_config", None) kv_cache_quant_algo = getattr(quant_config, "kv_cache_quant_algo", None) diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 8aa3f65510b1..ac4879b340e4 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -1283,6 +1283,7 @@ def configure_kv_cache_dtype(self): else False ), speculative_draft_attention_backend=self.draft_attention_backend, + speculative_draft_kv_cache_dtype=self.server_args.speculative_draft_kv_cache_dtype, ) ) # This runner's OWN resolved dtype string (target or draft). Attention diff --git a/python/sglang/srt/model_executor/model_runner_components/spec_aux_hidden_state.py b/python/sglang/srt/model_executor/model_runner_components/spec_aux_hidden_state.py index 08dd87afddf4..f060d0fccfe3 100644 --- a/python/sglang/srt/model_executor/model_runner_components/spec_aux_hidden_state.py +++ b/python/sglang/srt/model_executor/model_runner_components/spec_aux_hidden_state.py @@ -146,6 +146,13 @@ def _resolve_dflash_aux_hidden_state( draft_num_layers=int(draft_num_layers), ) + # Native export uses HF layer-output ids; shift them. + draft_architectures = ( + getattr(draft_model_config.hf_config, "architectures", None) or [] + ) + if "MuseGlimmerAssistantModel" in draft_architectures: + target_layer_ids = [i + 1 for i in target_layer_ids] + if spec_algorithm.is_dspark(): from sglang.srt.speculative.dspark_components.dspark_config import ( parse_dspark_draft_config, @@ -190,6 +197,9 @@ def _resolve_dflash_draft_cell_size( try: _, draft_kv_cache_dtype = configure_kv_cache_dtype( server_args_kv_cache_dtype=server_args.kv_cache_dtype, + speculative_draft_kv_cache_dtype=( + server_args.speculative_draft_kv_cache_dtype + ), model=None, model_dtype=draft_model_config.dtype, is_draft_worker=True, diff --git a/python/sglang/srt/model_executor/runner/base_runner.py b/python/sglang/srt/model_executor/runner/base_runner.py index 1597344ca19b..265b6fe0d064 100644 --- a/python/sglang/srt/model_executor/runner/base_runner.py +++ b/python/sglang/srt/model_executor/runner/base_runner.py @@ -314,7 +314,9 @@ def forward_fn(): run_ctx=canary_run_ctx, ) - run_flashinfer_autotune_forward(self.model_runner, forward_fn, skip_logits=True) + run_flashinfer_autotune_forward( + self.model_runner, forward_fn, skip_logits=False + ) def _alloc_dummy_decode_buffers( self, diff --git a/python/sglang/srt/model_loader/gguf_name_maps.py b/python/sglang/srt/model_loader/gguf_name_maps.py new file mode 100644 index 000000000000..4aefaca681c7 --- /dev/null +++ b/python/sglang/srt/model_loader/gguf_name_maps.py @@ -0,0 +1,69 @@ +# 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 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Per-architecture GGUF -> HF tensor name maps. + +``GGUFModelLoader`` normally derives this map from ``gguf.get_tensor_name_map``, +which only covers architectures upstream gguf-py knows, and from a meta-device +``AutoModelForCausalLM.from_config`` to enumerate the HF parameter names. Neither +works for an architecture that lives outside transformers, so those are supplied +here instead. + +A builder returns the complete ``{gguf_tensor_name: hf_param_name}`` map. Any +GGUF tensor left out of the map is skipped by ``gguf_quant_weights_iterator``, +which is how dummy tensors are dropped. +""" + +from typing import Callable, Dict + +from transformers import PretrainedConfig + +# Sandwich naming: ffn_norm is the pre-FFN norm. +_MUSE_GLIMMER_LAYER_TENSORS = { + "attn_norm": "input_layernorm", + "post_attention_norm": "post_attn_norm", + "ffn_norm": "post_attention_layernorm", + "post_ffw_norm": "post_ffn_norm", + "attn_q": "self_attn.q_proj", + "attn_k": "self_attn.k_proj", + "attn_v": "self_attn.v_proj", + "attn_output": "self_attn.o_proj", + "attn_output_gate": "self_attn.output_gate_proj", + "ffn_gate": "mlp.gate_proj", + "ffn_up": "mlp.up_proj", + "ffn_down": "mlp.down_proj", +} + +_MUSE_GLIMMER_GLOBAL_TENSORS = { + "token_embd": "model.embed_tokens", + "output_norm": "model.norm", + "output": "lm_head", +} + +# attn_q_norm/attn_k_norm omitted: Muse Glimmer's QK-norm is non-parametric. + + +def build_muse_glimmer_name_map(config: PretrainedConfig) -> Dict[str, str]: + name_map = { + f"{gguf}.weight": f"{hf}.weight" + for gguf, hf in _MUSE_GLIMMER_GLOBAL_TENSORS.items() + } + for layer in range(config.num_hidden_layers): + for gguf, hf in _MUSE_GLIMMER_LAYER_TENSORS.items(): + name_map[f"blk.{layer}.{gguf}.weight"] = f"model.layers.{layer}.{hf}.weight" + return name_map + + +GGUF_HF_NAME_MAP_BUILDERS: Dict[str, Callable[[PretrainedConfig], Dict[str, str]]] = { + "muse-glimmer": build_muse_glimmer_name_map, +} diff --git a/python/sglang/srt/model_loader/loader.py b/python/sglang/srt/model_loader/loader.py index cb941a800495..ebe6c516fd87 100644 --- a/python/sglang/srt/model_loader/loader.py +++ b/python/sglang/srt/model_loader/loader.py @@ -3038,8 +3038,14 @@ def _get_gguf_weights_map(self, model_config: ModelConfig): "Please install gguf via `pip install gguf` to use gguf quantizer." ) from err + from sglang.srt.model_loader.gguf_name_maps import GGUF_HF_NAME_MAP_BUILDERS + config = model_config.hf_config model_type = config.model_type + name_map_builder = GGUF_HF_NAME_MAP_BUILDERS.get(model_type) + if name_map_builder is not None: + return name_map_builder(config) + # hack: ggufs have a different name than transformers if model_type == "cohere": model_type = "command-r" diff --git a/python/sglang/srt/models/dflash.py b/python/sglang/srt/models/dflash.py index 69c01c7adc9d..688bf12ce123 100644 --- a/python/sglang/srt/models/dflash.py +++ b/python/sglang/srt/models/dflash.py @@ -65,9 +65,12 @@ def _get_dflash_layer_attention_params( ) return -1, attention_type if layer_type == "sliding_attention": + # DFlash uses non-causal attention over the draft block on every layer -- + # the reference sets it on the whole draft context, so a sliding layer is + # windowed but still bidirectional (masking only p1 - p0 >= sliding_window). sliding_window_size = get_dflash_attention_sliding_window_size(config) assert sliding_window_size is not None - return sliding_window_size, AttentionType.DECODER + return sliding_window_size, AttentionType.ENCODER_ONLY raise ValueError( "Unsupported DFLASH draft layer type. " f"layer_types[{layer_id}]={layer_type!r}." @@ -359,11 +362,12 @@ def __init__(self, config, quant_config=None, prefix: str = "") -> None: # concat(K * hidden_size) -> hidden_size, where K is the number of target-layer # feature tensors concatenated per token (not necessarily equal to num_layers). draft_config = parse_dflash_draft_config(draft_hf_config=config) - target_num_layers = ( - int(draft_config.num_target_layers) - if draft_config.num_target_layers is not None - else num_layers - ) + if draft_config.num_target_layers is not None: + target_num_layers = int(draft_config.num_target_layers) + elif draft_config.target_layer_ids is not None: + target_num_layers = max(draft_config.target_layer_ids) + 1 + else: + target_num_layers = num_layers target_layer_ids = draft_config.resolve_target_layer_ids( target_num_layers=target_num_layers, draft_num_layers=num_layers ) @@ -448,6 +452,12 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): params_dict = dict(self.named_parameters()) + # Alias the native export's "encoder." names. + _VENDOR_ENCODER_ALIASES = { + "encoder.fc.weight": "fc.weight", + "encoder.output_norm_enc.weight": "hidden_norm.weight", + } + def resolve_param_name(name: str) -> Optional[str]: if name in params_dict: return name @@ -459,6 +469,9 @@ def resolve_param_name(name: str) -> Optional[str]: prefixed_name = f"model.{name}" if prefixed_name in params_dict: return prefixed_name + aliased_name = _VENDOR_ENCODER_ALIASES.get(name) + if aliased_name is not None and aliased_name in params_dict: + return aliased_name return None for name, loaded_weight in weights: @@ -586,4 +599,8 @@ def project_target_hidden(self, target_hidden: torch.Tensor) -> torch.Tensor: return self.hidden_norm(self.fc(fused)) -EntryClass = [DFlashDraftModel, DFlashLagunaForCausalLM] +class MuseGlimmerAssistantModel(DFlashDraftModel): + """Alias for checkpoints declaring architectures=["MuseGlimmerAssistantModel"].""" + + +EntryClass = [DFlashDraftModel, DFlashLagunaForCausalLM, MuseGlimmerAssistantModel] diff --git a/python/sglang/srt/models/muse_glimmer.py b/python/sglang/srt/models/muse_glimmer.py new file mode 100644 index 000000000000..a1ddf6fbcad2 --- /dev/null +++ b/python/sglang/srt/models/muse_glimmer.py @@ -0,0 +1,1027 @@ +# 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 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +import logging +import re +from typing import Iterable, List, Optional, Tuple + +import torch +from torch import nn +from transformers.activations import ACT2FN +from transformers.vision_utils import ( + get_vision_cu_seqlens, + get_vision_position_ids, + get_vision_window_index, +) + +from sglang.srt.layers.activation import SiluAndMul +from sglang.srt.layers.attention.vision import ( + VisionAttention, + VisionAttentionMetadata, + prepare_vision_attention_metadata, +) +from sglang.srt.layers.layernorm import RMSNorm +from sglang.srt.layers.linear import ( + ColumnParallelLinear, + MergedColumnParallelLinear, + QKVParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) +from sglang.srt.layers.logits_processor import LogitsProcessor +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.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from sglang.srt.managers.mm_utils import ( + MultiModalityDataPaddingPatternMultimodalTokens, + general_mm_embed_routine, +) +from sglang.srt.managers.schedule_batch import MultimodalDataItem, MultimodalInputs +from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.model_loader.weight_utils import ( + default_weight_loader, + maybe_remap_kv_scale_name, +) +from sglang.srt.models.utils import apply_qk_norm, permute_inv +from sglang.srt.runtime_context import get_parallel +from sglang.srt.utils import add_prefix, is_cuda + +_is_cuda = is_cuda() + +if _is_cuda: + from sglang.kernels.ops.elementwise.elementwise import fused_sigmoid_mul + +logger = logging.getLogger(__name__) + +# The four per-layer norms whose checkpoint weight is an offset from 1.0. +_OFFSET_NORM_SUFFIXES = ( + "input_layernorm.weight", + "post_attn_norm.weight", + "post_attention_layernorm.weight", + "post_ffn_norm.weight", +) + +_VISION_NAME_FRAGMENTS = ( + "vision_encoder", + "vision_tower", + "vision_adapter", + "vision_projection", + "perception_emb_norm", +) + +# Vendor tensor names -> this port's; applied simultaneously. +_VENDOR_RENAMES = { + "post_attention_layernorm": "post_attn_norm", + "pre_feedforward_layernorm": "post_attention_layernorm", + "post_feedforward_layernorm": "post_ffn_norm", + "self_attn.gate_proj": "self_attn.output_gate_proj", +} + +_VENDOR_RENAME_RE = re.compile("|".join(re.escape(key) for key in _VENDOR_RENAMES)) + + +def _vendor_weight_name(name: str) -> str: + name = name.replace("model.language_model.", "model.", 1) + # The vision modules hang off the entry class, not off ``model``. + name = name.replace("model.vision_", "vision_", 1) + return _VENDOR_RENAME_RE.sub(lambda m: _VENDOR_RENAMES[m.group(0)], name) + + +def get_attention_sliding_window_size(config) -> int: + return config.sliding_window - 1 + + +class MuseGlimmerMLP(nn.Module): + def __init__( + self, + config, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + self.gate_up_proj = MergedColumnParallelLinear( + config.hidden_size, + [config.intermediate_size] * 2, + bias=False, + quant_config=quant_config, + prefix=add_prefix("gate_up_proj", prefix), + ) + self.down_proj = RowParallelLinear( + config.intermediate_size, + config.hidden_size, + bias=False, + quant_config=quant_config, + prefix=add_prefix("down_proj", prefix), + ) + if config.hidden_act != "silu": + raise ValueError( + f"Muse Glimmer expects hidden_act=silu, got {config.hidden_act}" + ) + self.act_fn = SiluAndMul() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + gate_up, _ = self.gate_up_proj(x) + x = self.act_fn(gate_up) + x, _ = self.down_proj(x) + return x + + +class MuseGlimmerAttention(nn.Module): + def __init__( + self, + config, + layer_id: int, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + tp_size = get_parallel().tp_size + self.layer_id = layer_id + self.hidden_size = config.hidden_size + self.head_dim = config.head_dim + + self.total_num_heads = config.num_attention_heads + assert self.total_num_heads % tp_size == 0 + self.num_heads = self.total_num_heads // tp_size + + self.total_num_kv_heads = config.num_key_value_heads + if self.total_num_kv_heads >= tp_size: + assert self.total_num_kv_heads % tp_size == 0 + else: + assert tp_size % self.total_num_kv_heads == 0 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) + + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + + self.scaling = config.qk_scale_factor / self.head_dim + + self.use_rope = config.no_rope_layers[layer_id] == 1 + self.is_sliding = config.layer_types[layer_id] == "sliding_attention" + self.use_qk_norm = config.use_qk_norm + self.use_output_gate = config.use_attn_output_gate + + # Split q/k/v; one module holds one quant format. + self.unfused_qkv = quant_config is not None + if self.unfused_qkv: + if self.total_num_kv_heads < tp_size: + raise ValueError( + "Muse Glimmer unfused q/k/v needs num_key_value_heads >= tp_size " + f"(got {self.total_num_kv_heads} kv heads, tp_size={tp_size}); " + "the KV-head replication QKVParallelLinear does is not " + "reproduced here." + ) + self.q_proj = ColumnParallelLinear( + config.hidden_size, + self.total_num_heads * self.head_dim, + bias=False, + quant_config=quant_config, + prefix=add_prefix("q_proj", prefix), + ) + self.k_proj = ColumnParallelLinear( + config.hidden_size, + self.total_num_kv_heads * self.head_dim, + bias=False, + quant_config=quant_config, + prefix=add_prefix("k_proj", prefix), + ) + self.v_proj = ColumnParallelLinear( + config.hidden_size, + self.total_num_kv_heads * self.head_dim, + bias=False, + quant_config=quant_config, + prefix=add_prefix("v_proj", prefix), + ) + else: + self.qkv_proj = QKVParallelLinear( + config.hidden_size, + self.head_dim, + self.total_num_heads, + self.total_num_kv_heads, + bias=False, + quant_config=quant_config, + prefix=add_prefix("qkv_proj", prefix), + ) + self.o_proj = RowParallelLinear( + self.total_num_heads * self.head_dim, + config.hidden_size, + bias=False, + quant_config=quant_config, + prefix=add_prefix("o_proj", prefix), + ) + if self.use_output_gate: + self.output_gate_proj = ColumnParallelLinear( + config.hidden_size, + self.total_num_heads * self.head_dim, + bias=False, + quant_config=quant_config, + prefix=add_prefix("output_gate_proj", prefix), + ) + + self.qk_norm = ( + RMSNorm( + hidden_size=self.head_dim, + eps=config.rms_norm_eps, + has_weight=False, + ) + if self.use_qk_norm + else None + ) + + # Reference RoPE pairs adjacent elements (GPT-J), not NeoX. + self.rotary_emb = ( + get_rope( + self.head_dim, + rotary_dim=self.head_dim, + max_position=config.max_position_embeddings, + base=config.rope_theta, + is_neox_style=config.rope_is_neox_style, + ) + if self.use_rope + else None + ) + + self.attn = RadixAttention( + self.num_heads, + self.head_dim, + self.scaling, + num_kv_heads=self.num_kv_heads, + layer_id=layer_id, + sliding_window_size=( + get_attention_sliding_window_size(config) if self.is_sliding else -1 + ), + quant_config=quant_config, + prefix=add_prefix("attn", prefix), + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + if self.unfused_qkv: + q, _ = self.q_proj(hidden_states) + k, _ = self.k_proj(hidden_states) + v, _ = self.v_proj(hidden_states) + else: + qkv, _ = self.qkv_proj(hidden_states) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + + if self.qk_norm is not None: + q, k = apply_qk_norm( + q, + k, + q_norm=self.qk_norm, + k_norm=self.qk_norm, + head_dim=self.head_dim, + ) + + if self.rotary_emb is not None: + q, k = self.rotary_emb(positions, q, k) + + attn_out = self.attn(q, k, v, forward_batch) + + if self.use_output_gate: + gate, _ = self.output_gate_proj(hidden_states) + if _is_cuda: + attn_out = fused_sigmoid_mul(attn_out, gate, inplace=True) + else: + attn_out = torch.sigmoid(gate) * attn_out + + out, _ = self.o_proj(attn_out) + return out + + +class MuseGlimmerDecoderLayer(nn.Module): + def __init__( + self, + config, + layer_id: int, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.self_attn = MuseGlimmerAttention( + config, + layer_id, + quant_config=quant_config, + prefix=add_prefix("self_attn", prefix), + ) + self.post_attn_norm = RMSNorm(config.hidden_size, eps=config.post_norm_eps) + self.post_attention_layernorm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + self.mlp = MuseGlimmerMLP( + config, quant_config=quant_config, prefix=add_prefix("mlp", prefix) + ) + self.post_ffn_norm = RMSNorm(config.hidden_size, eps=config.post_norm_eps) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + hidden_states = self.self_attn(positions, hidden_states, forward_batch) + hidden_states = residual + self.post_attn_norm(hidden_states) + + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + self.post_ffn_norm(hidden_states) + return hidden_states + + +class MuseGlimmerModel(nn.Module): + def __init__( + self, + config, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + self.config = config + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=add_prefix("embed_tokens", prefix), + ) + self.embed_norm = ( + RMSNorm(config.hidden_size, eps=config.rms_norm_eps, has_weight=False) + if config.normalize_tok_embeddings + else None + ) + self.layers = nn.ModuleList( + [ + MuseGlimmerDecoderLayer( + config, + i, + quant_config=quant_config, + prefix=add_prefix(f"layers.{i}", prefix), + ) + for i in range(config.num_hidden_layers) + ] + ) + # Reference MuseGlimmerFinalRMSNorm: weight is the scale, not an offset. + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + self.layers_to_capture: List[int] = [] + + def get_input_embeddings(self) -> nn.Module: + return self.embed_tokens + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + hidden_states = ( + self.embed_tokens(input_ids) if input_embeds is None else input_embeds + ) + if self.embed_norm is not None: + hidden_states = self.embed_norm(hidden_states) + + aux_hidden_states = [] + for i, layer in enumerate(self.layers): + if i in self.layers_to_capture: + aux_hidden_states.append(hidden_states) + hidden_states = layer(positions, hidden_states, forward_batch) + + hidden_states = self.norm(hidden_states) + if aux_hidden_states: + return hidden_states, aux_hidden_states + return hidden_states + + +def muse_glimmer_bilinear_pos_embed_lookup( + grid_thw: torch.Tensor, num_grid_per_side: int +) -> Tuple[torch.Tensor, torch.Tensor]: + """Bilinear resample of the position-embedding grid; returns (indices, weights) to gather-and-sum.""" + side = num_grid_per_side + device = grid_thw.device + index_parts: List[List[torch.Tensor]] = [[] for _ in range(4)] + weight_parts: List[List[torch.Tensor]] = [[] for _ in range(4)] + + for t, h, w in grid_thw.tolist(): + t, h, w = int(t), int(h), int(w) + h_grid = (torch.arange(h, device=device).float() + 0.5) * (side / h) - 0.5 + w_grid = (torch.arange(w, device=device).float() + 0.5) * (side / w) - 0.5 + + h_floor = torch.floor(h_grid).long() + w_floor = torch.floor(w_grid).long() + h_ceil = h_floor + 1 + w_ceil = w_floor + 1 + h_frac = h_grid - h_floor.float() + w_frac = w_grid - w_floor.float() + + h_floor_valid = (h_floor >= 0) & (h_floor <= side - 1) + h_ceil_valid = (h_ceil >= 0) & (h_ceil <= side - 1) + w_floor_valid = (w_floor >= 0) & (w_floor <= side - 1) + w_ceil_valid = (w_ceil >= 0) & (w_ceil <= side - 1) + h_floor = h_floor.clamp(0, side - 1) + h_ceil = h_ceil.clamp(0, side - 1) + w_floor = w_floor.clamp(0, side - 1) + w_ceil = w_ceil.clamp(0, side - 1) + + h_floor_offset = h_floor * side + h_ceil_offset = h_ceil * side + corner_indices = [ + (h_floor_offset[:, None] + w_floor[None, :]).flatten(), + (h_floor_offset[:, None] + w_ceil[None, :]).flatten(), + (h_ceil_offset[:, None] + w_floor[None, :]).flatten(), + (h_ceil_offset[:, None] + w_ceil[None, :]).flatten(), + ] + corner_weights = [ + ( + (1 - h_frac)[:, None] + * (1 - w_frac)[None, :] + * (h_floor_valid[:, None] & w_floor_valid[None, :]) + ).flatten(), + ( + (1 - h_frac)[:, None] + * w_frac[None, :] + * (h_floor_valid[:, None] & w_ceil_valid[None, :]) + ).flatten(), + ( + h_frac[:, None] + * (1 - w_frac)[None, :] + * (h_ceil_valid[:, None] & w_floor_valid[None, :]) + ).flatten(), + ( + h_frac[:, None] + * w_frac[None, :] + * (h_ceil_valid[:, None] & w_ceil_valid[None, :]) + ).flatten(), + ] + for corner in range(4): + index_parts[corner].append(corner_indices[corner].repeat(t)) + weight_parts[corner].append(corner_weights[corner].repeat(t)) + + indices = torch.stack([torch.cat(part) for part in index_parts]) + weights = torch.stack([torch.cat(part) for part in weight_parts]) + return indices, weights + + +class MuseGlimmerVisionPatchEmbedder(nn.Module): + """Linear patch projection plus a bilinearly resampled learned position grid.""" + + def __init__( + self, + config, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + self.num_grid_per_side = config.pos_emb_height + self.patch_embedding = ReplicatedLinear( + config.patch_temporal * 3 * config.patch_size**2, + config.hidden_size, + bias=False, + quant_config=quant_config, + prefix=add_prefix("patch_embedding", prefix), + ) + self.position_embedding_table = nn.Embedding( + config.pos_emb_height * config.pos_emb_width, config.hidden_size + ) + + def forward( + self, pixel_values: torch.Tensor, grid_thw: torch.Tensor + ) -> torch.Tensor: + embeddings, _ = self.patch_embedding(pixel_values) + indices, weights = muse_glimmer_bilinear_pos_embed_lookup( + grid_thw, self.num_grid_per_side + ) + indices = indices.to(embeddings.device) + weights = weights.to(device=embeddings.device, dtype=torch.float32) + pos_embeds = ( + self.position_embedding_table(indices).float() * weights[:, :, None] + ).sum(0) + return embeddings + pos_embeds.to(embeddings.dtype) + + +class MuseGlimmerVisionRotaryEmbedding(nn.Module): + + def __init__(self, head_dim: int, theta: float): + super().__init__() + spatial_dim = head_dim // 2 + inv_freq = 1.0 / ( + theta ** (torch.arange(0, spatial_dim, 2, dtype=torch.float) / spatial_dim) + ) + self.register_buffer("inv_freq", inv_freq, persistent=False) + + def forward(self, position_ids: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + # position_ids: (total_patches, 2) laid out as (width, height). + freqs = position_ids.float()[:, :, None] * self.inv_freq[None, None, :] + freq = torch.cat([freqs[:, 0], freqs[:, 1]], dim=-1) + return freq.cos(), freq.sin() + + +class MuseGlimmerVisionMLP(nn.Module): + def __init__( + self, + config, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + self.fc1 = ColumnParallelLinear( + config.hidden_size, + config.intermediate_size, + bias=True, + quant_config=quant_config, + prefix=add_prefix("fc1", prefix), + ) + self.fc2 = RowParallelLinear( + config.intermediate_size, + config.hidden_size, + bias=True, + quant_config=quant_config, + prefix=add_prefix("fc2", prefix), + ) + self.act = ACT2FN[config.hidden_act] + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x, _ = self.fc1(x) + x = self.act(x) + x, _ = self.fc2(x) + return x + + +class MuseGlimmerVisionEncoderLayer(nn.Module): + def __init__( + self, + config, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + self.norm1 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.norm2 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.attn = VisionAttention( + embed_dim=config.hidden_size, + num_heads=config.num_attention_heads, + projection_size=config.hidden_size, + use_qkv_parallel=True, + qkv_bias=True, + proj_bias=True, + flatten_batch=True, + quant_config=quant_config, + prefix=add_prefix("attn", prefix), + ) + self.mlp = MuseGlimmerVisionMLP( + config, quant_config=quant_config, prefix=add_prefix("mlp", prefix) + ) + + def forward( + self, + hidden_states: torch.Tensor, + cu_seqlens: torch.Tensor, + position_embeddings: Tuple[torch.Tensor, torch.Tensor], + forward_metadata: VisionAttentionMetadata, + ) -> torch.Tensor: + hidden_states = hidden_states + self.attn( + self.norm1(hidden_states), + cu_seqlens=cu_seqlens, + position_embeddings=position_embeddings, + forward_metadata=forward_metadata, + ) + return hidden_states + self.mlp(self.norm2(hidden_states)) + + +class MuseGlimmerVisionModel(nn.Module): + + def __init__( + self, + config, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + self.config = config + self.merge_size = config.merge_size + self.patch_size = config.patch_size + # Square position grid; the window is the full extent of that grid. + self.window_size = config.pos_emb_height * config.patch_size + self.attention_types = config.attention_types + + self.patch_embedder = MuseGlimmerVisionPatchEmbedder( + config, + quant_config=quant_config, + prefix=add_prefix("patch_embedder", prefix), + ) + self.rotary_emb = MuseGlimmerVisionRotaryEmbedding( + config.hidden_size // config.num_attention_heads, config.rope_theta + ) + self.ln_pre = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.layers = nn.ModuleList( + [ + MuseGlimmerVisionEncoderLayer( + config, + quant_config=quant_config, + prefix=add_prefix(f"layers.{i}", prefix), + ) + for i in range(config.num_hidden_layers) + ] + ) + self.ln_post = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + + @property + def dtype(self) -> torch.dtype: + return self.ln_pre.weight.dtype + + @property + def device(self) -> torch.device: + return self.ln_pre.weight.device + + def pixel_shuffle( + self, hidden_states: torch.Tensor, grid_thw: torch.Tensor + ) -> torch.Tensor: + """Fold each ``merge_size x merge_size`` patch block into one wider token.""" + factor = self.merge_size + dim = hidden_states.shape[-1] + outputs = [] + offset = 0 + + for t, h, w in grid_thw.tolist(): + t, h, w = int(t), int(h), int(w) + num_tokens = t * h * w + chunk = hidden_states[offset : offset + num_tokens] + offset += num_tokens + + num_out_per_frame = (h // factor) * (w // factor) + block_perm = torch.arange(h * w, device=hidden_states.device) + block_perm = ( + block_perm.view(h // factor, factor, w // factor, factor) + .permute(0, 2, 1, 3) + .reshape(-1) + ) + if t > 1: + frame_offsets = ( + torch.arange(t, device=hidden_states.device) * h * w + ).view(t, 1) + block_perm = (block_perm.unsqueeze(0) + frame_offsets).reshape(-1) + + merged = chunk[block_perm].view(t * num_out_per_frame, factor * factor, dim) + merged = ( + merged.permute(0, 2, 1) + .contiguous() + .view(t * num_out_per_frame, dim * factor * factor) + ) + outputs.append(merged) + + return torch.cat(outputs, dim=0) + + def forward( + self, pixel_values: torch.Tensor, grid_thw: torch.Tensor + ) -> torch.Tensor: + pixel_values = pixel_values.to(device=self.device, dtype=self.dtype) + hidden_states = self.patch_embedder(pixel_values, grid_thw) + hidden_states = self.ln_pre(hidden_states) + + cu_seqlens = get_vision_cu_seqlens(grid_thw).to(self.device) + window_index, cu_window_seqlens = get_vision_window_index( + grid_thw, + spatial_merge_size=1, + window_size=self.window_size, + patch_size=self.patch_size, + ) + window_index = window_index.to(self.device) + cu_window_seqlens = cu_window_seqlens.to(self.device) + + # The reference offsets patch coordinates by 1 and orders them (w, h). + position_ids = get_vision_position_ids(grid_thw, spatial_merge_size=1) + position_ids = (position_ids.flip(-1) + 1).to(self.device) + + hidden_states = hidden_states[window_index] + cos, sin = self.rotary_emb(position_ids[window_index]) + position_embeddings = ( + cos.to(self.dtype), + sin.to(self.dtype), + ) + + metadata_by_type = { + "full_attention": prepare_vision_attention_metadata( + cu_seqlens, device=self.device + ), + "window_attention": prepare_vision_attention_metadata( + cu_window_seqlens, device=self.device + ), + } + cu_seqlens_by_type = { + "full_attention": cu_seqlens, + "window_attention": cu_window_seqlens, + } + + hidden_states = hidden_states.unsqueeze(0) + for i, layer in enumerate(self.layers): + attention_type = self.attention_types[i] + hidden_states = layer( + hidden_states, + cu_seqlens=cu_seqlens_by_type[attention_type], + position_embeddings=position_embeddings, + forward_metadata=metadata_by_type[attention_type], + ) + hidden_states = hidden_states.squeeze(0) + + hidden_states = hidden_states[permute_inv(window_index)] + hidden_states = self.ln_post(hidden_states) + return self.pixel_shuffle(hidden_states, grid_thw) + + +class MuseGlimmerVisionAdapter(nn.Module): + """Two-layer projector between the shuffled ViT output and the decoder width.""" + + def __init__( + self, + config, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + self.fc1 = ColumnParallelLinear( + config.out_hidden_size, + config.projector_hidden_size, + bias=False, + quant_config=quant_config, + prefix=add_prefix("fc1", prefix), + ) + self.fc2 = RowParallelLinear( + config.projector_hidden_size, + config.projector_hidden_size, + bias=False, + quant_config=quant_config, + prefix=add_prefix("fc2", prefix), + ) + self.act = ACT2FN[config.projector_hidden_act] + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x, _ = self.fc1(x) + x = self.act(x) + x, _ = self.fc2(x) + return self.act(x) + + +class MuseGlimmerForCausalLM(nn.Module): + packed_modules_mapping = { + "qkv_proj": ["q_proj", "k_proj", "v_proj"], + "gate_up_proj": ["gate_proj", "up_proj"], + } + + checkpoint_uses_vendor_names = False + + builds_vision_tower = False + + def __init__( + self, + config, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + self.config = config + self.quant_config = quant_config + self.checkpoint_norms_are_absolute = ( + quant_config is not None and quant_config.get_name() == "gguf" + ) + + if config.output_soft_cap_temp is not None: + config.final_logit_softcapping = config.output_soft_cap_temp + + self.model = MuseGlimmerModel( + config, quant_config=quant_config, prefix=add_prefix("model", prefix) + ) + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + # Pad vocab to 128-align for the MXFP8 kernel. + padding_size=128, + quant_config=quant_config, + prefix=add_prefix("lm_head", prefix), + ) + self.logits_processor = LogitsProcessor( + config, logit_scale=config.output_multiplier + ) + self.capture_aux_hidden_states = False + + def get_attention_sliding_window_size(self): + return get_attention_sliding_window_size(self.config) + + def get_input_embeddings(self) -> nn.Embedding: + return self.model.embed_tokens + + @torch.no_grad() + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: Optional[torch.Tensor] = None, + ): + hidden_states = self.model(input_ids, positions, forward_batch, input_embeds) + + aux_hidden_states = None + if self.capture_aux_hidden_states: + hidden_states, aux_hidden_states = hidden_states + + return self.logits_processor( + input_ids, hidden_states, self.lm_head, forward_batch, aux_hidden_states + ) + + def set_dflash_layers_to_capture(self, layer_ids: List[int]): + if layer_ids is None: + raise ValueError( + "DFLASH requires explicit layer_ids for aux hidden capture." + ) + num_layers = len(self.model.layers) + bad = [i for i in layer_ids if not 0 <= i < num_layers] + if bad: + raise ValueError( + f"DFLASH target layer ids {bad} are out of range for a " + f"{num_layers}-layer Muse Glimmer target." + ) + self.capture_aux_hidden_states = True + self.model.layers_to_capture = list(layer_ids) + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + ("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), + ] + params_dict = dict(self.named_parameters()) + loaded = 0 + + for name, loaded_weight in weights: + # SGLang derives RoPE itself; the checkpoint ships a cached freqs buffer. + if "rotary_emb.freqs" in name or "rotary_emb.inv_freq" in name: + continue + if not self.builds_vision_tower and any( + fragment in name for fragment in _VISION_NAME_FRAGMENTS + ): + continue + + if self.checkpoint_uses_vendor_names: + name = _vendor_weight_name(name) + + # Fold the reference's +1; model.norm is exempt. + if not self.checkpoint_norms_are_absolute and name.endswith( + _OFFSET_NORM_SUFFIXES + ): + loaded_weight = loaded_weight + 1.0 + + # ModelOpt holds KV scales on self_attn, not self_attn.attn. + if name.endswith((".k_scale", ".v_scale")): + remapped = maybe_remap_kv_scale_name(name, params_dict) + if remapped is None: + continue + name = remapped + + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + if "output_gate_proj" in name: + continue + mapped = name.replace(weight_name, param_name) + if mapped not in params_dict: + continue + param = params_dict[mapped] + param.weight_loader(param, loaded_weight, shard_id) + loaded += 1 + break + else: + if name not in params_dict: + logger.warning( + "Muse Glimmer: unexpected checkpoint weight %s", name + ) + continue + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + loaded += 1 + + logger.info("Muse Glimmer: loaded %d weight tensors", loaded) + + +class MuseGlimmerForConditionalGeneration(MuseGlimmerForCausalLM): + """Vendor multimodal HF export: the MuseGlimmerForCausalLM decoder plus the image tower.""" + + checkpoint_uses_vendor_names = True + builds_vision_tower = True + + def __init__( + self, + config, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__(config, quant_config=quant_config, prefix=prefix) + self.language_model_only = bool(config.language_model_only) + if self.language_model_only: + self.builds_vision_tower = False + self.vision_tower = None + self.vision_adapter = None + self.vision_projection = None + self.perception_emb_norm = None + return + if config.vision_config is None: + raise ValueError( + "MuseGlimmerForConditionalGeneration needs a vision_config; a checkpoint " + "with no vision tower should declare MuseGlimmerForCausalLM instead." + ) + self.vision_tower = MuseGlimmerVisionModel( + config.vision_config, + quant_config=quant_config, + prefix=add_prefix("vision_tower", prefix), + ) + self.vision_adapter = MuseGlimmerVisionAdapter( + config, + quant_config=quant_config, + prefix=add_prefix("vision_adapter", prefix), + ) + self.vision_projection = ReplicatedLinear( + config.projector_hidden_size, + config.hidden_size, + bias=False, + quant_config=quant_config, + prefix=add_prefix("vision_projection", prefix), + ) + self.perception_emb_norm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps, has_weight=False + ) + + def pad_input_ids(self, input_ids: List[int], mm_inputs: MultimodalInputs): + return MultiModalityDataPaddingPatternMultimodalTokens().pad_input_tokens( + input_ids, mm_inputs + ) + + def get_image_feature(self, items: List[MultimodalDataItem]) -> torch.Tensor: + pixel_values = torch.cat([item.feature for item in items], dim=0) + image_grid_thw = torch.cat([item.image_grid_thw for item in items], dim=0).cpu() + assert pixel_values.dim() == 2, pixel_values.dim() + assert image_grid_thw.dim() == 2, image_grid_thw.dim() + + features = self.vision_tower(pixel_values, image_grid_thw) + features = self.vision_adapter(features) + features, _ = self.vision_projection(features) + return self.perception_emb_norm(features) + + @torch.no_grad() + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: Optional[torch.Tensor] = None, + ): + if self.language_model_only: + # Embed here: prefill BCG captures input_embeds. + if input_embeds is None: + input_embeds = self.get_input_embeddings()(input_ids) + if forward_batch.input_embeds is not None: + forward_batch.input_embeds.copy_(input_embeds) + input_embeds = forward_batch.input_embeds + hidden_states = self.model(None, positions, forward_batch, input_embeds) + else: + hidden_states = general_mm_embed_routine( + input_ids=input_ids, + forward_batch=forward_batch, + language_model=self.model, + multimodal_model=self, + positions=positions, + ) + + aux_hidden_states = None + if self.capture_aux_hidden_states: + hidden_states, aux_hidden_states = hidden_states + + return self.logits_processor( + input_ids, hidden_states, self.lm_head, forward_batch, aux_hidden_states + ) + + +EntryClass = [MuseGlimmerForCausalLM, MuseGlimmerForConditionalGeneration] diff --git a/python/sglang/srt/multimodal/processors/muse_glimmer.py b/python/sglang/srt/multimodal/processors/muse_glimmer.py new file mode 100644 index 000000000000..36d78bb3b527 --- /dev/null +++ b/python/sglang/srt/multimodal/processors/muse_glimmer.py @@ -0,0 +1,61 @@ +# 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 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""SGLang multimodal processor for Muse Glimmer (images).""" + +from typing import Dict, List, Union + +from sglang.srt.managers.schedule_batch import MultimodalProcessorOutput +from sglang.srt.models.muse_glimmer import MuseGlimmerForConditionalGeneration +from sglang.srt.multimodal.processors.base_processor import ( + BaseMultimodalProcessor as SGLangBaseProcessor, +) +from sglang.srt.multimodal.processors.base_processor import ( + MultimodalSpecialTokens, +) + + +class MuseGlimmerMultimodalProcessor(SGLangBaseProcessor): + models = [MuseGlimmerForConditionalGeneration] + + def __init__(self, hf_config, server_args, _processor, *args, **kwargs): + super().__init__(hf_config, server_args, _processor, *args, **kwargs) + + self.image_token_id = _processor.image_token_id + + self.mm_tokens = MultimodalSpecialTokens( + image_token=_processor.image_token, + image_token_id=self.image_token_id, + ).build(_processor) + + async def process_mm_data_async( + self, + image_data: List[Union[str, bytes, Dict]], + input_text, + request_obj, + *args, + **kwargs, + ): + base_output = await self.load_mm_data( + prompt=input_text, + image_data=image_data, + multimodal_tokens=self.mm_tokens, + ) + mm_items, input_ids, _ = await self.process_and_combine_mm_data_async( + base_output, self.mm_tokens + ) + return MultimodalProcessorOutput( + input_ids=input_ids.tolist(), + mm_items=mm_items, + im_token_id=self.image_token_id, + ) diff --git a/python/sglang/srt/parser/reasoning_parser.py b/python/sglang/srt/parser/reasoning_parser.py index 36bd69cf849d..215208a727c4 100644 --- a/python/sglang/srt/parser/reasoning_parser.py +++ b/python/sglang/srt/parser/reasoning_parser.py @@ -1614,6 +1614,206 @@ def finish(self) -> StreamingParseResult: return self._maybe_apply_force_nonempty_content(ret) +class MuseGlimmerDetector(BaseReasoningFormatDetector): + """Detector for Muse Glimmer's recipient-channel format. + + The chat template ends the generation prompt at ``<|start|>assistant`` with no + recipient and no ``<|message|>``, so the model itself emits the channel header as + ordinary text. A full turn looks like:: + + " to=self<|message|>" "<|eom|>" + "<|start|>assistant to=user<|message|>" "<|eot|>" + + Reasoning is the ``to=self`` channel; the answer is ``to=user``. Any other recipient + is a tool call (``to=functions.get_weather``), whose body is an ATEM block that must + reach the function-call detector with its markers intact — so those channels are + emitted as normal text including their header, following GptOssDetector's precedent + of preserving raw structural text for tool calls. + + When a tool-call parser consumes this detector's normal text + (``tool_call_parser_active=True``), the ``to=user`` channel keeps its framing too, + so the downstream detector sees every channel boundary and can tell a real tool + channel from one merely *quoted* inside the answer — unwrapping here would make a + quoted ``<|start|>assistant to=<|message|>`` indistinguishable from a real + header and turn quoted markup into a live call. The tool detector unwraps + ``to=user`` itself, so nothing framed leaks to the client. Non-streaming + additionally requires that a turn *without* any ATEM block come out unwrapped, + because serving bypasses the tool detector entirely when ``has_tool_call()`` is + false — hence the ATEM-presence branch in ``detect_and_parse``, mirroring the + vendor's reference reasoning parser. + + Keying on ``<|message|>`` rather than the literal " to=self" mirrors the vendor's own + reference implementation (which slices past the last ``<|message|>`` token), + and is robust to the header varying with + the recipient. It does require the delimiters to survive detokenization, which is why + ``muse`` is registered in ``_patch_reasoning_skip_special_tokens``. + + A single channel may also be cut short by the token cap, in which case there is no + terminator and the partial body is still attributed to whichever channel was open. + """ + + MESSAGE = "<|message|>" + EOM = "<|eom|>" + EOT = "<|eot|>" + START = "<|start|>" + _MAX_MARKER = max(len(MESSAGE), len(EOM), len(EOT), len(START)) + _RECIPIENT_RE = re.compile(r"to=([^\s<]+)") + + # ATEM markers that identify a tool-call turn in the non-reasoning remainder. + _ATEM_MARKERS = ("") + + def __init__( + self, + stream_reasoning: bool = True, + force_reasoning: bool = False, + continue_final_message: bool = False, + previous_content: str = "", + force_nonempty_content: bool = False, + tool_call_parser_active: bool = False, + ): + super().__init__( + " to=self" + self.MESSAGE, + self.EOM, + force_reasoning=force_reasoning, + stream_reasoning=stream_reasoning, + continue_final_message=continue_final_message, + previous_content=previous_content, + force_nonempty_content=force_nonempty_content, + ) + self._recipient: Optional[str] = None + self._in_body = False + self._pending_reasoning = "" + self._tool_call_parser_active = tool_call_parser_active + self._saw_reasoning_block = False + + def _sink(self, recipient: Optional[str]) -> str: + return "reasoning" if recipient == "self" else "normal" + + @classmethod + def _partial_marker_len(cls, buf: str) -> int: + """Length of the longest suffix of ``buf`` that could still become a marker. + + Returns 0 when nothing is held back, so ordinary text streams out immediately + instead of waiting for a terminator that may never arrive. + """ + markers = (cls.EOM, cls.EOT, cls.START) + for k in range(min(len(buf), cls._MAX_MARKER - 1), 0, -1): + tail = buf[-k:] + if any(m.startswith(tail) for m in markers): + return k + return 0 + + def _consume(self, flush: bool, preserve_channels: bool = False) -> Tuple[str, str]: + """Drain self._buffer into (reasoning, normal). + + With flush=False, holds back a short tail that could be the prefix of a marker + split across chunk boundaries; with flush=True, emits everything. + + With preserve_channels=True, the ``to=user`` channel keeps its header and + terminator like tool channels do (see the class docstring for why the + function-call detector needs the framing intact); reasoning is always + extracted and never framed. + """ + reasoning_parts: List[str] = [] + normal_parts: List[str] = [] + + while self._buffer: + if not self._in_body: + idx = self._buffer.find(self.MESSAGE) + if idx == -1: + if flush: + normal_parts.append(self._buffer) + self._buffer = "" + break + header = self._buffer[:idx] + m = self._RECIPIENT_RE.search(header) + self._recipient = m.group(1) if m else "user" + self._buffer = self._buffer[idx + len(self.MESSAGE) :] + self._in_body = True + if self._sink(self._recipient) == "reasoning": + if self._saw_reasoning_block: + reasoning_parts.append("\n") + self._saw_reasoning_block = True + elif self._recipient != "user" or preserve_channels: + # Keep the header so the function-call detector sees it. + normal_parts.append(header + self.MESSAGE) + continue + + end_idx, end_tok = -1, "" + for tok in (self.EOM, self.EOT): + i = self._buffer.find(tok) + if i != -1 and (end_idx == -1 or i < end_idx): + end_idx, end_tok = i, tok + + if end_idx != -1: + body = self._buffer[:end_idx] + self._buffer = self._buffer[end_idx + len(end_tok) :] + self._in_body = False + if self._sink(self._recipient) == "reasoning": + reasoning_parts.append(body) + else: + normal_parts.append(body) + if self._recipient != "user" or preserve_channels: + normal_parts.append(end_tok) + self._recipient = None + continue + + # Hold back only a genuine marker prefix. + if flush: + body, self._buffer = self._buffer, "" + else: + keep = self._partial_marker_len(self._buffer) + if keep == len(self._buffer): + break + body = self._buffer[: len(self._buffer) - keep] + self._buffer = self._buffer[len(self._buffer) - keep :] + if not body: + break + if self._sink(self._recipient) == "reasoning": + reasoning_parts.append(body) + else: + normal_parts.append(body) + + return "".join(reasoning_parts), "".join(normal_parts) + + def detect_and_parse(self, text: str) -> StreamingParseResult: + self._buffer += text + raw = self._buffer + reasoning, normal = self._consume(flush=True) + if self._tool_call_parser_active and any( + m in normal for m in self._ATEM_MARKERS + ): + self._buffer = raw + self._recipient = None + self._in_body = False + self._saw_reasoning_block = False + reasoning, normal = self._consume(flush=True, preserve_channels=True) + return self._maybe_apply_force_nonempty_content( + StreamingParseResult(normal_text=normal, reasoning_text=reasoning) + ) + + def parse_streaming_increment(self, new_text: str) -> StreamingParseResult: + self._buffer += new_text + reasoning, normal = self._consume( + flush=False, preserve_channels=self._tool_call_parser_active + ) + if not self.stream_reasoning: + self._pending_reasoning += reasoning + reasoning = "" + if not self._in_body and self._pending_reasoning: + reasoning, self._pending_reasoning = self._pending_reasoning, "" + return StreamingParseResult(normal_text=normal, reasoning_text=reasoning) + + def finish(self) -> StreamingParseResult: + reasoning, normal = self._consume( + flush=True, preserve_channels=self._tool_call_parser_active + ) + if self._pending_reasoning: + reasoning = self._pending_reasoning + reasoning + self._pending_reasoning = "" + return StreamingParseResult(normal_text=normal, reasoning_text=reasoning) + + class ReasoningParser: """ Parser that handles both streaming and non-streaming scenarios for extracting @@ -1623,6 +1823,10 @@ class ReasoningParser: model_type (str): Type of model to parse reasoning from stream_reasoning (bool): If False, accumulates reasoning content until complete. If True, streams reasoning content as it arrives. + tool_call_parser_active (bool): True when this parser's normal text feeds a + function-call parser rather than going straight to the client. Passed on + to detectors that accept it (channel-framed formats keep tool framing + intact for the downstream detector). """ DetectorMap: Dict[str, Type[BaseReasoningFormatDetector]] = { @@ -1637,6 +1841,7 @@ class ReasoningParser: "kimi_k2": KimiK2Detector, "kimi_k3": KimiK3Detector, "mimo": _MimoDetector, + "muse": MuseGlimmerDetector, "poolside_v1": _PoolsideV1Detector, "qwen3": Qwen3Detector, "qwen3-thinking": Qwen3Detector, @@ -1660,6 +1865,7 @@ def __init__( force_reasoning: Optional[bool] = None, request: ChatCompletionRequest = None, tokenizer=None, + tool_call_parser_active: bool = False, ): if not model_type: raise ValueError("Model type must be specified") @@ -1705,6 +1911,11 @@ def __init__( if "tokenizer" in sig.parameters: kwargs["tokenizer"] = tokenizer + if tool_call_parser_active: + sig = inspect.signature(detector_class) + if "tool_call_parser_active" in sig.parameters: + kwargs["tool_call_parser_active"] = True + self.detector = detector_class(**kwargs) def parse_non_stream(self, full_text: str) -> Tuple[Optional[str], Optional[str]]: diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 3cd4f8cec06e..8904d6f4b167 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -1746,6 +1746,7 @@ class ServerArgs: help="Choose the runner backend for NVFP4 GEMM operations. Options: 'auto' (default; selects flashinfer_cutedsl on SM100, marlin on SM80-SM90, flashinfer_cutlass otherwise (including SM120)), 'flashinfer_cutlass' (FlashInfer CUTLASS backend), 'flashinfer_cudnn' (FlashInfer cuDNN backend, optimal on CUDA 13+ with cuDNN 9.15+), 'flashinfer_cutedsl' (FlashInfer CuTe DSL backend), 'flashinfer_trtllm' (FlashInfer TensorRT-LLM backend, requires different weight preparation with shuffling), 'marlin' (weight-only W4A16 fallback for SM80+). ", cli_name="--fp4-gemm-backend", choices=FP4_GEMM_RUNNER_BACKEND_CHOICES, + resolvable=True, ), NS("exec.kernel"), ] = "auto" @@ -2157,6 +2158,20 @@ class ServerArgs: "Attention backend for speculative decoding drafting.", NS("spec"), ] = None + speculative_draft_kv_cache_dtype: A[ + Optional[str], + Arg( + help="KV cache dtype for the speculative draft model only. The draft pool is " + "allocated with one slot per target token (draft and target share a slot index " + "space), so for a small draft it can still rival the target pool: a 5-layer " + "DFLASH draft costs 10240 bytes/token in bf16. Setting fp8_e4m3 halves the draft " + "pool; the saving shows up as free device memory, so raise " + "--mem-fraction-static to convert it into KV capacity. Default follows " + "--kv-cache-dtype.", + choices=["auto", "fp8_e5m2", "fp8_e4m3", "bf16", "bfloat16"], + ), + NS("spec"), + ] = None speculative_draft_window_size: A[ Optional[int], "Sliding window size for the draft model. Honored by Llama EAGLE-3 (`LlamaForCausalLMEagle3`) and DFLASH only; other EAGLE-3 backends (e.g. MLA-based drafters) silently ignore it. For Llama EAGLE-3, the drafter only attends to the most recent N keys (verifier hidden states + its own outputs); the verifier is unaffected. For DFLASH, the draft worker keeps a recent target-token window in its local KV cache (paged backends may retain up to one extra page on the left for alignment). Default is full attention/context.", @@ -3075,6 +3090,14 @@ class ServerArgs: language_only: A[ bool, "For VLM, load weights for the language model only.", NS("disagg") ] = False + language_model_only: A[ + bool, + "Skip the multimodal encoder entirely: its weights are never loaded and the " + "tower is never built, freeing that GPU memory for KV cache. Multimodal " + "requests are rejected. Unlike --language-only this is a standalone mode, " + "not part of encoder/decoder disaggregation.", + NS("disagg"), + ] = False encoder_transfer_backend: A[ str, Arg( @@ -3530,6 +3553,7 @@ def __post_init__(self): # resolution (the declarative registry materializes too late to affect # it). Inkling opts into full-graph prefill capture here. self._apply_inkling_prefill_cuda_graph_default() + self._apply_muse_glimmer_prefill_cuda_graph_max_bs_default() # must run before _handle_cuda_graph_config and _handle_data_parallelism self._handle_dwdp() @@ -3819,6 +3843,8 @@ def _handle_model_capability_adjustments(self): def _handle_model_source_paths(self): """Prepare metadata for model paths backed by remote object stores.""" + self._resolve_hf_gguf_model_path() + seen_paths = set() for model_path in ( self.model_path, @@ -4258,6 +4284,16 @@ def _apply_inkling_prefill_cuda_graph_default(self): ): self.cuda_graph_backend_prefill = Backend.FULL + def _apply_muse_glimmer_prefill_cuda_graph_max_bs_default(self): + if ( + self.cuda_graph_max_bs_prefill is not None + or parse_connector_type(self.model_path) == ConnectorType.INSTANCE + ): + return + arch = self.get_model_config().hf_config.architectures[0] + if arch in ("MuseGlimmerForCausalLM", "MuseGlimmerForConditionalGeneration"): + self.cuda_graph_max_bs_prefill = 512 + def _handle_cuda_graph_config(self): from sglang.srt.arg_groups.kimi_k3_hook import disable_kimi_k3_symm_mem @@ -4816,6 +4852,7 @@ def _handle_gpu_memory_settings(self, gpu_mem): if ( model_config.is_multimodal and not self.language_only + and not self.language_model_only and self.disaggregation_mode != "decode" ): self.adjust_mem_fraction_for_vlm(model_config) @@ -6347,7 +6384,11 @@ def _handle_context_parallelism(self): raise ValueError( "MiMo V2 CP-v2 only supports --cp-strategy zigzag." ) - if model_config.is_multimodal and not self.language_only: + if ( + model_config.is_multimodal + and not self.language_only + and not self.language_model_only + ): raise ValueError( "MiMo V2 CP-v2 only supports text inference; add " "--language-only." @@ -7275,6 +7316,18 @@ def _resolve_storage_layout_compatibility(self): f"switching to {new_layout} layout for {self.hicache_io_backend} io backend" ) + def _resolve_hf_gguf_model_path(self): + """Turn a Hub reference to a .gguf into a local file path.""" + from sglang.srt.utils.hf_transformers_utils import resolve_hf_gguf_reference + + resolved = resolve_hf_gguf_reference(self.model_path, revision=self.revision) + if resolved is None: + return + logger.info("Resolved GGUF %s -> %s", self.model_path, resolved) + if self.tokenizer_path == self.model_path: + self.tokenizer_path = resolved + self.model_path = resolved + def _handle_load_format(self): # The quantization side of the gguf coupling moved to the pipeline # (arg_groups/overrides.py: _gguf_quantization); load_format itself is @@ -7429,7 +7482,39 @@ def _check_format(has_params, has_consolidated, has_hf_weights) -> bool: except Exception: return False + LANGUAGE_MODEL_ONLY_ARCHITECTURES = ("MuseGlimmerForConditionalGeneration",) + + def _handle_language_model_only(self): + if not self.language_model_only: + return + for flag, name in ( + (self.encoder_only, "--encoder-only"), + (self.language_only, "--language-only"), + (self.enable_prefix_mm_cache, "--enable-prefix-mm-cache"), + ( + self.enable_broadcast_mm_inputs_process, + "--enable-broadcast-mm-inputs-process", + ), + (self.mm_enable_dp_encoder, "--mm-enable-dp-encoder"), + ): + if flag: + raise ValueError( + f"--language-model-only cannot be combined with {name}" + ) + if self.disaggregation_mode != "null": + raise ValueError( + "--language-model-only is incompatible with --disaggregation-mode " + "prefill/decode" + ) + architectures = self.get_model_config().hf_config.architectures + if not any(a in self.LANGUAGE_MODEL_ONLY_ARCHITECTURES for a in architectures): + raise ValueError( + f"--language-model-only does not support {architectures}. " + f"Supported: {list(self.LANGUAGE_MODEL_ONLY_ARCHITECTURES)}." + ) + def _handle_encoder_disaggregation(self): + self._handle_language_model_only() if self.enable_prefix_mm_cache and not self.encoder_only: raise ValueError( "--enable-prefix-mm-cache requires --encoder-only to be enabled" diff --git a/python/sglang/srt/speculative/dflash_utils.py b/python/sglang/srt/speculative/dflash_utils.py index e4c13a6bc342..b6a6132c78ad 100644 --- a/python/sglang/srt/speculative/dflash_utils.py +++ b/python/sglang/srt/speculative/dflash_utils.py @@ -513,7 +513,9 @@ def parse_dflash_draft_config(*, draft_hf_config: Any) -> DFlashDraftConfig: f"Got len(target_layer_ids)={len(parsed_target_layer_ids)}." ) - mask_token = dflash_cfg.get("mask_token", None) + mask_token = dflash_cfg.get( + "mask_token", _cfg_get(draft_hf_config, "mask_token", None) + ) if mask_token is None: mask_token = DEFAULT_DFLASH_MASK_TOKEN if not isinstance(mask_token, str) or not mask_token: @@ -522,7 +524,9 @@ def parse_dflash_draft_config(*, draft_hf_config: Any) -> DFlashDraftConfig: f"got {mask_token!r}." ) - mask_token_id = dflash_cfg.get("mask_token_id", None) + mask_token_id = dflash_cfg.get( + "mask_token_id", _cfg_get(draft_hf_config, "mask_token_id", None) + ) if mask_token_id is not None: if not isinstance(mask_token_id, Integral) or isinstance(mask_token_id, bool): raise ValueError( diff --git a/python/sglang/srt/speculative/dflash_worker_v2.py b/python/sglang/srt/speculative/dflash_worker_v2.py index d8d94b1d4c0f..d0c3888aed33 100644 --- a/python/sglang/srt/speculative/dflash_worker_v2.py +++ b/python/sglang/srt/speculative/dflash_worker_v2.py @@ -79,6 +79,14 @@ def _get_fused_kv_materialize_helper(): return _FusedKVMaterializeHelper +# is_floating_point() is True for fp8; list dtypes explicitly. +_DENSE_HEAD_DTYPES = (torch.float16, torch.bfloat16, torch.float32) + + +def _is_dense_head_weight(weight) -> bool: + return weight is not None and weight.dtype in _DENSE_HEAD_DTYPES + + class _DflashDraftSampler: """Capture-safe greedy argmax over the target LM head, run inside the draft cuda graph so the draft sampling is captured and counted in fwd_occupancy. @@ -231,6 +239,12 @@ def __init__( mask_token=self._mask_token, mask_token_id=self._mask_token_id_override, ) + target_model = self._target_worker.model_runner.model + self._noise_embed_scale = ( + float(target_model.get_dflash_noise_embedding_scale()) + if hasattr(target_model, "get_dflash_noise_embedding_scale") + else 1.0 + ) if self.ps.tp_rank == 0: logger.info( "Initialized DFLASH draft runner. attention_backend=%s, model=%s, block_size=%s, draft_window_size=%s, compact_cache=%s", @@ -241,10 +255,11 @@ def __init__( self.use_compact_draft_cache, ) logger.info( - "DFLASH draft runner ready. mask_token=%s, mask_token_id=%s, mask_token_id_override=%s", + "DFLASH draft runner ready. mask_token=%s, mask_token_id=%s, mask_token_id_override=%s, noise_embed_scale=%s", self._mask_token, self._mask_token_id, self._mask_token_id_override, + self._noise_embed_scale, ) self._block_pos_offsets = build_block_pos_offsets( @@ -375,9 +390,11 @@ def _eager(reason): return _eager("block_size<=1") target_model = self._target_worker.model_runner.model lm_head = getattr(target_model, "lm_head", None) - if lm_head is None or not hasattr(lm_head, "weight"): + if lm_head is None: return _eager("no target lm_head") - if not torch.is_floating_point(lm_head.weight): + if not hasattr(lm_head, "weight"): + return _eager("quantized lm_head has no dense weight") + if not _is_dense_head_weight(lm_head.weight): # Quantized lm_head (FP8/INT) would break the static matmul. return _eager("quantized lm_head") tp_group = get_tp_group() @@ -805,6 +822,42 @@ def _resolve_mask_token_id( return int(resolved_id) + def _greedy_sample_from_quantized_head( + self, + *, + hidden_states: torch.Tensor, + lm_head, + chunk_size: int, + ) -> torch.Tensor: + """Greedy argmax over a target LM head that has no dense ``weight``. + + A GGUF head stores packed ``qweight`` plus a type tag, so the dense path's + ``weight[:num_org]`` slicing has nothing to slice. Logits come from the + layer's own kernel instead -- the same call ``LogitsProcessor._get_logits`` + makes for GGUF models. Padding rows are excluded so argmax cannot return + an id outside the real vocabulary. + """ + tp_size = int(get_tp_group().world_size) + if tp_size != 1: + raise RuntimeError( + "DFLASH with a quantized target lm_head is only supported at " + f"tp=1, got tp_size={tp_size}." + ) + + num_tokens = int(hidden_states.shape[0]) + out_tokens = torch.empty( + (num_tokens,), dtype=torch.long, device=hidden_states.device + ) + num_org = int(getattr(lm_head, "org_vocab_size", 0)) or None + + for start in range(0, num_tokens, int(chunk_size)): + end = min(num_tokens, start + int(chunk_size)) + logits = lm_head.quant_method.apply(lm_head, hidden_states[start:end], None) + if num_org is not None and logits.shape[-1] > num_org: + logits = logits[:, :num_org] + out_tokens[start:end] = torch.argmax(logits, dim=-1).to(torch.long) + return out_tokens + def _greedy_sample_from_vocab_parallel_head( self, *, @@ -822,6 +875,11 @@ def _greedy_sample_from_vocab_parallel_head( if hidden_states.numel() == 0: return torch.empty((0,), dtype=torch.long, device=hidden_states.device) + if not _is_dense_head_weight(getattr(lm_head, "weight", None)): + return self._greedy_sample_from_quantized_head( + hidden_states=hidden_states, lm_head=lm_head, chunk_size=chunk_size + ) + weight = lm_head.weight # [local_vocab_padded, hidden] weight_dtype = weight.dtype num_tokens = int(hidden_states.shape[0]) @@ -1487,9 +1545,13 @@ def forward_batch_generation( target_model = self.target_worker.model_runner.model embed_module = target_model.get_input_embeddings() lm_head = getattr(target_model, "lm_head", None) - if lm_head is None or not hasattr(lm_head, "weight"): + if lm_head is None or not ( + hasattr(lm_head, "weight") + or callable(getattr(getattr(lm_head, "quant_method", None), "apply", None)) + ): raise RuntimeError( - "DFLASH requires the target model to expose `lm_head` with `weight`." + "DFLASH requires the target model to expose `lm_head` with either " + "`weight` or a `quant_method` that can produce logits." ) block_size = int(self.block_size) @@ -1562,6 +1624,8 @@ def forward_batch_generation( verify_out_cache_loc_2d.copy_(verify_out_cache_loc.view(bs, block_size)) noise_embedding = embed_module(block_ids) + if self._noise_embed_scale != 1.0: + noise_embedding = noise_embedding * self._noise_embed_scale input_embeds = noise_embedding.view(-1, noise_embedding.shape[-1]) positions = positions_2d.reshape(-1) diff --git a/python/sglang/srt/speculative/draft_worker_common.py b/python/sglang/srt/speculative/draft_worker_common.py index b393ce8d03b2..28a150f7375f 100644 --- a/python/sglang/srt/speculative/draft_worker_common.py +++ b/python/sglang/srt/speculative/draft_worker_common.py @@ -101,6 +101,10 @@ def build_draft_tp_worker( draft_model_runner = draft_worker.model_runner draft_worker.draft_runner = draft_model_runner + + # DFlash drafts have no vocab; borrow the target's. + if draft_model_runner.model_config.vocab_size is None: + draft_model_runner.model_config.vocab_size = target_model_config.vocab_size return DraftWorkerBundle( draft_worker=draft_worker, draft_model_runner=draft_model_runner, diff --git a/python/sglang/srt/utils/hf_transformers/__init__.py b/python/sglang/srt/utils/hf_transformers/__init__.py index 3e6b3fa78845..a6b48c93c617 100644 --- a/python/sglang/srt/utils/hf_transformers/__init__.py +++ b/python/sglang/srt/utils/hf_transformers/__init__.py @@ -35,6 +35,7 @@ get_rope_config, get_sparse_attention_config, get_tokenizer_from_processor, + resolve_hf_gguf_reference, ) from .config import get_config from .processor import get_processor @@ -51,6 +52,7 @@ "_fix_v5_add_bos_eos_token", "attach_additional_stop_token_ids", "check_gguf_file", + "resolve_hf_gguf_reference", "download_from_hf", "get_config", "get_context_length", diff --git a/python/sglang/srt/utils/hf_transformers/common.py b/python/sglang/srt/utils/hf_transformers/common.py index db584861be60..efc8dee39f64 100644 --- a/python/sglang/srt/utils/hf_transformers/common.py +++ b/python/sglang/srt/utils/hf_transformers/common.py @@ -52,6 +52,8 @@ MiniCPMV4_6VisionConfig, MiniMaxM3VLConfig, MultiModalityConfig, + MuseGlimmerAssistantConfig, + MuseGlimmerConfig, NemotronH_Nano_Omni_Reasoning_V3_Config, NemotronH_Nano_VL_V2_Config, NemotronHConfig, @@ -101,6 +103,8 @@ Step3VLConfig, LongcatFlashConfig, Olmo3Config, + MuseGlimmerConfig, + MuseGlimmerAssistantConfig, KimiK3Config, KimiLinearConfig, Qwen3NextConfig, @@ -284,6 +288,64 @@ def check_gguf_file(model: Union[str, os.PathLike]) -> bool: return header == b"GGUF" +def resolve_hf_gguf_reference( + model: str, revision: Optional[str] = None +) -> Optional[str]: + """Download a .gguf named by Hub reference and return its local path. + + owner/repo/path/inside/repo.gguf -> exactly that file + owner/repo -> the only .gguf in the repo + """ + from sglang.srt.utils import is_remote_url + + if not model or os.path.exists(model) or is_remote_url(model): + return None + + parts = model.strip("/").split("/") + if len(parts) < 2: + return None + + from huggingface_hub import hf_hub_download + + if len(parts) > 2 and model.endswith(".gguf"): + repo_id = "/".join(parts[:2]) + filename = "/".join(parts[2:]) + return hf_hub_download(repo_id, filename, revision=revision) + + if len(parts) != 2: + return None + + from huggingface_hub import HfApi + + try: + files = [ + s.rfilename for s in HfApi().repo_info(model, revision=revision).siblings + ] + except Exception: + return None + if any(f == "config.json" for f in files): + return None + + candidates = [f for f in files if f.endswith(".gguf")] + if not candidates: + return None + if len(candidates) > 1: + listing = "\n ".join(f"{model}/{f}" for f in sorted(candidates)) + raise ValueError( + f"{model} contains {len(candidates)} .gguf files; name the one to " + f"serve:\n {listing}" + ) + return hf_hub_download(model, candidates[0], revision=revision) + + +def gguf_sidecar_dir( + gguf_path: Union[str, os.PathLike], sentinel: str +) -> Optional[Path]: + """Directory containing *sentinel* next to a .gguf file, if there is one.""" + directory = Path(gguf_path).parent + return directory if (directory / sentinel).is_file() else None + + # --------------------------------------------------------------------------- # Rope / text config helpers # --------------------------------------------------------------------------- @@ -491,6 +553,19 @@ def get_generation_config( revision: Optional[str] = None, **kwargs, ): + if check_gguf_file(model): + sidecar = gguf_sidecar_dir(model, "generation_config.json") + if sidecar is not None: + model = str(sidecar) + else: + from .gguf_native import ( + build_gguf_generation_config, + has_native_gguf_support, + ) + + if has_native_gguf_support(model): + return build_gguf_generation_config(model) + try: return GenerationConfig.from_pretrained( model, trust_remote_code=trust_remote_code, revision=revision, **kwargs diff --git a/python/sglang/srt/utils/hf_transformers/config.py b/python/sglang/srt/utils/hf_transformers/config.py index 1c26fbeebfb7..c0b258997e17 100644 --- a/python/sglang/srt/utils/hf_transformers/config.py +++ b/python/sglang/srt/utils/hf_transformers/config.py @@ -37,8 +37,10 @@ _override_v_head_dim_if_zero, check_gguf_file, get_hf_text_config, + gguf_sidecar_dir, resolve_runai_obj_uri, ) +from .gguf_native import build_gguf_config, has_native_gguf_support from .mistral_utils import is_mistral_model, load_mistral_config @@ -224,6 +226,7 @@ def get_config( **kwargs, ): is_gguf = check_gguf_file(model) + gguf_has_sidecar_config = False if is_gguf: if model_config_parser not in ("auto", "hf"): raise ValueError( @@ -231,7 +234,14 @@ def get_config( "with GGUF inputs; only 'hf' (or 'auto') is supported." ) _ensure_gguf_version() - kwargs["gguf_file"] = model + gguf_has_sidecar_config = gguf_sidecar_dir(model, "config.json") is not None + if not gguf_has_sidecar_config and has_native_gguf_support(model): + config = build_gguf_config(model) + if model_override_args: + config.update(model_override_args) + return config + if not gguf_has_sidecar_config: + kwargs["gguf_file"] = model model = Path(model).parent # Skip auto-resolution for GGUF: the name-based Mistral heuristic # would misfire on the rewritten parent dir. @@ -264,9 +274,13 @@ def get_config( else: setattr(config, key, value) - if is_gguf: + if is_gguf and not gguf_has_sidecar_config: if config.model_type not in MODEL_FOR_CAUSAL_LM_MAPPING_NAMES: - raise RuntimeError(f"Can't get gguf config for {config.model_type}.") + raise RuntimeError( + f"Can't get gguf config for {config.model_type}. Place a " + "config.json next to the .gguf file to load the config from " + "there instead." + ) _set_architectures(config, MODEL_FOR_CAUSAL_LM_MAPPING_NAMES[config.model_type]) return config diff --git a/python/sglang/srt/utils/hf_transformers/gguf_native.py b/python/sglang/srt/utils/hf_transformers/gguf_native.py new file mode 100644 index 000000000000..a9a95b6cd1c6 --- /dev/null +++ b/python/sglang/srt/utils/hf_transformers/gguf_native.py @@ -0,0 +1,258 @@ +# 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 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Reading config and tokenizer from a GGUF whose architecture transformers lacks. + +``load_gguf_checkpoint`` refuses any architecture outside its own +``GGUF_SUPPORTED_ARCHITECTURES``, and it does so before touching a single field, +so both the config and the tokenizer are unreachable for such a checkpoint -- +even though the tokenizer half of that reader is entirely architecture-agnostic +(it dispatches on ``tokenizer.ggml.model``, not on the model architecture). + +This module carries SGLang's own path for those checkpoints: + +* ``GGUF_NATIVE_CONFIG_BUILDERS`` maps a GGUF ``general.architecture`` to a + builder returning a fully populated config. +* ``build_gguf_tokenizer`` reuses transformers' own converters, which work fine + once they are reached directly instead of through the gated loader. + +Reaching for these is a last resort: a config.json next to the .gguf still wins, +because the checkpoint author's own config outranks anything reconstructed. +""" + +from typing import Any, Callable, Dict, Optional + +from transformers import PretrainedConfig + +from sglang.srt.configs.muse_glimmer import MuseGlimmerConfig + +GGUF_NATIVE_CONFIG_BUILDERS: Dict[str, Callable[[str], PretrainedConfig]] = { + "muse-glimmer": MuseGlimmerConfig.from_gguf, +} + + +def read_gguf_architecture(gguf_path: str) -> Optional[str]: + """The ``general.architecture`` string, or None if it cannot be read.""" + try: + from gguf import GGUFReader + + reader = GGUFReader(gguf_path) + field = reader.fields.get("general.architecture") + if field is None: + return None + value = field.contents() + return value if isinstance(value, str) else None + except Exception: + return None + + +def has_native_gguf_support(gguf_path: str) -> bool: + return read_gguf_architecture(gguf_path) in GGUF_NATIVE_CONFIG_BUILDERS + + +def build_gguf_config(gguf_path: str) -> PretrainedConfig: + arch = read_gguf_architecture(gguf_path) + return GGUF_NATIVE_CONFIG_BUILDERS[arch](gguf_path) + + +_GPT4O_SPLIT_REGEX = ( + r"[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]*" + r"[\p{Ll}\p{Lm}\p{Lo}\p{M}]+(?i:'s|'t|'re|'ve|'m|'ll|'d)?|" + r"[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]+" + r"[\p{Ll}\p{Lm}\p{Lo}\p{M}]*(?i:'s|'t|'re|'ve|'m|'ll|'d)?|" + r"\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n/]*|\s*[\r\n]+|\s+(?!\S)|\s+" +) + +_PRE_TOKENIZER_REGEX = { + # LLAMA_VOCAB_PRE_TYPE_LLAMA3 + "llama-bpe": ( + r"(?:'[sS]|'[tT]|'[rR][eE]|'[vV][eE]|'[mM]|'[lL][lL]|'[dD])|" + r"[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|" + r"\s*[\r\n]+|\s+(?!\S)|\s+" + ), + "gpt-4o": _GPT4O_SPLIT_REGEX, + "llama4": _GPT4O_SPLIT_REGEX, +} + +_GGML_TOKEN_TYPE_CONTROL = 3 + + +def build_gguf_generation_config(gguf_path: str): + """GenerationConfig from GGUF metadata, or None if there is nothing to say. + + llama.cpp records the end-of-generation ids explicitly, and for a + Harmony-style model the distinction matters: ``eos_token_id`` ends the + sequence and ``eot_token_id`` ends a turn, so both must stop generation while + an end-of-*message* id must not -- stopping on that truncates the model + mid-reasoning, before it answers. + """ + from gguf import GGUFReader + from transformers import GenerationConfig + + reader = GGUFReader(gguf_path) + meta = {key: field.contents() for key, field in reader.fields.items()} + + stop_ids = [] + for key in ("tokenizer.ggml.eos_token_id", "tokenizer.ggml.eot_token_id"): + if key in meta: + value = int(meta[key]) + if value not in stop_ids: + stop_ids.append(value) + if not stop_ids: + return None + + fields: Dict[str, Any] = { + "eos_token_id": stop_ids if len(stop_ids) > 1 else stop_ids[0] + } + if "tokenizer.ggml.bos_token_id" in meta: + fields["bos_token_id"] = int(meta["tokenizer.ggml.bos_token_id"]) + if "tokenizer.ggml.padding_token_id" in meta: + fields["pad_token_id"] = int(meta["tokenizer.ggml.padding_token_id"]) + return GenerationConfig(**fields) + + +def build_gguf_tokenizer(gguf_path: str, **kwargs: Any): + """Build a fast tokenizer from GGUF metadata alone. + + transformers' own GGUF tokenizer path is unreachable for an architecture its + checkpoint loader rejects, and its converters key on ``tokenizer.ggml.model`` + (here "gpt2") which loses both the special-token block and the pre-tokenizer + regex. So the tokenizers spec is assembled directly instead: a byte-level BPE + over the NORMAL tokens, the CONTROL tokens registered as added specials, and + the split regex named by ``tokenizer.ggml.pre``. + """ + import json + + from gguf import GGUFReader + from tokenizers import Tokenizer + from transformers import PreTrainedTokenizerFast + + reader = GGUFReader(gguf_path) + meta = {key: field.contents() for key, field in reader.fields.items()} + + tokens = list(meta["tokenizer.ggml.tokens"]) + token_types = [int(t) for t in meta["tokenizer.ggml.token_type"]] + merges = [tuple(m.split(" ", 1)) for m in meta["tokenizer.ggml.merges"]] + + pre_name = meta.get("tokenizer.ggml.pre") + if pre_name not in _PRE_TOKENIZER_REGEX: + raise ValueError( + f"No pre-tokenizer regex known for tokenizer.ggml.pre={pre_name!r}; " + f"known: {sorted(_PRE_TOKENIZER_REGEX)}" + ) + + control_ids = [ + i for i, t in enumerate(token_types) if t == _GGML_TOKEN_TYPE_CONTROL + ] + control = set(control_ids) + vocab = {tok: i for i, tok in enumerate(tokens) if i not in control} + + def token_of(key): + idx = meta.get(f"tokenizer.ggml.{key}") + return None if idx is None else tokens[int(idx)] + + bos = token_of("bos_token_id") + + spec = { + "version": "1.0", + "truncation": None, + "padding": None, + "added_tokens": [ + { + "id": i, + "content": tokens[i], + "single_word": False, + "lstrip": False, + "rstrip": False, + "normalized": False, + "special": True, + } + for i in control_ids + ], + "normalizer": None, + "pre_tokenizer": { + "type": "Sequence", + "pretokenizers": [ + { + "type": "Split", + "pattern": {"Regex": _PRE_TOKENIZER_REGEX[pre_name]}, + "behavior": "Isolated", + "invert": False, + }, + { + "type": "ByteLevel", + "add_prefix_space": False, + "trim_offsets": True, + "use_regex": False, + }, + ], + }, + "post_processor": None, + "decoder": { + "type": "ByteLevel", + "add_prefix_space": True, + "trim_offsets": True, + "use_regex": True, + }, + "model": { + "type": "BPE", + "dropout": None, + "unk_token": None, + "continuing_subword_prefix": None, + "end_of_word_suffix": None, + "fuse_unk": False, + "byte_fallback": False, + "ignore_merges": True, + "vocab": vocab, + "merges": [list(m) for m in merges], + }, + } + + if meta.get("tokenizer.ggml.add_bos_token") and bos is not None: + bos_id = int(meta["tokenizer.ggml.bos_token_id"]) + spec["post_processor"] = { + "type": "TemplateProcessing", + "single": [ + {"SpecialToken": {"id": bos, "type_id": 0}}, + {"Sequence": {"id": "A", "type_id": 0}}, + ], + "pair": [ + {"SpecialToken": {"id": bos, "type_id": 0}}, + {"Sequence": {"id": "A", "type_id": 0}}, + {"Sequence": {"id": "B", "type_id": 0}}, + ], + "special_tokens": { + bos: {"id": bos, "ids": [bos_id], "tokens": [bos]}, + }, + } + + backend = Tokenizer.from_str(json.dumps(spec)) + + named = { + bos, + token_of("eos_token_id"), + token_of("padding_token_id"), + token_of("unknown_token_id"), + } + additional = [tokens[i] for i in control_ids if tokens[i] not in named] + + return PreTrainedTokenizerFast( + tokenizer_object=backend, + bos_token=bos, + eos_token=token_of("eos_token_id"), + unk_token=token_of("unknown_token_id"), + pad_token=token_of("padding_token_id"), + additional_special_tokens=additional, + chat_template=meta.get("tokenizer.chat_template"), + **kwargs, + ) diff --git a/python/sglang/srt/utils/hf_transformers/tokenizer.py b/python/sglang/srt/utils/hf_transformers/tokenizer.py index def8d25814e2..b6d25a1d8f61 100644 --- a/python/sglang/srt/utils/hf_transformers/tokenizer.py +++ b/python/sglang/srt/utils/hf_transformers/tokenizer.py @@ -34,8 +34,10 @@ _resolve_local_or_cached_file, attach_additional_stop_token_ids, check_gguf_file, + gguf_sidecar_dir, resolve_runai_obj_uri, ) +from .gguf_native import build_gguf_tokenizer, has_native_gguf_support from .mistral_utils import ( _MISTRAL_TOKENIZER_REDIRECTS, is_bare_tekken_checkpoint, @@ -144,7 +146,8 @@ def _resolve_tokenizer_name(tokenizer_name, kwargs): if check_gguf_file(tokenizer_name): _ensure_gguf_version() - kwargs["gguf_file"] = tokenizer_name + if gguf_sidecar_dir(tokenizer_name, "tokenizer_config.json") is None: + kwargs["gguf_file"] = tokenizer_name tokenizer_name = Path(tokenizer_name).parent tokenizer_name = resolve_runai_obj_uri(tokenizer_name) @@ -487,6 +490,17 @@ def get_tokenizer( if "use_fast" not in kwargs: kwargs["use_fast"] = True + if ( + check_gguf_file(tokenizer_name) + and gguf_sidecar_dir(tokenizer_name, "tokenizer_config.json") is None + and has_native_gguf_support(tokenizer_name) + ): + _ensure_gguf_version() + tokenizer = build_gguf_tokenizer(tokenizer_name) + _fix_special_tokens_pattern(tokenizer) + attach_additional_stop_token_ids(tokenizer) + return patch_tokenizer(tokenizer) + tokenizer_name = _resolve_tokenizer_name(tokenizer_name, kwargs) common_kwargs = dict( diff --git a/test/registered/mlx/models_e2e/test_muse_glimmer_mlx_correctness.py b/test/registered/mlx/models_e2e/test_muse_glimmer_mlx_correctness.py new file mode 100644 index 000000000000..6e644fee4a27 --- /dev/null +++ b/test/registered/mlx/models_e2e/test_muse_glimmer_mlx_correctness.py @@ -0,0 +1,490 @@ +"""Correctness tests for Muse Glimmer served on the SGLang MLX backend. + +Muse Glimmer interleaves sliding-window (window=2048) and NoPE full-attention +layers, gates attention output through a sigmoid projection fused into +q_proj, and ships as an mlx-lm ``model_file`` artifact — so it exercises +the MLX backend's remote-code gate, container-level window discovery, and +banded-mask paths end to end. Two guards: + +1. ``TestMuseGlimmerMlxServing`` — black-box serving smoke against a running + server, including a prompt long enough to engage the sliding window. +2. ``TestMuseGlimmerMlxReferenceCorrectness`` — token-for-token equivalence of + SGLang greedy decoding against raw, unpatched mlx_lm greedy generation + on identical ``input_ids``. Both sides keep full KV history (the Muse Glimmer + model file's ``make_cache`` deliberately avoids ``RotatingKVCache``), + so exact equality holds even past the window and across prefill + chunkings. + +The artifact is weight-derived and private: its path comes exclusively from the +``SGLANG_MLX_TEST_ONYX_ARTIFACT`` environment variable and the whole module SKIPS +cleanly when it is unset or missing — CI has no dependency on private +assets. Window-engagement prompt lengths are derived from the artifact's +own config, so the suite also runs against tiny packaged fixtures. +""" + +from __future__ import annotations + +import concurrent.futures +import importlib.util +import json +import os +import unittest +from pathlib import Path + +import requests + +from sglang.srt.utils import kill_process_tree +from sglang.test.ci.ci_register import register_cpu_ci, register_mlx_ci +from sglang.test.test_utils import ( + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + DEFAULT_URL_FOR_TEST, + CustomTestCase, + popen_launch_server, +) + +# Registered on the CPU suite as an import check; skipped wherever mlx or +# the private artifact is absent. stage-b-e2e-mlx is manual-dispatch only. +register_cpu_ci(est_time=1, suite="base-a-test-cpu") +register_mlx_ci(est_time=420, suite="stage-b-e2e-mlx") + +_HAS_MLX = ( + importlib.util.find_spec("mlx") is not None + and importlib.util.find_spec("mlx_lm") is not None +) + +ARTIFACT = os.environ.get("SGLANG_MLX_TEST_ONYX_ARTIFACT") +MEM_FRACTION_STATIC = os.environ.get("SGLANG_MLX_TEST_MEM_FRACTION", "0.85") +MIN_FREE_GB = float(os.environ.get("SGLANG_MLX_TEST_MIN_FREE_GB", "20")) +MAX_NEW_TOKENS = 32 + + +def _artifact_or_skip() -> Path: + if not _HAS_MLX: + raise unittest.SkipTest("requires mlx + mlx_lm (Apple Silicon only)") + if not ARTIFACT: + raise unittest.SkipTest( + "SGLANG_MLX_TEST_ONYX_ARTIFACT is not set; the Muse Glimmer artifact is private " + "and must be supplied via the environment" + ) + path = Path(ARTIFACT).expanduser() + if not (path / "config.json").is_file(): + raise unittest.SkipTest(f"no packaged artifact at {path}") + config = json.loads((path / "config.json").read_text()) + if config.get("onyx_mlx_format") != 1: + raise unittest.SkipTest( + f"{path} is not a packaged Muse Glimmer MLX artifact (missing marker)" + ) + return path + + +def _artifact_config(path: Path) -> dict: + return json.loads((path / "config.json").read_text()) + + +def _available_gb(): + try: + import psutil + + return psutil.virtual_memory().available / 1024**3 + except Exception: + return None + + +def _check_memory(path: Path): + # Tiny fixtures need no headroom; only guard for real-size artifacts. + weights_gb = ( + sum(p.stat().st_size for p in path.glob("model*.safetensors")) / 1024**3 + ) + if weights_gb <= 1: + return + # A previous test class's just-killed server can hold memory for a few + # seconds; wait for the release instead of skipping on a transient dip. + import time + + deadline = time.monotonic() + 90 + avail = _available_gb() + while avail is not None and avail < MIN_FREE_GB and time.monotonic() < deadline: + time.sleep(5) + avail = _available_gb() + if avail is not None and avail < MIN_FREE_GB: + raise unittest.SkipTest( + f"insufficient free memory: {avail:.1f} GB < {MIN_FREE_GB} GB " + f"needed to safely serve a {weights_gb:.0f} GB artifact" + ) + + +def _case_lengths(config: dict) -> list[int]: + window = int(config.get("sliding_window", 2048)) + max_pos = int(config.get("max_position_embeddings", 16384)) + return sorted( + { + max(4, window - 1), + min(window + 64, max_pos - MAX_NEW_TOKENS - 1), + min(2 * window, max_pos - MAX_NEW_TOKENS - 1), + } + ) + + +def _prompt_id_cases(config: dict) -> list[list[int]]: + """Deterministic random input_ids around the artifact's window size. + + Serving-smoke material only: random tokens produce flat next-token + distributions, fine for shape/length assertions but useless for exact + greedy matching. + """ + import numpy as np + + vocab = int(config["vocab_size"]) + hi = min(vocab, 200000) + lo = min(1000, hi // 2) + rng = np.random.default_rng(20260731) + return [rng.integers(lo, hi, size=n).tolist() for n in _case_lengths(config)] + + +def _natural_prompt_cases(config: dict, artifact: Path) -> list[list[int]]: + """Natural-text prompts at the same window-straddling lengths. + + Real text gives the model real top-1 margins, so exact greedy + equivalence across different batching shapes is meaningful — a benign + kernel-shape difference cannot flip a decisive argmax, while a masking + or batching bug still can. + """ + from mlx_lm.utils import load_tokenizer + + tok = load_tokenizer(artifact) + filler = ( + "Day %d: we walked along the ridge, catalogued mosses and lichens, " + "measured stream flow, and noted the weather turning. " + ) + cases = [] + for n in _case_lengths(config): + # Prompts end MID-NARRATIVE: continuing a strictly cyclic journal is + # near-deterministic, so greedy continuations sit on wide top-1 + # margins (an open-ended question would put the decision point on a + # near-tie, where benign kernel-shape noise can flip the argmax). + text, day = "The following is a field journal. ", 1 + ids: list[int] = [] + while len(ids) < n: + text += filler % day + day += 1 + ids = tok.encode(text) + cases.append(ids[:n]) + return cases + + +def _reference_greedy(model, prompt_ids, max_tokens, prefill_step_size=None): + import mlx.core as mx + from mlx_lm.generate import generate_step + + kwargs = {} + if prefill_step_size is not None: + # Match the server's chunked-prefill size: Metal matmuls accumulate + # at reduced precision and are shape-dependent, so exact greedy + # equality requires both sides to prefill in identical chunk shapes. + kwargs["prefill_step_size"] = prefill_step_size + out = [] + for token, _ in generate_step( + mx.array(prompt_ids), model, max_tokens=max_tokens, **kwargs + ): + out.append(int(token)) + if len(out) >= max_tokens: + break + return out + + +class TestMuseGlimmerMlxServing(CustomTestCase): + @classmethod + def setUpClass(cls): + cls.artifact = _artifact_or_skip() + _check_memory(cls.artifact) + cls.config = _artifact_config(cls.artifact) + cls.base_url = DEFAULT_URL_FOR_TEST + + env = os.environ.copy() + env["SGLANG_USE_MLX"] = "1" + cls.process = popen_launch_server( + str(cls.artifact), + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=[ + "--trust-remote-code", + "--disable-radix-cache", + "--disable-cuda-graph", + "--mem-fraction-static", + MEM_FRACTION_STATIC, + ], + env=env, + ) + + @classmethod + def tearDownClass(cls): + if hasattr(cls, "process") and cls.process is not None: + kill_process_tree(cls.process.pid) + + def _generate(self, payload): + resp = requests.post(f"{self.base_url}/generate", json=payload, timeout=600) + resp.raise_for_status() + return resp.json() + + def test_solo_generation(self): + ids = _prompt_id_cases(self.config)[0] + out = self._generate( + { + "input_ids": ids, + "sampling_params": { + "max_new_tokens": 16, + "temperature": 0, + "ignore_eos": True, + }, + } + ) + self.assertEqual(len(out["output_ids"]), 16) + + def test_window_engaging_prompt(self): + # Longest case exceeds the sliding window: prefill and decode both + # run with banded masks engaged on the sliding layers. + ids = _prompt_id_cases(self.config)[-1] + self.assertGreater(len(ids), int(self.config["sliding_window"])) + out = self._generate( + { + "input_ids": ids, + "sampling_params": { + "max_new_tokens": 16, + "temperature": 0, + "ignore_eos": True, + }, + } + ) + self.assertEqual(len(out["output_ids"]), 16) + + def test_batch_generation(self): + cases = _prompt_id_cases(self.config) + out = self._generate( + { + "input_ids": cases, + "sampling_params": { + "max_new_tokens": 12, + "temperature": 0, + "ignore_eos": True, + }, + } + ) + self.assertEqual(len(out), len(cases)) + for r in out: + self.assertEqual(len(r["output_ids"]), 12) + + def test_over_context_request_rejected(self): + # A prompt beyond the context window must produce a clear error, + # not unbounded cache growth or a process kill. + max_pos = int(self.config.get("max_position_embeddings", 16384)) + too_long = [1000 + (i % 1000) for i in range(max_pos + 64)] + resp = requests.post( + f"{self.base_url}/generate", + json={ + "input_ids": too_long, + "sampling_params": {"max_new_tokens": 8, "temperature": 0}, + }, + timeout=600, + ) + body = ( + resp.json() + if resp.headers.get("content-type", "").startswith("application/json") + else {} + ) + rejected = resp.status_code >= 400 or ( + isinstance(body, dict) + and body.get("meta_info", {}).get("finish_reason", {}).get("type") + == "abort" + ) + self.assertTrue( + rejected, + f"over-context request was not rejected: {resp.status_code} {body}", + ) + # The server must still be healthy afterwards. + ok = requests.get(f"{self.base_url}/health", timeout=60) + self.assertEqual(ok.status_code, 200) + + def test_abort_then_rerun_is_isolated(self): + # Closing a stream mid-decode aborts the request server-side; a + # fresh identical request afterwards must produce the same output + # as one issued before the abort (no cache/state leakage). + ids = _prompt_id_cases(self.config)[0] + params = {"max_new_tokens": 24, "temperature": 0, "ignore_eos": True} + + def full_run(): + r = requests.post( + f"{self.base_url}/generate", + json={"input_ids": ids, "sampling_params": params}, + timeout=600, + ) + r.raise_for_status() + return r.json()["output_ids"] + + before = full_run() + with requests.post( + f"{self.base_url}/generate", + json={ + "input_ids": ids, + "sampling_params": params, + "stream": True, + }, + stream=True, + timeout=600, + ) as resp: + for line in resp.iter_lines(): + if line and line.startswith(b"data:") and b"[DONE]" not in line: + break # first token seen -> drop the connection mid-decode + after = full_run() + self.assertEqual(before, after, "post-abort rerun diverged") + + def test_explicit_eom_stop_honored(self): + # 200007 must not be a default stop, but the stop machinery must + # still be ABLE to stop on it when a request asks — proving its + # absence from the default set is configuration, not inability. + # The prompt pre-opens the reasoning channel, so the only way for + # greedy decoding to close it is <|eom|> — whether the model would + # have chosen to reason on its own varies with server config. + if self.config["vocab_size"] < 200008: + self.skipTest("tiny fixture vocab has no <|eom|> token") + out = self._generate( + { + "text": ( + "<|start|>user<|message|>Hi, who are you?<|eot|>" + "<|start|>assistant to=self<|message|>" + ), + "sampling_params": { + "max_new_tokens": 512, + "temperature": 0, + "stop_token_ids": [200007], + }, + } + ) + finish = out["meta_info"]["finish_reason"] + self.assertEqual(finish["type"], "stop") + self.assertEqual(finish["matched"], 200007) + + def test_eom_not_terminal_for_real_artifact(self): + # <|eom|> (200007) closes a message, not a turn; the packaged + # generation_config must not list it as EOS, or every response + # truncates at end-of-thinking. Only meaningful on the real vocab. + if self.config["vocab_size"] < 200008: + self.skipTest("tiny fixture vocab has no <|eom|> token") + gen = json.loads((self.artifact / "generation_config.json").read_text()) + self.assertNotIn(200007, gen["eos_token_id"]) + self.assertEqual(sorted(gen["eos_token_id"]), [200001, 200008]) + + +class TestMuseGlimmerMlxReferenceCorrectness(CustomTestCase): + """SGLang vs raw mlx_lm greedy equivalence on identical input_ids. + + Real-weight artifacts only: random-weight fixtures produce softcapped + logits with near-zero top-1 margins, so ANY benign kernel-shape + difference (e.g. the wrapper's trailing-window KV truncation vs the + reference's full-history banded mask, both mathematically equivalent) + flips argmax ties and exact matching carries no signal. + """ + + @classmethod + def setUpClass(cls): + cls.artifact = _artifact_or_skip() + _check_memory(cls.artifact) + cls.config = _artifact_config(cls.artifact) + if cls.config["vocab_size"] < 200008: + raise unittest.SkipTest( + "exact greedy equivalence needs real Muse Glimmer weights; " + "tiny random-weight fixtures have no top-1 margin" + ) + + from mlx_lm.utils import load_model + + cls.chunk = max(256, int(cls.config.get("sliding_window", 2048)) // 2) + ref_model, _ = load_model(cls.artifact) + cls.cases = [] + for prompt_ids in _natural_prompt_cases(cls.config, cls.artifact): + ref_ids = _reference_greedy( + ref_model, + prompt_ids, + MAX_NEW_TOKENS, + prefill_step_size=cls.chunk, + ) + cls.cases.append((prompt_ids, ref_ids)) + del ref_model + import gc + + gc.collect() + + cls.base_url = DEFAULT_URL_FOR_TEST + env = os.environ.copy() + env["SGLANG_USE_MLX"] = "1" + # Chunk size below the sliding window forces the longest prompt to + # prefill in several chunks that cross the window boundary — exact + # equality then also covers chunked-prefill mask stitching. The + # reference used the SAME chunk size (set in the class attribute + # above) so both sides prefill in identical kernel shapes. + chunk = cls.chunk + cls.process = popen_launch_server( + str(cls.artifact), + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=[ + "--trust-remote-code", + "--disable-radix-cache", + "--disable-cuda-graph", + "--mem-fraction-static", + MEM_FRACTION_STATIC, + "--chunked-prefill-size", + str(chunk), + ], + env=env, + ) + + @classmethod + def tearDownClass(cls): + if hasattr(cls, "process") and cls.process is not None: + kill_process_tree(cls.process.pid) + + def _sglang_greedy(self, prompt_ids, max_tokens): + resp = requests.post( + f"{self.base_url}/generate", + json={ + "input_ids": prompt_ids, + "sampling_params": { + "max_new_tokens": max_tokens, + "temperature": 0, + "ignore_eos": True, + }, + }, + timeout=600, + ) + resp.raise_for_status() + return resp.json()["output_ids"] + + def test_solo_greedy_matches_reference(self): + for prompt_ids, ref_ids in self.cases: + got = self._sglang_greedy(prompt_ids, MAX_NEW_TOKENS) + self.assertEqual( + got, + ref_ids, + f"greedy divergence at prompt length {len(prompt_ids)}", + ) + + def test_concurrent_ragged_batch_matches_reference(self): + # All cases in flight together: ragged lengths force mixed + # prefill/decode batching; every stream must still match its solo + # mlx_lm reference exactly (no cross-request contamination). + with concurrent.futures.ThreadPoolExecutor(len(self.cases)) as pool: + futures = [ + pool.submit(self._sglang_greedy, prompt_ids, MAX_NEW_TOKENS) + for prompt_ids, _ in self.cases + ] + results = [f.result() for f in futures] + for (prompt_ids, ref_ids), got in zip(self.cases, results): + self.assertEqual( + got, + ref_ids, + f"batched divergence at prompt length {len(prompt_ids)}", + ) + + +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/test/registered/spec/dflash/test_muse_glimmer_dflash_assistant_gsm8k.py b/test/registered/spec/dflash/test_muse_glimmer_dflash_assistant_gsm8k.py new file mode 100644 index 000000000000..1af9ebd630b6 --- /dev/null +++ b/test/registered/spec/dflash/test_muse_glimmer_dflash_assistant_gsm8k.py @@ -0,0 +1,94 @@ +import unittest + +from sglang.srt.utils import kill_process_tree +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.kits.eval_accuracy_kit import GSM8KMixin +from sglang.test.test_utils import ( + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + DEFAULT_URL_FOR_TEST, + CustomTestCase, + popen_launch_server, +) + +register_cuda_ci(est_time=900, stage="nightly", runner_config="1-gpu-large") + +TARGET_MODEL = "meta-models/Muse-Glimmer-30B" +DRAFT_MODEL = "meta-models/Muse-Glimmer-30B-assistant" + + +class TestMuseGlimmerDflashAssistantGSM8K(CustomTestCase, GSM8KMixin): + """GSM8K + DFlash accept-length regression test for the native + MuseGlimmerAssistantModel draft (``meta-models/Muse-Glimmer-30B-assistant``). + + This checkpoint loads through sglang's own native ``models/dflash.py`` / + ``configs/muse_glimmer.py::MuseGlimmerAssistantConfig`` with no extra wheel + dependency. Integrating it surfaced two bugs that fail *silently*, never raise, and do not move + GSM8K accuracy at temperature 0 (speculative decoding always falls back + to the target's own correct token on a draft miss, so a broken draft + still produces exactly the target's answers -- just slower): + + 1. The vendor's weight names (``encoder.fc.weight`` / + ``encoder.output_norm_enc.weight``) didn't match what + ``DFlashDraftModel.load_weights`` expected (``fc.weight`` / + ``hidden_norm.weight``), so those two tensors silently stayed at + random init. + 2. The vendor's ``target_layer_ids`` are in the HF "output of layer k" + convention; ``models/muse_glimmer.py::set_dflash_layers_to_capture`` + uses ids as-is (Muse Glimmer's own draft configs carry llama.cpp's + layer-*input* convention), so every captured layer was off by one. + + Both together collapsed real (non-simulated) accept_length to ~1.00 at + ``--speculative-dflash-block-size 5`` -- effectively no speculation, only + the mandatory bonus token -- with GSM8K accuracy unaffected throughout. + ``gsm8k_accept_length_thres`` is therefore the actual regression guard + here; the accuracy threshold alone would not have caught this. + + Measured after both fixes, same block size, same target: median + accept_length 3.12 (this draft) vs 3.09 (our own GGUF-converted draft, + the previous ground truth) across batch sizes 1/4/8 at 512in/256out -- + statistically indistinguishable. 2.5 leaves headroom below that for + workload variance while sitting well above the ~1.0 broken-state value. + """ + + model = TARGET_MODEL + gsm8k_backend = ( + "sgl_eval" # chat completions API, not /generate or raw /completions + ) + gsm8k_score_threshold = 0.85 + gsm8k_num_examples = 200 + gsm8k_accept_length_thres = 2.5 + + @classmethod + def setUpClass(cls): + cls.base_url = DEFAULT_URL_FOR_TEST + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=[ + "--reasoning-parser", + "muse", + "--tool-call-parser", + "muse", + "--language-model-only", + "--speculative-algorithm", + "DFLASH", + "--speculative-draft-model-path", + DRAFT_MODEL, + "--speculative-draft-load-format", + "auto", + "--speculative-dflash-block-size", + "5", + "--mem-fraction-static", + "0.85", + ], + ) + + @classmethod + def tearDownClass(cls): + if hasattr(cls, "process") and cls.process: + kill_process_tree(cls.process.pid) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/entrypoints/openai/test_serving_responses.py b/test/registered/unit/entrypoints/openai/test_serving_responses.py index 9324569988ae..9b68fad2cc3a 100644 --- a/test/registered/unit/entrypoints/openai/test_serving_responses.py +++ b/test/registered/unit/entrypoints/openai/test_serving_responses.py @@ -311,6 +311,78 @@ async def fake_generate( self.assertFalse(parser_cls.call_args.kwargs["force_reasoning"]) +class SkipSpecialTokensForwardingTestCase(CustomTestCase): + """The skip_special_tokens override from _process_messages must reach the + engine sampling params; muse's channel markers die in detok otherwise.""" + + def _create_responses_sampling_params(self, serving): + serving.default_chat_template_kwargs = None + rendered = MessageProcessingResult( + prompt="prompt", + prompt_ids=[1, 2, 3], + image_data=None, + audio_data=None, + video_data=None, + modalities=[], + stop=[], + ) + captured = {} + + async def fake_generate( + request_id, + request_prompt, + adapted_request, + sampling_params, + context, + **kwargs, + ): + captured["sampling_params"] = sampling_params + context.append_output( + { + "text": "done", + "meta_info": { + "prompt_tokens": 3, + "completion_tokens": 1, + "cached_tokens": 0, + }, + } + ) + yield context + + serving._generate_with_builtin_tools = fake_generate + request = ResponsesRequest( + model="x", + input="answer", + request_id="resp_skip_special", + store=False, + ) + + with ( + patch.object( + serving, "_apply_conversation_template", return_value=rendered + ), + patch( + "sglang.srt.entrypoints.openai.serving_responses.ReasoningParser" + ) as parser_cls, + ): + parser_cls.return_value.parse_non_stream.return_value = (None, "done") + response = asyncio.run(serving.create_responses(request)) + + self.assertEqual(response.status, "completed") + return captured["sampling_params"] + + def test_marker_preserving_parser_disables_skip_special_tokens(self): + serving = make_serving() + serving.reasoning_parser = "muse" + params = self._create_responses_sampling_params(serving) + self.assertFalse(params["skip_special_tokens"]) + + def test_default_parser_keeps_skip_special_tokens(self): + serving = make_serving() + params = self._create_responses_sampling_params(serving) + self.assertTrue(params["skip_special_tokens"]) + + class InputItemNormalizationTestCase(CustomTestCase): def test_function_call_becomes_assistant_tool_call(self): normalized = OpenAIServingResponses._normalize_response_message_for_chat( diff --git a/test/registered/unit/entrypoints/openai/test_serving_responses_stream.py b/test/registered/unit/entrypoints/openai/test_serving_responses_stream.py index 5c00ce3c5d2a..c199a7aa6863 100644 --- a/test/registered/unit/entrypoints/openai/test_serving_responses_stream.py +++ b/test/registered/unit/entrypoints/openai/test_serving_responses_stream.py @@ -167,6 +167,7 @@ def fake_parse_stream_chunk(delta): parser_cls.return_value.parse_stream_chunk.side_effect = ( fake_parse_stream_chunk ) + parser_cls.return_value.parse_stream_end.return_value = ("", []) fixture = StreamFixture(serving, request) events = fixture.run(chunks) @@ -178,6 +179,35 @@ def fake_parse_stream_chunk(delta): self.assertEqual(output[1]["name"], "get_weather") self.assertEqual(output[2]["content"][0]["text"], "It's sunny.") + def test_reasoning_parser_flushed_at_stream_end(self): + """Bug regression: the stream loop never drained text the reasoning + parser held back as a possible marker prefix, so a response whose text + genuinely ends with e.g. "<|e" lost that tail on /v1/responses (chat + flushes via parse_stream_end; responses did not).""" + serving = make_serving() + serving.reasoning_parser = "muse" + serving.tool_call_parser = None + + request = ResponsesRequest(model="x", input="hi", stream=True, store=False) + text = ( + " to=self<|message|>think<|eom|>" + "<|start|>assistant to=user<|message|>Answer<|e" + ) + fixture = StreamFixture(serving, request) + events = fixture.run( + [ + engine_chunk(text[:30], 4), + engine_chunk(text, 9, finish=True), + ] + ) + + streamed = "".join( + p["delta"] + for ev, p in zip(event_types(events), event_payloads(events)) + if ev == "response.output_text.delta" + ) + self.assertEqual(streamed, "Answer<|e") + class MultiToolCallStreamingOrderTestCase(CustomTestCase): """The wire order of message / function_call items across tool-call deltas.""" diff --git a/test/registered/unit/function_call/test_muse_glimmer_detector.py b/test/registered/unit/function_call/test_muse_glimmer_detector.py new file mode 100644 index 000000000000..8e75ceb1b8c0 --- /dev/null +++ b/test/registered/unit/function_call/test_muse_glimmer_detector.py @@ -0,0 +1,432 @@ +"""Unit tests for the Muse Glimmer ATEM tool-call detector — no server, no model loading. + +The expectations here are pinned to the checkpoint's own ``response_template`` +(``MUSE_GLIMMER_RESPONSE_SCHEMA`` in ``tokenizer_config.json``) and to the vendor's +reference parser, with particular attention to channel scoping: an +```` that only appears inside a reasoning block or a final answer +must never become a real tool call. +""" + +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.muse_glimmer_detector import MuseGlimmerDetector +from sglang.srt.parser.reasoning_parser import ReasoningParser +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(1.0, "base-a-test-cpu") + +DOUBLED = "get_weather.get_weather" + + +def atem(name: str, **params: str) -> str: + body = "".join( + f'{v}\n' for k, v in params.items() + ) + return ( + f'\n\n{body}' + f"\n" + ) + + +class TestMuseGlimmerDetector(CustomTestCase): + def setUp(self): + self.tools = [ + Tool( + type="function", + function=Function( + name="get_weather", + description="Get weather", + parameters={ + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + ), + ), + ] + + # ---- helpers ---------------------------------------------------------- + + def parse(self, text): + """Non-streaming parse -> (normal_text, [(name, args), ...]).""" + result = MuseGlimmerDetector().detect_and_parse(text, self.tools) + return result.normal_text, [ + (c.name, json.loads(c.parameters)) for c in result.calls if c.name + ] + + def parse_streaming(self, text, chunk_size): + detector = MuseGlimmerDetector() + normal, calls = [], [] + for i in range(0, len(text), chunk_size): + result = detector.parse_streaming_increment( + text[i : i + chunk_size], self.tools + ) + normal.append(result.normal_text) + calls.extend( + (c.name, json.loads(c.parameters)) for c in result.calls if c.name + ) + return "".join(normal), calls + + def assert_streaming_matches(self, text): + """Streaming must agree with one-shot parsing at every chunk boundary.""" + expected = self.parse(text) + for chunk_size in (1, 2, 3, 5, 7, 13, 29, 100): + self.assertEqual( + self.parse_streaming(text, chunk_size), + expected, + f"streaming diverged at chunk_size={chunk_size}", + ) + + # ---- tool extraction -------------------------------------------------- + + def test_single_tool_call(self): + text = ( + f" to=self<|message|>Need weather.<|eom|>" + f"<|start|>assistant to={DOUBLED}<|message|>{atem(DOUBLED, city='Paris')}" + ) + normal, calls = self.parse(text) + self.assertEqual(calls, [("get_weather", {"city": "Paris"})]) + self.assertEqual(normal, "Need weather.") + self.assert_streaming_matches(text) + + def test_parallel_tool_calls(self): + text = ( + f" to=self<|message|>Two cities.<|eom|>" + f"<|start|>assistant to={DOUBLED}<|message|>" + f"{atem(DOUBLED, city='Paris')}<|eom|>" + f"<|start|>assistant to={DOUBLED}<|message|>{atem(DOUBLED, city='Tokyo')}" + ) + _, calls = self.parse(text) + self.assertEqual( + calls, + [("get_weather", {"city": "Paris"}), ("get_weather", {"city": "Tokyo"})], + ) + self.assert_streaming_matches(text) + + def test_tool_call_then_final_answer(self): + text = ( + f" to=self<|message|>r<|eom|>" + f"<|start|>assistant to={DOUBLED}<|message|>" + f"{atem(DOUBLED, city='Paris')}<|eom|>" + f"<|start|>assistant to=user<|message|>It is sunny." + ) + normal, calls = self.parse(text) + self.assertEqual(calls, [("get_weather", {"city": "Paris"})]) + self.assertIn("It is sunny.", normal) + self.assert_streaming_matches(text) + + def test_namespaced_name_passes_through(self): + tools = [ + Tool( + type="function", + function=Function( + name="weather.get", + description="d", + parameters={"type": "object", "properties": {}}, + ), + ) + ] + text = ( + f" to=self<|message|>r<|eom|><|start|>assistant to=weather.get<|message|>" + f"{atem('weather.get', city='Paris')}" + ) + result = MuseGlimmerDetector().detect_and_parse(text, tools) + self.assertEqual([c.name for c in result.calls], ["weather.get"]) + + def test_parameter_value_typing(self): + """``allow_non_json: True`` — JSON literals decode, bare strings do not.""" + invoke = ( + '\n\n' + 'hello world\n' + '42\n' + 'true\n' + 'null\n' + '{"a": 1}\n' + '[1, 2]\n' + "\n" + ) + text = f"<|start|>assistant to=get_weather<|message|>{invoke}" + _, calls = self.parse(text) + self.assertEqual( + calls[0][1], + { + "s": "hello world", + "i": 42, + "b": True, + "n": None, + "o": {"a": 1}, + "l": [1, 2], + }, + ) + + def test_multiline_parameter_value(self): + value = 'line1\nline2\n"quoted"\n' + text = ( + f"<|start|>assistant to=get_weather<|message|>" + f'\n\n' + f'{value}\n' + f"\n" + ) + _, calls = self.parse(text) + self.assertEqual(calls[0][1], {"code": value}) + self.assert_streaming_matches(text) + + # ---- channel scoping (safety) ----------------------------------------- + + def test_invoke_inside_reasoning_is_not_a_call(self): + text = ( + f" to=self<|message|>Maybe I call {atem(DOUBLED, city='X')} — no.<|eom|>" + f"<|start|>assistant to=user<|message|>I will not call it." + ) + _, calls = self.parse(text) + self.assertEqual(calls, []) + self.assert_streaming_matches(text) + + def test_invoke_inside_final_answer_is_not_a_call(self): + text = ( + f" to=self<|message|>r<|eom|><|start|>assistant to=user<|message|>" + f"You would write:\n{atem(DOUBLED, city='X')}" + ) + _, calls = self.parse(text) + self.assertEqual(calls, []) + self.assert_streaming_matches(text) + + def test_invoke_inside_truncated_reasoning_is_not_a_call(self): + """Generation cut mid-CoT leaves no closing ``<|eom|>`` to anchor on.""" + text = f" to=self<|message|>I could call {atem(DOUBLED, city='X')} but" + _, calls = self.parse(text) + self.assertEqual(calls, []) + self.assert_streaming_matches(text) + + def test_truncated_tool_channel_drops_partial_invoke(self): + """A token cap mid-invoke must not fabricate a call from partial + arguments, and ATEM scaffolding must not leak into content.""" + text = ( + f" to=self<|message|>Need weather.<|eom|>" + f"<|start|>assistant to={DOUBLED}<|message|>" + f'\n\n' + f'Par' + ) + normal, calls = self.parse(text) + self.assertEqual(calls, []) + self.assertEqual(normal, "Need weather.") + self.assert_streaming_matches(text) + + def test_prose_opening_with_to_equals_does_not_stall(self): + """A bare ``to=...`` header opens the stream, so prose that happens to + start the same way is ambiguous. Mis-reading it as a header parks the + parser waiting for a ``<|message|>`` that never arrives and strands the + whole response in the buffer.""" + for text in ("to=x is the syntax.", "to= is an assignment", "to=aNeed weather.<|eom|>" + f"<|start|>assistant to={DOUBLED}<|message|>" + f"{atem(DOUBLED, city='Paris')}<|eom|>" + f"<|start|>assistant to=user<|message|>It is sunny in Paris." + ) + reasoning, remainder = ReasoningParser("muse").parse_non_stream(raw) + content, calls = FunctionCallParser(self.tools, "muse").parse_non_stream( + remainder + ) + self.assertEqual(reasoning, "Need weather.") + self.assertEqual(content, "It is sunny in Paris.") + self.assertEqual( + [(c.name, json.loads(c.parameters)) for c in calls], + [("get_weather", {"city": "Paris"})], + ) + + def test_pipeline_quoted_invoke_stays_content(self): + """A quoted ATEM block must survive as text, not become a call.""" + raw = ( + f" to=self<|message|>r<|eom|><|start|>assistant to=user<|message|>" + f"Example:\n{atem(DOUBLED, city='X')}" + ) + _, remainder = ReasoningParser("muse").parse_non_stream(raw) + content, calls = FunctionCallParser(self.tools, "muse").parse_non_stream( + remainder + ) + self.assertEqual(calls, []) + self.assertIn("`` search.""" + for quoted_at in ("after prose:\n", ""): + raw = ( + f" to=self<|message|>r<|eom|><|start|>assistant to=user<|message|>" + f"{quoted_at}<|start|>assistant to={DOUBLED}<|message|>" + f"{atem(DOUBLED, city='X')}" + ) + _, remainder = ReasoningParser( + "muse", tool_call_parser_active=True + ).parse_non_stream(raw) + content, calls = FunctionCallParser(self.tools, "muse").parse_non_stream( + remainder + ) + self.assertEqual( + calls, [], f"quoted header parsed as a call ({quoted_at!r})" + ) + self.assertIn("r<|eom|><|start|>assistant to=user<|message|>" + f"Let me check.<|eom|><|start|>assistant to={DOUBLED}<|message|>" + f"{atem(DOUBLED, city='Paris')}" + ) + want_calls = [("get_weather", {"city": "Paris"})] + _, remainder = ReasoningParser( + "muse", tool_call_parser_active=True + ).parse_non_stream(raw) + content, calls = FunctionCallParser(self.tools, "muse").parse_non_stream( + remainder + ) + self.assertEqual( + [(c.name, json.loads(c.parameters)) for c in calls], want_calls + ) + self.assertIn("Let me check.", content) + for chunk_size in (1, 7, 100): + _, s_content, s_calls = self.pipeline_stream(raw, chunk_size) + self.assertEqual(s_calls, want_calls, f"chunk_size={chunk_size}") + self.assertIn("Let me check.", s_content) + + def test_pipeline_plain_answer_stays_clean_without_tool_parse(self): + """Serving skips the tool detector when ``has_tool_call()`` is false, so + a turn with no ATEM block must come out of the reasoning parser already + unwrapped. Goes red if framing is preserved unconditionally.""" + raw = ( + " to=self<|message|>r<|eom|>" + "<|start|>assistant to=user<|message|>Hello there." + ) + reasoning, remainder = ReasoningParser( + "muse", tool_call_parser_active=True + ).parse_non_stream(raw) + self.assertFalse( + FunctionCallParser(self.tools, "muse").has_tool_call(remainder) + ) + self.assertEqual(reasoning, "r") + self.assertEqual(remainder, "Hello there.") + + def test_reasoning_finish_flushes_unframed_text(self): + """An unframed turn never emits ``<|message|>``, so nothing leaves the + buffer until the end-of-stream flush. Goes red if the Muse Glimmer reasoning + detector loses its ``finish()`` override.""" + rp = ReasoningParser("muse") + _, streamed = rp.parse_stream_chunk("Just plain text.") + _, flushed = rp.parse_stream_end() + self.assertEqual((streamed or "") + (flushed or ""), "Just plain text.") + + def test_interleaved_reasoning_blocks_join_with_newline(self): + """A turn may reason, call a tool, then reason again; the reference + schema joins the blocks with a newline. Goes red if the reasoning + detector concatenates the bodies directly, gluing two thoughts into + one word ("...thoughtsecond...").""" + raw = ( + f" to=self<|message|>first thought<|eom|>" + f"<|start|>assistant to={DOUBLED}<|message|>" + f"{atem(DOUBLED, city='Paris')}<|eom|>" + f"<|start|>assistant to=self<|message|>second thought<|eom|>" + f"<|start|>assistant to=user<|message|>done" + ) + reasoning, _ = ReasoningParser( + "muse", tool_call_parser_active=True + ).parse_non_stream(raw) + self.assertEqual(reasoning, "first thought\nsecond thought") + for chunk_size in (1, 7, 100): + s_reasoning, s_content, s_calls = self.pipeline_stream(raw, chunk_size) + self.assertEqual(s_reasoning, "first thought\nsecond thought") + self.assertEqual(s_calls, [("get_weather", {"city": "Paris"})]) + self.assertIn("done", s_content) + + def test_stream_end_flushes_partial_marker(self): + """An answer ending in a marker prefix (``<|st``) is held back while + streaming in case it grows into ``<|start|>``; the stream's end proves + it never will. Goes red if the tool parser loses its stream-end flush + (``parse_stream_end`` / detector ``finish``).""" + raw = ( + " to=self<|message|>r<|eom|>" + "<|start|>assistant to=user<|message|>answer<|st" + ) + for chunk_size in (1, 7, 100): + _, content, calls = self.pipeline_stream(raw, chunk_size) + self.assertEqual(calls, []) + self.assertEqual(content, "answer<|st", f"chunk_size={chunk_size}") + + def test_whitespace_before_header_is_tolerated(self): + """The model may put whitespace between ``<|eom|>`` and the next + ``<|start|>``; it travels with the header text. Goes red if the header + state requires ``<|start|>`` at exactly the first byte again.""" + text = ( + f" to=self<|message|>r<|eom|>\n<|start|>assistant to={DOUBLED}" + f"<|message|>{atem(DOUBLED, city='Paris')}" + ) + _, calls = self.parse(text) + self.assertEqual(calls, [("get_weather", {"city": "Paris"})]) + self.assert_streaming_matches(text) + + def test_registered_in_parser_enum(self): + self.assertIs( + FunctionCallParser.ToolCallParserEnum["muse"], MuseGlimmerDetector + ) + + +if __name__ == "__main__": + import unittest + + unittest.main() diff --git a/test/registered/unit/hardware_backend/mlx/test_mlx_remote_code_gate.py b/test/registered/unit/hardware_backend/mlx/test_mlx_remote_code_gate.py new file mode 100644 index 000000000000..50571ab961fa --- /dev/null +++ b/test/registered/unit/hardware_backend/mlx/test_mlx_remote_code_gate.py @@ -0,0 +1,178 @@ +"""Unit tests for the MLX remote-code gate. + +mlx-lm executes ``config.json``'s ``model_file`` unconditionally at load +time, so SGLang refuses such checkpoints before any checkpoint Python runs +unless the server was started with ``--trust-remote-code``. The refusal +tests prove non-execution with a sentinel ``model_file`` whose import would +leave an observable marker. +""" + +from __future__ import annotations + +import importlib.util +import json +import tempfile +import unittest +from pathlib import Path +from unittest.mock import patch + +from sglang.test.ci.ci_register import register_cpu_ci, register_mlx_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") +register_mlx_ci(est_time=5, suite="stage-a-unit-test-mlx") + +_HAS_MLX = ( + importlib.util.find_spec("mlx") is not None + and importlib.util.find_spec("mlx_lm") is not None +) + +from sglang.srt.hardware_backend.mlx.remote_code_gate import ( # noqa: E402 + RemoteCodeGateError, + ensure_remote_code_allowed, +) + +_SENTINEL = "GATE FAILED: checkpoint python executed" + + +def _make_checkpoint(tmp: Path, config: dict, *, with_sentinel: bool = True) -> Path: + (tmp / "config.json").write_text(json.dumps(config)) + if with_sentinel: + # Importing this file would create marker.txt — the refusal tests + # assert it never appears. + (tmp / "evil.py").write_text( + "from pathlib import Path\n" + f"Path(__file__).parent.joinpath('marker.txt').write_text({_SENTINEL!r})\n" + ) + return tmp + + +class TestEnsureRemoteCodeAllowed(CustomTestCase): + def setUp(self): + self._tmp = tempfile.TemporaryDirectory() + self.dir = Path(self._tmp.name) + + def tearDown(self): + self._tmp.cleanup() + + def _assert_sentinel_not_executed(self): + self.assertFalse( + (self.dir / "marker.txt").exists(), + "checkpoint python executed despite gate refusal", + ) + + def test_refuses_model_file_without_trust(self): + _make_checkpoint(self.dir, {"model_type": "onyx", "model_file": "evil.py"}) + with self.assertRaisesRegex(RemoteCodeGateError, "--trust-remote-code"): + ensure_remote_code_allowed(self.dir, trust_remote_code=False) + self._assert_sentinel_not_executed() + + def test_allows_model_file_with_trust(self): + _make_checkpoint(self.dir, {"model_type": "onyx", "model_file": "evil.py"}) + ensure_remote_code_allowed(self.dir, trust_remote_code=True) + # The gate itself never imports the file either way. + self._assert_sentinel_not_executed() + + def test_builtin_checkpoint_passes_without_trust(self): + _make_checkpoint(self.dir, {"model_type": "qwen3"}, with_sentinel=False) + ensure_remote_code_allowed(self.dir, trust_remote_code=False) + + def test_missing_config_rejected(self): + with self.assertRaisesRegex(RemoteCodeGateError, "no config.json"): + ensure_remote_code_allowed(self.dir, trust_remote_code=True) + + def test_malformed_config_rejected(self): + (self.dir / "config.json").write_text("{not json") + with self.assertRaisesRegex(RemoteCodeGateError, "not valid JSON"): + ensure_remote_code_allowed(self.dir, trust_remote_code=True) + + def test_non_object_config_rejected(self): + (self.dir / "config.json").write_text('["a", "b"]') + with self.assertRaisesRegex(RemoteCodeGateError, "JSON object"): + ensure_remote_code_allowed(self.dir, trust_remote_code=True) + + def test_missing_model_file_target_rejected(self): + _make_checkpoint( + self.dir, + {"model_type": "onyx", "model_file": "nope.py"}, + with_sentinel=False, + ) + with self.assertRaisesRegex(RemoteCodeGateError, "does not exist"): + ensure_remote_code_allowed(self.dir, trust_remote_code=True) + + def test_absolute_model_file_rejected(self): + _make_checkpoint( + self.dir, + {"model_type": "onyx", "model_file": "/etc/anything.py"}, + with_sentinel=False, + ) + with self.assertRaisesRegex(RemoteCodeGateError, "relative path"): + ensure_remote_code_allowed(self.dir, trust_remote_code=True) + + def test_traversal_model_file_rejected(self): + _make_checkpoint( + self.dir, + {"model_type": "onyx", "model_file": "../outside.py"}, + with_sentinel=False, + ) + with self.assertRaisesRegex(RemoteCodeGateError, "relative path"): + ensure_remote_code_allowed(self.dir, trust_remote_code=True) + + def test_non_string_model_file_rejected(self): + _make_checkpoint( + self.dir, + {"model_type": "onyx", "model_file": 42}, + with_sentinel=False, + ) + with self.assertRaisesRegex(RemoteCodeGateError, "non-string"): + ensure_remote_code_allowed(self.dir, trust_remote_code=True) + + +@unittest.skipUnless(_HAS_MLX, "requires mlx + mlx_lm") +class TestModelRunnerGateWiring(CustomTestCase): + """The runner must gate BEFORE calling mlx_lm's loader, on the same + resolved directory it then loads from.""" + + class _StopInit(Exception): + pass + + def setUp(self): + self._tmp = tempfile.TemporaryDirectory() + self.dir = Path(self._tmp.name) + + def tearDown(self): + self._tmp.cleanup() + + def test_refusal_precedes_loader_call(self): + from sglang.srt.hardware_backend.mlx.model_runner import MlxModelRunner + + _make_checkpoint(self.dir, {"model_type": "onyx", "model_file": "evil.py"}) + with patch( + "sglang.srt.hardware_backend.mlx.model_runner.mlx_lm_load" + ) as loader: + with self.assertRaisesRegex(RemoteCodeGateError, "--trust-remote-code"): + MlxModelRunner(model_path=str(self.dir), trust_remote_code=False) + loader.assert_not_called() + self.assertFalse((self.dir / "marker.txt").exists()) + + def test_trusted_load_uses_resolved_directory(self): + from sglang.srt.hardware_backend.mlx.model_runner import MlxModelRunner + + _make_checkpoint(self.dir, {"model_type": "onyx", "model_file": "evil.py"}) + with patch( + "sglang.srt.hardware_backend.mlx.model_runner.mlx_lm_load", + side_effect=self._StopInit, + ) as loader: + with self.assertRaises(self._StopInit): + MlxModelRunner(model_path=str(self.dir), trust_remote_code=True) + loader.assert_called_once() + called_path = loader.call_args.args[0] + self.assertEqual( + Path(called_path).resolve(), + self.dir.resolve(), + "loader must receive the same directory the gate inspected", + ) + + +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/test/registered/unit/hardware_backend/mlx/test_muse_glimmer_mlx_model.py b/test/registered/unit/hardware_backend/mlx/test_muse_glimmer_mlx_model.py new file mode 100644 index 000000000000..5f1524c3ad97 --- /dev/null +++ b/test/registered/unit/hardware_backend/mlx/test_muse_glimmer_mlx_model.py @@ -0,0 +1,279 @@ +"""Unit tests for the Muse Glimmer MLX model file's load path — no weights, no server. + +The e2e suite (``test/registered/mlx/models_e2e/test_muse_glimmer_mlx_correctness.py``) +needs the private packaged artifact, so the checkpoint-schema logic it relies +on is pinned here with tiny synthetic weights instead: + +1. ``flatten_rc_config`` — RC nested-config translation, including the two + convention conversions (qk_scale_factor gains sqrt(head_dim); NoPE layers + come from zeros in ``layer_rope_theta``). +2. ``ModelArgs`` validation — derived ``no_rope_layers``/``layer_types``, + rejection of inconsistent or malformed lists, format-version check. +3. ``sanitize`` — all three accepted weight layouts (raw HF, RC multimodal, + packaged) plus rejection of incomplete or mislabeled checkpoints. The + positional RC norm renames and the per-head q/gate interleave are verified + numerically, since getting either silently wrong still yields a model that + runs but computes garbage. +""" + +from __future__ import annotations + +import importlib.util +import unittest + +from sglang.test.ci.ci_register import register_cpu_ci, register_mlx_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=6, suite="base-a-test-cpu") +register_mlx_ci(est_time=6, suite="stage-a-unit-test-mlx") + +_HAS_MLX = ( + importlib.util.find_spec("mlx") is not None + and importlib.util.find_spec("mlx_lm") is not None +) +_SKIP_REASON = "requires mlx + mlx_lm" + +if _HAS_MLX: + import mlx.core as mx + + from sglang.srt.hardware_backend.mlx.models.muse_glimmer_mlx import ( + Model, + ModelArgs, + flatten_rc_config, + ) + +# Tiny architecture: 4 layers so the derived NoPE pattern (last layer NoPE +# with every_n_layers_nope=4) exercises both layer types. +_TINY = dict( + hidden_size=8, + num_hidden_layers=4, + num_attention_heads=2, + num_key_value_heads=1, + head_dim=4, + intermediate_size=16, + vocab_size=32, + every_n_layers_nope=4, + sliding_window=4, + max_position_embeddings=64, +) + +_RC_TEXT_CONFIG = dict( + hidden_size=8, + num_hidden_layers=4, + num_attention_heads=2, + num_key_value_heads=1, + head_dim=4, + intermediate_size=16, + vocab_size=32, + rms_norm_eps=1e-5, + post_norm_eps=1e-8, + max_position_embeddings=64, + qk_scale_factor=0.5, + output_multiplier=0.2, + final_logit_softcapping=20.0, + sliding_window=4, + layer_rope_theta=[500000.0, 500000.0, 500000.0, 0], +) + + +def _raw_weights(args): + """A complete raw HF export with deterministic values.""" + mx.random.seed(0) + H, D, hid = args.num_attention_heads, args.head_dim, args.hidden_size + kv = args.num_key_value_heads * D + weights = { + "model.embed_tokens.weight": mx.random.normal((args.vocab_size, hid)), + "model.norm.weight": mx.random.normal((hid,)), + "lm_head.weight": mx.random.normal((args.vocab_size, hid)), + } + for i in range(args.num_hidden_layers): + p = f"model.layers.{i}." + weights.update( + { + p + "self_attn.q_proj.weight": mx.random.normal((H * D, hid)), + p + "self_attn.k_proj.weight": mx.random.normal((kv, hid)), + p + "self_attn.v_proj.weight": mx.random.normal((kv, hid)), + p + "self_attn.o_proj.weight": mx.random.normal((hid, H * D)), + p + "self_attn.output_gate_proj.weight": mx.random.normal((H * D, hid)), + p + "input_layernorm.weight": mx.full((hid,), 0.10), + p + "post_attn_norm.weight": mx.full((hid,), 0.20), + p + "post_attention_layernorm.weight": mx.full((hid,), 0.30), + p + "post_ffn_norm.weight": mx.full((hid,), 0.40), + p + + "mlp.gate_proj.weight": mx.random.normal( + (args.intermediate_size, hid) + ), + p + + "mlp.up_proj.weight": mx.random.normal((args.intermediate_size, hid)), + p + + "mlp.down_proj.weight": mx.random.normal( + (hid, args.intermediate_size) + ), + } + ) + return weights + + +def _rc_weights(args): + """The same export in the RC multimodal layout (nested prefix, RC names, + a vision tower to be dropped, no output-gate fusion).""" + raw = _raw_weights(args) + rc = {} + renames = { + "self_attn.output_gate_proj.weight": "self_attn.gate_proj.weight", + "post_attn_norm.weight": "post_attention_layernorm.weight", + "post_attention_layernorm.weight": "pre_feedforward_layernorm.weight", + "post_ffn_norm.weight": "post_feedforward_layernorm.weight", + } + for name, w in raw.items(): + if not name.startswith("model."): + rc[name] = w # lm_head stays top-level in the RC layout too + continue + rest = name[len("model.") :] + for raw_suffix, rc_suffix in renames.items(): + if rest.endswith(raw_suffix): + rest = rest[: -len(raw_suffix)] + rc_suffix + break + rc["model.language_model." + rest] = w + rc["model.vision_tower.patch_embed.weight"] = mx.zeros((4, 4)) + return rc + + +@unittest.skipUnless(_HAS_MLX, _SKIP_REASON) +class TestFlattenRcConfig(CustomTestCase): + def test_field_mapping_and_conversions(self): + flat = flatten_rc_config({"text_config": dict(_RC_TEXT_CONFIG)}) + self.assertEqual(flat["model_type"], "onyx") + # RC qk_scale_factor is against SDPA's 1/sqrt(head_dim). + self.assertAlmostEqual(flat["qk_scale_factor"], 0.5 * 4**0.5) + self.assertEqual(flat["output_soft_cap_temp"], 20.0) + # The current vendor export stores q/k in the NeoX rotary + # layout; the RC path always reads that convention. + self.assertIs(flat["rope_is_neox_style"], True) + # normalize_tok_embeddings must stay at the ModelArgs default (True): + # the 20260806 export ships the raw table (needs the runtime norm), + # and on older baked-table exports the scaleless RMS norm is + # idempotent, so always-on covers both generations. + self.assertNotIn("normalize_tok_embeddings", flat) + self.assertIs(ModelArgs.normalize_tok_embeddings, True) + # Zeros in layer_rope_theta mark NoPE layers. + self.assertEqual(flat["no_rope_layers"], [1, 1, 1, 0]) + + def test_non_silu_activation_rejected(self): + cfg = dict(_RC_TEXT_CONFIG, hidden_activation="gelu") + with self.assertRaisesRegex(ValueError, "silu"): + flatten_rc_config({"text_config": cfg}) + + def test_model_args_from_dict_accepts_rc_schema(self): + args = ModelArgs.from_dict({"text_config": dict(_RC_TEXT_CONFIG)}) + self.assertEqual(args.no_rope_layers, [1, 1, 1, 0]) + self.assertEqual( + args.layer_types, + ["sliding_attention"] * 3 + ["full_attention"], + ) + + +@unittest.skipUnless(_HAS_MLX, _SKIP_REASON) +class TestModelArgsValidation(CustomTestCase): + def test_derives_nope_and_layer_types(self): + args = ModelArgs(**_TINY) + self.assertEqual(args.no_rope_layers, [1, 1, 1, 0]) + self.assertEqual( + args.layer_types, + ["sliding_attention"] * 3 + ["full_attention"], + ) + + def test_layer_types_must_match_no_rope_layers(self): + with self.assertRaisesRegex(ValueError, "disagrees"): + ModelArgs(**_TINY, layer_types=["full_attention"] * 4) + + def test_wrong_length_no_rope_layers_rejected(self): + with self.assertRaisesRegex(ValueError, "entries"): + ModelArgs(**_TINY, no_rope_layers=[1, 0]) + + def test_non_binary_no_rope_flags_rejected(self): + with self.assertRaisesRegex(ValueError, "non-binary"): + ModelArgs(**_TINY, no_rope_layers=[1, 1, 2, 0]) + + def test_unknown_format_version_rejected(self): + with self.assertRaisesRegex(ValueError, "onyx_mlx_format"): + ModelArgs(**_TINY, onyx_mlx_format=99) + + +@unittest.skipUnless(_HAS_MLX, _SKIP_REASON) +class TestSanitize(CustomTestCase): + def _model(self, **overrides): + return Model(ModelArgs(**dict(_TINY, **overrides))) + + def test_raw_export_folds_norms_and_fuses_gate(self): + model = self._model() + raw = _raw_weights(model.args) + out = model.sanitize(dict(raw)) + + # Offset norms gain +1.0; the final norm does not. + norm = out["model.layers.0.input_layernorm.weight"] + self.assertTrue(mx.allclose(norm, mx.full(norm.shape, 1.10))) + self.assertTrue(mx.allclose(out["model.norm.weight"], raw["model.norm.weight"])) + + # q/gate interleave is per-head [q_head; gate_head]. + H, D = model.args.num_attention_heads, model.args.head_dim + hid = model.args.hidden_size + fused = out["model.layers.0.self_attn.q_proj.weight"] + self.assertEqual(tuple(fused.shape), (2 * H * D, hid)) + per_head = fused.reshape(H, 2 * D, hid) + q = raw["model.layers.0.self_attn.q_proj.weight"].reshape(H, D, hid) + g = raw["model.layers.0.self_attn.output_gate_proj.weight"].reshape(H, D, hid) + self.assertTrue(mx.allclose(per_head[:, :D, :], q)) + self.assertTrue(mx.allclose(per_head[:, D:, :], g)) + self.assertNotIn("model.layers.0.self_attn.output_gate_proj.weight", out) + + def test_sanitized_raw_weights_load_and_forward(self): + model = self._model() + out = model.sanitize(_raw_weights(model.args)) + model.load_weights(list(out.items())) + logits = model(mx.array([[1, 2, 3]], dtype=mx.int32), cache=model.make_cache()) + self.assertEqual(tuple(logits.shape), (1, 3, model.args.vocab_size)) + self.assertTrue(bool(mx.all(mx.isfinite(logits)))) + + def test_rc_layout_positional_renames(self): + model = self._model() + out = model.sanitize(_rc_weights(model.args)) + # Distinct per-norm constants prove each RC name landed in its + # positional slot (+1 folded): RC post_attention_layernorm -> + # raw post_attn_norm (0.20), RC pre_feedforward_layernorm -> + # raw post_attention_layernorm (0.30). + for raw_name, value in ( + ("input_layernorm", 1.10), + ("post_attn_norm", 1.20), + ("post_attention_layernorm", 1.30), + ("post_ffn_norm", 1.40), + ): + w = out[f"model.layers.0.{raw_name}.weight"] + self.assertTrue( + mx.allclose(w, mx.full(w.shape, value)), + f"{raw_name} expected {value}", + ) + self.assertFalse(any("vision" in k for k in out)) + self.assertFalse(any("language_model" in k for k in out)) + + def test_packaged_artifact_passes_through(self): + model = self._model(onyx_mlx_format=1) + packaged = {"model.embed_tokens.weight": mx.zeros((32, 8))} + self.assertIs(model.sanitize(packaged), packaged) + + def test_packaged_marker_with_raw_keys_rejected(self): + model = self._model(onyx_mlx_format=1) + raw = _raw_weights(ModelArgs(**_TINY)) + with self.assertRaisesRegex(ValueError, "raw-checkpoint keys"): + model.sanitize(raw) + + def test_incomplete_raw_checkpoint_rejected(self): + model = self._model() + raw = _raw_weights(model.args) + del raw["model.layers.0.mlp.up_proj.weight"] + with self.assertRaisesRegex(ValueError, "missing"): + model.sanitize(raw) + + +if __name__ == "__main__": + unittest.main() From 894815f94fc4eef7efbf5a17a86063f504c9e7f1 Mon Sep 17 00:00:00 2001 From: sglang-bot <232288953+sglang-bot@users.noreply.github.com> Date: Mon, 10 Aug 2026 03:23:11 -0700 Subject: [PATCH 02/18] misc fix --- .../mlx/models/muse_glimmer_mlx.py | 18 +++++++------- .../test_muse_glimmer_mlx_correctness.py | 6 ++--- .../mlx/test_mlx_remote_code_gate.py | 24 ++++++++++++------- .../mlx/test_muse_glimmer_mlx_model.py | 2 +- 4 files changed, 29 insertions(+), 21 deletions(-) diff --git a/python/sglang/srt/hardware_backend/mlx/models/muse_glimmer_mlx.py b/python/sglang/srt/hardware_backend/mlx/models/muse_glimmer_mlx.py index 31cba44d0f50..0b0ec10031a8 100644 --- a/python/sglang/srt/hardware_backend/mlx/models/muse_glimmer_mlx.py +++ b/python/sglang/srt/hardware_backend/mlx/models/muse_glimmer_mlx.py @@ -116,7 +116,7 @@ from mlx_lm.models.cache import KVCache # Version of the packaged (fused/folded) weight layout this file understands. -ONYX_MLX_FORMAT_VERSION = 1 +MUSE_GLIMMER_MLX_FORMAT_VERSION = 1 # The four per-layer norms whose checkpoint weight is an offset from 1.0. # model.norm (MuseGlimmerFinalRMSNorm) is NOT in this list and must not be offset. @@ -191,7 +191,7 @@ def flatten_rc_config(config: dict) -> dict: layer_rope_theta = text.get("layer_rope_theta") flat = { - "model_type": "onyx", + "model_type": "muse_glimmer_text", "hidden_size": text["hidden_size"], "num_hidden_layers": text["num_hidden_layers"], "num_attention_heads": text["num_attention_heads"], @@ -218,7 +218,7 @@ def flatten_rc_config(config: dict) -> dict: @dataclass class ModelArgs(BaseModelArgs): - model_type: str = "onyx" + model_type: str = "muse_glimmer_text" hidden_size: int = 6656 num_hidden_layers: int = 52 num_attention_heads: int = 32 @@ -321,11 +321,11 @@ def __post_init__(self): ) if self.onyx_mlx_format is not None and ( - self.onyx_mlx_format != ONYX_MLX_FORMAT_VERSION + self.onyx_mlx_format != MUSE_GLIMMER_MLX_FORMAT_VERSION ): raise ValueError( f"onyx_mlx_format {self.onyx_mlx_format} is not supported by " - f"this model file (expected {ONYX_MLX_FORMAT_VERSION}); " + f"this model file (expected {MUSE_GLIMMER_MLX_FORMAT_VERSION}); " "regenerate the artifact with a matching packager" ) @@ -611,7 +611,7 @@ def _expected_raw_keys(self) -> set: return keys def sanitize(self, weights: dict) -> dict: - if self.args.onyx_mlx_format == ONYX_MLX_FORMAT_VERSION: + if self.args.onyx_mlx_format == MUSE_GLIMMER_MLX_FORMAT_VERSION: # Packaged artifact: weights are already fused/folded. A raw-only # key here means the marker was stamped on the wrong directory. stray = sorted( @@ -622,7 +622,7 @@ def sanitize(self, weights: dict) -> dict: if stray: raise ValueError( "config.json claims a packaged Muse Glimmer MLX artifact " - f"(onyx_mlx_format={ONYX_MLX_FORMAT_VERSION}) but the " + f"(onyx_mlx_format={MUSE_GLIMMER_MLX_FORMAT_VERSION}) but the " f"weights contain raw-checkpoint keys {stray[:4]}" f"{'...' if len(stray) > 4 else ''}; the marker belongs " "on packaged artifacts only — repackage from the raw HF export" @@ -650,7 +650,7 @@ def sanitize(self, weights: dict) -> dict: hint = ( " (weights look already fused: if this is a packaged " 'artifact, its config.json must carry "onyx_mlx_format": ' - f"{ONYX_MLX_FORMAT_VERSION})" + f"{MUSE_GLIMMER_MLX_FORMAT_VERSION})" ) raise ValueError( "not a complete raw Muse Glimmer HF checkpoint: " @@ -676,7 +676,7 @@ def sanitize(self, weights: dict) -> dict: f"raw q_proj.weight has shape {raw_q_shape}, expected " f"({H * D}, {hidden}); a width of {2 * H * D} means the gate " "is already fused — such artifacts must carry " - f'"onyx_mlx_format": {ONYX_MLX_FORMAT_VERSION} in config.json' + f'"onyx_mlx_format": {MUSE_GLIMMER_MLX_FORMAT_VERSION} in config.json' ) new_weights = {} diff --git a/test/registered/mlx/models_e2e/test_muse_glimmer_mlx_correctness.py b/test/registered/mlx/models_e2e/test_muse_glimmer_mlx_correctness.py index 6e644fee4a27..aa993f4928d6 100644 --- a/test/registered/mlx/models_e2e/test_muse_glimmer_mlx_correctness.py +++ b/test/registered/mlx/models_e2e/test_muse_glimmer_mlx_correctness.py @@ -16,7 +16,7 @@ chunkings. The artifact is weight-derived and private: its path comes exclusively from the -``SGLANG_MLX_TEST_ONYX_ARTIFACT`` environment variable and the whole module SKIPS +``SGLANG_MLX_TEST_MUSE_GLIMMER_ARTIFACT`` environment variable and the whole module SKIPS cleanly when it is unset or missing — CI has no dependency on private assets. Window-engagement prompt lengths are derived from the artifact's own config, so the suite also runs against tiny packaged fixtures. @@ -52,7 +52,7 @@ and importlib.util.find_spec("mlx_lm") is not None ) -ARTIFACT = os.environ.get("SGLANG_MLX_TEST_ONYX_ARTIFACT") +ARTIFACT = os.environ.get("SGLANG_MLX_TEST_MUSE_GLIMMER_ARTIFACT") MEM_FRACTION_STATIC = os.environ.get("SGLANG_MLX_TEST_MEM_FRACTION", "0.85") MIN_FREE_GB = float(os.environ.get("SGLANG_MLX_TEST_MIN_FREE_GB", "20")) MAX_NEW_TOKENS = 32 @@ -63,7 +63,7 @@ def _artifact_or_skip() -> Path: raise unittest.SkipTest("requires mlx + mlx_lm (Apple Silicon only)") if not ARTIFACT: raise unittest.SkipTest( - "SGLANG_MLX_TEST_ONYX_ARTIFACT is not set; the Muse Glimmer artifact is private " + "SGLANG_MLX_TEST_MUSE_GLIMMER_ARTIFACT is not set; the Muse Glimmer artifact is private " "and must be supplied via the environment" ) path = Path(ARTIFACT).expanduser() diff --git a/test/registered/unit/hardware_backend/mlx/test_mlx_remote_code_gate.py b/test/registered/unit/hardware_backend/mlx/test_mlx_remote_code_gate.py index 50571ab961fa..b3cd636cd050 100644 --- a/test/registered/unit/hardware_backend/mlx/test_mlx_remote_code_gate.py +++ b/test/registered/unit/hardware_backend/mlx/test_mlx_remote_code_gate.py @@ -62,13 +62,17 @@ def _assert_sentinel_not_executed(self): ) def test_refuses_model_file_without_trust(self): - _make_checkpoint(self.dir, {"model_type": "onyx", "model_file": "evil.py"}) + _make_checkpoint( + self.dir, {"model_type": "muse_glimmer_text", "model_file": "evil.py"} + ) with self.assertRaisesRegex(RemoteCodeGateError, "--trust-remote-code"): ensure_remote_code_allowed(self.dir, trust_remote_code=False) self._assert_sentinel_not_executed() def test_allows_model_file_with_trust(self): - _make_checkpoint(self.dir, {"model_type": "onyx", "model_file": "evil.py"}) + _make_checkpoint( + self.dir, {"model_type": "muse_glimmer_text", "model_file": "evil.py"} + ) ensure_remote_code_allowed(self.dir, trust_remote_code=True) # The gate itself never imports the file either way. self._assert_sentinel_not_executed() @@ -94,7 +98,7 @@ def test_non_object_config_rejected(self): def test_missing_model_file_target_rejected(self): _make_checkpoint( self.dir, - {"model_type": "onyx", "model_file": "nope.py"}, + {"model_type": "muse_glimmer_text", "model_file": "nope.py"}, with_sentinel=False, ) with self.assertRaisesRegex(RemoteCodeGateError, "does not exist"): @@ -103,7 +107,7 @@ def test_missing_model_file_target_rejected(self): def test_absolute_model_file_rejected(self): _make_checkpoint( self.dir, - {"model_type": "onyx", "model_file": "/etc/anything.py"}, + {"model_type": "muse_glimmer_text", "model_file": "/etc/anything.py"}, with_sentinel=False, ) with self.assertRaisesRegex(RemoteCodeGateError, "relative path"): @@ -112,7 +116,7 @@ def test_absolute_model_file_rejected(self): def test_traversal_model_file_rejected(self): _make_checkpoint( self.dir, - {"model_type": "onyx", "model_file": "../outside.py"}, + {"model_type": "muse_glimmer_text", "model_file": "../outside.py"}, with_sentinel=False, ) with self.assertRaisesRegex(RemoteCodeGateError, "relative path"): @@ -121,7 +125,7 @@ def test_traversal_model_file_rejected(self): def test_non_string_model_file_rejected(self): _make_checkpoint( self.dir, - {"model_type": "onyx", "model_file": 42}, + {"model_type": "muse_glimmer_text", "model_file": 42}, with_sentinel=False, ) with self.assertRaisesRegex(RemoteCodeGateError, "non-string"): @@ -146,7 +150,9 @@ def tearDown(self): def test_refusal_precedes_loader_call(self): from sglang.srt.hardware_backend.mlx.model_runner import MlxModelRunner - _make_checkpoint(self.dir, {"model_type": "onyx", "model_file": "evil.py"}) + _make_checkpoint( + self.dir, {"model_type": "muse_glimmer_text", "model_file": "evil.py"} + ) with patch( "sglang.srt.hardware_backend.mlx.model_runner.mlx_lm_load" ) as loader: @@ -158,7 +164,9 @@ def test_refusal_precedes_loader_call(self): def test_trusted_load_uses_resolved_directory(self): from sglang.srt.hardware_backend.mlx.model_runner import MlxModelRunner - _make_checkpoint(self.dir, {"model_type": "onyx", "model_file": "evil.py"}) + _make_checkpoint( + self.dir, {"model_type": "muse_glimmer_text", "model_file": "evil.py"} + ) with patch( "sglang.srt.hardware_backend.mlx.model_runner.mlx_lm_load", side_effect=self._StopInit, diff --git a/test/registered/unit/hardware_backend/mlx/test_muse_glimmer_mlx_model.py b/test/registered/unit/hardware_backend/mlx/test_muse_glimmer_mlx_model.py index 5f1524c3ad97..f2b8b1cf7042 100644 --- a/test/registered/unit/hardware_backend/mlx/test_muse_glimmer_mlx_model.py +++ b/test/registered/unit/hardware_backend/mlx/test_muse_glimmer_mlx_model.py @@ -143,7 +143,7 @@ def _rc_weights(args): class TestFlattenRcConfig(CustomTestCase): def test_field_mapping_and_conversions(self): flat = flatten_rc_config({"text_config": dict(_RC_TEXT_CONFIG)}) - self.assertEqual(flat["model_type"], "onyx") + self.assertEqual(flat["model_type"], "muse_glimmer_text") # RC qk_scale_factor is against SDPA's 1/sqrt(head_dim). self.assertAlmostEqual(flat["qk_scale_factor"], 0.5 * 4**0.5) self.assertEqual(flat["output_soft_cap_temp"], 20.0) From caa91b7d01ca610bb2e37bc095af6e427aa7e358 Mon Sep 17 00:00:00 2001 From: sglang-bot <232288953+sglang-bot@users.noreply.github.com> Date: Mon, 10 Aug 2026 03:30:53 -0700 Subject: [PATCH 03/18] align mlx config keys with published artifact --- .../mlx/models/muse_glimmer_mlx.py | 24 +++++++++---------- .../test_muse_glimmer_mlx_correctness.py | 2 +- .../mlx/test_mlx_remote_code_gate.py | 16 ++++++------- .../mlx/test_muse_glimmer_mlx_model.py | 10 ++++---- 4 files changed, 26 insertions(+), 26 deletions(-) diff --git a/python/sglang/srt/hardware_backend/mlx/models/muse_glimmer_mlx.py b/python/sglang/srt/hardware_backend/mlx/models/muse_glimmer_mlx.py index 0b0ec10031a8..398088fe2498 100644 --- a/python/sglang/srt/hardware_backend/mlx/models/muse_glimmer_mlx.py +++ b/python/sglang/srt/hardware_backend/mlx/models/muse_glimmer_mlx.py @@ -14,7 +14,7 @@ """Muse Glimmer (dense, text-only) for mlx-lm. Loaded via mlx-lm's custom-architecture path: ship this file in the checkpoint -directory as ``mlx_onyx.py``, set ``"model_file": "mlx_onyx.py"`` in +directory as ``muse_glimmer_mlx.py``, set ``"model_file": "muse_glimmer_mlx.py"`` in ``config.json``. This copy under ``sglang/srt/hardware_backend/mlx/models/`` is the maintained source; artifacts ship a byte-identical copy. It must stay importable standalone (mlx / mlx-lm imports only — no sglang imports), @@ -90,7 +90,7 @@ but an older-generation export served through this path gets the wrong rope layout (their configs are byte-identical; prefer repackaging). * **Packaged MLX artifact**: already fused/folded, marked by - ``"onyx_mlx_format": 1`` in ``config.json`` (stamped at packaging time + ``"muse_glimmer_mlx_format": 1`` in ``config.json`` (stamped at packaging time only, never present on raw HF exports). Passed through untouched. Config schemas. ``ModelArgs.from_dict`` accepts the flat schema written at @@ -191,7 +191,7 @@ def flatten_rc_config(config: dict) -> dict: layer_rope_theta = text.get("layer_rope_theta") flat = { - "model_type": "muse_glimmer_text", + "model_type": "muse_glimmer", "hidden_size": text["hidden_size"], "num_hidden_layers": text["num_hidden_layers"], "num_attention_heads": text["num_attention_heads"], @@ -218,7 +218,7 @@ def flatten_rc_config(config: dict) -> dict: @dataclass class ModelArgs(BaseModelArgs): - model_type: str = "muse_glimmer_text" + model_type: str = "muse_glimmer" hidden_size: int = 6656 num_hidden_layers: int = 52 num_attention_heads: int = 32 @@ -243,7 +243,7 @@ class ModelArgs(BaseModelArgs): layer_types: Optional[List[str]] = None # Set on saved MLX artifacts at packaging time (never on raw HF # exports); tells sanitize() the weights are already fused/folded. - onyx_mlx_format: Optional[int] = None + muse_glimmer_mlx_format: Optional[int] = None @classmethod def from_dict(cls, params): @@ -320,11 +320,11 @@ def __post_init__(self): f"{mismatches}" ) - if self.onyx_mlx_format is not None and ( - self.onyx_mlx_format != MUSE_GLIMMER_MLX_FORMAT_VERSION + if self.muse_glimmer_mlx_format is not None and ( + self.muse_glimmer_mlx_format != MUSE_GLIMMER_MLX_FORMAT_VERSION ): raise ValueError( - f"onyx_mlx_format {self.onyx_mlx_format} is not supported by " + f"muse_glimmer_mlx_format {self.muse_glimmer_mlx_format} is not supported by " f"this model file (expected {MUSE_GLIMMER_MLX_FORMAT_VERSION}); " "regenerate the artifact with a matching packager" ) @@ -611,7 +611,7 @@ def _expected_raw_keys(self) -> set: return keys def sanitize(self, weights: dict) -> dict: - if self.args.onyx_mlx_format == MUSE_GLIMMER_MLX_FORMAT_VERSION: + if self.args.muse_glimmer_mlx_format == MUSE_GLIMMER_MLX_FORMAT_VERSION: # Packaged artifact: weights are already fused/folded. A raw-only # key here means the marker was stamped on the wrong directory. stray = sorted( @@ -622,7 +622,7 @@ def sanitize(self, weights: dict) -> dict: if stray: raise ValueError( "config.json claims a packaged Muse Glimmer MLX artifact " - f"(onyx_mlx_format={MUSE_GLIMMER_MLX_FORMAT_VERSION}) but the " + f"(muse_glimmer_mlx_format={MUSE_GLIMMER_MLX_FORMAT_VERSION}) but the " f"weights contain raw-checkpoint keys {stray[:4]}" f"{'...' if len(stray) > 4 else ''}; the marker belongs " "on packaged artifacts only — repackage from the raw HF export" @@ -649,7 +649,7 @@ def sanitize(self, weights: dict) -> dict: if gate_missing and not unexpected: hint = ( " (weights look already fused: if this is a packaged " - 'artifact, its config.json must carry "onyx_mlx_format": ' + 'artifact, its config.json must carry "muse_glimmer_mlx_format": ' f"{MUSE_GLIMMER_MLX_FORMAT_VERSION})" ) raise ValueError( @@ -676,7 +676,7 @@ def sanitize(self, weights: dict) -> dict: f"raw q_proj.weight has shape {raw_q_shape}, expected " f"({H * D}, {hidden}); a width of {2 * H * D} means the gate " "is already fused — such artifacts must carry " - f'"onyx_mlx_format": {MUSE_GLIMMER_MLX_FORMAT_VERSION} in config.json' + f'"muse_glimmer_mlx_format": {MUSE_GLIMMER_MLX_FORMAT_VERSION} in config.json' ) new_weights = {} diff --git a/test/registered/mlx/models_e2e/test_muse_glimmer_mlx_correctness.py b/test/registered/mlx/models_e2e/test_muse_glimmer_mlx_correctness.py index aa993f4928d6..47b6cf3448b8 100644 --- a/test/registered/mlx/models_e2e/test_muse_glimmer_mlx_correctness.py +++ b/test/registered/mlx/models_e2e/test_muse_glimmer_mlx_correctness.py @@ -70,7 +70,7 @@ def _artifact_or_skip() -> Path: if not (path / "config.json").is_file(): raise unittest.SkipTest(f"no packaged artifact at {path}") config = json.loads((path / "config.json").read_text()) - if config.get("onyx_mlx_format") != 1: + if config.get("muse_glimmer_mlx_format") != 1: raise unittest.SkipTest( f"{path} is not a packaged Muse Glimmer MLX artifact (missing marker)" ) diff --git a/test/registered/unit/hardware_backend/mlx/test_mlx_remote_code_gate.py b/test/registered/unit/hardware_backend/mlx/test_mlx_remote_code_gate.py index b3cd636cd050..8d44d12d0d2b 100644 --- a/test/registered/unit/hardware_backend/mlx/test_mlx_remote_code_gate.py +++ b/test/registered/unit/hardware_backend/mlx/test_mlx_remote_code_gate.py @@ -63,7 +63,7 @@ def _assert_sentinel_not_executed(self): def test_refuses_model_file_without_trust(self): _make_checkpoint( - self.dir, {"model_type": "muse_glimmer_text", "model_file": "evil.py"} + self.dir, {"model_type": "muse_glimmer", "model_file": "evil.py"} ) with self.assertRaisesRegex(RemoteCodeGateError, "--trust-remote-code"): ensure_remote_code_allowed(self.dir, trust_remote_code=False) @@ -71,7 +71,7 @@ def test_refuses_model_file_without_trust(self): def test_allows_model_file_with_trust(self): _make_checkpoint( - self.dir, {"model_type": "muse_glimmer_text", "model_file": "evil.py"} + self.dir, {"model_type": "muse_glimmer", "model_file": "evil.py"} ) ensure_remote_code_allowed(self.dir, trust_remote_code=True) # The gate itself never imports the file either way. @@ -98,7 +98,7 @@ def test_non_object_config_rejected(self): def test_missing_model_file_target_rejected(self): _make_checkpoint( self.dir, - {"model_type": "muse_glimmer_text", "model_file": "nope.py"}, + {"model_type": "muse_glimmer", "model_file": "nope.py"}, with_sentinel=False, ) with self.assertRaisesRegex(RemoteCodeGateError, "does not exist"): @@ -107,7 +107,7 @@ def test_missing_model_file_target_rejected(self): def test_absolute_model_file_rejected(self): _make_checkpoint( self.dir, - {"model_type": "muse_glimmer_text", "model_file": "/etc/anything.py"}, + {"model_type": "muse_glimmer", "model_file": "/etc/anything.py"}, with_sentinel=False, ) with self.assertRaisesRegex(RemoteCodeGateError, "relative path"): @@ -116,7 +116,7 @@ def test_absolute_model_file_rejected(self): def test_traversal_model_file_rejected(self): _make_checkpoint( self.dir, - {"model_type": "muse_glimmer_text", "model_file": "../outside.py"}, + {"model_type": "muse_glimmer", "model_file": "../outside.py"}, with_sentinel=False, ) with self.assertRaisesRegex(RemoteCodeGateError, "relative path"): @@ -125,7 +125,7 @@ def test_traversal_model_file_rejected(self): def test_non_string_model_file_rejected(self): _make_checkpoint( self.dir, - {"model_type": "muse_glimmer_text", "model_file": 42}, + {"model_type": "muse_glimmer", "model_file": 42}, with_sentinel=False, ) with self.assertRaisesRegex(RemoteCodeGateError, "non-string"): @@ -151,7 +151,7 @@ def test_refusal_precedes_loader_call(self): from sglang.srt.hardware_backend.mlx.model_runner import MlxModelRunner _make_checkpoint( - self.dir, {"model_type": "muse_glimmer_text", "model_file": "evil.py"} + self.dir, {"model_type": "muse_glimmer", "model_file": "evil.py"} ) with patch( "sglang.srt.hardware_backend.mlx.model_runner.mlx_lm_load" @@ -165,7 +165,7 @@ def test_trusted_load_uses_resolved_directory(self): from sglang.srt.hardware_backend.mlx.model_runner import MlxModelRunner _make_checkpoint( - self.dir, {"model_type": "muse_glimmer_text", "model_file": "evil.py"} + self.dir, {"model_type": "muse_glimmer", "model_file": "evil.py"} ) with patch( "sglang.srt.hardware_backend.mlx.model_runner.mlx_lm_load", diff --git a/test/registered/unit/hardware_backend/mlx/test_muse_glimmer_mlx_model.py b/test/registered/unit/hardware_backend/mlx/test_muse_glimmer_mlx_model.py index f2b8b1cf7042..6c69c40738cf 100644 --- a/test/registered/unit/hardware_backend/mlx/test_muse_glimmer_mlx_model.py +++ b/test/registered/unit/hardware_backend/mlx/test_muse_glimmer_mlx_model.py @@ -143,7 +143,7 @@ def _rc_weights(args): class TestFlattenRcConfig(CustomTestCase): def test_field_mapping_and_conversions(self): flat = flatten_rc_config({"text_config": dict(_RC_TEXT_CONFIG)}) - self.assertEqual(flat["model_type"], "muse_glimmer_text") + self.assertEqual(flat["model_type"], "muse_glimmer") # RC qk_scale_factor is against SDPA's 1/sqrt(head_dim). self.assertAlmostEqual(flat["qk_scale_factor"], 0.5 * 4**0.5) self.assertEqual(flat["output_soft_cap_temp"], 20.0) @@ -196,8 +196,8 @@ def test_non_binary_no_rope_flags_rejected(self): ModelArgs(**_TINY, no_rope_layers=[1, 1, 2, 0]) def test_unknown_format_version_rejected(self): - with self.assertRaisesRegex(ValueError, "onyx_mlx_format"): - ModelArgs(**_TINY, onyx_mlx_format=99) + with self.assertRaisesRegex(ValueError, "muse_glimmer_mlx_format"): + ModelArgs(**_TINY, muse_glimmer_mlx_format=99) @unittest.skipUnless(_HAS_MLX, _SKIP_REASON) @@ -257,12 +257,12 @@ def test_rc_layout_positional_renames(self): self.assertFalse(any("language_model" in k for k in out)) def test_packaged_artifact_passes_through(self): - model = self._model(onyx_mlx_format=1) + model = self._model(muse_glimmer_mlx_format=1) packaged = {"model.embed_tokens.weight": mx.zeros((32, 8))} self.assertIs(model.sanitize(packaged), packaged) def test_packaged_marker_with_raw_keys_rejected(self): - model = self._model(onyx_mlx_format=1) + model = self._model(muse_glimmer_mlx_format=1) raw = _raw_weights(ModelArgs(**_TINY)) with self.assertRaisesRegex(ValueError, "raw-checkpoint keys"): model.sanitize(raw) From d3c5493ae20d7b4ddd939643160d5da1ef7af5db Mon Sep 17 00:00:00 2001 From: Brayden Zhong Date: Mon, 10 Aug 2026 10:56:34 +0000 Subject: [PATCH 04/18] Skip mxfp8_gemm in FlashInfer autotune --- python/sglang/srt/model_executor/runner/flashinfer_autotune.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/model_executor/runner/flashinfer_autotune.py b/python/sglang/srt/model_executor/runner/flashinfer_autotune.py index f01256fd4124..170c51b63eb7 100644 --- a/python/sglang/srt/model_executor/runner/flashinfer_autotune.py +++ b/python/sglang/srt/model_executor/runner/flashinfer_autotune.py @@ -33,7 +33,8 @@ logger = logging.getLogger(__name__) -FLASHINFER_AUTOTUNE_WORKAROUND_SKIPS = frozenset() +# TODO: Remove after FlashInfer fixes the mxfp8_gemm autotuning IMA. +FLASHINFER_AUTOTUNE_WORKAROUND_SKIPS = frozenset({"mxfp8_gemm"}) def get_flashinfer_autotune_skip_ops(model_runner: ModelRunner) -> set[str]: From 561afe856f30e89571cee273196ff61dde6c3b6c Mon Sep 17 00:00:00 2001 From: Brayden Zhong Date: Mon, 10 Aug 2026 11:04:57 +0000 Subject: [PATCH 05/18] Revert "Skip mxfp8_gemm in FlashInfer autotune" This reverts commit d3c5493ae20d7b4ddd939643160d5da1ef7af5db. --- python/sglang/srt/model_executor/runner/flashinfer_autotune.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/python/sglang/srt/model_executor/runner/flashinfer_autotune.py b/python/sglang/srt/model_executor/runner/flashinfer_autotune.py index 170c51b63eb7..f01256fd4124 100644 --- a/python/sglang/srt/model_executor/runner/flashinfer_autotune.py +++ b/python/sglang/srt/model_executor/runner/flashinfer_autotune.py @@ -33,8 +33,7 @@ logger = logging.getLogger(__name__) -# TODO: Remove after FlashInfer fixes the mxfp8_gemm autotuning IMA. -FLASHINFER_AUTOTUNE_WORKAROUND_SKIPS = frozenset({"mxfp8_gemm"}) +FLASHINFER_AUTOTUNE_WORKAROUND_SKIPS = frozenset() def get_flashinfer_autotune_skip_ops(model_runner: ModelRunner) -> set[str]: From 10c0a709d49e696a58b0b259e8ac12c18acb1cae Mon Sep 17 00:00:00 2001 From: Brayden Zhong Date: Mon, 10 Aug 2026 11:11:03 +0000 Subject: [PATCH 06/18] Keep MXFP8 dense GEMM on cutlass for SM120 flashinfer_mxfp8_blockscaled_linear swaps cutlass -> cute-dsl at M <= 64 for the small-M speedup, but CuTe-DSL has no mm_mxfp8 kernel at capability 120, so the swap raises BackendSupportedError. Prefill CUDA graph capture uses small batch shapes, so an MXFP8-containing checkpoint (e.g. the MIXED_PRECISION NVFP4 build, whose down_proj is MXFP8) fails at startup on SM120. Running FlashInfer autotune over mxfp8_gemm happened to mask this, so it only surfaces when that op is skipped or autotune is off. Gate the swap on SM100, which keeps the speedup where CuTe-DSL has the kernel. --- python/sglang/srt/layers/quantization/fp8_utils.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/layers/quantization/fp8_utils.py b/python/sglang/srt/layers/quantization/fp8_utils.py index e45d6a44034c..f5a044969655 100755 --- a/python/sglang/srt/layers/quantization/fp8_utils.py +++ b/python/sglang/srt/layers/quantization/fp8_utils.py @@ -1250,7 +1250,8 @@ def flashinfer_mxfp8_blockscaled_linear( # At small M the persistent CUTLASS kernel is 2-5x slower than the # CuTe-DSL swap-AB/split-K kernels (both consume the same swizzled # 1D scales). - if backend == "cutlass" and q_input.shape[0] <= 64: + # CuTe-DSL has no mm_mxfp8 kernel on SM120, so the swap is SM100-only there. + if backend == "cutlass" and q_input.shape[0] <= 64 and _is_sm100_supported: backend = "cute-dsl" if backend == "trtllm": From 9798994498f71b1c5381af1bc646d4f01f40ad9a Mon Sep 17 00:00:00 2001 From: Brayden Zhong Date: Mon, 10 Aug 2026 11:43:34 +0000 Subject: [PATCH 07/18] Load the published Muse Glimmer GGUFs The published GGUF comes from a newer converter than the reader assumed: * attention.post_norm_rms_epsilon, attention.scale and the bos/eos token ids are no longer emitted, and output_multiplier is now logit_scale. Each is an architecture constant MuseGlimmerConfig already defaults to, so treat them as optional instead of required. * the attention output gate ships as attn_gate, not attn_output_gate. Keying on the old name left use_attn_output_gate False and dropped 52 tensors -- no error, just wrong logits. * GGUF_HF_NAME_MAP_BUILDERS is looked up by HF config.model_type, not by the GGUF architecture. Those strings were both "onyx" before the rename and now differ ("muse_glimmer" vs "muse-glimmer"), so the table needs the HF spelling. * a speculative draft given as a Hub .gguf reference was never resolved to a local path, unlike the target. Checked against meta-models/Muse-Glimmer-30B's config.json: all 16 config values match, and all 731 tensors resolve (627 mapped, 104 non-parametric qk-norm skips, 0 unmapped). --- python/sglang/srt/configs/muse_glimmer.py | 42 +++++++++++++++---- .../sglang/srt/model_loader/gguf_name_maps.py | 7 +++- python/sglang/srt/server_args.py | 26 +++++++++--- 3 files changed, 58 insertions(+), 17 deletions(-) diff --git a/python/sglang/srt/configs/muse_glimmer.py b/python/sglang/srt/configs/muse_glimmer.py index ffdddfb76f06..2266fbd6be2d 100644 --- a/python/sglang/srt/configs/muse_glimmer.py +++ b/python/sglang/srt/configs/muse_glimmer.py @@ -209,6 +209,18 @@ def muse_glimmer_vision_config_kwargs_from_hf( return kwargs +def _f(v): + return None if v is None else float(v) + + +def _i(v): + return None if v is None else int(v) + + +def _mul_sqrt(v, head_dim): + return None if v is None else float(v) * math.sqrt(head_dim) + + def muse_glimmer_config_kwargs_from_gguf(gguf_path: str) -> Dict[str, Any]: from gguf import GGUFReader @@ -220,6 +232,10 @@ def muse_glimmer_config_kwargs_from_gguf(gguf_path: str) -> Dict[str, Any]: def get(suffix): return meta[f"{_ARCH}.{suffix}"] + def opt(suffix): + """None when this converter generation did not emit the key.""" + return meta.get(f"{_ARCH}.{suffix}") + head_dim = int(get("attention.key_length")) swa_pattern = [bool(x) for x in get("attention.sliding_window_pattern")] @@ -234,23 +250,31 @@ def get(suffix): head_dim=head_dim, max_position_embeddings=int(get("context_length")), rms_norm_eps=float(get("attention.layer_norm_rms_epsilon")), - post_norm_eps=float(get("attention.post_norm_rms_epsilon")), rope_theta=float(get("rope.freq_base")), sliding_window=int(get("attention.sliding_window")), layer_types=[ "sliding_attention" if s else "full_attention" for s in swa_pattern ], no_rope_layers=[1 if s else 0 for s in swa_pattern], - qk_scale_factor=float(get("attention.scale")) * math.sqrt(head_dim), - output_multiplier=float(get("output_multiplier")), - output_soft_cap_temp=float(get("final_logit_softcapping")), use_qk_norm=any(n.endswith("attn_q_norm.weight") for n in tensor_names), - use_attn_output_gate=any( - n.endswith("attn_output_gate.weight") for n in tensor_names - ), + use_attn_output_gate=any(n.endswith("attn_gate.weight") for n in tensor_names), tie_word_embeddings="output.weight" not in tensor_names, - bos_token_id=int(meta["tokenizer.ggml.bos_token_id"]), - eos_token_id=int(meta["tokenizer.ggml.eos_token_id"]), architectures=["MuseGlimmerForCausalLM"], dtype="bfloat16", + # Converter generations differ in which of these they emit, and every one + # is an architecture constant that MuseGlimmerConfig already defaults to, + # so an absent key falls back rather than raising. attention.scale is + # stored pre-divided by sqrt(head_dim); the class stores it before that. + **{ + k: v + for k, v in ( + ("post_norm_eps", _f(opt("attention.post_norm_rms_epsilon"))), + ("qk_scale_factor", _mul_sqrt(opt("attention.scale"), head_dim)), + ("output_multiplier", _f(opt("logit_scale"))), + ("output_soft_cap_temp", _f(opt("final_logit_softcapping"))), + ("bos_token_id", _i(meta.get("tokenizer.ggml.bos_token_id"))), + ("eos_token_id", _i(meta.get("tokenizer.ggml.eos_token_id"))), + ) + if v is not None + }, ) diff --git a/python/sglang/srt/model_loader/gguf_name_maps.py b/python/sglang/srt/model_loader/gguf_name_maps.py index 4aefaca681c7..434f337af363 100644 --- a/python/sglang/srt/model_loader/gguf_name_maps.py +++ b/python/sglang/srt/model_loader/gguf_name_maps.py @@ -38,7 +38,7 @@ "attn_k": "self_attn.k_proj", "attn_v": "self_attn.v_proj", "attn_output": "self_attn.o_proj", - "attn_output_gate": "self_attn.output_gate_proj", + "attn_gate": "self_attn.output_gate_proj", "ffn_gate": "mlp.gate_proj", "ffn_up": "mlp.up_proj", "ffn_down": "mlp.down_proj", @@ -64,6 +64,9 @@ def build_muse_glimmer_name_map(config: PretrainedConfig) -> Dict[str, str]: return name_map +# Keyed by HF ``config.model_type`` (loader.py looks it up with that), which is +# not the GGUF ``general.architecture`` that GGUF_NATIVE_CONFIG_BUILDERS uses: +# llama.cpp spells the arch "muse-glimmer" while the HF config says "muse_glimmer". GGUF_HF_NAME_MAP_BUILDERS: Dict[str, Callable[[PretrainedConfig], Dict[str, str]]] = { - "muse-glimmer": build_muse_glimmer_name_map, + "muse_glimmer": build_muse_glimmer_name_map, } diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 8904d6f4b167..85e03cd172bb 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -7321,12 +7321,26 @@ def _resolve_hf_gguf_model_path(self): from sglang.srt.utils.hf_transformers_utils import resolve_hf_gguf_reference resolved = resolve_hf_gguf_reference(self.model_path, revision=self.revision) - if resolved is None: - return - logger.info("Resolved GGUF %s -> %s", self.model_path, resolved) - if self.tokenizer_path == self.model_path: - self.tokenizer_path = resolved - self.model_path = resolved + if resolved is not None: + logger.info("Resolved GGUF %s -> %s", self.model_path, resolved) + if self.tokenizer_path == self.model_path: + self.tokenizer_path = resolved + self.model_path = resolved + + # A speculative draft can be a .gguf too, and it is loaded by path, so it + # needs the same Hub-reference resolution as the target. + if self.speculative_draft_model_path: + resolved_draft = resolve_hf_gguf_reference( + self.speculative_draft_model_path, + revision=self.speculative_draft_model_revision, + ) + if resolved_draft is not None: + logger.info( + "Resolved draft GGUF %s -> %s", + self.speculative_draft_model_path, + resolved_draft, + ) + self.speculative_draft_model_path = resolved_draft def _handle_load_format(self): # The quantization side of the gguf coupling moved to the pipeline From 82505b051f42e1aeaf137646ec4d0bd1245f42e5 Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Mon, 10 Aug 2026 21:58:30 -0700 Subject: [PATCH 08/18] drop duplicated skip_special_tokens field and write --- python/sglang/srt/entrypoints/openai/protocol.py | 1 - python/sglang/srt/entrypoints/openai/serving_responses.py | 5 ----- 2 files changed, 6 deletions(-) diff --git a/python/sglang/srt/entrypoints/openai/protocol.py b/python/sglang/srt/entrypoints/openai/protocol.py index a3ea690656b6..ecb458fc603a 100644 --- a/python/sglang/srt/entrypoints/openai/protocol.py +++ b/python/sglang/srt/entrypoints/openai/protocol.py @@ -1937,7 +1937,6 @@ class MessageProcessingResult: tool_call_constraint: Optional[ToolCallConstraint] = None skip_special_tokens: bool = True require_reasoning: bool = False - skip_special_tokens: bool = True class ToolCallProcessingResult(NamedTuple): diff --git a/python/sglang/srt/entrypoints/openai/serving_responses.py b/python/sglang/srt/entrypoints/openai/serving_responses.py index ee00e6a9f81b..fbd83586ea58 100644 --- a/python/sglang/srt/entrypoints/openai/serving_responses.py +++ b/python/sglang/srt/entrypoints/openai/serving_responses.py @@ -394,11 +394,6 @@ async def create_responses( else None ), ) - if processed_messages is not None: - sampling_params["skip_special_tokens"] = ( - processed_messages.skip_special_tokens - ) - # _process_messages set skip_special_tokens on a chat_request # we then discard, so re-apply it to the engine sampling dict. if processed_messages is not None and ( From e322e0337320a2f22ea6ee420732e424aef1265f Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Mon, 10 Aug 2026 22:35:53 -0700 Subject: [PATCH 09/18] derive dflash sliding-attention causality from config --- python/sglang/srt/models/dflash.py | 29 ++++++++++++++++++----------- 1 file changed, 18 insertions(+), 11 deletions(-) diff --git a/python/sglang/srt/models/dflash.py b/python/sglang/srt/models/dflash.py index bbd70c1ebf18..85198c4defcb 100644 --- a/python/sglang/srt/models/dflash.py +++ b/python/sglang/srt/models/dflash.py @@ -43,6 +43,20 @@ logger = logging.getLogger(__name__) +def _get_dflash_attention_type(config) -> AttentionType: + """Causality of a DFLASH draft, from the checkpoint. + + DFLASH drafts attend over the whole draft block, so they are bidirectional + unless the checkpoint declares itself causal. + """ + text_config = getattr(config, "text_config", None) or config + return ( + AttentionType.DECODER + if getattr(text_config, "is_causal", False) + else AttentionType.ENCODER_ONLY + ) + + def _get_dflash_layer_attention_params( config, layer_id: int ) -> Tuple[int, AttentionType]: @@ -57,20 +71,13 @@ def _get_dflash_layer_attention_params( layer_type = layer_types[layer_id] if layer_type == "full_attention": - text_config = getattr(config, "text_config", None) or config - attention_type = ( - AttentionType.DECODER - if getattr(text_config, "is_causal", False) - else AttentionType.ENCODER_ONLY - ) - return -1, attention_type + return -1, _get_dflash_attention_type(config) if layer_type == "sliding_attention": - # DFlash uses non-causal attention over the draft block on every layer -- - # the reference sets it on the whole draft context, so a sliding layer is - # windowed but still bidirectional (masking only p1 - p0 >= sliding_window). + # Windowing is orthogonal to causality: the layer masks only + # p1 - p0 >= sliding_window and keeps the checkpoint's causality. sliding_window_size = get_dflash_attention_sliding_window_size(config) assert sliding_window_size is not None - return sliding_window_size, AttentionType.ENCODER_ONLY + return sliding_window_size, _get_dflash_attention_type(config) raise ValueError( "Unsupported DFLASH draft layer type. " f"layer_types[{layer_id}]={layer_type!r}." From 3e6c2e88b4ed734b5b00e7f47518775576252efa Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Mon, 10 Aug 2026 22:53:40 -0700 Subject: [PATCH 10/18] move draft vocab_size default into MuseGlimmerAssistantConfig --- python/sglang/srt/configs/model_config.py | 12 +++++------- python/sglang/srt/configs/muse_glimmer.py | 3 +++ 2 files changed, 8 insertions(+), 7 deletions(-) diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index 6c0ac8c53145..c34cdd2215f5 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -1043,13 +1043,11 @@ def _derive_model_shapes(self): self.num_nextn_predict_layers = getattr( self.hf_text_config, "num_nextn_predict_layers", None ) - # DFlash drafts have no vocab of their own. - if self.is_draft_model and not hasattr(self.hf_text_config, "vocab_size"): - self.vocab_size = None - else: - self.vocab_size = self.hf_text_config.vocab_size - if _hf_arch(self.hf_config) == "GlmImageForConditionalGeneration": - self.vocab_size = self.hf_text_config.vision_vocab_size + self.vocab_size = self.hf_text_config.vocab_size + # GLM-Image is the only model here whose output head predicts vision tokens. + # Use vision_vocab_size for lm_head, LogitsProcessor, and graph-mode logits buffers. + if _hf_arch(self.hf_config) == "GlmImageForConditionalGeneration": + self.vocab_size = self.hf_text_config.vision_vocab_size def _init_mla_scaling(self, rope_scaling: Optional[dict]) -> None: """Base MLA attention scale from the head dims, then the rope mscale.""" diff --git a/python/sglang/srt/configs/muse_glimmer.py b/python/sglang/srt/configs/muse_glimmer.py index 2266fbd6be2d..f07de5ea6267 100644 --- a/python/sglang/srt/configs/muse_glimmer.py +++ b/python/sglang/srt/configs/muse_glimmer.py @@ -28,6 +28,9 @@ class MuseGlimmerAssistantConfig(PretrainedConfig): model_type = "muse_glimmer_assistant" + # The DFlash draft has no head of its own; the target's vocab is borrowed + # when the draft worker is built. + vocab_size = None class MuseGlimmerVisionConfig(PretrainedConfig): From da64661d3abe7c55f2366e98aa1cbc50c8a1fbb2 Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Mon, 10 Aug 2026 23:10:42 -0700 Subject: [PATCH 11/18] sync kv dtype tag on fa4 override; tolerate mlx-lm resolver rename; flush stream parsers once; drop redundant write --- .../entrypoints/openai/serving_responses.py | 9 ++++++-- .../srt/function_call/base_format_detector.py | 4 +++- .../hardware_backend/mlx/remote_code_gate.py | 21 ++++++++++++++++--- python/sglang/srt/mem_cache/kv_cache_dtype.py | 4 ++++ 4 files changed, 32 insertions(+), 6 deletions(-) diff --git a/python/sglang/srt/entrypoints/openai/serving_responses.py b/python/sglang/srt/entrypoints/openai/serving_responses.py index fbd83586ea58..aede733587eb 100644 --- a/python/sglang/srt/entrypoints/openai/serving_responses.py +++ b/python/sglang/srt/entrypoints/openai/serving_responses.py @@ -591,7 +591,6 @@ async def _make_request( is_multimodal = self.tokenizer_manager.model_config.is_multimodal processed_messages = self._process_messages(chat_request, is_multimodal) - processed_messages.skip_special_tokens = chat_request.skip_special_tokens if is_multimodal: request_prompts = [processed_messages.prompt] @@ -2008,6 +2007,7 @@ def _sanitize_response_dict(d: dict) -> dict: total_tokens_meta = 0 reasoning_tokens_meta = 0 finish_reason: Optional[dict[str, Any]] = None + flushed = False stream_offset = 0 incremental = self.tokenizer_manager.server_args.incremental_streaming_output @@ -2215,9 +2215,14 @@ def _close_tool_call_state(tool_index: int): stream_offset = len(text) if not delta and finish_reason is None: continue + # finish_reason is sticky, so gate on `flushed` to drain the + # parsers exactly once no matter how many terminal chunks arrive. flush = ( - finish_reason is not None and finish_reason.get("type") != "abort" + not flushed + and finish_reason is not None + and finish_reason.get("type") != "abort" ) + flushed = flushed or flush if reasoning_parser_obj is not None: reasoning_chunk, delta = reasoning_parser_obj.parse_stream_chunk( diff --git a/python/sglang/srt/function_call/base_format_detector.py b/python/sglang/srt/function_call/base_format_detector.py index dd848402efb6..48b6b9dea04f 100644 --- a/python/sglang/srt/function_call/base_format_detector.py +++ b/python/sglang/srt/function_call/base_format_detector.py @@ -351,10 +351,12 @@ def has_tool_call(self, text: str) -> bool: raise NotImplementedError() def finish(self, tools: List[Tool]) -> StreamingParseResult: - """Called once when the stream ends; flush any buffered state. + """Called when the stream ends; flush any buffered state. Detectors that hold text back while waiting for a marker that can no longer arrive (the stream is over) override this to release it. + Overrides should be idempotent: there is nothing left to release the + second time, and callers see several chunks carrying a finish reason. """ return StreamingParseResult() diff --git a/python/sglang/srt/hardware_backend/mlx/remote_code_gate.py b/python/sglang/srt/hardware_backend/mlx/remote_code_gate.py index 51c7cc1df966..49e42cf57035 100644 --- a/python/sglang/srt/hardware_backend/mlx/remote_code_gate.py +++ b/python/sglang/srt/hardware_backend/mlx/remote_code_gate.py @@ -45,11 +45,26 @@ def resolve_model_directory(model_path: str, revision: Optional[str] = None) -> Uses mlx-lm's resolver so the directory is byte-identical to what a direct ``mlx_lm.load`` call would consume; existing local paths are returned as-is (no network access). mlx-lm 0.31.x exposes this as - ``mlx_lm.utils._download`` (formerly ``get_model_path``). + ``mlx_lm.utils._download`` (formerly ``get_model_path``); mlx-lm is an + unpinned dependency, so accept either name and fail with the reason + rather than a bare ImportError if it is renamed again. """ - from mlx_lm.utils import _download + from mlx_lm import utils as mlx_lm_utils - return Path(_download(model_path, revision=revision)) + resolver = getattr(mlx_lm_utils, "_download", None) or getattr( + mlx_lm_utils, "get_model_path", None + ) + if resolver is None: + raise RemoteCodeGateError( + "this mlx-lm exposes neither mlx_lm.utils._download nor " + "mlx_lm.utils.get_model_path, so the checkpoint directory cannot " + "be resolved for inspection before mlx-lm loads it" + ) + resolved = resolver(model_path, revision=revision) + # get_model_path returned (path, config) in some releases. + if isinstance(resolved, tuple): + resolved = resolved[0] + return Path(resolved) def ensure_remote_code_allowed(model_dir: Path, trust_remote_code: bool) -> None: diff --git a/python/sglang/srt/mem_cache/kv_cache_dtype.py b/python/sglang/srt/mem_cache/kv_cache_dtype.py index fc71a081f6bc..f3e174c5ee33 100644 --- a/python/sglang/srt/mem_cache/kv_cache_dtype.py +++ b/python/sglang/srt/mem_cache/kv_cache_dtype.py @@ -95,5 +95,9 @@ def configure_kv_cache_dtype( model_dtype, ) kv_cache_dtype = model_dtype + # The pool is no longer quantized, so the returned tag must stop saying it + # is: "auto" is what an unquantized pool reports, and attention backends + # gate their descale paths on it (model_runner.kv_cache_dtype_str). + resolved_kv_cache_dtype = "auto" return resolved_kv_cache_dtype, kv_cache_dtype From b30ca236342a750875e2a7aae67c84cb21513d5b Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Mon, 10 Aug 2026 23:31:20 -0700 Subject: [PATCH 12/18] fix cpu suite: kv dtype stub field; fp4 resolvable whitelist; responses skip_special_tokens assertion --- .../unit/entrypoints/openai/test_serving_responses.py | 6 +++++- .../layers/quantization/test_fp4_kv_cache_quant_method.py | 2 +- test/registered/unit/test_model_overrides.py | 1 + 3 files changed, 7 insertions(+), 2 deletions(-) diff --git a/test/registered/unit/entrypoints/openai/test_serving_responses.py b/test/registered/unit/entrypoints/openai/test_serving_responses.py index 9b68fad2cc3a..b87f1f643903 100644 --- a/test/registered/unit/entrypoints/openai/test_serving_responses.py +++ b/test/registered/unit/entrypoints/openai/test_serving_responses.py @@ -380,7 +380,11 @@ def test_marker_preserving_parser_disables_skip_special_tokens(self): def test_default_parser_keeps_skip_special_tokens(self): serving = make_serving() params = self._create_responses_sampling_params(serving) - self.assertTrue(params["skip_special_tokens"]) + # ResponsesRequest has no skip_special_tokens of its own, so the chat + # request's True is a synthesized default, not user intent. Leave the + # key unset rather than overriding --preferred-sampling-params with it; + # only the marker-preserving parsers force it off. + self.assertNotIn("skip_special_tokens", params) class InputItemNormalizationTestCase(CustomTestCase): diff --git a/test/registered/unit/layers/quantization/test_fp4_kv_cache_quant_method.py b/test/registered/unit/layers/quantization/test_fp4_kv_cache_quant_method.py index 0ebf64efa170..87f530d3d0a5 100644 --- a/test/registered/unit/layers/quantization/test_fp4_kv_cache_quant_method.py +++ b/test/registered/unit/layers/quantization/test_fp4_kv_cache_quant_method.py @@ -75,7 +75,7 @@ def test_model_runner_rejects_legacy_fp4_alias(self): from sglang.srt.runtime_context import get_context runner = object.__new__(ModelRunner) - runner.server_args = SimpleNamespace() + runner.server_args = SimpleNamespace(speculative_draft_kv_cache_dtype=None) runner.draft_attention_backend = None # The runner reads the requested dtype off the model bag, so the double # publishes it rather than carrying it on a stand-in config. diff --git a/test/registered/unit/test_model_overrides.py b/test/registered/unit/test_model_overrides.py index c1411e7a8689..73c502b62047 100644 --- a/test/registered/unit/test_model_overrides.py +++ b/test/registered/unit/test_model_overrides.py @@ -89,6 +89,7 @@ def test_server_args_whitelist_is_exactly_the_migrated_fields(self): "decode_attention_backend", "flashinfer_allreduce_fusion_backend", "fp8_gemm_runner_backend", + "fp4_gemm_runner_backend", "disable_custom_all_reduce", "enable_aiter_allreduce_fusion", "enable_symm_mem", From 47bdced145c30c6e0d89ba8de47d1ef2d5f5f254 Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Mon, 10 Aug 2026 23:36:52 -0700 Subject: [PATCH 13/18] trim comments; revert stale finish() docstring note --- python/sglang/srt/configs/muse_glimmer.py | 3 +-- .../sglang/srt/entrypoints/openai/serving_responses.py | 4 ++-- python/sglang/srt/function_call/base_format_detector.py | 4 +--- .../sglang/srt/hardware_backend/mlx/remote_code_gate.py | 5 ++--- python/sglang/srt/mem_cache/kv_cache_dtype.py | 5 ++--- python/sglang/srt/models/dflash.py | 9 +++------ .../unit/entrypoints/openai/test_serving_responses.py | 5 ++--- 7 files changed, 13 insertions(+), 22 deletions(-) diff --git a/python/sglang/srt/configs/muse_glimmer.py b/python/sglang/srt/configs/muse_glimmer.py index f07de5ea6267..40866764cfc1 100644 --- a/python/sglang/srt/configs/muse_glimmer.py +++ b/python/sglang/srt/configs/muse_glimmer.py @@ -28,8 +28,7 @@ class MuseGlimmerAssistantConfig(PretrainedConfig): model_type = "muse_glimmer_assistant" - # The DFlash draft has no head of its own; the target's vocab is borrowed - # when the draft worker is built. + # The DFlash draft has no head; draft_worker_common borrows the target's. vocab_size = None diff --git a/python/sglang/srt/entrypoints/openai/serving_responses.py b/python/sglang/srt/entrypoints/openai/serving_responses.py index aede733587eb..d25653989fe3 100644 --- a/python/sglang/srt/entrypoints/openai/serving_responses.py +++ b/python/sglang/srt/entrypoints/openai/serving_responses.py @@ -2215,8 +2215,8 @@ def _close_tool_call_state(tool_index: int): stream_offset = len(text) if not delta and finish_reason is None: continue - # finish_reason is sticky, so gate on `flushed` to drain the - # parsers exactly once no matter how many terminal chunks arrive. + # finish_reason is sticky, so flush exactly once however many + # terminal chunks arrive. flush = ( not flushed and finish_reason is not None diff --git a/python/sglang/srt/function_call/base_format_detector.py b/python/sglang/srt/function_call/base_format_detector.py index 48b6b9dea04f..dd848402efb6 100644 --- a/python/sglang/srt/function_call/base_format_detector.py +++ b/python/sglang/srt/function_call/base_format_detector.py @@ -351,12 +351,10 @@ def has_tool_call(self, text: str) -> bool: raise NotImplementedError() def finish(self, tools: List[Tool]) -> StreamingParseResult: - """Called when the stream ends; flush any buffered state. + """Called once when the stream ends; flush any buffered state. Detectors that hold text back while waiting for a marker that can no longer arrive (the stream is over) override this to release it. - Overrides should be idempotent: there is nothing left to release the - second time, and callers see several chunks carrying a finish reason. """ return StreamingParseResult() diff --git a/python/sglang/srt/hardware_backend/mlx/remote_code_gate.py b/python/sglang/srt/hardware_backend/mlx/remote_code_gate.py index 49e42cf57035..fca8f32bdcc5 100644 --- a/python/sglang/srt/hardware_backend/mlx/remote_code_gate.py +++ b/python/sglang/srt/hardware_backend/mlx/remote_code_gate.py @@ -45,9 +45,8 @@ def resolve_model_directory(model_path: str, revision: Optional[str] = None) -> Uses mlx-lm's resolver so the directory is byte-identical to what a direct ``mlx_lm.load`` call would consume; existing local paths are returned as-is (no network access). mlx-lm 0.31.x exposes this as - ``mlx_lm.utils._download`` (formerly ``get_model_path``); mlx-lm is an - unpinned dependency, so accept either name and fail with the reason - rather than a bare ImportError if it is renamed again. + ``mlx_lm.utils._download`` (formerly ``get_model_path``); mlx-lm is + unpinned, so accept either name. """ from mlx_lm import utils as mlx_lm_utils diff --git a/python/sglang/srt/mem_cache/kv_cache_dtype.py b/python/sglang/srt/mem_cache/kv_cache_dtype.py index f3e174c5ee33..e4ba5587c949 100644 --- a/python/sglang/srt/mem_cache/kv_cache_dtype.py +++ b/python/sglang/srt/mem_cache/kv_cache_dtype.py @@ -95,9 +95,8 @@ def configure_kv_cache_dtype( model_dtype, ) kv_cache_dtype = model_dtype - # The pool is no longer quantized, so the returned tag must stop saying it - # is: "auto" is what an unquantized pool reports, and attention backends - # gate their descale paths on it (model_runner.kv_cache_dtype_str). + # Unquantized pool now, and "auto" is the tag for that; attention + # backends gate their descale paths on it (kv_cache_dtype_str). resolved_kv_cache_dtype = "auto" return resolved_kv_cache_dtype, kv_cache_dtype diff --git a/python/sglang/srt/models/dflash.py b/python/sglang/srt/models/dflash.py index 85198c4defcb..a7fab2328762 100644 --- a/python/sglang/srt/models/dflash.py +++ b/python/sglang/srt/models/dflash.py @@ -44,11 +44,8 @@ def _get_dflash_attention_type(config) -> AttentionType: - """Causality of a DFLASH draft, from the checkpoint. - - DFLASH drafts attend over the whole draft block, so they are bidirectional - unless the checkpoint declares itself causal. - """ + """DFLASH drafts attend over the whole draft block, so they are + bidirectional unless the checkpoint declares itself causal.""" text_config = getattr(config, "text_config", None) or config return ( AttentionType.DECODER @@ -74,7 +71,7 @@ def _get_dflash_layer_attention_params( return -1, _get_dflash_attention_type(config) if layer_type == "sliding_attention": # Windowing is orthogonal to causality: the layer masks only - # p1 - p0 >= sliding_window and keeps the checkpoint's causality. + # p1 - p0 >= sliding_window. sliding_window_size = get_dflash_attention_sliding_window_size(config) assert sliding_window_size is not None return sliding_window_size, _get_dflash_attention_type(config) diff --git a/test/registered/unit/entrypoints/openai/test_serving_responses.py b/test/registered/unit/entrypoints/openai/test_serving_responses.py index b87f1f643903..08c8fd96949c 100644 --- a/test/registered/unit/entrypoints/openai/test_serving_responses.py +++ b/test/registered/unit/entrypoints/openai/test_serving_responses.py @@ -381,9 +381,8 @@ def test_default_parser_keeps_skip_special_tokens(self): serving = make_serving() params = self._create_responses_sampling_params(serving) # ResponsesRequest has no skip_special_tokens of its own, so the chat - # request's True is a synthesized default, not user intent. Leave the - # key unset rather than overriding --preferred-sampling-params with it; - # only the marker-preserving parsers force it off. + # request's True is a synthesized default, not user intent -- leave the + # key unset instead of overriding --preferred-sampling-params with it. self.assertNotIn("skip_special_tokens", params) From ca772dd335c159e660125ac84b1c7bdcb33d8048 Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Mon, 10 Aug 2026 23:38:38 -0700 Subject: [PATCH 14/18] compress comments further --- python/sglang/srt/entrypoints/openai/serving_responses.py | 3 +-- python/sglang/srt/mem_cache/kv_cache_dtype.py | 3 +-- python/sglang/srt/models/dflash.py | 6 ++---- .../unit/entrypoints/openai/test_serving_responses.py | 5 ++--- 4 files changed, 6 insertions(+), 11 deletions(-) diff --git a/python/sglang/srt/entrypoints/openai/serving_responses.py b/python/sglang/srt/entrypoints/openai/serving_responses.py index d25653989fe3..81f9c3e7709e 100644 --- a/python/sglang/srt/entrypoints/openai/serving_responses.py +++ b/python/sglang/srt/entrypoints/openai/serving_responses.py @@ -2215,8 +2215,7 @@ def _close_tool_call_state(tool_index: int): stream_offset = len(text) if not delta and finish_reason is None: continue - # finish_reason is sticky, so flush exactly once however many - # terminal chunks arrive. + # finish_reason is sticky, so it would otherwise re-flush. flush = ( not flushed and finish_reason is not None diff --git a/python/sglang/srt/mem_cache/kv_cache_dtype.py b/python/sglang/srt/mem_cache/kv_cache_dtype.py index e4ba5587c949..457fe2f9501f 100644 --- a/python/sglang/srt/mem_cache/kv_cache_dtype.py +++ b/python/sglang/srt/mem_cache/kv_cache_dtype.py @@ -95,8 +95,7 @@ def configure_kv_cache_dtype( model_dtype, ) kv_cache_dtype = model_dtype - # Unquantized pool now, and "auto" is the tag for that; attention - # backends gate their descale paths on it (kv_cache_dtype_str). + # "auto" is the tag for an unquantized pool; backends gate descale on it. resolved_kv_cache_dtype = "auto" return resolved_kv_cache_dtype, kv_cache_dtype diff --git a/python/sglang/srt/models/dflash.py b/python/sglang/srt/models/dflash.py index a7fab2328762..4adc9990d438 100644 --- a/python/sglang/srt/models/dflash.py +++ b/python/sglang/srt/models/dflash.py @@ -44,8 +44,7 @@ def _get_dflash_attention_type(config) -> AttentionType: - """DFLASH drafts attend over the whole draft block, so they are - bidirectional unless the checkpoint declares itself causal.""" + """Bidirectional over the draft block unless the checkpoint says causal.""" text_config = getattr(config, "text_config", None) or config return ( AttentionType.DECODER @@ -70,8 +69,7 @@ def _get_dflash_layer_attention_params( if layer_type == "full_attention": return -1, _get_dflash_attention_type(config) if layer_type == "sliding_attention": - # Windowing is orthogonal to causality: the layer masks only - # p1 - p0 >= sliding_window. + # Windowing is orthogonal to causality (mask is p1 - p0 >= window). sliding_window_size = get_dflash_attention_sliding_window_size(config) assert sliding_window_size is not None return sliding_window_size, _get_dflash_attention_type(config) diff --git a/test/registered/unit/entrypoints/openai/test_serving_responses.py b/test/registered/unit/entrypoints/openai/test_serving_responses.py index 08c8fd96949c..ab6c561ab5e3 100644 --- a/test/registered/unit/entrypoints/openai/test_serving_responses.py +++ b/test/registered/unit/entrypoints/openai/test_serving_responses.py @@ -380,9 +380,8 @@ def test_marker_preserving_parser_disables_skip_special_tokens(self): def test_default_parser_keeps_skip_special_tokens(self): serving = make_serving() params = self._create_responses_sampling_params(serving) - # ResponsesRequest has no skip_special_tokens of its own, so the chat - # request's True is a synthesized default, not user intent -- leave the - # key unset instead of overriding --preferred-sampling-params with it. + # The chat request's True is a synthesized default (ResponsesRequest has + # no such field), so leave it unset for --preferred-sampling-params. self.assertNotIn("skip_special_tokens", params) From 605ba0e344b4975f5d553ec7f2ba56d41aae0cbc Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Mon, 10 Aug 2026 23:49:54 -0700 Subject: [PATCH 15/18] extract shared muse glimmer wire format --- .../function_call/muse_glimmer_detector.py | 44 +++++++-------- .../srt/function_call/muse_glimmer_format.py | 39 +++++++++++++ python/sglang/srt/parser/reasoning_parser.py | 56 +++++++------------ 3 files changed, 79 insertions(+), 60 deletions(-) create mode 100644 python/sglang/srt/function_call/muse_glimmer_format.py diff --git a/python/sglang/srt/function_call/muse_glimmer_detector.py b/python/sglang/srt/function_call/muse_glimmer_detector.py index aab5d52640a8..83bcdce5f938 100644 --- a/python/sglang/srt/function_call/muse_glimmer_detector.py +++ b/python/sglang/srt/function_call/muse_glimmer_detector.py @@ -12,21 +12,23 @@ ToolCallItem, _GetInfoFunc, ) +from sglang.srt.function_call.muse_glimmer_format import ( + EOM, + EOT, + FUNCTION_CALLS_CLOSE, + FUNCTION_CALLS_OPEN, + INVOKE_CLOSE, + INVOKE_OPEN, + MAX_MARKER, + MESSAGE, + RECIPIENT_RE, + START, + has_atem_markers, + partial_marker_len, +) logger = logging.getLogger(__name__) -# Channel framing, shared with the reasoning-side MuseGlimmerDetector. -MESSAGE = "<|message|>" -EOM = "<|eom|>" -EOT = "<|eot|>" -START = "<|start|>" - -# ATEM payload markers. -FUNCTION_CALLS_OPEN = "" -FUNCTION_CALLS_CLOSE = "" -INVOKE_CLOSE = "" - -_RECIPIENT_RE = re.compile(r"to=([^\s<]+)") _INVOKE_OPEN_RE = re.compile(r']*?\bname="(?P[^"]+)"[^>]*?>') _PARAM_RE = re.compile( r']*?\bname="(?P[^"]+)"[^>]*?>(?P.*?)' @@ -43,10 +45,6 @@ def _is_tool_channel(recipient: Optional[str]) -> bool: return recipient is not None and recipient not in _NON_TOOL_RECIPIENTS -# Longest marker that could straddle a chunk boundary while streaming. -_MAX_MARKER = max(len(m) for m in (MESSAGE, EOM, EOT, START, FUNCTION_CALLS_OPEN)) - - def _decode_value(raw: str): try: return json.loads(raw) @@ -98,7 +96,7 @@ def __init__(self): self._open_invoke: Optional[str] = None def has_tool_call(self, text: str) -> bool: - return FUNCTION_CALLS_OPEN in text or " Set[str]: return {t.function.name for t in tools or [] if t.function and t.function.name} @@ -150,12 +148,8 @@ def finish(self, tools: List[Tool]) -> StreamingParseResult: return StreamingParseResult(normal_text="".join(normal_parts), calls=calls) def _held_back(self, text: str) -> int: - """Length of the trailing suffix that could still grow into a marker.""" - markers = (MESSAGE, EOM, EOT, START, FUNCTION_CALLS_OPEN, " bool: def structure_info(self) -> _GetInfoFunc: return lambda name: StructureInfo( - begin=f'{FUNCTION_CALLS_OPEN}\n', + begin=f'{FUNCTION_CALLS_OPEN}\n{INVOKE_OPEN} name="{name}">', end=f"{INVOKE_CLOSE}\n{FUNCTION_CALLS_CLOSE}", trigger=FUNCTION_CALLS_OPEN, ) diff --git a/python/sglang/srt/function_call/muse_glimmer_format.py b/python/sglang/srt/function_call/muse_glimmer_format.py new file mode 100644 index 000000000000..fbd91097721e --- /dev/null +++ b/python/sglang/srt/function_call/muse_glimmer_format.py @@ -0,0 +1,39 @@ +"""Muse Glimmer wire format, shared by its reasoning and function-call detectors.""" + +import re +from typing import Sequence + +# Channel framing. +MESSAGE = "<|message|>" +EOM = "<|eom|>" +EOT = "<|eot|>" +START = "<|start|>" + +# ATEM payload markers. +FUNCTION_CALLS_OPEN = "" +FUNCTION_CALLS_CLOSE = "" +INVOKE_OPEN = " bool: + return INVOKE_OPEN in text or FUNCTION_CALLS_OPEN in text + + +def partial_marker_len(text: str, markers: Sequence[str], max_len: int) -> int: + """Length of the longest suffix of ``text`` that could still become a marker. + + Returns 0 when nothing is held back, so ordinary text streams out immediately + instead of waiting for a terminator that may never arrive. + """ + for k in range(min(len(text), max_len - 1), 0, -1): + tail = text[-k:] + if any(m.startswith(tail) for m in markers): + return k + return 0 diff --git a/python/sglang/srt/parser/reasoning_parser.py b/python/sglang/srt/parser/reasoning_parser.py index 215208a727c4..0ff158fc655e 100644 --- a/python/sglang/srt/parser/reasoning_parser.py +++ b/python/sglang/srt/parser/reasoning_parser.py @@ -23,6 +23,16 @@ strip_partial_marker_suffix, strip_response_wrappers, ) +from sglang.srt.function_call.muse_glimmer_format import ( + EOM, + EOT, + MAX_CHANNEL_MARKER, + MESSAGE, + RECIPIENT_RE, + START, + has_atem_markers, + partial_marker_len, +) from sglang.srt.parser.harmony_parser import HarmonyParser from sglang.srt.parser.inkling_tokenizer import ( CONTENT_INVOKE_TOOL_JSON, @@ -1652,16 +1662,6 @@ class MuseGlimmerDetector(BaseReasoningFormatDetector): terminator and the partial body is still attributed to whichever channel was open. """ - MESSAGE = "<|message|>" - EOM = "<|eom|>" - EOT = "<|eot|>" - START = "<|start|>" - _MAX_MARKER = max(len(MESSAGE), len(EOM), len(EOT), len(START)) - _RECIPIENT_RE = re.compile(r"to=([^\s<]+)") - - # ATEM markers that identify a tool-call turn in the non-reasoning remainder. - _ATEM_MARKERS = ("") - def __init__( self, stream_reasoning: bool = True, @@ -1672,8 +1672,8 @@ def __init__( tool_call_parser_active: bool = False, ): super().__init__( - " to=self" + self.MESSAGE, - self.EOM, + " to=self" + MESSAGE, + EOM, force_reasoning=force_reasoning, stream_reasoning=stream_reasoning, continue_final_message=continue_final_message, @@ -1689,20 +1689,6 @@ def __init__( def _sink(self, recipient: Optional[str]) -> str: return "reasoning" if recipient == "self" else "normal" - @classmethod - def _partial_marker_len(cls, buf: str) -> int: - """Length of the longest suffix of ``buf`` that could still become a marker. - - Returns 0 when nothing is held back, so ordinary text streams out immediately - instead of waiting for a terminator that may never arrive. - """ - markers = (cls.EOM, cls.EOT, cls.START) - for k in range(min(len(buf), cls._MAX_MARKER - 1), 0, -1): - tail = buf[-k:] - if any(m.startswith(tail) for m in markers): - return k - return 0 - def _consume(self, flush: bool, preserve_channels: bool = False) -> Tuple[str, str]: """Drain self._buffer into (reasoning, normal). @@ -1719,16 +1705,16 @@ def _consume(self, flush: bool, preserve_channels: bool = False) -> Tuple[str, s while self._buffer: if not self._in_body: - idx = self._buffer.find(self.MESSAGE) + idx = self._buffer.find(MESSAGE) if idx == -1: if flush: normal_parts.append(self._buffer) self._buffer = "" break header = self._buffer[:idx] - m = self._RECIPIENT_RE.search(header) + m = RECIPIENT_RE.search(header) self._recipient = m.group(1) if m else "user" - self._buffer = self._buffer[idx + len(self.MESSAGE) :] + self._buffer = self._buffer[idx + len(MESSAGE) :] self._in_body = True if self._sink(self._recipient) == "reasoning": if self._saw_reasoning_block: @@ -1736,11 +1722,11 @@ def _consume(self, flush: bool, preserve_channels: bool = False) -> Tuple[str, s self._saw_reasoning_block = True elif self._recipient != "user" or preserve_channels: # Keep the header so the function-call detector sees it. - normal_parts.append(header + self.MESSAGE) + normal_parts.append(header + MESSAGE) continue end_idx, end_tok = -1, "" - for tok in (self.EOM, self.EOT): + for tok in (EOM, EOT): i = self._buffer.find(tok) if i != -1 and (end_idx == -1 or i < end_idx): end_idx, end_tok = i, tok @@ -1762,7 +1748,9 @@ def _consume(self, flush: bool, preserve_channels: bool = False) -> Tuple[str, s if flush: body, self._buffer = self._buffer, "" else: - keep = self._partial_marker_len(self._buffer) + keep = partial_marker_len( + self._buffer, (EOM, EOT, START), MAX_CHANNEL_MARKER + ) if keep == len(self._buffer): break body = self._buffer[: len(self._buffer) - keep] @@ -1780,9 +1768,7 @@ def detect_and_parse(self, text: str) -> StreamingParseResult: self._buffer += text raw = self._buffer reasoning, normal = self._consume(flush=True) - if self._tool_call_parser_active and any( - m in normal for m in self._ATEM_MARKERS - ): + if self._tool_call_parser_active and has_atem_markers(normal): self._buffer = raw self._recipient = None self._in_body = False From c33475208bf2d176d87e75320ff91a5218c27952 Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Tue, 11 Aug 2026 00:18:56 -0700 Subject: [PATCH 16/18] drop mlx reference-correctness test; it never runs in ci --- .../test_muse_glimmer_mlx_correctness.py | 490 ------------------ .../mlx/test_muse_glimmer_mlx_model.py | 5 +- 2 files changed, 2 insertions(+), 493 deletions(-) delete mode 100644 test/registered/mlx/models_e2e/test_muse_glimmer_mlx_correctness.py diff --git a/test/registered/mlx/models_e2e/test_muse_glimmer_mlx_correctness.py b/test/registered/mlx/models_e2e/test_muse_glimmer_mlx_correctness.py deleted file mode 100644 index 47b6cf3448b8..000000000000 --- a/test/registered/mlx/models_e2e/test_muse_glimmer_mlx_correctness.py +++ /dev/null @@ -1,490 +0,0 @@ -"""Correctness tests for Muse Glimmer served on the SGLang MLX backend. - -Muse Glimmer interleaves sliding-window (window=2048) and NoPE full-attention -layers, gates attention output through a sigmoid projection fused into -q_proj, and ships as an mlx-lm ``model_file`` artifact — so it exercises -the MLX backend's remote-code gate, container-level window discovery, and -banded-mask paths end to end. Two guards: - -1. ``TestMuseGlimmerMlxServing`` — black-box serving smoke against a running - server, including a prompt long enough to engage the sliding window. -2. ``TestMuseGlimmerMlxReferenceCorrectness`` — token-for-token equivalence of - SGLang greedy decoding against raw, unpatched mlx_lm greedy generation - on identical ``input_ids``. Both sides keep full KV history (the Muse Glimmer - model file's ``make_cache`` deliberately avoids ``RotatingKVCache``), - so exact equality holds even past the window and across prefill - chunkings. - -The artifact is weight-derived and private: its path comes exclusively from the -``SGLANG_MLX_TEST_MUSE_GLIMMER_ARTIFACT`` environment variable and the whole module SKIPS -cleanly when it is unset or missing — CI has no dependency on private -assets. Window-engagement prompt lengths are derived from the artifact's -own config, so the suite also runs against tiny packaged fixtures. -""" - -from __future__ import annotations - -import concurrent.futures -import importlib.util -import json -import os -import unittest -from pathlib import Path - -import requests - -from sglang.srt.utils import kill_process_tree -from sglang.test.ci.ci_register import register_cpu_ci, register_mlx_ci -from sglang.test.test_utils import ( - DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, - DEFAULT_URL_FOR_TEST, - CustomTestCase, - popen_launch_server, -) - -# Registered on the CPU suite as an import check; skipped wherever mlx or -# the private artifact is absent. stage-b-e2e-mlx is manual-dispatch only. -register_cpu_ci(est_time=1, suite="base-a-test-cpu") -register_mlx_ci(est_time=420, suite="stage-b-e2e-mlx") - -_HAS_MLX = ( - importlib.util.find_spec("mlx") is not None - and importlib.util.find_spec("mlx_lm") is not None -) - -ARTIFACT = os.environ.get("SGLANG_MLX_TEST_MUSE_GLIMMER_ARTIFACT") -MEM_FRACTION_STATIC = os.environ.get("SGLANG_MLX_TEST_MEM_FRACTION", "0.85") -MIN_FREE_GB = float(os.environ.get("SGLANG_MLX_TEST_MIN_FREE_GB", "20")) -MAX_NEW_TOKENS = 32 - - -def _artifact_or_skip() -> Path: - if not _HAS_MLX: - raise unittest.SkipTest("requires mlx + mlx_lm (Apple Silicon only)") - if not ARTIFACT: - raise unittest.SkipTest( - "SGLANG_MLX_TEST_MUSE_GLIMMER_ARTIFACT is not set; the Muse Glimmer artifact is private " - "and must be supplied via the environment" - ) - path = Path(ARTIFACT).expanduser() - if not (path / "config.json").is_file(): - raise unittest.SkipTest(f"no packaged artifact at {path}") - config = json.loads((path / "config.json").read_text()) - if config.get("muse_glimmer_mlx_format") != 1: - raise unittest.SkipTest( - f"{path} is not a packaged Muse Glimmer MLX artifact (missing marker)" - ) - return path - - -def _artifact_config(path: Path) -> dict: - return json.loads((path / "config.json").read_text()) - - -def _available_gb(): - try: - import psutil - - return psutil.virtual_memory().available / 1024**3 - except Exception: - return None - - -def _check_memory(path: Path): - # Tiny fixtures need no headroom; only guard for real-size artifacts. - weights_gb = ( - sum(p.stat().st_size for p in path.glob("model*.safetensors")) / 1024**3 - ) - if weights_gb <= 1: - return - # A previous test class's just-killed server can hold memory for a few - # seconds; wait for the release instead of skipping on a transient dip. - import time - - deadline = time.monotonic() + 90 - avail = _available_gb() - while avail is not None and avail < MIN_FREE_GB and time.monotonic() < deadline: - time.sleep(5) - avail = _available_gb() - if avail is not None and avail < MIN_FREE_GB: - raise unittest.SkipTest( - f"insufficient free memory: {avail:.1f} GB < {MIN_FREE_GB} GB " - f"needed to safely serve a {weights_gb:.0f} GB artifact" - ) - - -def _case_lengths(config: dict) -> list[int]: - window = int(config.get("sliding_window", 2048)) - max_pos = int(config.get("max_position_embeddings", 16384)) - return sorted( - { - max(4, window - 1), - min(window + 64, max_pos - MAX_NEW_TOKENS - 1), - min(2 * window, max_pos - MAX_NEW_TOKENS - 1), - } - ) - - -def _prompt_id_cases(config: dict) -> list[list[int]]: - """Deterministic random input_ids around the artifact's window size. - - Serving-smoke material only: random tokens produce flat next-token - distributions, fine for shape/length assertions but useless for exact - greedy matching. - """ - import numpy as np - - vocab = int(config["vocab_size"]) - hi = min(vocab, 200000) - lo = min(1000, hi // 2) - rng = np.random.default_rng(20260731) - return [rng.integers(lo, hi, size=n).tolist() for n in _case_lengths(config)] - - -def _natural_prompt_cases(config: dict, artifact: Path) -> list[list[int]]: - """Natural-text prompts at the same window-straddling lengths. - - Real text gives the model real top-1 margins, so exact greedy - equivalence across different batching shapes is meaningful — a benign - kernel-shape difference cannot flip a decisive argmax, while a masking - or batching bug still can. - """ - from mlx_lm.utils import load_tokenizer - - tok = load_tokenizer(artifact) - filler = ( - "Day %d: we walked along the ridge, catalogued mosses and lichens, " - "measured stream flow, and noted the weather turning. " - ) - cases = [] - for n in _case_lengths(config): - # Prompts end MID-NARRATIVE: continuing a strictly cyclic journal is - # near-deterministic, so greedy continuations sit on wide top-1 - # margins (an open-ended question would put the decision point on a - # near-tie, where benign kernel-shape noise can flip the argmax). - text, day = "The following is a field journal. ", 1 - ids: list[int] = [] - while len(ids) < n: - text += filler % day - day += 1 - ids = tok.encode(text) - cases.append(ids[:n]) - return cases - - -def _reference_greedy(model, prompt_ids, max_tokens, prefill_step_size=None): - import mlx.core as mx - from mlx_lm.generate import generate_step - - kwargs = {} - if prefill_step_size is not None: - # Match the server's chunked-prefill size: Metal matmuls accumulate - # at reduced precision and are shape-dependent, so exact greedy - # equality requires both sides to prefill in identical chunk shapes. - kwargs["prefill_step_size"] = prefill_step_size - out = [] - for token, _ in generate_step( - mx.array(prompt_ids), model, max_tokens=max_tokens, **kwargs - ): - out.append(int(token)) - if len(out) >= max_tokens: - break - return out - - -class TestMuseGlimmerMlxServing(CustomTestCase): - @classmethod - def setUpClass(cls): - cls.artifact = _artifact_or_skip() - _check_memory(cls.artifact) - cls.config = _artifact_config(cls.artifact) - cls.base_url = DEFAULT_URL_FOR_TEST - - env = os.environ.copy() - env["SGLANG_USE_MLX"] = "1" - cls.process = popen_launch_server( - str(cls.artifact), - cls.base_url, - timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, - other_args=[ - "--trust-remote-code", - "--disable-radix-cache", - "--disable-cuda-graph", - "--mem-fraction-static", - MEM_FRACTION_STATIC, - ], - env=env, - ) - - @classmethod - def tearDownClass(cls): - if hasattr(cls, "process") and cls.process is not None: - kill_process_tree(cls.process.pid) - - def _generate(self, payload): - resp = requests.post(f"{self.base_url}/generate", json=payload, timeout=600) - resp.raise_for_status() - return resp.json() - - def test_solo_generation(self): - ids = _prompt_id_cases(self.config)[0] - out = self._generate( - { - "input_ids": ids, - "sampling_params": { - "max_new_tokens": 16, - "temperature": 0, - "ignore_eos": True, - }, - } - ) - self.assertEqual(len(out["output_ids"]), 16) - - def test_window_engaging_prompt(self): - # Longest case exceeds the sliding window: prefill and decode both - # run with banded masks engaged on the sliding layers. - ids = _prompt_id_cases(self.config)[-1] - self.assertGreater(len(ids), int(self.config["sliding_window"])) - out = self._generate( - { - "input_ids": ids, - "sampling_params": { - "max_new_tokens": 16, - "temperature": 0, - "ignore_eos": True, - }, - } - ) - self.assertEqual(len(out["output_ids"]), 16) - - def test_batch_generation(self): - cases = _prompt_id_cases(self.config) - out = self._generate( - { - "input_ids": cases, - "sampling_params": { - "max_new_tokens": 12, - "temperature": 0, - "ignore_eos": True, - }, - } - ) - self.assertEqual(len(out), len(cases)) - for r in out: - self.assertEqual(len(r["output_ids"]), 12) - - def test_over_context_request_rejected(self): - # A prompt beyond the context window must produce a clear error, - # not unbounded cache growth or a process kill. - max_pos = int(self.config.get("max_position_embeddings", 16384)) - too_long = [1000 + (i % 1000) for i in range(max_pos + 64)] - resp = requests.post( - f"{self.base_url}/generate", - json={ - "input_ids": too_long, - "sampling_params": {"max_new_tokens": 8, "temperature": 0}, - }, - timeout=600, - ) - body = ( - resp.json() - if resp.headers.get("content-type", "").startswith("application/json") - else {} - ) - rejected = resp.status_code >= 400 or ( - isinstance(body, dict) - and body.get("meta_info", {}).get("finish_reason", {}).get("type") - == "abort" - ) - self.assertTrue( - rejected, - f"over-context request was not rejected: {resp.status_code} {body}", - ) - # The server must still be healthy afterwards. - ok = requests.get(f"{self.base_url}/health", timeout=60) - self.assertEqual(ok.status_code, 200) - - def test_abort_then_rerun_is_isolated(self): - # Closing a stream mid-decode aborts the request server-side; a - # fresh identical request afterwards must produce the same output - # as one issued before the abort (no cache/state leakage). - ids = _prompt_id_cases(self.config)[0] - params = {"max_new_tokens": 24, "temperature": 0, "ignore_eos": True} - - def full_run(): - r = requests.post( - f"{self.base_url}/generate", - json={"input_ids": ids, "sampling_params": params}, - timeout=600, - ) - r.raise_for_status() - return r.json()["output_ids"] - - before = full_run() - with requests.post( - f"{self.base_url}/generate", - json={ - "input_ids": ids, - "sampling_params": params, - "stream": True, - }, - stream=True, - timeout=600, - ) as resp: - for line in resp.iter_lines(): - if line and line.startswith(b"data:") and b"[DONE]" not in line: - break # first token seen -> drop the connection mid-decode - after = full_run() - self.assertEqual(before, after, "post-abort rerun diverged") - - def test_explicit_eom_stop_honored(self): - # 200007 must not be a default stop, but the stop machinery must - # still be ABLE to stop on it when a request asks — proving its - # absence from the default set is configuration, not inability. - # The prompt pre-opens the reasoning channel, so the only way for - # greedy decoding to close it is <|eom|> — whether the model would - # have chosen to reason on its own varies with server config. - if self.config["vocab_size"] < 200008: - self.skipTest("tiny fixture vocab has no <|eom|> token") - out = self._generate( - { - "text": ( - "<|start|>user<|message|>Hi, who are you?<|eot|>" - "<|start|>assistant to=self<|message|>" - ), - "sampling_params": { - "max_new_tokens": 512, - "temperature": 0, - "stop_token_ids": [200007], - }, - } - ) - finish = out["meta_info"]["finish_reason"] - self.assertEqual(finish["type"], "stop") - self.assertEqual(finish["matched"], 200007) - - def test_eom_not_terminal_for_real_artifact(self): - # <|eom|> (200007) closes a message, not a turn; the packaged - # generation_config must not list it as EOS, or every response - # truncates at end-of-thinking. Only meaningful on the real vocab. - if self.config["vocab_size"] < 200008: - self.skipTest("tiny fixture vocab has no <|eom|> token") - gen = json.loads((self.artifact / "generation_config.json").read_text()) - self.assertNotIn(200007, gen["eos_token_id"]) - self.assertEqual(sorted(gen["eos_token_id"]), [200001, 200008]) - - -class TestMuseGlimmerMlxReferenceCorrectness(CustomTestCase): - """SGLang vs raw mlx_lm greedy equivalence on identical input_ids. - - Real-weight artifacts only: random-weight fixtures produce softcapped - logits with near-zero top-1 margins, so ANY benign kernel-shape - difference (e.g. the wrapper's trailing-window KV truncation vs the - reference's full-history banded mask, both mathematically equivalent) - flips argmax ties and exact matching carries no signal. - """ - - @classmethod - def setUpClass(cls): - cls.artifact = _artifact_or_skip() - _check_memory(cls.artifact) - cls.config = _artifact_config(cls.artifact) - if cls.config["vocab_size"] < 200008: - raise unittest.SkipTest( - "exact greedy equivalence needs real Muse Glimmer weights; " - "tiny random-weight fixtures have no top-1 margin" - ) - - from mlx_lm.utils import load_model - - cls.chunk = max(256, int(cls.config.get("sliding_window", 2048)) // 2) - ref_model, _ = load_model(cls.artifact) - cls.cases = [] - for prompt_ids in _natural_prompt_cases(cls.config, cls.artifact): - ref_ids = _reference_greedy( - ref_model, - prompt_ids, - MAX_NEW_TOKENS, - prefill_step_size=cls.chunk, - ) - cls.cases.append((prompt_ids, ref_ids)) - del ref_model - import gc - - gc.collect() - - cls.base_url = DEFAULT_URL_FOR_TEST - env = os.environ.copy() - env["SGLANG_USE_MLX"] = "1" - # Chunk size below the sliding window forces the longest prompt to - # prefill in several chunks that cross the window boundary — exact - # equality then also covers chunked-prefill mask stitching. The - # reference used the SAME chunk size (set in the class attribute - # above) so both sides prefill in identical kernel shapes. - chunk = cls.chunk - cls.process = popen_launch_server( - str(cls.artifact), - cls.base_url, - timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, - other_args=[ - "--trust-remote-code", - "--disable-radix-cache", - "--disable-cuda-graph", - "--mem-fraction-static", - MEM_FRACTION_STATIC, - "--chunked-prefill-size", - str(chunk), - ], - env=env, - ) - - @classmethod - def tearDownClass(cls): - if hasattr(cls, "process") and cls.process is not None: - kill_process_tree(cls.process.pid) - - def _sglang_greedy(self, prompt_ids, max_tokens): - resp = requests.post( - f"{self.base_url}/generate", - json={ - "input_ids": prompt_ids, - "sampling_params": { - "max_new_tokens": max_tokens, - "temperature": 0, - "ignore_eos": True, - }, - }, - timeout=600, - ) - resp.raise_for_status() - return resp.json()["output_ids"] - - def test_solo_greedy_matches_reference(self): - for prompt_ids, ref_ids in self.cases: - got = self._sglang_greedy(prompt_ids, MAX_NEW_TOKENS) - self.assertEqual( - got, - ref_ids, - f"greedy divergence at prompt length {len(prompt_ids)}", - ) - - def test_concurrent_ragged_batch_matches_reference(self): - # All cases in flight together: ragged lengths force mixed - # prefill/decode batching; every stream must still match its solo - # mlx_lm reference exactly (no cross-request contamination). - with concurrent.futures.ThreadPoolExecutor(len(self.cases)) as pool: - futures = [ - pool.submit(self._sglang_greedy, prompt_ids, MAX_NEW_TOKENS) - for prompt_ids, _ in self.cases - ] - results = [f.result() for f in futures] - for (prompt_ids, ref_ids), got in zip(self.cases, results): - self.assertEqual( - got, - ref_ids, - f"batched divergence at prompt length {len(prompt_ids)}", - ) - - -if __name__ == "__main__": - unittest.main(verbosity=2) diff --git a/test/registered/unit/hardware_backend/mlx/test_muse_glimmer_mlx_model.py b/test/registered/unit/hardware_backend/mlx/test_muse_glimmer_mlx_model.py index 6c69c40738cf..35dd77bb0558 100644 --- a/test/registered/unit/hardware_backend/mlx/test_muse_glimmer_mlx_model.py +++ b/test/registered/unit/hardware_backend/mlx/test_muse_glimmer_mlx_model.py @@ -1,8 +1,7 @@ """Unit tests for the Muse Glimmer MLX model file's load path — no weights, no server. -The e2e suite (``test/registered/mlx/models_e2e/test_muse_glimmer_mlx_correctness.py``) -needs the private packaged artifact, so the checkpoint-schema logic it relies -on is pinned here with tiny synthetic weights instead: +End-to-end coverage needs the private packaged artifact, so the checkpoint-schema +logic is pinned here with tiny synthetic weights instead: 1. ``flatten_rc_config`` — RC nested-config translation, including the two convention conversions (qk_scale_factor gains sqrt(head_dim); NoPE layers From 5e86984391dc985eca24db6b6b0c8dbf1c865e40 Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Tue, 11 Aug 2026 00:51:22 -0700 Subject: [PATCH 17/18] unframed prose escape hatch; force_nonempty_content in streaming; skip token oracle on draft --- .../function_call/muse_glimmer_detector.py | 18 ++-------- .../srt/function_call/muse_glimmer_format.py | 15 ++++++++ .../sglang/srt/model_executor/model_runner.py | 6 ++++ python/sglang/srt/parser/reasoning_parser.py | 35 +++++++++++++++++++ 4 files changed, 58 insertions(+), 16 deletions(-) diff --git a/python/sglang/srt/function_call/muse_glimmer_detector.py b/python/sglang/srt/function_call/muse_glimmer_detector.py index 83bcdce5f938..6c9cb02d3015 100644 --- a/python/sglang/srt/function_call/muse_glimmer_detector.py +++ b/python/sglang/srt/function_call/muse_glimmer_detector.py @@ -23,6 +23,7 @@ MESSAGE, RECIPIENT_RE, START, + could_start_header, has_atem_markers, partial_marker_len, ) @@ -68,21 +69,6 @@ def _normalize_name(emitted: str, registered: Set[str]) -> str: return emitted -def _could_start_header(text: str) -> bool: - """Whether the tail could still grow into a header.""" - stripped = text.lstrip() - if not stripped: - return True - if not (stripped.startswith("to=") or "to=".startswith(stripped[:3])): - return False - if MESSAGE in stripped: - return True - recipient, angle, marker = stripped[3:].partition("<") - if any(c.isspace() for c in recipient): - return False - return not angle or MESSAGE.startswith("<" + marker) - - class MuseGlimmerDetector(BaseFormatDetector): """Format detector for Muse Glimmer's ATEM tool-call blocks.""" @@ -162,7 +148,7 @@ def parse_streaming_increment( while self._buffer: if not self._in_body: # Resolve the channel header before anything can be emitted. - if self._at_stream_start and _could_start_header(self._buffer): + if self._at_stream_start and could_start_header(self._buffer): pass else: ws = len(self._buffer) - len(self._buffer.lstrip()) diff --git a/python/sglang/srt/function_call/muse_glimmer_format.py b/python/sglang/srt/function_call/muse_glimmer_format.py index fbd91097721e..b33135e6a6c4 100644 --- a/python/sglang/srt/function_call/muse_glimmer_format.py +++ b/python/sglang/srt/function_call/muse_glimmer_format.py @@ -22,6 +22,21 @@ MAX_MARKER = max(MAX_CHANNEL_MARKER, len(FUNCTION_CALLS_OPEN)) +def could_start_header(text: str) -> bool: + """Whether the tail could still grow into a header.""" + stripped = text.lstrip() + if not stripped: + return True + if not (stripped.startswith("to=") or "to=".startswith(stripped[:3])): + return False + if MESSAGE in stripped: + return True + recipient, angle, marker = stripped[3:].partition("<") + if any(c.isspace() for c in recipient): + return False + return not angle or MESSAGE.startswith("<" + marker) + + def has_atem_markers(text: str) -> bool: return INVOKE_OPEN in text or FUNCTION_CALLS_OPEN in text diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 6b30e3e19701..d8005cca806b 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -722,6 +722,12 @@ def maybe_init_elastic_ep(self): ElasticEPStateManager.init(self.server_args) def init_token_oracle(self): + # The oracle sampler is installed process-wide, so only the target + # publishes it; a draft would overwrite it with its own vocab (and a + # DFlash draft has none of its own until the worker borrows one). + if self.is_draft_worker: + self._token_oracle_manager = None + return self._token_oracle_manager = install_token_oracle_from_env( server_args=self.server_args, vocab_size=self.model_config.vocab_size, diff --git a/python/sglang/srt/parser/reasoning_parser.py b/python/sglang/srt/parser/reasoning_parser.py index 0ff158fc655e..2141cf2b5827 100644 --- a/python/sglang/srt/parser/reasoning_parser.py +++ b/python/sglang/srt/parser/reasoning_parser.py @@ -30,6 +30,7 @@ MESSAGE, RECIPIENT_RE, START, + could_start_header, has_atem_markers, partial_marker_len, ) @@ -1682,6 +1683,7 @@ def __init__( ) self._recipient: Optional[str] = None self._in_body = False + self._at_stream_start = True self._pending_reasoning = "" self._tool_call_parser_active = tool_call_parser_active self._saw_reasoning_block = False @@ -1705,12 +1707,31 @@ def _consume(self, flush: bool, preserve_channels: bool = False) -> Tuple[str, s while self._buffer: if not self._in_body: + # Decide whether a header is still coming before buffering for + # one; otherwise unframed prose never streams (the <|message|> + # it waits for never arrives). Mirrors the function-call + # detector's own header resolution. + if not (self._at_stream_start and could_start_header(self._buffer)): + ws = len(self._buffer) - len(self._buffer.lstrip()) + head = self._buffer[ws : ws + len(START)] + if not START.startswith(head): + self._in_body = True + self._recipient = None + self._at_stream_start = False + continue + if ws: + normal_parts.append(self._buffer[:ws]) + self._buffer = self._buffer[ws:] + if len(head) < len(START): + break + idx = self._buffer.find(MESSAGE) if idx == -1: if flush: normal_parts.append(self._buffer) self._buffer = "" break + self._at_stream_start = False header = self._buffer[:idx] m = RECIPIENT_RE.search(header) self._recipient = m.group(1) if m else "user" @@ -1772,6 +1793,7 @@ def detect_and_parse(self, text: str) -> StreamingParseResult: self._buffer = raw self._recipient = None self._in_body = False + self._at_stream_start = True self._saw_reasoning_block = False reasoning, normal = self._consume(flush=True, preserve_channels=True) return self._maybe_apply_force_nonempty_content( @@ -1788,6 +1810,14 @@ def parse_streaming_increment(self, new_text: str) -> StreamingParseResult: reasoning = "" if not self._in_body and self._pending_reasoning: reasoning, self._pending_reasoning = self._pending_reasoning, "" + if self._force_nonempty_content: + # Base contract: keep a copy of the reasoning so finish() can promote + # it to content if the turn produces none. Drop it when real content + # arrives, NOT when the reasoning channel closes -- <|eom|> lands in + # the same chunk as the last reasoning text. + self._accumulated_reasoning += reasoning + if normal: + self._accumulated_reasoning = "" return StreamingParseResult(normal_text=normal, reasoning_text=reasoning) def finish(self) -> StreamingParseResult: @@ -1797,6 +1827,11 @@ def finish(self) -> StreamingParseResult: if self._pending_reasoning: reasoning = self._pending_reasoning + reasoning self._pending_reasoning = "" + if self._force_nonempty_content: + promoted = self._accumulated_reasoning + reasoning + self._accumulated_reasoning = "" + if not normal and promoted: + return StreamingParseResult(normal_text=promoted) return StreamingParseResult(normal_text=normal, reasoning_text=reasoning) From 38a1bc5d2f21bddff8865c54c8f7e0b12b87c4f2 Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Tue, 11 Aug 2026 00:55:48 -0700 Subject: [PATCH 18/18] let checkpoints keep their language_model_only declaration; trim comments --- python/sglang/srt/configs/model_config.py | 5 ++++- python/sglang/srt/model_executor/model_runner.py | 5 ++--- python/sglang/srt/parser/reasoning_parser.py | 13 +++++-------- 3 files changed, 11 insertions(+), 12 deletions(-) diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index c34cdd2215f5..62f6bfef5d44 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -536,7 +536,10 @@ def __init__( self.hf_config.encoder_only = encoder_only self.hf_config.language_only = language_only - self.hf_config.language_model_only = language_model_only + # Checkpoints declare this one themselves (hf_transformers/processor.py), + # so the flag may only turn it on: writing the default back would build a + # vision tower with no weights to fill. + self.hf_config.language_model_only = language_model_only or self.is_lm_only # matryoshka embeddings self.matryoshka_dimensions = getattr( diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index d8005cca806b..82ce80201df2 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -722,9 +722,8 @@ def maybe_init_elastic_ep(self): ElasticEPStateManager.init(self.server_args) def init_token_oracle(self): - # The oracle sampler is installed process-wide, so only the target - # publishes it; a draft would overwrite it with its own vocab (and a - # DFlash draft has none of its own until the worker borrows one). + # The oracle sampler is process-wide, so a draft would overwrite the + # target's with its own vocab -- which a DFlash draft does not have. if self.is_draft_worker: self._token_oracle_manager = None return diff --git a/python/sglang/srt/parser/reasoning_parser.py b/python/sglang/srt/parser/reasoning_parser.py index 2141cf2b5827..50b60f543438 100644 --- a/python/sglang/srt/parser/reasoning_parser.py +++ b/python/sglang/srt/parser/reasoning_parser.py @@ -1707,10 +1707,8 @@ def _consume(self, flush: bool, preserve_channels: bool = False) -> Tuple[str, s while self._buffer: if not self._in_body: - # Decide whether a header is still coming before buffering for - # one; otherwise unframed prose never streams (the <|message|> - # it waits for never arrives). Mirrors the function-call - # detector's own header resolution. + # Without this, unframed prose never streams: it buffers + # forever waiting for a <|message|> that never arrives. if not (self._at_stream_start and could_start_header(self._buffer)): ws = len(self._buffer) - len(self._buffer.lstrip()) head = self._buffer[ws : ws + len(START)] @@ -1811,10 +1809,9 @@ def parse_streaming_increment(self, new_text: str) -> StreamingParseResult: if not self._in_body and self._pending_reasoning: reasoning, self._pending_reasoning = self._pending_reasoning, "" if self._force_nonempty_content: - # Base contract: keep a copy of the reasoning so finish() can promote - # it to content if the turn produces none. Drop it when real content - # arrives, NOT when the reasoning channel closes -- <|eom|> lands in - # the same chunk as the last reasoning text. + # Kept so finish() can promote it to content if the turn produces + # none. Dropped on real content, NOT when the channel closes -- + # <|eom|> lands in the same chunk as the last reasoning text. self._accumulated_reasoning += reasoning if normal: self._accumulated_reasoning = ""