A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 16:00 UTC
[Misc] Human-readable `max-model-len` cli arg (#16181)
Signed-off-by: NickLucche <nlucches@redhat.com> Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
Nicolò Lucchesi committed
Apr 7, 2025 at 20:40 UTC
090c856d7681f65143fece96f9dfd555c4b7d59b
2 files changed
+85
-3
tests/engine/test_arg_utils.py
+37
-1
@@ -1,6 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
3
-from argparse import ArgumentTypeError
3
+from argparse import ArgumentError, ArgumentTypeError
4
5
import pytest
6
@@ -142,3 +142,39 @@ def test_composite_arg_parser(arg, expected, option):
142
else:
143
args = parser.parse_args([f"--{option}", arg])
144
assert getattr(args, option.replace("-", "_")) == expected
145
+
146
+
147
+def test_human_readable_model_len():
148
+ # `exit_on_error` disabled to test invalid values below
149
+ parser = EngineArgs.add_cli_args(
150
+ FlexibleArgumentParser(exit_on_error=False))
151
+
152
+ args = parser.parse_args([])
153
+ assert args.max_model_len is None
154
+
155
+ args = parser.parse_args(["--max-model-len", "1024"])
156
+ assert args.max_model_len == 1024
157
+
158
+ # Lower
159
+ args = parser.parse_args(["--max-model-len", "1m"])
160
+ assert args.max_model_len == 1_000_000
161
+ args = parser.parse_args(["--max-model-len", "10k"])
162
+ assert args.max_model_len == 10_000
163
+
164
+ # Capital
165
+ args = parser.parse_args(["--max-model-len", "3K"])
166
+ assert args.max_model_len == 1024 * 3
167
+ args = parser.parse_args(["--max-model-len", "10M"])
168
+ assert args.max_model_len == 2**20 * 10
169
+
170
+ # Decimal values
171
+ args = parser.parse_args(["--max-model-len", "10.2k"])
172
+ assert args.max_model_len == 10200
173
+ # ..truncated to the nearest int
174
+ args = parser.parse_args(["--max-model-len", "10.212345k"])
175
+ assert args.max_model_len == 10212
176
+
177
+ # Invalid (do not allow decimals with binary multipliers)
178
+ for invalid in ["1a", "pwd", "10.24", "1.23M"]:
179
+ with pytest.raises(ArgumentError):
180
+ args = parser.parse_args(["--max-model-len", invalid])
vllm/engine/arg_utils.py
+48
-2
@@ -3,6 +3,7 @@
3
import argparse
4
import dataclasses
5
import json
6
+import re
7
import threading
8
from dataclasses import dataclass
9
from typing import (TYPE_CHECKING, Any, Dict, List, Literal, Mapping, Optional,
@@ -368,10 +369,14 @@ class EngineArgs:
369
'data type. CUDA 11.8+ supports fp8 (=fp8_e4m3) and fp8_e5m2. '
370
'ROCm (AMD GPU) supports fp8 (=fp8_e4m3)')
371
parser.add_argument('--max-model-len',
371
- type=int,
372
+ type=human_readable_int,
373
default=EngineArgs.max_model_len,
374
help='Model context length. If unspecified, will '
374
- 'be automatically derived from the model config.')
375
+ 'be automatically derived from the model config. '
376
+ 'Supports k/m/g/K/M/G in human-readable format.\n'
377
+ 'Examples:\n'
378
+ '- 1k → 1000\n'
379
+ '- 1K → 1024\n')
380
parser.add_argument(
381
'--guided-decoding-backend',
382
type=str,
@@ -1740,6 +1745,47 @@ def _warn_or_fallback(feature_name: str) -> bool:
1745
return should_exit
1746
1747
1748
+def human_readable_int(value):
1749
+ """Parse human-readable integers like '1k', '2M', etc.
1750
+ Including decimal values with decimal multipliers.
1751
+
1752
+ Examples:
1753
+ - '1k' -> 1,000
1754
+ - '1K' -> 1,024
1755
+ - '25.6k' -> 25,600
1756
+ """
1757
+ value = value.strip()
1758
+ match = re.fullmatch(r'(\d+(?:\.\d+)?)([kKmMgGtT])', value)
1759
+ if match:
1760
+ decimal_multiplier = {
1761
+ 'k': 10**3,
1762
+ 'm': 10**6,
1763
+ 'g': 10**9,
1764
+ }
1765
+ binary_multiplier = {
1766
+ 'K': 2**10,
1767
+ 'M': 2**20,
1768
+ 'G': 2**30,
1769
+ }
1770
+
1771
+ number, suffix = match.groups()
1772
+ if suffix in decimal_multiplier:
1773
+ mult = decimal_multiplier[suffix]
1774
+ return int(float(number) * mult)
1775
+ elif suffix in binary_multiplier:
1776
+ mult = binary_multiplier[suffix]
1777
+ # Do not allow decimals with binary multipliers
1778
+ try:
1779
+ return int(number) * mult
1780
+ except ValueError as e:
1781
+ raise argparse.ArgumentTypeError("Decimals are not allowed " \
1782
+ f"with binary suffixes like {suffix}. Did you mean to use " \
1783
+ f"{number}{suffix.lower()} instead?") from e
1784
+
1785
+ # Regular plain number.
1786
+ return int(value)
1787
+
1788
+
1789
# These functions are used by sphinx to build the documentation
1790
def _engine_args_parser():
1791
return EngineArgs.add_cli_args(FlexibleArgumentParser())