From 7ace64d1d88de1dfd3cfb6622f3fcbba4fe7eef8 Mon Sep 17 00:00:00 2001 From: Ke Bao Date: Fri, 23 Jan 2026 11:02:32 +0800 Subject: [PATCH] Update mamba env setting (#17566) --- python/sglang/srt/configs/mamba_utils.py | 10 ++----- python/sglang/srt/environ.py | 5 ++++ python/sglang/srt/server_args.py | 10 +++---- .../models/test_generation_models.py | 19 +++++++----- .../radix_cache/test_mamba_unittest.py | 29 ++++++++++--------- 5 files changed, 39 insertions(+), 34 deletions(-) diff --git a/python/sglang/srt/configs/mamba_utils.py b/python/sglang/srt/configs/mamba_utils.py index 9e64d7524..d953bcd3b 100644 --- a/python/sglang/srt/configs/mamba_utils.py +++ b/python/sglang/srt/configs/mamba_utils.py @@ -12,7 +12,6 @@ # limitations under the License. """Common config utils for mamba2 - NemotronH, FalconH1, Qwen3Next, LFM2, etc.""" -import os from abc import ABC from dataclasses import dataclass, field from typing import List, Optional @@ -21,6 +20,7 @@ import numpy as np import torch from sglang.srt.distributed.utils import divide +from sglang.srt.environ import envs def extra_groups_for_head_shards(ngroups: int, tp_size: int): @@ -47,12 +47,8 @@ def mamba2_state_dtype() -> Mamba2StateDType: "bfloat16": torch.bfloat16, "float16": torch.float16, } - conv_dtype = dtype_map.get( - os.environ.get("SGLANG_MAMBA_CONV_DTYPE", "bfloat16"), torch.bfloat16 - ) - ssm_dtype = dtype_map.get( - os.environ.get("SGLANG_MAMBA_SSM_DTYPE", "float32"), torch.float32 - ) + 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) return Mamba2StateDType(conv=conv_dtype, temporal=ssm_dtype) diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index acde01bcd..4bc71df40 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -174,6 +174,7 @@ class Envs: # Constrained Decoding (Grammar) SGLANG_GRAMMAR_POLL_INTERVAL = EnvFloat(0.005) SGLANG_GRAMMAR_MAX_POLL_ITERATIONS = EnvInt(10000) + SGLANG_DISABLE_OUTLINES_DISK_CACHE = EnvBool(False) # CuTe DSL GDN Decode SGLANG_USE_CUTEDSL_GDN_DECODE = EnvBool(False) @@ -408,6 +409,10 @@ class Envs: # MM splitting behavior control SGLANG_ENABLE_MM_SPLITTING = EnvBool(False) + # Mamba + SGLANG_MAMBA_CONV_DTYPE = EnvStr("bfloat16") + SGLANG_MAMBA_SSM_DTYPE = EnvStr("float32") + # Release & Resume Memory SGLANG_MEMORY_SAVER_CUDA_GRAPH = EnvBool(False) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 1e39fad37..aea0c65bf 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -2412,14 +2412,12 @@ class ServerArgs: self.enable_dynamic_batch_tokenizer = False def _handle_environment_variables(self): - os.environ["SGLANG_ENABLE_TORCH_COMPILE"] = ( - "1" if self.enable_torch_compile else "0" - ) - os.environ["SGLANG_MAMBA_SSM_DTYPE"] = self.mamba_ssm_dtype - os.environ["SGLANG_DISABLE_OUTLINES_DISK_CACHE"] = ( + envs.SGLANG_ENABLE_TORCH_COMPILE.set("1" if self.enable_torch_compile else "0") + 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" ) - os.environ["SGLANG_ENABLE_DETERMINISTIC_INFERENCE"] = ( + envs.SGLANG_ENABLE_DETERMINISTIC_INFERENCE.set( "1" if self.enable_deterministic_inference else "0" ) # Set the highest strict level for Kimi K2 tool calls diff --git a/test/registered/models/test_generation_models.py b/test/registered/models/test_generation_models.py index a180102ec..d1041ea0d 100644 --- a/test/registered/models/test_generation_models.py +++ b/test/registered/models/test_generation_models.py @@ -29,10 +29,12 @@ import dataclasses import multiprocessing as mp import os import unittest +from contextlib import nullcontext from typing import List, Optional import torch +from sglang.srt.environ import envs from sglang.srt.utils import is_hip from sglang.test.runners import ( DEFAULT_PROMPTS, @@ -115,6 +117,10 @@ ALL_MODELS = [ ), ] +MAMBA_MODEL_PATHS = [ + "LiquidAI/LFM2.5-1.2B-Instruct", +] + TORCH_DTYPES = [torch.float16] @@ -131,18 +137,17 @@ class TestGenerationModels(CustomTestCase): torch_dtype: torch.dtype, ) -> None: model_path = model_case.model_path - prefill_tolerance, decode_tolerance, rouge_l_tolerance = ( - model_case.prefill_tolerance, - model_case.decode_tolerance, - model_case.rouge_l_tolerance, - ) max_new_tokens = 32 # Set conv dtype for hybrid models to match inference dtype dtype_str = {torch.float16: "float16", torch.bfloat16: "bfloat16"}.get( torch_dtype, "bfloat16" ) - os.environ["SGLANG_MAMBA_CONV_DTYPE"] = dtype_str + + if model_case.model_path in MAMBA_MODEL_PATHS: + env_ctx = envs.SGLANG_MAMBA_CONV_DTYPE.override(dtype_str) + else: + env_ctx = nullcontext() with HFRunner( model_path, @@ -152,7 +157,7 @@ class TestGenerationModels(CustomTestCase): ) as hf_runner: hf_outputs = hf_runner.forward(prompts, max_new_tokens=max_new_tokens) - with SRTRunner( + with env_ctx, SRTRunner( model_path, tp_size=model_case.tp_size, torch_dtype=torch_dtype, diff --git a/test/registered/radix_cache/test_mamba_unittest.py b/test/registered/radix_cache/test_mamba_unittest.py index 7e66dadcf..4cc231095 100644 --- a/test/registered/radix_cache/test_mamba_unittest.py +++ b/test/registered/radix_cache/test_mamba_unittest.py @@ -1,9 +1,9 @@ -import os import unittest import torch from sglang.srt.configs.mamba_utils import Mamba2CacheParams, Mamba2StateShape +from sglang.srt.environ import envs from sglang.srt.managers.schedule_batch import Req from sglang.srt.mem_cache.allocator import TokenToKVPoolAllocator from sglang.srt.mem_cache.base_prefix_cache import ( @@ -85,8 +85,9 @@ class TestMamba(unittest.TestCase): state_size=128, conv_kernel=4, ) - os.environ["SGLANG_MAMBA_SSM_DTYPE"] = "bfloat16" - mamba2_cache_params = Mamba2CacheParams(shape=shape, layers=mamba_layers) + + with envs.SGLANG_MAMBA_SSM_DTYPE.override("bfloat16"): + mamba2_cache_params = Mamba2CacheParams(shape=shape, layers=mamba_layers) req_to_token_pool = HybridReqToTokenPool( size=max_num_reqs, @@ -159,17 +160,17 @@ class TestMamba(unittest.TestCase): mamba_layers = [ i for i in range(num_layers) if i not in full_attention_layer_ids ] - os.environ["SGLANG_MAMBA_SSM_DTYPE"] = "bfloat16" - shape = Mamba2StateShape.create( - tp_world_size=1, - intermediate_size=4096, - n_groups=16, - num_heads=32, - head_dim=128, - state_size=128, - conv_kernel=4, - ) - mamba2_cache_params = Mamba2CacheParams(shape=shape, layers=mamba_layers) + with envs.SGLANG_MAMBA_SSM_DTYPE.override("bfloat16"): + shape = Mamba2StateShape.create( + tp_world_size=1, + intermediate_size=4096, + n_groups=16, + num_heads=32, + head_dim=128, + state_size=128, + conv_kernel=4, + ) + mamba2_cache_params = Mamba2CacheParams(shape=shape, layers=mamba_layers) req_to_token_pool = HybridReqToTokenPool( size=max_num_reqs,