diff --git a/python/sglang/srt/configs/hybrid_arch.py b/python/sglang/srt/configs/hybrid_arch.py new file mode 100644 index 000000000..16a42d6e7 --- /dev/null +++ b/python/sglang/srt/configs/hybrid_arch.py @@ -0,0 +1,116 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +from sglang.srt.configs import ( + BailingHybridConfig, + FalconH1Config, + GraniteMoeHybridConfig, + InternS2PreviewConfig, + JetNemotronConfig, + JetVLMConfig, + KimiLinearConfig, + Lfm2Config, + Lfm2MoeConfig, + Lfm2VlConfig, + NemotronH_Nano_VL_V2_Config, + NemotronHConfig, + Qwen3_5Config, + Qwen3_5MoeConfig, + Qwen3NextConfig, + ZayaConfig, +) + +if TYPE_CHECKING: + from sglang.srt.configs.model_config import ModelConfig + + +def _get_linear_attn_registry_result(model_config: ModelConfig) -> Any: + return model_config.linear_attn_registry_result + + +def qwen3_next_config(model_config: ModelConfig): + config = model_config.hf_config + if isinstance(config, Qwen3NextConfig): + return config + return None + + +def hybrid_lightning_config(model_config: ModelConfig): + config = model_config.hf_config + if isinstance(config, BailingHybridConfig): + return config + return None + + +def hybrid_gdn_config(model_config: ModelConfig): + config = model_config.hf_config.get_text_config() + if isinstance( + config, + Qwen3NextConfig + | Qwen3_5Config + | Qwen3_5MoeConfig + | InternS2PreviewConfig + | JetNemotronConfig + | JetVLMConfig, + ): + return config + return None + + +def mamba2_config(model_config: ModelConfig): + config = model_config.hf_config + if isinstance(config, NemotronHConfig) and model_config.is_draft_model: + # NemotronH MTP draft models have no Mamba layers (pattern like "*E") + # so they shouldn't use HybridLinearAttnBackend + pattern = getattr(config, "mtp_hybrid_override_pattern", None) + if pattern is not None and "M" not in pattern: + return None + if isinstance( + config, + FalconH1Config + | NemotronHConfig + | Lfm2Config + | Lfm2MoeConfig + | Lfm2VlConfig + | ZayaConfig, + ): + return config + if isinstance(config, NemotronH_Nano_VL_V2_Config): + return config.llm_config + + if isinstance(config, GraniteMoeHybridConfig): + has_mamba = any( + layer_type == "mamba" for layer_type in getattr(config, "layer_types", []) + ) + if not has_mamba: + return None + else: + return config + + return None + + +def kimi_linear_config(model_config: ModelConfig): + config = model_config.hf_config + if isinstance(config, KimiLinearConfig): + return config + return None + + +def linear_attn_model_spec(model_config: ModelConfig): + result = _get_linear_attn_registry_result(model_config) + return result[0] if result else None + + +def mambaish_config(model_config: ModelConfig): + existing = ( + mamba2_config(model_config) + or hybrid_gdn_config(model_config) + or kimi_linear_config(model_config) + or hybrid_lightning_config(model_config) + ) + if existing: + return existing + result = _get_linear_attn_registry_result(model_config) + return result[1] if result else None diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index 523472315..9f3ea8f7c 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -18,12 +18,14 @@ 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.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.server_args import ServerArgs @@ -287,6 +289,14 @@ class ModelConfig: ) ) self.hf_text_config = get_hf_text_config(self.hf_config) + + rope_scaling = getattr(self.hf_text_config, "rope_parameters", None) or getattr( + self.hf_text_config, "rope_scaling", {} + ) + self.model_is_mrope = ( + 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, @@ -685,6 +695,10 @@ class ModelConfig: "Gemma4UnifiedForConditionalGeneration", ] + @cached_property + def linear_attn_registry_result(self) -> Any: + return get_linear_attn_config(self.hf_config) + def _detect_attention_sinks(self) -> bool: """Check whether the model uses learned attention sinks. diff --git a/python/sglang/srt/hardware_backend/mlx/model_runner_stub.py b/python/sglang/srt/hardware_backend/mlx/model_runner_stub.py index 44f103f79..6f0a7b75f 100644 --- a/python/sglang/srt/hardware_backend/mlx/model_runner_stub.py +++ b/python/sglang/srt/hardware_backend/mlx/model_runner_stub.py @@ -9,6 +9,7 @@ from typing import Tuple import torch +from sglang.srt.configs.hybrid_arch import mambaish_config from sglang.srt.hardware_backend.mlx.kv_cache.auxiliary_state import ( MlxAuxiliaryStateReqToTokenPool, ) @@ -160,7 +161,7 @@ class MlxModelRunnerStub(ModelRunner): self.is_hybrid_swa = False # Create minimal pools - if self.mambaish_config is not None: + if mambaish_config(self.model_config) is not None: auxiliary_state_size = self.server_args.max_mamba_cache_size if auxiliary_state_size is None: auxiliary_state_size = self.max_running_requests * 4 diff --git a/python/sglang/srt/layers/attention/attention_registry.py b/python/sglang/srt/layers/attention/attention_registry.py index 99e430809..9060ecc68 100644 --- a/python/sglang/srt/layers/attention/attention_registry.py +++ b/python/sglang/srt/layers/attention/attention_registry.py @@ -2,6 +2,13 @@ import logging import warnings from typing import TYPE_CHECKING +from sglang.srt.configs.hybrid_arch import ( + hybrid_gdn_config, + hybrid_lightning_config, + kimi_linear_config, + mamba2_config, + mambaish_config, +) from sglang.srt.configs.linear_attn_model_registry import ( get_linear_attn_config, import_backend_class, @@ -255,7 +262,7 @@ def attn_backend_wrapper(runner: "ModelRunner", full_attn_backend: "AttentionBac need to change the code of the original attention backend. """ assert not ( - runner.hybrid_gdn_config is not None and runner.use_mla_backend + hybrid_gdn_config(runner.model_config) is not None and runner.use_mla_backend ), "hybrid_gdn can only be used with non-MLA models." from sglang.srt.configs.model_config import is_minimax_sparse @@ -271,7 +278,7 @@ def attn_backend_wrapper(runner: "ModelRunner", full_attn_backend: "AttentionBac full_attn_backend, sparse_backend, sparse_backend.sparse_layer_ids ) - if cfg := runner.mambaish_config: + if cfg := mambaish_config(runner.model_config): from sglang.srt.layers.attention.fla.utils import check_environments from sglang.srt.layers.attention.linear.kda_backend import KDAAttnBackend from sglang.srt.layers.attention.linear.lightning_backend import ( @@ -303,11 +310,11 @@ def attn_backend_wrapper(runner: "ModelRunner", full_attn_backend: "AttentionBac ) check_environments() - if runner.hybrid_gdn_config is not None and not is_npu(): + if hybrid_gdn_config(runner.model_config) is not None and not is_npu(): maybe_set_default_flashinfer_gdn_prefill(runner) initialize_linear_attn_config(runner.server_args) hybrid_backend_cls = HybridLinearAttnBackend - if runner.hybrid_gdn_config is not None: + if hybrid_gdn_config(runner.model_config) is not None: if is_blackwell(): assert ( runner.server_args.attention_backend == "triton" @@ -321,7 +328,7 @@ def attn_backend_wrapper(runner: "ModelRunner", full_attn_backend: "AttentionBac ), "ascend backend is the only supported backend on NPU for hybrid GDN models, use --attention-backend ascend to specify the backend." logger.info(f"Using hybrid linear attention backend for hybrid GDN models.") linear_attn_backend = GDNAttnBackend(runner) - elif runner.mamba2_config is not None: + elif mamba2_config(runner.model_config) is not None: from sglang.srt.configs.lfm2 import Lfm2Config from sglang.srt.configs.lfm2_moe import Lfm2MoeConfig from sglang.srt.configs.lfm2_vl import Lfm2VlConfig @@ -337,7 +344,7 @@ def attn_backend_wrapper(runner: "ModelRunner", full_attn_backend: "AttentionBac Lfm2MoeConfig, Lfm2VlConfig, ) - if isinstance(runner.mamba2_config, short_conv_cfgs): + if isinstance(mamba2_config(runner.model_config), short_conv_cfgs): if is_npu(): # The model conv layers call # get_attn_backend().conv_state_metadata() unconditionally, @@ -362,9 +369,9 @@ def attn_backend_wrapper(runner: "ModelRunner", full_attn_backend: "AttentionBac hybrid_backend_cls = ShortConvHybridAttnBackend else: linear_attn_backend = Mamba2AttnBackend(runner) - elif runner.kimi_linear_config is not None: + elif kimi_linear_config(runner.model_config) is not None: linear_attn_backend = KDAAttnBackend(runner) - elif runner.hybrid_lightning_config is not None: + elif hybrid_lightning_config(runner.model_config) is not None: linear_attn_backend = LightningAttentionBackend(runner) else: spec_result = get_linear_attn_config(runner.model_config.hf_config) diff --git a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py index 5e2f46ba5..58cc14837 100644 --- a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py +++ b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py @@ -3,6 +3,7 @@ from typing import Optional, Union import torch +from sglang.srt.configs.hybrid_arch import mamba2_config from sglang.srt.layers.attention.base_attn_backend import AttentionBackend from sglang.srt.layers.attention.mamba.causal_conv1d_triton import PAD_SLOT_ID from sglang.srt.layers.attention.mamba.mamba import MambaMixer2 @@ -684,7 +685,7 @@ class Mamba2AttnBackend(MambaAttnBackendBase): def __init__(self, model_runner: ModelRunner): super().__init__(model_runner) - config = model_runner.mamba2_config + config = mamba2_config(model_runner.model_config) assert config is not None self.mamba_chunk_size = config.mamba_chunk_size self.conv_states_shape = ( diff --git a/python/sglang/srt/layers/attention/linear/gdn_backend.py b/python/sglang/srt/layers/attention/linear/gdn_backend.py index 5e3d9d09d..1d73b9f65 100644 --- a/python/sglang/srt/layers/attention/linear/gdn_backend.py +++ b/python/sglang/srt/layers/attention/linear/gdn_backend.py @@ -2,6 +2,7 @@ from typing import Optional, Tuple, Union import torch +from sglang.srt.configs.hybrid_arch import hybrid_gdn_config from sglang.srt.layers.attention.fla.fused_gdn_gating import fused_gdn_gating from sglang.srt.layers.attention.hybrid_linear_attn_backend import MambaAttnBackendBase from sglang.srt.layers.attention.linear.kernels.gdn_triton import TritonGDNKernel @@ -73,7 +74,7 @@ def maybe_set_default_flashinfer_gdn_prefill(model_runner: ModelRunner) -> None: cuda_version = torch.version.cuda chunk_size = args.chunked_prefill_size - config = model_runner.hybrid_gdn_config + config = hybrid_gdn_config(model_runner.model_config) if ( cuda_version is None or int(cuda_version.split(".", 1)[0]) < 13 diff --git a/python/sglang/srt/layers/attention/triton_backend.py b/python/sglang/srt/layers/attention/triton_backend.py index 25f2357d4..55091f114 100644 --- a/python/sglang/srt/layers/attention/triton_backend.py +++ b/python/sglang/srt/layers/attention/triton_backend.py @@ -10,6 +10,11 @@ from sglang.kernels.ops.attention.metadata import get_num_kv_splits_triton from sglang.kernels.ops.kvcache.kv_indices import ( create_flashinfer_kv_indices_triton, ) +from sglang.srt.configs.hybrid_arch import ( + hybrid_gdn_config, + kimi_linear_config, + linear_attn_model_spec, +) from sglang.srt.configs.model_config import AttentionArch from sglang.srt.distributed.device_communicators.pynccl_allocator import ( use_symmetric_memory, @@ -182,9 +187,9 @@ class TritonAttnBackend(AttentionBackend): self.v_head_dim = full_v_head_dim self.swa_v_head_dim = swa_v_head_dim elif ( - model_runner.hybrid_gdn_config is not None - or model_runner.kimi_linear_config is not None - or model_runner.linear_attn_model_spec is not None + hybrid_gdn_config(model_runner.model_config) is not None + or kimi_linear_config(model_runner.model_config) is not None + or linear_attn_model_spec(model_runner.model_config) is not None ): # For hybrid linear models, layer_id = 0 may not be full attention self.v_head_dim = model_runner.token_to_kv_pool.get_v_head_dim() diff --git a/python/sglang/srt/mem_cache/kv_cache_builder.py b/python/sglang/srt/mem_cache/kv_cache_builder.py index 09c1b74ff..3e6ac0639 100644 --- a/python/sglang/srt/mem_cache/kv_cache_builder.py +++ b/python/sglang/srt/mem_cache/kv_cache_builder.py @@ -23,6 +23,13 @@ class KVCacheBuildResult: from typing import TYPE_CHECKING +from sglang.srt.configs.hybrid_arch import ( + hybrid_gdn_config, + hybrid_lightning_config, + kimi_linear_config, + linear_attn_model_spec, + mamba2_config, +) from sglang.srt.configs.model_config import ModelImpl, is_deepseek_dsa from sglang.srt.environ import envs from sglang.srt.managers.mm_utils import init_mm_embedding_cache @@ -146,14 +153,14 @@ def build_kv_cache( # Hybrid memory pool is_hybrid_swa = tp_worker.is_hybrid_swa - _spec = tp_worker.model_runner.linear_attn_model_spec + _spec = linear_attn_model_spec(tp_worker.model_runner.model_config) _registry_needs_mamba = _spec.uses_mamba_radix_cache if _spec is not None else False is_hybrid_ssm = ( - tp_worker.model_runner.hybrid_gdn_config is not None - or tp_worker.model_runner.mamba2_config is not None + hybrid_gdn_config(tp_worker.model_runner.model_config) is not None + or mamba2_config(tp_worker.model_runner.model_config) is not None or _registry_needs_mamba - or tp_worker.model_runner.kimi_linear_config is not None - or tp_worker.model_runner.hybrid_lightning_config is not None + or kimi_linear_config(tp_worker.model_runner.model_config) is not None + or hybrid_lightning_config(tp_worker.model_runner.model_config) is not None ) is_dsa = is_deepseek_dsa(model_config.hf_config) diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index a30b31c94..3104b7fcf 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -37,6 +37,7 @@ from typing import TYPE_CHECKING, Dict, List, Optional, Set, Tuple, Union import torch from sglang.kernels.ops.attention.position import compute_position_triton +from sglang.srt.configs.hybrid_arch import mambaish_config from sglang.srt.environ import envs from sglang.srt.kv_canary.req_to_expected_token_ids_manager import ( compute_req_all_ids_info, @@ -850,7 +851,7 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin): if model_runner.use_ngram_embedding: ret._init_ngram_embedding_info(batch, device) - if model_runner.model_is_mrope: + if model_runner.model_config.model_is_mrope: if ( ret.spec_info is not None and getattr(ret.spec_info, "positions", None) is not None @@ -1229,7 +1230,7 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin): # Mamba-hybrid families need the fabricated-row idle conversion # below; this includes their MTP draft workers, whose mamba-less # "*E" pattern makes mambaish_config return None. - hybrid_ssm = model_runner.mambaish_config is not None or ( + hybrid_ssm = mambaish_config(model_runner.model_config) is not None or ( model_runner.is_draft_worker and getattr( model_runner.model_config.hf_config, diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index a6b95349b..18c1cb993 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -25,31 +25,12 @@ import threading import time from collections import defaultdict from dataclasses import dataclass -from typing import Any, Optional, Union +from typing import Optional, Union import torch import torch.distributed as dist -from sglang.srt.configs import ( - BailingHybridConfig, - FalconH1Config, - GraniteMoeHybridConfig, - InternS2PreviewConfig, - JetNemotronConfig, - JetVLMConfig, - KimiLinearConfig, - Lfm2Config, - Lfm2MoeConfig, - Lfm2VlConfig, - NemotronH_Nano_VL_V2_Config, - NemotronHConfig, - Qwen3_5Config, - Qwen3_5MoeConfig, - Qwen3NextConfig, - ZayaConfig, -) from sglang.srt.configs.device_config import DeviceConfig -from sglang.srt.configs.linear_attn_model_registry import get_linear_attn_config from sglang.srt.configs.load_config import LoadConfig, LoadFormat from sglang.srt.configs.model_config import ( AttentionArch, @@ -255,8 +236,6 @@ UNBALANCED_MODEL_LOADING_TIMEOUT_S = 480 # leave more time for post data proces logger = logging.getLogger(__name__) -_UNSET: Any = object() - @dataclass class ModelRunnerOutput: @@ -334,17 +313,9 @@ class ModelRunner(ModelRunnerKVCacheMixin): self.req_to_token_pool = req_to_token_pool self.token_to_kv_pool_allocator = token_to_kv_pool_allocator self.is_hybrid_swa = model_config.is_hybrid_swa - self.is_hybrid_swa_compress = getattr( - model_config, "is_hybrid_swa_compress", False - ) + self.is_hybrid_swa_compress = model_config.is_hybrid_swa_compress self.use_mla_backend = self.model_config.attention_arch == AttentionArch.MLA self.attention_chunk_size = model_config.attention_chunk_size - rope_scaling = getattr( - model_config.hf_text_config, "rope_parameters", None - ) or getattr(model_config.hf_text_config, "rope_scaling", {}) - self.model_is_mrope = ( - rope_scaling is not None and "mrope_section" in rope_scaling - ) self.enable_elastic_ep = server_args.elastic_ep_backend is not None self.forward_pass_id = 0 self.init_new_workspace = False @@ -535,8 +506,6 @@ class ModelRunner(ModelRunnerKVCacheMixin): # For hisparse (must be set before initialize() so CUDA graph capture can see it) self.hisparse_coordinator = None - self._linear_attn_registry_cache: Any = _UNSET - # Load model weights and configure self.initialize() self.check_quantized_moe_compatibility() @@ -1796,69 +1765,6 @@ class ModelRunner(ModelRunnerKVCacheMixin): return result - @property - def qwen3_next_config(self): - config = self.model_config.hf_config - if isinstance(config, Qwen3NextConfig): - return config - return None - - @property - def hybrid_lightning_config(self): - config = self.model_config.hf_config - if isinstance(config, BailingHybridConfig): - return config - return None - - @property - def hybrid_gdn_config(self): - config = self.model_config.hf_config.get_text_config() - if isinstance( - config, - Qwen3NextConfig - | Qwen3_5Config - | Qwen3_5MoeConfig - | InternS2PreviewConfig - | JetNemotronConfig - | JetVLMConfig, - ): - return config - return None - - @property - def mamba2_config(self): - config = self.model_config.hf_config - if isinstance(config, NemotronHConfig) and self.is_draft_worker: - # NemotronH MTP draft models have no Mamba layers (pattern like "*E") - # so they shouldn't use HybridLinearAttnBackend - pattern = getattr(config, "mtp_hybrid_override_pattern", None) - if pattern is not None and "M" not in pattern: - return None - if isinstance( - config, - FalconH1Config - | NemotronHConfig - | Lfm2Config - | Lfm2MoeConfig - | Lfm2VlConfig - | ZayaConfig, - ): - return config - if isinstance(config, NemotronH_Nano_VL_V2_Config): - return config.llm_config - - if isinstance(config, GraniteMoeHybridConfig): - has_mamba = any( - layer_type == "mamba" - for layer_type in getattr(config, "layer_types", []) - ) - if not has_mamba: - return None - else: - return config - - return None - @property def effective_max_total_num_tokens(self): """Return the max token pool size considering hybrid swa settings.""" @@ -1867,38 +1773,6 @@ class ModelRunner(ModelRunnerKVCacheMixin): else: return self.max_total_num_tokens - @property - def kimi_linear_config(self): - config = self.model_config.hf_config - if isinstance(config, KimiLinearConfig): - return config - return None - - def _get_linear_attn_registry_result(self): - if self._linear_attn_registry_cache is _UNSET: - self._linear_attn_registry_cache = get_linear_attn_config( - self.model_config.hf_config - ) - return self._linear_attn_registry_cache - - @property - def linear_attn_model_spec(self): - result = self._get_linear_attn_registry_result() - return result[0] if result else None - - @property - def mambaish_config(self): - existing = ( - self.mamba2_config - or self.hybrid_gdn_config - or self.kimi_linear_config - or self.hybrid_lightning_config - ) - if existing: - return existing - result = self._get_linear_attn_registry_result() - return result[1] if result else None - def _record_kv_cache_dtype(self, resolved: str) -> None: # Load-time resolution transition: the weight-resolved kv-cache dtype # is declared into the flags tier; the dual-apply inside the helper diff --git a/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py b/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py index 8c9bff1e2..05fd92f05 100644 --- a/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py +++ b/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py @@ -6,6 +6,10 @@ from typing import TYPE_CHECKING, Optional import torch +from sglang.srt.configs.hybrid_arch import ( + hybrid_gdn_config, + mambaish_config, +) from sglang.srt.configs.model_config import ( get_dsa_index_head_dim, get_minimax_sparse_attention_config, @@ -113,7 +117,10 @@ class ModelRunnerKVCacheMixin: ) slack_gb = pre_model_load_memory * (1 - self.mem_fraction_static) - if self.mambaish_config is not None and self.post_capture_kv_active: + if ( + mambaish_config(self.model_config) is not None + and self.post_capture_kv_active + ): # Mamba state is a fixed pre-capture allocation, so it can't ride the ~0 post-capture slack. slack_gb = max( slack_gb, @@ -123,7 +130,7 @@ class ModelRunnerKVCacheMixin: / 1024, ) rest_memory = available_gpu_memory - slack_gb - if self.mambaish_config is not None: + if mambaish_config(self.model_config) is not None: rest_memory = self.handle_max_mamba_cache(rest_memory) # Loaded weights (target + draft) can exceed the static budget @@ -147,7 +154,7 @@ class ModelRunnerKVCacheMixin: return int(rest_memory * (1 << 30)) # return in bytes def handle_max_mamba_cache(self: ModelRunner, total_rest_memory): - config = self.mambaish_config + config = mambaish_config(self.model_config) server_args = self.server_args assert config is not None @@ -341,19 +348,21 @@ class ModelRunnerKVCacheMixin: unsupported_pool_family = None if is_dsv4_model: unsupported_pool_family = "DeepSeekV4TokenToKVPool" - elif current_platform.is_out_of_tree() and not self.mambaish_config: + elif current_platform.is_out_of_tree() and not mambaish_config( + self.model_config + ): unsupported_pool_family = "out-of-tree platform KV pool" - elif ( - self.server_args.attention_backend == "ascend" and not self.mambaish_config + elif self.server_args.attention_backend == "ascend" and not mambaish_config( + self.model_config ): unsupported_pool_family = "NPU/Ascend KV pool" elif self.use_mla_backend and is_dsa_model: unsupported_pool_family = "DSA/MLA KV pool" - elif self.use_mla_backend and not self.mambaish_config: + elif self.use_mla_backend and not mambaish_config(self.model_config): unsupported_pool_family = "MLA KV pool" elif self.is_hybrid_swa: unsupported_pool_family = "SWA KV pool" - elif self.mambaish_config: + elif mambaish_config(self.model_config): unsupported_pool_family = "hybrid linear/Mamba KV pool" elif is_float4_e2m1fn_x2(self.kv_cache_dtype): unsupported_pool_family = "FP4 MHA KV pool" @@ -401,7 +410,7 @@ class ModelRunnerKVCacheMixin: decode_max_bs, running_requests, ) - if eager_decode_gap or self.mambaish_config is not None: + if eager_decode_gap or mambaish_config(self.model_config) is not None: headroom_gb = max( headroom_gb, self.server_args.mamba_pre_capture_reserve_mb( @@ -460,7 +469,7 @@ class ModelRunnerKVCacheMixin: allocator.""" from sglang.srt.mem_cache.unified_memory_pool import init_unified_mamba_pools - config = self.mambaish_config + config = mambaish_config(self.model_config) assert config is not None assert ( not self.use_mla_backend @@ -616,7 +625,7 @@ class ModelRunnerKVCacheMixin: and self.server_args.disaggregation_mode == "null" and self.req_to_token_pool is None ): - if self.mambaish_config is not None: + if mambaish_config(self.model_config) is not None: self._init_unified_mamba_pools(max_num_reqs) return if self.is_hybrid_swa and not is_deepseek_v4(self.model_config.hf_config): @@ -645,7 +654,7 @@ class ModelRunnerKVCacheMixin: # Extra slots for pre-allocated requests pre_alloc_size = self.server_args.disaggregation_decode_extra_slots - if config := self.mambaish_config: + if config := mambaish_config(self.model_config): self.req_to_token_pool = HybridMambaDecodeReqToTokenPool( size=max_num_reqs, max_context_len=self.model_config.context_len @@ -677,7 +686,7 @@ class ModelRunnerKVCacheMixin: enable_memory_saver=self.server_args.enable_memory_saver, pre_alloc_size=pre_alloc_size, ) - elif config := self.mambaish_config: + elif config := mambaish_config(self.model_config): self.req_to_token_pool = HybridReqToTokenPool( size=max_num_reqs, mamba_size=self.server_args.max_mamba_cache_size, @@ -821,7 +830,9 @@ class ModelRunnerKVCacheMixin: self.server_args.max_speculative_num_draft_tokens or 0 ), ) - elif current_platform.is_out_of_tree() and not self.mambaish_config: + elif current_platform.is_out_of_tree() and not mambaish_config( + self.model_config + ): if self.use_mla_backend and is_dsa_model: PoolCls = current_platform.get_dsa_kv_pool_cls() self.token_to_kv_pool = PoolCls( @@ -871,8 +882,8 @@ class ModelRunnerKVCacheMixin: start_layer=self.start_layer, end_layer=self.end_layer, ) - elif ( - self.server_args.attention_backend == "ascend" and not self.mambaish_config + elif self.server_args.attention_backend == "ascend" and not mambaish_config( + self.model_config ): if self.is_hybrid_swa: from sglang.srt.hardware_backend.npu.memory_pool_npu import ( @@ -987,7 +998,7 @@ class ModelRunnerKVCacheMixin: index_head_dim=get_dsa_index_head_dim(self.model_config.hf_config), **pool_kwargs, ) - elif self.use_mla_backend and not self.mambaish_config: + elif self.use_mla_backend and not mambaish_config(self.model_config): assert not is_dsa_model if is_float4_e2m1fn_x2(self.kv_cache_dtype): self.token_to_kv_pool = MLATokenToKVPoolFP4( @@ -1074,7 +1085,7 @@ class ModelRunnerKVCacheMixin: start_layer=self.start_layer, end_layer=self.end_layer, ) - elif config := self.mambaish_config: + elif config := mambaish_config(self.model_config): extra_args = {} if self.use_mla_backend: extra_args = { @@ -1178,7 +1189,7 @@ class ModelRunnerKVCacheMixin: elif _is_npu and ( self.server_args.attention_backend == "ascend" or is_dsv4_model - or self.hybrid_gdn_config is not None + or hybrid_gdn_config(self.model_config) is not None ): if self.is_hybrid_swa: # DSV4 on NPU: SWA allocator subclass that also drives the @@ -1359,7 +1370,7 @@ class ModelRunnerKVCacheMixin: requested_per_worker = None max_num_reqs = min(estimated, token_capacity // 2) - if self.mambaish_config is not None: + if mambaish_config(self.model_config) is not None: ratio = self._calculate_mamba_ratio() max_num_reqs = min( max_num_reqs, self.server_args.max_mamba_cache_size // ratio diff --git a/python/sglang/srt/model_executor/pool_configurator.py b/python/sglang/srt/model_executor/pool_configurator.py index 18a8cb272..145dc1f82 100644 --- a/python/sglang/srt/model_executor/pool_configurator.py +++ b/python/sglang/srt/model_executor/pool_configurator.py @@ -19,6 +19,7 @@ from typing import TYPE_CHECKING, Optional import torch +from sglang.srt.configs.hybrid_arch import mambaish_config from sglang.srt.configs.model_config import ( get_dsa_index_head_dim, get_minimax_sparse_attention_config, @@ -124,7 +125,7 @@ class DefaultPoolConfigurator(MemoryPoolConfigurator): def __init__(self, mr: ModelRunner): # Determine effective number of layers for KV cache - if mambaish := mr.mambaish_config: + if mambaish := mambaish_config(mr.model_config): effective_layer_ids = [ i for i in mambaish.full_attention_layer_ids diff --git a/python/sglang/srt/speculative/dflash_worker_v2.py b/python/sglang/srt/speculative/dflash_worker_v2.py index d9c82b360..06015f0d6 100644 --- a/python/sglang/srt/speculative/dflash_worker_v2.py +++ b/python/sglang/srt/speculative/dflash_worker_v2.py @@ -9,6 +9,7 @@ from sglang.kernels.ops.speculative.dflash import ( _compute_dflash_accept_bonus_triton_unchecked, _prepare_dflash_draft_block_unchecked, ) +from sglang.srt.configs.hybrid_arch import mambaish_config from sglang.srt.distributed import get_tp_group from sglang.srt.managers.schedule_batch import ScheduleBatch from sglang.srt.managers.scheduler import GenerationBatchResult @@ -277,12 +278,11 @@ class DFlashWorkerV2(BaseSpecWorker): def init_attention_backends(self): self._draft_worker.init_attention_backends() - self._need_mamba_verify_commit = ( - self.model_runner.mambaish_config is not None - and hasattr( - self.model_runner.attn_backend, - "update_mamba_state_after_mtp_verify", - ) + self._need_mamba_verify_commit = mambaish_config( + self.model_runner.model_config + ) is not None and hasattr( + self.model_runner.attn_backend, + "update_mamba_state_after_mtp_verify", ) def init_cuda_graphs(self): diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index 0ce9dd001..f853c0980 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/python/sglang/srt/speculative/eagle_worker_v2.py @@ -1555,7 +1555,7 @@ class EAGLEWorkerV2(BaseSpecWorker): ) if ( _is_npu - and self._target_worker.model_runner.model_is_mrope + and self._target_worker.model_runner.model_config.model_is_mrope and batch.spec_info is not None and getattr(batch.spec_info, "positions", None) is not None and not batch.forward_mode.is_idle() diff --git a/python/sglang/srt/speculative/spec_utils.py b/python/sglang/srt/speculative/spec_utils.py index d170dc634..f83514673 100644 --- a/python/sglang/srt/speculative/spec_utils.py +++ b/python/sglang/srt/speculative/spec_utils.py @@ -37,6 +37,7 @@ from sglang.kernels.ops.speculative.cache_locs import ( from sglang.kernels.ops.speculative.eagle import ( fill_accept_out_cache_loc_func as fill_accept_out_cache_loc_func, ) +from sglang.srt.configs.hybrid_arch import mambaish_config from sglang.srt.distributed.parallel_state import ( GroupCoordinator, patch_tensor_parallel_group, @@ -649,7 +650,7 @@ def commit_mamba_states_after_verify( commit hook. """ model_runner = target_worker.model_runner - if model_runner.mambaish_config is None: + if mambaish_config(model_runner.model_config) is None: return attn_backend = model_runner.attn_backend if not hasattr(attn_backend, "update_mamba_state_after_mtp_verify"): diff --git a/test/registered/unit/model_executor/test_pool_configurator.py b/test/registered/unit/model_executor/test_pool_configurator.py index 12088beb6..f8921facd 100644 --- a/test/registered/unit/model_executor/test_pool_configurator.py +++ b/test/registered/unit/model_executor/test_pool_configurator.py @@ -94,6 +94,7 @@ def _make_model_runner( mc.get_num_kv_heads = lambda tp_size: num_kv_heads mc.get_swa_num_kv_heads = lambda tp_size: swa_num_kv_heads or num_kv_heads mc.hf_config = SimpleNamespace(architectures=["LlamaForCausalLM"]) + mc.hf_config.get_text_config = lambda: mc.hf_config mr.model_config = mc mr.kv_cache_dtype = "fake_bf16"