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

Update mamba env setting (#17566)

Browse Source
This commit is contained in:
Ke Bao
2026-01-23 11:02:32 +08:00
committed by GitHub
parent 62e6a749b0
commit 7ace64d1d8
5 changed files with 39 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -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)

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

@@ -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)

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

@@ -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

19
test/registered/models/test_generation_models.py
View File

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

29
test/registered/radix_cache/test_mamba_unittest.py
View File

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