TensorRT-LLMs/tensorrt_llm/_torch/flashinfer_utils.py
Robin Kobus b95cab2a7c
[None][ci] move unittests to sub-directories (#6635)
Signed-off-by: Robin Kobus <19427718+Funatiq@users.noreply.github.com>
2025-08-20 05:42:22 -04:00

28 lines
645 B
Python

import os
import platform
import traceback
from ..logger import logger
IS_FLASHINFER_AVAILABLE = False
def get_env_enable_pdl():
return os.environ.get("TRTLLM_ENABLE_PDL", "0") == "1"
ENABLE_PDL = get_env_enable_pdl()
if ENABLE_PDL:
logger.info("PDL is enabled")
if platform.system() != "Windows":
try:
import flashinfer
logger.info(f"flashinfer is available: {flashinfer.__version__}")
IS_FLASHINFER_AVAILABLE = True
except ImportError:
traceback.print_exc()
print(
"flashinfer is not installed properly, please try pip install or building from source codes"
)