Merge remote-tracking branch 'origin/aryan/open-h3-regenerate-2k' into aryan/regen2k-nvfp4-decoder

docs-site/lib/cookbook-pages.ts

@@ -48,6 +48,15 @@ export const recipePages = [
     tags: ["Reference to video", "Audio", "NVFP4"],
     featured: true,
   },
+  {
+    id: "open-h3-regenerate-2k",
+    family: "MiniMax H3",
+    title: "Open-H3-Regenerate-2K",
+    mark: "H3",
+    path: "/cookbook/minimax-h3/open-h3-regenerate-2k",
+    description: "Regenerate a 768P MiniMax-H3 draft at 2560x1440 with a three-step refiner.",
+    tags: ["Video to video", "2K", "3-step"],
+  },
   {
     id: "fastwan21",
     family: "Wan",

docs/assets/cookbook-recipes.json

@@ -760,6 +760,22 @@
       ],
       "limitations": ["Pass --reference-audio for an additional audio reference. The checked-in source defaults to four GPUs."]
     },
+    {
+      "id": "open-h3-regenerate-2k",
+      "family": "minimax_h3",
+      "stage": "inference",
+      "task": "Video to 2K video",
+      "label": "Open-H3-Regenerate-2K",
+      "summary": "Regenerate a 768P MiniMax-H3 draft at 2560x1440 with NVlabs' three-step refiner. The output is silent.",
+      "model": "Efficient-Large-Model/Open-H3-Regenerate-2K",
+      "source": "examples/inference/basic/basic_open_h3_regenerate_2k.py",
+      "command": ": \"${INPUT_VIDEO:?Set INPUT_VIDEO to a 768P MiniMax-H3 draft}\"\npython examples/inference/basic/basic_open_h3_regenerate_2k.py --video \"$INPUT_VIDEO\" --prompt \"A red fox trots through fresh snow in a pine forest at golden hour.\"",
+      "gpu_types": ["NVIDIA"],
+      "hardware": {"platform": "cuda", "gpu_count": 1, "evidence": "source-configured"},
+      "evidence": "Source-backed",
+      "expected_artifact": "2560x1440 MP4 at outputs/open_h3_regenerate_2k/regen2k_0.mp4",
+      "related": ["minimax-h3-t2v", "fasth3-8step-v2-cuda"]
+    },
     {
       "id": "fasth3-lora-preview",
       "family": "minimax_h3",

docs/cookbook/options.yaml

@@ -44,7 +44,7 @@ runtimes:
         host: 0.0.0.0
         port: 8000
         served_model_name: null
-    workloads: [t2v, i2v]
+    workloads: [t2v, i2v, v2v]
     options: &cuda_options
       server.served_model_name:
         anyOf:

docs/cookbook/recipes/open-h3-regenerate-2k.yaml

@@ -0,0 +1,18 @@
+title: Open-H3-Regenerate-2K
+summary: Regenerate a 768P MiniMax-H3 draft at 2560x1440 with three refine steps.
+
+deployments:
+  cuda-rest:
+    runtime: fastvideo-cuda-rest
+    guide: docs/inference/open_h3_regenerate_2k.md
+    workload: v2v
+    defaults: examples/serving/openai_open_h3_regenerate_2k.yaml
+    requirements:
+      - 'UV_TORCH_BACKEND=cu130 uv pip install -e ".[fasth3]"'
+      - "One GPU with about 150 GB free (B200 or GB200). Send the draft as video_reference; local paths must sit under server.media_roots."
+      - "Optional NVFP4 decoder (Blackwell only): hf download FastVideo/FastH3-NVFP4-VAE-Decoder --include 'quality-36L/decoder-nvfp4.safetensors' --local-dir fasth3-nvfp4-vae"
+    options:
+      - server.*
+      - generator.engine.offload.text_encoder
+      - generator.pipeline.experimental.h3_vae_nvfp4_decoder
+      - default_request.sampling.seed

docs/design/inference_schema_parity_inventory.yaml

@@ -153,6 +153,12 @@ surfaces:
 
   pipeline_config_extensions:
     preset_owned:
+      draft_size:
+        sources:
+          - fastvideo.configs.pipelines.open_h3_regenerate_2k.OpenH3Regenerate2KPipelineConfig
+      refine_start_sigma:
+        sources:
+          - fastvideo.configs.pipelines.open_h3_regenerate_2k.OpenH3Regenerate2KPipelineConfig
       audio_encoder_precision:
         sources:
           - fastvideo.configs.pipelines.wan.WanS2V14BConfig

docs/inference/open_h3_regenerate_2k.md

@@ -87,6 +87,8 @@ The response headers carry `X-Request-Id`, `X-Model`, `X-Inference-Time-S` and `
 - `size`, `width` and `height` may be omitted or `2560x1440`; `num_frames` may be omitted or 124 (the frame count of the
   VAE grid; the video has 121 frames); `fps` must be 24. `seconds` is not accepted.
 - `num_inference_steps` is 3 (default) or 2; `seed` sets the refine noise.
+- Serve with `attention_backend: TORCH_SDPA`: the refine stage selects cuDNN attention, which at about 133k tokens is
+  about 3x faster than PyTorch's default flash dispatch. The `FLASH_ATTN` backend bypasses that selection.
 
 ## NVFP4 video decoder
 
@@ -153,7 +155,7 @@ keeping structure; use it only when the decode budget is tight.
 
 ## Parity with the reference
 
-Measured on one GB200 (tray hpc-rack-1-5, BF16, dense attention) against the reference `standard_fwd` with the released
+Measured on one GB200 (BF16, `TORCH_SDPA` with cuDNN attention) against the reference `standard_fwd` with the released
 VAE, same draft, seed (1766580932), prompt and VAE posterior seed (42). The reference's noise is a CUDA generator draw
 that FastVideo reproduces (the same `torch.Generator(device="cuda")` seed), so no injected noise was needed.
 

docs/js/cookbook-workloads.mjs

@@ -32,6 +32,14 @@ export const WORKLOADS = Object.freeze({
     extension: "wav",
     media_type: "audio/wav",
   }),
+  v2v: Object.freeze({
+    input: "video",
+    output: "video",
+    endpoint: "/v1/videos/sync",
+    extension: "mp4",
+    media_type: "video/mp4",
+    reference: "video_reference",
+  }),
   v2a: Object.freeze({
     input: "video",
     output: "audio",

docs/tests/cookbook-config.test.mjs

@@ -488,6 +488,7 @@ test("sample requests follow shared workload media", () => {
     t2i: { endpoint: /\/v1\/images\/generations/, output: /--output output\.json/, reference: null },
     t2a: { endpoint: /\/v1\/videos\/sync/, output: /--output output\.wav/, reference: null },
     v2a: { endpoint: /\/v1\/videos\/sync/, output: /--output output\.wav/, reference: "video_reference" },
+    v2v: { endpoint: /\/v1\/videos\/sync/, output: /--output output\.mp4/, reference: "video_reference" },
   };
   for (const [workload, expected] of Object.entries(expectations)) {
     const data = catalog();

examples/inference/basic/basic_open_h3_regenerate_2k.py

@@ -38,7 +38,7 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--output", type=Path, default=Path("outputs/open_h3_regenerate_2k"))
     parser.add_argument("--steps", type=int, choices=(2, 3), default=3, help="refine steps (3 is the released setting)")
     parser.add_argument("--seed", type=int, default=1766580932)
-    parser.add_argument("--attention-backend", default="FLASH_ATTN")
+    parser.add_argument("--attention-backend", default="TORCH_SDPA")
     parser.add_argument("--draft-size", choices=("768p", "576p"), default="768p",
                         help="size the draft is encoded at; 576p for drafts from other video models")
     parser.add_argument("--vae-decoder", default=None,

examples/serving/openai_open_h3_regenerate_2k.yaml

@@ -17,7 +17,7 @@ generator:
       vae_enabled: false
   pipeline:
     experimental:
-      attention_backend: FLASH_ATTN
+      attention_backend: TORCH_SDPA
       # 768p for MiniMax-H3 drafts; 576p encodes drafts from other video models at 1024x576.
       draft_size: 768p
       # NVFP4 video decoder (Blackwell only). Download one tier from FastVideo/FastH3-NVFP4-VAE-Decoder and point

fastvideo-kernel/benchmarks/bench_vsa_tile_layout.py

@@ -5,7 +5,7 @@ Benchmark the SM100 fused VSA64 tile layout against the legacy scatter/transpose
 Times Wan VSA inference preprocessing through the backend
 (``VideoSparseAttentionImpl.preprocess_qkv``), toggling
 ``FASTVIDEO_DISABLE_VSA64_FUSED_LAYOUT``:
-  - legacy: tile scatter into BSHD + four BSHD->BHSD transposes
+  - legacy: tile scatter into BSHD + three contiguous Q/K/V transposes + a gate view
   - fused:  one Triton kernel writing tiled BHSD directly
 
 With ``--include_attn`` the timed region also covers ``forward`` (the
@@ -59,20 +59,20 @@ def parse_arguments() -> argparse.Namespace:
 
 
 def make_step(impl: VideoSparseAttentionImpl, qkvg: torch.Tensor, builder: VideoSparseAttentionMetadataBuilder,
-              build_kwargs: dict, fused: bool, include_attn: bool) -> Callable[[], torch.Tensor]:
+              build_kwargs: dict, fused: bool, include_attn: bool) -> Callable[[], tuple[torch.Tensor, ...]]:
     metadata = builder.build(**build_kwargs)
 
-    def step() -> torch.Tensor:
+    def step() -> tuple[torch.Tensor, ...]:
         # The routing gate reads the switch on every call.
         with envs.FASTVIDEO_DISABLE_VSA64_FUSED_LAYOUT.override(not fused):
             tiled = impl.preprocess_qkv(qkvg, metadata)
         q, k, v, gate = tiled.chunk(4)
         if include_attn:
-            return impl.forward(q, k, v, gate, metadata)
+            return (impl.forward(q, k, v, gate, metadata), )
         if metadata.fused_layout_active:
-            return tiled
-        # Legacy route: the transposes forward() applies before video_sparse_attn.
-        return torch.cat([t.transpose(1, 2).contiguous() for t in (q, k, v, gate)])
+            return q, k, v, gate
+        return (q.transpose(1, 2).contiguous(), k.transpose(1, 2).contiguous(), v.transpose(1, 2).contiguous(),
+                gate.transpose(1, 2))
 
     with torch.no_grad():
         step()
@@ -100,7 +100,7 @@ def main() -> None:
     legacy = make_step(impl, qkvg, builder, build_kwargs, fused=False, include_attn=args.include_attn)
     fused = make_step(impl, qkvg, builder, build_kwargs, fused=True, include_attn=args.include_attn)
     with torch.no_grad():
-        if not torch.equal(legacy(), fused()):
+        if not all(torch.equal(old, new) for old, new in zip(legacy(), fused(), strict=True)):
             raise RuntimeError("Fused and legacy outputs differ.")
         legacy_ms = do_bench(legacy, warmup=args.warmup, rep=args.rep, return_mode="median")
         fused_ms = do_bench(fused, warmup=args.warmup, rep=args.rep, return_mode="median")

fastvideo-kernel/tests/test_vsa_tile_layout.py

@@ -5,6 +5,9 @@ Production only routes SM100 here, but the kernel is plain Triton, so parity is
 checked on any CUDA device.
 """
 
+import importlib.util
+from pathlib import Path
+
 import pytest
 import torch
 
@@ -40,6 +43,63 @@ def test_tile_to_bhsd_matches_scatter_and_four_transposes(dit_shape):
     source_index[metadata.non_pad_index] = metadata.tile_partition_indices.to(torch.int32)
     out = torch.empty((4, 2, padded_sequence, 128), device="cuda", dtype=torch.bfloat16)
     actual = tile_to_bhsd(qkvg, source_index, out, sequence, padded_sequence, 2, 128).chunk(4)
-    for name, old, new in zip(("q", "k", "v", "gate"), expected, actual):
+    for name, old, new in zip(("q", "k", "v", "gate"), expected, actual, strict=True):
         assert old.shape == new.shape, name
         assert torch.equal(old, new), name
+
+
+@pytest.mark.parametrize("fused", [False, True])
+def test_layout_benchmark_matches_forward_preprocessing(monkeypatch, fused):
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA device required")
+
+    from fastvideo.attention.backends import video_sparse_attn as vsa
+
+    path = Path(__file__).resolve().parents[1] / "benchmarks/bench_vsa_tile_layout.py"
+    spec = importlib.util.spec_from_file_location("bench_vsa_tile_layout", path)
+    benchmark = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(benchmark)
+    original_gate = vsa._can_fuse_vsa64_layout
+
+    def allow_test_device(x, metadata):
+        with monkeypatch.context() as gate_patch:
+            gate_patch.setattr(torch.cuda, "get_device_capability", lambda *args: (10, 0))
+            return original_gate(x, metadata)
+
+    monkeypatch.setattr(vsa, "_can_fuse_vsa64_layout", allow_test_device)
+    impl = object.__new__(vsa.VideoSparseAttentionImpl)
+    builder = vsa.VideoSparseAttentionMetadataBuilder()
+    kwargs = dict(current_timestep=0,
+                  raw_latent_shape=(5, 5, 6),
+                  patch_size=(1, 1, 1),
+                  VSA_sparsity=0.8,
+                  device=torch.device("cuda"),
+                  cache_tile_buf=True)
+    qkvg = torch.randn((4, 157, 2, 128), device="cuda", dtype=torch.bfloat16)[:, :150]
+    captured = []
+
+    def capture_attention_inputs(q, k, v, *args, **kwargs):
+        captured[:] = [q, k, v, kwargs["compress_attn_weight"]]
+        return q
+
+    monkeypatch.setattr(vsa, "video_sparse_attn", capture_attention_inputs)
+    with torch.no_grad():
+        layout_step = benchmark.make_step(impl, qkvg, builder, kwargs, fused=fused, include_attn=False)
+        attention_step = benchmark.make_step(impl, qkvg, builder, kwargs, fused=fused, include_attn=True)
+
+        def reject_cat(*args, **kwargs):
+            pytest.fail("The timed layout step must not concatenate its outputs")
+
+        with monkeypatch.context() as timed_patch:
+            timed_patch.setattr(torch, "cat", reject_cat)
+            layout = layout_step()
+            attention = attention_step()
+
+    assert isinstance(layout, tuple) and len(layout) == 4
+    for actual, expected in zip(layout, captured, strict=True):
+        assert torch.equal(actual, expected)
+        assert actual.stride() == expected.stride()
+    assert all(t.is_contiguous() for t in layout[:3])
+    assert layout[3].is_contiguous() is fused
+    assert isinstance(attention, tuple) and len(attention) == 1
+    assert torch.equal(attention[0], captured[0].transpose(1, 2))

fastvideo/layers/lora/linear.py

@@ -254,13 +254,7 @@ class BaseLayerWithLoRA(nn.Module):
             current_device = self.base_layer.weight.data.device
             data = self.base_layer.weight.data.to(get_local_torch_device()).full_tensor()
 
-            # Apply LoRA with alpha scaling
-            lora_delta = (
-                self.slice_lora_b_weights(self.lora_B).to(data) @ self.slice_lora_a_weights(self.lora_A).to(data))
-            scale = ((self.lora_alpha / self.lora_rank) if (self.lora_alpha and self.lora_rank) else 1.0)
-            lora_delta *= scale * self.lora_strength
-            data += lora_delta
-            unsharded_base_layer.weight = nn.Parameter(data.to(current_device))
+            unsharded_base_layer.weight = nn.Parameter(self._merged_weight(data).to(current_device))
             if isinstance(getattr(self.base_layer, "bias", None), DTensor):
                 unsharded_base_layer.bias = nn.Parameter(
                     self.base_layer.bias.to(get_local_torch_device(),
@@ -277,18 +271,22 @@ class BaseLayerWithLoRA(nn.Module):
             current_device = self.base_layer.weight.data.device
             data = self.base_layer.weight.data.to(get_local_torch_device())
 
-            # Apply LoRA with alpha scaling. An adapter stored wider than the base (an FP32 adapter on a BF16 DiT)
-            # is multiplied and added at its own precision and rounded once, as a merge-then-cast reference does.
-            accumulate_dtype = torch.promote_types(data.dtype, self.lora_A.dtype)
-            lora_delta = (self.slice_lora_b_weights(self.lora_B.to(accumulate_dtype)) @ self.slice_lora_a_weights(
-                self.lora_A.to(accumulate_dtype)))
-            scale = ((self.lora_alpha / self.lora_rank) if (self.lora_alpha and self.lora_rank) else 1.0)
-            lora_delta *= scale * self.lora_strength
-            data = (data.to(accumulate_dtype) + lora_delta).to(data.dtype)
-            self.base_layer.weight.data = data.to(current_device, non_blocking=True)
+            self.base_layer.weight.data = self._merged_weight(data).to(current_device, non_blocking=True)
 
         self.merged = True
 
+    def _merged_weight(self, weight: torch.Tensor) -> torch.Tensor:
+        """``weight`` plus the scaled adapter delta, in ``weight``'s dtype.
+
+        An adapter stored wider than the base (an FP32 adapter on a BF16 DiT) is multiplied and added at its own
+        precision and rounded once, as a merge-then-cast reference does; with FSDP or without.
+        """
+        accumulate_dtype = torch.promote_types(weight.dtype, self.lora_A.dtype)
+        lora_B = self.slice_lora_b_weights(self.lora_B.to(device=weight.device, dtype=accumulate_dtype))
+        lora_A = self.slice_lora_a_weights(self.lora_A.to(device=weight.device, dtype=accumulate_dtype))
+        scale = ((self.lora_alpha / self.lora_rank) if (self.lora_alpha and self.lora_rank) else 1.0)
+        return (weight.to(accumulate_dtype) + (lora_B @ lora_A) * (scale * self.lora_strength)).to(weight.dtype)
+
     @torch.no_grad()
     # @torch.compile(dynamic=True)
     def unmerge_lora_weights(self, *, if_merged_only: bool = True) -> None:

fastvideo/models/loader/component_loader.py

@@ -1655,7 +1655,6 @@ class PipelineComponentLoader:
             component_model_path,
         )
 
-        # A component under a subfolder (e.g. "base_model/vae") dispatches on its own name.
         module_name = module_name.rsplit("/", 1)[-1]
         # Get the appropriate loader for this module type
         loader = ComponentLoader.for_module_type(module_name, transformers_or_diffusers)

fastvideo/pipelines/basic/open_h3_regenerate_2k/open_h3_regenerate_2k_pipeline.py

@@ -54,7 +54,6 @@ class OpenH3Regenerate2KPipeline(MiniMaxH3Pipeline):
                                                revision=revision,
                                                allow_patterns=self.get_hf_download_allow_patterns())
         if self.fastvideo_args.lora_path is None:
-            # The merged refine LoRA is part of the release; it is a constructor-time adapter like any H3 LoRA.
             self.fastvideo_args.lora_path = os.path.join(self.model_path, LORA_FILE)
         return verify_model_config_and_directory(os.path.join(self.model_path, BASE_MODEL_DIR),
                                                  required_component_dirs=())

fastvideo/pipelines/basic/open_h3_regenerate_2k/sharp_curve.py

@@ -34,6 +34,9 @@ def interpolate_curve(table: torch.Tensor, timesteps: torch.Tensor) -> torch.Ten
     return torch.lerp(table[low], table[low + 1], (position - low).unsqueeze(-1))
 
 
+_ATTACHED_ATTR = "_open_h3_sharp_curve_branches"
+
+
 class SharpCurveBranches:
 
     def __init__(self, weights_dir: str | Path) -> None:
@@ -50,15 +53,17 @@ class SharpCurveBranches:
         if len(self.factors) != NUM_CURVE_BRANCHES:
             raise ValueError(f"Expected {NUM_CURVE_BRANCHES} Sharp curve branches, found {len(self.factors)}.")
         self.coords: torch.Tensor | None = None
-        self._attached: set[int] = set()
         self._device_factors: dict[str, tuple[torch.Tensor, torch.Tensor]] = {}
 
     def set_timesteps(self, unique_timesteps: torch.Tensor) -> None:
         self.coords = interpolate_curve(self.table, unique_timesteps)
 
     def attach(self, transformer: torch.nn.Module) -> None:
-        """Register the branch hooks on ``transformer`` once."""
-        if id(transformer) in self._attached:
+        """Register the branch hooks on ``transformer`` unless this instance already carries them.
+
+        The mark lives on the module, so a reloaded lazy transformer, a new instance, gets its own hooks.
+        """
+        if getattr(transformer, _ATTACHED_ATTR, None) is self:
             return
         modules = dict(transformer.named_modules())
         missing = sorted(set(self.factors) - set(modules))
@@ -66,7 +71,7 @@ class SharpCurveBranches:
             raise ValueError(f"Sharp curve branches match no transformer module: {missing[:5]}")
         for name in self.factors:
             modules[name].register_forward_hook(self._hook(name))
-        self._attached.add(id(transformer))
+        setattr(transformer, _ATTACHED_ATTR, self)
 
     def _hook(self, name: str):
 

fastvideo/pipelines/basic/open_h3_regenerate_2k/stages/denoising.py

@@ -53,9 +53,7 @@ class Regenerate2KDenoisingStage(PipelineStage):
     @torch.no_grad()
     def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
         transformer = self.transformer
-        if is_lazy_module(transformer):
-            transformer.materialize()
-        self.curve.attach(transformer)
+        self.curve.attach(transformer.materialize() if is_lazy_module(transformer) else transformer)
         layout = batch.extra.get(MINIMAX_H3_LAYOUT_KEY)
         if not isinstance(layout, MiniMaxH3PackedLayout):
             raise ValueError("The MiniMax-H3 packed layout is missing before denoising.")

fastvideo/tests/layers/test_lora_merge_precision.py

@@ -42,3 +42,20 @@ def test_fp32_adapter_is_added_in_fp32_and_rounded_once() -> None:
 def test_bf16_adapter_keeps_the_bf16_merge() -> None:
     merged, base, a, b = _merged(torch.bfloat16)
     assert torch.equal(merged, base + b @ a)
+
+
+def test_merge_helper_used_with_and_without_fsdp_rounds_once() -> None:
+    torch.manual_seed(1)
+    layer = ReplicatedLinear(IN, OUT, bias=False, params_dtype=torch.bfloat16)
+    with torch.no_grad():
+        layer.weight.copy_(torch.randn(OUT, IN))
+    wrapped = get_lora_layer(layer)
+    a, b = torch.randn(RANK, IN), torch.randn(OUT, RANK) * 0.1
+    wrapped.lora_A, wrapped.lora_B = torch.nn.Parameter(a), torch.nn.Parameter(b)
+    wrapped.lora_rank, wrapped.lora_alpha, wrapped.lora_strength = RANK, RANK, 0.5
+    weight = layer.weight.detach().clone()
+
+    merged = wrapped._merged_weight(weight)
+
+    assert merged.dtype == torch.bfloat16
+    assert torch.equal(merged, (weight.float() + 0.5 * (b @ a)).to(torch.bfloat16))

fastvideo/tests/pipelines/test_open_h3_regenerate_2k.py

@@ -124,7 +124,7 @@ def test_sharp_branch_adds_curve_term_to_projection_output(tmp_path):
         block.proj = _Projection(6)
         model.blocks.append(block)
     curve.attach(model)
-    curve.attach(model)  # idempotent: a second attach must not double the term
+    curve.attach(model)
     timesteps = torch.tensor([0.1, 0.9])
     x = torch.randn(2, 4)
     base, _ = model.blocks[0].proj.linear(x), None
@@ -199,15 +199,51 @@ class _RecordingCurve:
     def __init__(self):
         self.timesteps = []
         self.attached = 0
+        self.targets = []
 
     def attach(self, transformer):
         self.attached += 1
+        self.targets.append(transformer)
 
     def set_timesteps(self, unique):
         self.timesteps.append(unique.clone())
 
 
-def _denoise(monkeypatch, steps, velocity=0.5):
+def _curve_model(out_features):
+    model = torch.nn.Module()
+    model.blocks = torch.nn.ModuleList()
+    for _ in range(NUM_CURVE_BRANCHES):
+        block = torch.nn.Module()
+        block.proj = _Projection(out_features)
+        model.blocks.append(block)
+    return model
+
+
+def test_sharp_branches_reach_a_reloaded_transformer_instance(tmp_path):
+    _write_curve_release(tmp_path, 6, NUM_CURVE_BRANCHES)
+    curve = SharpCurveBranches(tmp_path)
+    curve.set_timesteps(torch.tensor([0.1, 0.9]))
+    x = torch.randn(2, 4)
+    for model in (_curve_model(6), _curve_model(6)):
+        curve.attach(model)
+        out, _ = model.blocks[0].proj(x)
+        assert not torch.allclose(out, model.blocks[0].proj.linear(x))
+
+
+def test_stage_attaches_the_branches_to_each_materialized_lazy_transformer(monkeypatch):
+    from fastvideo.pipelines.lazy_module import LazyModule
+
+    lazy = LazyModule("transformer", lambda: _ConstantVelocityTransformer(0.5))
+    first = _denoise(monkeypatch, steps=2, transformer=lazy)[2]
+    loaded_first = first.targets[-1]
+    assert lazy.release()
+    second = _denoise(monkeypatch, steps=2, transformer=lazy, curve=first)[2]
+
+    assert not isinstance(loaded_first, LazyModule) and not isinstance(second.targets[-1], LazyModule)
+    assert second.targets[-1] is not loaded_first
+
+
+def _denoise(monkeypatch, steps, velocity=0.5, transformer=None, curve=None):
     from fastvideo.pipelines.basic.minimax_h3.packing import build_packed_sequence
     from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_latent_preparation import MINIMAX_H3_LAYOUT_KEY
     from fastvideo.pipelines.basic.open_h3_regenerate_2k.stages import denoising
@@ -221,7 +257,8 @@ def _denoise(monkeypatch, steps, velocity=0.5):
     batch.latents = torch.zeros(layout.video_indices.numel(), 8)
     batch.audio_latents = torch.ones(layout.audio_indices.numel(), 4)
     batch.extra[MINIMAX_H3_LAYOUT_KEY] = layout
-    transformer, curve = _ConstantVelocityTransformer(velocity), _RecordingCurve()
+    transformer = transformer if transformer is not None else _ConstantVelocityTransformer(velocity)
+    curve = curve if curve is not None else _RecordingCurve()
     stage = denoising.Regenerate2KDenoisingStage(transformer, curve, None)
     return stage.forward(batch, None), transformer, curve