diff --git a/tests/engine/test_arg_utils.py b/tests/engine/test_arg_utils.py index 3e3a82f80f..81587c652d 100644 --- a/tests/engine/test_arg_utils.py +++ b/tests/engine/test_arg_utils.py @@ -1054,3 +1054,34 @@ class TestDpDeviceIdSharding: get_physical_gpu_ids_for_local_dp_rank( evar, local_dp_rank=2, world_size=2, user_assigned_gpu_ids=[4, 5, 6, 7] ) + + +@pytest.mark.parametrize( + ("name", "mem_gib", "capability", "expected_api_tokens", "expected_api_seqs"), + [ + ("nvidia gb10", 121, (12, 1), 2048, 256), + ("nvidia h100", 80, (9, 0), 8192, 1024), + ("nvidia rtx pro 6000 blackwell", 96, (12, 0), 8192, 1024), + ], +) +def test_get_batch_defaults_gb10_stays_on_small_tier( + monkeypatch, name, mem_gib, capability, expected_api_tokens, expected_api_seqs +): + """GB10 reports >= 70 GiB (unified memory) but not the compute of a + large discrete GPU, so it keeps the small-tier defaults.""" + from types import SimpleNamespace + + from vllm.usage.usage_lib import UsageContext + from vllm.utils.mem_constants import GiB_bytes + + stub = SimpleNamespace( + get_device_total_memory=lambda: mem_gib * GiB_bytes, + get_device_name=lambda: name, + is_device_capability=lambda cap: cap == capability, + is_tpu=lambda: False, + is_cpu=lambda: False, + ) + monkeypatch.setattr("vllm.engine.arg_utils.current_platform", stub) + tokens, seqs = EngineArgs.get_batch_defaults(1) + assert tokens[UsageContext.OPENAI_API_SERVER] == expected_api_tokens + assert seqs[UsageContext.OPENAI_API_SERVER] == expected_api_seqs diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index e0cfb3b6d2..8f285317e0 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -2937,10 +2937,14 @@ class EngineArgs: try: device_memory = current_platform.get_device_total_memory() device_name = current_platform.get_device_name().lower() + # GB10 (DGX Spark, cc 12.1) has the memory but not the compute of + # a large discrete GPU, keep it on the small tier. + is_gb10 = current_platform.is_device_capability((12, 1)) except Exception: # This is only used to set default_max_num_batched_tokens device_memory = 0 device_name = "" + is_gb10 = False # NOTE(Kuntai): Setting large `max_num_batched_tokens` for A100 reduces # throughput, see PR #17885 for more details. @@ -2955,7 +2959,11 @@ class EngineArgs: UsageContext.LLM_CLASS: 1024, UsageContext.OPENAI_API_SERVER: 1024, } - elif device_memory >= 70 * GiB_bytes and "a100" not in device_name: + elif ( + device_memory >= 70 * GiB_bytes + and "a100" not in device_name + and not is_gb10 + ): # For GPUs like H100 and H200, use larger offline defaults. default_max_num_batched_tokens = { UsageContext.LLM_CLASS: 16384, diff --git a/vllm/platforms/cuda.py b/vllm/platforms/cuda.py index 7a28dcbc96..462d1aa18d 100644 --- a/vllm/platforms/cuda.py +++ b/vllm/platforms/cuda.py @@ -27,6 +27,7 @@ with contextlib.suppress(ImportError): import vllm.envs as envs from vllm.logger import init_logger from vllm.utils.import_utils import import_pynvml +from vllm.utils.platform_utils import cuda_get_device_properties from vllm.v1.attention.backends.registry import AttentionBackendEnum from .interface import DeviceCapability, Platform, PlatformEnum, in_wsl @@ -802,7 +803,14 @@ class NvmlCudaPlatform(CudaPlatformBase): def get_device_total_memory(cls, device_id: int = 0) -> int: physical_device_id = cls.device_id_to_physical_device_id(device_id) handle = pynvml.nvmlDeviceGetHandleByIndex(physical_device_id) - return int(pynvml.nvmlDeviceGetMemoryInfo(handle).total) + try: + return int(pynvml.nvmlDeviceGetMemoryInfo(handle).total) + except pynvml.NVMLError: + # Integrated parts (e.g. GB10 / DGX Spark) do not implement the + # device-level NVML memory query - nvidia-smi reports [N/A] for + # memory.total there. Fall back to the same source + # NonNvmlCudaPlatform reads. + return int(cuda_get_device_properties(device_id, ("total_memory",))[0]) @classmethod @with_nvml_context