| |
|
|
| """Unified FlashAttention logging controlled by a single ``FA_LOG_LEVEL`` env var. |
| |
| Host-side messages go through Python ``logging`` (logger name ``flash_attn``). |
| A default ``StreamHandler`` is attached automatically when ``FA_LOG_LEVEL >= 1`` |
| so that standalone scripts get output without extra setup; applications that |
| configure their own logging can remove or replace it via the standard API. |
| |
| FA_LOG_LEVEL mapping:: |
| |
| 0 off nothing logged |
| 1 host host-side summaries only (no kernel printf) |
| 2 kernel host + curated kernel traces |
| 3 max host + all kernel traces (noisy, perf hit) |
| |
| Set via environment variable:: |
| |
| FA_LOG_LEVEL=1 python train.py |
| |
| Device-side ``cute.printf`` calls are compile-time eliminated via |
| ``cutlass.const_expr`` when the log level is below the callsite threshold, |
| so there is zero performance cost when device logging is off. |
| Changing the log level after kernel compilation requires a recompile |
| (the level participates in the forward compile key). |
| """ |
|
|
| import logging |
| import os |
| import sys |
|
|
| import cutlass.cute as cute |
| from cutlass import const_expr |
|
|
| _LOG_LEVEL_NAMES = {"off": 0, "host": 1, "kernel": 2, "max": 3} |
|
|
|
|
| def _parse_log_level(raw: str) -> int: |
| if raw in _LOG_LEVEL_NAMES: |
| return _LOG_LEVEL_NAMES[raw] |
| try: |
| level = int(raw) |
| except ValueError: |
| return 0 |
| return max(0, min(level, 3)) |
|
|
|
|
| _fa_log_level: int = _parse_log_level(os.environ.get("FA_LOG_LEVEL", "0")) |
|
|
| _logger = logging.getLogger("flashrt_fa4") |
| _logger.addHandler(logging.NullHandler()) |
| _default_handler: logging.Handler | None = None |
|
|
|
|
| def _configure_default_handler() -> None: |
| global _default_handler |
| if _fa_log_level >= 1: |
| if _default_handler is None: |
| _default_handler = logging.StreamHandler(sys.stdout) |
| _default_handler.setFormatter(logging.Formatter("[FA] %(message)s")) |
| _logger.addHandler(_default_handler) |
| _logger.setLevel(logging.DEBUG) |
| else: |
| if _default_handler is not None: |
| _logger.removeHandler(_default_handler) |
| _default_handler = None |
| _logger.setLevel(logging.WARNING) |
|
|
|
|
| _configure_default_handler() |
|
|
|
|
| def get_fa_log_level() -> int: |
| return _fa_log_level |
|
|
|
|
| def set_fa_log_level(level: int | str) -> None: |
| """Set the FA log level programmatically. |
| |
| Host logging takes effect immediately. Device logging changes only |
| affect kernels compiled after this call (new compile-key selection). |
| """ |
| global _fa_log_level |
| if isinstance(level, str): |
| level = _parse_log_level(level) |
| _fa_log_level = max(0, min(int(level), 3)) |
| _configure_default_handler() |
|
|
|
|
| def fa_log(level: int, msg: str): |
| if _fa_log_level >= level: |
| _logger.info(msg) |
|
|
|
|
| def fa_printf(level: int, fmt, *args): |
| if const_expr(_fa_log_level >= level): |
| cute.printf(fmt, *args) |
|
|
