TensorRT-LLMs/tensorrt_llm/_torch/flashinfer_utils.py
Venky 639c939a4f
[TRTC-1943][feat] Env vars override support in LLM API (#9104)
Signed-off-by: Venky Ganesh <23023424+venkywonka@users.noreply.github.com>
2025-12-01 10:04:49 -08:00

28 lines
743 B
Python

import os
import platform
import traceback
from ..logger import logger
IS_FLASHINFER_AVAILABLE = False
def get_env_enable_pdl():
enabled = os.environ.get("TRTLLM_ENABLE_PDL", "0") == "1"
if enabled and not getattr(get_env_enable_pdl, "_printed", False):
logger.info("PDL enabled")
setattr(get_env_enable_pdl, "_printed", True)
return 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"
)