From 68e6ffca9ed4d7171ac2ecb5c5d3f5776dd32f36 Mon Sep 17 00:00:00 2001 From: Will Lin Date: Fri, 21 Aug 2026 10:33:06 +0000 Subject: [PATCH] [bugfix] H3 VAE: clone the reduce-overhead stitched canvas at the tile-driver returns (#1734 review F1) Carried blocker fix for PR #1734 on this integration stack, per pr1734.md finding F1 (reproduced on GB200/torch 2.12): with @torch.compile(mode="reduce-overhead") on _stitch_tiles, the stitched canvas is a CUDA-graph static buffer, and the collect-then-torch.cat consumers - _decode (via _decode_chunks), _encode, _encode_pixels, encode_keyframe - hold each chunk/clip result across the next _stitch_tiles replay, which overwrites the pooled storage. First tiled decode() with >=2 temporal chunks (any real video) and tiled encode() of >17 frames raised: RuntimeError: Error: accessing tensor output of CUDAGraphs that has been overwritten by a subsequent run. ... line ..., in _stitch_tiles: return torch.cat(result_rows, dim=-2) Fix: .clone() the stitch output at the two eager call sites (_encode_clip tiled return, _decode_clip tiled return) so no cudagraph-owned storage escapes the tile driver; the clone is read before the next replay, so it is race-free by construction. The streaming _decode_to_pixels path was already safe (full copy-out per chunk before the next decode) and stays correct at one extra D2D copy per chunk (~tens of microseconds vs the decode compute). This follows the option (a) recommendation in the review; reduce-overhead is retained on _stitch_tiles/_project_decoder_tile. Tests: - test_decode_clip_emits_tiled_stage_ranges updated: the tile driver now returns a caller-owned copy (is-not + equal), pinning the ownership contract at the mock level. - New CUDA regression gate (the review's must-add test): test_tiled_decode_and_encode_survive_cudagraph_buffer_reuse_on_cuda - real tiled decode() (2 temporal chunks, 2x2 spatial tiles) and encode()/ encode_pixels() (2 clips) with _stitch_tiles unmocked, asserting bitwise repeat-consistency plus value parity against the fully eager tile helpers (via _torchdynamo_orig_callable). Red on the unfixed merge (cudagraph overwrite RuntimeError on the first decode), green with this fix. Co-Authored-By: Claude Fable 5 --- fastvideo/models/vaes/minimax_h3_video.py | 15 ++- .../tests/vaes/test_minimax_h3_vae_compile.py | 105 +++++++++++++++++- 2 files changed, 115 insertions(+), 5 deletions(-) diff --git a/fastvideo/models/vaes/minimax_h3_video.py b/fastvideo/models/vaes/minimax_h3_video.py index f0e1eef83..d29f6ca5a 100644 --- a/fastvideo/models/vaes/minimax_h3_video.py +++ b/fastvideo/models/vaes/minimax_h3_video.py @@ -750,7 +750,12 @@ 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) + # Under mode="reduce-overhead" 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. + return self._stitch_tiles(rows, latent_y_overlaps, latent_x_overlaps).clone() def _decode_clip(self, z: torch.Tensor) -> torch.Tensor: """Decode one temporal clip, with optional overlapping spatial tiles.""" @@ -796,7 +801,13 @@ class AutoencoderKLMiniMaxH3(nn.Module): rows.append(row) with nvtx_range("minimax_h3.vae.decode_clip.stitch_tiles"): - return self._stitch_tiles(rows, y_overlaps, x_overlaps) + # 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 and never held stale storage; the + # clone keeps that path correct too at one D2D copy per chunk.) + return self._stitch_tiles(rows, y_overlaps, x_overlaps).clone() def _encode(self, x: torch.Tensor) -> torch.Tensor: clip_length = self.config.clip_length diff --git a/fastvideo/tests/vaes/test_minimax_h3_vae_compile.py b/fastvideo/tests/vaes/test_minimax_h3_vae_compile.py index 6de6fca13..613c466eb 100644 --- a/fastvideo/tests/vaes/test_minimax_h3_vae_compile.py +++ b/fastvideo/tests/vaes/test_minimax_h3_vae_compile.py @@ -1,11 +1,16 @@ # SPDX-License-Identifier: Apache-2.0 -"""CPU contract tests for MiniMax H3 VAE compilation and profiling ranges.""" +"""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 types import MethodType, 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 @@ -297,7 +302,11 @@ def test_decode_clip_emits_tiled_stage_ranges() -> None: 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 + # The tile driver must hand back a caller-owned copy: under + # mode="reduce-overhead" the stitched canvas is CUDA-graph pooled storage + # that the next replay overwrites, so returning it by identity is a bug. + assert decoded_clip is not stitched_clip + assert torch.equal(decoded_clip, stitched_clip) assert vae._split_tiles.call_count == 2 assert vae._project_decoder_tile.call_count == 4 assert vae._stitch_tiles.call_count == 1 @@ -327,3 +336,93 @@ def test_decode_clip_emits_tiled_stage_ranges() -> None: ("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() + + +def _dynamo_original(compiled_function: Any) -> Any: + """Return the eager callable behind a ``torch.compile``-decorated function.""" + original = getattr(compiled_function, "_torchdynamo_orig_callable", None) + if original is None: + original = getattr(compiled_function, "__wrapped__", None) + assert original is not None, "cannot recover the eager tile helpers" + return original + + +@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: 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) + + # 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 the fully eager tile helpers (same weights, same math; + # 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) + eager_vae._stitch_tiles = MethodType(_dynamo_original(AutoencoderKLMiniMaxH3._stitch_tiles), eager_vae) + eager_vae._project_decoder_tile = MethodType(_dynamo_original(AutoencoderKLMiniMaxH3._project_decoder_tile), + eager_vae) + 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)