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/envs.py b/fastvideo/envs.py index 62a002089..e68b319ac 100644 --- a/fastvideo/envs.py +++ b/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": diff --git a/fastvideo/models/dits/minimax_h3.py b/fastvideo/models/dits/minimax_h3.py index a36d2bc5f..84949fa59 100644 --- a/fastvideo/models/dits/minimax_h3.py +++ b/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, diff --git a/fastvideo/models/loader/component_loader.py b/fastvideo/models/loader/component_loader.py index a595bd338..3820376ef 100644 --- a/fastvideo/models/loader/component_loader.py +++ b/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={ diff --git a/fastvideo/models/vaes/minimax_h3_audio.py b/fastvideo/models/vaes/minimax_h3_audio.py index 6aefc411c..92d378622 100644 --- a/fastvideo/models/vaes/minimax_h3_audio.py +++ b/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 diff --git a/fastvideo/models/vaes/minimax_h3_video.py b/fastvideo/models/vaes/minimax_h3_video.py index e2d671578..f0e1eef83 100644 --- a/fastvideo/models/vaes/minimax_h3_video.py +++ b/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,34 @@ 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, + ) 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 +355,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 +468,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 +518,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 +530,7 @@ 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] def __init__(self, config: MiniMaxH3VideoVAEConfig) -> None: super().__init__() @@ -655,12 +696,15 @@ class AutoencoderKLMiniMaxH3(nn.Module): slice_rest[dim] = slice(blend_extent, None) return torch.cat([blended, b[tuple(slice_rest)]], dim=dim) + # The fixed spatial tile grid reuses one compiled blend-and-concatenate graph. + @torch.compile(backend="inductor", mode="reduce-overhead", dynamic=False) 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 +721,12 @@ class AutoencoderKLMiniMaxH3(nn.Module): result_rows.append(torch.cat(result_row, dim=-1)) return torch.cat(result_rows, dim=-2) + # Each fixed-shape latent tile reuses one compiled decoder-input projection. + @torch.compile(backend="inductor", mode="reduce-overhead", dynamic=False) + 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)) @@ -703,33 +753,50 @@ class AutoencoderKLMiniMaxH3(nn.Module): return self._stitch_tiles(rows, latent_y_overlaps, latent_x_overlaps) 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"): + return self._stitch_tiles(rows, y_overlaps, x_overlaps) def _encode(self, x: torch.Tensor) -> torch.Tensor: clip_length = self.config.clip_length @@ -809,21 +876,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] diff --git a/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_conditioning.py b/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_conditioning.py index f77b7d230..d39128e91 100644 --- a/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_conditioning.py +++ b/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") diff --git a/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_decoding.py b/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_decoding.py index 06bd1b88e..05bf3da0b 100644 --- a/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_decoding.py +++ b/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)}.") diff --git a/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_denoising.py b/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_denoising.py index 5881bf98f..f2a95840b 100644 --- a/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_denoising.py +++ b/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.") @@ -147,7 +148,9 @@ class MiniMaxH3DenoisingStage(PipelineStage): vsa_dense_first_n = int(batch.extra.get("vsa_dense_first_n_steps", 0)) 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] diff --git a/fastvideo/profiler.py b/fastvideo/profiler.py index 3303d4600..d75879d7b 100644 --- a/fastvideo/profiler.py +++ b/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.""" diff --git a/fastvideo/tests/contract/test_profiler_regions.py b/fastvideo/tests/contract/test_profiler_regions.py index de18fa7a5..e9b706eba 100644 --- a/fastvideo/tests/contract/test_profiler_regions.py +++ b/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), + ] diff --git a/fastvideo/tests/vaes/test_minimax_h3_vae_compile.py b/fastvideo/tests/vaes/test_minimax_h3_vae_compile.py new file mode 100644 index 000000000..6de6fca13 --- /dev/null +++ b/fastvideo/tests/vaes/test_minimax_h3_vae_compile.py @@ -0,0 +1,329 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CPU contract tests for MiniMax H3 VAE compilation and profiling ranges.""" + +from contextlib import contextmanager +from types import SimpleNamespace +from typing import Any +from unittest.mock import Mock, patch + +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, + } + 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_project_decoder_tile_uses_reduce_overhead_compile() -> None: + """Compile the per-tile decoder-input projection with CUDA Graph replay.""" + _assert_reduce_overhead_compile(AutoencoderKLMiniMaxH3._project_decoder_tile) + + +def test_stitch_tiles_uses_reduce_overhead_compile() -> None: + """Compile spatial tile blending and concatenation with CUDA Graph replay.""" + _assert_reduce_overhead_compile(AutoencoderKLMiniMaxH3._stitch_tiles) + + +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))) + + 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"), + ] diff --git a/fastvideo/tests/worker/test_gpu_worker.py b/fastvideo/tests/worker/test_gpu_worker.py index 2c278d767..77850595d 100644 --- a/fastvideo/tests/worker/test_gpu_worker.py +++ b/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: 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