[ServerArgs] allow --mamba-ssm-dtype extend (#12481)

This commit is contained in:
Hanming Lu
2025-11-02 11:50:04 -08:00
committed by GitHub
parent 819fc59123
commit 66fb9b1307

View File

@@ -162,6 +162,8 @@ MOE_RUNNER_BACKEND_CHOICES = [
"cutlass",
]
MAMBA_SSM_DTYPE_CHOICES = ["float32", "bfloat16"]
# Allow external code to add more choices
def add_load_format_choices(choices):
@@ -200,6 +202,10 @@ def add_radix_eviction_policy_choices(choices):
RADIX_EVICTION_POLICY_CHOICES.extend(choices)
def add_mamba_ssm_dtype_choices(choices):
MAMBA_SSM_DTYPE_CHOICES.extend(choices)
@dataclasses.dataclass
class ServerArgs:
"""
@@ -2902,7 +2908,7 @@ class ServerArgs:
"--mamba-ssm-dtype",
type=str,
default=ServerArgs.mamba_ssm_dtype,
choices=["float32", "bfloat16"],
choices=MAMBA_SSM_DTYPE_CHOICES,
help="The data type of the SSM states in mamba cache.",
)
parser.add_argument(