From b2b11177fe706d0e5c4cfb4f77015f7421030eae Mon Sep 17 00:00:00 2001 From: Samuel Mendoza-Jonas Date: Fri, 12 Jun 2026 15:49:47 +1000 Subject: [PATCH] [https://nvbugs/6506990][fix] Don't treat NVLE-only as CC enabled NVLE can be enabled independently of CC, so CC restrictions such as H2D bounce buffering do not apply to NVLE-only systems. Query and cache the CC and NVLE states together as one coherent NVML snapshot. Keep confidential_compute_enabled() as the existing boolean interface so current call sites remain unchanged, and cover modern and legacy NVML states with CPU-only tests. Signed-off-by: Samuel Mendoza-Jonas Signed-off-by: Dan Hansen <1+dhansen-nvidia@users.noreply.github.com> --- tensorrt_llm/_utils.py | 38 +++-- .../utils/test_confidential_compute.py | 136 ++++++++++++++++++ 2 files changed, 161 insertions(+), 13 deletions(-) create mode 100644 tests/unittest/utils/test_confidential_compute.py diff --git a/tensorrt_llm/_utils.py b/tensorrt_llm/_utils.py index 71574f3e1aa5..9572c3760cec 100644 --- a/tensorrt_llm/_utils.py +++ b/tensorrt_llm/_utils.py @@ -1171,18 +1171,22 @@ def set_prometheus_multiproc_dir() -> object: f"PROMETHEUS_MULTIPROC_DIR: {os.environ['PROMETHEUS_MULTIPROC_DIR']}") -def confidential_compute_enabled() -> bool: - """ - Query NVML for the confidential compute state +@lru_cache(maxsize=1) +def get_cc_and_nvle_status() -> tuple[bool, bool]: + """Query NVML for the confidential compute and NVLink encryption state. + + Returns: + A tuple of ``(cc_enabled, nvle_enabled)``. """ try: import pynvml except ImportError: - logger.error("pynvml not available; assuming CC=off") - return False + logger.error("pynvml not available; assuming CC and NVLE are off") + return False, False cc_enabled = False + nvle_enabled = False try: pynvml.nvmlInit() @@ -1192,11 +1196,13 @@ def confidential_compute_enabled() -> bool: cc_settings = pynvml.c_nvmlSystemConfComputeSettings_v1_t() ret = pynvml.nvmlSystemGetConfComputeSettings(byref(cc_settings)) pynvml._nvmlCheckReturn(ret) - cc_enabled = ( - cc_settings.ccFeature == pynvml.NVML_CC_SYSTEM_FEATURE_ENABLED - or cc_settings.multiGpuMode - == pynvml.NVML_CC_SYSTEM_MULTIGPU_PROTECTED_PCIE - or cc_settings.multiGpuMode == pynvml.NVML_CC_SYSTEM_MULTIGPU_NVLE) + # PPCIE implies CC, but NVLE does not necessarily + cc_enabled = (cc_settings.ccFeature + == pynvml.NVML_CC_SYSTEM_FEATURE_ENABLED + or cc_settings.multiGpuMode + == pynvml.NVML_CC_SYSTEM_MULTIGPU_PROTECTED_PCIE) + nvle_enabled = ( + cc_settings.multiGpuMode == pynvml.NVML_CC_SYSTEM_MULTIGPU_NVLE) except pynvml.NVMLError_NotSupported: # Simple query for older GPUs try: @@ -1204,17 +1210,23 @@ def confidential_compute_enabled() -> bool: cc_enabled = ( cc_state.ccFeature == pynvml.NVML_CC_SYSTEM_FEATURE_ENABLED) except Exception as e: - logger.error(f"Error querying confidential compute state: {str(e)}") + logger.error(f"Error querying CC and NVLE state: {str(e)}") except Exception as e: - logger.error(f"Error querying confidential compute state: {str(e)}") + logger.error(f"Error querying CC and NVLE state: {str(e)}") finally: # Shutdown try: pynvml.nvmlShutdown() - except: + except pynvml.NVMLError: # Ignore shutdown errors pass + return cc_enabled, nvle_enabled + + +def confidential_compute_enabled() -> bool: + """Return whether confidential compute restrictions are enabled.""" + cc_enabled, _ = get_cc_and_nvle_status() return cc_enabled diff --git a/tests/unittest/utils/test_confidential_compute.py b/tests/unittest/utils/test_confidential_compute.py new file mode 100644 index 000000000000..6693ba2bb91d --- /dev/null +++ b/tests/unittest/utils/test_confidential_compute.py @@ -0,0 +1,136 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import ctypes +import sys +import types + +import pytest + +from tensorrt_llm._utils import confidential_compute_enabled, get_cc_and_nvle_status + + +class _NvmlError(Exception): + pass + + +class _NvmlErrorNotSupported(_NvmlError): + pass + + +class _ConfComputeSettings(ctypes.Structure): + _fields_ = [ + ("ccFeature", ctypes.c_uint), + ("multiGpuMode", ctypes.c_uint), + ] + + +def _make_pynvml( + *, + cc_feature: int, + multi_gpu_mode: int, + settings_supported: bool = True, + legacy_cc_feature: int = 0, +) -> types.ModuleType: + pynvml = types.ModuleType("pynvml") + pynvml.NVMLError = _NvmlError + pynvml.NVMLError_NotSupported = _NvmlErrorNotSupported + pynvml.NVML_CC_SYSTEM_FEATURE_ENABLED = 1 + pynvml.NVML_CC_SYSTEM_MULTIGPU_PROTECTED_PCIE = 1 + pynvml.NVML_CC_SYSTEM_MULTIGPU_NVLE = 2 + pynvml.c_nvmlSystemConfComputeSettings_v1_t = _ConfComputeSettings + pynvml.nvmlInit = lambda: None + pynvml.nvmlShutdown = lambda: None + pynvml._nvmlCheckReturn = lambda _: None + pynvml.settings_queries = 0 + + def get_settings(settings_ptr): + pynvml.settings_queries += 1 + if not settings_supported: + raise _NvmlErrorNotSupported + settings_ptr._obj.ccFeature = cc_feature + settings_ptr._obj.multiGpuMode = multi_gpu_mode + return 0 + + pynvml.nvmlSystemGetConfComputeSettings = get_settings + pynvml.nvmlSystemGetConfComputeState = lambda: types.SimpleNamespace( + ccFeature=legacy_cc_feature + ) + return pynvml + + +@pytest.fixture(autouse=True) +def clear_confidential_compute_status_cache(): + get_cc_and_nvle_status.cache_clear() + yield + get_cc_and_nvle_status.cache_clear() + + +@pytest.mark.parametrize( + "cc_feature,multi_gpu_mode,expected", + [ + (0, 0, (False, False)), + (0, 1, (True, False)), + (0, 2, (False, True)), + (1, 2, (True, True)), + ], +) +def test_get_cc_and_nvle_status(monkeypatch, cc_feature, multi_gpu_mode, expected): + pynvml = _make_pynvml( + cc_feature=cc_feature, + multi_gpu_mode=multi_gpu_mode, + ) + monkeypatch.setitem(sys.modules, "pynvml", pynvml) + + assert get_cc_and_nvle_status() == expected + + +def test_get_cc_and_nvle_status_is_cached(monkeypatch): + pynvml = _make_pynvml(cc_feature=0, multi_gpu_mode=2) + monkeypatch.setitem(sys.modules, "pynvml", pynvml) + + assert get_cc_and_nvle_status() == (False, True) + assert get_cc_and_nvle_status() == (False, True) + assert pynvml.settings_queries == 1 + + +@pytest.mark.parametrize( + "cc_feature,multi_gpu_mode,expected", + [ + (0, 0, False), + (0, 1, True), + (0, 2, False), + (1, 2, True), + ], +) +def test_confidential_compute_enabled(monkeypatch, cc_feature, multi_gpu_mode, expected): + pynvml = _make_pynvml( + cc_feature=cc_feature, + multi_gpu_mode=multi_gpu_mode, + ) + monkeypatch.setitem(sys.modules, "pynvml", pynvml) + + assert confidential_compute_enabled() is expected + + +def test_get_cc_and_nvle_status_uses_legacy_cc_fallback(monkeypatch): + pynvml = _make_pynvml( + cc_feature=0, + multi_gpu_mode=0, + settings_supported=False, + legacy_cc_feature=1, + ) + monkeypatch.setitem(sys.modules, "pynvml", pynvml) + + assert get_cc_and_nvle_status() == (True, False)