diff --git a/tensorrt_llm/_torch/flashinfer_utils.py b/tensorrt_llm/_torch/flashinfer_utils.py index a7eafaaf9282..61384e21401f 100644 --- a/tensorrt_llm/_torch/flashinfer_utils.py +++ b/tensorrt_llm/_torch/flashinfer_utils.py @@ -1,16 +1,41 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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 os import platform import traceback +from .._utils import get_sm_version from ..logger import logger IS_FLASHINFER_AVAILABLE = False def get_env_enable_pdl() -> bool: - enabled = os.environ.get("TRTLLM_ENABLE_PDL", "1") == "1" - if enabled and not getattr(get_env_enable_pdl, "_printed", False): - logger.info("PDL enabled") + requested = os.environ.get("TRTLLM_ENABLE_PDL", "1") == "1" + if not requested: + return False + + sm_version = get_sm_version() + enabled = sm_version >= 90 + if not getattr(get_env_enable_pdl, "_printed", False): + if enabled: + logger.info("PDL enabled") + else: + logger.info( + f"PDL disabled on SM{sm_version}: requires SM90 or newer") setattr(get_env_enable_pdl, "_printed", True) return enabled diff --git a/tests/unittest/_torch/test_flashinfer_utils.py b/tests/unittest/_torch/test_flashinfer_utils.py new file mode 100644 index 000000000000..db7008a39f10 --- /dev/null +++ b/tests/unittest/_torch/test_flashinfer_utils.py @@ -0,0 +1,41 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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 pytest + +from tensorrt_llm._torch import flashinfer_utils + + +@pytest.mark.parametrize( + ("env_value", "sm_version", "expected"), + [ + (None, 89, False), + ("1", 89, False), + ("0", 89, False), + (None, 90, True), + ("1", 90, True), + ("0", 90, False), + (None, 100, True), + ], +) +def test_get_env_enable_pdl_hardware_gate(monkeypatch, env_value, sm_version, expected): + if env_value is None: + monkeypatch.delenv("TRTLLM_ENABLE_PDL", raising=False) + else: + monkeypatch.setenv("TRTLLM_ENABLE_PDL", env_value) + monkeypatch.setattr(flashinfer_utils, "get_sm_version", lambda: sm_version) + monkeypatch.delattr(flashinfer_utils.get_env_enable_pdl, "_printed", raising=False) + + assert flashinfer_utils.get_env_enable_pdl() is expected