diff --git a/python/sglang/benchmark/utils.py b/python/sglang/benchmark/utils.py index c3cbebf0c..cbcf24d92 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 5e6f5a579..62f9fdb21 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 3b94061c8..54b5fa46c 100644 --- a/python/sglang/srt/configs/__init__.py +++ b/python/sglang/srt/configs/__init__.py @@ -37,6 +37,10 @@ from sglang.srt.configs.locate_anything import LocateAnythingConfig 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 @@ __all__ = [ "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 432f74ef0..62f6bfef5 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -264,6 +264,7 @@ class ModelConfig: 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 @@ class ModelConfig: ) # 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,10 @@ class ModelConfig: self.hf_config.encoder_only = encoder_only self.hf_config.language_only = language_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( @@ -582,6 +588,7 @@ class ModelConfig: 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=( @@ -1830,6 +1837,7 @@ multimodal_model_archs = [ "MossVLForConditionalGeneration", "NemotronH_Nano_VL_V2", "NemotronH_Nano_Omni_Reasoning_V3", + "MuseGlimmerForConditionalGeneration", "PixtralForConditionalGeneration", "Qwen2AudioForConditionalGeneration", "Qwen2VLForConditionalGeneration", @@ -1893,6 +1901,7 @@ multimodal_breakable_cuda_graph_supported_model_archs = [ "InternS2MobiusForConditionalGeneration", "Qwen3_5ForConditionalGeneration", "Qwen3_5MoeForConditionalGeneration", + "MuseGlimmerForConditionalGeneration", ] if external_mm_model_arch := envs.SGLANG_EXTERNAL_MM_MODEL_ARCH.get(): @@ -2036,6 +2045,8 @@ def is_hybrid_swa_model( "Gemma4UnifiedForConditionalGeneration", "LagunaForCausalLM", "MellumForCausalLM", + "MuseGlimmerForCausalLM", + "MuseGlimmerForConditionalGeneration", "InklingForConditionalGeneration", "InklingForConditionalGenerationMTP", "UnlimitedOCRForCausalLM", @@ -2111,6 +2122,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 000000000..40866764c --- /dev/null +++ b/python/sglang/srt/configs/muse_glimmer.py @@ -0,0 +1,282 @@ +# 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" + # The DFlash draft has no head; draft_worker_common borrows the target's. + vocab_size = None + + +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 _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 + + 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}"] + + 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")] + + 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")), + 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], + use_qk_norm=any(n.endswith("attn_q_norm.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, + 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/configs/muse_glimmer_processing.py b/python/sglang/srt/configs/muse_glimmer_processing.py new file mode 100644 index 000000000..433d35ceb --- /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 a2f48d9b7..5aa6fcca1 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/serving_chat.py b/python/sglang/srt/entrypoints/openai/serving_chat.py index 34dceaa2b..6919dadf6 100644 --- a/python/sglang/srt/entrypoints/openai/serving_chat.py +++ b/python/sglang/srt/entrypoints/openai/serving_chat.py @@ -727,11 +727,7 @@ class OpenAIServingChat(OpenAIServingBase): 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 @@ class OpenAIServingChat(OpenAIServingBase): 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 @@ class OpenAIServingChat(OpenAIServingBase): 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 @@ class OpenAIServingChat(OpenAIServingBase): 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 @@ class OpenAIServingChat(OpenAIServingBase): 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 @@ class OpenAIServingChat(OpenAIServingBase): # 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 @@ class OpenAIServingChat(OpenAIServingBase): 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 @@ class OpenAIServingChat(OpenAIServingBase): 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 @@ class OpenAIServingChat(OpenAIServingBase): 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 @@ class OpenAIServingChat(OpenAIServingBase): 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 516467359..81f9c3e77 100644 --- a/python/sglang/srt/entrypoints/openai/serving_responses.py +++ b/python/sglang/srt/entrypoints/openai/serving_responses.py @@ -394,7 +394,6 @@ class OpenAIServingResponses(OpenAIServingChat): else None ), ) - # _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 ( @@ -592,7 +591,6 @@ class OpenAIServingResponses(OpenAIServingChat): 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] @@ -808,6 +806,7 @@ class OpenAIServingResponses(OpenAIServingChat): *, 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 +818,11 @@ class OpenAIServingResponses(OpenAIServingChat): ), 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 +855,6 @@ class OpenAIServingResponses(OpenAIServingChat): ) 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 +1980,7 @@ class OpenAIServingResponses(OpenAIServingChat): ), request=request, tokenizer=self.tokenizer_manager.tokenizer, + tool_call_parser_active=isinstance(tool_parser, FunctionCallParser), ) current_output_index = -1 @@ -2003,6 +2007,7 @@ class OpenAIServingResponses(OpenAIServingChat): 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 @@ -2210,11 +2215,26 @@ class OpenAIServingResponses(OpenAIServingChat): stream_offset = len(text) if not delta and finish_reason is None: continue + # finish_reason is sticky, so it would otherwise re-flush. + flush = ( + 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( 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 +2298,7 @@ class OpenAIServingResponses(OpenAIServingChat): ) ) - if not delta: + if not delta and not flush: continue if isinstance(tool_parser, JsonArrayParser): @@ -2286,6 +2306,10 @@ class OpenAIServingResponses(OpenAIServingChat): 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 83a512ad3..1a1ca91bb 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -652,6 +652,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 405623006..dd848402e 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 @@ class BaseFormatDetector(ABC): """ 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 669228a81..011114b46 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.minicpm5_detector import MiniCPM5Detector 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 @@ class FunctionCallParser: 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 000000000..6c9cb02d3 --- /dev/null +++ b/python/sglang/srt/function_call/muse_glimmer_detector.py @@ -0,0 +1,259 @@ +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, +) +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, + could_start_header, + has_atem_markers, + partial_marker_len, +) + +logger = logging.getLogger(__name__) + +_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 + + +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 + + +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 has_atem_markers(text) + + def _registered_names(self, tools: Optional[List[Tool]]) -> 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: + markers = (MESSAGE, EOM, EOT, START, FUNCTION_CALLS_OPEN, INVOKE_OPEN) + return partial_marker_len(text, markers, MAX_MARKER) + + def parse_streaming_increment( + self, new_text: str, tools: List[Tool] + ) -> 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{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 000000000..b33135e6a --- /dev/null +++ b/python/sglang/srt/function_call/muse_glimmer_format.py @@ -0,0 +1,54 @@ +"""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: + """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 + + +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/hardware_backend/mlx/model_runner.py b/python/sglang/srt/hardware_backend/mlx/model_runner.py index c69c0b073..1afa3fe1f 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 @@ from sglang.srt.hardware_backend.mlx.kv_cache import ( 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 @@ class MlxModelRunner: 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 @@ class MlxModelRunner: # 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 @@ -481,10 +500,18 @@ class MlxModelRunner: 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 000000000..398088fe2 --- /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 ``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), +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 + ``"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 +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. +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. +_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": "muse_glimmer", + "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 = "muse_glimmer" + 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. + muse_glimmer_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.muse_glimmer_mlx_format is not None and ( + self.muse_glimmer_mlx_format != MUSE_GLIMMER_MLX_FORMAT_VERSION + ): + raise ValueError( + 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" + ) + + +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.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( + 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"(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" + ) + 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 "muse_glimmer_mlx_format": ' + f"{MUSE_GLIMMER_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'"muse_glimmer_mlx_format": {MUSE_GLIMMER_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 000000000..fca8f32bd --- /dev/null +++ b/python/sglang/srt/hardware_backend/mlx/remote_code_gate.py @@ -0,0 +1,126 @@ +# 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``); mlx-lm is + unpinned, so accept either name. + """ + from mlx_lm import utils as mlx_lm_utils + + 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: + """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 2766e03f5..8f654c457 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 @@ class MlxTpModelWorker(TpModelWorker): 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/layers/quantization/fp8_utils.py b/python/sglang/srt/layers/quantization/fp8_utils.py index e45d6a440..f5a044969 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": diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index c74ec817e..a91433dd3 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -829,7 +829,11 @@ class Scheduler( # 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 19d43d26e..d0b893bf5 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -467,7 +467,7 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): 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) @@ -1024,6 +1024,11 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): ) 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 2bbdb499a..457fe2f95 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) @@ -90,5 +95,7 @@ def configure_kv_cache_dtype( model_dtype, ) kv_cache_dtype = model_dtype + # "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/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 84bd42ffe..b5756935b 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -722,6 +722,11 @@ class ModelRunner: ElasticEPStateManager.init(self.server_args) def init_token_oracle(self): + # 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 self._token_oracle_manager = install_token_oracle_from_env( server_args=self.server_args, vocab_size=self.model_config.vocab_size, @@ -1301,6 +1306,7 @@ class ModelRunner: 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 08dd87afd..f060d0fcc 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_loader/gguf_name_maps.py b/python/sglang/srt/model_loader/gguf_name_maps.py new file mode 100644 index 000000000..434f337af --- /dev/null +++ b/python/sglang/srt/model_loader/gguf_name_maps.py @@ -0,0 +1,72 @@ +# 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_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 + + +# 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, +} diff --git a/python/sglang/srt/model_loader/loader.py b/python/sglang/srt/model_loader/loader.py index cb941a800..ebe6c516f 100644 --- a/python/sglang/srt/model_loader/loader.py +++ b/python/sglang/srt/model_loader/loader.py @@ -3038,8 +3038,14 @@ class GGUFModelLoader(BaseModelLoader): "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 17c11cd24..4adc9990d 100644 --- a/python/sglang/srt/models/dflash.py +++ b/python/sglang/srt/models/dflash.py @@ -43,6 +43,16 @@ if _is_npu: logger = logging.getLogger(__name__) +def _get_dflash_attention_type(config) -> AttentionType: + """Bidirectional over the draft block unless the checkpoint says 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,17 +67,12 @@ 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": + # 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, AttentionType.DECODER + return sliding_window_size, _get_dflash_attention_type(config) raise ValueError( "Unsupported DFLASH draft layer type. " f"layer_types[{layer_id}]={layer_type!r}." @@ -381,11 +386,12 @@ class DFlashDraftModel(nn.Module): # 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 ) @@ -470,6 +476,12 @@ class DFlashDraftModel(nn.Module): 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 @@ -481,6 +493,9 @@ class DFlashDraftModel(nn.Module): 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: @@ -608,4 +623,8 @@ class DFlashLagunaForCausalLM(DFlashDraftModel): 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 000000000..a1ddf6fbc --- /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 000000000..36d78bb3b --- /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 36bd69cf8..50b60f543 100644 --- a/python/sglang/srt/parser/reasoning_parser.py +++ b/python/sglang/srt/parser/reasoning_parser.py @@ -23,6 +23,17 @@ from sglang.srt.function_call.kimik3_format import ( 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, + could_start_header, + 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, @@ -1614,6 +1625,213 @@ class CohereCommand4Detector(BaseReasoningFormatDetector): 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. + """ + + 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" + MESSAGE, + 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._at_stream_start = True + 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" + + 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: + # 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)] + 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" + self._buffer = self._buffer[idx + len(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 + MESSAGE) + continue + + end_idx, end_tok = -1, "" + 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 + + 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 = partial_marker_len( + self._buffer, (EOM, EOT, START), MAX_CHANNEL_MARKER + ) + 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 has_atem_markers(normal): + 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( + 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, "" + if self._force_nonempty_content: + # 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 = "" + 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 = "" + 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) + + class ReasoningParser: """ Parser that handles both streaming and non-streaming scenarios for extracting @@ -1623,6 +1841,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 +1859,7 @@ class ReasoningParser: "kimi_k2": KimiK2Detector, "kimi_k3": KimiK3Detector, "mimo": _MimoDetector, + "muse": MuseGlimmerDetector, "poolside_v1": _PoolsideV1Detector, "qwen3": Qwen3Detector, "qwen3-thinking": Qwen3Detector, @@ -1660,6 +1883,7 @@ class ReasoningParser: 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 +1929,11 @@ class ReasoningParser: 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 a5046681e..ad34ea51b 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -1747,6 +1747,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" @@ -2158,6 +2159,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.", @@ -3093,6 +3108,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( @@ -3549,6 +3572,7 @@ class ServerArgs: # 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() @@ -3852,6 +3876,8 @@ class ServerArgs: 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, @@ -4303,6 +4329,16 @@ class ServerArgs: ): 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 @@ -4861,6 +4897,7 @@ class ServerArgs: 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) @@ -6392,7 +6429,11 @@ class ServerArgs: 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." @@ -7320,6 +7361,32 @@ class ServerArgs: 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 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 # (arg_groups/overrides.py: _gguf_quantization); load_format itself is @@ -7474,7 +7541,39 @@ class ServerArgs: 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 6cfe26f5a..70c747fe6 100644 --- a/python/sglang/srt/speculative/dflash_utils.py +++ b/python/sglang/srt/speculative/dflash_utils.py @@ -515,7 +515,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: @@ -524,7 +526,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 d8d94b1d4..d0c3888ae 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 @@ class DFlashWorkerV2(BaseSpecWorker): 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 @@ class DFlashWorkerV2(BaseSpecWorker): 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 @@ class DFlashWorkerV2(BaseSpecWorker): 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 @@ class DFlashWorkerV2(BaseSpecWorker): 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 @@ class DFlashWorkerV2(BaseSpecWorker): 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 @@ class DFlashWorkerV2(BaseSpecWorker): 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 @@ class DFlashWorkerV2(BaseSpecWorker): 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 b393ce8d0..28a150f73 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 2ad7d53c3..56d1c150e 100644 --- a/python/sglang/srt/utils/hf_transformers/__init__.py +++ b/python/sglang/srt/utils/hf_transformers/__init__.py @@ -35,6 +35,7 @@ from .common import ( 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, resolve_image_processor_backend @@ -51,6 +52,7 @@ __all__ = [ "_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 db584861b..efc8dee39 100644 --- a/python/sglang/srt/utils/hf_transformers/common.py +++ b/python/sglang/srt/utils/hf_transformers/common.py @@ -52,6 +52,8 @@ from sglang.srt.configs import ( MiniCPMV4_6VisionConfig, MiniMaxM3VLConfig, MultiModalityConfig, + MuseGlimmerAssistantConfig, + MuseGlimmerConfig, NemotronH_Nano_Omni_Reasoning_V3_Config, NemotronH_Nano_VL_V2_Config, NemotronHConfig, @@ -101,6 +103,8 @@ _CONFIG_REGISTRY: Dict[str, Type[PretrainedConfig]] = { 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 1c26fbeeb..c0b258997 100644 --- a/python/sglang/srt/utils/hf_transformers/config.py +++ b/python/sglang/srt/utils/hf_transformers/config.py @@ -37,8 +37,10 @@ from .common import ( _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 000000000..a9a95b6cd --- /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 def8d2581..b6d25a1d8 100644 --- a/python/sglang/srt/utils/hf_transformers/tokenizer.py +++ b/python/sglang/srt/utils/hf_transformers/tokenizer.py @@ -34,8 +34,10 @@ from .common import ( _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/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 000000000..1af9ebd63 --- /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 932456998..ab6c561ab 100644 --- a/test/registered/unit/entrypoints/openai/test_serving_responses.py +++ b/test/registered/unit/entrypoints/openai/test_serving_responses.py @@ -311,6 +311,80 @@ class ReasoningRequestForwardingTestCase(unittest.TestCase): 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) + # 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) + + 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 5c00ce3c5..c199a7aa6 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 @@ class NonHarmonyStreamTestCase(CustomTestCase): 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 @@ class NonHarmonyStreamTestCase(CustomTestCase): 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 000000000..8e75ceb1b --- /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 000000000..8d44d12d0 --- /dev/null +++ b/test/registered/unit/hardware_backend/mlx/test_mlx_remote_code_gate.py @@ -0,0 +1,186 @@ +"""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": "muse_glimmer", "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": "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. + 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": "muse_glimmer", "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": "muse_glimmer", "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": "muse_glimmer", "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": "muse_glimmer", "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": "muse_glimmer", "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": "muse_glimmer", "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 000000000..35dd77bb0 --- /dev/null +++ b/test/registered/unit/hardware_backend/mlx/test_muse_glimmer_mlx_model.py @@ -0,0 +1,278 @@ +"""Unit tests for the Muse Glimmer MLX model file's load path — no weights, no server. + +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 + 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"], "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) + # 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, "muse_glimmer_mlx_format"): + ModelArgs(**_TINY, muse_glimmer_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(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(muse_glimmer_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() 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 0ebf64efa..87f530d3d 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 @@ class TestKVCacheQuantRegistry(CustomTestCase): 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 c1411e7a8..73c502b62 100644 --- a/test/registered/unit/test_model_overrides.py +++ b/test/registered/unit/test_model_overrides.py @@ -89,6 +89,7 @@ class TestModelOverridableWhitelist(CustomTestCase): "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",