Hugging Face's logo

gitcommit90
/
SenseNova-Loop-9B-Preview-NVFP4

Image-Text-to-Text
Transformers
Safetensors
English
Chinese
sensenova_loop
text-generation
nvfp4
modelopt
quantized
multimodal
looped-transformer
mhc
mtp
custom_code
conversational
8-bit precision
Model card Files
xet
 Community
SenseNova-Loop-9B-Preview-NVFP4 / runtime /overlay /python /sglang /srt /configs /model_config.py
gitcommit90's picture
gitcommit90
NVFP4 (ModelOpt) quantization of sensenova/SenseNova-Loop-9B-Preview @27abea77 with BF16 MTP head
a57aa9f verified
Raw History Blame Contribute Delete
123 kB
# Copyright 2023-2024 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 copy
import json
import logging
import math
import os
from enum import Enum, IntEnum, auto
from functools import cached_property, lru_cache
from pathlib import Path
from typing import Any, Callable, List, Optional, Set, Union
import torch
from transformers import PretrainedConfig
from transformers.utils import cached_file
from sglang.srt.arg_groups.overrides import resolving_view
from sglang.srt.configs.bailing_hybrid import is_bailing_multi_gate_enabled
from sglang.srt.configs.embedding_model_spec import resolve_embedding_model_spec
from sglang.srt.configs.linear_attn_model_registry import get_linear_attn_config
from sglang.srt.environ import envs
from sglang.srt.layers.quantization import QUANTIZATION_METHODS
from sglang.srt.runtime_context import get_platform
from sglang.srt.server_args import ServerArgs
from sglang.srt.utils import is_hip, retry
from sglang.srt.utils.hf_transformers.config import POINTER_ADAPTER_BASES
from sglang.srt.utils.hf_transformers_utils import (
get_config,
get_context_length,
get_generation_config,
get_hf_text_config,
get_sparse_attention_config,
)
from sglang.srt.utils.runai_utils import ObjectStorageModel, is_runai_obj_uri
from sglang.utils import is_in_ci
logger = logging.getLogger(__name__)
_MODEL_CONFIG_FACTORIES: dict[type[ServerArgs], Callable[..., "ModelConfig"]] = {}
def register_model_config_factory(
server_args_type: type[ServerArgs], factory: Callable[..., "ModelConfig"]
) -> None:
"""Register a factory with the same signature as ModelConfig.from_server_args.
Register before resolution or model construction in each process. The nearest
registered type in the argument record's MRO wins, so registrations cover
subclasses too. Repeating the same registration is harmless; replacing a
different factory for the same type is an error.
"""
if not issubclass(server_args_type, ServerArgs):
raise TypeError("model-config factories require a ServerArgs subclass")
previous = _MODEL_CONFIG_FACTORIES.get(server_args_type)
if previous is not None and previous is not factory:
raise ValueError(
f"A model-config factory is already registered for {server_args_type.__qualname__}"
)
_MODEL_CONFIG_FACTORIES[server_args_type] = factory
MIMO_V2_MODEL_ARCHS = (
"MiMoV2ForCausalLM",
"MiMoV2FlashForCausalLM",
)
MIMO_V2_MULTIMODAL_ARCHS = ("MiMoV2ForCausalLM",)
BAILING_MULTI_GATE_MM_ARCHS = frozenset(
{
"BailingMMNativeForConditionalGeneration",
"BailingMM2NativeForConditionalGeneration",
"BailingMoeV3VLForConditionalGeneration",
}
)
SWA_SINK_ARCHS = frozenset(
{
"GptOssForCausalLM",
"GraniteSWAForCausalLM",
"GraniteMoeSWAForCausalLM",
"IQuestQ1ForCausalLM",
"IQuestQ1MTP",
}
)
def _quant_config_to_dict(quant_config):
if quant_config is not None and not isinstance(quant_config, dict):
return quant_config.to_dict()
return quant_config
def requires_mm_token_modalities(
model_architectures: Optional[List[str]], hf_text_config: PretrainedConfig
) -> bool:
"""Whether a Bailing multimodal wrapper uses modality-specific routers."""
return bool(
model_architectures
and any(arch in BAILING_MULTI_GATE_MM_ARCHS for arch in model_architectures)
and is_bailing_multi_gate_enabled(hf_text_config)
)
def unwrap_modelopt_quantization_config(quant_config: dict) -> dict:
quantization = quant_config.get("quantization", quant_config)
if not isinstance(quantization, dict):
return {}
nested = quantization.get("quantization")
return nested if isinstance(nested, dict) else quantization
def get_mimo_v2_fused_qkv_expected_tp_size(hf_config):
layout = getattr(hf_config, "attention_projection_layout", None)
if layout in (None, "split"):
return None
if layout != "fused_qkv":
raise ValueError(
"MiMoV2 hf_config has unsupported "
f"attention_projection_layout={layout!r}; expected 'fused_qkv', "
"'split' or unset."
)
num_key_value_heads = getattr(hf_config, "num_key_value_heads", None)
text_config = getattr(hf_config, "text_config", None)
if num_key_value_heads is None and text_config is not None:
num_key_value_heads = getattr(text_config, "num_key_value_heads", None)
if num_key_value_heads is None:
raise ValueError(
"MiMoV2 hf_config has attention_projection_layout='fused_qkv' "
"but num_key_value_heads is missing; this value is required to "
"derive the fused qkv_proj TP size."
)
return num_key_value_heads
class AttentionArch(IntEnum):
MLA = auto()
MHA = auto()
SSM = auto() # State Space Models (Mamba, Mamba2)
# Pure Mamba-1 (selective-scan) archs; same mixer/state layout, differing only
# in cosmetic details handled in their model files.
PURE_MAMBA1_ARCHITECTURES = (
"FalconMambaForCausalLM",
"MambaForCausalLM",
)
# Pure state-space (SSM) causal-LMs: no attention, so no num_attention_heads /
# head_dim in their HF config. Used for head-dim derivation and attention-arch
# detection below.
PURE_SSM_ARCHITECTURES = ("Mamba2ForCausalLM",) + PURE_MAMBA1_ARCHITECTURES
class ModelImpl(str, Enum):
AUTO = "auto"
SGLANG = "sglang"
TRANSFORMERS = "transformers"
MINDSPORE = "mindspore"
def _hf_arch(config) -> Optional[str]:
"""First architecture from a HF config dict or PretrainedConfig (or None)."""
archs = (
config.get("architectures")
if isinstance(config, dict)
else getattr(config, "architectures", None)
)
return archs[0] if archs else None
def _hf_attr(config, name):
"""Read an arbitrary field from a HF config dict or PretrainedConfig."""
if isinstance(config, dict):
return config.get(name)
return getattr(config, name, None)
def is_deepseek_dsa(config) -> bool:
return (
_hf_arch(config)
in (
"DeepseekV3ForCausalLM",
"DeepseekV32ForCausalLM",
"DeepseekV3ForCausalLMNextN",
"MistralLarge3ForCausalLM",
"PixtralForConditionalGeneration",
"GlmMoeDsaForCausalLM",
"GlmMoeDsaForCausalLMNextN",
"Glm5NextForConditionalGenerationNextN",
"Glm5NextForConditionalGeneration",
"LongcatFlashForCausalLM",
"LongcatFlashForCausalLMNextN",
"Dots3NoteForCausalLM",
"Dots3NoteForCausalLMNextN",
"HYV4ForCausalLM",
"HYV4ForCausalLMNextN",
)
and _hf_attr(config, "index_topk") is not None
)
def dsa_draft_shards_under_dcp(config, dcp_size: int) -> bool:
"""Whether a DSA EAGLE/NEXTN draft shards its KV like the DCP target.
Only on ROCm; elsewhere the draft keeps a replicated pool and attends it
with dcp_size 1.
"""
return dcp_size > 1 and is_hip() and is_deepseek_dsa(config)
def is_glm_moe_dsa(config) -> bool:
"""True for GLM-5.2, both the main arch and the NextN draft head."""
return _hf_arch(config) in (
"GlmMoeDsaForCausalLM",
"GlmMoeDsaForCausalLMNextN",
)
def is_kimi_k3(config) -> bool:
return _hf_arch(config) in (
"KimiK3ForConditionalGeneration",
"KimiK3LinearForCausalLM",
)
def uses_kda_attention(config) -> bool:
configs = [config]
get_text_config = getattr(config, "get_text_config", None)
if callable(get_text_config):
configs.append(get_text_config())
else:
text_config = _hf_attr(config, "text_config")
if text_config is not None:
configs.append(text_config)
for config in configs:
linear_attn_config = _hf_attr(config, "linear_attn_config")
if isinstance(linear_attn_config, dict) and linear_attn_config.get(
"kda_layers"
):
return True
layer_types = _hf_attr(config, "layer_types") or []
if (
"linear_attention" in layer_types
and _hf_attr(config, "linear_num_heads") is not None
and _hf_attr(config, "linear_head_dim") is not None
):
return True
return False
def is_dspark_draft(config) -> bool:
return _hf_arch(config) in ("DSparkDraftModel", "DSparkMiniMaxDraftModel")
# The decision_config.json attention mode whose full-attention layers read the whole prompt.
NONCAUSAL_FULL_ATTENTION = "noncausal_full_attention"
@lru_cache
def load_decision_config(model_path: str, revision: Optional[str]) -> Optional[dict]:
"""decision_config.json of a checkpoint whose LM head is a decision readout."""
path = cached_file(
model_path,
"decision_config.json",
revision=revision,
_raise_exceptions_for_missing_entries=False,
)
if path is None:
return None
with open(path) as f:
config = json.load(f)
if config["format_version"] != 1:
raise ValueError(
f"decision_config.json format_version {config['format_version']} "
"is not supported, only version 1 is"
)
# Serving a mode this loader does not honor would silently change the answers.
attention_mode = config.get("attention_mode", "causal")
if attention_mode not in ("causal", NONCAUSAL_FULL_ATTENTION):
raise ValueError(
f"decision_config.json attention_mode {attention_mode!r} is not supported"
)
pooling = config.get("pooling", "last")
if pooling != "last":
raise ValueError(
f"decision_config.json pooling {pooling!r} is not supported, only 'last' is"
)
return config
@lru_cache
def load_joint_head_config(model_path: str, revision: Optional[str]) -> Optional[dict]:
"""joint_head_config.json of a Clef checkpoint, whose joint schema head
answers /v1/systemone decisions from the backbone's hidden states."""
path = cached_file(
model_path,
"joint_head_config.json",
revision=revision,
_raise_exceptions_for_missing_entries=False,
_raise_exceptions_for_connection_errors=False,
)
if path is None:
return None
with open(path) as f:
return json.load(f)
# The values the Kev reference reads for head.pt fields that older checkpoints
# did not write.
_POINTER_HEAD_DEFAULTS = {
"base_revision": None,
"head_dim": 256,
"option_isolation": False,
"special_embeddings": False,
"weights_dtype": "fp32",
"temperature": 1.0,
}
@lru_cache
def load_pointer_head(model_path: str, revision: Optional[str]) -> Optional[dict]:
"""head.pt of a Kev checkpoint, whose pointer head answers /v1/systemone
decisions from the backbone's hidden states."""
path = cached_file(
model_path,
"head.pt",
revision=revision,
_raise_exceptions_for_missing_entries=False,
_raise_exceptions_for_connection_errors=False,
)
if path is None:
return None
meta = {
**_POINTER_HEAD_DEFAULTS,
**torch.load(path, map_location="cpu", weights_only=True),
}
# Both change the prompt or the attention mask the head was trained with.
for name in ("option_isolation", "special_embeddings"):
if meta[name]:
raise ValueError(f"head.pt {name} is not supported")
return meta
@lru_cache
def load_pointer_adapter(model_path: str, revision: Optional[str]) -> Optional[dict]:
"""head.pt of a Kev LoRA checkpoint with its merge scale. Such a checkpoint has
no model config and is served as the base backbone its head.pt pins, with the
adapter merged into the base weights at load."""
meta = load_pointer_head(model_path, revision)
if meta is None:
return None
path = cached_file(
model_path,
"adapter_config.json",
revision=revision,
_raise_exceptions_for_missing_entries=False,
_raise_exceptions_for_connection_errors=False,
)
if path is None:
return None
# Kev applies an adapter trained on a BF16 backbone unmerged, which a merge
# into the base weights does not reproduce.
if meta["weights_dtype"] != "fp32":
raise ValueError(
f"{model_path} is a Kev adapter trained on a {meta['weights_dtype']} "
"backbone, which SGLang does not serve. Serve its full weight release."
)
with open(path) as f:
adapter = json.load(f)
return {**meta, "lora_scale": adapter["lora_alpha"] / adapter["r"]}
def is_qwen3_5(config) -> bool:
return _hf_arch(config) in (
"Qwen3_5ForConditionalGeneration",
"Qwen3_5MoeForConditionalGeneration",
"Qwen3_5ForCausalLM",
"Qwen3_5MoeForCausalLM",
"SenseNovaLoopForConditionalGeneration",
)
def is_qwen3_5_mtp_draft(config) -> bool:
"""The Qwen3.5 MoE MTP draft: _config_draft_model rewrites architectures[0] to
Qwen3_5ForCausalLMMTP before quantization is resolved."""
return (
_hf_arch(config) == "Qwen3_5ForCausalLMMTP"
and _hf_attr(config, "model_type") == "qwen3_5_moe"
)
def is_deepseek_v4(config) -> bool:
return _hf_arch(config) in (
"DeepseekV4ForCausalLM",
"DeepseekV4ForCausalLMNextN",
"DeepseekV4ForCausalLMDSpark",
)
def is_qwen4_exp(config) -> bool:
return _hf_arch(config) in (
"Qwen4ExpForConditionalGeneration",
"Qwen4ExpForCausalLMMTP",
)
def resolve_spec_hidden_size(
hf_config, hidden_size: int, hc_mult: int
) -> tuple[int, Optional[int]]:
# DSV4 and Qwen4-Exp carry the hc-flattened stream across the target->draft
# boundary; other hc models (hy_v4) collapse to hidden_size first.
if hc_mult <= 1 or not (is_deepseek_v4(hf_config) or is_qwen4_exp(hf_config)):
return hidden_size, None
hc_hidden_size = hidden_size * hc_mult
return hc_hidden_size, hc_hidden_size
def get_dsa_index_head_dim(config: PretrainedConfig) -> int:
assert is_deepseek_dsa(config) or is_deepseek_v4(config)
return config.index_head_dim
def is_minimax_sparse(config: PretrainedConfig) -> bool:
arch = (config.architectures or [None])[0]
return arch in (
"MiniMaxM3SparseForCausalLM",
"MiniMaxM3SparseForConditionalGeneration",
)
def uses_minimax_m3_nvfp4_native_stack(quantization: Optional[str]) -> bool:
"""Whether an NVFP4 checkpoint runs its native Blackwell stack.
The attention backend, the KV pool class and the KV cell-size budget must
all agree on this answer; sizing the pool for one layout while allocating
the other runs out of memory at startup.
"""
from sglang.srt.runtime_context import get_platform
return quantization == "minimax_m3_nvfp4" and get_platform().is_blackwell
def get_minimax_sparse_attention_config(config: PretrainedConfig) -> dict:
text_cfg = getattr(config, "text_config", None)
cfg = (
getattr(text_cfg, "sparse_attention_config", None)
if text_cfg is not None
else None
)
if cfg is None:
cfg = getattr(config, "sparse_attention_config", None)
if cfg is None:
raise ValueError("Could not find sparse config. Is it MiniMax M3 Sparse model?")
return cfg
def get_minimax_sparse_layer_ids(sparse_cfg: dict) -> tuple[list[int], list[int]]:
sparse_freq = sparse_cfg["sparse_attention_freq"]
dense_layer_ids = [i for i, f in enumerate(sparse_freq) if f == 0]
sparse_layer_ids = [i for i, f in enumerate(sparse_freq) if f != 0]
return dense_layer_ids, sparse_layer_ids
def get_minimax_sparse_disable_value_layer_ids(sparse_cfg: dict) -> list[int]:
flags = sparse_cfg.get("sparse_disable_index_value")
if flags is None:
return []
return [i for i, f in enumerate(flags) if f != 0]
def get_minimax_sparse_score_type(sparse_cfg: dict) -> str:
score_type = sparse_cfg.get("sparse_score_type", "max")
assert score_type in (
"max",
"lse",
), f"sparse_score_type must be 'max' or 'lse', got {score_type!r}"
return score_type
def get_dsa_index_topk(config: PretrainedConfig) -> int:
assert is_deepseek_dsa(config)
return config.index_topk
def dsa_layer_skips_topk(config: PretrainedConfig, layer_id: int) -> bool:
"""Return whether a DSA layer reuses the previous layer's top-k indices."""
assert is_deepseek_dsa(config)
indexer_types = getattr(config, "indexer_types", None)
if indexer_types is not None:
return (
0 <= layer_id < len(indexer_types) and indexer_types[layer_id] == "shared"
)
# LongCat computes fresh top-k indices every cli_factor layers.
cli_factor = getattr(config, "cli_factor", 1)
if cli_factor is None:
cli_factor = 1
assert cli_factor > 0, f"cli_factor must be positive, got {cli_factor}"
if cli_factor > 1:
return layer_id % cli_factor != 0
pattern = getattr(config, "index_topk_pattern", None)
if pattern is not None:
return layer_id < len(pattern) and pattern[layer_id] == "S"
freq = getattr(config, "index_topk_freq", 1)
if freq is None:
freq = 1
assert freq > 0, f"index_topk_freq must be positive, got {freq}"
offset = getattr(config, "index_skip_topk_offset", None)
if offset is not None:
assert offset > 0, (
"index_skip_topk_offset must be positive; offset <= 0 "
"marks layer 0 as skip_topk with no prior topk to reuse"
)
return max(layer_id - offset + 1, 0) % freq != 0
return max(layer_id - 1, 0) % freq != 0
def get_dsa_index_n_heads(config: PretrainedConfig) -> int:
assert is_deepseek_dsa(config)
return config.index_n_heads
def get_dsa_index_kpool(config: PretrainedConfig) -> int:
return getattr(config, "index_kpool", 1)
def get_dsa_mtp_topk_width(config: PretrainedConfig) -> int:
"""MTP seeds include index_topk pooled tokens plus up to index_kpool - 1 tail tokens."""
index_kpool = get_dsa_index_kpool(config)
assert index_kpool >= 1, f"index_kpool must be positive, got {index_kpool}"
return config.index_topk + index_kpool - 1
def get_dsa_index_kpool_compress(config: PretrainedConfig) -> bool:
return getattr(config, "index_kpool_compress", False)
REQUANTIZATION_METHODS = ["quark_mxfp4"]
def get_num_indexer_layers(config) -> int:
"""Layer count for the global indexer-topk capturer's host buffer.
DSA models (V3.2) expose one capturer slot per transformer layer. With
index_topk_freq > 1 some layers reuse prev layer's topk; those still get a
slot mirrored at the MLA call site even if no Indexer module is built.
DSv4 has C4 indexers only on layers whose compress_ratio == 4. Other
architectures: set num_indexer_layers on hf_text_config; 0 disables the
capturer.
"""
if is_deepseek_dsa(config):
return config.num_hidden_layers
if is_deepseek_v4(config):
compress_ratios = getattr(config, "compress_ratios", None) or []
return sum(1 for r in compress_ratios if r == 4)
return getattr(config, "num_indexer_layers", 0)
class ModelConfig:
def __init__(
self,
model_path: str,
trust_remote_code: bool = True,
revision: Optional[str] = None,
context_length: Optional[int] = None,
model_override_args: str = "{}",
is_embedding: Optional[bool] = None,
enable_multimodal: Optional[bool] = None,
dtype: str = "auto",
quantization: Optional[str] = None,
override_config_file: Optional[str] = None,
is_draft_model: bool = False,
model_impl: Union[str, ModelImpl] = ModelImpl.AUTO,
sampling_defaults: str = "openai",
quantize_and_serve: bool = False,
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,
is_draft_quantization_explicit: bool = False,
) -> None:
# Parse args
self.model_path = model_path
self.revision = revision
self.quantization = quantization
self.is_draft_model = is_draft_model
self.is_draft_quantization_explicit = is_draft_quantization_explicit
self.speculative_algorithm = speculative_algorithm
self.model_impl = model_impl
self.sampling_defaults = sampling_defaults
self.quantize_and_serve = quantize_and_serve
self.is_multi_layer_eagle = is_multi_layer_eagle
self.disable_hybrid_swa_memory = disable_hybrid_swa_memory
self.model_config_parser = model_config_parser
# Validate quantize_and_serve configuration
self._validate_quantize_and_serve_config()
# Get hf config
self._maybe_pull_model_for_runai(self.model_path)
self._maybe_pull_model_tokenizer_from_remote()
self.model_override_args = json.loads(model_override_args)
kwargs = {}
if override_config_file and override_config_file.strip():
kwargs["_configuration_file"] = override_config_file.strip()
# get_config() is cached. ModelConfig mutates hf_config for draft-model
# remapping and architecture-specific normalization, so each instance
# must own an isolated copy.
self.hf_config = copy.deepcopy(
get_config(
self.model_path,
trust_remote_code=trust_remote_code,
revision=revision,
model_override_args=self.model_override_args,
model_config_parser=model_config_parser,
**kwargs,
)
)
self.pointer_head = None
base = POINTER_ADAPTER_BASES.get((self.model_path, revision))
if base is not None:
# A Kev LoRA checkpoint, whose base weights load with its adapter merged.
self.pointer_head = load_pointer_adapter(self.model_path, revision)
self.model_path, self.revision = base
revision = self.revision
self.hf_text_config = get_hf_text_config(self.hf_config)
self.decision_config = None
if self.hf_config.architectures == ["Qwen3_5Model"]:
self.decision_config = load_decision_config(self.model_path, revision)
if self.decision_config is not None:
# A bare backbone whose readout the loader places in the LM head.
self.hf_config.architectures = ["Qwen3_5ForConditionalGeneration"]
if (
self.decision_config.get("attention_mode")
== NONCAUSAL_FULL_ATTENTION
):
# Full attention sees the whole prompt, while linear attention stays causal.
self.hf_text_config.is_causal = False
self.joint_head_config = None
if self.hf_config.architectures == ["Qwen3_5ForConditionalGeneration"]:
self.joint_head_config = load_joint_head_config(self.model_path, revision)
if self.hf_config.architectures == ["Qwen3_5TextModel"]:
self.pointer_head = load_pointer_head(self.model_path, revision)
if self.pointer_head is not None:
# A Kev backbone without an LM head, served by the text-only class.
self.hf_config.architectures = ["Qwen3_5ForCausalLM"]
self.requires_mm_token_modalities = requires_mm_token_modalities(
self.hf_config.architectures, self.hf_text_config
)
self.is_embedding_gemma = is_embedding_gemma(self.hf_text_config)
self.embedding_model_spec = resolve_embedding_model_spec(
self.hf_config.architectures,
is_embedding_requested=bool(is_embedding),
is_embedding_gemma=self.is_embedding_gemma,
)
rope_scaling = getattr(self.hf_text_config, "rope_parameters", None) or getattr(
self.hf_text_config, "rope_scaling", {}
)
# Text-only serving comes from either the checkpoint's own declaration
# or the --language-model-only flag; every capability flag below
# (is_multimodal, is_*_understandable_model) must see both sources,
# since /model_info advertises them and drives media warmup requests.
self.is_lm_only = language_model_only or getattr(
self.hf_config, "language_model_only", False
)
self.model_is_mrope = (
not self.is_lm_only
and rope_scaling is not None
and "mrope_section" in rope_scaling
)
self.hf_generation_config = get_generation_config(
self.model_path,
trust_remote_code=trust_remote_code,
revision=revision,
**kwargs,
)
# Set enable_multimodal
if enable_multimodal is None:
mm_disabled_models = [
"Gemma3ForConditionalGeneration",
"Llama4ForConditionalGeneration",
"Step3VLForConditionalGeneration",
"InklingForConditionalGeneration",
]
if (
self.hf_config.architectures[0] in mm_disabled_models
and self.model_impl != ModelImpl.TRANSFORMERS
):
enable_multimodal = False
logger.info(
f"Multimodal is disabled for {self.hf_config.model_type}. To enable it, set --enable-multimodal."
)
elif self.hf_config.architectures[0] in MIMO_V2_MULTIMODAL_ARCHS and not (
hasattr(self.hf_config, "vision_config")
and hasattr(self.hf_config, "audio_config")
):
enable_multimodal = False
logger.info(
"Multimodal is disabled for this MiMoV2 checkpoint: "
"vision_config/audio_config not found in the model config "
"(likely a text-only MiMoV2 variant)."
)
else:
enable_multimodal = True
# Config draft model
self._config_draft_model()
# Mixed FP8/MXFP4 ckpts mark mxfp4 routed experts via this key.
quantization_config = (
_quant_config_to_dict(getattr(self.hf_config, "quantization_config", None))
or {}
)
self.hf_quant_config: dict = quantization_config
routed_experts_quant_method = quantization_config.get(
"routed_experts_quant_method"
)
self.is_fp4_experts: bool = routed_experts_quant_method == "mxfp4"
if self.is_fp4_experts:
logger.info("Detected mixed checkpoint layout: routed experts are MXFP4.")
# MiMo-V2 mxfp4 ckpts declare the routed-expert layout via store_dtype.
if (
not self.is_fp4_experts
and _hf_arch(self.hf_config) in MIMO_V2_MODEL_ARCHS
and str(quantization_config.get("store_dtype") or "").lower() == "mxfp4"
):
self.is_fp4_experts = True
logger.info("Detected MiMo-V2 mxfp4 routed-expert layout.")
# DSV4 mxfp4 layout applies only when the ckpt does not opt in above.
if is_deepseek_v4(self.hf_config) and routed_experts_quant_method is None:
self.is_fp4_experts = envs.SGLANG_DSV4_FP4_EXPERTS.get()
if (
not envs.SGLANG_DSV4_FP4_EXPERTS.is_set()
or envs.SGLANG_DSV4_FP4_DEQUANT.is_set()
):
from sglang.srt.configs.deepseek_v4 import try_detect_fp4_experts
detected = try_detect_fp4_experts(self.model_path)
if detected is not None:
self.is_fp4_experts = detected
logger.info(
"Auto-detected DSV4 routed-expert layout: is_fp4_experts=%s",
self.is_fp4_experts,
)
if envs.SGLANG_DSV4_FP4_DEQUANT.get():
envs.SGLANG_DSV4_FP4_DEQUANT.set(self.is_fp4_experts is not None)
# HF config.json inherits topk_group=4 from the V3 template, but
# DSV4 trains with no group limiting (sqrtsoftplus + full-expert
# top-k). Force topk_group == n_group so deepseek_v2.py:531's
# `n_group > topk_group` evaluates False and routes to the
# ungrouped sqrtsoftplus path. The grouped impl only supports
# sigmoid scoring (topk.py:722) and would silently corrupt expert
# weights if hit.
n_group = getattr(self.hf_config, "n_group", None)
if n_group is not None:
self.hf_config.topk_group = n_group
# Handle hybrid NVFP4 moe (nvidia/DeepSeek-V4-Pro-NVFP4)
self.nvfp4_moe_meta: Optional[dict] = None
hybrid_quant_cfg = _quant_config_to_dict(
getattr(self.hf_config, "quantization_config", None)
)
if (
hybrid_quant_cfg is not None
and str(hybrid_quant_cfg.get("quant_algo", "")).upper() == "MIXED_PRECISION"
and str(hybrid_quant_cfg.get("moe_quant_algo", "")).upper() == "NVFP4"
and hybrid_quant_cfg.get("group_size") is not None
):
self.nvfp4_moe_meta = {
"group_size": int(hybrid_quant_cfg["group_size"]),
"exclude_modules": list(hybrid_quant_cfg.get("ignore") or []),
}
logger.info(
"Auto-detected hybrid FP8+NVFP4 checkpoint "
"(NVFP4 MoE group_size=%d, %d exclude_modules)",
self.nvfp4_moe_meta["group_size"],
len(self.nvfp4_moe_meta["exclude_modules"]),
)
# Check model type
self.attention_chunk_size = getattr(
self.hf_text_config, "attention_chunk_size", None
)
self.sliding_window_size = self._get_sliding_window_size()
self.is_generation = not self.is_embedding_gemma and is_generation_model(
self.hf_config.architectures, is_embedding
)
# The vision_config/audio_config attribute heuristic is only applied when
# the transformers backend is explicitly requested. Some text-only models
# (e.g. xai-org/grok-2 with model_type="git") would otherwise be
# false-positively detected because their HF config auto-populates a
# `vision_config` in __post_init__.
has_multimodal_subconfig = self.hf_config is not self.hf_text_config or (
self.model_impl == ModelImpl.TRANSFORMERS
and (
hasattr(self.hf_config, "vision_config")
or hasattr(self.hf_config, "audio_config")
)
)
has_dsv41_vision = (
self.hf_config.model_type == "deepseek_v41"
and self.hf_config.vision_n_layers > 0
)
self.is_multimodal = (
enable_multimodal
and not self.is_lm_only
and (
is_multimodal_model(self.hf_config.architectures)
or has_multimodal_subconfig
or has_dsv41_vision
)
)
self.is_audio_model = enable_multimodal and is_audio_model(
self.hf_config.architectures
)
# Gated on `is_multimodal` because this flag is advertised via /model_info
# and drives the VLM warmup request, while the OpenAI serving layer rejects
# media input for models that are not `is_multimodal`. A text-only model
# with an auto-populated `vision_config` (see above) would otherwise warm up
# with an image request that its own serving layer answers with 400.
# Key on the tower, not the attribute: several config classes default
# vision_config to None, which presence alone would read as image-capable
# (MuseGlimmerConfig's text-only layouts are one such case).
# TODO: requires further polishing
self.is_image_understandable_model = (
self.is_multimodal
and getattr(self.hf_config, "vision_config", None) is not None
)
if self.is_multimodal and has_dsv41_vision:
self.is_image_understandable_model = True
# Models expose audio_config at different nesting levels:
# - top-level audio_config: e.g. Qwen2Audio
# - thinker_config.audio_config: Qwen3-Omni, Qwen3-ASR (nested thinker arch)
# - sound_config: Nemotron AVLM with Parakeet audio encoder
# - is_audio_model(): Whisper, Qwen3-ASR (architecture-based fallback)
# TODO: Handle this more robustly by standardizing the config structure in the future
self.is_audio_understandable_model = (
enable_multimodal
and not self.is_lm_only
and (
hasattr(self.hf_config, "audio_config")
or hasattr(
getattr(self.hf_config, "thinker_config", None), "audio_config"
)
or getattr(self.hf_config, "sound_config", None) is not None
or is_audio_model(self.hf_config.architectures)
)
)
self.is_multimodal_chunked_prefill_supported = (
enable_multimodal
and is_multimodal_chunked_prefill_supported(self.hf_config.architectures)
)
self.is_encoder_decoder = is_encoder_decoder_model(self.hf_config.architectures)
self.is_local_attention_model = is_local_attention_model(
self.hf_config.architectures
)
self.use_ngram_embedding = getattr(self.hf_config, "use_ngram_embedding", False)
self.ngram_embedding_n = (
self.hf_config.ngram_embedding_n if self.use_ngram_embedding else 0
)
self.use_engram = bool(getattr(self.hf_config, "engram_layer_ids", ()))
self._derive_multimodal_cuda_graph_support(enable_multimodal)
self.dtype = _get_and_verify_dtype(self.hf_text_config, dtype)
# Derive context length and model shapes
self._derive_context_length(context_length)
self._derive_model_shapes()
# Update hybrid model
self._derive_hybrid_model()
# Verify quantization
self._verify_quantization()
# Verify dual-chunk attention config
self._verify_dual_chunk_attention_config()
# Cache attributes
self.hf_eos_token_id = self._get_hf_eos_token_id()
# Set by scheduler when reasoning_parser is enabled
self.think_end_ids: Optional[List[int]] = None
self.self_opened_think_start_ids: Optional[List[int]] = None
self.request_selectable_think_end_id_sequences: Optional[List[List[int]]] = None
# multimodal
self.image_token_id = getattr(
self.hf_config, "image_token_id", None
) or getattr(self.hf_config, "image_token_index", None)
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 = self.is_lm_only
# matryoshka embeddings
self.matryoshka_dimensions = getattr(
self.hf_config, "matryoshka_dimensions", None
)
self.is_matryoshka = self.matryoshka_dimensions or getattr(
self.hf_config, "is_matryoshka", False
)
@staticmethod
def from_server_args(
server_args: ServerArgs,
model_path: str = None,
model_revision: str = None,
is_draft_model: bool = False,
context_length: Optional[int] = None,
**kwargs,
):
for record_type in type(server_args).__mro__:
factory = _MODEL_CONFIG_FACTORIES.get(record_type)
if factory is not None:
return factory(
server_args,
model_path=model_path,
model_revision=model_revision,
is_draft_model=is_draft_model,
context_length=context_length,
**kwargs,
)
cfg = resolving_view(server_args)
quantization = (
cfg.speculative_draft_model_quantization
if is_draft_model
else cfg.quantization
)
override_config_file = (
cfg.decrypted_draft_config_file
if is_draft_model
else cfg.decrypted_config_file
)
return ModelConfig(
model_path=model_path or cfg.model_path,
trust_remote_code=cfg.trust_remote_code,
revision=model_revision or cfg.revision,
context_length=(
context_length if context_length is not None else cfg.context_length
),
model_override_args=cfg.json_model_override_args,
is_embedding=cfg.is_embedding,
enable_multimodal=cfg.enable_multimodal,
dtype=cfg.dtype,
quantization=quantization,
model_impl=cfg.model_impl,
sampling_defaults=cfg.sampling_defaults,
quantize_and_serve=cfg.quantize_and_serve,
override_config_file=override_config_file,
is_multi_layer_eagle=cfg.enable_multi_layer_eagle,
language_only=cfg.language_only,
language_model_only=cfg.language_model_only,
encoder_only=cfg.encoder_only,
is_draft_model=is_draft_model,
is_draft_quantization_explicit=(
is_draft_model and cfg._speculative_draft_quantization_explicitly_set
),
disable_hybrid_swa_memory=cfg.disable_hybrid_swa_memory,
model_config_parser=cfg.model_config_parser,
speculative_algorithm=cfg.speculative_algorithm,
**kwargs,
)
def _derive_multimodal_cuda_graph_support(self, enable_multimodal: bool) -> None:
"""Declare graph capabilities before deriving shapes and validating config.
External ModelConfig subclasses can extend this without modifying global
architecture tables or repairing the config after construction.
"""
# A multimodal arch is piecewise-incompatible until its LM prefill is validated.
self.is_piecewise_cuda_graph_disabled_model = (
is_piecewise_cuda_graph_disabled_model(self.hf_config.architectures)
or (
self.is_multimodal
and not is_multimodal_piecewise_cuda_graph_supported(
self.hf_config.architectures
)
)
)
# The vision encoder still runs eagerly; this opt-in captures only the
# language-model prefill of architectures validated for piecewise graphs.
self.is_multimodal_piecewise_cuda_graph_supported = enable_multimodal and (
is_multimodal_piecewise_cuda_graph_supported(self.hf_config.architectures)
)
self.is_multimodal_breakable_cuda_graph_supported = enable_multimodal and (
is_multimodal_breakable_cuda_graph_supported(self.hf_config.architectures)
)
def _config_draft_model(self):
is_draft_model = self.is_draft_model
from sglang.srt.configs.dots3 import Dots3Config
if is_draft_model and isinstance(self.hf_text_config, Dots3Config):
self.hf_config.architectures[0] = (
self.hf_text_config.configure_draft_model()
)
if (
is_draft_model
and self.speculative_algorithm == "DSPARK"
and self.hf_config.architectures[0] == "DSparkMiniMaxDraftModel"
):
from sglang.srt.configs.minimax_vl import configure_minimax_m3_dspark
self.hf_config = configure_minimax_m3_dspark(self.hf_config)
self.hf_text_config = self.hf_config
if self.is_draft_quantization_explicit and self.quantization is None:
self.hf_config.quantization_config = None
if is_draft_model and self.hf_config.architectures[0] in [
"DeepseekV3ForCausalLM",
"DeepseekV32ForCausalLM",
]:
self.hf_config.architectures[0] = "DeepseekV3ForCausalLMNextN"
if is_draft_model and self.hf_config.architectures[0] == "GlmMoeDsaForCausalLM":
self.hf_config.architectures[0] = "GlmMoeDsaForCausalLMNextN"
if (
is_draft_model
and self.hf_config.architectures[0] == "DeepseekV4ForCausalLM"
):
from sglang.srt.speculative.dspark_components.dspark_config import (
checkpoint_bundles_dspark_draft,
)
# A dspark-bundled checkpoint may also carry MTP layers; the
# selected algorithm decides which draft arch to load.
if checkpoint_bundles_dspark_draft(self.hf_config) and (
self.speculative_algorithm in (None, "DSPARK")
):
self.hf_config.architectures[0] = "DeepseekV4ForCausalLMDSpark"
logger.info(
"Draft checkpoint bundles a DSpark head; loading draft arch "
"DeepseekV4ForCausalLMDSpark."
)
else:
self.hf_config.architectures[0] = "DeepseekV4ForCausalLMNextN"
self.hf_config.num_nextn_predict_layers = 1
if is_draft_model and self.hf_config.architectures[0] == "Glm4MoeForCausalLM":
self.hf_config.architectures[0] = "Glm4MoeForCausalLMNextN"
if (
is_draft_model
and self.hf_config.architectures[0] == "Glm4MoeLiteForCausalLM"
):
self.hf_config.architectures[0] = "Glm4MoeLiteForCausalLMNextN"
if (
is_draft_model
and self.hf_config.architectures[0] == "Glm5NextForConditionalGeneration"
):
self.hf_config.architectures[0] = "Glm5NextForConditionalGenerationNextN"
self.hf_text_config.architectures = list(self.hf_config.architectures)
self.hf_text_config.num_nextn_predict_layers = 1
self.hf_text_config.linear_attn_config = None
if is_draft_model and self.hf_config.architectures[0] in [
"GlmOcrForConditionalGeneration",
]:
self.hf_config.architectures[0] = "GlmOcrForConditionalGenerationNextN"
if (
is_draft_model
and self.hf_config.architectures[0] == "LongcatFlashForCausalLM"
):
self.hf_config.architectures[0] = "LongcatFlashForCausalLMNextN"
self.hf_config.num_hidden_layers = self.hf_config.num_nextn_predict_layers
if is_draft_model and self.hf_config.architectures[0] == "MiMoForCausalLM":
self.hf_config.architectures[0] = "MiMoMTP"
if is_draft_model and self.hf_config.architectures[0] in MIMO_V2_MODEL_ARCHS:
self.hf_config.architectures[0] = "MiMoV2MTP"
self.hf_config.num_nextn_predict_layers = 1
if is_draft_model and self.hf_config.architectures[0] == "Step3p5ForCausalLM":
self.hf_config.architectures[0] = "Step3p5MTP"
if (
is_draft_model
and self.hf_config.architectures[0] == "InklingForConditionalGeneration"
):
self.hf_config.architectures[0] = "InklingForConditionalGenerationMTP"
if self.hf_text_config.mtp_local_layer_ids:
# A local MTP depth keeps the trunk's window unless it declares
# its own: the draft backend builds its window rail only if set.
self.sliding_window_size = (
self.hf_text_config.mtp_local_extent or self.sliding_window_size
)
if (
is_draft_model
and self.hf_config.architectures[0] == "GigaChat35ForCausalLM"
):
self.hf_config.architectures[0] = "GigaChat35ForCausalLMNextN"
if (
is_draft_model
and self.hf_config.architectures[0] == "Step3p7ForConditionalGeneration"
):
self.hf_config = self.hf_text_config
self.hf_config.architectures = ["Step3p5MTP"]
if is_draft_model and self.hf_config.architectures[0] in [
"BailingMoeV2ForCausalLM",
"BailingMoeForCausalLM",
"BailingMoeV2_5ForCausalLM",
"BailingMoeV3ForCausalLM",
]:
self.hf_config.architectures[0] = "BailingMoeForCausalLMNextN"
if (
is_draft_model
and self.hf_config.architectures[0] == "Ernie4_5_MoeForCausalLM"
):
self.hf_config.architectures[0] = "Ernie4_5_MoeForCausalLMMTP"
if is_draft_model and self.hf_config.architectures[0] == "Qwen3NextForCausalLM":
self.hf_config.architectures[0] = "Qwen3NextForCausalLMMTP"
self.hf_config.num_nextn_predict_layers = 1
if (
is_draft_model
and self.hf_config.architectures[0] == "Qwen4ExpForConditionalGeneration"
):
# The target's ModelConfig shares this hf_config object; deep-copy
# before the MTP rewrites below so the target keeps its full depth.
self.hf_config = copy.deepcopy(self.hf_config)
self.hf_text_config = get_hf_text_config(self.hf_config)
self.hf_config.architectures[0] = "Qwen4ExpForCausalLMMTP"
text_config = self.hf_text_config
text_config.num_nextn_predict_layers = 1
# layers_block_type follows layer_types, not num_hidden_layers,
# so both must shrink for the draft's full_attention_layer_ids to be [0].
text_config.num_hidden_layers = 1
text_config.layer_types = ["full_attention"]
text_config.full_attention_interval = 1
if is_draft_model and self.hf_config.architectures[0] == "Qwen3MoeForCausalLM":
self.hf_config.architectures[0] = "Qwen3MoeForCausalLMMTP"
self.hf_config.num_nextn_predict_layers = 1
if (
is_draft_model
and self.hf_config.architectures[0]
== "SenseNovaLoopForConditionalGeneration"
):
# MTP patch: the bundled mtp.* head is a plain Qwen3.5 MTP layer (one
# full-attention layer, BF16). Shrink the looped execution layout to
# that single layer; deep-copy so the target keeps its 56-slot plan.
self.hf_config = copy.deepcopy(self.hf_config)
self.hf_text_config = get_hf_text_config(self.hf_config)
text_config = self.hf_text_config
text_config.num_hidden_layers = 1
text_config.layer_types = ["full_attention"]
text_config.full_attention_interval = 1
text_config.execution_plan = (0,)
text_config.mhc = None
text_config.num_nextn_predict_layers = 1
self.hf_config.num_nextn_predict_layers = 1
self.hf_config.architectures[0] = "Qwen3_5ForCausalLMMTP"
if is_draft_model and self.hf_config.architectures[0] in [
"Qwen3_5ForConditionalGeneration",
"Qwen3_5MoeForConditionalGeneration",
"Qwen3_5ForCausalLM",
"Qwen3_5MoeForCausalLM",
"InternS2PreviewForConditionalGeneration",
"InternS2MobiusForConditionalGeneration",
]:
if (
self.hf_config.architectures[0]
== "InternS2MobiusForConditionalGeneration"
):
# The target owns 2,560 experts through four shared physical
# banks, while its bundled MTP layer is an ordinary Qwen3.5
# MoE layer with the checkpoint-declared smaller expert set.
self.hf_text_config.model_type = "qwen3_5_moe_text"
self.hf_text_config.num_experts = self.hf_text_config.mtp_num_experts
self.hf_text_config.num_experts_per_tok = (
self.hf_text_config.mtp_num_experts_per_tok
)
self.hf_config.architectures[0] = "Qwen3_5ForCausalLMMTP"
self.hf_config.num_nextn_predict_layers = 1
self.hf_text_config.num_nextn_predict_layers = 1
if is_draft_model and self.hf_config.architectures[0] in (
"ExaoneMoEForCausalLM",
"ExaoneMoeForCausalLM",
):
self.hf_config.architectures[0] = "ExaoneMoEForCausalLMMTP"
self.hf_config.num_nextn_predict_layers = 1
if (
is_draft_model
and self.hf_config.architectures[0] == "NemotronH_Omni_Reasoning_V3"
):
self.hf_config = self.hf_text_config
self.hf_config.architectures = ["NemotronHForCausalLMMTP"]
self.hf_config.num_nextn_predict_layers = 1
if is_draft_model and self.hf_config.architectures[0] in [
"NemotronHForCausalLM",
"NemotronHPuzzleForCausalLM",
]:
self.hf_config.architectures[0] = "NemotronHForCausalLMMTP"
self.hf_config.num_nextn_predict_layers = 1
if is_draft_model and self.hf_config.architectures[0] == "HYV3ForCausalLM":
self.hf_config.architectures[0] = "HYV3ForCausalLMNextN"
self.hf_config.num_nextn_predict_layers = 1
if is_draft_model and self.hf_config.architectures[0] == "HYV4ForCausalLM":
self.hf_config.architectures[0] = "HYV4ForCausalLMNextN"
self.hf_config.num_nextn_predict_layers = 1
def _derive_hybrid_model(self):
# Use self.context_len after it has been initialized to prevent using context_len which may be None.
self.is_hybrid_swa = (
is_hybrid_swa_model(self.hf_config.architectures, self.hf_text_config)
and not self.disable_hybrid_swa_memory
)
# Whole-model split, read-only; per-runner slices live on ModelLayerInfo.
self.swa_attention_layer_ids: Optional[List[int]] = None
self.full_attention_layer_ids: Optional[List[int]] = None
if self.is_hybrid_swa:
logger.debug(f"Hybrid swa model: {self.hf_config.architectures=}")
self.is_deepseek_v4_arch = any(
arch
in [
"DeepseekV4ForCausalLM",
"DeepseekV4ForCausalLMNextN",
"DeepseekV4ForCausalLMDSpark",
]
for arch in self.hf_config.architectures
)
if not self.is_deepseek_v4_arch:
self.swa_attention_layer_ids, self.full_attention_layer_ids = (
get_hybrid_layer_ids(
self.hf_config.architectures,
self.hf_text_config,
)
)
self.has_attention_sinks = self._detect_attention_sinks()
self.is_hybrid_swa_compress = self.hf_config.architectures[0] in [
*MIMO_V2_MODEL_ARCHS,
"MiMoV2MTP",
"Gemma4ForCausalLM",
"Gemma4ForConditionalGeneration",
"EmbeddingGemma2Model",
"InklingForConditionalGeneration",
"InklingForConditionalGenerationMTP",
"Gemma4UnifiedForConditionalGeneration",
"DiffusionGemmaForBlockDiffusion",
]
@cached_property
def linear_attn_registry_result(self) -> Any:
return get_linear_attn_config(self.hf_config)
@property
def has_asymmetric_kv(self) -> bool:
"""Whether K and V rows differ in width (MiMoV2 is 192 / 128).
Not an ``__init__`` field because the MLA special-casing below still
rewrites ``v_head_dim``.
"""
return (
self.head_dim != self.v_head_dim or self.swa_head_dim != self.swa_v_head_dim
)
def _detect_attention_sinks(self) -> bool:
"""Check whether the model uses learned attention sinks.
Attention sinks are per-head scalars added to the softmax denominator
to compensate for evicted KV-cache entries under sliding-window
attention. Not every hybrid-SWA model uses them.
"""
archs = self.hf_config.architectures or []
if any(a in SWA_SINK_ARCHS for a in archs):
return True
# MiMoV2 creates sinks only when the config flags are set.
if any(a in archs for a in (*MIMO_V2_MODEL_ARCHS, "MiMoV2MTP")):
return getattr(
self.hf_text_config, "add_swa_attention_sink_bias", False
) or getattr(self.hf_text_config, "add_full_attention_sink_bias", False)
return bool(getattr(self.hf_text_config, "learnable_sink", False))
def _derive_context_length(self, context_length: int):
is_draft_model = self.is_draft_model
derived_context_len = get_context_length(self.hf_text_config)
if context_length is not None:
if context_length > derived_context_len:
reason = "Target model's" if is_draft_model else "User-specified"
msg = (
f"Warning: {reason} context_length ({context_length}) is greater than the derived context_length ({derived_context_len}). "
f"This may lead to incorrect model outputs or CUDA errors. Note that the derived context_length may differ from max_position_embeddings in the model's config."
)
if (
envs.SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN.get()
or is_in_ci() # FIXME: fix this special case
):
logger.warning(msg)
self.context_len = context_length
if is_draft_model:
self.hf_text_config.max_position_embeddings = context_length
logger.warning(
f"Overriding the draft model's max_position_embeddings to {context_length}."
)
else:
raise ValueError(
f"{msg} To allow overriding this maximum, set the env var SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN=1"
)
else:
self.context_len = context_length
else:
if (
self.hf_config.architectures
and self.hf_config.architectures[0] == "EmbeddingGemma2Model"
):
self.context_len = min(derived_context_len, 8192)
else:
self.context_len = derived_context_len
# Transfer context_len to HuggingFace config so models can access it
self.hf_config.context_len = self.context_len
def _derive_model_shapes(self):
from sglang.srt.configs.dots3 import Dots3Config
is_dflash_mla_draft = False
if getattr(self.hf_config, "dflash_config", None):
from sglang.srt.speculative.dflash_utils import (
parse_dflash_draft_config,
validate_dflash_mla_draft_config,
)
is_dflash_mla_draft = parse_dflash_draft_config(
draft_hf_config=self.hf_config
).is_mla
if is_dflash_mla_draft:
# Validate before the generic defaults below fill missing fields.
validate_dflash_mla_draft_config(self.hf_config)
# Pure SSM models have no attention heads; use head_dim == 0 so the
# KV-cell size is 0 rather than a division on a missing head count.
is_pure_ssm = any(
arch in self.hf_config.architectures for arch in PURE_SSM_ARCHITECTURES
)
# Unify the config keys for hf_text_config
self.head_dim = getattr(self.hf_text_config, "head_dim", None)
if self.head_dim is None:
if is_pure_ssm:
self.head_dim = 0
else:
self.head_dim = (
self.hf_text_config.hidden_size
// self.hf_text_config.num_attention_heads
)
setattr(self.hf_text_config, "head_dim", self.head_dim)
self.v_head_dim = getattr(self.hf_text_config, "v_head_dim", None)
if self.v_head_dim is None or self.v_head_dim == 0:
self.v_head_dim = self.head_dim
setattr(self.hf_text_config, "v_head_dim", self.v_head_dim)
self.swa_head_dim = getattr(self.hf_text_config, "swa_head_dim", None)
if self.swa_head_dim is None:
self.swa_head_dim = self.head_dim
setattr(self.hf_text_config, "swa_head_dim", self.swa_head_dim)
self.swa_v_head_dim = getattr(self.hf_text_config, "swa_v_head_dim", None)
if self.swa_v_head_dim is None:
self.swa_v_head_dim = self.swa_head_dim
setattr(self.hf_text_config, "swa_v_head_dim", self.swa_v_head_dim)
# FIXME: temporary special judge for MLA architecture
if (
is_dflash_mla_draft
or "DeepseekV2ForCausalLM" in self.hf_config.architectures
or "DeepseekV32ForCausalLM" in self.hf_config.architectures
or "DeepseekV3ForCausalLM" in self.hf_config.architectures
or "DeepseekV3ForCausalLMNextN" in self.hf_config.architectures
or "Glm4MoeLiteForCausalLM" in self.hf_config.architectures
or "Glm4MoeLiteForCausalLMNextN" in self.hf_config.architectures
or "GlmMoeDsaForCausalLM" in self.hf_config.architectures
or "GlmMoeDsaForCausalLMNextN" in self.hf_config.architectures
or "Glm5NextForConditionalGeneration" in self.hf_config.architectures
or "Glm5NextForConditionalGenerationNextN" in self.hf_config.architectures
or "LongcatFlashForCausalLM" in self.hf_config.architectures
or "LongcatFlashForCausalLMNextN" in self.hf_config.architectures
or "HYV4ForCausalLM" in self.hf_config.architectures
or "HYV4ForCausalLMNextN" in self.hf_config.architectures
or "DotsVLMForCausalLM" in self.hf_config.architectures
or "Dots3NoteForCausalLM" in self.hf_config.architectures
or "Dots3NoteForCausalLMNextN" in self.hf_config.architectures
or "MistralLarge3ForCausalLM" in self.hf_config.architectures
or (
"PixtralForConditionalGeneration" in self.hf_config.architectures
and getattr(self.hf_text_config, "kv_lora_rank", None) is not None
)
or "MistralLarge3ForCausalLMEagle" in self.hf_config.architectures
or "KimiK25ForConditionalGeneration" in self.hf_config.architectures
or "Eagle3DeepseekV2ForCausalLM" in self.hf_config.architectures
or "GigaChat35ForCausalLM" in self.hf_config.architectures
or "GigaChat35ForCausalLMNextN" in self.hf_config.architectures
):
self.head_dim = 256
self.attention_arch = AttentionArch.MLA
self.kv_lora_rank = self.hf_text_config.kv_lora_rank
self.qk_nope_head_dim = self.hf_text_config.qk_nope_head_dim
self.qk_rope_head_dim = self.hf_text_config.qk_rope_head_dim
self.v_head_dim = self.hf_text_config.v_head_dim
if isinstance(self.hf_text_config, Dots3Config):
self.swa_kv_lora_rank = self.hf_text_config.swa_kv_lora_rank
self.swa_qk_rope_head_dim = self.hf_text_config.swa_qk_rope_head_dim
else:
self.swa_kv_lora_rank = self.kv_lora_rank
self.swa_qk_rope_head_dim = self.qk_rope_head_dim
if is_dflash_mla_draft:
# Sliding draft layers share the latent rows, not a GQA head_dim.
self.swa_v_head_dim = self.v_head_dim
self.index_head_dim = (
get_dsa_index_head_dim(self.hf_text_config)
if is_deepseek_dsa(self.hf_text_config)
else None
)
# In transformers v5, rope_scaling is just rope_parameters.
rope_scaling = getattr(
self.hf_text_config, "rope_parameters", None
) or getattr(self.hf_text_config, "rope_scaling", None)
self._init_mla_scaling(rope_scaling)
elif (
"DeepseekV4ForCausalLM" in self.hf_config.architectures
or "DeepseekV4ForCausalLMNextN" in self.hf_config.architectures
or "DeepseekV4ForCausalLMDSpark" in self.hf_config.architectures
):
self.qk_rope_head_dim = self.hf_config.qk_rope_head_dim
self.qk_nope_head_dim = self.hf_config.head_dim - self.qk_rope_head_dim
self.window_size = self.hf_config.sliding_window
self.head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim
self.v_head_dim = self.head_dim
self.index_head_dim = self.hf_config.index_head_dim
self.compress_ratios = self.hf_config.compress_ratios
self.attention_arch = AttentionArch.MHA
self._init_mla_scaling(self.hf_config.rope_scaling)
elif "Glm4MoeForCausalLMNextN" in self.hf_config.architectures:
if self.head_dim is None:
self.head_dim = (
self.hf_text_config.hidden_size
// self.hf_text_config.num_attention_heads
)
if self.swa_head_dim is None:
self.swa_head_dim = self.head_dim
self.v_head_dim = self.head_dim
self.swa_v_head_dim = self.swa_head_dim
self.attention_arch = AttentionArch.MHA
elif "MiniCPM3ForCausalLM" in self.hf_config.architectures:
self.head_dim = 128
self.attention_arch = AttentionArch.MLA
self.kv_lora_rank = self.hf_config.kv_lora_rank
self.qk_rope_head_dim = self.hf_config.qk_rope_head_dim
elif "DeepseekVL2ForCausalLM" in self.hf_config.architectures and getattr(
self.hf_text_config, "use_mla", True
):
self.head_dim = 256
self.attention_arch = AttentionArch.MLA
self.kv_lora_rank = self.hf_text_config.kv_lora_rank
self.qk_rope_head_dim = self.hf_text_config.qk_rope_head_dim
self.qk_nope_head_dim = self.hf_text_config.qk_nope_head_dim
elif "KimiVLForConditionalGeneration" in self.hf_config.architectures:
self.head_dim = 256
self.attention_arch = AttentionArch.MLA
self.kv_lora_rank = self.hf_text_config.kv_lora_rank
self.qk_rope_head_dim = self.hf_text_config.qk_rope_head_dim
self.v_head_dim = self.hf_text_config.v_head_dim
self.qk_nope_head_dim = self.hf_text_config.qk_nope_head_dim
elif (
"KimiLinearForCausalLM" in self.hf_config.architectures
or "KimiK3LinearForCausalLM" in self.hf_config.architectures
or "KimiK3ForConditionalGeneration" in self.hf_config.architectures
):
tc = self.hf_text_config
self.head_dim = 72
self.attention_arch = AttentionArch.MLA
self.kv_lora_rank = tc.kv_lora_rank
self.qk_rope_head_dim = tc.qk_rope_head_dim
self.v_head_dim = tc.v_head_dim
self.qk_nope_head_dim = tc.qk_nope_head_dim
self._init_mla_scaling(getattr(tc, "rope_scaling", None))
elif "BailingMoeV2_5ForCausalLM" in self.hf_config.architectures or (
# Every Bailing draft is renamed to BailingMoeForCausalLMNextN; only
# the MLA ones (V2.5 and V3) carry kv_lora_rank, V2 keeps its GQA
# shapes.
"BailingMoeForCausalLMNextN" in self.hf_config.architectures
and getattr(self.hf_text_config, "kv_lora_rank", None) is not None
):
self.head_dim = self.hf_text_config.head_dim
self.attention_arch = AttentionArch.MLA
self.kv_lora_rank = self.hf_text_config.kv_lora_rank
self.qk_nope_head_dim = self.hf_text_config.qk_nope_head_dim
self.qk_rope_head_dim = self.hf_text_config.qk_rope_head_dim
self.v_head_dim = self.hf_config.v_head_dim
self._init_mla_scaling(self.hf_config.rope_scaling)
elif (
"BailingMoeV3ForCausalLM" in self.hf_config.architectures
or "BailingMoeV3VLForConditionalGeneration" in self.hf_config.architectures
):
self.head_dim = 128
self.attention_arch = AttentionArch.MLA
self.kv_lora_rank = self.hf_text_config.kv_lora_rank
self.qk_rope_head_dim = (
0
if getattr(self.hf_text_config, "use_mla_nope", False)
else self.hf_text_config.qk_rope_head_dim
)
self.v_head_dim = self.hf_text_config.v_head_dim
self.qk_nope_head_dim = self.hf_text_config.qk_nope_head_dim
self.scaling = 1 / math.sqrt(self.qk_nope_head_dim + self.qk_rope_head_dim)
elif (
"SarvamMLAForCausalLM" in self.hf_config.architectures
or "Glm5NextForConditionalGeneration" in self.hf_config.architectures
):
self.head_dim = (
self.hf_config.qk_nope_head_dim + self.hf_config.qk_rope_head_dim
)
self.attention_arch = AttentionArch.MLA
self.kv_lora_rank = self.hf_config.kv_lora_rank
self.qk_rope_head_dim = self.hf_config.qk_rope_head_dim
self.qk_nope_head_dim = self.hf_config.qk_nope_head_dim
self.v_head_dim = self.hf_config.v_head_dim
self._init_mla_scaling(self.hf_config.rope_scaling)
else:
if (
"MistralModel" in self.hf_config.architectures
or "MixtralForCausalLM" in self.hf_config.architectures
or "MistralForCausalLM" in self.hf_config.architectures
):
if getattr(self, "head_dim", None) is None:
self.head_dim = (
self.hf_config.hidden_size // self.hf_config.num_attention_heads
)
# In transformers==4.52.3, the head_dim is null in MistralConfig
if (
not hasattr(self.hf_text_config, "head_dim")
or self.hf_text_config.head_dim is None
):
setattr(self.hf_text_config, "head_dim", self.head_dim)
elif "BaichuanForCausalLM" in self.hf_config.architectures:
self.use_alibi = self.hf_config.hidden_size != 4096
# Pure Mamba SSMs have no attention (head_dim set to 0 above).
if is_pure_ssm:
self.attention_arch = AttentionArch.SSM
else:
self.attention_arch = AttentionArch.MHA
# Mamba2 has no num_attention_heads.
self.num_attention_heads = getattr(
self.hf_text_config, "num_attention_heads", None
)
self.num_key_value_heads = getattr(
self.hf_text_config, "num_key_value_heads", None
)
self.first_k_dense_replace = getattr(
self.hf_text_config, "first_k_dense_replace", None
)
self.full_attention_interval = getattr(
self.hf_text_config, "full_attention_interval", None
)
# for Dbrx and MPT models
if self.hf_config.model_type in ["dbrx", "mpt"]:
self.num_key_value_heads = getattr(
self.hf_config.attn_config, "kv_n_heads", None
)
if self.num_key_value_heads is None:
self.num_key_value_heads = self.num_attention_heads
self.hidden_size = self.hf_text_config.hidden_size
hc_mult = getattr(self.hf_text_config, "hc_mult", 1)
is_glm5_next = getattr(self.hf_config, "model_type", None) == "glm5_next" or (
getattr(self.hf_text_config, "model_type", None) == "glm5_next_text"
)
if is_glm5_next and not getattr(self.hf_text_config, "mhc", False):
hc_mult = 1
if is_glm5_next:
# mHC-flattened hidden size; None when not running an mHC model.
self.hc_hidden_size = self.hidden_size * hc_mult if hc_mult > 1 else None
self.spec_hidden_size = self.hidden_size
else:
self.spec_hidden_size, self.hc_hidden_size = resolve_spec_hidden_size(
self.hf_config, self.hidden_size, hc_mult
)
self.num_hidden_layers = self.hf_text_config.num_hidden_layers
self.num_attention_layers = self.num_hidden_layers
if "LongcatFlashForCausalLM" in self.hf_config.architectures:
self.num_attention_layers = self.num_hidden_layers * 2
if "IQuestLoopCoderForCausalLM" in self.hf_config.architectures:
loop_num = getattr(self.hf_text_config, "loop_num", 1)
self.num_attention_layers = int(self.num_hidden_layers * int(loop_num))
if "NanbeigeForCausalLM" in self.hf_config.architectures:
num_loops = getattr(self.hf_text_config, "num_loops", 1)
self.num_attention_layers = int(self.num_hidden_layers * int(num_loops))
if "WhisperForConditionalGeneration" in self.hf_config.architectures:
# Whisper has unique layer ID scheme:
# - Encoder self-attention: 0 to encoder_layers-1 (no KV cache)
# - Decoder self-attention: encoder_layers to encoder_layers+decoder_layers-1 (uses KV cache)
# - Decoder cross-attention: encoder_layers+decoder_layers to encoder_layers+2*decoder_layers-1
# Even though cross-attention doesn't save KV cache, attention backend needs buffer to exist
encoder_layers = getattr(self.hf_text_config, "encoder_layers", 0)
decoder_layers = getattr(
self.hf_text_config, "decoder_layers", self.num_hidden_layers
)
self.num_attention_layers = encoder_layers + 2 * decoder_layers
self.num_nextn_predict_layers = getattr(
self.hf_text_config, "num_nextn_predict_layers", None
)
self.vocab_size = self.hf_text_config.vocab_size
# GLM-Image is the only model here whose output head predicts vision tokens.
# Use vision_vocab_size for lm_head, LogitsProcessor, and graph-mode logits buffers.
if _hf_arch(self.hf_config) == "GlmImageForConditionalGeneration":
self.vocab_size = self.hf_text_config.vision_vocab_size
def _init_mla_scaling(self, rope_scaling: Optional[dict]) -> None:
"""Base MLA attention scale from the head dims, then the rope mscale."""
self.scaling = 1 / math.sqrt(self.qk_nope_head_dim + self.qk_rope_head_dim)
if rope_scaling:
self.scaling = compute_mla_mscale_scaling(rope_scaling, self.scaling)
def get_total_num_attention_heads(self) -> int:
return self.num_attention_heads
def get_num_attention_heads(self, tensor_parallel_size) -> int:
# Pure-SSM (Mamba) models have no attention; num_attention_heads is None.
total_num_attention_heads = self.num_attention_heads or 0
return max(1, total_num_attention_heads // tensor_parallel_size)
# adapted from https://github.com/vllm-project/vllm/blob/main/vllm/config.py#L289
def get_total_num_kv_heads(self) -> int:
"""Returns the total number of KV heads."""
# For GPTBigCode & Falcon:
# NOTE: for falcon, when new_decoder_architecture is True, the
# multi_query flag is ignored and we use n_head_kv for the number of
# KV heads.
falcon_model_types = ["falcon", "RefinedWeb", "RefinedWebModel"]
new_decoder_arch_falcon = (
self.hf_config.model_type in falcon_model_types
and getattr(self.hf_config, "new_decoder_architecture", False)
)
if not new_decoder_arch_falcon and getattr(
self.hf_text_config, "multi_query", False
):
# Multi-query attention, only one KV head.
# Currently, tensor parallelism is not supported in this case.
return 1
# For DBRX and MPT
if self.hf_config.model_type in ["mpt"]:
if "kv_n_heads" in self.hf_config.attn_config:
return self.hf_config.attn_config["kv_n_heads"]
return self.hf_config.num_attention_heads
if self.hf_config.model_type in ["dbrx"]:
return getattr(
self.hf_config.attn_config,
"kv_n_heads",
self.hf_config.num_attention_heads,
)
if self.hf_config.model_type in ["nemotron-nas"]:
nkvh = {
self.hf_config.num_attention_heads // block.attention.n_heads_in_group
for block in self.hf_config.block_configs
if not block.attention.no_op
}
if len(nkvh) == 0:
raise RuntimeError("Couldn't determine number of kv heads")
if len(nkvh) > 1:
raise ValueError(
"Variable GQA (VGQA) is not yet supported for nemotron-nas in sglang"
)
return next(iter(nkvh))
attributes = [
# For Falcon:
"n_head_kv",
"num_kv_heads",
# For LLaMA-2:
"num_key_value_heads",
# For ChatGLM:
"multi_query_group_num",
# For Step3
"num_attention_groups",
]
for attr in attributes:
num_kv_heads = getattr(self.hf_text_config, attr, None)
if num_kv_heads is not None:
return num_kv_heads
# Mamba SSMs have no attention, so no KV heads.
if self.attention_arch == AttentionArch.SSM:
return 0
# For non-grouped-query attention models, the number of KV heads is
# equal to the number of attention heads.
return self.hf_text_config.num_attention_heads
def get_max_num_attention_heads(self) -> int:
"""Max per-layer query head count; num_attention_heads unless the
model sets num_attention_heads_per_layer."""
per_layer = getattr(self.hf_text_config, "num_attention_heads_per_layer", None)
if per_layer:
return max(per_layer)
return self.num_attention_heads
def get_num_kv_heads(self, tensor_parallel_size: int, dcp_size: int = 1) -> int:
"""Number of KV heads per GPU.
DCP ranks replicate KV, so heads shard across ``tp // dcp`` groups.
Drafts never join the group and ignore ``dcp_size``. With fewer heads
than groups, each GPU keeps one.
"""
total_num_kv_heads = self.get_total_num_kv_heads()
if self.is_draft_model:
dcp_size = 1
kv_tensor_parallel_size = tensor_parallel_size // dcp_size
return max(1, total_num_kv_heads // kv_tensor_parallel_size)
def get_swa_num_kv_heads(self, tensor_parallel_size) -> int:
"""Similar to get_num_kv_heads(), but for SWA."""
if hasattr(self.hf_text_config, "swa_num_key_value_heads"):
total_num_kv_heads = self.hf_text_config.swa_num_key_value_heads
return max(1, total_num_kv_heads // tensor_parallel_size)
elif hasattr(self.hf_text_config, "attention_other_setting"): # For step3p5
total_num_kv_heads = self.hf_text_config.attention_other_setting.get(
"num_attention_groups"
)
return max(1, total_num_kv_heads // tensor_parallel_size)
else:
return self.get_num_kv_heads(tensor_parallel_size)
# adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/config.py
def _parse_quant_hf_config(self):
quant_cfg = _quant_config_to_dict(
getattr(self.hf_config, "quantization_config", None)
)
if quant_cfg is not None:
# Identify modelopt quantization
if (
"quant_method" not in quant_cfg
or quant_cfg["quant_method"] == "modelopt"
):
parsed_cfg = self._parse_modelopt_quant_config(
{"quantization": quant_cfg}
)
if parsed_cfg:
quant_cfg.update(parsed_cfg)
if quant_cfg is None:
# compressed-tensors uses a "compression_config" key
quant_cfg = getattr(self.hf_config, "compression_config", None)
if quant_cfg is None:
# check if is modelopt or mixed-precision model -- Both of them don't have corresponding field
# in hf `config.json` but has a standalone `hf_quant_config.json` in the root directory
# example: https://huggingface.co/nvidia/Llama-3.1-8B-Instruct-FP8/tree/main
# example: https://huggingface.co/Barrrrry/DeepSeek-R1-W4AFP8/tree/main
is_local = os.path.exists(self.model_path)
if not is_local:
# Conditional import based on SGLANG_USE_MODELSCOPE environment variable
if envs.SGLANG_USE_MODELSCOPE.get():
from modelscope import HubApi, model_file_download
hf_api = HubApi()
else:
import huggingface_hub
from huggingface_hub import HfApi, hf_hub_download
hf_api = HfApi()
try:
# In offline mode, skip file_exists check to avoid OfflineModeIsEnabled error
# Instead, directly try to download/read from cache with local_files_only
file_exists = False # Initialize to avoid UnboundLocalError
if not huggingface_hub.constants.HF_HUB_OFFLINE:
# Online mode: check if file exists before attempting download (optimization)
file_exists = retry(
lambda: hf_api.file_exists(
self.model_path, "hf_quant_config.json"
),
max_retry=2,
initial_delay=1.0,
max_delay=5.0,
)
if not file_exists:
# File doesn't exist on hub, no need to try downloading
return quant_cfg # None
# Download (online mode) or read from cache (offline mode)
if envs.SGLANG_USE_MODELSCOPE.get():
quant_config_file = model_file_download(
model_id=self.model_path,
file_path="hf_quant_config.json",
revision=self.revision,
)
else:
quant_config_file = hf_hub_download(
repo_id=self.model_path,
filename="hf_quant_config.json",
revision=self.revision,
local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE,
)
with open(quant_config_file) as f:
quant_config_dict = json.load(f)
quant_cfg = self._parse_modelopt_quant_config(quant_config_dict)
except huggingface_hub.errors.LocalEntryNotFoundError:
# Offline mode and file not in cache - this is normal for non-quantized models
logger.debug(
f"hf_quant_config.json not found in cache for {self.model_path} "
"(offline mode, normal for non-quantized models)"
)
except huggingface_hub.errors.OfflineModeIsEnabled:
# Should not reach here after our changes, but keep for safety
logger.warning(
"Offline mode is enabled, skipping hf_quant_config.json check"
)
except Exception as e:
logger.warning(
"Failed to load hf_quant_config.json for model %s: %s",
self.model_path,
e,
)
elif os.path.exists(os.path.join(self.model_path, "hf_quant_config.json")):
quant_config_file = os.path.join(
self.model_path, "hf_quant_config.json"
)
with open(quant_config_file) as f:
quant_config_dict = json.load(f)
quant_cfg = self._parse_modelopt_quant_config(quant_config_dict)
return quant_cfg
def _find_quant_modelslim_config(self):
if self.is_draft_model:
return None
quant_config_file = Path(self.model_path, "quant_model_description.json")
quant_cfg = None
if quant_config_file.is_file():
with open(quant_config_file) as f:
quant_cfg = json.load(f)
# This field is required for flagless model loading but is not present in
# modelslim model description, so we're adding it here manually.
quant_cfg["quant_method"] = "modelslim"
return quant_cfg
def _parse_modelopt_quant_config(self, quant_config_dict: dict) -> Optional[dict]:
"""Parse ModelOpt quantization config and return the appropriate quant_method.
Supports both nested LLM ``hf_quant_config.json``::
{"quantization": {"quant_algo": "FP8", ...}}
and flat formats (``config.json`` ``quantization_config``, or diffusion /
unified ModelOpt exports such as Cosmos3)::
{"quant_algo": "FP8", "quant_method": "modelopt", ...}
"""
json_quant_configs = unwrap_modelopt_quantization_config(quant_config_dict)
if not json_quant_configs:
return None
quant_algo = json_quant_configs.get("quant_algo", None)
if quant_algo == "MIXED_PRECISION":
quantized_layers = json_quant_configs.get("quantized_layers") or {}
has_modelopt_nvfp4_layers = any(
str(layer_info.get("quant_algo", "")).upper()
in ("NVFP4", "W4A16_NVFP4")
for layer_info in quantized_layers.values()
if isinstance(layer_info, dict)
)
if has_modelopt_nvfp4_layers:
return {"quant_method": "modelopt_mixed", "quant_algo": quant_algo}
return {"quant_method": "w4afp8", "quant_algo": quant_algo}
elif quant_algo and ("FP4" in quant_algo or "NVFP4" in quant_algo):
return {"quant_method": "modelopt_fp4", "quant_algo": quant_algo}
elif quant_algo == "FP8":
return {"quant_method": "modelopt_fp8", "quant_algo": quant_algo}
elif quant_algo == "MXFP8":
group_size = json_quant_configs.get("group_size", 32)
ignored_layers = json_quant_configs.get(
"exclude_modules", json_quant_configs.get("ignore")
)
kv_cache_quant_algo = json_quant_configs.get("kv_cache_quant_algo")
if kv_cache_quant_algo is None:
kv_cache_scheme = json_quant_configs.get("kv_cache_scheme")
if (
isinstance(kv_cache_scheme, dict)
and kv_cache_scheme.get("type") == "float"
and kv_cache_scheme.get("num_bits") == 8
):
kv_cache_quant_algo = "FP8"
parsed = {
"quant_method": "mxfp8",
"quant_algo": quant_algo,
"activation_scheme": "dynamic",
"weight_block_size": [1, group_size],
"scale_fmt": "ue8m0",
}
if ignored_layers is not None:
parsed["modules_to_not_convert"] = ignored_layers
if kv_cache_quant_algo is not None:
parsed["kv_cache_quant_algo"] = kv_cache_quant_algo
return parsed
else:
return None
def get_quantization_config_log_str(self) -> Optional[str]:
"""
Get a concise string representation of the quantization config for logging.
Returns something like "quant=fp8, fmt=e4m3" or "quant=gptq, bits=4".
"""
try:
quant_cfg = self._parse_quant_hf_config()
if not quant_cfg:
return None
quant_method = quant_cfg.get("quant_method", "quantized")
log_str = f"quant={quant_method}"
# Append interesting fields if they exist
for field in ["bits", "quant_algo", "fmt", "requantization_method"]:
if field in quant_cfg:
log_str += f", {field}={quant_cfg[field]}"
return log_str
except Exception:
return None
def _is_already_quantized(self) -> bool:
"""Check if the model is already quantized based on config files."""
# Check for quantization in hf_config (config.json)
if getattr(self.hf_config, "quantization_config", None) or getattr(
self.hf_config, "compression_config", None
):
return True
# Check for HuggingFace quantization config
quant_cfg = getattr(self.hf_config, "quantization_config", None)
if quant_cfg is None:
from sglang.srt.utils import has_hf_quant_config
return has_hf_quant_config(self.model_path)
return True
def _get_modelopt_quant_type(self) -> str:
"""Extract ModelOpt quantization type from unified quantization flag."""
if self.quantization == "modelopt_fp8":
return "fp8"
elif self.quantization == "modelopt_fp4":
return "nvfp4"
elif self.quantization == "modelopt_mixed":
raise ValueError(
"modelopt_mixed is only supported for pre-quantized checkpoints."
)
elif self.quantization == "modelopt":
# Auto-detect from model config
quant_cfg = self._parse_quant_hf_config()
if quant_cfg:
quant_method = quant_cfg.get("quant_method", "").lower()
if "fp4" in quant_method:
return "fp4"
elif "fp8" in quant_method:
return "fp8"
# Default to fp8 if can't detect
return "fp8"
else:
return "fp8" # Default fallback
def _get_sliding_window_size(self) -> Optional[int]:
for key in ("sliding_window_size", "sliding_window", "window_size"):
value = getattr(self.hf_text_config, key, None)
if value is not None:
return value
return None
def _validate_quantize_and_serve_config(self):
"""Validate quantize_and_serve configuration."""
if not self.quantize_and_serve:
return
# Check if ModelOpt quantization is specified
_MODELOPT_QUANTIZATION_METHODS = [
"modelopt",
"modelopt_fp8",
"modelopt_fp4",
"nvfp4_online",
"modelopt_mixed",
]
modelopt_quantization_specified = (
self.quantization in _MODELOPT_QUANTIZATION_METHODS
)
if not modelopt_quantization_specified:
raise ValueError(
"quantize_and_serve requires ModelOpt quantization (set with --quantization "
f"{{{', '.join(sorted(_MODELOPT_QUANTIZATION_METHODS))}}})"
)
# quantize_and_serve is disabled due to compatibility issues
raise NotImplementedError(
"quantize_and_serve functionality is currently disabled due to compatibility issues. "
"Please use the separate quantize-then-deploy workflow instead. "
"Step 1: Quantize and export model. "
"Step 2: Deploy the exported model."
)
# adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/config.py
def _verify_quantization(self) -> None:
supported_quantization = [*QUANTIZATION_METHODS]
rocm_supported_quantization = [
"awq",
"fp8",
"compressed_tensors",
"compressed-tensors",
"w8a8_fp8",
"petit_nvfp4",
"quark",
"mxfp4",
"mxfp8",
"auto-round",
"auto-round-int8",
"quark_int4fp8_moe",
"quark_mxfp4",
]
optimized_quantization_methods = [
"fp8",
"marlin",
"modelopt_fp8",
"modelopt_fp4",
"modelopt_mixed",
"nvfp4_online",
"gptq_marlin_24",
"gptq_marlin",
"awq_marlin",
"compressed_tensors",
"compressed-tensors",
"experts_int8",
"w8a8_int8",
"w8a8_fp8",
"moe_wna16",
"w4afp8",
"petit_nvfp4",
"quark",
"modelslim",
"humming",
"quark_mxfp4",
"auto-round",
"minimax_m3_nvfp4",
]
compatible_quantization_methods = {
"modelopt_fp8": ["modelopt"],
# Keep explicit or inherited modelopt_fp4 for literal FP8 checkpoints
# so eligible MoE experts are requantized online.
"modelopt_fp4": ["modelopt", "fp8"],
"modelopt_mixed": ["modelopt"],
"nvfp4_online": ["fp8", "modelopt_fp8"],
# Resolved by the MiniMax override from the MoE fields.
"minimax_m3_nvfp4": ["mxfp8"],
"petit_nvfp4": ["modelopt"],
"w8a8_int8": ["compressed-tensors", "compressed_tensors"],
"w8a8_fp8": ["compressed-tensors", "compressed_tensors"],
"auto-round-int8": ["compressed-tensors", "compressed_tensors"],
}
if self.quantization is not None:
self.quantization = self.quantization.lower()
# Parse quantization method from the HF and ModelSlim model config, if available.
# Only one function should return config, other should return None.
cfg_list = []
hf_config = self._parse_quant_hf_config()
modelslim_config = self._find_quant_modelslim_config()
quant_config = modelslim_config or hf_config
if quant_config is not None:
cfg_list.append(quant_config)
# Filter out None values
cfg_list = [item for item in cfg_list if item is not None]
if len(cfg_list) > 1:
raise ValueError(
"Config list contains configs from 2 methods, must be only 1"
)
quant_cfg = cfg_list[0] if cfg_list else None
if quant_cfg is not None:
quant_method = quant_cfg.get(
"quant_method", "" if not self.quantization else self.quantization
).lower()
# ModelOpt FP4 and mixed checkpoints can quantize only the target
# model; an embedded MTP draft may stay unquantized, so an explicit
# nvfp4_online opt-in for the draft wins over checkpoint detection.
# The online loader rejects already-packed weights at load time.
preserve_online_draft_quantization = (
self.is_draft_model
and self.quantization == "nvfp4_online"
and quant_method in ("modelopt_fp4", "modelopt_mixed")
)
# An explicit online-requantization request (e.g. quark_mxfp4 on top
# of an NVFP4/mixed checkpoint) must not be overridden back to the
# source format
if self.quantization not in REQUANTIZATION_METHODS:
# Detect which checkpoint is it
if not preserve_online_draft_quantization:
for _, method in QUANTIZATION_METHODS.items():
quantization_override = method.override_quantization_method(
quant_cfg, self.quantization
)
if quantization_override:
quant_method = quantization_override
self.quantization = quantization_override
break
# Verify quantization configurations.
if self.quantization is None:
self.quantization = quant_method
elif self.quantization != quant_method:
# Check if the CLI-specified quantization is compatible with HF config's quant_method
is_compatible = preserve_online_draft_quantization or (
self.quantization in compatible_quantization_methods
and quant_method
in compatible_quantization_methods[self.quantization]
)
if is_compatible:
# Keep the CLI-specified quantization (e.g., modelopt_fp4) even if
# HF config says "modelopt" - they are compatible
logger.info(
f"Using CLI-specified quantization ({self.quantization}) which is "
f"compatible with HF config quant_method ({quant_method})."
)
elif self.is_draft_model and not (
self.is_draft_quantization_explicit
and self.quantization in REQUANTIZATION_METHODS
and is_hip()
and is_qwen3_5_mtp_draft(self.hf_config)
):
# Allow auto-detection of quantization from checkpoint for draft model
# only if the CLI quantization is not compatible. An explicit
# online-requantization request for the draft (e.g. quark_mxfp4
# for an MTP stack the checkpoint left in bf16) is honored below.
logger.info(
f"Draft model quantization ({quant_method}) differs from "
f"main model quantization ({self.quantization}). "
f"Using draft model's detected quantization: {quant_method}"
)
self.quantization = quant_method
elif self.quantization in REQUANTIZATION_METHODS:
logger.info_once(
f"Requantizing from quant_method='{quant_method}' to the requested online quantization='{self.quantization}'. Beware that requantization may incur a loss in accuracy, the requantized model should be re-validated/re-evaluated. More details at https://docs.sglang.io/advanced_features/quantization.html#online-quantization."
)
else:
raise ValueError(
"Quantization method specified in the model config "
f"({quant_method}) does not match the quantization "
f"method specified in the `quantization` argument "
f"({self.quantization})."
)
if self.quantization is not None:
if self.quantization not in supported_quantization:
raise ValueError(
f"Unknown quantization method: {self.quantization}. Must "
f"be one of {supported_quantization}."
)
if is_hip() and self.quantization not in rocm_supported_quantization:
raise ValueError(
f"{self.quantization} quantization is currently not "
f"supported in ROCm."
)
if self.quantization not in optimized_quantization_methods:
# Don't warn for MXFP4/MXFP8 on SM100 since they have optimized kernels
if not (
self.quantization in ["mxfp4", "mxfp8"] and get_platform().is_sm100
):
logger.warning(
"%s quantization is not fully "
"optimized yet. The speed can be slower than "
"non-quantized models.",
self.quantization,
)
def _verify_dual_chunk_attention_config(self) -> None:
if hasattr(self.hf_config, "dual_chunk_attention_config"):
# Try loading the sparse attention config
sparse_attn_config = get_sparse_attention_config(self.model_path)
if not sparse_attn_config:
return
self.hf_config.dual_chunk_attention_config["sparse_attention_config"] = (
sparse_attn_config
)
if (
"sparse_attention_enabled"
not in self.hf_config.dual_chunk_attention_config
):
self.hf_config.dual_chunk_attention_config[
"sparse_attention_enabled"
] = True
def _get_hf_eos_token_id(self) -> Optional[Set[int]]:
eos_ids = getattr(self.hf_config, "eos_token_id", None)
if eos_ids is not None:
# it can be either int or list of int
eos_ids = {eos_ids} if isinstance(eos_ids, int) else set(eos_ids)
if eos_ids is None:
eos_ids = set()
if self.hf_generation_config:
generation_eos_ids = getattr(
self.hf_generation_config, "eos_token_id", None
)
if generation_eos_ids:
generation_eos_ids = (
{generation_eos_ids}
if isinstance(generation_eos_ids, int)
else set(generation_eos_ids)
)
eos_ids = eos_ids | generation_eos_ids
return eos_ids
def get_default_sampling_params(self) -> dict[str, Any]:
"""
Get default sampling parameters from the model's generation config.
This method returns non-default sampling parameters from the model's
generation_config.json when sampling_defaults is set to "model".
Returns:
A dictionary containing the non-default sampling parameters.
"""
if self.sampling_defaults != "model":
return {}
if self.hf_generation_config is None:
return {}
config = self.hf_generation_config.to_dict()
available_params = [
"repetition_penalty",
"temperature",
"top_k",
"top_p",
"min_p",
]
default_sampling_params = {
p: config.get(p) for p in available_params if config.get(p) is not None
}
return default_sampling_params
def _maybe_pull_model_for_runai(self, model: str) -> None:
if is_runai_obj_uri(model):
# local path for loading the config
self.model_path = ObjectStorageModel.get_path(model)
# remote path for loading the weights
self.model_weights = model
def _maybe_pull_model_tokenizer_from_remote(self) -> None:
"""
Pull the model config files to a temporary
directory in case of remote.
Args:
model: The model name or path.
"""
from sglang.srt.connector import create_remote_connector
from sglang.srt.utils import is_remote_url
if is_remote_url(self.model_path):
logger.info("Pulling model configs from remote...")
# BaseConnector implements __del__() to clean up the local dir.
# Since config files need to exist all the time, so we DO NOT use
# with statement to avoid closing the client.
client = create_remote_connector(self.model_path)
if is_remote_url(self.model_path):
client.pull_files(allow_pattern=["*config.json"])
self.model_weights = self.model_path
self.model_path = client.get_local_dir()
# adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/config.py
_STR_DTYPE_TO_TORCH_DTYPE = {
"half": torch.float16,
"float16": torch.float16,
"float": torch.float32,
"float32": torch.float32,
"bfloat16": torch.bfloat16,
}
# adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/config.py
def _get_and_verify_dtype(
config: PretrainedConfig,
dtype: Union[str, torch.dtype],
) -> torch.dtype:
# NOTE: getattr(config, "torch_dtype", torch.float32) is not correct
# because config.torch_dtype can be None.
if isinstance(config, dict):
config_dtype = config.get("dtype", None) or config.get("torch_dtype", None)
model_type = config.get("model_type", "")
else:
config_dtype = getattr(config, "dtype", None)
model_type = getattr(config, "model_type", "")
if isinstance(config_dtype, str):
config_dtype = _STR_DTYPE_TO_TORCH_DTYPE.get(config_dtype, None)
if config_dtype is None:
config_dtype = torch.float32
if isinstance(dtype, str):
dtype = dtype.lower()
if dtype == "auto":
if config_dtype == torch.float32:
if model_type.startswith("gemma"):
if model_type == "gemma":
gemma_version = ""
else:
gemma_version = model_type[5]
logger.info(
f"For Gemma {gemma_version}, we downcast float32 to bfloat16 instead "
"of float16 by default. Please specify `dtype` if you "
"want to use float16."
)
torch_dtype = torch.bfloat16
else:
# Following the common practice, we use float16 for float32
# models.
torch_dtype = torch.float16
else:
torch_dtype = config_dtype
else:
if dtype not in _STR_DTYPE_TO_TORCH_DTYPE:
raise ValueError(f"Unknown dtype: {dtype}")
torch_dtype = _STR_DTYPE_TO_TORCH_DTYPE[dtype]
elif isinstance(dtype, torch.dtype):
torch_dtype = dtype
else:
raise ValueError(f"Unknown dtype: {dtype}")
# Verify the dtype.
if torch_dtype != config_dtype:
if torch_dtype == torch.float32:
# Upcasting to float32 is allowed.
logger.debug("Upcasting %s to %s.", config_dtype, torch_dtype)
pass
elif config_dtype == torch.float32:
# Downcasting from float32 to float16 or bfloat16 is allowed.
logger.debug("Downcasting %s to %s.", config_dtype, torch_dtype)
pass
else:
# Casting between float16 and bfloat16 is allowed with a warning.
logger.debug("Casting %s to %s.", config_dtype, torch_dtype)
return torch_dtype
def is_embedding_gemma(config) -> bool:
"""Whether ``config`` is Google's bidirectional EmbeddingGemma checkpoint.
EmbeddingGemma uses the otherwise generative ``Gemma3TextModel``
architecture, so its model type alone is insufficient for dispatch. The
upstream ``use_bidirectional_attention`` flag is the defining distinction.
"""
return getattr(config, "model_type", None) == "gemma3_text" and getattr(
config, "use_bidirectional_attention", False
)
def is_generation_model(model_architectures: List[str], is_embedding: bool = False):
# We have two ways to determine whether a model is a generative model.
# 1. Check the model architecture
# 2. check the `is_embedding` server args
if (
"LlamaEmbeddingModel" in model_architectures
or "MistralModel" in model_architectures
or "LlamaForSequenceClassification" in model_architectures
or "LlamaForSequenceClassificationWithNormal_Weights" in model_architectures
or "InternLM2ForRewardModel" in model_architectures
or "Qwen2ForRewardModel" in model_architectures
or "Qwen3ForRewardModel" in model_architectures
or "Qwen2ForSequenceClassification" in model_architectures
or "Qwen3ForSequenceClassification" in model_architectures
or "Qwen3Model" in model_architectures
or "CLIPModel" in model_architectures
or "BertModel" in model_architectures
or "Contriever" in model_architectures
or "BertForSequenceClassification" in model_architectures
or "XLMRobertaModel" in model_architectures
or "XLMRobertaForSequenceClassification" in model_architectures
or "Gemma2ForSequenceClassification" in model_architectures
or "EmbeddingGemma2Model" in model_architectures
or "Lfm2BidirectionalModel" in model_architectures
):
return False
else:
return not is_embedding
multimodal_model_archs = [
"CLIPModel",
"Cohere2VisionForConditionalGeneration",
"Cosmos3EdgeForConditionalGeneration",
"DeepseekVL2ForCausalLM",
"Ernie4_5_VLMoeForConditionalGeneration",
"MiniMaxM3SparseForConditionalGeneration",
"Gemma3ForConditionalGeneration",
"Gemma3nForConditionalGeneration",
"Gemma4ForConditionalGeneration",
"EmbeddingGemma2Model",
"Gemma4UnifiedForConditionalGeneration",
"DiffusionGemmaForBlockDiffusion",
"Glm4vForConditionalGeneration",
"Glm4vMoeForConditionalGeneration",
"Glm5NextForConditionalGeneration",
"GlmOcrForConditionalGeneration",
"GlmAsrForConditionalGeneration",
"GlmImageForConditionalGeneration",
"Grok1VForCausalLM",
"Grok1AForCausalLM",
"LlavaLlamaForCausalLM",
"Llama4ForConditionalGeneration",
"LlavaMistralForCausalLM",
"LlavaQwenForCausalLM",
"LlavaForConditionalGeneration",
"LlavaVidForCausalLM",
"Lfm2VlForConditionalGeneration",
"LightOnOCRForConditionalGeneration",
*MIMO_V2_MULTIMODAL_ARCHS,
"MiMoV2ASRForCausalLM",
"MiniCPMO",
"MiniCPMV",
"Mistral3ForConditionalGeneration",
"MultiModalityCausalLM",
"MllamaForConditionalGeneration",
"MossVLForConditionalGeneration",
"NemotronH_Nano_VL_V2",
"NemotronH_Nano_Omni_Reasoning_V3",
"NemotronH_Omni_Reasoning_V3",
"MuseGlimmerForConditionalGeneration",
"MuseSparkForConditionalGeneration",
"PixtralForConditionalGeneration",
"Qwen2AudioForConditionalGeneration",
"Qwen2VLForConditionalGeneration",
"Qwen2_5_VLForConditionalGeneration",
"Qwen3VLForConditionalGeneration",
"Qwen3VLMoeForConditionalGeneration",
"Qwen3_5ForConditionalGeneration",
"SenseNovaLoopForConditionalGeneration",
"Qwen3_5MoeForConditionalGeneration",
"Qwen4ExpForConditionalGeneration",
"InternS2PreviewForConditionalGeneration",
"InternS2MobiusForConditionalGeneration",
"Qwen3ASRForConditionalGeneration",
"Qwen3OmniMoeForConditionalGeneration",
"KimiVLForConditionalGeneration",
"LocateAnythingForConditionalGeneration",
"InternVLChatModel",
"InternS1ForConditionalGeneration",
"InternS1ProForConditionalGeneration",
"Phi4MMForCausalLM",
"VoxtralForConditionalGeneration",
"WhisperForConditionalGeneration",
"Step3VLForConditionalGeneration",
"POINTSV15ChatModel",
"DotsVLMForCausalLM",
"Dots3NoteForCausalLM",
"DotsOCRForCausalLM",
"Sarashina2VisionForCausalLM",
"NVILAForConditionalGeneration",
"NVILALiteForConditionalGeneration",
"DeepseekOCRForCausalLM",
"UnlimitedOCRForCausalLM",
"JetVLMForConditionalGeneration",
"PaddleOCRVLForConditionalGeneration",
"MiDashengLMModel",
"StepVLForConditionalGeneration",
"Step3p7ForConditionalGeneration",
"KimiK25ForConditionalGeneration",
"BailingMMNativeForConditionalGeneration",
"BailingMM2NativeForConditionalGeneration",
"BailingMoeV3VLForConditionalGeneration",
]
piecewise_cuda_graph_disabled_model_archs = [
"DeepseekV4ForCausalLM",
"DeepseekV4ForCausalLMNextN",
"DeepseekV4ForCausalLMDSpark",
"Qwen3NextForCausalLM",
"Glm5NextForConditionalGeneration",
"BailingMoeV2_5ForCausalLM",
"LLaDAModelLM",
]
# Multimodal archs allowed to keep prefill piecewise CUDA graph enabled. The
# generic "multimodal model" rule in ServerArgs disables prefill piecewise CG for
# all multimodal models; archs here opt back in because their LM prefill captures
# cleanly (vision encoder runs eagerly outside the graph via general_mm_embed_routine).
multimodal_piecewise_cuda_graph_supported_model_archs = [
"KimiK25ForConditionalGeneration",
"MiniMaxM3SparseForCausalLM",
"MiniMaxM3SparseForConditionalGeneration",
]
# Multimodal archs whose LM prefill is validated under breakable CUDA graph;
# replay eligibility for embed-carrying batches is checked by can_run_graph.
# The Kimi archs are structurally multimodal -- their configs always carry a
# vision_config, so is_multimodal is True even for text-only serving -- and the
# generic multimodal rule disabled prefill CG for them despite the LM prefill
# capturing cleanly.
multimodal_breakable_cuda_graph_supported_model_archs = [
"MiniMaxM3SparseForConditionalGeneration",
"Cohere2VisionForConditionalGeneration",
"Glm5NextForConditionalGeneration",
"InternS2MobiusForConditionalGeneration",
"PaddleOCRVLForConditionalGeneration",
"Qwen3_5ForConditionalGeneration",
"Qwen3_5MoeForConditionalGeneration",
"Qwen4ExpForConditionalGeneration",
"MuseGlimmerForConditionalGeneration",
"MuseSparkForConditionalGeneration",
"KimiK3ForConditionalGeneration",
"KimiK25ForConditionalGeneration",
"MiniMaxM3SparseForCausalLM",
"MiniMaxM3SparseForConditionalGeneration",
]
if external_mm_model_arch := envs.SGLANG_EXTERNAL_MM_MODEL_ARCH.get():
multimodal_model_archs.append(external_mm_model_arch)
def is_multimodal_model(model_architectures: List[str]):
if any(
multi_model_arch in model_architectures
for multi_model_arch in multimodal_model_archs
):
return True
else:
return False
def is_audio_model(model_architectures: List[str]):
models = [
"WhisperForConditionalGeneration",
"Qwen3ASRForConditionalGeneration",
]
return any(model in model_architectures for model in models)
def is_encoder_decoder_model(model_architectures: List[str]):
models = [
"WhisperForConditionalGeneration",
"MllamaForConditionalGeneration",
"MossVLForConditionalGeneration",
]
return any(model in model_architectures for model in models)
def is_local_attention_model(model_architectures: List[str]):
return "Llama4ForConditionalGeneration" in model_architectures
def is_multimodal_chunked_prefill_supported(model_architectures: List[str]):
"""Check if chunked prefill is supported for a MultiModal model."""
unsupported = [
"Grok1VForCausalLM",
"Grok1AForCausalLM",
"LlavaLlamaForCausalLM",
"MllamaForConditionalGeneration",
"MossVLForConditionalGeneration",
"CLIPModel",
]
if any(multi_model_arch in unsupported for multi_model_arch in model_architectures):
return False
else:
return True
def is_piecewise_cuda_graph_disabled_model(model_architectures: List[str]):
return any(
arch in piecewise_cuda_graph_disabled_model_archs
for arch in model_architectures
)
def is_multimodal_piecewise_cuda_graph_supported(model_architectures: List[str]):
"""Whether a multimodal arch may keep prefill piecewise CUDA graph enabled."""
return any(
arch in multimodal_piecewise_cuda_graph_supported_model_archs
for arch in model_architectures
)
def is_multimodal_breakable_cuda_graph_supported(model_architectures: List[str]):
"""Whether a multimodal arch may keep prefill breakable CUDA graph enabled."""
return any(
arch in multimodal_breakable_cuda_graph_supported_model_archs
for arch in model_architectures
)
# SequenceClassification models that use CrossEncodingPooler
_cross_encoding_pooler_archs = [
"BertForSequenceClassification",
"XLMRobertaForSequenceClassification",
]
def is_cross_encoding_pooler_model(model_architectures: List[str]) -> bool:
return any(arch in _cross_encoding_pooler_archs for arch in model_architectures)
# SequenceClassification models whose forward routes the head through
# score_and_pool (per-position pooling); only these support setwise readout
# (token_indices_to_pool). Keep in sync with callers of layers.pooler.score_and_pool.
_score_and_pool_archs = [
"LlamaForSequenceClassification",
"Qwen2ForSequenceClassification",
"Qwen3ForSequenceClassification",
]
def is_score_and_pool_model(model_architectures: List[str]) -> bool:
return any(arch in _score_and_pool_archs for arch in model_architectures)
def yarn_get_mscale(scale: float = 1, mscale: float = 1) -> float:
if scale <= 1:
return 1.0
return 0.1 * mscale * math.log(scale) + 1.0
def compute_mla_mscale_scaling(rope_scaling: dict, base_scaling: float) -> float:
"""Compute MLA attention scaling factor from rope_scaling with mscale.
Used by DeepSeek, BailingMoe, SarvamMLA and similar MLA models.
Transformers v5 also exposes the default RoPE parameters through
``rope_scaling``. Those parameters do not request any scaling.
Warns if 'factor' is missing from a scaling request (common in v5 configs).
"""
# v5 uses "rope_type", v4 uses "type"
rope_type = rope_scaling.get("rope_type") or rope_scaling.get("type")
if rope_type == "default":
return base_scaling
if not rope_scaling.get("apply_yarn_scaling", True) or not rope_scaling.get(
"apply_scale", True
):
return base_scaling
mscale_all_dim = rope_scaling.get("mscale_all_dim", False)
if "factor" not in rope_scaling:
logger.warning(
"rope_scaling missing 'factor', defaulting to 1.0. Check model accuracy.",
)
scaling_factor = rope_scaling.get("factor", 1.0)
mscale = yarn_get_mscale(scaling_factor, float(mscale_all_dim))
return base_scaling * mscale * mscale
def is_hybrid_swa_model(
model_architectures: List[str],
hf_text_config: Optional[PretrainedConfig] = None,
):
hybrid_swa_archs = {
"BeamForCausalLM",
"Llama4ForConditionalGeneration",
"DeepseekV4ForCausalLM",
"DeepseekV4ForCausalLMNextN",
"DeepseekV4ForCausalLMDSpark",
*SWA_SINK_ARCHS,
*MIMO_V2_MODEL_ARCHS,
"MiMoV2MTP",
"Step3p5ForCausalLM",
"Step3p5MTP",
"Step3p7ForConditionalGeneration",
"Gemma4ForCausalLM",
"Gemma4ForConditionalGeneration",
"EmbeddingGemma2Model",
"Gemma4UnifiedForConditionalGeneration",
"DiffusionGemmaForBlockDiffusion",
"LagunaForCausalLM",
"MellumForCausalLM",
"MuseGlimmerForCausalLM",
"MuseGlimmerForConditionalGeneration",
"InklingForConditionalGeneration",
"InklingForConditionalGenerationMTP",
"UnlimitedOCRForCausalLM",
}
if any(arch in hybrid_swa_archs for arch in model_architectures):
if "BeamForCausalLM" in model_architectures and hf_text_config is not None:
# A schedule without sliding layers uses the ordinary KV pool; the
# SWA pool also serves all-sliding schedules.
return "sliding_attention" in hf_text_config.layer_types
# Only treat Laguna as hybrid SWA when it actually has a sliding window.
if (
"LagunaForCausalLM" in model_architectures
and hf_text_config is not None
and not getattr(hf_text_config, "sliding_window", 0)
):
return False
return True
# Also recognize models that explicitly opt-in via their HF text config,
# so custom hybrid-SWA architectures don't need to be added to the allowlist.
if hf_text_config is not None and getattr(hf_text_config, "is_hybrid_swa", False):
return True
return False
def get_hybrid_layer_ids(
model_architectures: List[str],
hf_text_config: PretrainedConfig,
):
num_hidden_layers = hf_text_config.num_hidden_layers
if "Llama4ForConditionalGeneration" in model_architectures:
swa_attention_layer_ids = [
i for i in range(num_hidden_layers) if (i + 1) % 4 != 0
]
full_attention_layer_ids = [
i for i in range(num_hidden_layers) if (i + 1) % 4 == 0
]
elif any(arch in SWA_SINK_ARCHS for arch in model_architectures) or any(
arch in ("Dots3NoteForCausalLM", "Dots3NoteForCausalLMNextN")
for arch in model_architectures
):
layer_types = getattr(hf_text_config, "layer_types", [])
swa_attention_layer_ids = [
i for i, x in enumerate(layer_types) if x == "sliding_attention"
]
full_attention_layer_ids = [
i for i, x in enumerate(layer_types) if x == "full_attention"
]
elif any(arch in MIMO_V2_MODEL_ARCHS for arch in model_architectures):
hybrid_layer_pattern = getattr(hf_text_config, "hybrid_layer_pattern", None)
swa_attention_layer_ids = [
i for i in range(num_hidden_layers) if hybrid_layer_pattern[i] == 1
]
full_attention_layer_ids = [
i for i in range(num_hidden_layers) if hybrid_layer_pattern[i] == 0
]
elif "MiMoV2MTP" in model_architectures:
swa_attention_layer_ids = [0]
full_attention_layer_ids = []
elif (
"Step3p5ForCausalLM" in model_architectures
or "Step3p7ForConditionalGeneration" in model_architectures
):
layer_types = hf_text_config.layer_types
swa_attention_layer_ids = [
i
for i, x in enumerate(layer_types)
if x == "sliding_attention" and i < num_hidden_layers
]
full_attention_layer_ids = [
i
for i, x in enumerate(layer_types)
if x == "full_attention" and i < num_hidden_layers
]
elif "Step3p5MTP" in model_architectures:
swa_attention_layer_ids = [0]
full_attention_layer_ids = []
elif (
"BeamForCausalLM" in model_architectures
or "Gemma4ForCausalLM" in model_architectures
or "Gemma4ForConditionalGeneration" in model_architectures
or "EmbeddingGemma2Model" in model_architectures
or "Gemma4UnifiedForConditionalGeneration" in model_architectures
or "DiffusionGemmaForBlockDiffusion" 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 = [
i for i, x in enumerate(layer_types) if x == "sliding_attention"
]
full_attention_layer_ids = [
i for i, x in enumerate(layer_types) if x == "full_attention"
]
elif "InklingForConditionalGenerationMTP" in model_architectures:
# Runner-local layer setup selects one depth from these shared lists,
# including when every draft depth uses full attention.
local_set = set(hf_text_config.mtp_local_layer_ids)
swa_attention_layer_ids = sorted(local_set)
full_attention_layer_ids = [
i
for i in range(hf_text_config.num_nextn_predict_layers)
if i not in local_set
]
elif "InklingForConditionalGeneration" in model_architectures:
local_layer_ids = hf_text_config.local_layer_ids
local_layer_id_set = set(local_layer_ids)
assert len(local_layer_id_set) == len(local_layer_ids), (
f"Inkling local_layer_ids must be unique: {local_layer_ids}"
)
assert all(
0 <= layer_id < num_hidden_layers for layer_id in local_layer_id_set
), (
f"Inkling local_layer_ids must be in [0, {num_hidden_layers}): {local_layer_ids}"
)
swa_attention_layer_ids = [
i for i in range(num_hidden_layers) if i in local_layer_id_set
]
full_attention_layer_ids = [
i for i in range(num_hidden_layers) if i not in local_layer_id_set
]
elif "UnlimitedOCRForCausalLM" in model_architectures:
swa_attention_layer_ids = list(range(num_hidden_layers))
full_attention_layer_ids = []
elif "MuseSparkForConditionalGeneration" in model_architectures:
swa_attention_layer_ids = [
i
for i, kind in enumerate(hf_text_config.layer_types)
if kind == "sliding_attention"
]
full_attention_layer_ids = [
i
for i, kind in enumerate(hf_text_config.layer_types)
if kind == "full_attention"
]
elif getattr(hf_text_config, "hybrid_layer_pattern", None) is not None:
# Generic fallback for custom hybrid SWA models that opt in via
# hf_text_config.is_hybrid_swa and expose a hybrid_layer_pattern
# (1 = SWA, 0 = full) without needing to be added to the allowlist.
hybrid_layer_pattern = hf_text_config.hybrid_layer_pattern
swa_attention_layer_ids = [
i for i in range(num_hidden_layers) if hybrid_layer_pattern[i] == 1
]
full_attention_layer_ids = [
i for i in range(num_hidden_layers) if hybrid_layer_pattern[i] == 0
]
else:
swa_attention_layer_ids = None
full_attention_layer_ids = None
return swa_attention_layer_ids, full_attention_layer_ids