feat: update flashinfer-cli (#1613)

<!-- .github/pull_request_template.md -->

## 📌 Description

- add functionality of clear cache/cubin, and show config
- using click library
- show progress bar when downloading the cubins

Usage:
```
python3 -m flashinfer show-config
python3 -m flashinfer list-cubin
python3 -m flashinfer clear-cache
python3 -m flashinfer clear-cubin
python3 -m flashinfer download-cubin
```

More features will be coming in the future.

## 🔍 Related Issues

<!-- Link any related issues here -->

## 🚀 Pull Request Checklist

Thank you for contributing to FlashInfer! Before we review your pull
request, please make sure the following items are complete.

### ✅ Pre-commit Checks

- [x] I have installed `pre-commit` by running `pip install pre-commit`
(or used your preferred method).
- [x] I have installed the hooks with `pre-commit install`.
- [x] I have run the hooks manually with `pre-commit run --all-files`
and fixed any reported issues.

> If you are unsure about how to set up `pre-commit`, see [the
pre-commit documentation](https://pre-commit.com/).

## 🧪 Tests

- [x] Tests have been added or updated as needed.
- [x] All tests are passing (`unittest`, etc.).

## Reviewer Notes

<!-- Optional: anything you'd like reviewers to focus on, concerns, etc.
-->

---------

Co-authored-by: Yaxing Cai <[email protected]>

Author: yzh119

flashinfer/__main__.py

@@ -15,20 +15,131 @@
 """
 
 # flashinfer-cli
-import argparse
+import click
+from tabulate import tabulate  # type: ignore[import-untyped]
+
+from .artifacts import (
+    ArtifactPath,
+    download_artifacts,
+    clear_cubin,
+    get_artifacts_status,
+)
+from .jit import clear_cache_dir
+from .jit.cubin_loader import FLASHINFER_CUBINS_REPOSITORY
+from .jit.env import FLASHINFER_CACHE_DIR, FLASHINFER_CUBIN_DIR
+from .jit.core import current_compilation_context
+from .jit.cpp_ext import get_cuda_path, get_cuda_version
+
+
+def _download_cubin():
+    """Helper function to download cubin"""
+    try:
+        download_artifacts()
+        click.secho("✅ All cubin download tasks completed successfully.", fg="green")
+    except Exception as e:
+        click.secho(f"❌ Cubin download failed: {e}", fg="red")
+
+
+@click.group(invoke_without_command=True)
+@click.option(
+    "--download-cubin", "download_cubin_flag", is_flag=True, help="Download artifacts"
+)
+@click.pass_context
+def cli(ctx, download_cubin_flag):
+    """FlashInfer CLI"""
+    if download_cubin_flag:
+        _download_cubin()
+    elif ctx.invoked_subcommand is None:
+        click.echo(ctx.get_help())
+
+
+# list of environment variables
+env_variables = {
+    "FLASHINFER_CACHE_DIR": FLASHINFER_CACHE_DIR,
+    "FLASHINFER_CUBIN_DIR": FLASHINFER_CUBIN_DIR,
+    "CUDA_HOME": get_cuda_path(),
+    "CUDA_VERSION": get_cuda_version(),
+    "FLASHINFER_CUDA_ARCH_LIST": current_compilation_context.TARGET_CUDA_ARCHS,
+    "FLASHINFER_CUBINS_REPOSITORY": FLASHINFER_CUBINS_REPOSITORY,
+}
+
+
+@cli.command("show-config")
+def show_config_cmd():
+    """Show configuration"""
+    import torch
+
+    # Section: Torch Version Info
+    click.secho("=== Torch Version Info ===", fg="yellow")
+    click.secho("Torch version:", fg="magenta", nl=False)
+    click.secho(f" {torch.__version__}", fg="cyan")
+    click.secho("", fg="white")
+
+    # Section: Environment Variables
+    click.secho("=== Environment Variables ===", fg="yellow")
+    for name, value in env_variables.items():
+        click.secho(f"{name}:", fg="magenta", nl=False)
+        click.secho(f" {value}", fg="cyan")
+    click.secho("", fg="white")
+
+    # Section: Artifact path
+    click.secho("=== Artifact Path ===", fg="yellow")
+    # list all artifact paths
+    for name, path in ArtifactPath.__dict__.items():
+        if not name.startswith("__"):
+            click.secho(f"{name}:", fg="magenta", nl=False)
+            click.secho(f" {path}", fg="cyan")
+    click.secho("", fg="white")
+
+    # Section: Downloaded Cubins
+    click.secho("=== Downloaded Cubins ===", fg="yellow")
+
+    status = get_artifacts_status()
+    num_downloaded = sum(1 for _, _, exists in status if exists)
+    total_cubins = len(status)
+
+    click.secho(f"Downloaded {num_downloaded}/{total_cubins} cubins", fg="cyan")
+
+
+@cli.command("list-cubins")
+def list_cubins_cmd():
+    """List downloaded cubins"""
+    status = get_artifacts_status()
+    table_data = []
+    for name, extension, exists in status:
+        status_str = "Downloaded" if exists else "Missing"
+        color = "green" if exists else "red"
+        table_data.append([f"{name}{extension}", click.style(status_str, fg=color)])
+
+    click.echo(tabulate(table_data, headers=["Cubin", "Status"], tablefmt="github"))
+    click.secho("", fg="white")
+
+
+@cli.command("download-cubin")
+def download_cubin_cmd():
+    """Download artifacts"""
+    _download_cubin()
+
+
+@cli.command("clear-cache")
+def clear_cache_cmd():
+    """Clear cache"""
+    try:
+        clear_cache_dir()
+        click.secho("✅ Cache cleared successfully.", fg="green")
+    except Exception as e:
+        click.secho(f"❌ Cache clear failed: {e}", fg="red")
+
+
+@cli.command("clear-cubin")
+def clear_cubin_cmd():
+    """Clear cubin"""
+    try:
+        clear_cubin()
+        click.secho("✅ Cubin cleared successfully.", fg="green")
+    except Exception as e:
+        click.secho(f"❌ Cubin clear failed: {e}", fg="red")
 
-from .artifacts import download_artifacts
 
 if __name__ == "__main__":
-    parser = argparse.ArgumentParser("FlashInfer CLI")
-    parser.add_argument(
-        "--download-cubin", action="store_true", help="Download artifacts"
-    )
-
-    args = parser.parse_args()
-
-    if args.download_cubin:
-        if download_artifacts():
-            print("✅ All cubin download tasks completed successfully.")
-        else:
-            print("❌ Some cubin download tasks failed.")
+    cli()

flashinfer/artifacts.py

@@ -20,9 +20,70 @@
 from concurrent.futures import ThreadPoolExecutor, as_completed
 
 import requests  # type: ignore[import-untyped]
+import shutil
 
 from .jit.core import logger
-from .jit.cubin_loader import FLASHINFER_CUBINS_REPOSITORY, get_cubin
+from .jit.cubin_loader import (
+    FLASHINFER_CUBINS_REPOSITORY,
+    get_cubin,
+    FLASHINFER_CUBIN_DIR,
+)
+
+
+import logging
+from contextlib import contextmanager
+
+
+@contextmanager
+def temp_env_var(key, value):
+    old_value = os.environ.get(key, None)
+    os.environ[key] = value
+    try:
+        yield
+    finally:
+        if old_value is None:
+            os.environ.pop(key, None)
+        else:
+            os.environ[key] = old_value
+
+
+@contextmanager
+def patch_logger_for_tqdm(logger):
+    """
+    Context manager to patch the logger so that log messages are displayed using tqdm.write,
+    preventing interference with tqdm progress bars.
+    """
+    import tqdm
+
+    class TqdmLoggingHandler(logging.Handler):
+        def emit(self, record):
+            try:
+                msg = self.format(record)
+                tqdm.write(msg, end="\n")
+            except Exception:
+                self.handleError(record)
+
+    # Save original handlers and level
+    original_handlers = logger.handlers[:]
+    original_level = logger.level
+
+    # Remove all existing handlers to prevent duplicate output
+    for h in original_handlers:
+        logger.removeHandler(h)
+
+    # Add our tqdm-aware handler
+    handler = TqdmLoggingHandler()
+    handler.setFormatter(logging.Formatter("%(levelname)s: %(message)s"))
+    logger.addHandler(handler)
+    logger.setLevel(logging.INFO)
+    try:
+        yield
+    finally:
+        # Remove tqdm handler and restore original handlers and level
+        logger.removeHandler(handler)
+        for h in original_handlers:
+            logger.addHandler(h)
+        logger.setLevel(original_level)
 
 
 def get_available_cubin_files(source, retries=3, delay=5, timeout=10):
@@ -72,11 +133,9 @@ class MetaInfoHash:
     )
 
 
-def download_artifacts() -> bool:
-    env_backup = os.environ.get("FLASHINFER_CUBIN_CHECKSUM_DISABLED", None)
-    os.environ["FLASHINFER_CUBIN_CHECKSUM_DISABLED"] = "1"
+def get_cubin_file_list():
     cubin_files = [
-        (ArtifactPath.TRTLLM_GEN_FMHA + "flashInferMetaInfo", ".h"),
+        (ArtifactPath.TRTLLM_GEN_FMHA + "include/flashInferMetaInfo", ".h"),
         (ArtifactPath.TRTLLM_GEN_GEMM + "include/flashinferMetaInfo", ".h"),
         (ArtifactPath.TRTLLM_GEN_BMM + "include/flashinferMetaInfo", ".h"),
     ]
@@ -92,19 +151,54 @@ def download_artifacts() -> bool:
                 FLASHINFER_CUBINS_REPOSITORY + "/" + kernel
             )
         ]
-    pool = ThreadPoolExecutor(4)
-    futures = []
+    return cubin_files
+
+
+def download_artifacts():
+    import tqdm
+
+    with temp_env_var("FLASHINFER_CUBIN_CHECKSUM_DISABLED", "1"):
+        cubin_files = get_cubin_file_list()
+        num_threads = int(os.environ.get("FLASHINFER_CUBIN_DOWNLOAD_THREADS", "4"))
+        pool = ThreadPoolExecutor(num_threads)
+        futures = []
+        for name, extension in cubin_files:
+            ret = pool.submit(get_cubin, name, "", extension)
+            futures.append(ret)
+        results = []
+        with (
+            patch_logger_for_tqdm(logger),
+            tqdm(total=len(futures), desc="Downloading cubins") as pbar,
+        ):
+            for ret in as_completed(futures):
+                result = ret.result()
+                results.append(result)
+                pbar.update(1)
+        all_success = all(results)
+    if not all_success:
+        raise RuntimeError("Failed to download cubins")
+
+
+def get_artifacts_status():
+    """
+    Check which cubins are already downloaded and return (num_downloaded, total).
+    Does not download any cubins.
+    """
+    cubin_files = get_cubin_file_list()
+    status = []
     for name, extension in cubin_files:
-        ret = pool.submit(get_cubin, name, "", extension)
-        futures.append(ret)
-    results = []
-    for ret in as_completed(futures):
-        result = ret.result()
-        results.append(result)
-    all_success = all(results)
-    if not env_backup:
-        os.environ.pop("FLASHINFER_CUBIN_CHECKSUM_DISABLED")
+        # get_cubin stores cubins in FLASHINFER_CUBIN_DIR with the same relative path
+        # Remove any leading slashes from name
+        rel_path = name.lstrip("/")
+        local_path = os.path.join(FLASHINFER_CUBIN_DIR, rel_path)
+        exists = os.path.isfile(local_path + extension)
+        status.append((name, extension, exists))
+    return status
+
+
+def clear_cubin():
+    if os.path.exists(FLASHINFER_CUBIN_DIR):
+        print(f"Clearing cubin directory: {FLASHINFER_CUBIN_DIR}")
+        shutil.rmtree(FLASHINFER_CUBIN_DIR)
     else:
-        os.environ["FLASHINFER_CUBIN_CHECKSUM_DISABLED"] = env_backup
-
-    return all_success
+        print(f"Cubin directory does not exist: {FLASHINFER_CUBIN_DIR}")

flashinfer/jit/cpp_ext.py

@@ -22,6 +22,18 @@
 from ..compilation_context import CompilationContext
 
 
+@functools.cache
+def get_cuda_path() -> str:
+    if CUDA_HOME is None:
+        # get output of "which nvcc"
+        result = subprocess.run(["which", "nvcc"], capture_output=True)
+        if result.returncode != 0:
+            raise RuntimeError("Could not find nvcc")
+        return result.stdout.decode("utf-8").strip()
+    else:
+        return CUDA_HOME
+
+
 @functools.cache
 def get_cuda_version() -> Version:
     if CUDA_HOME is None:

setup.py

@@ -60,6 +60,9 @@ def generate_build_meta(aot_build_meta: dict) -> None:
     "requests",
     "pynvml",
     "einops",
+    "click",
+    "tqdm",
+    "tabulate",
     "packaging>=24.2",
     "nvidia-cudnn-frontend>=1.13.0",
 ]