diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index a51168c71..16e21b73c 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -43,6 +43,7 @@ from sglang.srt.utils.common import ( get_device_sm, get_int_env_var, get_quantization_config, + human_readable_int, is_blackwell_supported, is_cpu, is_cuda, @@ -3384,9 +3385,10 @@ class ServerArgs: ) parser.add_argument( "--context-length", - type=int, + type=human_readable_int, default=ServerArgs.context_length, - help="The model's maximum context length. Defaults to None (will use the value from the model's config.json instead).", + help="The model's maximum context length. Defaults to None (will use the value from the model's config.json instead)." + + f"\n\n{human_readable_int.__doc__}", ) parser.add_argument( "--is-embedding", @@ -3613,10 +3615,11 @@ class ServerArgs: ) parser.add_argument( "--max-total-tokens", - type=int, + type=human_readable_int, default=ServerArgs.max_total_tokens, help="The maximum number of tokens in the memory pool. If not specified, it will be automatically calculated based on the memory usage fraction. " - "This option is typically used for development and debugging purposes.", + "This option is typically used for development and debugging purposes." + + f"\n\n{human_readable_int.__doc__}", ) parser.add_argument( "--chunked-prefill-size", @@ -3638,9 +3641,10 @@ class ServerArgs: ) parser.add_argument( "--max-prefill-tokens", - type=int, + type=human_readable_int, default=ServerArgs.max_prefill_tokens, - help="The maximum number of tokens in a prefill batch. The real bound will be the maximum of this value and the model's maximum context length.", + help="The maximum number of tokens in a prefill batch. The real bound will be the maximum of this value and the model's maximum context length." + + f"\n\n{human_readable_int.__doc__}", ) parser.add_argument( "--schedule-policy", diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index aceb0830a..cd504fc77 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py @@ -47,6 +47,7 @@ import warnings from collections import OrderedDict, defaultdict from contextlib import contextmanager from dataclasses import dataclass +from decimal import Decimal from functools import lru_cache, partial from importlib.metadata import PackageNotFoundError, version from importlib.util import find_spec @@ -2308,6 +2309,44 @@ def nullable_str(val: str): return val +def human_readable_int(value: str) -> int: + """Supports standard SI suffixes (k, M, G, T) and IEC suffixes + (Ki, Mi, Gi, Ti). Suffixes are case-sensitive. + + Decimals are allowed for SI suffixes only. + + Examples: + '1k' -> 1000 '1M' -> 1000000 '25.6k' -> 25600 + '1Ki' -> 1024 '1Mi' -> 1048576 + """ + value = value.strip() + + si_multiplier = {"k": 10**3, "M": 10**6, "G": 10**9, "T": 10**12} + iec_multiplier = {"Ki": 2**10, "Mi": 2**20, "Gi": 2**30, "Ti": 2**40} + + match = re.fullmatch(r"(\d+(?:\.\d+)?)(Ki|Mi|Gi|Ti|k|M|G|T)", value) + if match: + number, suffix = match.groups() + if suffix in iec_multiplier: + if "." in number: + raise argparse.ArgumentTypeError( + f"Decimals are not allowed with IEC suffixes like '{suffix}'. " + f"Use an integer IEC value such as '{int(Decimal(number))}{suffix}', " + f"or an SI value such as '{number}{suffix[0]}'." + ) + return int(number) * iec_multiplier[suffix] + return int(Decimal(number) * si_multiplier[suffix]) + + try: + return int(value) + except ValueError: + raise argparse.ArgumentTypeError( + f"Invalid integer value: '{value}'. " + "Use a plain integer, SI suffixes (1k, 1M), or IEC suffixes (1Ki, 1Mi). " + "Suffixes are case-sensitive." + ) + + def pyspy_dump_schedulers(): """py-spy dump on all scheduler in a local node.""" try: