3bac084d4e
Co-authored-by: Xiaoyu Zhang <1182563586@qq.com> Co-authored-by: Xinyuan Tong <xinyuantong.cs@gmail.com>
2277 lines
95 KiB
Python
2277 lines
95 KiB
Python
# 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
|
|
from pathlib import Path
|
|
from typing import Any, List, Optional, Set, Union
|
|
|
|
import torch
|
|
from transformers import PretrainedConfig
|
|
|
|
from sglang.srt.arg_groups.overrides import resolving_view
|
|
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_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__)
|
|
|
|
MIMO_V2_MODEL_ARCHS = (
|
|
"MiMoV2ForCausalLM",
|
|
"MiMoV2FlashForCausalLM",
|
|
)
|
|
MIMO_V2_MULTIMODAL_ARCHS = ("MiMoV2ForCausalLM",)
|
|
|
|
SWA_SINK_ARCHS = frozenset(
|
|
{
|
|
"GptOssForCausalLM",
|
|
"GraniteSWAForCausalLM",
|
|
"GraniteMoeSWAForCausalLM",
|
|
}
|
|
)
|
|
|
|
|
|
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 get_mimo_v2_fused_qkv_expected_tp_size(hf_config):
|
|
layout = getattr(hf_config, "attention_projection_layout", None)
|
|
if layout is None:
|
|
return None
|
|
if layout != "fused_qkv":
|
|
raise ValueError(
|
|
"MiMoV2 hf_config has unsupported "
|
|
f"attention_projection_layout={layout!r}; expected 'fused_qkv' "
|
|
"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()
|
|
|
|
|
|
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",
|
|
"LongcatFlashForCausalLM",
|
|
"LongcatFlashForCausalLMNextN",
|
|
"Dots3NoteForCausalLM",
|
|
"Dots3NoteForCausalLMNextN",
|
|
)
|
|
and _hf_attr(config, "index_topk") is not None
|
|
)
|
|
|
|
|
|
def is_kimi_k3(config) -> bool:
|
|
return _hf_arch(config) in (
|
|
"KimiK3ForConditionalGeneration",
|
|
"KimiK3LinearForCausalLM",
|
|
)
|
|
|
|
|
|
def is_dspark_draft(config) -> bool:
|
|
return _hf_arch(config) == "DSparkDraftModel"
|
|
|
|
|
|
def is_qwen3_5(config) -> bool:
|
|
return _hf_arch(config) in (
|
|
"Qwen3_5ForConditionalGeneration",
|
|
"Qwen3_5MoeForConditionalGeneration",
|
|
"Qwen3_5ForCausalLM",
|
|
"Qwen3_5MoeForCausalLM",
|
|
)
|
|
|
|
|
|
def is_deepseek_v4(config) -> bool:
|
|
return _hf_arch(config) in (
|
|
"DeepseekV4ForCausalLM",
|
|
"DeepseekV4ForCausalLMNextN",
|
|
"DeepseekV4ForCausalLMDSpark",
|
|
)
|
|
|
|
|
|
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 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)
|
|
|
|
# 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
|
|
|
|
|
|
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.hf_text_config = get_hf_text_config(self.hf_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", {}
|
|
)
|
|
self.is_lm_only = 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 {}
|
|
)
|
|
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.")
|
|
|
|
# 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")
|
|
)
|
|
)
|
|
self.is_multimodal = (
|
|
enable_multimodal
|
|
and not self.is_lm_only
|
|
and (
|
|
is_multimodal_model(self.hf_config.architectures)
|
|
or has_multimodal_subconfig
|
|
)
|
|
)
|
|
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
|
|
)
|
|
|
|
# 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)
|
|
# 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
|
|
)
|
|
)
|
|
)
|
|
# Multimodal archs whose language-model prefill is verified safe to capture
|
|
# under piecewise CUDA graph. ServerArgs otherwise disables prefill piecewise
|
|
# CG for every multimodal model; this opt-in re-enables it for listed archs
|
|
# (the vision encoder still runs eagerly via general_mm_embed_routine, only the
|
|
# LM forward is captured).
|
|
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)
|
|
)
|
|
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.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 = language_model_only or 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,
|
|
):
|
|
|
|
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 _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.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] 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 (
|
|
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] == "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] 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] == "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] 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
|
|
|
|
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
|
|
)
|
|
|
|
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",
|
|
"InklingForConditionalGeneration",
|
|
"InklingForConditionalGenerationMTP",
|
|
"Gemma4UnifiedForConditionalGeneration",
|
|
]
|
|
|
|
@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 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:
|
|
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
|
|
|
|
# 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:
|
|
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 (
|
|
"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 "LongcatFlashForCausalLM" in self.hf_config.architectures
|
|
or "LongcatFlashForCausalLMNextN" 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
|
|
):
|
|
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
|
|
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.
|
|
self._init_mla_scaling(self.hf_text_config.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 "BailingMoeForCausalLMNextN" in self.hf_config.architectures
|
|
):
|
|
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:
|
|
self.head_dim = 128
|
|
self.attention_arch = AttentionArch.MLA
|
|
self.kv_lora_rank = self.hf_config.kv_lora_rank
|
|
self.qk_rope_head_dim = (
|
|
0 if self.hf_config.use_mla_nope else self.hf_config.qk_rope_head_dim
|
|
)
|
|
self.v_head_dim = self.hf_config.v_head_dim
|
|
self.qk_nope_head_dim = self.hf_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:
|
|
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
|
|
|
|
self.attention_arch = AttentionArch.MHA
|
|
|
|
self.num_attention_heads = self.hf_text_config.num_attention_heads
|
|
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)
|
|
self.spec_hidden_size = (
|
|
self.hidden_size * hc_mult if hc_mult > 1 else self.hidden_size
|
|
)
|
|
# mHC-flattened hidden size; None when not running an mHC model
|
|
# (e.g. non-DeepSeek-V4 configs without ``hc_mult``).
|
|
self.hc_hidden_size = self.spec_hidden_size if hc_mult > 1 else None
|
|
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 "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:
|
|
total_num_attention_heads = self.num_attention_heads
|
|
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
|
|
|
|
# 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."""
|
|
json_quant_configs = quant_config_dict["quantization"]
|
|
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",
|
|
"gptq",
|
|
"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",
|
|
]
|
|
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"],
|
|
"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:
|
|
# Allow auto-detection of quantization from checkpoint for draft model
|
|
# only if the CLI quantization is not compatible
|
|
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 "Lfm2BidirectionalModel" in model_architectures
|
|
):
|
|
return False
|
|
else:
|
|
return not is_embedding
|
|
|
|
|
|
multimodal_model_archs = [
|
|
"CLIPModel",
|
|
"Cohere2VisionForConditionalGeneration",
|
|
"DeepseekVL2ForCausalLM",
|
|
"Ernie4_5_VLMoeForConditionalGeneration",
|
|
"MiniMaxM3SparseForConditionalGeneration",
|
|
"Gemma3ForConditionalGeneration",
|
|
"Gemma3nForConditionalGeneration",
|
|
"Gemma4ForConditionalGeneration",
|
|
"Gemma4UnifiedForConditionalGeneration",
|
|
"Glm4vForConditionalGeneration",
|
|
"Glm4vMoeForConditionalGeneration",
|
|
"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",
|
|
"MuseGlimmerForConditionalGeneration",
|
|
"PixtralForConditionalGeneration",
|
|
"Qwen2AudioForConditionalGeneration",
|
|
"Qwen2VLForConditionalGeneration",
|
|
"Qwen2_5_VLForConditionalGeneration",
|
|
"Qwen3VLForConditionalGeneration",
|
|
"Qwen3VLMoeForConditionalGeneration",
|
|
"Qwen3_5ForConditionalGeneration",
|
|
"Qwen3_5MoeForConditionalGeneration",
|
|
"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",
|
|
]
|
|
|
|
piecewise_cuda_graph_disabled_model_archs = [
|
|
"DeepseekV4ForCausalLM",
|
|
"DeepseekV4ForCausalLMNextN",
|
|
"DeepseekV4ForCausalLMDSpark",
|
|
"Qwen3NextForCausalLM",
|
|
"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;
|
|
# embed-carrying batches are rejected at replay (can_run_graph) and run eager.
|
|
# 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 = [
|
|
"Cohere2VisionForConditionalGeneration",
|
|
"InternS2MobiusForConditionalGeneration",
|
|
"PaddleOCRVLForConditionalGeneration",
|
|
"Qwen3_5ForConditionalGeneration",
|
|
"Qwen3_5MoeForConditionalGeneration",
|
|
"MuseGlimmerForConditionalGeneration",
|
|
"KimiK3ForConditionalGeneration",
|
|
"KimiK25ForConditionalGeneration",
|
|
]
|
|
|
|
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)
|
|
|
|
|
|
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 = {
|
|
"Llama4ForConditionalGeneration",
|
|
"DeepseekV4ForCausalLM",
|
|
"DeepseekV4ForCausalLMNextN",
|
|
"DeepseekV4ForCausalLMDSpark",
|
|
*SWA_SINK_ARCHS,
|
|
*MIMO_V2_MODEL_ARCHS,
|
|
"MiMoV2MTP",
|
|
"Step3p5ForCausalLM",
|
|
"Step3p5MTP",
|
|
"Step3p7ForConditionalGeneration",
|
|
"Gemma4ForCausalLM",
|
|
"Gemma4ForConditionalGeneration",
|
|
"Gemma4UnifiedForConditionalGeneration",
|
|
"LagunaForCausalLM",
|
|
"MellumForCausalLM",
|
|
"MuseGlimmerForCausalLM",
|
|
"MuseGlimmerForConditionalGeneration",
|
|
"InklingForConditionalGeneration",
|
|
"InklingForConditionalGenerationMTP",
|
|
"UnlimitedOCRForCausalLM",
|
|
}
|
|
if any(arch in hybrid_swa_archs for arch in model_architectures):
|
|
# 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 (
|
|
"Gemma4ForCausalLM" in model_architectures
|
|
or "Gemma4ForConditionalGeneration" in model_architectures
|
|
or "Gemma4UnifiedForConditionalGeneration" in model_architectures
|
|
or "LagunaForCausalLM" in model_architectures
|
|
or "MellumForCausalLM" in model_architectures
|
|
or "MuseGlimmerForCausalLM" in model_architectures
|
|
or "MuseGlimmerForConditionalGeneration" in model_architectures
|
|
):
|
|
layer_types = getattr(hf_text_config, "layer_types", [])
|
|
swa_attention_layer_ids = [
|
|
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:
|
|
# One block per MTP depth; a banded head marks its sliding-window depths
|
|
# in mtp_local_layer_ids. The per-depth pool routing in the KV-cache
|
|
# mixin is authoritative; this keeps model_config's swa/full lists
|
|
# self-consistent for other consumers.
|
|
mtp_local_layer_ids = hf_text_config.mtp_local_layer_ids
|
|
if mtp_local_layer_ids:
|
|
num_depths = hf_text_config.num_nextn_predict_layers
|
|
local_set = set(mtp_local_layer_ids)
|
|
swa_attention_layer_ids = sorted(local_set)
|
|
full_attention_layer_ids = [
|
|
i for i in range(num_depths) if i not in local_set
|
|
]
|
|
else:
|
|
swa_attention_layer_ids = []
|
|
full_attention_layer_ids = [0]
|
|
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 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
|