From 628591b6207d38bf10d70ca527bb41fe2acdd012 Mon Sep 17 00:00:00 2001
From: Davids048 <jundasu@ucsd.edu>
Date: Fri, 21 Aug 2026 05:18:44 +0000
Subject: [PATCH] Add some logging.

---
 fastvideo/attention/backends/flash_attn.py  |  7 ++++++-
 fastvideo/models/loader/component_loader.py |  7 ++++++-
 fastvideo/worker/gpu_worker.py              | 11 +++++++++++
 3 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/fastvideo/attention/backends/flash_attn.py b/fastvideo/attention/backends/flash_attn.py
index 8aad6a5b0..5f7dce691 100644
--- a/fastvideo/attention/backends/flash_attn.py
+++ b/fastvideo/attention/backends/flash_attn.py
@@ -19,7 +19,12 @@ from fastvideo.attention.backends.abstract import (
 from fastvideo.logger import init_logger
 
 logger = init_logger(__name__)
-logger.info("Using FlashAttention-%s backend", fa_version)
+# Every worker records the loaded FlashAttention implementation so a
+# distributed profiling log contains one backend receipt per rank.
+logger.info("Worker %s Using FlashAttention-%s backend",
+            os.environ.get("RANK", "0"),
+            fa_version,
+            local_main_process_only=False)
 
 # FP4 FA4 support: quantize Q/K to NVFP4 E2M1 for block-scaled MMA on Blackwell.
 # Requires: flash-attention-fp4, flashinfer, cutlass-dsl. Enable via nvfp4_fa4=True kwarg.
diff --git a/fastvideo/models/loader/component_loader.py b/fastvideo/models/loader/component_loader.py
index f11f15255..7f7ff2bf0 100644
--- a/fastvideo/models/loader/component_loader.py
+++ b/fastvideo/models/loader/component_loader.py
@@ -1057,7 +1057,12 @@ class TransformerLoader(ComponentLoader):
             # so recording here makes the decision readable from the loaded
             # transformer — and records the narrowed one for teacher/critic.
             resolved = record_resolved_attention_backend(dit_config)
-            logger.info("transformer attention backend: %s", resolved.name if resolved else "automatic selection")
+            # Every worker records its resolved backend so distributed profile
+            # snapshots can prove that all ranks use the requested kernels.
+            logger.info("Worker %s transformer attention backend: %s",
+                        os.environ.get("RANK", "0"),
+                        resolved.name if resolved else "automatic selection",
+                        local_main_process_only=False)
             model = maybe_load_fsdp_model(
                 model_cls=model_cls,
                 init_params={
diff --git a/fastvideo/worker/gpu_worker.py b/fastvideo/worker/gpu_worker.py
index b4f8e6227..dcaac6e43 100644
--- a/fastvideo/worker/gpu_worker.py
+++ b/fastvideo/worker/gpu_worker.py
@@ -4,6 +4,7 @@ from typing import Any, cast
 
 import torch
 
+import fastvideo.envs as envs
 from fastvideo.distributed import (cleanup_dist_env_and_memory, maybe_init_distributed_environment_and_model_parallel)
 from fastvideo.distributed.parallel_state import get_local_torch_device
 from fastvideo.fastvideo_args import FastVideoArgs
@@ -13,6 +14,14 @@ from fastvideo.pipelines import ForwardBatch, LoRAPipeline, build_pipeline
 logger = init_logger(__name__)
 
 
+def _log_cuda_device_uuid(rank: int, device: torch.device) -> None:
+    """Record an NVIDIA worker UUID when external NVTX profiling is enabled."""
+    if not envs.FASTVIDEO_NVTX_PROFILE:
+        return
+    device_uuid = torch.cuda.get_device_properties(device).uuid
+    logger.info("Worker %d CUDA device UUID: GPU-%s", rank, device_uuid, local_main_process_only=False)
+
+
 class Worker:
 
     def __init__(self, fastvideo_args: FastVideoArgs, local_rank: int, rank: int, distributed_init_method: str):
@@ -61,6 +70,8 @@ class Worker:
         if current_platform.is_cuda_alike():
             torch.cuda.set_device(self.device)
             self.init_gpu_memory = torch.cuda.mem_get_info(self.device)[0]
+            if current_platform.is_cuda():
+                _log_cuda_device_uuid(self.rank, self.device)
         else:
             # For MPS, we can't get memory info the same way
             self.init_gpu_memory = 0