[Bugfix][Platform] Fall back to torch when device-level NVML memory is unsupported (GB10) (#57378)

Signed-off-by: hclsys <chenglunhu@gmail.com>
Signed-off-by: Chenglun Hu <chenglunhu@gmail.com>
Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>

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

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,

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