[perf] Optimize MiniMax H3 VAE decoding (#1734)

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.
@@ -223,7 +228,15 @@ class FlashAttentionImpl(AttentionImpl):
     ) -> None:
         self.causal = causal
         self.softmax_scale = softmax_scale
-        self.nvfp4_fa4 = extra_impl_args.get("nvfp4_fa4", False) or os.environ.get("FASTVIDEO_NVFP4_FA4", "0") == "1"
+        # An explicit ``nvfp4_fa4`` impl arg wins over the process-wide
+        # FASTVIDEO_NVFP4_FA4 env opt-in, so precision-sensitive layers (e.g.
+        # the FP32-pinned H3 VAE attention) can force-disable FP4 Q/K
+        # quantization while the DiT keeps it. When the arg is absent the env
+        # keeps its previous semantics.
+        nvfp4_fa4 = extra_impl_args.get("nvfp4_fa4")
+        if nvfp4_fa4 is None:
+            nvfp4_fa4 = os.environ.get("FASTVIDEO_NVFP4_FA4", "0") == "1"
+        self.nvfp4_fa4 = bool(nvfp4_fa4)
         if self.nvfp4_fa4:
             cap = torch.cuda.get_device_capability()
             assert cap in [(10, 0), (10, 3)], (f"NVFP4 FA4 requires Blackwell (sm100a/sm103a), got sm{cap[0]}{cap[1]}")

fastvideo/envs.py

@@ -28,6 +28,7 @@ if TYPE_CHECKING:
     NVCC_THREADS: str | None = None
     CMAKE_BUILD_TYPE: str | None = None
     VERBOSE: bool = False
+    FASTVIDEO_NVTX_PROFILE: bool = False
     FASTVIDEO_TORCH_PROFILER_DIR: str | None = None
     FASTVIDEO_TORCH_PROFILER_RECORD_SHAPES: bool = False
     FASTVIDEO_TORCH_PROFILER_WITH_PROFILE_MEMORY: bool = False
@@ -229,6 +230,10 @@ environment_variables: dict[str, Callable[[], Any]] = {
     "FASTVIDEO_WORKER_MULTIPROC_METHOD":
     lambda: os.getenv("FASTVIDEO_WORKER_MULTIPROC_METHOD", "spawn"),
 
+    # Emit lightweight NVTX ranges for external profilers such as Nsight Systems.
+    "FASTVIDEO_NVTX_PROFILE":
+    lambda: os.getenv("FASTVIDEO_NVTX_PROFILE", "0") != "0",
+
     # Enables torch profiler if set. Path to the directory where torch profiler
     # traces are saved. Note that it must be an absolute path.
     "FASTVIDEO_TORCH_PROFILER_DIR":

fastvideo/models/dits/minimax_h3.py

@@ -34,6 +34,7 @@ from fastvideo.models.dits.minimax_h3_fusions import (
     minimax_h3_swiglu,
 )
 from fastvideo.platforms import AttentionBackendEnum
+from fastvideo.profiler import nvtx_range
 from fastvideo.utils import get_compute_dtype
 
 logger = init_logger(__name__)
@@ -840,14 +841,17 @@ class MiniMaxH3Transformer3DModel(BaseDiT):
             local_timestep_indices, _ = sequence_model_parallel_shard(local_timestep_indices, dim=0)
             rotary_emb = (rotary_cos, rotary_sin)
 
-        for block in self.transformer_blocks:
-            packed_hidden_states = block(
-                packed_hidden_states,
-                temb,
-                adaln_indices,
-                rotary_emb,
-                original_seq_len,
-            )
+        # The eager driver owns profiling markers while each block's compiled
+        # forward owns the graph that the marker surrounds.
+        for block_index, block in enumerate(self.transformer_blocks):
+            with nvtx_range(f"minimax_h3.transformer_block.{block_index}"):
+                packed_hidden_states = block(
+                    packed_hidden_states,
+                    temb,
+                    adaln_indices,
+                    rotary_emb,
+                    original_seq_len,
+                )
 
         packed_hidden_states = self.norm_out(
             packed_hidden_states,

fastvideo/models/loader/component_loader.py

@@ -1098,7 +1098,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={

fastvideo/models/vaes/minimax_h3_audio.py

@@ -392,9 +392,16 @@ class MiniMaxH3AudioBigVGANDecoder(nn.Module):
         return torch.clamp(hidden_states, min=-1.0, max=1.0)
 
 
+def _is_minimax_h3_audio_vae_decoder(name: str, submodule: nn.Module) -> bool:
+    """Select the audio decoder that serves the H3 VAE ``decode`` path."""
+    return name == "decoder" and isinstance(submodule, MiniMaxH3AudioBigVGANDecoder)
+
+
 class MiniMaxH3AudioVAE(nn.Module):
     """DAC encoder plus BigVGAN decoder for mono 32 kHz waveforms."""
 
+    _compile_conditions = [_is_minimax_h3_audio_vae_decoder]
+
     def __init__(self, config: MiniMaxH3AudioVAEConfig):
         super().__init__()
         self.config = config

fastvideo/models/vaes/minimax_h3_video.py

@@ -15,7 +15,10 @@ import torch.nn as nn
 import torch.nn.functional as F
 from torch.utils.checkpoint import checkpoint
 
+from fastvideo.attention import get_attn_backend
 from fastvideo.configs.models.vaes.minimax_h3_video import MiniMaxH3VideoVAEConfig
+from fastvideo.platforms import AttentionBackendEnum
+from fastvideo.profiler import nvtx_range
 
 
 class DiagonalGaussianDistribution:
@@ -292,6 +295,7 @@ class MiniMaxH3VideoRotaryPosEmbed(nn.Module):
 class MiniMaxH3VideoAttention(nn.Module):
 
     def __init__(self, dim: int, heads: int, dim_head: int, eps: float = 1e-5, bias: bool = True) -> None:
+        """Build projections and the selected dense FastVideo attention implementation."""
         super().__init__()
         self.heads = heads
         self.dim_head = dim_head
@@ -303,12 +307,38 @@ class MiniMaxH3VideoAttention(nn.Module):
         self.to_k = nn.Linear(dim, inner_dim, bias=bias)
         self.to_v = nn.Linear(dim, inner_dim, bias=bias)
         self.to_out = nn.ModuleList([nn.Linear(inner_dim, dim, bias=bias), nn.Dropout(0.0)])
+        self.attn_impl = None
+        from fastvideo.platforms import current_platform
+
+        if current_platform.is_cuda_alike():
+            attention_backend = get_attn_backend(
+                dim_head,
+                # FlashAttention executes the FP32 VAE activations in BF16 and
+                # restores FP32 output, so resolve against the kernel dtype.
+                torch.bfloat16,
+                supported_attention_backends=(
+                    AttentionBackendEnum.TORCH_SDPA,
+                    AttentionBackendEnum.FLASH_ATTN,
+                ),
+            )
+            self.attn_impl = attention_backend.get_impl_cls()(
+                num_heads=heads,
+                head_size=dim_head,
+                softmax_scale=dim_head**-0.5,
+                num_kv_heads=heads,
+                causal=False,
+                # The FASTVIDEO_NVFP4_FA4 env opt-in targets the DiT; this VAE
+                # is FP32-pinned (_keep_in_fp32_modules), so force-disable FP4
+                # Q/K quantization for its attention regardless of the env.
+                nvfp4_fa4=False,
+            )
 
     def forward(
         self,
         hidden_states: torch.Tensor,
         rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
     ) -> torch.Tensor:
+        """Apply dense self-attention to one spatial VAE token sequence."""
         query = self.to_q(hidden_states).unflatten(2, (self.heads, -1))
         key = self.to_k(hidden_states).unflatten(2, (self.heads, -1))
         value = self.to_v(hidden_states).unflatten(2, (self.heads, -1))
@@ -329,9 +359,17 @@ class MiniMaxH3VideoAttention(nn.Module):
             query = torch.cat([query_rotary * cos + query_rotated * sin, query_pass], dim=-1)
             key = torch.cat([key_rotary * cos + key_rotated * sin, key_pass], dim=-1)
 
-        query, key, value = (tensor.permute(0, 2, 1, 3) for tensor in (query, key, value))
-        hidden_states = F.scaled_dot_product_attention(query, key, value)
-        hidden_states = hidden_states.permute(0, 2, 1, 3).flatten(2, 3)
+        if self.attn_impl is not None and query.device.type != "cpu":
+            # VAE decoding has no diffusion-step metadata, so call the selected
+            # backend implementation directly with dense BSHD tensors.
+            hidden_states = self.attn_impl.forward(query, key, value, None)
+            hidden_states = hidden_states.flatten(2, 3)
+        else:
+            # Keep CPU construction and execution available without requiring
+            # an accelerator attention backend.
+            query, key, value = (tensor.permute(0, 2, 1, 3) for tensor in (query, key, value))
+            hidden_states = F.scaled_dot_product_attention(query, key, value)
+            hidden_states = hidden_states.permute(0, 2, 1, 3).flatten(2, 3)
         return self.to_out[0](hidden_states)
 
 
@@ -434,6 +472,7 @@ class MiniMaxH3VideoViTDecoder3d(nn.Module):
         self.gradient_checkpointing = False
 
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        """Decode one latent spatial input through the H3 video transformer."""
         batch_size, num_channels, num_frames, height, width = hidden_states.shape
         hidden_states = hidden_states.permute(0, 2, 3, 4, 1).reshape(
             batch_size,
@@ -483,6 +522,11 @@ class MiniMaxH3VideoViTDecoder3d(nn.Module):
         )
 
 
+def _is_minimax_h3_video_vae_decoder(name: str, submodule: nn.Module) -> bool:
+    """Select the video decoder that serves the H3 VAE ``decode`` path."""
+    return name == "decoder" and isinstance(submodule, MiniMaxH3VideoViTDecoder3d)
+
+
 class AutoencoderKLMiniMaxH3(nn.Module):
     """MiniMax-H3 causal encoder and ViT decoder with exact release geometry."""
 
@@ -490,6 +534,10 @@ class AutoencoderKLMiniMaxH3(nn.Module):
     _no_split_modules = ["MiniMaxH3VideoResnetBlock3d", "MiniMaxH3VideoTransformerBlock"]
     _repeated_blocks = ["MiniMaxH3VideoTransformerBlock"]
     _keep_in_fp32_modules = ["encoder", "decoder", "quant_conv", "post_quant_conv"]
+    _compile_conditions = [_is_minimax_h3_video_vae_decoder]
+    # ``prepare_for_compile`` flips this per instance when the pipeline opts
+    # into ``enable_torch_compile_vae``; default instances stay fully eager.
+    _tile_helpers_compiled = False
 
     def __init__(self, config: MiniMaxH3VideoVAEConfig) -> None:
         super().__init__()
@@ -655,12 +703,40 @@ class AutoencoderKLMiniMaxH3(nn.Module):
         slice_rest[dim] = slice(blend_extent, None)
         return torch.cat([blended, b[tuple(slice_rest)]], dim=dim)
 
+    def prepare_for_compile(self) -> None:
+        """Compile the fixed-shape tile helpers for the opt-in VAE compile path.
+
+        ``ComposedPipelineBase._maybe_compile_pipeline_module`` calls this hook
+        only when ``enable_torch_compile_vae`` is set, right before the decoder
+        is compiled through ``_compile_conditions``. The spatial tile grid and
+        the per-tile decoder-input projection have fixed shapes, so
+        ``mode="reduce-overhead"`` records one CUDA graph per geometry and
+        replays it for every tile and temporal chunk. Keeping this behind the
+        opt-in means default (eager) users pay neither the inductor/triton
+        toolchain requirement and first-decode compile latency nor the
+        permanent cudagraph memory pools, and multi-resolution callers never
+        churn ``dynamic=False`` recompiles they did not ask for.
+        """
+        if self._tile_helpers_compiled:
+            return
+        # The fixed spatial tile grid reuses one compiled blend-and-concatenate graph.
+        self._stitch_tiles = torch.compile(self._stitch_tiles, backend="inductor", mode="reduce-overhead", dynamic=False)
+        # Each fixed-shape latent tile reuses one compiled decoder-input projection.
+        self._project_decoder_tile = torch.compile(
+            self._project_decoder_tile,
+            backend="inductor",
+            mode="reduce-overhead",
+            dynamic=False,
+        )
+        self._tile_helpers_compiled = True
+
     def _stitch_tiles(
         self,
         tiles: list[list[torch.Tensor]],
         height_overlaps: list[int],
         width_overlaps: list[int],
     ) -> torch.Tensor:
+        """Blend decoded tile overlaps and concatenate the spatial canvas."""
         result_rows = []
         for row_index, row in enumerate(tiles):
             result_row = []
@@ -677,6 +753,10 @@ class AutoencoderKLMiniMaxH3(nn.Module):
             result_rows.append(torch.cat(result_row, dim=-1))
         return torch.cat(result_rows, dim=-2)
 
+    def _project_decoder_tile(self, tile: torch.Tensor) -> torch.Tensor:
+        """Project one spatial latent tile into the decoder input channels."""
+        return self.post_quant_conv(tile)
+
     def _encode_clip(self, x: torch.Tensor) -> torch.Tensor:
         if not self.use_tiling:
             return self.quant_conv(self.encoder(x))
@@ -700,36 +780,73 @@ class AutoencoderKLMiniMaxH3(nn.Module):
             rows.append(row)
         latent_y_overlaps = [overlap // self.spatial_compression_ratio for overlap in y_overlaps]
         latent_x_overlaps = [overlap // self.spatial_compression_ratio for overlap in x_overlaps]
-        return self._stitch_tiles(rows, latent_y_overlaps, latent_x_overlaps)
+        stitched = self._stitch_tiles(rows, latent_y_overlaps, latent_x_overlaps)
+        if self._tile_helpers_compiled:
+            # Under the opt-in mode="reduce-overhead" compile the stitched
+            # canvas is a CUDA-graph static buffer that the next _stitch_tiles
+            # replay overwrites. Callers (_encode/_encode_pixels/
+            # encode_keyframe) collect per-clip results across replays before
+            # concatenating, so hand them a caller-owned tensor instead of
+            # cudagraph-pooled storage. Eager instances return the fresh
+            # torch.cat result directly.
+            stitched = stitched.clone()
+        return stitched
 
     def _decode_clip(self, z: torch.Tensor) -> torch.Tensor:
-        if not self.use_tiling:
-            return self.decoder(self.post_quant_conv(z))
-        height = z.shape[-2] * self.spatial_compression_ratio
-        width = z.shape[-1] * self.spatial_compression_ratio
-        y_indices, y_lengths, y_overlaps = self._split_tiles(
-            height,
-            self.tile_sample_min_height,
-            self.tile_sample_min_overlap_height,
-        )
-        x_indices, x_lengths, x_overlaps = self._split_tiles(
-            width,
-            self.tile_sample_min_width,
-            self.tile_sample_min_overlap_width,
-        )
-        ratio = self.spatial_compression_ratio
-        rows = []
-        for y_position, y_length in zip(y_indices, y_lengths):
-            row = []
-            for x_position, x_length in zip(x_indices, x_lengths):
-                tile = z[
-                    ...,
-                    y_position // ratio:y_position // ratio + y_length // ratio,
-                    x_position // ratio:x_position // ratio + x_length // ratio,
-                ]
-                row.append(self.decoder(self.post_quant_conv(tile)))
-            rows.append(row)
-        return self._stitch_tiles(rows, y_overlaps, x_overlaps)
+        """Decode one temporal clip, with optional overlapping spatial tiles."""
+        with nvtx_range("minimax_h3.vae.decode_clip"):
+            if not self.use_tiling:
+                with nvtx_range("minimax_h3.vae.decode_clip.no_s_tile.post_quant_conv"):
+                    projected_clip = self.post_quant_conv(z)
+                with nvtx_range("minimax_h3.vae.decode_clip.no_s_tile.decoder_forward"):
+                    return self.decoder(projected_clip)
+
+            height = z.shape[-2] * self.spatial_compression_ratio
+            width = z.shape[-1] * self.spatial_compression_ratio
+            with nvtx_range("minimax_h3.vae.decode_clip.split_tiles"):
+                y_indices, y_lengths, y_overlaps = self._split_tiles(
+                    height,
+                    self.tile_sample_min_height,
+                    self.tile_sample_min_overlap_height,
+                )
+                x_indices, x_lengths, x_overlaps = self._split_tiles(
+                    width,
+                    self.tile_sample_min_width,
+                    self.tile_sample_min_overlap_width,
+                )
+
+            ratio = self.spatial_compression_ratio
+            rows = []
+            # The eager tile driver owns NVTX so each marker remains outside
+            # the compiled decoder graph.
+            with nvtx_range("minimax_h3.vae.decode_clip.decode_tiles"):
+                for row_index, (y_position, y_length) in enumerate(zip(y_indices, y_lengths)):
+                    row = []
+                    for column_index, (x_position, x_length) in enumerate(zip(x_indices, x_lengths)):
+                        with nvtx_range(f"minimax_h3.vae.decode_clip.tile.{row_index}.{column_index}"):
+                            tile = z[
+                                ...,
+                                y_position // ratio:y_position // ratio + y_length // ratio,
+                                x_position // ratio:x_position // ratio + x_length // ratio,
+                            ]
+                            projected_tile = self._project_decoder_tile(tile)
+                            with nvtx_range("minimax_h3.vae.decode_clip.tile.decoder_forward"):
+                                decoded_tile = self.decoder(projected_tile)
+                            row.append(decoded_tile)
+                    rows.append(row)
+
+            with nvtx_range("minimax_h3.vae.decode_clip.stitch_tiles"):
+                stitched = self._stitch_tiles(rows, y_overlaps, x_overlaps)
+                if self._tile_helpers_compiled:
+                    # Same CUDA-graph output-ownership contract as
+                    # _encode_clip: _decode collects chunks across
+                    # _stitch_tiles replays before torch.cat, so the pooled
+                    # canvas must not escape this driver. (The streaming
+                    # _decode_to_pixels path copies each chunk out before the
+                    # next decode; the clone keeps it correct too at one D2D
+                    # copy per chunk.)
+                    stitched = stitched.clone()
+                return stitched
 
     def _encode(self, x: torch.Tensor) -> torch.Tensor:
         clip_length = self.config.clip_length
@@ -809,21 +926,27 @@ class AutoencoderKLMiniMaxH3(nn.Module):
 
         output_frame_start = 0
         overlap = None
-        for index in range(num_chunks):
-            start = index * tokens_chunk_size
-            clip = self._decode_clip(z[:, :, start:start + tokens_chunk_size + self.token_overlap])
-            chunk = clip[:, :, self.frame_pre_padding:chunk_num_frames]
-            next_overlap = None
-            if self.config.token_drop > 0:
-                next_overlap = clip[:, :, chunk_num_frames + self.frame_pre_padding:].clone()
-            if overlap is not None:
-                chunk = self._blend(overlap, chunk, self.frame_overlap, dim=-3)
+        for chunk_index in range(num_chunks):
+            with nvtx_range(f"minimax_h3.vae.temporal_chunk.{chunk_index}"):
+                start = chunk_index * tokens_chunk_size
+                clip = self._decode_clip(z[:, :, start:start + tokens_chunk_size + self.token_overlap])
+                with nvtx_range(f"minimax_h3.vae.temporal_chunk.{chunk_index}.frame_segment.0"):
+                    chunk = clip[:, :, self.frame_pre_padding:chunk_num_frames]
+                    if overlap is not None:
+                        chunk = self._blend(overlap, chunk, self.frame_overlap, dim=-3)
+                    num_frames = min(chunk.shape[2], output_num_frames - output_frame_start)
+                    chunk = chunk[:, :, :num_frames]
 
-            num_frames = min(chunk.shape[2], output_num_frames - output_frame_start)
-            if num_frames > 0:
-                yield chunk[:, :, :num_frames]
-                output_frame_start += num_frames
+                next_overlap = None
+                if self.config.token_drop > 0:
+                    with nvtx_range(f"minimax_h3.vae.temporal_chunk.{chunk_index}.frame_segment.1"):
+                        next_overlap = clip[:, :, chunk_num_frames + self.frame_pre_padding:].clone()
+
+            # Yield after the ranges close so consumer-side CPU copies do not inflate decoder timing.
             overlap = next_overlap
+            if num_frames > 0:
+                output_frame_start += num_frames
+                yield chunk
 
         if overlap is not None and output_frame_start < output_num_frames:
             yield overlap[:, :, :output_num_frames - output_frame_start]

fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_conditioning.py

@@ -12,6 +12,7 @@ from torch.distributed.tensor import DTensor
 from fastvideo.distributed import get_local_torch_device
 from fastvideo.fastvideo_args import FastVideoArgs
 from fastvideo.models.encoders.minimax_h3_qwen3_vl import MiniMaxH3Qwen3VLConditioner
+from fastvideo.profiler import nvtx_range
 from fastvideo.pipelines.basic.minimax_h3.packing import (
     MINIMAX_H3_IMAGE_PAD_TOKEN,
     MINIMAX_H3_TEXT_TAG,
@@ -256,6 +257,7 @@ class MiniMaxH3ConditioningStage(PipelineStage):
 
     @torch.no_grad()
     def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        """Encode one H3 prompt presentation and attach its packed text features."""
         device = get_local_torch_device()
         first_param = next(self.conditioner.parameters(), None)
         moved_for_forward = (fastvideo_args.text_encoder_cpu_offload and first_param is not None
@@ -263,10 +265,13 @@ class MiniMaxH3ConditioningStage(PipelineStage):
         if moved_for_forward:
             self.conditioner.to(device)
         try:
-            if self.ref2va:
-                prompt_embeds, text_token_tags = self._encode_ref2va(batch, device)
-            else:
-                prompt_embeds, text_token_tags = self._encode_fl2va(batch, device)
+            # Keep both H3 prompt-presentation modes under one text-encoding
+            # range so Nsight Systems exposes their complete conditioning cost.
+            with nvtx_range("minimax_h3.text_encoding"):
+                if self.ref2va:
+                    prompt_embeds, text_token_tags = self._encode_ref2va(batch, device)
+                else:
+                    prompt_embeds, text_token_tags = self._encode_fl2va(batch, device)
         finally:
             if moved_for_forward:
                 self.conditioner.to("cpu")

fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_decoding.py

@@ -11,6 +11,7 @@ from fastvideo.distributed import get_local_torch_device, get_world_group, model
 from fastvideo.fastvideo_args import FastVideoArgs
 from fastvideo.models.vaes.minimax_h3_audio import MiniMaxH3AudioVAE
 from fastvideo.models.vaes.minimax_h3_video import AutoencoderKLMiniMaxH3
+from fastvideo.profiler import nvtx_range
 from fastvideo.pipelines.basic.minimax_h3.packing import (
     MiniMaxH3PackedLayout,
     unpack_audio_tokens,
@@ -55,6 +56,7 @@ class MiniMaxH3VideoDecodingStage(PipelineStage):
 
     @torch.no_grad()
     def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        """Decode H3 video latents into normalized CPU pixels."""
         if model_parallel_is_initialized() and not get_world_group().is_first_rank:
             # Distributed executors consume rank 0's ForwardBatch. Keep a
             # verifier-compatible placeholder on other ranks and avoid
@@ -88,8 +90,12 @@ class MiniMaxH3VideoDecodingStage(PipelineStage):
                 dtype=torch.float32,
                 pin_memory=fastvideo_args.pin_cpu_memory and is_pin_memory_available(),
             )
-            # The published decode recipe uses FP16 autocast over FP32 weights.
-            with torch.autocast(device_type=device.type, dtype=torch.float16, enabled=device.type == "cuda"):
+            # Attribute the streamed decoder computation while retaining
+            # per-chunk device-to-host transfer and pinned-buffer reuse.
+            with (
+                    nvtx_range("minimax_h3.vae"),
+                    torch.autocast(device_type=device.type, dtype=torch.float16, enabled=device.type == "cuda"),
+            ):
                 self.vae.decode_to_pixels(latents, output)
             batch.output = output
             return batch
@@ -121,6 +127,7 @@ class MiniMaxH3AudioDecodingStage(PipelineStage):
 
     @torch.no_grad()
     def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        """Decode H3 audio latents into a stereo CPU waveform."""
         if model_parallel_is_initialized() and not get_world_group().is_first_rank:
             batch.extra["audio"] = torch.empty((0, 2), device="cpu", dtype=torch.float32)
             batch.extra["audio_sample_rate"] = self.audio_vae.sampling_rate
@@ -144,7 +151,10 @@ class MiniMaxH3AudioDecodingStage(PipelineStage):
                 self._clear_runtime(batch)
                 return batch
 
-            decoded = self.audio_vae.decode(latents).sample.float()
+            # The range isolates waveform synthesis from packing and runtime
+            # cleanup so the audio decoder has one stable timeline boundary.
+            with nvtx_range("minimax_h3.audio_vae"):
+                decoded = self.audio_vae.decode(latents).sample.float()
             if decoded.ndim != 3 or decoded.shape[0] != 2 or decoded.shape[1] != 1:
                 raise ValueError("MiniMax-H3 audio VAE must decode stereo channels as two mono batch items; "
                                  f"got {tuple(decoded.shape)}.")

fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_denoising.py

@@ -11,8 +11,8 @@ from fastvideo.attention.selector import component_attention_backend, get_attn_b
 from fastvideo.distributed import get_local_torch_device
 from fastvideo.fastvideo_args import FastVideoArgs
 from fastvideo.forward_context import set_forward_context
-from fastvideo.profiler import profiler_region
 from fastvideo.hooks.activation_trace import trace_step
+from fastvideo.profiler import nvtx_range, profiler_region
 from fastvideo.pipelines.basic.minimax_h3.packing import (
     MINIMAX_H3_KEYFRAME_NOISE_AUG,
     MiniMaxH3PackedLayout,
@@ -89,6 +89,7 @@ class MiniMaxH3DenoisingStage(PipelineStage):
 
     @torch.no_grad()
     def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        """Denoise the packed H3 video and audio streams over one shared schedule."""
         layout = batch.extra.get(MINIMAX_H3_LAYOUT_KEY)
         if not isinstance(layout, MiniMaxH3PackedLayout):
             raise ValueError("MiniMax-H3 packed layout is missing before denoising.")
@@ -151,7 +152,9 @@ class MiniMaxH3DenoisingStage(PipelineStage):
             vsa_tile_size = int(fastvideo_args.VSA_tile_size)
 
         try:
-            with profiler_region("inference_denoising"):
+            # The stage range groups the complete denoising loop while the
+            # indexed model ranges retain timing detail for every H3 block.
+            with profiler_region("inference_denoising"), nvtx_range("minimax_h3.dit"):
                 for index, (video_timestep,
                             audio_timestep) in enumerate(zip(video_timesteps, audio_timesteps, strict=True)):
                     unique_timesteps, timestep_indices = row_timestep_plan[index]

fastvideo/profiler.py

@@ -38,6 +38,25 @@ logger = init_logger(__name__)
 _GLOBAL_CONTROLLER: TorchProfilerController | None = None
 
 
+@contextlib.contextmanager
+def nvtx_range(name: str):
+    """Emit one optional NVTX range for an external CUDA profiler.
+
+    ``FASTVIDEO_NVTX_PROFILE=1`` enables the marker. The context manager stays
+    a no-op without CUDA so call sites can remain shared with CPU tests.
+    """
+    enabled = envs.FASTVIDEO_NVTX_PROFILE and torch.cuda.is_available()
+    if not enabled:
+        yield
+        return
+
+    torch.cuda.nvtx.range_push(name)
+    try:
+        yield
+    finally:
+        torch.cuda.nvtx.range_pop()
+
+
 @dataclass(frozen=True)
 class ProfilerRegion:
     """Metadata describing a profiler region."""

fastvideo/tests/attention/test_flash_attn_nvfp4_opt_out.py

@@ -0,0 +1,66 @@
+# SPDX-License-Identifier: Apache-2.0
+"""NVFP4 FA4 opt-in/opt-out precedence for ``FlashAttentionImpl``.
+
+The process-wide ``FASTVIDEO_NVFP4_FA4=1`` env opt-in targets the DiT. An
+explicit ``nvfp4_fa4`` impl arg must win over the env in both directions so
+precision-sensitive layers (e.g. the FP32-pinned MiniMax-H3 VAE attention)
+can force-disable FP4 Q/K quantization while the DiT keeps it.
+"""
+
+from unittest.mock import patch
+
+import pytest
+
+# The backend module needs a flash-attention package (FA2/FA3, or FA4 with
+# FASTVIDEO_FA4=1) at import. exc_type=ImportError also covers partial
+# installs (e.g. an FA4-only environment without the opt-in env set raises
+# ImportError rather than ModuleNotFoundError).
+flash_attn_module = pytest.importorskip(
+    "fastvideo.attention.backends.flash_attn",
+    reason="no usable flash-attention package installed",
+    exc_type=ImportError,
+)
+FlashAttentionImpl = flash_attn_module.FlashAttentionImpl
+
+
+def _build_impl(**extra_impl_args) -> FlashAttentionImpl:
+    return FlashAttentionImpl(
+        num_heads=2,
+        head_size=64,
+        causal=False,
+        softmax_scale=0.125,
+        num_kv_heads=2,
+        **extra_impl_args,
+    )
+
+
+@pytest.fixture
+def fa4_runtime_available():
+    """Satisfy the Blackwell + flash-attention-fp4 asserts without hardware."""
+    with (
+        patch.object(flash_attn_module, "_FA4_FP4_AVAILABLE", True),
+        patch("torch.cuda.get_device_capability", return_value=(10, 0)),
+    ):
+        yield
+
+
+def test_nvfp4_disabled_by_default(monkeypatch) -> None:
+    monkeypatch.delenv("FASTVIDEO_NVFP4_FA4", raising=False)
+    assert _build_impl().nvfp4_fa4 is False
+
+
+def test_nvfp4_env_opt_in_enables_when_arg_absent(monkeypatch, fa4_runtime_available) -> None:
+    monkeypatch.setenv("FASTVIDEO_NVFP4_FA4", "1")
+    assert _build_impl().nvfp4_fa4 is True
+
+
+def test_explicit_disable_wins_over_env_opt_in(monkeypatch) -> None:
+    """The H3 VAE constructs its impl with ``nvfp4_fa4=False``; the DiT env
+    opt-in must not FP4-quantize Q/K inside the FP32-pinned VAE."""
+    monkeypatch.setenv("FASTVIDEO_NVFP4_FA4", "1")
+    assert _build_impl(nvfp4_fa4=False).nvfp4_fa4 is False
+
+
+def test_explicit_enable_works_without_env(monkeypatch, fa4_runtime_available) -> None:
+    monkeypatch.delenv("FASTVIDEO_NVFP4_FA4", raising=False)
+    assert _build_impl(nvfp4_fa4=True).nvfp4_fa4 is True

fastvideo/tests/contract/test_profiler_regions.py

@@ -15,6 +15,12 @@ import json
 import os
 import subprocess
 import sys
+from unittest.mock import Mock
+
+import pytest
+import torch
+
+from fastvideo.profiler import nvtx_range
 
 # Five-window child: ops before any region, inside a region, between regions,
 # inside a second (short-named) region, after the last region. Exits without
@@ -105,3 +111,73 @@ def test_noop_without_profiler_dir(tmp_path):
     proc = subprocess.run([sys.executable, "-c", child], env=env,
                           capture_output=True, text=True, timeout=300)
     assert proc.returncode == 0, proc.stderr
+
+
+def test_nvtx_range_disabled_is_noop(monkeypatch):
+    """Keep CUDA NVTX untouched when external profiling is disabled."""
+    range_push = Mock()
+    range_pop = Mock()
+    monkeypatch.setenv("FASTVIDEO_NVTX_PROFILE", "0")
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
+    monkeypatch.setattr(torch.cuda.nvtx, "range_push", range_push)
+    monkeypatch.setattr(torch.cuda.nvtx, "range_pop", range_pop)
+
+    with nvtx_range("disabled"):
+        body_executed = True
+
+    assert body_executed is True
+    range_push.assert_not_called()
+    range_pop.assert_not_called()
+
+
+def test_nvtx_range_without_cuda_is_noop(monkeypatch):
+    """Keep NVTX untouched when profiling is enabled on a CPU-only process."""
+    range_push = Mock()
+    range_pop = Mock()
+    monkeypatch.setenv("FASTVIDEO_NVTX_PROFILE", "1")
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: False)
+    monkeypatch.setattr(torch.cuda.nvtx, "range_push", range_push)
+    monkeypatch.setattr(torch.cuda.nvtx, "range_pop", range_pop)
+
+    with nvtx_range("cpu-only"):
+        body_executed = True
+
+    assert body_executed is True
+    range_push.assert_not_called()
+    range_pop.assert_not_called()
+
+
+def test_nvtx_range_enabled_orders_push_body_pop(monkeypatch):
+    """Place the profiled body between one matching NVTX push and pop."""
+    events = []
+    monkeypatch.setenv("FASTVIDEO_NVTX_PROFILE", "1")
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
+    monkeypatch.setattr(torch.cuda.nvtx, "range_push", lambda name: events.append(("push", name)))
+    monkeypatch.setattr(torch.cuda.nvtx, "range_pop", lambda: events.append(("pop", None)))
+
+    with nvtx_range("minimax_h3.test"):
+        events.append(("body", None))
+
+    assert events == [
+        ("push", "minimax_h3.test"),
+        ("body", None),
+        ("pop", None),
+    ]
+
+
+def test_nvtx_range_body_exception_pops_and_propagates(monkeypatch):
+    """Balance the NVTX stack while preserving a body exception."""
+    events = []
+    monkeypatch.setenv("FASTVIDEO_NVTX_PROFILE", "1")
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
+    monkeypatch.setattr(torch.cuda.nvtx, "range_push", lambda name: events.append(("push", name)))
+    monkeypatch.setattr(torch.cuda.nvtx, "range_pop", lambda: events.append(("pop", None)))
+
+    with pytest.raises(RuntimeError, match="profile body failed"):
+        with nvtx_range("minimax_h3.failure"):
+            raise RuntimeError("profile body failed")
+
+    assert events == [
+        ("push", "minimax_h3.failure"),
+        ("pop", None),
+    ]

fastvideo/tests/vaes/test_minimax_h3_vae_compile.py

@@ -0,0 +1,502 @@
+# SPDX-License-Identifier: Apache-2.0
+"""CPU contract tests for MiniMax H3 VAE compilation and profiling ranges.
+
+The final test is a CUDA regression gate for the reduce-overhead tile path
+(real tiled decode/encode with an unmocked ``_stitch_tiles``).
+"""
+
+from contextlib import contextmanager
+from types import SimpleNamespace
+from typing import Any
+from unittest.mock import Mock, patch
+
+import pytest
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from fastvideo.models.vaes.minimax_h3_audio import (
+    MiniMaxH3AudioBigVGANDecoder,
+    MiniMaxH3AudioVAE,
+)
+from fastvideo.models.vaes.minimax_h3_video import (
+    AutoencoderKLMiniMaxH3,
+    MiniMaxH3VideoAttention,
+    MiniMaxH3VideoViTDecoder3d,
+)
+from fastvideo.platforms import AttentionBackendEnum
+from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase
+
+
+def _empty_typed_module(module_type: type[nn.Module]) -> nn.Module:
+    """Create a weightless instance that retains its production module type."""
+    module = object.__new__(module_type)
+    nn.Module.__init__(module)
+    return module
+
+
+def _assert_dynamic_compile_selects_decoder(
+    vae_type: type[nn.Module],
+    decoder_type: type[nn.Module],
+) -> None:
+    """Verify one H3 VAE compiles only its top-level decoder in place."""
+    vae = _empty_typed_module(vae_type)
+    decoder = _empty_typed_module(decoder_type)
+    same_type_under_another_name = _empty_typed_module(decoder_type)
+    unrelated_submodule = nn.Identity()
+    vae.decoder = decoder
+    vae.same_type_under_another_name = same_type_under_another_name
+    vae.unrelated_submodule = unrelated_submodule
+
+    compiled_forward = Mock(name="compiled_forward")
+    compile_kwargs = {"backend": "inductor", "dynamic": False}
+    with patch(
+        "fastvideo.pipelines.composed_pipeline_base.torch.compile",
+        return_value=compiled_forward,
+    ) as compile_mock:
+        compiled_count = ComposedPipelineBase._compile_with_conditions(vae, compile_kwargs)
+
+    assert compiled_count == 1
+    compile_mock.assert_called_once()
+    selected_forward = compile_mock.call_args.args[0]
+    assert selected_forward.__self__ is decoder
+    assert selected_forward.__func__ is decoder_type.forward
+    assert compile_mock.call_args.kwargs == compile_kwargs
+    assert decoder.forward is compiled_forward
+    assert "forward" not in same_type_under_another_name.__dict__
+    assert "forward" not in unrelated_submodule.__dict__
+
+    wrong_type_vae = _empty_typed_module(vae_type)
+    wrong_type_vae.decoder = nn.Identity()
+    with patch("fastvideo.pipelines.composed_pipeline_base.torch.compile") as wrong_type_compile:
+        wrong_type_count = ComposedPipelineBase._compile_with_conditions(wrong_type_vae, compile_kwargs)
+    assert wrong_type_count == 0
+    wrong_type_compile.assert_not_called()
+
+
+def _assert_reduce_overhead_compile(compiled_function: Any) -> None:
+    """Verify a class-owned compile boundary enables CUDA Graph replay."""
+    assert hasattr(compiled_function, "get_compiler_config")
+    assert compiled_function.get_compiler_config()["triton.cudagraphs"] is True
+
+
+def test_video_attention_uses_selected_fastvideo_backend() -> None:
+    """Pass BSHD tensors to the selected dense backend without forward metadata."""
+    backend_call: dict[str, Any] = {}
+
+    class RecordingAttentionImpl:
+        """Record the backend construction and forward contracts."""
+
+        def __init__(self, **kwargs: Any) -> None:
+            backend_call["init"] = kwargs
+
+        def forward(
+            self,
+            query: torch.Tensor,
+            key: torch.Tensor,
+            value: torch.Tensor,
+            attention_metadata: Any,
+        ) -> torch.Tensor:
+            """Return values unchanged after recording the backend inputs."""
+            backend_call["shapes"] = (query.shape, key.shape, value.shape)
+            backend_call["metadata"] = attention_metadata
+            return value
+
+    class RecordingAttentionBackend:
+        """Supply the recording implementation through the backend API."""
+
+        @staticmethod
+        def get_impl_cls() -> type[RecordingAttentionImpl]:
+            return RecordingAttentionImpl
+
+    with (
+        patch("fastvideo.platforms.current_platform") as current_platform,
+        patch(
+            "fastvideo.models.vaes.minimax_h3_video.get_attn_backend",
+            return_value=RecordingAttentionBackend,
+        ) as get_backend,
+    ):
+        current_platform.is_cuda_alike.return_value = True
+        attention = MiniMaxH3VideoAttention(dim=8, heads=2, dim_head=4)
+
+    attention.to_q = nn.Identity()
+    attention.to_k = nn.Identity()
+    attention.to_v = nn.Identity()
+    attention.norm_q = nn.Identity()
+    attention.norm_k = nn.Identity()
+    attention.to_out[0] = nn.Identity()
+    output = attention(torch.empty((1, 3, 8), device="meta"))
+
+    get_backend.assert_called_once_with(
+        4,
+        torch.bfloat16,
+        supported_attention_backends=(
+            AttentionBackendEnum.TORCH_SDPA,
+            AttentionBackendEnum.FLASH_ATTN,
+        ),
+    )
+    assert backend_call["init"] == {
+        "num_heads": 2,
+        "head_size": 4,
+        "softmax_scale": 0.5,
+        "num_kv_heads": 2,
+        "causal": False,
+        # The FP32-pinned VAE opts out of the FASTVIDEO_NVFP4_FA4 env opt-in.
+        "nvfp4_fa4": False,
+    }
+    assert backend_call["shapes"] == ((1, 3, 2, 4), ) * 3
+    assert backend_call["metadata"] is None
+    assert output.shape == (1, 3, 8)
+
+
+def test_video_attention_cpu_uses_torch_sdpa() -> None:
+    """Use PyTorch SDPA when H3 VAE attention receives CPU tensors."""
+    with (
+        patch("fastvideo.platforms.current_platform") as current_platform,
+        patch("fastvideo.models.vaes.minimax_h3_video.get_attn_backend") as get_backend,
+    ):
+        current_platform.is_cuda_alike.return_value = False
+        attention = MiniMaxH3VideoAttention(dim=8, heads=2, dim_head=4)
+
+    get_backend.assert_not_called()
+    assert attention.attn_impl is None
+
+    attention.to_q = nn.Identity()
+    attention.to_k = nn.Identity()
+    attention.to_v = nn.Identity()
+    attention.norm_q = nn.Identity()
+    attention.norm_k = nn.Identity()
+    attention.to_out[0] = nn.Identity()
+    hidden_states = torch.randn(1, 3, 8)
+    query = hidden_states.unflatten(2, (2, 4)).permute(0, 2, 1, 3)
+    expected = F.scaled_dot_product_attention(query, query, query).permute(0, 2, 1, 3).flatten(2, 3)
+
+    torch.testing.assert_close(attention(hidden_states), expected)
+
+
+def test_compile_with_conditions_selects_minimax_h3_video_decoder() -> None:
+    """Compile the registered video decoder with the VAE runtime kwargs."""
+    assert not hasattr(MiniMaxH3VideoViTDecoder3d.forward, "get_compiler_config")
+    _assert_dynamic_compile_selects_decoder(AutoencoderKLMiniMaxH3, MiniMaxH3VideoViTDecoder3d)
+
+
+def test_tile_helpers_stay_eager_by_default() -> None:
+    """Leave the tile helpers uncompiled unless the VAE compile opt-in runs.
+
+    Default users must keep pre-PR eager behavior: no inductor/triton
+    requirement, no first-decode compile latency, and no permanent cudagraph
+    memory pools unless ``enable_torch_compile_vae`` was requested.
+    """
+    assert AutoencoderKLMiniMaxH3._tile_helpers_compiled is False
+    for helper in (AutoencoderKLMiniMaxH3._stitch_tiles, AutoencoderKLMiniMaxH3._project_decoder_tile):
+        assert not hasattr(helper, "get_compiler_config")
+        assert not hasattr(helper, "_torchdynamo_orig_callable")
+
+
+def test_prepare_for_compile_installs_reduce_overhead_tile_helpers() -> None:
+    """Compile both tile helpers with CUDA Graph replay under the opt-in hook.
+
+    ``ComposedPipelineBase._maybe_compile_pipeline_module`` invokes
+    ``prepare_for_compile`` only when ``enable_torch_compile_vae`` is set,
+    right before the decoder is compiled through ``_compile_conditions``.
+    """
+    vae = _empty_typed_module(AutoencoderKLMiniMaxH3)
+    assert vae._tile_helpers_compiled is False
+
+    vae.prepare_for_compile()
+
+    assert vae._tile_helpers_compiled is True
+    _assert_reduce_overhead_compile(vae._stitch_tiles)
+    _assert_reduce_overhead_compile(vae._project_decoder_tile)
+
+    # Idempotent: the pipeline hook may run more than once per instance.
+    compiled_stitch = vae._stitch_tiles
+    compiled_project = vae._project_decoder_tile
+    vae.prepare_for_compile()
+    assert vae._stitch_tiles is compiled_stitch
+    assert vae._project_decoder_tile is compiled_project
+
+    # The opt-in mutates only this instance; the class (and therefore every
+    # default instance) stays eager.
+    assert AutoencoderKLMiniMaxH3._tile_helpers_compiled is False
+    assert not hasattr(AutoencoderKLMiniMaxH3._stitch_tiles, "get_compiler_config")
+    assert not hasattr(AutoencoderKLMiniMaxH3._project_decoder_tile, "get_compiler_config")
+
+
+def test_tile_drivers_return_caller_owned_tensors_when_compiled() -> None:
+    """Clone the stitched canvas out of cudagraph-pooled storage under the opt-in.
+
+    With ``mode="reduce-overhead"`` the ``_stitch_tiles`` output is a
+    CUDA-graph static buffer that the next replay overwrites, while the
+    collect-then-``torch.cat`` consumers (``_decode``/``_encode``/
+    ``_encode_pixels``/``encode_keyframe``) hold each chunk or clip result
+    across replays. The tile drivers therefore hand back a caller-owned copy
+    when the helpers are compiled — and return the stitched tensor by identity
+    when eager, keeping default behavior and peak memory unchanged.
+    """
+
+    def _mock_tiled_vae(stitched: torch.Tensor) -> nn.Module:
+        vae = _empty_typed_module(AutoencoderKLMiniMaxH3)
+        vae.use_tiling = True
+        vae.spatial_compression_ratio = 1
+        vae.tile_sample_min_height = 1
+        vae.tile_sample_min_width = 1
+        vae.tile_sample_min_overlap_height = 0
+        vae.tile_sample_min_overlap_width = 0
+        vae._split_tiles = Mock(return_value=([0, 1], [1, 1], [0]))
+        vae.post_quant_conv = nn.Identity()
+        vae._project_decoder_tile = vae.post_quant_conv
+        vae.quant_conv = nn.Identity()
+        vae.encoder = nn.Identity()
+        vae.decoder = nn.Identity()
+        vae._stitch_tiles = Mock(return_value=stitched)
+        return vae
+
+    latent_clip = torch.zeros((1, 1, 1, 2, 2))
+    for driver in ("_decode_clip", "_encode_clip"):
+        stitched = torch.zeros((1, 1, 1, 2, 2))
+
+        # Default instances return the stitched tensor without copying.
+        eager_vae = _mock_tiled_vae(stitched)
+        assert getattr(eager_vae, driver)(latent_clip) is stitched
+
+        compiled_vae = _mock_tiled_vae(stitched)
+        compiled_vae._tile_helpers_compiled = True
+        owned = getattr(compiled_vae, driver)(latent_clip)
+        assert owned is not stitched
+        assert torch.equal(owned, stitched)
+
+
+def test_compile_with_conditions_selects_minimax_h3_audio_decoder() -> None:
+    """Compile the audio VAE decoder that the H3 waveform decode path calls."""
+    _assert_dynamic_compile_selects_decoder(MiniMaxH3AudioVAE, MiniMaxH3AudioBigVGANDecoder)
+
+
+def test_decode_emits_indexed_temporal_chunk_ranges() -> None:
+    """Nest frame-segment ranges under each temporal decoder chunk range."""
+    vae = _empty_typed_module(AutoencoderKLMiniMaxH3)
+    vae.tokens_chunk_size = 1
+    vae.token_overlap = 1
+    vae.temporal_compression_ratio = 1
+    vae.frame_pre_padding = 0
+    vae.frame_overlap = 1
+    vae.config = SimpleNamespace(token_drop=1)
+    vae._decode_clip = Mock(return_value=torch.zeros((1, 1, 2, 1, 1)))
+    range_events = []
+
+    @contextmanager
+    def record_range(name: str):
+        range_events.append(("enter", name))
+        try:
+            yield
+        finally:
+            range_events.append(("exit", name))
+
+    with patch("fastvideo.models.vaes.minimax_h3_video.nvtx_range", record_range):
+        decoded = vae._decode(torch.zeros((1, 1, 2, 1, 1)))
+
+    assert decoded.shape == (1, 1, 3, 1, 1)
+    assert vae._decode_clip.call_count == 2
+    assert range_events == [
+        ("enter", "minimax_h3.vae.temporal_chunk.0"),
+        ("enter", "minimax_h3.vae.temporal_chunk.0.frame_segment.0"),
+        ("exit", "minimax_h3.vae.temporal_chunk.0.frame_segment.0"),
+        ("enter", "minimax_h3.vae.temporal_chunk.0.frame_segment.1"),
+        ("exit", "minimax_h3.vae.temporal_chunk.0.frame_segment.1"),
+        ("exit", "minimax_h3.vae.temporal_chunk.0"),
+        ("enter", "minimax_h3.vae.temporal_chunk.1"),
+        ("enter", "minimax_h3.vae.temporal_chunk.1.frame_segment.0"),
+        ("exit", "minimax_h3.vae.temporal_chunk.1.frame_segment.0"),
+        ("enter", "minimax_h3.vae.temporal_chunk.1.frame_segment.1"),
+        ("exit", "minimax_h3.vae.temporal_chunk.1.frame_segment.1"),
+        ("exit", "minimax_h3.vae.temporal_chunk.1"),
+    ]
+
+
+def test_decode_clip_no_spatial_tiling_stage_ranges() -> None:
+    """Separate untiled latent projection and decoder ranges."""
+    vae = _empty_typed_module(AutoencoderKLMiniMaxH3)
+    vae.use_tiling = False
+    range_events = []
+    vae.post_quant_conv = nn.Identity()
+    vae.post_quant_conv.register_forward_hook(
+        lambda _module, _args, _output: range_events.append(("call", "post_quant_conv")))
+    vae.decoder = nn.Identity()
+    vae.decoder.register_forward_hook(
+        lambda _module, _args, _output: range_events.append(("call", "decoder_forward")))
+    latent_clip = torch.zeros((1, 1, 1, 2, 2))
+
+    @contextmanager
+    def record_range(name: str):
+        range_events.append(("enter", name))
+        try:
+            yield
+        finally:
+            range_events.append(("exit", name))
+
+    with patch("fastvideo.models.vaes.minimax_h3_video.nvtx_range", record_range):
+        decoded_clip = vae._decode_clip(latent_clip)
+
+    assert decoded_clip is latent_clip
+    assert range_events == [
+        ("enter", "minimax_h3.vae.decode_clip"),
+        ("enter", "minimax_h3.vae.decode_clip.no_s_tile.post_quant_conv"),
+        ("call", "post_quant_conv"),
+        ("exit", "minimax_h3.vae.decode_clip.no_s_tile.post_quant_conv"),
+        ("enter", "minimax_h3.vae.decode_clip.no_s_tile.decoder_forward"),
+        ("call", "decoder_forward"),
+        ("exit", "minimax_h3.vae.decode_clip.no_s_tile.decoder_forward"),
+        ("exit", "minimax_h3.vae.decode_clip"),
+    ]
+
+
+def test_decode_clip_emits_tiled_stage_ranges() -> None:
+    """Nest indexed decoder tiles between tile-splitting and stitching ranges."""
+    vae = _empty_typed_module(AutoencoderKLMiniMaxH3)
+    vae.use_tiling = True
+    vae.spatial_compression_ratio = 1
+    vae.tile_sample_min_height = 1
+    vae.tile_sample_min_width = 1
+    vae.tile_sample_min_overlap_height = 0
+    vae.tile_sample_min_overlap_width = 0
+    vae._split_tiles = Mock(side_effect=[
+        ([0, 1], [1, 1], [0]),
+        ([0, 1], [1, 1], [0]),
+    ])
+    vae.post_quant_conv = nn.Identity()
+    vae._project_decoder_tile = Mock(side_effect=vae.post_quant_conv)
+    vae.decoder = nn.Identity()
+    stitched_clip = torch.zeros((1, 1, 1, 2, 2))
+    vae._stitch_tiles = Mock(return_value=stitched_clip)
+    range_events = []
+
+    @contextmanager
+    def record_range(name: str):
+        range_events.append(("enter", name))
+        try:
+            yield
+        finally:
+            range_events.append(("exit", name))
+
+    with patch("fastvideo.models.vaes.minimax_h3_video.nvtx_range", record_range):
+        decoded_clip = vae._decode_clip(torch.zeros((1, 1, 1, 2, 2)))
+
+    # Default (eager) instances return the stitched canvas by identity; the
+    # caller-owned copy under the compile opt-in is pinned by
+    # test_tile_drivers_return_caller_owned_tensors_when_compiled.
+    assert decoded_clip is stitched_clip
+    assert vae._split_tiles.call_count == 2
+    assert vae._project_decoder_tile.call_count == 4
+    assert vae._stitch_tiles.call_count == 1
+    assert range_events == [
+        ("enter", "minimax_h3.vae.decode_clip"),
+        ("enter", "minimax_h3.vae.decode_clip.split_tiles"),
+        ("exit", "minimax_h3.vae.decode_clip.split_tiles"),
+        ("enter", "minimax_h3.vae.decode_clip.decode_tiles"),
+        ("enter", "minimax_h3.vae.decode_clip.tile.0.0"),
+        ("enter", "minimax_h3.vae.decode_clip.tile.decoder_forward"),
+        ("exit", "minimax_h3.vae.decode_clip.tile.decoder_forward"),
+        ("exit", "minimax_h3.vae.decode_clip.tile.0.0"),
+        ("enter", "minimax_h3.vae.decode_clip.tile.0.1"),
+        ("enter", "minimax_h3.vae.decode_clip.tile.decoder_forward"),
+        ("exit", "minimax_h3.vae.decode_clip.tile.decoder_forward"),
+        ("exit", "minimax_h3.vae.decode_clip.tile.0.1"),
+        ("enter", "minimax_h3.vae.decode_clip.tile.1.0"),
+        ("enter", "minimax_h3.vae.decode_clip.tile.decoder_forward"),
+        ("exit", "minimax_h3.vae.decode_clip.tile.decoder_forward"),
+        ("exit", "minimax_h3.vae.decode_clip.tile.1.0"),
+        ("enter", "minimax_h3.vae.decode_clip.tile.1.1"),
+        ("enter", "minimax_h3.vae.decode_clip.tile.decoder_forward"),
+        ("exit", "minimax_h3.vae.decode_clip.tile.decoder_forward"),
+        ("exit", "minimax_h3.vae.decode_clip.tile.1.1"),
+        ("exit", "minimax_h3.vae.decode_clip.decode_tiles"),
+        ("enter", "minimax_h3.vae.decode_clip.stitch_tiles"),
+        ("exit", "minimax_h3.vae.decode_clip.stitch_tiles"),
+        ("exit", "minimax_h3.vae.decode_clip"),
+    ]
+
+
+def _tiny_real_vae() -> AutoencoderKLMiniMaxH3:
+    """Random-weight VAE small enough for a real tiled decode/encode on GPU."""
+    from fastvideo.configs.models.vaes.minimax_h3_video import (
+        MiniMaxH3VideoVAEArchConfig,
+        MiniMaxH3VideoVAEConfig,
+    )
+
+    arch = MiniMaxH3VideoVAEArchConfig(
+        latent_channels=4,
+        block_out_channels=(32, 32),
+        layers_per_block=1,
+        spatial_downsample_factors=(2, 2),
+        temporal_downsample_factors=(2, 2),
+        decoder_num_layers=1,
+        decoder_num_attention_heads=1,
+        decoder_attention_head_dim=8,
+        decoder_num_register_tokens=2,
+        decoder_ffn_mult=1,
+        latents_mean=(0.0, ) * 4,
+        latents_std=(1.0, ) * 4,
+    )
+    return AutoencoderKLMiniMaxH3(
+        MiniMaxH3VideoVAEConfig(
+            arch_config=arch,
+            use_tiling=False,
+            use_temporal_tiling=False,
+            use_parallel_tiling=False,
+        )).eval()
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="reduce-overhead tile compile requires CUDA graphs")
+@torch.inference_mode()
+def test_tiled_decode_and_encode_survive_cudagraph_buffer_reuse_on_cuda() -> None:
+    """Real tiled decode()/encode() with an unmocked reduce-overhead ``_stitch_tiles``.
+
+    Regression gate for the CUDA-graph output-clobbering bug under the
+    ``enable_torch_compile_vae`` opt-in (``prepare_for_compile``): the
+    stitched canvas is a cudagraph static buffer, and the
+    collect-then-``torch.cat`` consumers (``_decode``/``_encode``/
+    ``_encode_pixels``) hold chunk/clip results across subsequent
+    ``_stitch_tiles`` replays. Without the eager ``.clone()`` at the
+    tile-driver returns, the first tiled ``decode()`` with >=2 temporal
+    chunks raises ``accessing tensor output of CUDAGraphs that has been
+    overwritten by a subsequent run``. This test needs >=2 chunks (decode),
+    >=2 clips (encode), and a >=2x2 spatial tile grid.
+    """
+    torch.manual_seed(20260821)
+    vae = _tiny_real_vae().to("cuda")
+    vae.enable_tiling(16, 16, 4, 4)
+    # The hook ComposedPipelineBase._maybe_compile_pipeline_module runs for
+    # the opt-in; it installs the reduce-overhead compiled tile helpers.
+    vae.prepare_for_compile()
+    assert vae._tile_helpers_compiled
+
+    # 8 latent tokens = 2 temporal chunks (tokens_chunk_size 5); 8x8 latents =
+    # 32x32 pixels = a 2x2 grid of 16px tiles.
+    z = torch.randn(1, 4, 8, 8, 8, device="cuda")
+    pad_tokens, num_chunks, _ = vae._temporal_decode_plan(z.shape[2])
+    assert num_chunks >= 2, "decode workload must span multiple stitch replays"
+
+    decoded_first = vae.decode(z).sample
+    decoded_second = vae.decode(z).sample
+    assert torch.equal(decoded_first, decoded_second)
+
+    # 34 frames = 2 encode clips of clip_length 17 -> 2 stitch replays.
+    pixels = torch.rand(1, 3, 34, 32, 32, device="cuda")
+    encoded_first = vae.encode(pixels).latent_dist.parameters
+    encoded_second = vae.encode(pixels).latent_dist.parameters
+    assert torch.equal(encoded_first, encoded_second)
+
+    uint8_pixels = torch.randint(0, 256, (1, 3, 34, 32, 32), dtype=torch.uint8)
+    streamed_first = vae.encode_pixels(uint8_pixels).latent_dist.parameters
+    streamed_second = vae.encode_pixels(uint8_pixels).latent_dist.parameters
+    assert torch.equal(streamed_first, streamed_second)
+
+    # Output parity vs a default (fully eager, never-prepared) instance with
+    # the same weights and tiling — the tolerance absorbs inductor fusion
+    # reassociation only.
+    eager_vae = _tiny_real_vae().to("cuda")
+    eager_vae.load_state_dict(vae.state_dict())
+    eager_vae.enable_tiling(16, 16, 4, 4)
+    assert not eager_vae._tile_helpers_compiled
+    torch.testing.assert_close(decoded_first, eager_vae.decode(z).sample, atol=2e-4, rtol=2e-4)
+    torch.testing.assert_close(encoded_first, eager_vae.encode(pixels).latent_dist.parameters, atol=2e-4, rtol=2e-4)

fastvideo/tests/worker/test_gpu_worker.py

@@ -1,11 +1,40 @@
 from types import SimpleNamespace
+from unittest.mock import Mock
 
 import pytest
 import torch
 
 from fastvideo.fastvideo_args import FastVideoArgs
 from fastvideo.pipelines import ForwardBatch
-from fastvideo.worker.gpu_worker import Worker
+from fastvideo.worker.gpu_worker import Worker, _log_cuda_device_uuid
+
+
+def test_cuda_device_uuid_receipt_is_disabled_without_nvtx_profiling(monkeypatch) -> None:
+    """Avoid NVIDIA property access during ordinary worker initialization."""
+    get_device_properties = Mock()
+    monkeypatch.setenv("FASTVIDEO_NVTX_PROFILE", "0")
+    monkeypatch.setattr(torch.cuda, "get_device_properties", get_device_properties)
+
+    _log_cuda_device_uuid(0, torch.device("cuda:0"))
+
+    get_device_properties.assert_not_called()
+
+
+def test_cuda_device_uuid_receipt_identifies_profiled_worker(monkeypatch) -> None:
+    """Bind one profiled worker rank to its NVIDIA device UUID in logs."""
+    log_info = Mock()
+    monkeypatch.setenv("FASTVIDEO_NVTX_PROFILE", "1")
+    monkeypatch.setattr(torch.cuda, "get_device_properties", lambda device: SimpleNamespace(uuid="device-uuid"))
+    monkeypatch.setattr("fastvideo.worker.gpu_worker.logger.info", log_info)
+
+    _log_cuda_device_uuid(2, torch.device("cuda:0"))
+
+    log_info.assert_called_once_with(
+        "Worker %d CUDA device UUID: GPU-%s",
+        2,
+        "device-uuid",
+        local_main_process_only=False,
+    )
 
 
 def _worker_returning(output_batch: ForwardBatch) -> Worker:

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