diff --git a/.file_mapping.json b/.file_mapping.json
index e95543c0..f40162a4 100644
--- a/.file_mapping.json
+++ b/.file_mapping.json
@@ -1,7 +1,7 @@
 {
-  "_source_commit": "97901d395925c227163a7370ec2763b4057783d5",
-  "_dest_commit": "058c8c046877f087e07b3472a76cac214f3ad39c",
-  "_generated_at": "2026-07-23T11:16:20Z",
+  "_source_commit": "a78ca94df154c804f1c32c0dff6be38a3d22abeb-dirty",
+  "_dest_commit": "2190a72a5799ab1bc8e4ddc19ae09cd1f14d7cbc",
+  "_generated_at": "2026-07-27T08:25:15Z",
   "files": {
     "imaginaire/__init__.py": "cosmos_framework/__init__.py",
     "imaginaire/attention/__init__.py": "cosmos_framework/model/attention/__init__.py",
@@ -334,6 +334,8 @@
     "projects/cosmos3/cosmos3/models/mot/cosmos3_vfm_network.py": "cosmos_framework/model/generator/mot/cosmos3_vfm_network.py",
     "projects/cosmos3/cosmos3/models/mot/domain_aware_linear.py": "cosmos_framework/model/generator/mot/domain_aware_linear.py",
     "projects/cosmos3/cosmos3/models/mot/dot_product_attention.py": "cosmos_framework/model/generator/mot/dot_product_attention.py",
+    "projects/cosmos3/cosmos3/models/mot/flex_attention.py": "cosmos_framework/model/generator/mot/flex_attention.py",
+    "projects/cosmos3/cosmos3/models/mot/flex_attention_test.py": "cosmos_framework/model/generator/mot/flex_attention_test.py",
     "projects/cosmos3/cosmos3/models/mot/inference_text_kv_memory.py": "cosmos_framework/model/generator/mot/inference_text_kv_memory.py",
     "projects/cosmos3/cosmos3/models/mot/modeling_utils.py": "cosmos_framework/model/generator/mot/modeling_utils.py",
     "projects/cosmos3/cosmos3/models/mot/parallelize_unified_mot.py": "cosmos_framework/model/generator/mot/parallelize_unified_mot.py",
diff --git a/cosmos_framework/checkpoint/dcp.py b/cosmos_framework/checkpoint/dcp.py
index 43145c88..ba5f996a 100644
--- a/cosmos_framework/checkpoint/dcp.py
+++ b/cosmos_framework/checkpoint/dcp.py
@@ -898,7 +898,6 @@ def load(
 
                     # Use rank-specific key for RNG state to support correct per-rank restoration
                     rng_key = f"rng_state_{dist.get_rank()}"
-                    current_rng_state = get_rand_state_dict()
                     _state_dict = {
                         "grad_scaler": grad_scaler.state_dict(),
                         "iteration": iteration,
@@ -909,7 +908,7 @@ def load(
                         k.startswith(f"{rng_key}.") or k == rng_key for k in metadata.state_dict_metadata.keys()
                     )
                     if rng_key_exists:
-                        _state_dict[rng_key] = current_rng_state
+                        _state_dict[rng_key] = get_rand_state_dict()
 
                     dcp.load(
                         _state_dict,
@@ -919,7 +918,8 @@ def load(
                     )
                     grad_scaler.load_state_dict(_state_dict["grad_scaler"])
                     iteration = _state_dict["iteration"]
-                    set_rand_state_dict(_state_dict.get(rng_key, current_rng_state))
+                    if rng_key_exists:
+                        set_rand_state_dict(_state_dict[rng_key])
 
                 elif key == "dataloader":
                     if not easy_io.exists(cur_key_ckpt_full_path, backend_key=self.load_s3_backend_key):
diff --git a/cosmos_framework/data/generator/action/domain_utils.py b/cosmos_framework/data/generator/action/domain_utils.py
index 66cc51cb..03204968 100644
--- a/cosmos_framework/data/generator/action/domain_utils.py
+++ b/cosmos_framework/data/generator/action/domain_utils.py
@@ -23,6 +23,7 @@
     "xdof_yam": 16,
     "molmoact2_yam": 16,  # MolmoAct2 uses the same YAM 20D FK action contract
     "fractal": 20,
+    "drawanything": 21,
 }
 
 
@@ -43,6 +44,7 @@
     "xdof_yam": 20,
     "molmoact2_yam": 20,
     "fractal": 10,
+    "drawanything": 3,
     # NOTE: ``libero`` (7/10/13 depending on ``rotation_space``) and ``hand_pose``
     # (variable with ``keypoint_option`` and ``rotation_format``) are absent
     # because their raw width is set per-dataset at construction time. Inference
diff --git a/cosmos_framework/data/generator/action/viewpoint_utils.py b/cosmos_framework/data/generator/action/viewpoint_utils.py
index f083e3df..445443d7 100644
--- a/cosmos_framework/data/generator/action/viewpoint_utils.py
+++ b/cosmos_framework/data/generator/action/viewpoint_utils.py
@@ -15,13 +15,14 @@
 from cosmos_framework.data.imaginaire.webdataset.augmentors.augmentor import Augmentor
 from cosmos_framework.utils import log
 
-Viewpoint = Literal["ego_view", "third_person_view", "wrist_view", "concat_view"]
+Viewpoint = Literal["ego_view", "third_person_view", "wrist_view", "concat_view", "top_down_2d_view"]
 
 DEFAULT_VIEWPOINT_TEMPLATES: dict[str, str] = {
     "ego_view": "This video is captured from a first-person perspective looking at the scene.",
     "third_person_view": "This video is captured from a third-person perspective looking towards the agent from the front.",
     "wrist_view": "This video is captured from a wrist-mounted camera.",
     "concat_view": "This video contains concatenated views from multiple camera perspectives.",
+    "top_down_2d_view": "This video is captured from a top-down view of a flat 2D scene.",
 }
 
 
diff --git a/cosmos_framework/inference/common/distillation_export.py b/cosmos_framework/inference/common/distillation_export.py
index 673d036d..e2029550 100644
--- a/cosmos_framework/inference/common/distillation_export.py
+++ b/cosmos_framework/inference/common/distillation_export.py
@@ -3,9 +3,45 @@
 
 """Portable student-only checkpoint export helpers."""
 
+import warnings
 from collections.abc import Callable
+from pathlib import Path, PurePath
 from typing import Any
 
+_PUBLIC_WAN_VAE_PATHS = {
+    "Wan2.2_VAE.pth": "pretrained/tokenizers/video/wan2pt2/Wan2.2_VAE.pth",
+}
+_PUBLIC_QWEN3_VL_CONFIG_PATHS = {
+    filename: f"cosmos_framework/model/generator/reasoner/qwen3_vl/configs/{filename}"
+    for filename in (
+        "Qwen3-VL-2B-Instruct.json",
+        "Qwen3-VL-4B-Instruct.json",
+        "Qwen3-VL-8B-Instruct.json",
+        "Qwen3-VL-32B-Instruct.json",
+    )
+}
+
+
+def _normalize_public_dependency_path(
+    value: Any,
+    *,
+    field_name: str,
+    public_paths: dict[str, str],
+) -> str:
+    if not isinstance(value, str):
+        raise TypeError(f"Expected {field_name} to be a string.")
+    filename = PurePath(value).name
+    if filename in public_paths:
+        return public_paths[filename]
+    if Path(value).is_absolute():
+        warnings.warn(
+            f"{field_name} contains an unrecognized absolute path and the exported artifact may not be portable: "
+            f"{value}",
+            UserWarning,
+            stacklevel=3,
+        )
+    return value
+
 
 def build_student_checkpoint_metadata(*, use_ema_weights: bool) -> dict[str, str | bool]:
     """Build portable metadata without source checkpoint or credential paths."""
@@ -59,11 +95,29 @@ def sanitize_student_public_model_config(
             tokenizer_config["bucket_name"] = "bucket"
         if "object_store_credential_path_pretrained" in tokenizer_config:
             tokenizer_config["object_store_credential_path_pretrained"] = ""
+        if tokenizer_key == "tokenizer" and "vae_path" in tokenizer_config:
+            tokenizer_config["vae_path"] = _normalize_public_dependency_path(
+                tokenizer_config["vae_path"],
+                field_name="tokenizer.vae_path",
+                public_paths=_PUBLIC_WAN_VAE_PATHS,
+            )
 
     vlm_config = config.get("vlm_config")
     if not isinstance(vlm_config, dict):
         return
 
+    model_instance = vlm_config.get("model_instance")
+    if isinstance(model_instance, dict):
+        model_instance_config = model_instance.get("config")
+        if isinstance(model_instance_config, dict):
+            base_config = model_instance_config.get("base_config")
+            if isinstance(base_config, dict) and "json_file" in base_config:
+                base_config["json_file"] = _normalize_public_dependency_path(
+                    base_config["json_file"],
+                    field_name="vlm_config.model_instance.config.base_config.json_file",
+                    public_paths=_PUBLIC_QWEN3_VL_CONFIG_PATHS,
+                )
+
     pretrained_weights = vlm_config.get("pretrained_weights")
     if isinstance(pretrained_weights, dict):
         pretrained_weights["enabled"] = False
diff --git a/cosmos_framework/inference/common/distillation_export_test.py b/cosmos_framework/inference/common/distillation_export_test.py
index 50007783..1b122bc6 100644
--- a/cosmos_framework/inference/common/distillation_export_test.py
+++ b/cosmos_framework/inference/common/distillation_export_test.py
@@ -1,6 +1,9 @@
 # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: OpenMDW-1.1
 
+import copy
+
+import pytest
 
 from cosmos_framework.inference.common import distillation_export
 from cosmos_framework.inference.common.distillation_export import (
@@ -75,7 +78,7 @@ def test_sanitize_student_public_model_config_removes_internal_loaders() -> None
             "tokenizer": {
                 "bucket_name": "internal-checkpoint-bucket",
                 "object_store_credential_path_pretrained": "/path/to/source.secret",
-                "vae_path": "pretrained/tokenizers/video/wan2pt2/Wan2.2_VAE.pth",
+                "vae_path": "/lustre/training/checkpoints/wan22_vae/Wan2.2_VAE.pth",
             },
             "sound_tokenizer": {
                 "bucket_name": "internal-checkpoint-bucket",
@@ -97,6 +100,16 @@ def test_sanitize_student_public_model_config_removes_internal_loaders() -> None
                     "config_variant": "gcp",
                     "pretrained_model_name": "Qwen/Qwen3-VL-32B-Instruct",
                 },
+                "model_instance": {
+                    "config": {
+                        "base_config": {
+                            "json_file": (
+                                "/checkout/cosmos-framework/cosmos_framework/model/generator/"
+                                "reasoner/qwen3_vl/configs/Qwen3-VL-32B-Instruct.json"
+                            )
+                        }
+                    }
+                },
             },
         }
     }
@@ -134,10 +147,141 @@ def test_sanitize_student_public_model_config_removes_internal_loaders() -> None
                     "config_variant": "hf",
                     "pretrained_model_name": "Qwen/Qwen3-VL-32B-Instruct",
                 },
+                "model_instance": {
+                    "config": {
+                        "base_config": {
+                            "json_file": (
+                                "cosmos_framework/model/generator/reasoner/qwen3_vl/configs/Qwen3-VL-32B-Instruct.json"
+                            )
+                        }
+                    }
+                },
+            },
+        }
+    }
+
+
+@pytest.mark.parametrize("model_size", ("2B", "4B", "8B", "32B"))
+def test_sanitize_student_public_model_config_preserves_qwen_variant(model_size: str) -> None:
+    filename = f"Qwen3-VL-{model_size}-Instruct.json"
+    model_dict = {
+        "config": {
+            "vlm_config": {
+                "model_instance": {
+                    "config": {"base_config": {"json_file": f"/checkout/private/qwen3_vl/configs/{filename}"}}
+                }
+            }
+        }
+    }
+
+    distillation_export.sanitize_student_public_model_config(model_dict)
+
+    assert (
+        model_dict["config"]["vlm_config"]["model_instance"]["config"]["base_config"]["json_file"]
+        == f"cosmos_framework/model/generator/reasoner/qwen3_vl/configs/{filename}"
+    )
+
+
+def test_sanitize_student_public_model_config_is_idempotent() -> None:
+    model_dict = {
+        "config": {
+            "tokenizer": {
+                "bucket_name": "bucket",
+                "vae_path": "pretrained/tokenizers/video/wan2pt2/Wan2.2_VAE.pth",
+            },
+            "vlm_config": {
+                "model_instance": {
+                    "config": {
+                        "base_config": {
+                            "json_file": (
+                                "cosmos_framework/model/generator/reasoner/qwen3_vl/configs/Qwen3-VL-32B-Instruct.json"
+                            )
+                        }
+                    }
+                }
             },
         }
     }
 
+    distillation_export.sanitize_student_public_model_config(model_dict)
+    first_result = copy.deepcopy(model_dict)
+    distillation_export.sanitize_student_public_model_config(model_dict)
+
+    assert model_dict == first_result
+
+
+@pytest.mark.parametrize(
+    ("field", "value"),
+    [
+        ("vae_path", "/custom/tokenizers/CustomVAE.pth"),
+        ("json_file", "/custom/reasoners/CustomReasoner.json"),
+    ],
+)
+def test_sanitize_student_public_model_config_warns_for_unknown_absolute_paths(
+    field: str,
+    value: str,
+) -> None:
+    model_dict = {
+        "config": {
+            "tokenizer": {"vae_path": "relative/custom-vae.pth"},
+            "vlm_config": {
+                "model_instance": {"config": {"base_config": {"json_file": "relative/custom-reasoner.json"}}}
+            },
+        }
+    }
+    if field == "vae_path":
+        model_dict["config"]["tokenizer"]["vae_path"] = value
+    else:
+        model_dict["config"]["vlm_config"]["model_instance"]["config"]["base_config"]["json_file"] = value
+
+    with pytest.warns(UserWarning, match="may not be portable") as warning_records:
+        distillation_export.sanitize_student_public_model_config(model_dict)
+
+    assert warning_records[0].filename == __file__
+    if field == "vae_path":
+        assert model_dict["config"]["tokenizer"]["vae_path"] == value
+    else:
+        assert model_dict["config"]["vlm_config"]["model_instance"]["config"]["base_config"]["json_file"] == value
+
+
+def test_sanitize_student_public_model_config_preserves_unknown_relative_paths() -> None:
+    model_dict = {
+        "config": {
+            "tokenizer": {"vae_path": "relative/custom-vae.pth"},
+            "vlm_config": {
+                "model_instance": {"config": {"base_config": {"json_file": "relative/custom-reasoner.json"}}}
+            },
+        }
+    }
+
+    distillation_export.sanitize_student_public_model_config(model_dict)
+
+    assert model_dict["config"]["tokenizer"]["vae_path"] == "relative/custom-vae.pth"
+    assert (
+        model_dict["config"]["vlm_config"]["model_instance"]["config"]["base_config"]["json_file"]
+        == "relative/custom-reasoner.json"
+    )
+
+
+@pytest.mark.parametrize(
+    ("model_dict", "field_name"),
+    [
+        ({"config": {"tokenizer": {"vae_path": None}}}, "tokenizer.vae_path"),
+        (
+            {"config": {"vlm_config": {"model_instance": {"config": {"base_config": {"json_file": None}}}}}},
+            "vlm_config.model_instance.config.base_config.json_file",
+        ),
+    ],
+)
+def test_sanitize_student_public_model_config_rejects_non_string_dependency_path(
+    model_dict: dict,
+    field_name: str,
+) -> None:
+    with pytest.raises(TypeError) as error:
+        distillation_export.sanitize_student_public_model_config(model_dict)
+
+    assert str(error.value) == f"Expected {field_name} to be a string."
+
 
 def test_resolve_vision_checkpoint_path_prefers_local_override() -> None:
     fallback_called = False
diff --git a/cosmos_framework/model/generator/mot/attention.py b/cosmos_framework/model/generator/mot/attention.py
index 45123e37..fee79382 100644
--- a/cosmos_framework/model/generator/mot/attention.py
+++ b/cosmos_framework/model/generator/mot/attention.py
@@ -95,6 +95,7 @@ def _is_split_info_compatible(attention_mask: object) -> bool:
 _dotproduct_attention_cache = {}
 
 
+from cosmos_framework.model.generator.mot.flex_attention import FlexMetadata, flex_attention_varlen
 from cosmos_framework.data.generator.sequence_packing.natten import (
     generate_natten_metadata,
     generate_temporal_causal_natten_metadata,
@@ -219,11 +220,17 @@ def three_way_attention(
     natten_metadata: dict | None,
     attention_meta: SplitInfo | None = None,
     packed_key_states_normalized: SequencePack | None = None,
+    flex_metadata: FlexMetadata | None = None,
 ):
     """
     Performs three-way attention, with understanding and generations attentions fully decomposed,
     and allows sparsity / multi-dimensional masking in the generation tower.
 
+    The generation-tower self-attention (``full_sa``) is computed by one of three
+    mutually exclusive paths: FlexAttention when ``flex_metadata`` is provided,
+    NATTEN when ``natten_metadata`` is provided, or dense self-attention when
+    neither is set. ``flex_metadata`` and ``natten_metadata`` must not both be set.
+
     When attention_meta is provided with null_action_supertokens=True, zeros V for the first
     num_action_tokens_per_supertoken tokens of each sample's GEN sequence (null action
     supertokens for temporal causal training). The metadata encodes is_causal=(True, False):
@@ -251,11 +258,9 @@ def three_way_attention(
         causal_k_normalized, causal_k_normalized_offsets = causal_k, causal_k_offsets
     causal_v, _ = get_causal_seq(packed_value_states)
     full_q, full_q_offsets = get_full_only_seq(packed_query_states)
-    full_k, full_k_offsets = get_full_only_seq(packed_key_states)
+    full_k, _ = get_full_only_seq(packed_key_states)
     full_v, _ = get_full_only_seq(packed_value_states)
 
-    sample_offsets = packed_query_states["sample_offsets"]
-
     if attention_meta is not None and attention_meta.null_action_supertokens:
         # Zero V for the first num_action_tokens_per_supertoken tokens of each
         # sample's GEN sequence (null action supertokens at t=0).
@@ -286,20 +291,23 @@ def three_way_attention(
     # [1,N_und,heads,head_dim] -> [N_und,heads,head_dim] -> [N_und,heads*head_dim]
     causal_out = causal_res.squeeze(0).flatten(-2, -1)  # type: ignore  # [N_und,heads*head_dim]
 
-    # If there's no metadata, it's a dense layer
-    if natten_metadata is None:
-        full_sa, full_sa_lse = attention(
+    # GEN-tower self-attention (full_sa) via one of three mutually exclusive
+    # paths. flex_metadata and natten_metadata cannot both be set.
+    assert not (flex_metadata is not None and natten_metadata is not None), (
+        "flex_metadata and natten_metadata are mutually exclusive; at most one may be set."
+    )
+    if flex_metadata is not None:
+        # FlexAttention: the multiview supertoken mask encoded in flex_metadata.
+        # Returns the heads-last (out, lse) that merge_attentions expects,
+        # matching the cosmos_framework.model.attention convention.
+        full_sa, full_sa_lse = flex_attention_varlen(
             full_q.unsqueeze(0),  # [1,N_full,heads,head_dim]
             full_k.unsqueeze(0),  # [1,N_full,heads,head_dim]
             full_v.unsqueeze(0),  # [1,N_full,heads,head_dim]
-            cumulative_seqlen_Q=full_q_offsets,
-            cumulative_seqlen_KV=full_k_offsets,
-            max_seqlen_Q=packed_query_states["max_full_len"],
-            max_seqlen_KV=packed_query_states["max_full_len"],
+            flex_metadata,
             return_lse=True,
         )  # full_sa: [1,N_full,heads,head_dim], full_sa_lse: [1,N_full,heads]
-    else:
-        assert natten_metadata is not None
+    elif natten_metadata is not None:
         full_sa, full_sa_lse = multi_dimensional_attention_varlen(
             full_q.unsqueeze(0),  # [1,N_full,heads,head_dim]
             full_k.unsqueeze(0),  # [1,N_full,heads,head_dim]
@@ -307,6 +315,20 @@ def three_way_attention(
             metadata=natten_metadata,
             return_lse=True,
         )  # full_sa: [1,N_full,heads,head_dim], full_sa_lse: [1,N_full,heads]
+    else:
+        # Dense layer: each GEN token attends to every GEN token within its own
+        # packed sample (block-diagonal, bidirectional). Self-attention, so the
+        # KV offsets equal the Q offsets.
+        full_sa, full_sa_lse = attention(
+            full_q.unsqueeze(0),  # [1,N_full,heads,head_dim]
+            full_k.unsqueeze(0),  # [1,N_full,heads,head_dim]
+            full_v.unsqueeze(0),  # [1,N_full,heads,head_dim]
+            cumulative_seqlen_Q=full_q_offsets,
+            cumulative_seqlen_KV=full_q_offsets,
+            max_seqlen_Q=packed_query_states["max_full_len"],
+            max_seqlen_KV=packed_query_states["max_full_len"],
+            return_lse=True,
+        )  # full_sa: [1,N_full,heads,head_dim], full_sa_lse: [1,N_full,heads]
 
     full_ca, full_ca_lse = attention(
         full_q.unsqueeze(0),  # [1,N_full,heads,head_dim]
@@ -547,7 +569,6 @@ def build_packed_sequence(
     num_layers: int,
     token_shapes: list[tuple[int, int, int]] | None = None,
     natten_parameter_list: list | None = None,
-    block_size: int = 128,
     is_image_batch: bool = False,
     cp_world_size: int = 1,
     video_temporal_causal: bool = False,
diff --git a/cosmos_framework/model/generator/mot/flex_attention.py b/cosmos_framework/model/generator/mot/flex_attention.py
new file mode 100644
index 00000000..8707694a
--- /dev/null
+++ b/cosmos_framework/model/generator/mot/flex_attention.py
@@ -0,0 +1,315 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: OpenMDW-1.1
+
+"""FlexAttention implementation of the generation-tower self-attention.
+
+``three_way_attention`` (see ``attention.py``) computes the generator's
+self-attention term (``full_sa``) over the packed GEN tokens and later merges it
+by log-sum-exp with the gen->und cross-attention (``full_ca``). In the dense
+case each GEN token attends to every GEN token *within its own sample*
+(block-diagonal, bidirectional).
+
+This reproduces that dense term with a single FlexAttention call in two explicit
+phases:
+
+1. :func:`build_flex_metadata` assembles the per-token :class:`FlexMetadata`
+   from caller-supplied fields (packed ``sample_id``, plus the multiview
+   supertoken fields ``frame_id`` / ``view_id`` / ``is_noisy`` /
+   ``cond_type_id``). The caller builds this and passes it into
+   :func:`flex_attention_varlen`.
+2. :func:`flex_attention_varlen` derives the ``BlockMask`` from that metadata
+   (via :func:`build_block_mask`) and runs the attention.
+
+When the multiview fields are populated, :func:`build_block_mask` enforces the
+supertoken rules: conditioning tokens attend only to conditioning tokens of the
+same type in the same (frame, view) and never to noisy tokens; noisy tokens
+attend to all noisy tokens within the sample and, additionally, to every
+conditioning token in the same (frame, view). Richer patterns can be added later
+by populating extra metadata fields and extending the ``mask_mod`` instead of
+hand-rolling new varlen bookkeeping.
+
+LSE convention
+--------------
+``flex_attention(..., return_lse=True)`` returns the log-sum-exp of the scaled
+scores in **natural log** with default scale ``1/sqrt(head_dim)`` and layout
+``[B, H, S]`` -- identical to ``cosmos_framework.model.attention(..., return_lse=True)`` once
+transposed to the heads-last ``[B, S, H]`` layout. This makes the output
+directly mergeable with ``full_ca`` via ``merge_attentions``.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass
+
+import torch
+from torch.nn.attention.flex_attention import BlockMask, create_block_mask, flex_attention
+
+# FlexAttention works at block granularity; the GEN sequence length is expected
+# to be pre-padded to a multiple of this by the caller so the compiled kernel
+# sees stable, block-aligned shapes.
+_FLEX_BLOCK_SIZE = 128
+
+# ``dynamic=False`` specialises one kernel per (block-aligned) shape and reuses
+# it across steps. Only the BlockMask *data* changes per step, which does not
+# trigger recompilation.
+_COMPILED_FLEX_ATTENTION = torch.compile(flex_attention, dynamic=False)
+
+# A FlexAttention mask predicate: (b, h, q_idx, kv_idx) -> bool tensor.
+MaskMod = Callable[[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor], torch.Tensor]
+
+
+@dataclass
+class FlexMetadata:
+    """Per-GEN-token metadata that drives the flex block mask.
+
+    Each field is a 1-D ``int64`` tensor of length ``seq_len`` (the block-padded
+    GEN sequence length), aligned with the packed token order. Padding positions
+    use the sentinel ``-1`` so real queries never attend to padding and padded
+    queries attend only to padding (no empty-softmax NaN).
+
+    All fields are required. They describe the multiview supertoken layout and
+    drive the ``mask_mod`` in :func:`build_block_mask`:
+
+    * ``sample_id``: packed sample index per token; yields the block-diagonal
+      same-sample constraint every rule requires.
+    * ``frame_id`` / ``view_id``: per-token frame and view indices.
+    * ``is_noisy``: ``bool`` tensor, ``True`` for noisy (visual) tokens and
+      ``False`` for conditioning tokens.
+    * ``cond_type_id``: conditioning-stream type per token (e.g. 0 = type A,
+      1 = type B); ``-1`` for noisy tokens.
+
+    The mask enforces (all four Q/K quadrants):
+
+    * conditioning Q -> conditioning K: same ``(frame, view, cond_type)``;
+    * conditioning Q -> noisy K: never;
+    * noisy Q -> noisy K: full (bidirectional) within the sample;
+    * noisy Q -> conditioning K: same ``(frame, view)`` (any cond type).
+    """
+
+    seq_len: int
+    sample_id: torch.Tensor
+    frame_id: torch.Tensor
+    view_id: torch.Tensor
+    is_noisy: torch.Tensor
+    cond_type_id: torch.Tensor
+
+
+def _to_flex_layout(x: torch.Tensor) -> torch.Tensor:
+    """Convert ``[1, S, H, D]`` to the FlexAttention layout ``[1, H, S, D]``.
+
+    ``S`` must already be a multiple of ``_FLEX_BLOCK_SIZE`` (the caller pre-pads).
+    """
+    return x.transpose(1, 2).contiguous()  # [1,H,S,D]
+
+
+def _from_flex_layout(x: torch.Tensor) -> torch.Tensor:
+    """Convert a FlexAttention output back to the heads-last layout.
+
+    Inverse of :func:`_to_flex_layout`: swaps the heads and sequence axes, so
+    ``[1, H, S, D] -> [1, S, H, D]`` (attention output) and ``[1, H, S] ->
+    [1, S, H]`` (LSE) both work.
+    """
+    return x.transpose(1, 2).contiguous()  # [1,S,H,...]
+
+
+def _build_gen_sample_ids(
+    full_q_offsets: torch.Tensor,
+    seq_len: int,
+    device: torch.device,
+) -> torch.Tensor:
+    """Per-GEN-token sample id in packed order (``-1`` for padding positions).
+
+    ``full_q_offsets`` is the cumulative per-sample offset array for the packed
+    GEN segment (shape ``[num_samples + 1]``), so ``searchsorted`` maps every
+    position to its sample. Positions at or beyond the last real token
+    (``full_q_offsets[-1]``) are marked ``-1`` so that (a) real queries never
+    attend to padding and (b) padded queries attend only to other padding,
+    avoiding an empty-softmax NaN. Uses only tensor ops (no host sync) so it
+    stays inside a compiled graph.
+    """
+    real_count = full_q_offsets[-1]  # 0-dim tensor; no .item() / host sync.
+    positions = torch.arange(seq_len, device=device)
+    sample_id = torch.searchsorted(full_q_offsets[1:].contiguous(), positions, right=True)
+    return torch.where(positions < real_count, sample_id, -1).to(torch.long)
+
+
+def _multiview_mask_mod_factory(metadata: FlexMetadata) -> MaskMod:
+    """Return the multiview supertoken ``mask_mod``.
+
+    All rules are gated on ``same_sample`` (block-diagonal packing). On top of
+    that, using the per-token frame/view/type metadata (all four Q/K quadrants):
+
+    * conditioning Q attends only to conditioning K of the **same conditioning
+      type** in the **same (frame, view)**;
+    * conditioning Q -> noisy K: never (no rule fires for this quadrant);
+    * noisy Q attends to **all** noisy K (bidirectional, within the sample) and,
+      in addition, to every conditioning K in the **same (frame, view)**
+      regardless of conditioning type.
+
+    Padding positions carry ``sample_id == -1`` (and ``-1`` in the other fields),
+    so real queries never see them and padded queries attend only to padding.
+    """
+    sample_id = metadata.sample_id
+    frame_id = metadata.frame_id
+    view_id = metadata.view_id
+    is_noisy = metadata.is_noisy
+    cond_type_id = metadata.cond_type_id
+
+    def mask_mod(
+        b: torch.Tensor,
+        h: torch.Tensor,
+        q_idx: torch.Tensor,
+        kv_idx: torch.Tensor,
+    ) -> torch.Tensor:
+        same_sample = sample_id[q_idx] == sample_id[kv_idx]
+        same_fv = (frame_id[q_idx] == frame_id[kv_idx]) & (view_id[q_idx] == view_id[kv_idx])
+        same_cond_type = cond_type_id[q_idx] == cond_type_id[kv_idx]
+        q_noisy = is_noisy[q_idx]
+        k_noisy = is_noisy[kv_idx]
+
+        # conditioning Q -> conditioning K: same (frame, view, cond type).
+        cond_to_cond = (~q_noisy) & (~k_noisy) & same_fv & same_cond_type
+        # noisy Q -> noisy K: full within the sample.
+        noisy_to_noisy = q_noisy & k_noisy
+        # noisy Q -> conditioning K: same (frame, view), any cond type.
+        noisy_to_cond = q_noisy & (~k_noisy) & same_fv
+
+        return same_sample & (cond_to_cond | noisy_to_noisy | noisy_to_cond)
+
+    return mask_mod
+
+
+def build_flex_metadata(
+    seq_len: int,
+    *,
+    sample_id: torch.Tensor,
+    frame_id: torch.Tensor,
+    view_id: torch.Tensor,
+    is_noisy: torch.Tensor,
+    cond_type_id: torch.Tensor,
+) -> FlexMetadata:
+    """Assemble per-GEN-token flex metadata from precomputed per-token fields.
+
+    The caller supplies ``sample_id`` (packed sample per token; e.g. via
+    :func:`_build_gen_sample_ids`) together with the multiview supertoken fields
+    ``frame_id`` / ``view_id`` / ``is_noisy`` / ``cond_type_id`` (each
+    ``[seq_len]``, ``-1`` for padding), which drive the multiview ``mask_mod``
+    in :func:`build_block_mask`.
+    """
+    return FlexMetadata(
+        seq_len=seq_len,
+        sample_id=sample_id,
+        frame_id=frame_id,
+        view_id=view_id,
+        is_noisy=is_noisy,
+        cond_type_id=cond_type_id,
+    )
+
+
+def build_block_mask(
+    metadata: FlexMetadata,
+    device: torch.device,
+) -> BlockMask:
+    """Build the GEN-tower :class:`BlockMask` from precomputed flex metadata.
+
+    Uses the multiview supertoken ``mask_mod``; the multiview fields
+    (``frame_id`` / ``view_id`` / ``is_noisy`` / ``cond_type_id``) must be
+    populated on ``metadata``.
+
+    The mask *data* depends on the per-step packing (``metadata.sample_id`` etc.)
+    and is rebuilt every call; because ``metadata.seq_len`` is block-aligned, the
+    compiled ``create_block_mask`` / attention kernels are still reused.
+    """
+    mask_mod = _multiview_mask_mod_factory(metadata)
+    return create_block_mask(
+        mask_mod,
+        B=None,
+        H=None,
+        Q_LEN=metadata.seq_len,
+        KV_LEN=metadata.seq_len,
+        device=device,
+        BLOCK_SIZE=_FLEX_BLOCK_SIZE,
+        _compile=True,
+    )
+
+
+def flex_attention_varlen(
+    full_q: torch.Tensor,
+    full_k: torch.Tensor,
+    full_v: torch.Tensor,
+    metadata: FlexMetadata,
+    return_lse: bool = False,
+) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+    """Dense (per-sample, bidirectional) GEN-tower self-attention via FlexAttention.
+
+    Drop-in replacement for the dense ``full_sa`` branch of
+    ``three_way_attention``: each GEN token attends to every GEN token within its
+    own packed sample.
+
+    The GEN sequence length ``N_full`` must already be a multiple of
+    ``_FLEX_BLOCK_SIZE`` (the caller pre-pads); this is asserted rather than
+    padded here.
+
+    Args:
+        full_q: GEN queries, ``[1, N_full, heads, head_dim]`` (may include
+            trailing pack padding beyond the last real token).
+        full_k: GEN keys, ``[1, N_full, kv_heads, head_dim]``.
+        full_v: GEN values, ``[1, N_full, kv_heads, head_dim]``.
+        metadata: precomputed per-token :class:`FlexMetadata` (see
+            :func:`build_flex_metadata`); its ``seq_len`` must match ``N_full``.
+        return_lse: when ``True`` also return the log-sum-exp, needed to merge
+            this term with ``full_ca`` via ``merge_attentions``.
+
+    Returns:
+        ``full_sa`` of shape ``[1, N_full, heads, head_dim]`` -- the heads-last
+        layout expected by ``merge_attentions``, with the sequence length
+        matching ``full_q`` (pack padding preserved) so the result lines up with
+        ``full_ca``. When ``return_lse`` is ``True``, returns the tuple
+        ``(full_sa, full_sa_lse)`` where ``full_sa_lse`` has shape
+        ``[1, N_full, heads]``.
+    """
+    seq_len = full_q.shape[1]
+    num_q_heads = full_q.shape[2]
+    num_kv_heads = full_k.shape[2]
+    device = full_q.device
+
+    # Padding to the flex block size is assumed to have been done upstream.
+    assert seq_len % _FLEX_BLOCK_SIZE == 0, (
+        f"flex_attention_varlen expects the GEN sequence length to be pre-padded to a multiple "
+        f"of _FLEX_BLOCK_SIZE={_FLEX_BLOCK_SIZE}, got {seq_len}."
+    )
+    assert metadata.seq_len == seq_len, (
+        f"FlexMetadata.seq_len ({metadata.seq_len}) must match the GEN sequence length ({seq_len})."
+    )
+
+    q = _to_flex_layout(full_q)  # [1,num_q_heads,N_full,head_dim]
+    k = _to_flex_layout(full_k)  # [1,num_kv_heads,N_full,head_dim]
+    v = _to_flex_layout(full_v)  # [1,num_kv_heads,N_full,head_dim]
+
+    # Build the block mask from the precomputed per-token flex metadata.
+    block_mask = build_block_mask(metadata, device)
+
+    if return_lse:
+        attn_out, lse = _COMPILED_FLEX_ATTENTION(
+            q,
+            k,
+            v,
+            block_mask=block_mask,
+            enable_gqa=num_q_heads != num_kv_heads,
+            return_lse=True,
+        )  # attn_out: [1,num_q_heads,N_full,head_dim], lse: [1,num_q_heads,N_full]
+        # Convert to the heads-last layout ([1,S,H,D] / [1,S,H]) that
+        # merge_attentions and from_mode_splits expect. Sequence length is unchanged.
+        return _from_flex_layout(attn_out), _from_flex_layout(lse)  # [1,N_full,heads,head_dim], [1,N_full,heads]
+
+    attn_out = _COMPILED_FLEX_ATTENTION(
+        q,
+        k,
+        v,
+        block_mask=block_mask,
+        enable_gqa=num_q_heads != num_kv_heads,
+        return_lse=False,
+    )  # attn_out: [1,num_q_heads,N_full,head_dim]
+    # Convert to the heads-last layout ([1,S,H,D]) that from_mode_splits expects.
+    return _from_flex_layout(attn_out)  # [1,N_full,heads,head_dim]
diff --git a/cosmos_framework/model/generator/mot/flex_attention_test.py b/cosmos_framework/model/generator/mot/flex_attention_test.py
new file mode 100644
index 00000000..d9b5e5d6
--- /dev/null
+++ b/cosmos_framework/model/generator/mot/flex_attention_test.py
@@ -0,0 +1,321 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: OpenMDW-1.1
+
+import math
+from typing import cast
+
+import pytest
+import torch
+
+from cosmos_framework.model.generator.mot.flex_attention import (
+    _FLEX_BLOCK_SIZE,
+    FlexMetadata,
+    _build_gen_sample_ids,
+    _from_flex_layout,
+    _multiview_mask_mod_factory,
+    _to_flex_layout,
+    build_flex_metadata,
+    flex_attention_varlen,
+)
+
+# Conditioning-stream type ids used throughout the tests (matches the layout in
+# assets/mv_attn_mask.py: 0 = cond type A, 1 = cond type B, -1 = noisy/visual).
+_COND_A = 0
+_COND_B = 1
+
+
+def _metadata_from_tokens(tokens: list[dict], seq_len: int | None = None, device: str = "cpu") -> FlexMetadata:
+    """Build a :class:`FlexMetadata` from an explicit list of token descriptors.
+
+    Each token dict has ``s`` (sample), ``t`` (frame), ``v`` (view), ``noisy``
+    (bool) and ``ct`` (cond type id, ``-1`` for noisy). Positions beyond
+    ``len(tokens)`` are padding and get the ``-1`` / ``False`` sentinels.
+    """
+    n = len(tokens)
+    if seq_len is None:
+        seq_len = n
+    pad = seq_len - n
+    assert pad >= 0
+
+    def col(key: str) -> torch.Tensor:
+        return torch.tensor([tok[key] for tok in tokens] + [-1] * pad, dtype=torch.long, device=device)
+
+    is_noisy = torch.tensor([tok["noisy"] for tok in tokens] + [False] * pad, dtype=torch.bool, device=device)
+    return FlexMetadata(
+        seq_len=seq_len,
+        sample_id=col("s"),
+        frame_id=col("t"),
+        view_id=col("v"),
+        is_noisy=is_noisy,
+        cond_type_id=col("ct"),
+    )
+
+
+def _make_multiview_tokens() -> list[dict]:
+    """A small multi-sample multiview layout with cond-A, cond-B and noisy tokens."""
+    tokens: list[dict] = []
+    # Sample 0: 2 frames x 2 views; per (t, v): 1 cond-A, 1 cond-B, 2 noisy.
+    for t in (0, 1):
+        for v in (0, 1):
+            tokens.append(dict(s=0, t=t, v=v, noisy=False, ct=_COND_A))
+            tokens.append(dict(s=0, t=t, v=v, noisy=False, ct=_COND_B))
+            tokens.append(dict(s=0, t=t, v=v, noisy=True, ct=-1))
+            tokens.append(dict(s=0, t=t, v=v, noisy=True, ct=-1))
+    # Sample 1: 1 frame, 1 view; 1 cond-A, 1 cond-B, 3 noisy.
+    tokens.append(dict(s=1, t=0, v=0, noisy=False, ct=_COND_A))
+    tokens.append(dict(s=1, t=0, v=0, noisy=False, ct=_COND_B))
+    for _ in range(3):
+        tokens.append(dict(s=1, t=0, v=0, noisy=True, ct=-1))
+    return tokens
+
+
+def _reference_visibility(tokens: list[dict], seq_len: int) -> torch.Tensor:
+    """Ground-truth ``[seq_len, seq_len]`` bool matrix ``M[q, k] = q attends to k``.
+
+    Encodes exactly the documented multiview rules; padding positions (index >=
+    len(tokens)) share the ``-1`` sample so they only attend to each other.
+    """
+
+    def desc(i: int) -> dict:
+        if i < len(tokens):
+            return tokens[i]
+        return dict(s=-1, t=-1, v=-1, noisy=False, ct=-1)
+
+    m = torch.zeros(seq_len, seq_len, dtype=torch.bool)
+    for q in range(seq_len):
+        dq = desc(q)
+        for k in range(seq_len):
+            dk = desc(k)
+            if dq["s"] != dk["s"]:
+                continue
+            same_fv = dq["t"] == dk["t"] and dq["v"] == dk["v"]
+            if not dq["noisy"] and not dk["noisy"]:
+                ok = same_fv and dq["ct"] == dk["ct"]
+            elif dq["noisy"] and dk["noisy"]:
+                ok = True
+            elif dq["noisy"] and not dk["noisy"]:
+                ok = same_fv
+            else:  # cond query -> noisy key: never
+                ok = False
+            m[q, k] = ok
+    return m
+
+
+def _mask_mod_to_dense(metadata: FlexMetadata) -> torch.Tensor:
+    """Evaluate the metadata's ``mask_mod`` on every (q, k) pair -> ``[S, S]`` bool."""
+    mask_mod = _multiview_mask_mod_factory(metadata)
+    s = metadata.seq_len
+    q_idx = torch.arange(s).view(-1, 1).expand(s, s)
+    kv_idx = torch.arange(s).view(1, -1).expand(s, s)
+    zero = torch.tensor(0)
+    return mask_mod(zero, zero, q_idx, kv_idx)
+
+
+@pytest.mark.L0
+def test_build_gen_sample_ids_marks_padding() -> None:
+    offsets = torch.tensor([0, 3, 7], dtype=torch.long)
+    sample_id = _build_gen_sample_ids(offsets, seq_len=10, device=torch.device("cpu"))
+    expected = torch.tensor([0, 0, 0, 1, 1, 1, 1, -1, -1, -1], dtype=torch.long)
+    assert torch.equal(sample_id, expected)
+
+
+@pytest.mark.L0
+def test_build_gen_sample_ids_no_padding() -> None:
+    offsets = torch.tensor([0, 2, 5], dtype=torch.long)
+    sample_id = _build_gen_sample_ids(offsets, seq_len=5, device=torch.device("cpu"))
+    assert torch.equal(sample_id, torch.tensor([0, 0, 1, 1, 1], dtype=torch.long))
+
+
+@pytest.mark.L0
+def test_build_flex_metadata_populates_all_fields() -> None:
+    seq_len = 4
+    sample_id = torch.tensor([0, 0, 1, 1], dtype=torch.long)
+    frame_id = torch.tensor([0, 1, 0, 1], dtype=torch.long)
+    view_id = torch.tensor([0, 0, 1, 1], dtype=torch.long)
+    is_noisy = torch.tensor([False, True, False, True])
+    cond_type_id = torch.tensor([0, -1, 1, -1], dtype=torch.long)
+
+    meta = build_flex_metadata(
+        seq_len,
+        sample_id=sample_id,
+        frame_id=frame_id,
+        view_id=view_id,
+        is_noisy=is_noisy,
+        cond_type_id=cond_type_id,
+    )
+
+    assert meta.seq_len == seq_len
+    assert torch.equal(meta.sample_id, sample_id)
+    assert torch.equal(meta.frame_id, frame_id)
+    assert torch.equal(meta.view_id, view_id)
+    assert torch.equal(meta.is_noisy, is_noisy)
+    assert torch.equal(meta.cond_type_id, cond_type_id)
+
+
+@pytest.mark.L0
+def test_flex_layout_roundtrip() -> None:
+    x4 = torch.randn(1, 6, 4, 8)  # [1,S,H,D]
+    assert _to_flex_layout(x4).shape == (1, 4, 6, 8)  # [1,H,S,D]
+    assert torch.equal(_from_flex_layout(_to_flex_layout(x4)), x4)
+
+    x3 = torch.randn(1, 4, 6)  # [1,H,S] (LSE layout)
+    assert _from_flex_layout(x3).shape == (1, 6, 4)  # [1,S,H]
+
+
+@pytest.mark.L0
+def test_multiview_mask_mod_matches_reference() -> None:
+    tokens = _make_multiview_tokens()
+    seq_len = len(tokens)
+    metadata = _metadata_from_tokens(tokens, seq_len=seq_len)
+
+    got = _mask_mod_to_dense(metadata)
+    expected = _reference_visibility(tokens, seq_len)
+    assert torch.equal(got, expected)
+
+
+@pytest.mark.L0
+def test_multiview_mask_mod_specific_rules() -> None:
+    # Two views, two frames, one sample. Layout index -> token:
+    #  0: condA (t0,v0)   1: condB (t0,v0)   2: noisy (t0,v0)
+    #  3: condA (t0,v1)   4: noisy (t0,v1)
+    #  5: condA (t1,v0)   6: noisy (t1,v0)
+    tokens = [
+        dict(s=0, t=0, v=0, noisy=False, ct=_COND_A),
+        dict(s=0, t=0, v=0, noisy=False, ct=_COND_B),
+        dict(s=0, t=0, v=0, noisy=True, ct=-1),
+        dict(s=0, t=0, v=1, noisy=False, ct=_COND_A),
+        dict(s=0, t=0, v=1, noisy=True, ct=-1),
+        dict(s=0, t=1, v=0, noisy=False, ct=_COND_A),
+        dict(s=0, t=1, v=0, noisy=True, ct=-1),
+    ]
+    m = _mask_mod_to_dense(_metadata_from_tokens(tokens))
+
+    # cond-A (t0,v0) attends only to itself among these (same type/frame/view).
+    assert m[0, 0]
+    assert not m[0, 1]  # different cond type
+    assert not m[0, 2]  # cond -> noisy never
+    assert not m[0, 3]  # different view
+    assert not m[0, 5]  # different frame
+
+    # noisy (t0,v0) attends to all noisy in the sample + cond of same (t,v).
+    assert m[2, 2] and m[2, 4] and m[2, 6]  # all noisy tokens
+    assert m[2, 0] and m[2, 1]  # cond A & B in same (t0,v0)
+    assert not m[2, 3]  # cond in different view
+    assert not m[2, 5]  # cond in different frame
+
+
+@pytest.mark.L0
+def test_multiview_mask_mod_block_diagonal_across_samples() -> None:
+    tokens = _make_multiview_tokens()
+    metadata = _metadata_from_tokens(tokens)
+    m = _mask_mod_to_dense(metadata)
+
+    sample_id = metadata.sample_id
+    cross = sample_id.view(-1, 1) != sample_id.view(1, -1)
+    assert not m[cross].any(), "attention must never cross sample boundaries"
+
+
+@pytest.mark.L0
+def test_multiview_mask_mod_padding_isolated() -> None:
+    tokens = _make_multiview_tokens()
+    n = len(tokens)
+    seq_len = n + 5  # add padding positions
+    metadata = _metadata_from_tokens(tokens, seq_len=seq_len)
+    m = _mask_mod_to_dense(metadata)
+
+    # Real queries never attend to padding keys, and vice versa.
+    assert not m[:n, n:].any()
+    assert not m[n:, :n].any()
+    # Padded queries attend only to padding (non-empty softmax -> no NaN).
+    assert m[n:, n:].all()
+
+
+@pytest.mark.L0
+def test_flex_attention_varlen_rejects_unaligned_seq_len() -> None:
+    seq_len = _FLEX_BLOCK_SIZE + 1  # not a multiple of the block size
+    q = torch.randn(1, seq_len, 2, 8)
+    meta = _metadata_from_tokens([dict(s=0, t=0, v=0, noisy=True, ct=-1)], seq_len=seq_len)
+    with pytest.raises(AssertionError, match="pre-padded to a multiple"):
+        flex_attention_varlen(q, q, q, meta)
+
+
+@pytest.mark.L0
+def test_flex_attention_varlen_rejects_seq_len_mismatch() -> None:
+    seq_len = _FLEX_BLOCK_SIZE
+    q = torch.randn(1, seq_len, 2, 8)
+    meta = _metadata_from_tokens([dict(s=0, t=0, v=0, noisy=True, ct=-1)], seq_len=seq_len - 1)
+    with pytest.raises(AssertionError, match="must match the GEN sequence length"):
+        flex_attention_varlen(q, q, q, meta)
+
+
+def _reference_attention(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    mask: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Dense masked attention reference. q/k/v: ``[1,S,H,D]`` / ``[1,S,Hkv,D]``.
+
+    Returns ``(out [1,S,H,D], lse [1,S,H])`` with the natural-log LSE and the
+    default ``1/sqrt(D)`` scale, matching the FlexAttention convention.
+    """
+    qh = q[0]  # [S,H,D]
+    kh = k[0]  # [S,Hkv,D]
+    vh = v[0]
+    num_q_heads = qh.shape[1]
+    num_kv_heads = kh.shape[1]
+    if num_q_heads != num_kv_heads:
+        factor = num_q_heads // num_kv_heads
+        kh = kh.repeat_interleave(factor, dim=1)
+        vh = vh.repeat_interleave(factor, dim=1)
+    scale = 1.0 / math.sqrt(qh.shape[-1])
+    scores = torch.einsum("shd,thd->hst", qh, kh) * scale  # [H,S,S]
+    neg_inf = torch.finfo(scores.dtype).min
+    scores = scores.masked_fill(~mask.view(1, *mask.shape), neg_inf)
+    weights = torch.softmax(scores, dim=-1)
+    out = torch.einsum("hst,thd->shd", weights, vh)  # [S,H,D]
+    lse = torch.logsumexp(scores, dim=-1)  # [H,S]
+    return out.unsqueeze(0), lse.transpose(0, 1).unsqueeze(0)
+
+
+@pytest.mark.L0
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="FlexAttention kernels require a GPU.")
+@pytest.mark.parametrize("num_kv_heads", [4, 1])
+@pytest.mark.parametrize("return_lse", [True, False])
+def test_flex_attention_varlen_matches_reference(num_kv_heads: int, return_lse: bool) -> None:
+    torch.manual_seed(0)
+    torch.compiler.reset()
+    device = "cuda"
+    dtype = torch.float32
+    num_q_heads = 4
+    head_dim = 32
+    seq_len = _FLEX_BLOCK_SIZE  # single block
+
+    tokens = _make_multiview_tokens()
+    n_real = len(tokens)
+    metadata = _metadata_from_tokens(tokens, seq_len=seq_len, device=device)
+
+    q = torch.randn(1, seq_len, num_q_heads, head_dim, device=device, dtype=dtype)
+    k = torch.randn(1, seq_len, num_kv_heads, head_dim, device=device, dtype=dtype)
+    v = torch.randn(1, seq_len, num_kv_heads, head_dim, device=device, dtype=dtype)
+
+    result = flex_attention_varlen(q, k, v, metadata, return_lse=return_lse)
+    lse: torch.Tensor | None = None
+    if return_lse:
+        assert isinstance(result, tuple)
+        out = cast(torch.Tensor, result[0])
+        lse = cast(torch.Tensor, result[1])
+        assert lse.shape == (1, seq_len, num_q_heads)
+    else:
+        assert isinstance(result, torch.Tensor)
+        out = cast(torch.Tensor, result)
+    assert out.shape == (1, seq_len, num_q_heads, head_dim)
+
+    mask = _mask_mod_to_dense(metadata).to(device)
+    ref_out, ref_lse = _reference_attention(q, k, v, mask)
+
+    # Only compare real (non-padding) token rows.
+    real = slice(0, n_real)
+    torch.testing.assert_close(out[:, real], ref_out[:, real], atol=2e-2, rtol=2e-2)
+    if lse is not None:
+        torch.testing.assert_close(lse[:, real], ref_lse[:, real], atol=2e-2, rtol=2e-2)
diff --git a/cosmos_framework/model/generator/tokenizers/dc_ae/dc_ae_4x32x32.py b/cosmos_framework/model/generator/tokenizers/dc_ae/dc_ae_4x32x32.py
index 9b93c259..e146a26f 100644
--- a/cosmos_framework/model/generator/tokenizers/dc_ae/dc_ae_4x32x32.py
+++ b/cosmos_framework/model/generator/tokenizers/dc_ae/dc_ae_4x32x32.py
@@ -2,6 +2,7 @@
 # SPDX-License-Identifier: OpenMDW-1.1
 
 from collections.abc import Sequence
+from typing import Optional
 
 import torch
 
@@ -34,9 +35,11 @@ def __init__(
         device: str = "cuda",
         compilable: bool = True,
         causal: bool = True,
+        encode_exact_durations: Optional[list[int]] = None,
     ):
         self._causal = causal
         assert self._causal, "DCAE4x32x32Interface is a causal tokenizer; causal must be True."
+        assert encode_exact_durations is None, "DCAE4x32x32Interface does not support encode_exact_durations."
         vae_path_full = f"s3://{bucket_name}/{vae_path}"
         self._spatial_compression_factor = spatial_compression_factor
         self._temporal_compression_factor = temporal_compression_factor
diff --git a/cosmos_framework/utils/device.py b/cosmos_framework/utils/device.py
index 7bc2f884..f9e5d53b 100644
--- a/cosmos_framework/utils/device.py
+++ b/cosmos_framework/utils/device.py
@@ -10,15 +10,17 @@
 from loguru import logger as logging
 
 
-def get_gpu_architecture():
+def get_gpu_architecture() -> str:
     """
     Retrieves the GPU architecture of the available GPUs.
 
     Returns:
         str: The GPU architecture, which can be "H100", "A100", or "Other".
     """
+    nvml_initialized = False
     try:
         pynvml.nvmlInit()
+        nvml_initialized = True
         device_count = pynvml.nvmlDeviceGetCount()
         for i in range(device_count):
             handle = pynvml.nvmlDeviceGetHandleByIndex(i)
@@ -39,7 +41,8 @@ def get_gpu_architecture():
     except pynvml.NVMLError as error:
         print(f"Failed to get GPU info: {error}")
     finally:
-        pynvml.nvmlShutdown()
+        if nvml_initialized:
+            pynvml.nvmlShutdown()
 
     # return "Other" incase of non hopper/ampere or error
     return "Other"
diff --git a/cosmos_framework/utils/distributed.py b/cosmos_framework/utils/distributed.py
index 10a148d0..87cdd39e 100644
--- a/cosmos_framework/utils/distributed.py
+++ b/cosmos_framework/utils/distributed.py
@@ -481,6 +481,7 @@ def broadcast_object_list(object_list, *args, **kwargs):
 class _TensorBroadcastMetadata:
     """Small placeholder used while broadcasting a nested object containing tensors."""
 
+    tensor_index: int
     shape: tuple[int, ...]
     dtype: torch.dtype
 
@@ -488,16 +489,17 @@ class _TensorBroadcastMetadata:
 def _extract_tensor_leaves(value: Any, min_tensor_bytes: int) -> tuple[Any, list[torch.Tensor]]:
     """Replace sufficiently large tensor leaves while preserving the surrounding containers."""
     tensor_leaves: list[torch.Tensor] = []
-    seen_direct_ids: set[int] = set()
+    tensor_metadata_by_id: dict[int, _TensorBroadcastMetadata] = {}
+    seen_container_ids: set[int] = set()
 
-    def _require_unique(current_value: Any) -> None:
+    def _require_unique_container(current_value: Any) -> None:
         object_id = id(current_value)
-        if object_id in seen_direct_ids:
+        if object_id in seen_container_ids:
             raise ValueError(
-                "Optimized object broadcast requires an acyclic tree without shared container or tensor references; "
+                "Optimized object broadcast requires an acyclic tree without shared container references; "
                 f"encountered {type(current_value).__name__} more than once."
             )
-        seen_direct_ids.add(object_id)
+        seen_container_ids.add(object_id)
 
     def _extract(current_value: Any) -> Any:
         if isinstance(current_value, torch.Tensor):
@@ -512,16 +514,25 @@ def _extract(current_value: Any) -> Any:
                     "Optimized object broadcast only supports plain torch.Tensor leaves, "
                     f"got {type(current_value).__name__}."
                 )
-            _require_unique(current_value)
+            object_id = id(current_value)
+            existing_metadata = tensor_metadata_by_id.get(object_id)
+            if existing_metadata is not None:
+                return existing_metadata
+            metadata = _TensorBroadcastMetadata(
+                tensor_index=len(tensor_leaves),
+                shape=tuple(current_value.shape),
+                dtype=current_value.dtype,
+            )
+            tensor_metadata_by_id[object_id] = metadata
             tensor_leaves.append(current_value)
-            return _TensorBroadcastMetadata(shape=tuple(current_value.shape), dtype=current_value.dtype)
+            return metadata
         if isinstance(current_value, dict):
             if type(current_value) is not dict:
                 raise TypeError(
                     "Optimized object broadcast only supports plain dict containers, "
                     f"got {type(current_value).__name__}."
                 )
-            _require_unique(current_value)
+            _require_unique_container(current_value)
             return {key: _extract(item) for key, item in current_value.items()}
         if isinstance(current_value, list):
             if type(current_value) is not list:
@@ -529,7 +540,7 @@ def _extract(current_value: Any) -> Any:
                     "Optimized object broadcast only supports plain list containers, "
                     f"got {type(current_value).__name__}."
                 )
-            _require_unique(current_value)
+            _require_unique_container(current_value)
             return [_extract(item) for item in current_value]
         if isinstance(current_value, tuple):
             is_namedtuple = hasattr(current_value, "_fields")
@@ -538,7 +549,7 @@ def _extract(current_value: Any) -> Any:
                     "Optimized object broadcast only supports plain tuple or namedtuple containers, "
                     f"got {type(current_value).__name__}."
                 )
-            _require_unique(current_value)
+            _require_unique_container(current_value)
             extracted_items = tuple(_extract(item) for item in current_value)
             # Namedtuples are tuple subclasses whose constructors expect each field as a positional argument.
             if is_namedtuple:
@@ -570,11 +581,12 @@ def broadcast_object_list_optimized(
     collective directly.
 
     Direct tensor extraction requires an acyclic tree of plain ``dict``, ``list``,
-    and ``tuple`` containers; namedtuples are also supported. The source raises
-    before entering the collective when it encounters a container subclass,
-    a tensor subclass selected for direct broadcast, a shared container/tensor
-    reference that would be rebuilt, or a cycle. Use the default threshold for
-    arbitrary picklable object graphs.
+    and ``tuple`` containers; namedtuples are also supported. Repeated references
+    to the same tensor are broadcast once and preserved in the rebuilt object
+    graph. The source raises before entering the collective when it encounters a
+    container subclass, a tensor subclass selected for direct broadcast, a shared
+    container reference, or a cycle. Use the default threshold for arbitrary
+    picklable object graphs.
 
     Args:
         object_list: List of objects to broadcast. Updated in place on every participating rank.
@@ -621,12 +633,23 @@ def broadcast_object_list_optimized(
         collective_device = device or torch.device("cuda", torch.cuda.current_device())
     else:
         collective_device = torch.device("cpu")
-    tensor_index = 0
+    rebuilt_tensors: list[torch.Tensor] = []
 
     def _rebuild(current_value: Any) -> Any:
-        nonlocal tensor_index
         if isinstance(current_value, _TensorBroadcastMetadata):
+            tensor_index = current_value.tensor_index
+            if 0 <= tensor_index < len(rebuilt_tensors):
+                return rebuilt_tensors[tensor_index]
+            if tensor_index != len(rebuilt_tensors):
+                raise RuntimeError(
+                    "Broadcast object-list skeleton contains an out-of-order tensor index: "
+                    f"expected {len(rebuilt_tensors)}, got {tensor_index}."
+                )
             if is_source:
+                if tensor_index >= len(tensor_leaves):
+                    raise RuntimeError(
+                        f"Broadcast object-list skeleton references missing source tensor index {tensor_index}."
+                    )
                 source_tensor = tensor_leaves[tensor_index]  # [*shape]
                 if (
                     source_tensor.device.type == "cuda"
@@ -647,8 +670,8 @@ def _rebuild(current_value: Any) -> Any:
                     dtype=current_value.dtype,
                     device=collective_device,
                 )
-            tensor_index += 1
             dist.broadcast(tensor, src=src, group=group, group_src=group_src)
+            rebuilt_tensors.append(tensor)
             return tensor
         if isinstance(current_value, dict):
             return {key: _rebuild(item) for key, item in current_value.items()}
@@ -663,8 +686,10 @@ def _rebuild(current_value: Any) -> Any:
         return current_value
 
     rebuilt_object_list = _rebuild(skeleton)
-    if is_source and tensor_index != len(tensor_leaves):
-        raise RuntimeError(f"Broadcast rebuilt {tensor_index} tensors but source provided {len(tensor_leaves)}.")
+    if is_source and len(rebuilt_tensors) != len(tensor_leaves):
+        raise RuntimeError(
+            f"Broadcast rebuilt {len(rebuilt_tensors)} tensors but source provided {len(tensor_leaves)}."
+        )
     if not isinstance(rebuilt_object_list, list):
         raise RuntimeError(f"Broadcast object-list skeleton must be a list, got {type(rebuilt_object_list).__name__}.")
     object_list[:] = rebuilt_object_list
diff --git a/cosmos_framework/utils/distributed_test.py b/cosmos_framework/utils/distributed_test.py
index 701a7ccf..56cd7575 100644
--- a/cosmos_framework/utils/distributed_test.py
+++ b/cosmos_framework/utils/distributed_test.py
@@ -96,7 +96,14 @@ def _broadcast_tensor(
 
     video = torch.arange(24, dtype=torch.float32).reshape(2, 3, 4)  # [2,3,4]
     frame_count = torch.tensor(197, dtype=torch.int64)  # []
-    source_objects = [{"video": [[video]], "metadata": (frame_count, "mads")}, False]
+    source_objects = [
+        {
+            "video": [[video]],
+            "video_alias": video,
+            "metadata": (frame_count, "mads"),
+        },
+        False,
+    ]
     result = distributed.broadcast_object_list_optimized(
         source_objects,
         src=0,
@@ -122,8 +129,10 @@ def _broadcast_tensor(
     assert receiver_objects[1] is False
     assert receiver_objects[0]["metadata"][1] == "mads"
     torch.testing.assert_close(receiver_objects[0]["video"][0][0], video)
+    assert receiver_objects[0]["video_alias"] is receiver_objects[0]["video"][0][0]
     torch.testing.assert_close(receiver_objects[0]["metadata"][0], frame_count)
     torch.testing.assert_close(source_objects[0]["video"][0][0], video)
+    assert source_objects[0]["video_alias"] is source_objects[0]["video"][0][0]
     assert not tensor_payloads
 
 
@@ -288,7 +297,7 @@ def _get_rank() -> int:
     monkeypatch.setattr(distributed.dist, "get_rank", _get_rank)
     monkeypatch.setattr(distributed.dist, "broadcast_object_list", object_broadcast)
 
-    with pytest.raises(ValueError, match="without shared container or tensor references"):
+    with pytest.raises(ValueError, match="without shared container references"):
         distributed.broadcast_object_list_optimized(
             [shared, shared],
             src=0,
diff --git a/cosmos_framework/utils/timer.py b/cosmos_framework/utils/timer.py
index 2c88b9b8..2e856e12 100644
--- a/cosmos_framework/utils/timer.py
+++ b/cosmos_framework/utils/timer.py
@@ -150,10 +150,11 @@ def __call__(self, func: Callable) -> Callable:
         @wraps(func)
         def wrapper(*args, **kwargs):  # noqa: ANN202
             self.start()
-            result = func(*args, **kwargs)
-            self.end()
-            self.report()
-            return result
+            try:
+                return func(*args, **kwargs)
+            finally:
+                self.end()
+                self.report()
 
         return wrapper  # type: ignore