[ServerArgs] allow --mamba-ssm-dtype extend (#12481)
This commit is contained in:
@@ -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(
|
||||||
|
|||||||
Reference in New Issue
Block a user