Update mamba env setting (#17566)

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(
+    conv_dtype = dtype_map.get(envs.SGLANG_MAMBA_CONV_DTYPE.get(), torch.bfloat16)
-        os.environ.get("SGLANG_MAMBA_CONV_DTYPE", "bfloat16"), torch.bfloat16
+    ssm_dtype = dtype_map.get(envs.SGLANG_MAMBA_SSM_DTYPE.get(), torch.float32)
-    )
-    ssm_dtype = dtype_map.get(
-        os.environ.get("SGLANG_MAMBA_SSM_DTYPE", "float32"), torch.float32
-    )
     return Mamba2StateDType(conv=conv_dtype, temporal=ssm_dtype)
 
 

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)
 

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"] = (
+        envs.SGLANG_ENABLE_TORCH_COMPILE.set("1" if self.enable_torch_compile else "0")
-            "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(
-        os.environ["SGLANG_MAMBA_SSM_DTYPE"] = self.mamba_ssm_dtype
-        os.environ["SGLANG_DISABLE_OUTLINES_DISK_CACHE"] = (
             "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

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,

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"
+        with envs.SGLANG_MAMBA_SSM_DTYPE.override("bfloat16"):
-        shape = Mamba2StateShape.create(
+            shape = Mamba2StateShape.create(
-            tp_world_size=1,
+                tp_world_size=1,
-            intermediate_size=4096,
+                intermediate_size=4096,
-            n_groups=16,
+                n_groups=16,
-            num_heads=32,
+                num_heads=32,
-            head_dim=128,
+                head_dim=128,
-            state_size=128,
+                state_size=128,
-            conv_kernel=4,
+                conv_kernel=4,
-        )
+            )
-        mamba2_cache_params = Mamba2CacheParams(shape=shape, layers=mamba_layers)
+            mamba2_cache_params = Mamba2CacheParams(shape=shape, layers=mamba_layers)
 
         req_to_token_pool = HybridReqToTokenPool(
             size=max_num_reqs,