stages ¶

@torch.no_grad()
def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
    """Decode H3 audio latents into a stereo CPU waveform."""
    # Audio decode is sub-second, so preserve the serial path's global
    # rank-zero ownership.
    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
        self._clear_runtime(batch)
        return batch

    layout = _layout(batch)
    if batch.audio_latents is None:
        raise ValueError("MiniMax-H3 audio latents are missing at decode.")
    latents = unpack_audio_tokens(
        batch.audio_latents[layout.num_condition_audio_rows:],
        layout.num_audio_latents,
    )
    device = get_local_torch_device()
    self.audio_vae.to(device)
    try:
        latents = self.audio_vae.denormalize_latents(latents.to(device=device, dtype=torch.float32))
        if fastvideo_args.output_type == "latent":
            batch.extra["audio"] = latents.detach().float().cpu()
            batch.extra["audio_sample_rate"] = self.audio_vae.sampling_rate
            self._clear_runtime(batch)
            return batch

        # 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)}.")
        batch.extra["audio"] = decoded[:, 0].transpose(0, 1).contiguous().cpu()
        batch.extra["audio_sample_rate"] = self.audio_vae.sampling_rate
        self._clear_runtime(batch)
        return batch
    finally:
        if fastvideo_args.vae_cpu_offload:
            self.audio_vae.to("cpu")

@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
                         and not isinstance(first_param, DTensor))
    if moved_for_forward:
        self.conditioner.to(device)
    try:
        # 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")
    batch.prompt_embeds = [prompt_embeds]
    batch.extra[MINIMAX_H3_TEXT_TOKEN_TAGS_KEY] = text_token_tags
    return batch

@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.")
    if not batch.prompt_embeds or batch.latents is None or batch.audio_latents is None:
        raise ValueError("MiniMax-H3 conditioning and packed latents must precede denoising.")

    full_cpu_offload = (fastvideo_args.dit_cpu_offload and not fastvideo_args.dit_layerwise_offload
                        and not fastvideo_args.use_fsdp_inference)
    device = get_local_torch_device()
    if full_cpu_offload:
        self.transformer.to(device)
        batch.latents = batch.latents.to(device)
        batch.audio_latents = batch.audio_latents.to(device)

    self.scheduler.set_timesteps(batch.num_inference_steps, device=device)
    self.audio_scheduler.set_timesteps(batch.num_inference_steps, device=device)
    video_timesteps = self.scheduler.timesteps
    audio_timesteps = self.audio_scheduler.timesteps
    if video_timesteps is None or audio_timesteps is None:
        raise ValueError("MiniMax-H3 schedulers did not produce timesteps.")
    if len(video_timesteps) != len(audio_timesteps):
        raise ValueError("MiniMax-H3 video and audio schedules must have the same number of intervals.")

    row_timestep_plan = []
    for video_timestep, audio_timestep in zip(video_timesteps, audio_timesteps, strict=True):
        video_value = float(video_timestep.item())
        audio_value = float(audio_timestep.item())
        unique, inverse = build_row_timesteps(
            layout,
            video_timestep=video_value,
            audio_timestep=audio_value,
            condition_video_timestep=max(video_value, MINIMAX_H3_KEYFRAME_NOISE_AUG),
            condition_audio_timestep=1.0,
        )
        row_timestep_plan.append((unique.to(device), inverse.to(device)))
    batch.timesteps = video_timesteps

    position_ids = layout.position_ids.to(device)
    token_tags = layout.token_tags.to(device)
    video_indices = layout.video_indices.to(device)
    audio_indices = layout.audio_indices.to(device)
    text_indices = layout.text_indices.to(device)
    prompt_embeds = batch.prompt_embeds[0].to(device)

    vsa_metadata_builder = _h3_vsa_metadata_builder(self.transformer, fastvideo_args)
    if vsa_metadata_builder is not None:
        vsa_patch_size = fastvideo_args.pipeline_config.dit_config.patch_size
        vsa_prefix_segments = _h3_vsa_prefix_segments(layout, vsa_patch_size)
        # Per-request knobs (sweeps flip these between generate_video calls
        # without respawning workers); mode None defers to the env default.
        vsa_mode = batch.extra.get("vsa_mode", "exempt")
        if vsa_mode not in ("exempt", "compete"):
            raise ValueError(f"vsa_mode must be 'exempt' or 'compete', got {vsa_mode!r}.")
        vsa_exempt = vsa_mode == "exempt"
        vsa_dense_layers = tuple(batch.extra.get("vsa_dense_layers", ()))
        vsa_dense_first_n = int(batch.extra.get("vsa_dense_first_n_steps", 0))
        # Run-level tile geometry (256 default, 64 = native Triton path),
        # plumbed like the run-level sparsity; the builder validates the
        # value against VSA_H3_TILE_SHAPES.
        vsa_tile_size = int(fastvideo_args.VSA_tile_size)

    try:
        # 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]
                attn_metadata = None
                if vsa_metadata_builder is not None:
                    # Optional schedule: run the first N steps dense (sparsity 0
                    # selects every tile — parity-proven ≡ dense ≤2e-4); early
                    # steps set global structure and are the most damage-prone.
                    vsa_sparsity = 0.0 if index < vsa_dense_first_n else float(batch.VSA_sparsity)
                    attn_metadata = vsa_metadata_builder.build(
                        current_timestep=index,
                        raw_latent_shape=(layout.num_video_latent_frames, layout.latent_height,
                                          layout.latent_width),
                        patch_size=vsa_patch_size,
                        VSA_sparsity=vsa_sparsity,
                        prefix_segments=vsa_prefix_segments,
                        device=device,
                        exempt=vsa_exempt,
                        dense_layers=vsa_dense_layers,
                        tile_size=vsa_tile_size,
                    )
                # Under torch.compile(mode="reduce-overhead") each denoising
                # step must be marked, or cudagraph trees flag cross-step
                # reuse of pooled outputs as "accessing tensor output of
                # CUDAGraphs that has been overwritten" (surfaces at sp=1;
                # sp>1 is masked by collective-induced graph breaks).
                torch.compiler.cudagraph_mark_step_begin()
                with trace_step(index), set_forward_context(
                        current_timestep=index,
                        attn_metadata=attn_metadata,
                        forward_batch=batch,
                ):
                    video_velocity, audio_velocity = self.transformer(
                        hidden_states=batch.latents[None],
                        audio_hidden_states=batch.audio_latents[None],
                        encoder_hidden_states=prompt_embeds,
                        timestep=unique_timesteps,
                        timestep_indices=timestep_indices,
                        token_tags=token_tags,
                        position_ids=position_ids,
                        video_indices=video_indices,
                        audio_indices=audio_indices,
                        text_indices=text_indices,
                    )

                video_start = layout.num_condition_video_rows
                audio_start = layout.num_condition_audio_rows
                batch.latents[video_start:] = self.scheduler.step(
                    video_velocity[0, video_start:].float(),
                    video_timestep,
                    batch.latents[video_start:],
                    return_dict=False,
                )[0]
                batch.audio_latents[audio_start:] = self.audio_scheduler.step(
                    audio_velocity[0, audio_start:].float(),
                    audio_timestep,
                    batch.audio_latents[audio_start:],
                    return_dict=False,
                )[0]
                batch.step_index = index
                batch.timestep = video_timestep
    finally:
        if bool(getattr(fastvideo_args, "dit_layerwise_offload", False)):
            manager = getattr(self.transformer, "_layerwise_offload_manager", None)
            if manager is not None and getattr(manager, "enabled", False):
                manager.release_all()
        if full_cpu_offload:
            self.transformer.to("cpu")
    return batch

@torch.no_grad()
def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
    """Decode H3 video latents into normalized CPU pixels."""
    placeholder = torch.empty((0, 3, 0, 0, 0), device="cpu", dtype=torch.float32)
    sp_group, is_output_rank, parallel = _decode_participation(fastvideo_args, fastvideo_args.vae_parallel_decode)
    if not is_output_rank and not parallel:
        # Consumers read the output rank's ForwardBatch. Keep a
        # verifier-compatible placeholder on other ranks and avoid
        # duplicating the full VAE decode and CPU output buffer.
        batch.output = placeholder
        return batch

    layout = _layout(batch)
    if batch.latents is None or batch.raw_latent_shape is None or len(batch.raw_latent_shape) != 5:
        raise ValueError("MiniMax-H3 video latents or raw geometry are missing at decode.")
    _, channels, num_frames, latent_height, latent_width = batch.raw_latent_shape
    latents = unpatchify_video_tokens(
        batch.latents[layout.num_condition_video_rows:],
        num_frames,
        latent_height,
        latent_width,
        channels,
        self.transformer.patch_size,
    )
    device = get_local_torch_device()
    self.vae.to(device)
    try:
        latents = self.vae.denormalize_latents(latents.to(device=device, dtype=torch.float32))
        if fastvideo_args.output_type == "latent":
            # No collectives on this path, so uniform participation is
            # trivial: every rank returns here.
            batch.output = latents.detach().float().cpu() if is_output_rank else placeholder
            return batch

        output = None
        if is_output_rank:
            output = torch.empty(
                self.vae.decoded_pixel_shape(latents.shape),
                device="cpu",
                dtype=torch.float32,
                pin_memory=fastvideo_args.pin_cpu_memory and is_pin_memory_available(),
            )
        # 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"),
        ):
            if parallel:
                strategy = fastvideo_args.vae_parallel_decode_strategy or DEFAULT_DECODE_GATHER_STRATEGY
                logger.info("MiniMax-H3 VAE decode: sequence-parallel chunks across %d ranks (%s)",
                            sp_group.world_size, strategy)
                decode_to_pixels_parallel(self.vae, latents, output, sp_group, strategy=strategy)
            else:
                self.vae.decode_to_pixels(latents, output)
        batch.output = output if is_output_rank else placeholder
        return batch
    finally:
        if fastvideo_args.vae_cpu_offload:
            self.vae.to("cpu")