[Bug]: trtllm-bench crashes on DGX Spark (GB10): nvmlDeviceGetMemoryInfo raises NVMLError_NotSupported in the batch-size heuristic #21317

System Info

• DGX Spark, 1x NVIDIA GB10 (SM 12.1), unified memory (128 GB LPDDR5x; /proc/meminfo MemTotal 127601452 kB), aarch64, Ubuntu 24.04 host
• Container nvcr.io/nvidia/tensorrt-llm/release:1.4.0rc2 (linux/arm64, TRT_LLM_GIT_COMMIT=b96ba434e8a0), CUDA 13.4.1, torch 2.14.0a0 (NGC 26.09), nvidia-ml-py from the image
• tensorrt_llm/bench/ is unchanged between v1.4.0rc2 and current main (7d8d94684), so main is affected too
• nvidia-smi --query-gpu=name,memory.total,memory.used --format=csv prints NVIDIA GB10, [N/A], [N/A]

Who can help?

trtllm-bench owners (NVIDIA/trt-llm-perf-devs per CODEOWNERS for /tensorrt_llm/bench)

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

NVML itself, inside the container (--gpus all):

import pynvml
pynvml.nvmlInit()
h = pynvml.nvmlDeviceGetHandleByIndex(0)
print(pynvml.nvmlDeviceGetName(h))       # NVIDIA GB10
pynvml.nvmlDeviceGetMemoryInfo(h)        # pynvml.NVMLError_NotSupported: Not Supported

trtllm-bench without --max_batch_size (the quick-start in docs/source/developer-guide/perf-benchmarking.md passes neither flag):

# dataset: 8 requests, {"task_id": i, "input_ids": [32 ids], "output_tokens": 32}
trtllm-bench --model Qwen/Qwen3-0.6B --model_path <local snapshot> \
  throughput --backend pytorch --dataset ds.jsonl \
  --num_requests 8 --warmup 0 --concurrency 4 --max_seq_len 256 \
  --max_num_tokens 1024 --kv_cache_free_gpu_mem_fraction 0.1 --config caps.yaml
# caps.yaml: kv_cache_config: {max_tokens: 8192, free_gpu_memory_fraction: 0.1}, cuda_graph_config: null

(--max_num_tokens, the KV fraction and caps.yaml only keep memory use small on a shared machine. The heuristic still runs because --max_batch_size is unset, and the crash happens before any of them take effect.)

Expected behavior

The benchmark sizes max_batch_size / max_num_tokens from the memory the GPU can use (on unified memory, the system memory) and runs, the same way it does on discrete GPUs.

actual behavior

It aborts in get_settings, before the model is loaded:

[10/29/2026-17:19:16] [TRT-LLM] [I] [bench] Number of attention layers: 28
Traceback (most recent call last):
  File "/usr/local/bin/trtllm-bench", line 6, in <module>
    sys.exit(main())
  ...  (click / telemetry frames)
  File "/usr/local/lib/python3.12/dist-packages/tensorrt_llm/bench/benchmark/throughput.py", line 394, in throughput_command
    exec_settings = get_settings(params, metadata, bench_env.model,
  File "/usr/local/lib/python3.12/dist-packages/tensorrt_llm/bench/benchmark/utils/general.py", line 102, in get_settings
    max_batch_size, max_num_tokens = get_benchmark_engine_settings(
  File "/usr/local/lib/python3.12/dist-packages/tensorrt_llm/bench/tuning/settings.py", line 68, in get_benchmark_engine_settings
    max_batch_size, max_num_tokens = calc_engine_setting(
  File "/usr/local/lib/python3.12/dist-packages/tensorrt_llm/bench/tuning/heuristics.py", line 112, in calc_engine_setting
    total_gpu_memory = get_device_memory() * n_gpus
  File "/usr/local/lib/python3.12/dist-packages/tensorrt_llm/bench/tuning/utils.py", line 30, in get_device_memory
    mem_info = pynvml.nvmlDeviceGetMemoryInfo(handle)
  File "/usr/local/lib/python3.12/dist-packages/pynvml.py", line 3831, in nvmlDeviceGetMemoryInfo
    _nvmlCheckReturn(ret)
  File "/usr/local/lib/python3.12/dist-packages/pynvml.py", line 1098, in _nvmlCheckReturn
    raise NVMLError(ret)
pynvml.NVMLError_NotSupported: Not Supported

additional notes

• Cause: tensorrt_llm/bench/tuning/utils.py::get_device_memory() calls pynvml.nvmlDeviceGetMemoryInfo with no error handling. It is reached from get_settings -> get_benchmark_engine_settings -> calc_engine_setting (heuristics.py) for unquantized, FP8, FP8 block-scale and NVFP4 checkpoints: on every trtllm-bench latency run (it has no --max_batch_size / --max_num_tokens options), and on throughput runs where either flag is not passed. When it raises, nvmlShutdown() is also skipped.
• CI misses it because the only GB10 test list, tests/integration/test_lists/test-db/l0_gb10.yml, has no trtllm-bench test.
• The integration-test helper tests/integration/defs/conftest.py::get_device_memory already falls back to /proc/meminfo "for systems with unified memory (e.g. DGX Spark)", but the product code has no such fallback.
• Related but different: [Bug]: quantize.py --device_map has no working value for GB10 (unified memory / NVML) #15021 (legacy quantize.py --device_map on GB10, closed as not planned) and [#15178][fix] Account for reclaimable memory in V2 KV cache sizing on integrated GPUs #20951 (open PR on runtime V2 KV cache sizing on integrated GPUs, _torch/pyexecutor only).
• I have a small fix ready (catch NVMLError_NotSupported and fall back to psutil.virtual_memory().total, nvmlShutdown() in finally, CPU unit tests) and will open a PR that links this issue.

Before submitting a new issue...

• Make sure you already searched for relevant issues, and checked the documentation and examples for answers to frequently asked questions.

AI assistance: I used Claude Code to help investigate this and draft the report; I ran the reproduction on my own DGX Spark.