Allow %i substitution for process ID for multi-GPU environments. Logging level above 3 now automatically include cudnn and cublas API logging

Author: bkryu

flashinfer/api_logging.py

@@ -25,9 +25,58 @@
 import torch
 
 
+# Helper function to substitute %i with process ID in file paths
+def _substitute_process_id(path: str) -> str:
+    """
+    Replace %i with the current process ID in a path.
+
+    This is useful for multi-process/multi-GPU environments where each process
+    needs its own log file.
+
+    Example: "flashinfer_log_%i.txt" -> "flashinfer_log_12345.txt"
+    """
+    if "%i" in path:
+        return path.replace("%i", str(os.getpid()))
+    return path
+
+
 # Read environment variables once at module load time
 _API_LOG_LEVEL = int(os.environ.get("FLASHINFER_LOGLEVEL_DBG", "0"))
-_API_LOG_DEST = os.environ.get("FLASHINFER_LOGDEST_DBG", "stdout")
+_API_LOG_DEST = _substitute_process_id(
+    os.environ.get("FLASHINFER_LOGDEST_DBG", "stdout")
+)
+
+# Enable cuDNN, cuBLAS, and cuBLASLt API logging when FlashInfer logging level >= 3
+# Only override if the user hasn't already configured the logging switch
+# If the switch is not set, we override both the switch and destination as a bundle
+if _API_LOG_LEVEL >= 3:
+    # cuBLAS logging: Check switch, set both switch and destination
+    if "CUBLAS_LOGINFO_DBG" not in os.environ:
+        os.environ["CUBLAS_LOGINFO_DBG"] = "1"
+        os.environ["CUBLAS_LOGDEST_DBG"] = _substitute_process_id(
+            "flashinfer_cublas_log_%i.txt"
+        )
+
+    # cuBLASLt logging: Check switch, set both switch and destination
+    if "CUBLASLT_LOG_LEVEL" not in os.environ:
+        os.environ["CUBLASLT_LOG_LEVEL"] = "2"
+        os.environ["CUBLASLT_LOG_FILE"] = _substitute_process_id(
+            "flashinfer_cublaslt_log_%i.txt"
+        )
+
+    # cuDNN backend logging: Check switch, set both switch and destination
+    if "CUDNN_LOGLEVEL_DBG" not in os.environ:
+        os.environ["CUDNN_LOGLEVEL_DBG"] = "2.5"
+        os.environ["CUDNN_LOGDEST_DBG"] = _substitute_process_id(
+            "flashinfer_cudnn_backend_log_%i.txt"
+        )
+
+    # cuDNN frontend logging: Check switch, set both switch and destination
+    if "CUDNN_FRONTEND_LOG_INFO" not in os.environ:
+        os.environ["CUDNN_FRONTEND_LOG_INFO"] = "1"
+        os.environ["CUDNN_FRONTEND_LOG_FILE"] = _substitute_process_id(
+            "flashinfer_cudnn_frontend_log_%i.txt"
+        )
 
 # Create logger using Python's logging library
 _logger = logging.getLogger("flashinfer.api")
@@ -129,6 +178,28 @@ def _log_system_info():
         # PyTorch version
         lines.append(f"PyTorch version: {torch.__version__}")
 
+        # cuDNN/cuBLAS/cuBLASLt logging status
+        if _API_LOG_LEVEL >= 3:
+            lines.append("")
+            lines.append("cuDNN/cuBLAS/cuBLASLt Logging: Enabled (Level 3)")
+            cublas_info = os.environ.get("CUBLAS_LOGINFO_DBG", "not set")
+            cublas_dest = os.environ.get("CUBLAS_LOGDEST_DBG", "not set")
+            cublaslt_level = os.environ.get("CUBLASLT_LOG_LEVEL", "not set")
+            cublaslt_file = os.environ.get("CUBLASLT_LOG_FILE", "not set")
+            cudnn_level = os.environ.get("CUDNN_LOGLEVEL_DBG", "not set")
+            cudnn_dest = os.environ.get("CUDNN_LOGDEST_DBG", "not set")
+            cudnn_fe_info = os.environ.get("CUDNN_FRONTEND_LOG_INFO", "not set")
+            cudnn_fe_file = os.environ.get("CUDNN_FRONTEND_LOG_FILE", "not set")
+
+            lines.append(f"  CUBLAS_LOGINFO_DBG={cublas_info}")
+            lines.append(f"  CUBLAS_LOGDEST_DBG={cublas_dest}")
+            lines.append(f"  CUBLASLT_LOG_LEVEL={cublaslt_level}")
+            lines.append(f"  CUBLASLT_LOG_FILE={cublaslt_file}")
+            lines.append(f"  CUDNN_LOGLEVEL_DBG={cudnn_level}")
+            lines.append(f"  CUDNN_LOGDEST_DBG={cudnn_dest}")
+            lines.append(f"  CUDNN_FRONTEND_LOG_INFO={cudnn_fe_info}")
+            lines.append(f"  CUDNN_FRONTEND_LOG_FILE={cudnn_fe_file}")
+
     except Exception as e:
         lines.append(f"Error gathering system information: {e}")
 
@@ -460,6 +531,7 @@ def flashinfer_api_log(func: Callable = None) -> Callable:
         - "stdout": Log to standard output
         - "stderr": Log to standard error
         - <path>: Log to specified file path
+        - Use %i in path for process ID substitution (e.g., "log_%i.txt" -> "log_12345.txt")
 
     Examples
     --------
@@ -482,6 +554,13 @@ def flashinfer_api_log(func: Callable = None) -> Callable:
     - **CUDA Graph Compatibility**: At level 3, tensor statistics (min/max/mean/nan_count)
       are automatically skipped during CUDA graph capture to avoid synchronization issues.
       The message "[statistics skipped: CUDA graph capture in progress]" will be logged.
+    - **cuDNN/cuBLAS/cuBLASLt Integration**: At level 3, if not already set by the user, the following
+      environment variables are automatically configured to enable cuDNN, cuBLAS, and cuBLASLt logging:
+      - CUBLAS_LOGINFO_DBG=1, CUBLAS_LOGDEST_DBG=flashinfer_cublas_log_%i.txt
+      - CUBLASLT_LOG_LEVEL=2, CUBLASLT_LOG_FILE=flashinfer_cublaslt_log_%i.txt
+      - CUDNN_LOGLEVEL_DBG=2.5, CUDNN_LOGDEST_DBG=flashinfer_cudnn_backend_log_%i.txt
+      - CUDNN_FRONTEND_LOG_INFO=1, CUDNN_FRONTEND_LOG_FILE=flashinfer_cudnn_frontend_log_%i.txt
+      The %i pattern is automatically replaced with the process ID for multi-process environments.
     - The logger does not propagate to the root logger to avoid duplicate logs.
     """
     # If logging is disabled, return original function with zero overhead