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