Release 2026-09-14 (i4 3a3e1e11)

.file_mapping.json

@@ -1,7 +1,7 @@
 {
-  "_source_commit": "116a1414e2fb10bc7416ca860ad67f517ef9371f-dirty",
-  "_dest_commit": "ffa9c6b60a6b04b2fae337577bc6cbd8a93c39f5",
-  "_generated_at": "2026-09-11T05:52:11Z",
+  "_source_commit": "3a3e1e118760abbc484d59807b797157c394635e-dirty",
+  "_dest_commit": "2b6c9a7061ae78dc83e29a4910ec5f8c9fe4b6ce",
+  "_generated_at": "2026-09-14T05:52:22Z",
   "files": {
     "imaginaire/__init__.py": "cosmos_framework/__init__.py",
     "imaginaire/attention/__init__.py": "cosmos_framework/model/attention/__init__.py",

cosmos_framework/configs/base/defaults/tokenizer.py

@@ -1,7 +1,6 @@
 # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: OpenMDW-1.1
 
-import torch
 from hydra.core.config_store import ConfigStore
 
 from cosmos_framework.utils.lazy_config import PLACEHOLDER, LazyDict
@@ -177,10 +176,6 @@ def register_tokenizer() -> None:
     # Wan2pt1 and Wan2pt2 tokenizers
     cs.store(group="tokenizer", package="model.config.tokenizer", name="wan2pt1_tokenizer", node=Wan2pt1VAEConfig)
     cs.store(group="tokenizer", package="model.config.tokenizer", name="wan2pt2_tokenizer", node=Wan2pt2VAEConfig)
-    # LiDAR VAEs are deliberately absent from this group: a range clip is its own modality
-    # with its own projections into the sequence, so it is registered under
-    # ``model.config.lidar_tokenizer`` by ``register_lidar_tokenizer`` below. Installing one
-    # here would displace the camera tokenizer and route rangemaps through the vision heads.
     # UniAE tokenizer
     cs.store(
         group="tokenizer",
@@ -199,34 +194,6 @@ def register_tokenizer() -> None:
     )
 
 
-def register_lidar_tokenizer() -> None:
-    """Register LiDAR tokenizers under ``model.config.lidar_tokenizer``.
-
-    A joint camera + LiDAR recipe holds two vision tokenizers at once, so the LiDAR one needs
-    its own package: registering it under ``model.config.tokenizer`` would displace the camera
-    tokenizer. With this group in place, the dataloader's per-sensor token accounting can
-    interpolate ``${model.config.lidar_tokenizer.temporal_compression_factor}`` the same way
-    it already reads the camera factors off ``model.config.tokenizer``.
-    """
-    cs = ConfigStore.instance()
-    cs.store(
-        group="lidar_tokenizer",
-        package="model.config.lidar_tokenizer",
-        name="lidar_tokenizer_v0",
-        node=LidarTokenizerV0Config,
-    )
-    cs.store(
-        group="lidar_tokenizer",
-        package="model.config.lidar_tokenizer",
-        name="lidar_tokenizer_v1_r105_b1800_symmetric",
-        node=LidarTokenizerV1R105B1800SymmetricConfig,
-    )
-    cs.store(
-        group="lidar_tokenizer",
-        package="model.config.lidar_tokenizer",
-        name="lidar_tokenizer_v1p2_r105_b1800",
-        node=LidarTokenizerV1P2R105B1800Config,
-    )
 
 
 def register_sound_tokenizer() -> None:

cosmos_framework/inference/common/distillation_export.py

@@ -86,6 +86,57 @@ def build_student_checkpoint_metadata(*, use_ema_weights: bool) -> dict[str, str
     }
 
 
+def _migrate_legacy_transfer_replay_config(config: dict[str, Any], *, base_config_field_names: set[str]) -> None:
+    """Preserve pre-replay-policy Transfer connectivity when projecting a student."""
+    legacy_key = "transfer_control_attention_mode"
+    if legacy_key not in config:
+        return
+
+    required_fields = {"teacher_forcing_replay_policy", "teacher_forcing_kv_implementation"}
+    if not required_fields <= base_config_field_names:
+        raise ValueError("Legacy Transfer attention requires a causal base config with teacher-forcing replay support.")
+
+    legacy_modes = {
+        "global_control": ("global", False),
+        "causal_control": ("causal", False),
+        "current_only_control": ("current", False),
+        "causal_control_with_rgb_history": ("causal", True),
+        "current_only_control_with_rgb_history": ("current", True),
+    }
+    legacy_mode = config[legacy_key]
+    if not isinstance(legacy_mode, str) or legacy_mode not in legacy_modes:
+        raise ValueError(f"Unsupported legacy {legacy_key}: {legacy_mode!r}.")
+    control_visibility, controls_read_rgb = legacy_modes[legacy_mode]
+    expected_policy = {
+        "control_visibility": control_visibility,
+        "controls_read_strict_past_clean_rgb": controls_read_rgb,
+        "clean_pass_causality": "frame",
+        "multiview_attention_scope": "all_views",
+        "decomposed_temporal_window_seconds": None,
+    }
+    policy = config.get("teacher_forcing_replay_policy", {})
+    if not isinstance(policy, dict):
+        raise TypeError("Expected teacher_forcing_replay_policy to be a dictionary during legacy Transfer migration.")
+    for key, expected in expected_policy.items():
+        if key in policy and policy[key] != expected:
+            raise ValueError(
+                f"Legacy {legacy_key}={legacy_mode!r} conflicts with teacher_forcing_replay_policy.{key}={policy[key]!r}; "
+                f"expected {expected!r}."
+            )
+
+    # The legacy Transfer implementation used three-way single-view attention.
+    # An explicit different implementation must not silently replace that kernel.
+    implementation = config.get("teacher_forcing_kv_implementation", "singleview_threeway_kv")
+    if implementation != "singleview_threeway_kv":
+        raise ValueError(
+            f"Legacy {legacy_key} conflicts with teacher_forcing_kv_implementation={implementation!r}; "
+            "expected 'singleview_threeway_kv'."
+        )
+    config["teacher_forcing_replay_policy"] = {**policy, **expected_policy}
+    config["teacher_forcing_kv_implementation"] = implementation
+    del config[legacy_key]
+
+
 def sanitize_student_model_config(
     model_dict: dict[str, Any],
     *,
@@ -98,6 +149,10 @@ def sanitize_student_model_config(
     if not isinstance(config, dict):
         raise TypeError("Expected model config to be a dictionary.")
 
+    # Migrate before filtering out training-only fields: dropping the legacy
+    # selector first would silently restore global controls without RGB history.
+    _migrate_legacy_transfer_replay_config(config, base_config_field_names=base_config_field_names)
+
     model_dict["_target_"] = base_model_target
     config["_type"] = base_config_type
 

cosmos_framework/inference/common/distillation_export_test.py

@@ -73,6 +73,144 @@ def test_build_student_checkpoint_metadata_omits_source_paths() -> None:
     }
 
 
+def _sanitize_causal_student(model_dict: dict) -> None:
+    sanitize_student_model_config(
+        model_dict,
+        base_model_target="omni_mot_causal_model",
+        base_config_type="omni_mot_causal_model_config",
+        base_config_field_names={
+            "video_temporal_causal",
+            "teacher_forcing_replay_policy",
+            "teacher_forcing_kv_implementation",
+            "teacher_forcing_frames_per_chunk",
+            "kv_cache_inference_size",
+            "attention_sink_size",
+        },
+    )
+
+
+@pytest.mark.parametrize(
+    ("legacy_mode", "control_visibility", "controls_read_rgb"),
+    [
+        ("global_control", "global", False),
+        ("causal_control", "causal", False),
+        ("current_only_control", "current", False),
+        ("causal_control_with_rgb_history", "causal", True),
+        ("current_only_control_with_rgb_history", "current", True),
+    ],
+)
+def test_sanitize_student_migrates_legacy_transfer_connectivity(
+    legacy_mode: str, control_visibility: str, controls_read_rgb: bool
+) -> None:
+    model_dict = {
+        "config": {
+            "video_temporal_causal": True,
+            "transfer_control_attention_mode": legacy_mode,
+            "teacher_forcing_frames_per_chunk": 1,
+            "kv_cache_inference_size": 51,
+            "attention_sink_size": 1,
+        }
+    }
+
+    _sanitize_causal_student(model_dict)
+
+    assert model_dict["config"] == {
+        "_type": "omni_mot_causal_model_config",
+        "video_temporal_causal": True,
+        "teacher_forcing_frames_per_chunk": 1,
+        "kv_cache_inference_size": 51,
+        "attention_sink_size": 1,
+        "teacher_forcing_kv_implementation": "singleview_threeway_kv",
+        "teacher_forcing_replay_policy": {
+            "control_visibility": control_visibility,
+            "controls_read_strict_past_clean_rgb": controls_read_rgb,
+            "clean_pass_causality": "frame",
+            "multiview_attention_scope": "all_views",
+            "decomposed_temporal_window_seconds": None,
+        },
+    }
+    migrated = copy.deepcopy(model_dict)
+    _sanitize_causal_student(model_dict)
+    assert model_dict == migrated
+
+
+def test_sanitize_student_merges_compatible_legacy_and_current_replay_settings() -> None:
+    model_dict = {
+        "config": {
+            "transfer_control_attention_mode": "causal_control_with_rgb_history",
+            "teacher_forcing_kv_implementation": "singleview_threeway_kv",
+            "teacher_forcing_replay_policy": {
+                "_type": "teacher_forcing_replay_policy_config",
+                "control_visibility": "causal",
+            },
+        }
+    }
+
+    _sanitize_causal_student(model_dict)
+
+    policy = model_dict["config"]["teacher_forcing_replay_policy"]
+    assert policy["_type"] == "teacher_forcing_replay_policy_config"
+    assert policy["control_visibility"] == "causal"
+    assert policy["controls_read_strict_past_clean_rgb"] is True
+
+
+@pytest.mark.parametrize(
+    "conflicting_policy",
+    [
+        {"control_visibility": "global"},
+        {"controls_read_strict_past_clean_rgb": False},
+        {"clean_pass_causality": "chunk"},
+        {"multiview_attention_scope": "same_view"},
+        {"decomposed_temporal_window_seconds": 0.1},
+    ],
+)
+def test_sanitize_student_rejects_conflicting_legacy_and_current_replay_settings(conflicting_policy: dict) -> None:
+    model_dict = {
+        "config": {
+            "transfer_control_attention_mode": "causal_control_with_rgb_history",
+            "teacher_forcing_replay_policy": conflicting_policy,
+        }
+    }
+    original = copy.deepcopy(model_dict)
+
+    with pytest.raises(ValueError, match="conflicts with teacher_forcing_replay_policy"):
+        _sanitize_causal_student(model_dict)
+
+    assert model_dict == original
+
+
+def test_sanitize_student_rejects_conflicting_legacy_kv_implementation() -> None:
+    model_dict = {
+        "config": {
+            "transfer_control_attention_mode": "causal_control_with_rgb_history",
+            "teacher_forcing_kv_implementation": "multiview_flex_kv",
+        }
+    }
+
+    with pytest.raises(ValueError, match="conflicts with teacher_forcing_kv_implementation"):
+        _sanitize_causal_student(model_dict)
+
+
+@pytest.mark.parametrize("legacy_mode", [None, "future_control", []])
+def test_sanitize_student_rejects_unknown_legacy_transfer_mode(legacy_mode: object) -> None:
+    model_dict = {"config": {"transfer_control_attention_mode": legacy_mode}}
+
+    with pytest.raises(ValueError, match="Unsupported legacy transfer_control_attention_mode"):
+        _sanitize_causal_student(model_dict)
+
+
+def test_sanitize_student_rejects_legacy_transfer_without_causal_base_support() -> None:
+    model_dict = {"config": {"transfer_control_attention_mode": "causal_control_with_rgb_history"}}
+
+    with pytest.raises(ValueError, match="requires a causal base config"):
+        sanitize_student_model_config(
+            model_dict,
+            base_model_target="omni_mot_model",
+            base_config_type="omni_mot_model_config",
+            base_config_field_names={"video_temporal_causal"},
+        )
+
+
 def test_sanitize_student_public_model_config_removes_internal_loaders() -> None:
     model_dict = {
         "config": {

cosmos_framework/model/generator/mot/flex_attention.py

@@ -1423,12 +1423,12 @@ def flex_attention(
         ``(out, lse)`` where ``lse`` has shape ``[1, N_full, heads]``.
 
     Raises:
+        RuntimeError: if ``return_lse`` is requested with a torch version that does not
+            provide ``torch.nn.attention.flex_attention.AuxRequest``.
         ValueError: if either length is not a multiple of the corresponding block size of
             the mask, if k and v disagree on length, if ``block_mask`` was built for
             different lengths, or if it was built at a block size other than ``backend``'s.
     """
-    from torch.nn.attention.flex_attention import AuxRequest
-
     q_seq_len = full_q.shape[1]
     kv_seq_len = full_k.shape[1]
     num_q_heads = full_q.shape[2]
@@ -1475,13 +1475,22 @@ def flex_attention(
         # return_aux rather than the deprecated return_lse: the latter records a
         # FutureWarning in a module-level set, which Dynamo rejects as an unsafe
         # side effect inside the activation-checkpointing HOP.
+        try:
+            from torch.nn.attention.flex_attention import AuxRequest as aux_request_cls
+        except ImportError:
+            aux_request_cls = None
+        if aux_request_cls is None:
+            raise RuntimeError(
+                "return_lse=True requires torch.nn.attention.flex_attention.AuxRequest, "
+                f"which is unavailable in torch {torch.__version__}."
+            )
         attn_out, aux = _COMPILED_FLEX_ATTENTION(
             q,
             k,
             v,
             block_mask=block_mask,
             enable_gqa=num_q_heads != num_kv_heads,
-            return_aux=AuxRequest(lse=True),
+            return_aux=aux_request_cls(lse=True),
             kernel_options=backend.kernel_options,
         )  # attn_out: [1,num_q_heads,N_full,head_dim], aux.lse: [1,num_q_heads,N_full]
         # Convert to the heads-last layout ([1,S,H,D] / [1,S,H]) that from_mode_splits