chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

fix(config): Support setting Mamba state dtype via config file (#18532)

Browse Source 
Co-authored-by: 瑀澈 <yuche.lz@alibaba-inc.com>
This commit is contained in:
Zheng Li
2026-02-11 00:20:06 +08:00
committed by GitHub
parent d0d387dea1
commit 44603764d6
8 changed files with 103 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
python/sglang/srt/configs/falcon_h1.py
View File

@@ -18,7 +18,11 @@
from transformers.configuration_utils import PretrainedConfig
from transformers.utils import logging
from sglang.srt.configs.mamba_utils import Mamba2CacheParams, Mamba2StateShape
from sglang.srt.configs.mamba_utils import (
Mamba2CacheParams,
Mamba2StateShape,
mamba2_state_dtype,
)
logger = logging.get_logger(__name__)
@@ -307,4 +311,6 @@ class FalconH1Config(PretrainedConfig):
state_size=self.mamba_d_state,
conv_kernel=self.mamba_d_conv,
)
return Mamba2CacheParams(shape=shape, layers=self.linear_layer_ids)
return Mamba2CacheParams(
shape=shape, layers=self.linear_layer_ids, dtype=mamba2_state_dtype(self)
)

10
python/sglang/srt/configs/jet_nemotron.py
View File

@@ -3,7 +3,11 @@ from typing import Any
from transformers.configuration_utils import PretrainedConfig
from sglang.srt.configs.mamba_utils import Mamba2CacheParams, Mamba2StateShape
from sglang.srt.configs.mamba_utils import (
Mamba2CacheParams,
Mamba2StateShape,
mamba2_state_dtype,
)
@dataclass
@@ -71,4 +75,6 @@ class JetNemotronConfig(PretrainedConfig):
conv_kernel=jet_block_config.conv_size,
)
return Mamba2CacheParams(shape=shape, layers=self.linear_layer_ids)
return Mamba2CacheParams(
shape=shape, layers=self.linear_layer_ids, dtype=mamba2_state_dtype(self)
)

9
python/sglang/srt/configs/lfm2.py
View File

@@ -20,7 +20,11 @@ from transformers import CONFIG_MAPPING
from transformers import Lfm2Config as HFLfm2Config
from transformers.utils import logging
from sglang.srt.configs.mamba_utils import Mamba2CacheParams, Mamba2StateShape
from sglang.srt.configs.mamba_utils import (
Mamba2CacheParams,
Mamba2StateShape,
mamba2_state_dtype,
)
logger = logging.get_logger(__name__)
@@ -87,11 +91,10 @@ class Lfm2Config(HFLfm2Config):
conv_kernel=conv_kernel,
)
# Uses default mamba2_state_dtype() which reads SGLANG_MAMBA_CONV_DTYPE env var
# (defaults to bfloat16). Set SGLANG_MAMBA_CONV_DTYPE=float16 for fp16 inference.
return Mamba2CacheParams(
shape=shape,
layers=conv_layer_ids,
dtype=mamba2_state_dtype(self),
)

61
python/sglang/srt/configs/mamba_utils.py
View File

@@ -12,6 +12,7 @@
# limitations under the License.
"""Common config utils for mamba2 - NemotronH, FalconH1, Qwen3Next, LFM2, etc."""
import logging
from abc import ABC
from dataclasses import dataclass, field
from typing import List, Optional
@@ -22,6 +23,8 @@ import torch
from sglang.srt.distributed.utils import divide
from sglang.srt.environ import envs
logger = logging.getLogger(__name__)
def extra_groups_for_head_shards(ngroups: int, tp_size: int):
"""Compute the increase in group numbers to account for
@@ -41,20 +44,72 @@ class Mamba2StateDType:
temporal: torch.dtype
def mamba2_state_dtype() -> Mamba2StateDType:
def mamba2_state_dtype(config=None) -> Mamba2StateDType:
"""
Get mamba2 state dtype from config or environment variable.
Priority (from highest to lowest):
1. Environment variable SGLANG_MAMBA_SSM_DTYPE
2. Config file (config.mamba_ssm_dtype or config.text_config.mamba_ssm_dtype)
3. Default "float32"
Args:
config: Optional config object (PretrainedConfig). If provided, will read
mamba_ssm_dtype from it. For VL models, reads from text_config.
Returns:
Mamba2StateDType with conv and temporal dtypes
"""
dtype_map = {
"float32": torch.float32,
"bfloat16": torch.bfloat16,
"float16": torch.float16,
}
conv_dtype = dtype_map.get(envs.SGLANG_MAMBA_CONV_DTYPE.get(), torch.bfloat16)
ssm_dtype = dtype_map.get(envs.SGLANG_MAMBA_SSM_DTYPE.get(), torch.float32)
# Get SSM dtype: default -> config -> env var
ssm_dtype = torch.float32 # Step 1: Default value
# Step 2: Try to read from config
if config is not None:
config_dtype = None
if hasattr(config, "text_config") and hasattr(
config.text_config, "mamba_ssm_dtype"
):
# VL model: read from text_config
config_dtype = config.text_config.mamba_ssm_dtype
elif hasattr(config, "mamba_ssm_dtype"):
# Text model: read from root config
config_dtype = config.mamba_ssm_dtype
if config_dtype is not None:
if config_dtype not in dtype_map:
logger.warning(
f"Invalid mamba_ssm_dtype '{config_dtype}' in config. "
f"Must be one of {list(dtype_map.keys())}. Using default 'float32'."
)
else:
ssm_dtype = dtype_map[config_dtype]
# Step 3: Check environment variable, if not None, override
env_ssm_dtype = envs.SGLANG_MAMBA_SSM_DTYPE.get()
if env_ssm_dtype is not None:
if env_ssm_dtype not in dtype_map:
logger.warning(
f"Invalid mamba_ssm_dtype '{env_ssm_dtype}' from environment variable. "
f"Must be one of {list(dtype_map.keys())}. Using default 'float32'."
)
else:
ssm_dtype = dtype_map[env_ssm_dtype]
logger.info(f"Mamba2 state dtype: conv_dtype={conv_dtype}, ssm_dtype={ssm_dtype}")
return Mamba2StateDType(conv=conv_dtype, temporal=ssm_dtype)
@dataclass(kw_only=True, frozen=True)
class BaseLinearStateParams(ABC):
dtype: Mamba2StateDType = field(default_factory=mamba2_state_dtype)
dtype: Mamba2StateDType = field(default_factory=lambda: mamba2_state_dtype(None))
layers: list[int]
@property

10
python/sglang/srt/configs/nemotron_h.py
View File

@@ -19,7 +19,11 @@ import regex as re
from transformers.configuration_utils import PretrainedConfig
from transformers.utils import logging
from sglang.srt.configs.mamba_utils import Mamba2CacheParams, Mamba2StateShape
from sglang.srt.configs.mamba_utils import (
Mamba2CacheParams,
Mamba2StateShape,
mamba2_state_dtype,
)
logger = logging.get_logger(__name__)
@@ -305,4 +309,6 @@ class NemotronHConfig(PretrainedConfig):
conv_kernel=self.conv_kernel,
)
return Mamba2CacheParams(shape=shape, layers=self.mamba_layer_ids)
return Mamba2CacheParams(
shape=shape, layers=self.mamba_layer_ids, dtype=mamba2_state_dtype(self)
)

10
python/sglang/srt/configs/qwen3_next.py
View File

@@ -19,7 +19,11 @@ import enum
from transformers.configuration_utils import PretrainedConfig
from transformers.utils import logging
from sglang.srt.configs.mamba_utils import Mamba2CacheParams, Mamba2StateShape
from sglang.srt.configs.mamba_utils import (
Mamba2CacheParams,
Mamba2StateShape,
mamba2_state_dtype,
)
from sglang.srt.configs.update_config import adjust_tp_num_heads_if_necessary
from sglang.srt.utils import is_cpu
@@ -293,4 +297,6 @@ class Qwen3NextConfig(PretrainedConfig):
conv_kernel=self.linear_conv_kernel_dim,
)
return Mamba2CacheParams(shape=shape, layers=self.linear_layer_ids)
return Mamba2CacheParams(
shape=shape, layers=self.linear_layer_ids, dtype=mamba2_state_dtype(self)
)

2
python/sglang/srt/environ.py
View File

@@ -418,7 +418,7 @@ class Envs:
# Mamba
SGLANG_MAMBA_CONV_DTYPE = EnvStr("bfloat16")
SGLANG_MAMBA_SSM_DTYPE = EnvStr("float32")
SGLANG_MAMBA_SSM_DTYPE = EnvStr(None)
# Release & Resume Memory
SGLANG_MEMORY_SAVER_CUDA_GRAPH = EnvBool(False)

10
python/sglang/srt/server_args.py
View File

@@ -513,7 +513,7 @@ class ServerArgs:
# Mamba cache
max_mamba_cache_size: Optional[int] = None
mamba_ssm_dtype: str = "float32"
mamba_ssm_dtype: Optional[str] = None
mamba_full_memory_ratio: float = 0.9
mamba_scheduler_strategy: str = "auto"
mamba_track_interval: int = 256
@@ -2600,7 +2600,8 @@ class ServerArgs:
def _handle_environment_variables(self):
envs.SGLANG_ENABLE_TORCH_COMPILE.set("1" if self.enable_torch_compile else "0")
envs.SGLANG_MAMBA_SSM_DTYPE.set(self.mamba_ssm_dtype)
if self.mamba_ssm_dtype is not None:
envs.SGLANG_MAMBA_SSM_DTYPE.set(self.mamba_ssm_dtype)
envs.SGLANG_DISABLE_OUTLINES_DISK_CACHE.set(
"1" if self.disable_outlines_disk_cache else "0"
)
@@ -4130,9 +4131,10 @@ class ServerArgs:
parser.add_argument(
"--mamba-ssm-dtype",
type=str,
default=ServerArgs.mamba_ssm_dtype,
default=None,
choices=MAMBA_SSM_DTYPE_CHOICES,
help="The data type of the SSM states in mamba cache.",
help="The data type of the SSM states in mamba cache. "
"If not set, will be read from model config (mamba_ssm_dtype).",
)
parser.add_argument(
"--mamba-full-memory-ratio",