Function nvtx_range_debug

tensorrt_llm/_utils.py:922–946

Creates an NVTX range annotation for debugging purposes. Similar to nvtx_range, but only creates the range if specific environment variables are set, making it suitable for debug profiling. Args: msg (str): The message/name for the NVTX range. color (str, optional)

(msg: str,
                     color: str = "grey",
                     domain: str = "TensorRT-LLM",
                     category: Optional[str] = None)

def nvtx_range_debug(msg: str,
                     color: str = "grey",
                     domain: str = "TensorRT-LLM",
                     category: Optional[str] = None):
    """
    Creates an NVTX range annotation for debugging purposes.

    Similar to nvtx_range, but only creates the range if specific environment
    variables are set, making it suitable for debug profiling.

    Args:
        msg (str): The message/name for the NVTX range.
        color (str, optional): The color to use for the range in the profiler. Defaults to "grey".
        domain (str, optional): The domain name for the range. Defaults to "TensorRT-LLM".
        category (str, optional): The category for the range. Defaults to None.

    Returns:
        contextmanager: A context manager that either marks the NVTX range if enabled,
                        or a null context manager that does nothing if disabled.
    """
    if os.getenv("TLLM_LLMAPI_ENABLE_NVTX", "0") == "1" or \
            os.getenv("TLLM_NVTX_DEBUG", "0") == "1":
        return nvtx_range(msg, color=color, domain=domain, category=category)
    else:
        return _null_context_manager()


def nvtx_mark_debug(msg: str,