diff --git a/fastvideo/attention/backends/flash_attn.py b/fastvideo/attention/backends/flash_attn.py index 993b5b4a0..c130b0dfd 100644 --- a/fastvideo/attention/backends/flash_attn.py +++ b/fastvideo/attention/backends/flash_attn.py @@ -291,9 +291,7 @@ class FlashAttentionImpl(AttentionImpl): # identical to before the knob; kwargs win over the # FASTVIDEO_FA4_PV_MODE env bridge (same pattern as nvfp4_fa4 above). self.fa4_pv_mode = _resolve_fa4_pv_mode(extra_impl_args) - # MiniMax-H3 dense attention may opt into cuDNN's fused SDPA. It is an - # explicit impl argument (threaded from FASTVIDEO_MINIMAX_H3_DENSE_ATTN - # by the H3 model) so no other model can inherit it by accident. + # Explicit impl argument (set by the H3 model) so other models cannot inherit the cuDNN route. self.cudnn_sdpa = bool(extra_impl_args.get("cudnn_sdpa", False)) if self.cudnn_sdpa and self.nvfp4_fa4: raise ValueError("cudnn_sdpa and nvfp4_fa4 select different Q/K precision; enable only one of " @@ -493,9 +491,7 @@ class FlashAttentionImpl(AttentionImpl): output = self._forward_nvfp4(query, key, value) elif self.cudnn_sdpa and cudnn_sdpa_supported(query, key, value): - # H3's dense, unmasked, inference-only route. The helper presents - # the BSHD inputs to cuDNN as BHSD strided views (contiguous BSHD - # is silently misread by the ATen op). + # The ATen cuDNN op silently misreads contiguous BSHD; the helper passes BHSD strided views. output = cudnn_sdpa_bshd(query, key, value, scale=self.softmax_scale, causal=self.causal) elif (self.fa4_packed_varlen and fa_version == "4" and not torch.is_grad_enabled() and query.shape[0] == 1 diff --git a/fastvideo/attention/utils/cudnn_sdpa.py b/fastvideo/attention/utils/cudnn_sdpa.py index 9af72430c..a9908e0cd 100644 --- a/fastvideo/attention/utils/cudnn_sdpa.py +++ b/fastvideo/attention/utils/cudnn_sdpa.py @@ -20,12 +20,10 @@ import torch # cuDNN's fused attention kernels take fp16/bf16 only; the H3 DiT runs bf16. _CUDNN_SDPA_DTYPES = (torch.float16, torch.bfloat16) -# cuDNN 9.x supports head sizes up to 128 on the fused SDPA path. Wider heads -# stay on FlashAttention instead of risking an engine-selection failure. +# cuDNN 9.x fused SDPA supports head sizes up to 128. _CUDNN_SDPA_MAX_HEAD_DIM = 128 -# Older torch builds may not register the op at all; keep the module importable -# (the guard then reports unsupported and the backend stays on FlashAttention). +# Older torch builds may not register the op. _ATEN_CUDNN_SDPA = getattr(torch.ops.aten, "_scaled_dot_product_cudnn_attention", None) if _ATEN_CUDNN_SDPA is not None: _ATEN_CUDNN_SDPA = _ATEN_CUDNN_SDPA.default @@ -70,7 +68,6 @@ def cudnn_sdpa_bshd( """Dense scaled-dot-product attention on contiguous ``[B, S, H, D]`` inputs.""" if _ATEN_CUDNN_SDPA is None: raise RuntimeError("torch build does not provide aten::_scaled_dot_product_cudnn_attention") - # compute_log_sumexp=False: inference needs only the forward result. output = _ATEN_CUDNN_SDPA( query.transpose(1, 2), key.transpose(1, 2), @@ -82,5 +79,4 @@ def cudnn_sdpa_bshd( False, scale=scale, )[0] - # cuDNN returns [B, H, S, D]; restore the caller's [B, S, H, D] contract. return output.transpose(1, 2).contiguous() diff --git a/fastvideo/entrypoints/video_generator.py b/fastvideo/entrypoints/video_generator.py index 27609443d..36eddc6ae 100644 --- a/fastvideo/entrypoints/video_generator.py +++ b/fastvideo/entrypoints/video_generator.py @@ -943,10 +943,6 @@ class VideoGenerator: needs_samples_out = batch.return_frames or batch.return_samples if (needs_samples_out and not batch.return_samples and not is_latent_output and envs.FASTVIDEO_H3_SKIP_FP32_SAMPLES_MIRROR.get()): - # ``frames`` below is built from ``output_batch.output`` directly, so the - # fp32 [0, 1] mirror only serves the ``samples`` result, which this - # request did not ask for. MiniMax-H3 hands back uint8 pixels, so that - # mirror is 4x the decoded data (1.5 GB for a 5 s 768p clip). needs_samples_out = False # ``output_shape_from_input`` pipelines (Kandinsky6 SR) decide the output geometry from the input video, so the # request's height / width / num_frames say nothing about the buffer to pre-allocate. diff --git a/fastvideo/models/vaes/minimax_h3_video.py b/fastvideo/models/vaes/minimax_h3_video.py index e6ea9669f..46f781fa8 100644 --- a/fastvideo/models/vaes/minimax_h3_video.py +++ b/fastvideo/models/vaes/minimax_h3_video.py @@ -444,12 +444,7 @@ def _run_decoder_block( hidden_states: torch.Tensor, rotary_emb: tuple[torch.Tensor, torch.Tensor], ) -> torch.Tensor: - """One decoder transformer block; ``FASTVIDEO_H3_VAE_COMPILE_BLOCKS`` compiles this region. - - The function takes the block as an argument instead of compiling ``block.forward`` - per block, so Dynamo traces one block-sized graph and every transformer block and - temporal chunk replays it (the parameters are graph inputs, not guards). - """ + # Takes the block as an argument so one compiled graph serves every block (parameters are graph inputs). return block(hidden_states, rotary_emb) @@ -499,24 +494,14 @@ class MiniMaxH3VideoViTDecoder3d(nn.Module): self.install_compiled_blocks() def install_compiled_blocks(self) -> int: - """Compile the decoder block region once, for every block and temporal chunk. - - Eager runs the blocks' pointwise fat as separate passes: an fp32 RMSNorm - kernel plus an activation-dtype cast, ``attn * scale1`` and the fp32 residual - add, ``silu`` plus ``mul`` for the SwiGLU, and ``ff * scale2`` plus the second - residual add. Inductor fuses each pair into one pass. Returns the number of - transformer blocks now served by the compiled region (0 when already installed). - """ + """Compile the shared decoder-block region; returns the number of blocks it now serves (0 if installed).""" if self._compiled_block_forward is not None: return 0 self._compiled_block_forward = torch.compile(_run_decoder_block, dynamic=False) return len(self.transformer_blocks) def set_compiled_blocks(self, enabled: bool) -> int: - """Switch the blocks between the compiled region and the eager loop. - - Returns the number of blocks whose forward changed (in either direction). - """ + """Switch between the compiled region and the eager loop; returns the number of blocks switched.""" if enabled: return self.install_compiled_blocks() switched = self._compiled_block_forward is not None diff --git a/fastvideo/pipelines/basic/minimax_h3/minimax_h3_pipeline.py b/fastvideo/pipelines/basic/minimax_h3/minimax_h3_pipeline.py index 14b1723ab..abceff7f9 100644 --- a/fastvideo/pipelines/basic/minimax_h3/minimax_h3_pipeline.py +++ b/fastvideo/pipelines/basic/minimax_h3/minimax_h3_pipeline.py @@ -806,8 +806,7 @@ class MiniMaxH3BasePipeline(LoRAPipeline, ComposedPipelineBase): raise box["error"] batch.extra["audio"] = box.get("audio") batch.extra["audio_sample_rate"] = box.get("audio_sample_rate") - # The audio stage bypasses PipelineStage.__call__, so release the lazy - # modules it owns here, exactly as __call__ would after it returns. + # forward() bypassed PipelineStage.__call__, which normally releases these. audio_stage._release_deferred_modules(f"{getattr(audio_stage, '_pipeline_stage_name', 'audio')}|" f"{type(audio_stage).__name__}") return batch