[feat] cosmos3: native MoT sequence-packing (video) + framework parity

fastvideo/pipelines/basic/cosmos3/sequence_packing.py

@@ -0,0 +1,332 @@
+# SPDX-License-Identifier: Apache-2.0
+"""FastVideo-native Cosmos3 sequence packing (video subset).
+
+Numerical-parity port of the official ``cosmos_framework`` data packer
+(``cosmos_framework.data.vfm.sequence_packing.pack_input_sequence``) restricted
+to the VIDEO generation path that the FastVideo Cosmos3 DiT consumes (T2V / I2V
+/ T2I). It builds, per sample, two splits:
+
+* a ``causal`` text split (prompt token ids, plus the trailing ``eos`` and
+  ``start_of_generation`` markers the framework appends when a generation
+  modality follows), and
+* a ``full`` vision split (VAE latent patch tokens).
+
+The 3D-MRoPE position ids ``[3, seq]`` are produced exactly like the framework:
+text tokens broadcast a single monotone id across the (t, h, w) axes, the
+temporal offset is bumped by ``temporal_modality_margin`` at the text->vision
+boundary, and vision tokens lay out a (T, H, W) grid with spatial ids reset per
+segment. Condition frames (I2V cond frame 0, T2I single conditioned frame, ...)
+are kept in the packed sequence and rope grid but excluded from the MSE-loss /
+timestep bookkeeping, mirroring the framework.
+
+The output ``Cosmos3PackedSequence`` maps 1:1 onto the
+``Cosmos3VFMTransformer.forward`` kwargs via :meth:`to_dit_kwargs`. This module
+is pure torch/python; it imports no diffusers/transformers model classes.
+
+Reference of record: ``cosmos_framework`` (NVIDIA), the parity oracle used by
+``tests/local_tests/cosmos3/test_cosmos3_packing_parity.py``.
+"""
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass, field
+from typing import Any
+
+import torch
+
+from fastvideo.models.dits.cosmos3 import (
+    compute_mrope_position_ids_text,
+    compute_mrope_position_ids_vision,
+)
+
+__all__ = [
+    "Cosmos3VisionItem",
+    "Cosmos3SampleInputs",
+    "Cosmos3PackedSequence",
+    "pack_cosmos3_video_sequence",
+]
+
+
+# ---------------------------------------------------------------------------
+# Inputs
+# ---------------------------------------------------------------------------
+@dataclass
+class Cosmos3VisionItem:
+    """One vision latent for a sample.
+
+    Args:
+        latent: VAE latent ``[C, T, H, W]`` (a leading batch axis of size 1 is
+            accepted and squeezed).
+        condition_frame_indexes: Latent-frame indices that are *conditioned*
+            (clean) rather than noisy. ``[]`` for T2V, ``[0]`` for I2V, and the
+            single conditioned frame for T2I.
+        fps: Frames-per-second for this clip; only used when
+            ``enable_fps_modulation`` is set.
+    """
+
+    latent: torch.Tensor
+    condition_frame_indexes: list[int] = field(default_factory=list)
+    fps: float | None = None
+
+
+@dataclass
+class Cosmos3SampleInputs:
+    """Per-sample packing inputs (one text prompt + its vision item)."""
+
+    text_ids: list[int]
+    vision: Cosmos3VisionItem
+    timestep: float
+
+
+# ---------------------------------------------------------------------------
+# Output
+# ---------------------------------------------------------------------------
+@dataclass
+class Cosmos3PackedSequence:
+    """Packed-sequence inputs consumed by ``Cosmos3VFMTransformer.forward``.
+
+    Field names mirror the framework ``PackedSequence`` (+ its ``vision``
+    ``ModalityData``) so the parity test can compare field-by-field.
+    """
+
+    # Sequence structure.
+    sample_lens: list[int]
+    split_lens: list[int]
+    attn_modes: list[str]
+    sequence_length: int
+    is_image_batch: bool
+
+    # Text modality.
+    text_ids: torch.Tensor
+    text_indexes: torch.Tensor
+    position_ids: torch.Tensor  # [3, sequence_length]
+
+    # Vision modality.
+    vision_tokens: list[torch.Tensor]
+    vision_token_shapes: list[tuple[int, int, int]]
+    vision_sequence_indexes: torch.Tensor
+    vision_timesteps: torch.Tensor
+    vision_mse_loss_indexes: torch.Tensor
+    vision_noisy_frame_indexes: list[torch.Tensor]
+    vision_condition_mask: list[torch.Tensor]
+    fps_vision: torch.Tensor | None = None
+
+    def to_dit_kwargs(self) -> dict[str, Any]:
+        """Return the kwargs dict for ``Cosmos3VFMTransformer.forward``."""
+        return dict(
+            text_ids=self.text_ids,
+            text_indexes=self.text_indexes,
+            position_ids=self.position_ids,
+            sequence_length=int(self.sequence_length),
+            split_lens=list(self.split_lens),
+            attn_modes=list(self.attn_modes),
+            vision_tokens=list(self.vision_tokens),
+            vision_token_shapes=list(self.vision_token_shapes),
+            vision_sequence_indexes=self.vision_sequence_indexes,
+            vision_timesteps=self.vision_timesteps,
+            vision_mse_loss_indexes=self.vision_mse_loss_indexes,
+            vision_noisy_frame_indexes=list(self.vision_noisy_frame_indexes),
+            fps_vision=self.fps_vision,
+        )
+
+
+# ---------------------------------------------------------------------------
+# Packing
+# ---------------------------------------------------------------------------
+def pack_cosmos3_video_sequence(
+    samples: list[Cosmos3SampleInputs],
+    special_tokens: dict[str, int],
+    *,
+    latent_patch_size: int = 2,
+    include_end_of_generation_token: bool = False,
+    temporal_modality_margin: int = 15_000,
+    reset_spatial_ids: bool = True,
+    enable_fps_modulation: bool = False,
+    base_fps: float = 24.0,
+    temporal_compression_factor: int = 4,
+    initial_mrope_temporal_offset: int | float = 0,
+) -> Cosmos3PackedSequence:
+    """Pack prompts + vision latents into the Cosmos3 DiT packed-sequence inputs.
+
+    Video subset of ``cosmos_framework`` ``pack_input_sequence`` under
+    ``unified_3d_mrope``: each sample is ``[causal text, full vision]``.
+
+    Args:
+        samples: Per-sample text prompt token ids + vision item + timestep.
+        special_tokens: Must contain ``eos_token_id`` and
+            ``start_of_generation`` (and ``end_of_generation`` if
+            ``include_end_of_generation_token``). ``bos_token_id`` is honored if
+            present (prepended) to match the framework.
+        latent_patch_size: Latent patch size used by the DiT.
+        include_end_of_generation_token: Append the framework's end-of-generation
+            marker after the vision split.
+        temporal_modality_margin: Temporal-offset bump applied at the
+            text->vision boundary (``unified_3d_mrope_temporal_modality_margin``).
+        reset_spatial_ids: Reset vision spatial ids to 0 per segment.
+        enable_fps_modulation: Use float, fps-scaled temporal positions.
+        base_fps: Base FPS used when ``enable_fps_modulation``.
+        temporal_compression_factor: VAE temporal compression factor.
+        initial_mrope_temporal_offset: Per-sample starting temporal offset.
+
+    Returns:
+        A :class:`Cosmos3PackedSequence`.
+    """
+    assert "eos_token_id" in special_tokens, "special_tokens must contain eos_token_id"
+    assert "start_of_generation" in special_tokens, "special_tokens must contain start_of_generation"
+    if latent_patch_size < 1:
+        raise ValueError(f"latent_patch_size must be >= 1, got {latent_patch_size}")
+
+    # Build-time accumulators (concatenated across samples).
+    sample_lens: list[int] = []
+    split_lens: list[int] = []
+    attn_modes: list[str] = []
+
+    text_ids: list[int] = []
+    text_indexes: list[int] = []
+    position_id_blocks: list[torch.Tensor] = []  # each [3, n]
+
+    vision_tokens: list[torch.Tensor] = []
+    vision_token_shapes: list[tuple[int, int, int]] = []
+    vision_sequence_indexes: list[int] = []
+    vision_timesteps: list[float] = []
+    vision_mse_loss_indexes: list[int] = []
+    vision_noisy_frame_indexes: list[torch.Tensor] = []
+    vision_condition_mask: list[torch.Tensor] = []
+    fps_values: list[float] = []
+
+    curr = 0  # running position in the packed sequence
+    is_image_batch = True
+
+    for sample in samples:
+        temporal_offset: int | float = initial_mrope_temporal_offset
+        sample_len = 0
+
+        # ---- 1. Text split (causal) ----
+        if "bos_token_id" in special_tokens:
+            shifted_text_ids = [special_tokens["bos_token_id"], *sample.text_ids]
+        else:
+            shifted_text_ids = list(sample.text_ids)
+        # The video path always has a following generation modality, so the
+        # framework appends eos + start_of_generation.
+        shifted_text_ids = [*shifted_text_ids, special_tokens["eos_token_id"], special_tokens["start_of_generation"]]
+        text_split_len = len(shifted_text_ids)
+
+        text_ids.extend(shifted_text_ids)
+        text_indexes.extend(range(curr, curr + text_split_len))
+
+        text_mrope, temporal_offset = compute_mrope_position_ids_text(
+            num_tokens=text_split_len,
+            temporal_offset=int(temporal_offset),
+        )
+        position_id_blocks.append(text_mrope)
+
+        attn_modes.append("causal")
+        split_lens.append(text_split_len)
+        curr += text_split_len
+        sample_len += text_split_len
+
+        # End of text modality: bump temporal offset before vision.
+        temporal_offset += temporal_modality_margin
+
+        # ---- 2. Vision split (full) ----
+        latent = sample.vision.latent
+        latent = latent.squeeze(0) if latent.dim() == 5 else latent  # [C, T, H, W]
+        _c, latent_t, latent_h, latent_w = latent.shape
+        patch_h = math.ceil(latent_h / latent_patch_size)
+        patch_w = math.ceil(latent_w / latent_patch_size)
+        num_vision_tokens = latent_t * patch_h * patch_w
+
+        vision_tokens.append(sample.vision.latent)
+        vision_token_shapes.append((latent_t, patch_h, patch_w))
+        vision_sequence_indexes.extend(range(curr, curr + num_vision_tokens))
+
+        condition_set = {idx for idx in sample.vision.condition_frame_indexes if 0 <= idx < latent_t}
+        cond_mask = torch.zeros((latent_t, 1, 1), device=latent.device, dtype=latent.dtype)
+        for frame_idx in condition_set:
+            cond_mask[frame_idx, 0, 0] = 1.0
+        vision_condition_mask.append(cond_mask)
+
+        noisy_frames = torch.tensor(
+            [idx for idx in range(latent_t) if idx not in condition_set],
+            device=latent.device,
+            dtype=torch.long,
+        )
+        vision_noisy_frame_indexes.append(noisy_frames)
+
+        # MSE-loss indices + per-token timesteps cover only the noisy frames.
+        frame_token_stride = patch_h * patch_w
+        for frame_idx in range(latent_t):
+            if frame_idx in condition_set:
+                continue
+            frame_start = curr + frame_idx * frame_token_stride
+            vision_mse_loss_indexes.extend(range(frame_start, frame_start + frame_token_stride))
+            vision_timesteps.extend([float(sample.timestep)] * frame_token_stride)
+
+        vision_fps = sample.vision.fps if enable_fps_modulation else None
+        if vision_fps is not None:
+            fps_values.append(float(vision_fps))
+        vision_mrope, temporal_offset = compute_mrope_position_ids_vision(
+            grid_t=latent_t,
+            grid_h=patch_h,
+            grid_w=patch_w,
+            temporal_offset=temporal_offset,
+            fps=vision_fps,
+            base_fps=base_fps,
+            temporal_compression_factor=temporal_compression_factor,
+            enable_fps_modulation=enable_fps_modulation,
+        )
+        position_id_blocks.append(vision_mrope)
+
+        curr += num_vision_tokens
+        sample_len += num_vision_tokens
+
+        # ---- 3. Optional end-of-generation marker ----
+        eov_len = 0
+        if include_end_of_generation_token:
+            assert "end_of_generation" in special_tokens, ("special_tokens must contain end_of_generation when "
+                                                           "include_end_of_generation_token=True")
+            text_ids.append(special_tokens["end_of_generation"])
+            text_indexes.append(curr)
+            eov_dtype = torch.float32 if enable_fps_modulation else torch.long
+            eov_ids = torch.full((3, 1), temporal_offset, dtype=eov_dtype)
+            position_id_blocks.append(eov_ids)
+            temporal_offset += 1
+            curr += 1
+            eov_len = 1
+            sample_len += 1
+
+        # The vision split and any trailing eov marker share the "full" split.
+        attn_modes.append("full")
+        split_lens.append(num_vision_tokens + eov_len)
+        sample_lens.append(sample_len)
+
+        if latent_t != 1:
+            is_image_batch = False
+
+    sequence_length = sum(sample_lens)
+
+    # position_ids: float iff any block is float (fps modulation path).
+    any_float = any(b.dtype.is_floating_point for b in position_id_blocks)
+    if any_float:
+        position_id_blocks = [b.to(torch.float32) for b in position_id_blocks]
+    position_ids = torch.cat(position_id_blocks, dim=1)  # [3, sequence_length]
+
+    timesteps_dtype = torch.float32
+    return Cosmos3PackedSequence(
+        sample_lens=sample_lens,
+        split_lens=split_lens,
+        attn_modes=attn_modes,
+        sequence_length=sequence_length,
+        is_image_batch=is_image_batch,
+        text_ids=torch.tensor(text_ids, dtype=torch.long),
+        text_indexes=torch.tensor(text_indexes, dtype=torch.long),
+        position_ids=position_ids,
+        vision_tokens=vision_tokens,
+        vision_token_shapes=vision_token_shapes,
+        vision_sequence_indexes=torch.tensor(vision_sequence_indexes, dtype=torch.long),
+        vision_timesteps=torch.tensor(vision_timesteps, dtype=timesteps_dtype),
+        vision_mse_loss_indexes=torch.tensor(vision_mse_loss_indexes, dtype=torch.long),
+        vision_noisy_frame_indexes=vision_noisy_frame_indexes,
+        vision_condition_mask=vision_condition_mask,
+        fps_vision=(torch.tensor(fps_values, dtype=torch.float32) if fps_values else None),
+    )

tests/local_tests/cosmos3/test_cosmos3_packing_parity.py

@@ -0,0 +1,312 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Numerical-parity test: FastVideo Cosmos3 sequence packing vs the OFFICIAL framework.
+
+FastVideo's native packer
+(``fastvideo.pipelines.basic.cosmos3.sequence_packing.pack_cosmos3_video_sequence``)
+builds the packed-sequence inputs the ``Cosmos3VFMTransformer`` consumes. This
+test asserts, for the SAME logical inputs (prompt token ids, vision latents,
+condition-frame indices, diffusion timestep, fps), that FastVideo's packing
+matches the official ``cosmos_framework.data.vfm.sequence_packing.pack_input_sequence``
+oracle field-by-field:
+
+  * ``position_ids`` (exact, ``[3, seq]``),
+  * ``text_ids`` / ``text_indexes``,
+  * ``split_lens`` / ``attn_modes`` / ``sample_lens`` / ``sequence_length``,
+  * vision ``sequence_indexes`` / ``token_shapes`` / ``timesteps`` /
+    ``mse_loss_indexes`` / ``noisy_frame_indexes`` / ``condition_mask``.
+
+Coverage spans T2V (no condition frames), I2V (condition frame 0), and T2I
+(single conditioned frame), across multiple grids, plus a multi-sample batch.
+
+Then BOTH the framework-packed and FastVideo-packed inputs are fed through the
+SAME tiny FastVideo DiT (framework weights copied in as in the existing DiT
+parity tests). Asserting bit-identical DiT output confirms FastVideo's own
+packing drives the DiT to the same result as the framework's packing.
+
+The official framework is the parity ORACLE; it runs on CPU / float32 via the
+SDPA monkey-patch from ``test_cosmos3_reference_forward``.
+
+Run:
+    cd <worktree> && <fv-cosmos3 python> -m pytest \
+        tests/local_tests/cosmos3/test_cosmos3_packing_parity.py -q
+"""
+from __future__ import annotations
+
+import pytest
+import torch
+
+# The official framework provides the parity oracle.
+cosmos_framework = pytest.importorskip(
+    "cosmos_framework",
+    reason="cosmos_framework not installed; run in fv-cosmos3 env.",
+)
+
+# Reuse the DiT parity helpers (weight copy + framework->DiT kwarg builder) and
+# the mRoPE tiny-model builders (real-checkpoint rope constants).
+from .test_cosmos3_dit_parity import (  # noqa: E402
+    _copy_weights,
+    _fastvideo_inputs_from_packed_seq,
+)
+from .test_cosmos3_dit_parity_mrope import (  # noqa: E402
+    _LATENT_CHANNEL,
+    _LATENT_PATCH_SIZE,
+    _MROPE_SECTION,
+    _RESET_SPATIAL_IDS,
+    _ROPE_THETA,
+    _TCF,
+    _TEMPORAL_MODALITY_MARGIN,
+    _build_tiny_cosmos3_mrope,
+    _build_tiny_fastvideo_dit_mrope,
+)
+from .test_cosmos3_reference_forward import _apply_sdpa_patches  # noqa: E402
+
+pytestmark = [pytest.mark.local]
+
+# Ensure the CPU/float32 SDPA patches are installed (idempotent).
+_apply_sdpa_patches()
+
+# Tiny special-token ids (kept < tiny vocab_size=64). The video path appends
+# eos + start_of_generation after the prompt tokens.
+_SPECIAL_TOKENS = {
+    "start_of_generation": 60,
+    "end_of_generation": 61,
+    "eos_token_id": 62,
+}
+
+
+# ---------------------------------------------------------------------------
+# Builders for the two packers from the SAME logical sample inputs.
+# ---------------------------------------------------------------------------
+def _make_vision(grid_t: int, latent_h: int, latent_w: int, seed: int) -> torch.Tensor:
+    """Deterministic VAE latent ``[1, C, T, H, W]``."""
+    torch.manual_seed(seed)
+    return torch.randn(1, _LATENT_CHANNEL, grid_t, latent_h, latent_w)
+
+
+def _framework_pack(
+    *,
+    text_ids_per_sample: list[list[int]],
+    visions: list[torch.Tensor],
+    cond_frames_per_sample: list[list[int]],
+    timesteps: list[float],
+    is_image_batch: bool,
+):
+    from cosmos_framework.data.vfm.sequence_packing import (
+        GenerationDataClean,
+        SequencePlan,
+        pack_input_sequence,
+    )
+
+    gen_data_clean = GenerationDataClean(
+        batch_size=len(visions),
+        is_image_batch=is_image_batch,
+        x0_tokens_vision=list(visions),
+        fps_vision=None,
+        num_vision_items_per_sample=[1] * len(visions),
+    )
+    plans = [
+        SequencePlan(has_text=True, has_vision=True, condition_frame_indexes_vision=list(cf))
+        for cf in cond_frames_per_sample
+    ]
+    return pack_input_sequence(
+        sequence_plans=plans,
+        input_text_indexes=[list(t) for t in text_ids_per_sample],
+        gen_data_clean=gen_data_clean,
+        input_timesteps=torch.tensor(timesteps, dtype=torch.float32),
+        special_tokens=_SPECIAL_TOKENS,
+        latent_patch_size=_LATENT_PATCH_SIZE,
+        include_end_of_generation_token=False,
+        position_embedding_type="unified_3d_mrope",
+        unified_3d_mrope_reset_spatial_ids=_RESET_SPATIAL_IDS,
+        unified_3d_mrope_temporal_modality_margin=_TEMPORAL_MODALITY_MARGIN,
+        enable_fps_modulation=False,
+        base_fps=24.0,
+        temporal_compression_factor=_TCF,
+    )
+
+
+def _fastvideo_pack(
+    *,
+    text_ids_per_sample: list[list[int]],
+    visions: list[torch.Tensor],
+    cond_frames_per_sample: list[list[int]],
+    timesteps: list[float],
+):
+    from fastvideo.pipelines.basic.cosmos3.sequence_packing import (
+        Cosmos3SampleInputs,
+        Cosmos3VisionItem,
+        pack_cosmos3_video_sequence,
+    )
+
+    samples = [
+        Cosmos3SampleInputs(
+            text_ids=list(t),
+            vision=Cosmos3VisionItem(latent=v, condition_frame_indexes=list(cf)),
+            timestep=float(ts),
+        )
+        for t, v, cf, ts in zip(text_ids_per_sample, visions, cond_frames_per_sample, timesteps)
+    ]
+    return pack_cosmos3_video_sequence(
+        samples,
+        _SPECIAL_TOKENS,
+        latent_patch_size=_LATENT_PATCH_SIZE,
+        include_end_of_generation_token=False,
+        temporal_modality_margin=_TEMPORAL_MODALITY_MARGIN,
+        reset_spatial_ids=_RESET_SPATIAL_IDS,
+        enable_fps_modulation=False,
+        base_fps=24.0,
+        temporal_compression_factor=_TCF,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Field-by-field comparison.
+# ---------------------------------------------------------------------------
+def _assert_packs_match(fw, fv) -> None:
+    """Assert the framework PackedSequence and FastVideo pack agree field-by-field."""
+    # Structure.
+    assert fv.split_lens == list(fw.split_lens), f"split_lens: fv={fv.split_lens} fw={list(fw.split_lens)}"
+    assert fv.attn_modes == list(fw.attn_modes), f"attn_modes: fv={fv.attn_modes} fw={list(fw.attn_modes)}"
+    assert fv.sample_lens == list(fw.sample_lens), f"sample_lens: fv={fv.sample_lens} fw={list(fw.sample_lens)}"
+    assert int(fv.sequence_length) == int(fw.sequence_length)
+
+    # Text.
+    torch.testing.assert_close(fv.text_ids, fw.text_ids.to(torch.long), rtol=0, atol=0)
+    torch.testing.assert_close(fv.text_indexes, fw.text_indexes.to(torch.long), rtol=0, atol=0)
+
+    # position_ids: exact, [3, seq], same dtype.
+    assert fv.position_ids.shape == fw.position_ids.shape, (
+        f"position_ids shape: fv={tuple(fv.position_ids.shape)} fw={tuple(fw.position_ids.shape)}")
+    assert fv.position_ids.dtype == fw.position_ids.dtype, (
+        f"position_ids dtype: fv={fv.position_ids.dtype} fw={fw.position_ids.dtype}")
+    torch.testing.assert_close(fv.position_ids, fw.position_ids, rtol=0, atol=0)
+
+    # Vision.
+    fwv = fw.vision
+    torch.testing.assert_close(fv.vision_sequence_indexes, fwv.sequence_indexes.to(torch.long), rtol=0, atol=0)
+    assert fv.vision_token_shapes == [tuple(s) for s in fwv.token_shapes], (
+        f"token_shapes: fv={fv.vision_token_shapes} fw={[tuple(s) for s in fwv.token_shapes]}")
+    torch.testing.assert_close(fv.vision_timesteps.to(torch.float32), fwv.timesteps.to(torch.float32))
+    torch.testing.assert_close(fv.vision_mse_loss_indexes, fwv.mse_loss_indexes.to(torch.long), rtol=0, atol=0)
+    assert len(fv.vision_noisy_frame_indexes) == len(fwv.noisy_frame_indexes)
+    for a, b in zip(fv.vision_noisy_frame_indexes, fwv.noisy_frame_indexes):
+        torch.testing.assert_close(a.to(torch.long), b.to(torch.long), rtol=0, atol=0)
+    assert len(fv.vision_condition_mask) == len(fwv.condition_mask)
+    for a, b in zip(fv.vision_condition_mask, fwv.condition_mask):
+        torch.testing.assert_close(a.flatten().to(torch.float32), b.flatten().to(torch.float32))
+
+
+# (grid_t, latent_h, latent_w, n_text, cond_frames, id) — single-sample cases.
+_CASES = [
+    pytest.param(1, 8, 8, 4, [], id="t2i_1x4x4"),
+    pytest.param(1, 4, 4, 5, [0], id="t2i_cond_1x2x2"),
+    pytest.param(2, 4, 4, 4, [], id="t2v_2x2x2"),
+    pytest.param(3, 8, 4, 6, [], id="t2v_3x4x2"),
+    pytest.param(2, 4, 4, 5, [0], id="i2v_2x2x2"),
+    pytest.param(3, 4, 4, 4, [0], id="i2v_3x2x2"),
+]
+
+
+class TestCosmos3PackingParity:
+
+    # -- Field-by-field packing parity -------------------------------------
+    @pytest.mark.parametrize(("grid_t", "latent_h", "latent_w", "n_text", "cond"), _CASES)
+    def test_packing_fields_match_framework(self, grid_t, latent_h, latent_w, n_text, cond):
+        torch.manual_seed(0)
+        text_ids = torch.randint(0, 60, (n_text,)).tolist()
+        vision = _make_vision(grid_t, latent_h, latent_w, seed=123)
+        timestep = 500.0
+
+        fw = _framework_pack(
+            text_ids_per_sample=[text_ids],
+            visions=[vision],
+            cond_frames_per_sample=[cond],
+            timesteps=[timestep],
+            is_image_batch=(grid_t == 1),
+        )
+        fv = _fastvideo_pack(
+            text_ids_per_sample=[text_ids],
+            visions=[vision],
+            cond_frames_per_sample=[cond],
+            timesteps=[timestep],
+        )
+        _assert_packs_match(fw, fv)
+
+    def test_packing_fields_match_framework_multi_sample(self):
+        """A batch of two samples (T2V + I2V) packs identically to the framework."""
+        torch.manual_seed(1)
+        t0 = torch.randint(0, 60, (3,)).tolist()
+        t1 = torch.randint(0, 60, (5,)).tolist()
+        v0 = _make_vision(2, 4, 4, seed=11)
+        v1 = _make_vision(2, 4, 4, seed=22)
+        kwargs = dict(
+            text_ids_per_sample=[t0, t1],
+            visions=[v0, v1],
+            cond_frames_per_sample=[[], [0]],
+            timesteps=[500.0, 250.0],
+        )
+        fw = _framework_pack(is_image_batch=False, **kwargs)
+        fv = _fastvideo_pack(**kwargs)
+        _assert_packs_match(fw, fv)
+
+    # -- End-to-end: FastVideo packing drives the DiT identically ----------
+    @pytest.mark.parametrize(("grid_t", "latent_h", "latent_w", "n_text", "cond"), _CASES)
+    def test_fastvideo_packing_drives_dit_like_framework(self, grid_t, latent_h, latent_w, n_text, cond):
+        """Feed BOTH the framework-packed and FastVideo-packed inputs through the
+        SAME FastVideo DiT (framework weights copied in); assert identical output.
+        """
+        num_layers = 2
+        torch.manual_seed(0)
+        text_ids = torch.randint(0, 60, (n_text,)).tolist()
+        vision = _make_vision(grid_t, latent_h, latent_w, seed=123)
+        timestep = 500.0
+
+        fw_pack = _framework_pack(
+            text_ids_per_sample=[text_ids],
+            visions=[vision],
+            cond_frames_per_sample=[cond],
+            timesteps=[timestep],
+            is_image_batch=(grid_t == 1),
+        )
+        fv_pack = _fastvideo_pack(
+            text_ids_per_sample=[text_ids],
+            visions=[vision],
+            cond_frames_per_sample=[cond],
+            timesteps=[timestep],
+        )
+        # Guard: the two packs must agree before we trust the DiT comparison.
+        _assert_packs_match(fw_pack, fv_pack)
+
+        # One DiT instance, framework weights copied in (parity oracle weights).
+        vfm = _build_tiny_cosmos3_mrope(seed=42, num_layers=num_layers)
+        dit = _build_tiny_fastvideo_dit_mrope(num_layers=num_layers)
+        _copy_weights(vfm, dit)
+
+        with torch.no_grad():
+            out_fw = dit(**_fastvideo_inputs_from_packed_seq(fw_pack))
+            out_fv = dit(**fv_pack.to_dit_kwargs())
+
+        # last_hidden_state must be bit-identical.
+        lhs_fw = out_fw["last_hidden_state"]
+        lhs_fv = out_fv["last_hidden_state"]
+        assert lhs_fw.shape == lhs_fv.shape
+        max_abs_lhs = (lhs_fw - lhs_fv).abs().max().item()
+        print(f"\n[packing->dit {grid_t}x{latent_h}x{latent_w} cond={cond}] "
+              f"last_hidden_state max abs diff = {max_abs_lhs:.3e}")
+        torch.testing.assert_close(lhs_fv, lhs_fw, rtol=0, atol=0)
+
+        # preds_vision must be bit-identical when there are noisy frames to
+        # predict. (A fully-conditioned clip has no noisy patches, so the DiT
+        # emits no "preds_vision" — both packs agree the mse-loss set is empty,
+        # already asserted by the field-parity guard above.)
+        has_preds = fv_pack.vision_mse_loss_indexes.numel() > 0
+        assert ("preds_vision" in out_fw) == has_preds
+        assert ("preds_vision" in out_fv) == has_preds
+        if has_preds:
+            pv_fw = out_fw["preds_vision"][0]
+            pv_fv = out_fv["preds_vision"][0]
+            assert pv_fw.shape == pv_fv.shape
+            max_abs_pv = (pv_fw - pv_fv).abs().max().item()
+            print(f"[packing->dit {grid_t}x{latent_h}x{latent_w} cond={cond}] "
+                  f"preds_vision max abs diff = {max_abs_pv:.3e}")
+            torch.testing.assert_close(pv_fv, pv_fw, rtol=0, atol=0)