feat: support human-readable suffixes (25.6k, 1M, 1Mi) for token CLI (#20577)

This commit is contained in:
shuwenn
2026-03-17 11:05:33 +08:00
committed by GitHub
parent 3cafb02e5b
commit 515b3a323d
2 changed files with 49 additions and 6 deletions

View File

@@ -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",

View File

@@ -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: