Update mamba env setting (#17566)

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"
-        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,