kvcache-ai/sglang
1
0
Fork 0
mirror of https://github.com/kvcache-ai/sglang.git synced 2026-07-13 02:46:57 +00:00
Code Issues 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