@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 8ee90c83f8
Fork 0

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 11, 00:00 UTC

Add `--max-model-len auto` to auto-fit context to available memory (#29431)

Signed-off-by: mgoin <mgoin64@gmail.com>

Michael Goin committed Dec 24, 2025 at 00:37 UTC 8ee90c83f8e8b53f35ee3df3a86377ea5a587eea
7 files changed +315 -33
tests/engine/test_arg_utils.py
+10
@@ -511,6 +511,16 @@ def test_human_readable_model_len():
511 args = parser.parse_args(["--max-model-len", "10.2123451234567t"])
512 assert args.max_model_len == 10212345123456
513
514 + # Special value -1 for auto-fit to GPU memory
515 + args = parser.parse_args(["--max-model-len", "-1"])
516 + assert args.max_model_len == -1
517 +
518 + # 'auto' is an alias for -1
519 + args = parser.parse_args(["--max-model-len", "auto"])
520 + assert args.max_model_len == -1
521 + args = parser.parse_args(["--max-model-len", "AUTO"])
522 + assert args.max_model_len == -1
523 +
524 # Invalid (do not allow decimals with binary multipliers)
525 for invalid in ["1a", "pwd", "10.24", "1.23M", "1.22T"]:
526 with pytest.raises(ArgumentError):
tests/v1/core/test_kv_cache_utils.py
+57
@@ -1798,3 +1798,60 @@ def test_request_with_prompt_embeds_and_mm_inputs(hash_fn: Callable[[Any], bytes
1798 )
1799 )
1800 assert block_hashes[1] == expected_hash2
1801 +
1802 +
1803 +def test_auto_fit_max_model_len():
1804 + """Test that max_model_len=-1 auto-fits to available GPU memory."""
1805 + # Create config with original_max_model_len=-1 to trigger auto-fit
1806 + model_config = ModelConfig(max_model_len=1024)
1807 + # Simulate the user passing -1 by setting original_max_model_len
1808 + model_config.original_max_model_len = -1
1809 + vllm_config = VllmConfig(model_config=model_config)
1810 +
1811 + mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2 # 16KB per block per layer
1812 + kv_cache_specs = {
1813 + "layer_1": new_kv_cache_spec(),
1814 + "layer_2": new_kv_cache_spec(),
1815 + }
1816 +
1817 + # With enough memory, max_model_len stays at the derived max
1818 + large_available_memory = mem_per_block_per_layer * 2 * 1024 # plenty of memory
1819 + _kv_cache_configs = get_kv_cache_configs(
1820 + vllm_config, [kv_cache_specs], [large_available_memory]
1821 + )
1822 + assert vllm_config.model_config.max_model_len == 1024
1823 +
1824 + # Reset for next test
1825 + model_config = ModelConfig(max_model_len=1024)
1826 + model_config.original_max_model_len = -1
1827 + vllm_config = VllmConfig(model_config=model_config)
1828 +
1829 + # With limited memory, max_model_len should be reduced
1830 + # Need memory for at least max_model_len tokens
1831 + # 32 blocks worth of memory for 2 layers = can fit 32*16=512 tokens
1832 + limited_memory = mem_per_block_per_layer * 2 * 32
1833 + _kv_cache_configs = get_kv_cache_configs(
1834 + vllm_config, [kv_cache_specs], [limited_memory]
1835 + )
1836 + # Should be reduced to fit in memory
1837 + assert vllm_config.model_config.max_model_len < 1024
1838 + assert vllm_config.model_config.max_model_len > 0
1839 +
1840 +
1841 +def test_auto_fit_max_model_len_not_triggered():
1842 + """Test that auto-fit is not triggered when original_max_model_len is not -1."""
1843 + model_config = ModelConfig(max_model_len=16)
1844 + # original_max_model_len should be None by default, not -1
1845 + vllm_config = VllmConfig(model_config=model_config)
1846 +
1847 + mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2
1848 + kv_cache_specs = {
1849 + "layer_1": new_kv_cache_spec(),
1850 + "layer_2": new_kv_cache_spec(),
1851 + }
1852 +
1853 + # This should work normally without auto-fit
1854 + _kv_cache_configs = get_kv_cache_configs(
1855 + vllm_config, [kv_cache_specs], [mem_per_block_per_layer * 2 * 32]
1856 + )
1857 + assert vllm_config.model_config.max_model_len == 16
vllm/config/model.py
+8 -4
@@ -172,7 +172,10 @@ class ModelConfig:
172 format. Examples:\n
173 - 1k -> 1000\n
174 - 1K -> 1024\n
175 - - 25.6k -> 25,600"""
175 + - 25.6k -> 25,600\n
176 + - -1 or 'auto' -> Automatically choose the maximum model length that fits in
177 + GPU memory. This will use the model's maximum context length if it fits,
178 + otherwise it will find the largest length that can be accommodated."""
179 spec_target_max_model_len: int | None = None
180 """Specify the maximum length for spec decoding draft models."""
181 quantization: QuantizationMethods | str | None = None
@@ -2151,9 +2154,10 @@ def _get_and_verify_max_len(
2154 if encoder_config and "max_seq_length" in encoder_config:
2155 derived_max_model_len = encoder_config["max_seq_length"]
2156
2154 - # If the user didn't specify `max_model_len`, then use that derived from
2155 - # the model config as a default value.
2156 - if max_model_len is None:
2157 + # If the user didn't specify `max_model_len` or specified -1 (auto-fit),
2158 + # then use that derived from the model config as a default value.
2159 + # When -1 is specified, the engine will later auto-fit to available memory.
2160 + if max_model_len is None or max_model_len == -1:
2161 # For LongRoPE, default to original_max_position_embeddings to avoid
2162 # performance degradation for shorter sequences
2163 if rope_parameters is not None and any(
vllm/engine/arg_utils.py
+7
@@ -2045,13 +2045,20 @@ def _raise_unsupported_error(feature_name: str):
2045 def human_readable_int(value):
2046 """Parse human-readable integers like '1k', '2M', etc.
2047 Including decimal values with decimal multipliers.
2048 + Also accepts -1 or 'auto' as a special value for auto-detection.
2049
2050 Examples:
2051 - '1k' -> 1,000
2052 - '1K' -> 1,024
2053 - '25.6k' -> 25,600
2054 + - '-1' or 'auto' -> -1 (special value for auto-detection)
2055 """
2056 value = value.strip()
2057 +
2058 + # Handle -1 or 'auto' as a special value for auto-detection
2059 + if value == "-1" or value.lower() == "auto":
2060 + return -1
2061 +
2062 match = re.fullmatch(r"(\d+(?:\.\d+)?)([kKmMgGtT])", value)
2063 if match:
2064 decimal_multiplier = {
vllm/v1/core/kv_cache_utils.py
+209 -29
@@ -624,6 +624,9 @@ def estimate_max_model_len(
624 Estimates the maximum model length that can fit in the available memory
625 using binary search.
626
627 + This function temporarily modifies max_model_len during estimation but
628 + restores the original value before returning, ensuring no side effects.
629 +
630 Args:
631 vllm_config: The global VllmConfig
632 kv_cache_spec: The kv cache spec of each attention layer in the model
@@ -632,33 +635,38 @@ def estimate_max_model_len(
635 Returns:
636 The estimated maximum model length that can fit in the available memory.
637 """
638 + # Save the original max_model_len to restore after estimation
639 + original_max_model_len = vllm_config.model_config.max_model_len
640
641 # Define a function to check if a given model length fits in memory
642 def fits_in_memory(model_len: int) -> bool:
638 - # Modify the max_model_len for this calculation
643 + # Temporarily modify the max_model_len for this calculation
644 vllm_config.model_config.max_model_len = model_len
645 # Calculate memory needed for the given model length
646 memory_needed = max_memory_usage_bytes(vllm_config, kv_cache_spec.values())
647 return memory_needed <= available_memory
648
644 - # Binary search for the maximum model length
645 - current_max = vllm_config.model_config.max_model_len
646 - left, right = 1, current_max
647 -
648 - # If even the smallest model length doesn't fit, return 0
649 - if not fits_in_memory(left):
650 - return 0
651 -
652 - # Binary search for the maximum model length that fits
653 - result = 1
654 - while left <= right:
655 - mid = (left + right) // 2
656 - if fits_in_memory(mid):
657 - result = mid
658 - left = mid + 1
659 - else:
660 - right = mid - 1
661 - return result
649 + try:
650 + # Binary search for the maximum model length
651 + left, right = 1, original_max_model_len
652 +
653 + # If even the smallest model length doesn't fit, return 0
654 + if not fits_in_memory(left):
655 + return 0
656 +
657 + # Binary search for the maximum model length that fits
658 + result = 1
659 + while left <= right:
660 + mid = (left + right) // 2
661 + if fits_in_memory(mid):
662 + result = mid
663 + left = mid + 1
664 + else:
665 + right = mid - 1
666 + return result
667 + finally:
668 + # Always restore the original max_model_len to avoid side effects
669 + vllm_config.model_config.max_model_len = original_max_model_len
670
671
672 def check_enough_kv_cache_memory(
@@ -1301,6 +1309,140 @@ def _report_kv_cache_config(
1309 )
1310
1311
1312 +def _max_memory_usage_bytes_from_groups(
1313 + vllm_config: VllmConfig,
1314 + kv_cache_groups: list[KVCacheGroupSpec],
1315 +) -> int:
1316 + """
1317 + Calculate maximum memory usage in bytes from KV cache groups.
1318 +
1319 + This correctly accounts for padding in hybrid models. For example, if a
1320 + model has 8 full attention layers and 9 sliding window layers, they will
1321 + be padded to 9 full + 9 sliding window for uniform group sizes.
1322 + """
1323 + if not kv_cache_groups:
1324 + return 0
1325 +
1326 + # UniformTypeKVCacheSpecs special case (single group, per-layer specs)
1327 + if len(kv_cache_groups) == 1 and isinstance(
1328 + kv_cache_groups[0].kv_cache_spec, UniformTypeKVCacheSpecs
1329 + ):
1330 + per_layer_specs = kv_cache_groups[0].kv_cache_spec.kv_cache_specs
1331 + return sum(
1332 + spec.max_memory_usage_bytes(vllm_config)
1333 + for spec in per_layer_specs.values()
1334 + )
1335 +
1336 + # General case: group_size pools, each shared by one layer per group
1337 + # Memory = group_size * page_size * blocks_for_max_len
1338 + group_size = max(len(group.layer_names) for group in kv_cache_groups)
1339 + page_size = get_uniform_page_size(
1340 + [group.kv_cache_spec for group in kv_cache_groups]
1341 + )
1342 + any_spec = kv_cache_groups[0].kv_cache_spec
1343 + blocks_needed = cdiv(any_spec.max_memory_usage_bytes(vllm_config), page_size)
1344 +
1345 + return group_size * page_size * blocks_needed
1346 +
1347 +
1348 +def _estimate_max_model_len_from_groups(
1349 + vllm_config: VllmConfig,
1350 + kv_cache_groups: list[KVCacheGroupSpec],
1351 + available_memory: int,
1352 +) -> int:
1353 + """
1354 + Binary search for the maximum model length that fits in available memory.
1355 + Returns 0 if even 1 token doesn't fit.
1356 + """
1357 + original_max = vllm_config.model_config.max_model_len
1358 +
1359 + def fits(model_len: int) -> bool:
1360 + vllm_config.model_config.max_model_len = model_len
1361 + return (
1362 + _max_memory_usage_bytes_from_groups(vllm_config, kv_cache_groups)
1363 + <= available_memory
1364 + )
1365 +
1366 + try:
1367 + left, right = 1, original_max
1368 + if not fits(left):
1369 + return 0
1370 + result = 1
1371 + while left <= right:
1372 + mid = (left + right) // 2
1373 + if fits(mid):
1374 + result = mid
1375 + left = mid + 1
1376 + else:
1377 + right = mid - 1
1378 + return result
1379 + finally:
1380 + vllm_config.model_config.max_model_len = original_max
1381 +
1382 +
1383 +def _auto_fit_max_model_len(
1384 + vllm_config: VllmConfig,
1385 + kv_cache_groups: list[KVCacheGroupSpec],
1386 + available_memory: list[int],
1387 +) -> None:
1388 + """
1389 + When max_model_len is set to -1, this function estimates the largest
1390 + context length that can be supported with the available GPU memory.
1391 + It uses binary search to find the maximum length that fits across all
1392 + workers.
1393 +
1394 + Args:
1395 + vllm_config: The global VllmConfig (will be modified in-place)
1396 + kv_cache_groups: The global KV cache groups (from get_kv_cache_groups).
1397 + This correctly accounts for padding in hybrid models.
1398 + available_memory: Memory available for KV cache in bytes for each
1399 + worker.
1400 + """
1401 + original_max = vllm_config.model_config.max_model_len
1402 +
1403 + if not kv_cache_groups:
1404 + # All workers have empty specs (attention-free model)
1405 + logger.info_once(
1406 + "Auto-fit max_model_len: attention-free model, "
1407 + "using derived max_model_len=%d",
1408 + original_max,
1409 + scope="local",
1410 + )
1411 + return
1412 +
1413 + # Use minimum available memory across all workers
1414 + min_available_memory = min(available_memory)
1415 + auto_fit_max = _estimate_max_model_len_from_groups(
1416 + vllm_config, kv_cache_groups, min_available_memory
1417 + )
1418 +
1419 + if auto_fit_max <= 0:
1420 + raise ValueError(
1421 + "Cannot auto-fit max_model_len: not enough GPU memory available "
1422 + "to serve even a single token. Try increasing `gpu_memory_utilization`."
1423 + )
1424 +
1425 + if auto_fit_max >= original_max:
1426 + # The model's full context length fits in memory
1427 + logger.info_once(
1428 + "Auto-fit max_model_len: full model context length %d fits in "
1429 + "available GPU memory",
1430 + original_max,
1431 + scope="local",
1432 + )
1433 + else:
1434 + # Need to reduce max_model_len to fit in memory
1435 + vllm_config.model_config.max_model_len = auto_fit_max
1436 + logger.info_once(
1437 + "Auto-fit max_model_len: reduced from %d to %d to fit in "
1438 + "available GPU memory (%.2f GiB available for KV cache)",
1439 + original_max,
1440 + auto_fit_max,
1441 + min_available_memory / GiB_bytes,
1442 + scope="local",
1443 + )
1444 +
1445 +
1446 def get_kv_cache_configs(
1447 vllm_config: VllmConfig,
1448 kv_cache_specs: list[dict[str, KVCacheSpec]],
@@ -1317,10 +1459,12 @@ def get_kv_cache_configs(
1459 1. Merge the KV cache specs of all workers to get the KVCacheSpecs for
1460 the whole model.
1461 2. Generate the KV cache groups based on the layer ratio of the whole model.
1320 - 3. Generate the KV cache configs for each worker based on the KV cache
1462 + This also handles spec unification for hybrid models.
1463 + 3. Handle auto-fit max_model_len and memory checks using the unified specs.
1464 + 4. Generate the KV cache configs for each worker based on the KV cache
1465 grouping strategy. (This is reasonable because the layer ratio of
1466 different PP stages are similar.)
1323 - 4. Change the num_blocks of each worker to the smallest among all workers
1467 + 5. Change the num_blocks of each worker to the smallest among all workers
1468 and shrink tensor sizes proportionally to avoid allocating unused memory.
1469
1470 Args:
@@ -1333,14 +1477,6 @@ def get_kv_cache_configs(
1477 The generated KVCacheConfigs for each worker.
1478 """
1479
1336 - # Check if the available memory is enough for each worker.
1337 - for kv_cache_spec_one_worker, available_memory_one_worker in zip(
1338 - kv_cache_specs, available_memory
1339 - ):
1340 - check_enough_kv_cache_memory(
1341 - vllm_config, kv_cache_spec_one_worker, available_memory_one_worker
1342 - )
1343 -
1480 # Merge the KV cache specs of all workers. Different PP stages may have
1481 # different layer names, and different TP ranks of the same PP stage should
1482 # have the same KV cache spec.
@@ -1354,8 +1490,52 @@ def get_kv_cache_configs(
1490 "The KV cache specs for the same layer are different "
1491 "across workers. This is not supported yet."
1492 )
1493 +
1494 + # Get global KV cache groups. This also handles spec unification for
1495 + # hybrid models when disable_hybrid_kv_cache_manager is enabled.
1496 + # After this call, merged_kv_cache_specs may be modified in-place.
1497 global_kv_cache_groups = get_kv_cache_groups(vllm_config, merged_kv_cache_specs)
1498
1499 + # If original_max_model_len was -1, automatically
1500 + # determine the maximum model length that fits in available GPU memory.
1501 + # We use the global groups here to correctly account for padding.
1502 + if vllm_config.model_config.original_max_model_len == -1:
1503 + _auto_fit_max_model_len(vllm_config, global_kv_cache_groups, available_memory)
1504 +
1505 + # Check if the available memory is enough (using min across all workers).
1506 + # We use the global groups to correctly account for padding.
1507 + if global_kv_cache_groups:
1508 + min_available_memory = min(available_memory)
1509 + if min_available_memory <= 0:
1510 + raise ValueError(
1511 + "No available memory for the cache blocks. "
1512 + "Try increasing `gpu_memory_utilization` when "
1513 + "initializing the engine."
1514 + )
1515 + max_model_len = vllm_config.model_config.max_model_len
1516 + needed_memory = _max_memory_usage_bytes_from_groups(
1517 + vllm_config, global_kv_cache_groups
1518 + )
1519 + if needed_memory > min_available_memory:
1520 + estimated_max_len = _estimate_max_model_len_from_groups(
1521 + vllm_config, global_kv_cache_groups, min_available_memory
1522 + )
1523 + estimated_msg = ""
1524 + if estimated_max_len > 0:
1525 + estimated_msg = (
1526 + f"Based on the available memory, the estimated maximum "
1527 + f"model length is {estimated_max_len}. "
1528 + )
1529 + raise ValueError(
1530 + f"To serve at least one request with the models's max seq len "
1531 + f"({max_model_len}), ({needed_memory / GiB_bytes:.2f} GiB KV "
1532 + f"cache is needed, which is larger than the available KV cache "
1533 + f"memory ({min_available_memory / GiB_bytes:.2f} GiB). "
1534 + f"{estimated_msg}"
1535 + f"Try increasing `gpu_memory_utilization` or decreasing "
1536 + f"`max_model_len` when initializing the engine."
1537 + )
1538 +
1539 kv_cache_configs: list[KVCacheConfig] = []
1540 for kv_cache_spec_one_worker, available_memory_one_worker in zip(
1541 kv_cache_specs, available_memory
vllm/v1/engine/core.py
+11
@@ -247,9 +247,20 @@ class EngineCore:
247
248 assert len(kv_cache_specs) == len(available_gpu_memory)
249
250 + # Track max_model_len before KV cache config to detect auto-fit changes
251 + max_model_len_before = vllm_config.model_config.max_model_len
252 +
253 kv_cache_configs = get_kv_cache_configs(
254 vllm_config, kv_cache_specs, available_gpu_memory
255 )
256 +
257 + # If auto-fit reduced max_model_len, sync the new value to workers.
258 + # This is needed because workers were spawned before memory profiling
259 + # and have the original (larger) max_model_len cached.
260 + max_model_len_after = vllm_config.model_config.max_model_len
261 + if max_model_len_after != max_model_len_before:
262 + self.collective_rpc("update_max_model_len", args=(max_model_len_after,))
263 +
264 scheduler_kv_cache_config = generate_scheduler_kv_cache_config(kv_cache_configs)
265 num_gpu_blocks = scheduler_kv_cache_config.num_blocks
266 num_cpu_blocks = 0
vllm/v1/worker/gpu_worker.py
+13
@@ -387,6 +387,19 @@ class Worker(WorkerBase):
387 def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]:
388 return self.model_runner.get_kv_cache_spec()
389
390 + def update_max_model_len(self, max_model_len: int) -> None:
391 + """Update max_model_len after auto-fit to GPU memory.
392 +
393 + This is called when max_model_len=-1 is used and the engine
394 + automatically determines the maximum context length that fits
395 + in GPU memory. Workers need to update their cached max_model_len
396 + to match the engine's decision.
397 + """
398 + self.model_config.max_model_len = max_model_len
399 + if self.model_runner is not None:
400 + self.model_runner.max_model_len = max_model_len
401 + logger.debug("Updated max_model_len to %d", max_model_len)
402 +
403 def initialize_from_config(self, kv_cache_config: KVCacheConfig) -> None:
404 """Allocate GPU KV cache with the specified kv_cache_config."""
405