feat(inference): cap consecutive SeaCache reuses (#225)

## Summary

Sync the SeaCache consecutive-reuse cap from
[imaginaire4!12068](https://gitlab-master.nvidia.com/cosmos-lab/imaginaire4/-/merge_requests/12068)
into the OSS inference stack.

- tune the default diffusion-cache threshold from `0.35` to `0.25`
- cap consecutive cached residual reuses at 2 by default, per CFG
pathway
- expose `diffusion_cache_max_consecutive_cached` through setup
overrides (`0` disables the cap)
- forward the override into `DiffusionCache.Config` and include the
effective value in startup logging
- add OSS-path coverage for config validation, state-machine behavior,
argument round-trip, and inference forwarding

The production behavior matches the i4 MR; paths and tests are adapted
to the OSS repository layout.

## Testing

Not run, per request. Static diff and control-flow review completed;
`git diff --check` passes.

Author: lfengad

cosmos_framework/inference/args_test.py

@@ -245,6 +245,21 @@ def check_model_equal(actual: pydantic.BaseModel, expected: pydantic.BaseModel):
     check_model_equal(OmniSetupOverrides.model_validate(args.model_dump()).build_setup(), args)
 
 
+def test_diffusion_cache_max_consecutive_cached_round_trip(tmp_path: Path) -> None:
+    overrides = OmniSetupOverrides(
+        checkpoint_path=DEFAULT_CHECKPOINT_NAME,
+        output_dir=tmp_path / "outputs",
+        diffusion_cache=True,
+        diffusion_cache_max_consecutive_cached=3,
+    )
+
+    args = overrides.build_setup()
+
+    assert args.diffusion_cache is True
+    assert args.diffusion_cache_max_consecutive_cached == 3
+    assert OmniSetupOverrides.model_validate(args.model_dump()).diffusion_cache_max_consecutive_cached == 3
+
+
 def test_sample_args(tmp_path: Path):
     setup_args = OmniSetupOverrides(
         checkpoint_path=DEFAULT_CHECKPOINT_NAME,

cosmos_framework/inference/common/args.py

@@ -885,6 +885,7 @@ class SetupArgs(ABC, CheckpointArgs, ParallelismArgs, QuantizationArgs, Guardrai
     diffusion_cache: bool
     diffusion_cache_thresh: float | None
     diffusion_cache_residual_order: int | None
+    diffusion_cache_max_consecutive_cached: pydantic.NonNegativeInt | None
 
     # Subclass must implement these fields/methods
     # ------------------------------------------------------------
@@ -949,11 +950,15 @@ class SetupOverrides(ABC, CheckpointOverrides, ParallelismOverrides, Quantizatio
     """Accumulated relative-L1 threshold (``diffusion_cache_thresh``), shared by the
     conditional and unconditional pathways. Larger values allow more skipping at
     the cost of lower fidelity. ``None`` uses the ``DiffusionCache.Config`` default
-    (0.35). Only takes effect when diffusion cache is enabled."""
+    (0.25). Only takes effect when diffusion cache is enabled."""
     diffusion_cache_residual_order: int | None = None
     """Polynomial order for extrapolating the generation residual on a skipped step via
     Newton divided differences: 0 = constant reuse, 1 = linear, 2 = quadratic.
     ``None`` uses the default (1).  Only used when diffusion cache is enabled."""
+    diffusion_cache_max_consecutive_cached: pydantic.NonNegativeInt | None = None
+    """Maximum consecutive residual reuses per CFG pathway before forcing a full
+    evaluation. ``0`` disables the limit and ``None`` uses the cache default (2).
+    Only used when diffusion cache is enabled."""
 
     def _build_setup(self):
         if self.num_iterations > 1 and not self.benchmark:

cosmos_framework/inference/inference.py

@@ -1366,6 +1366,8 @@ def _maybe_install_diffusion_cache(pipe: "OmniInference", setup_args: SetupArgs)
             config_overrides["diffusion_cache_thresh"] = setup_args.diffusion_cache_thresh
         if setup_args.diffusion_cache_residual_order is not None:
             config_overrides["residual_order"] = setup_args.diffusion_cache_residual_order
+        if setup_args.diffusion_cache_max_consecutive_cached is not None:
+            config_overrides["max_consecutive_cached"] = setup_args.diffusion_cache_max_consecutive_cached
         install_diffusion_cache(
             pipe=pipe,
             enabled=True,

cosmos_framework/inference/inference_test.py

@@ -10,6 +10,30 @@
 import pytest
 
 
+def test_diffusion_cache_max_consecutive_cached_override(monkeypatch: pytest.MonkeyPatch) -> None:
+    from cosmos_framework.inference import inference
+    from cosmos_framework.model.generator.mot import diffusion_cache
+
+    install = Mock()
+    monkeypatch.setattr(diffusion_cache, "install_diffusion_cache", install)
+    pipe = SimpleNamespace()
+    setup_args = SimpleNamespace(
+        diffusion_cache=True,
+        diffusion_cache_thresh=None,
+        diffusion_cache_residual_order=None,
+        diffusion_cache_max_consecutive_cached=3,
+    )
+
+    inference.OmniInference._maybe_install_diffusion_cache(pipe, setup_args)
+
+    install.assert_called_once_with(
+        pipe=pipe,
+        enabled=True,
+        sample_args_list=[],
+        config_overrides={"max_consecutive_cached": 3},
+    )
+
+
 def test_finalize_data_batch_does_not_mutate_reusable_video_list() -> None:
     torch = pytest.importorskip("torch")
 

cosmos_framework/model/generator/mot/diffusion_cache.py

@@ -413,7 +413,7 @@ class DiffusionCache:
 
     @dataclass(frozen=True, slots=True)
     class Config:
-        diffusion_cache_thresh: float = 0.35
+        diffusion_cache_thresh: float = 0.25
         """Accumulated relative-L1 budget before a full eval is forced (shared across
         the per-step CFG-pass pathways).  Larger ⇒ more skipping ⇒ faster but
         lower fidelity."""
@@ -425,6 +425,9 @@ class Config:
         """Retention: always run full for the first ``ret_steps`` steps (warmup)."""
         cutoff_from_end: int = 1
         """Always run full for the last ``cutoff_from_end`` steps (0 disables)."""
+        max_consecutive_cached: int = 2
+        """Maximum consecutive residual reuses per CFG pathway before forcing a
+        full evaluation. ``0`` disables the limit."""
         power_exp: float = 3.0
         """Exponent of the ``1/|f|^power_exp`` clean-signal power prior in the SEA filter."""
         timestep_max: float = 1000.0
@@ -433,6 +436,14 @@ class Config:
         def __post_init__(self) -> None:
             if self.residual_order < 0:
                 raise ValueError(f"residual_order must be >= 0, got {self.residual_order}")
+            if (
+                isinstance(self.max_consecutive_cached, bool)
+                or not isinstance(self.max_consecutive_cached, int)
+                or self.max_consecutive_cached < 0
+            ):
+                raise ValueError(
+                    f"max_consecutive_cached must be a nonnegative integer, got {self.max_consecutive_cached}"
+                )
 
         @classmethod
         def from_overrides(cls, overrides: dict[str, Any] | None) -> "DiffusionCache.Config":
@@ -457,6 +468,7 @@ class _PathwayState:
         accumulated: float = 0.0
         prev_indicator: list[torch.Tensor] | None = None
         history: list[tuple[int, LMCacheEntry]] = field(default_factory=list)
+        consecutive_cached: int = 0
 
     @dataclass(slots=True)
     class State:
@@ -580,6 +592,7 @@ def patched_lm_forward(self_lm: Any, pack: Any, *args: Any, **kwargs: Any) -> An
                 and _cache_entry_matches(ps.history[-1][1], in_causal, in_full, gen_only)
             )
             if reuse:
+                ps.consecutive_cached += 1
                 und_out = ps.history[-1][1][0]  # understanding: absolute reuse (not a residual)
                 gen_delta = _extrapolate_gen(ps.history, self.state.step, self.config.residual_order)
                 out_pack = dict(pack)
@@ -591,6 +604,7 @@ def patched_lm_forward(self_lm: Any, pack: Any, *args: Any, **kwargs: Any) -> An
 
             outputs = original_lm_forward(pack, *args, **kwargs)
             if caching:
+                ps.consecutive_cached = 0
                 out_pack = outputs[0] if isinstance(outputs, tuple) else outputs
                 entry = _lm_cache_entry(in_causal, in_full, out_pack["causal_seq"], out_pack["full_only_seq"])
                 ps.history.append((self.state.step, entry))
@@ -637,9 +651,13 @@ def _should_compute(self, pathway: str, indicator: list[torch.Tensor] | None) ->
         step = self.state.step
         cfg = self.config
 
+        max_consecutive_forced = bool(
+            cfg.max_consecutive_cached and ps.consecutive_cached >= cfg.max_consecutive_cached
+        )
         forced = (
             step < cfg.ret_steps
             or step >= self.num_steps - cfg.cutoff_from_end
+            or max_consecutive_forced
             or not ps.history
             or indicator is None
             or ps.prev_indicator is None
@@ -813,7 +831,8 @@ def install_diffusion_cache(
         f"Enabled diffusion cache (diffusion-time inference cache) "
         f"max_num_steps={max_num_steps} "
         f"diffusion_cache_thresh={cfg.diffusion_cache_thresh} ret_steps={cfg.ret_steps} "
-        f"cutoff_from_end={cfg.cutoff_from_end} power_exp={cfg.power_exp} "
+        f"cutoff_from_end={cfg.cutoff_from_end} "
+        f"max_consecutive_cached={cfg.max_consecutive_cached} power_exp={cfg.power_exp} "
         f"residual_order={cfg.residual_order}"
     )
     return cache

cosmos_framework/model/generator/mot/diffusion_cache_test.py

@@ -0,0 +1,57 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: OpenMDW-1.1
+
+import pytest
+import torch
+
+from cosmos_framework.model.generator.mot.diffusion_cache import DiffusionCache
+
+
+def _const_indicator(value: float = 1.0) -> list[torch.Tensor]:
+    return [torch.full((1, 2, 2, 1), value)]
+
+
+def _dummy_residual() -> tuple[torch.Tensor, torch.Tensor]:
+    return torch.zeros(1, 2), torch.zeros(1, 2)
+
+
+@pytest.mark.L0
+@pytest.mark.CPU
+def test_max_consecutive_cached_forces_full_per_pathway() -> None:
+    cache = DiffusionCache(
+        num_steps=20,
+        config={
+            "ret_steps": 0,
+            "cutoff_from_end": 0,
+            "diffusion_cache_thresh": 1.0,
+            "max_consecutive_cached": 3,
+        },
+    )
+    cache.state.step = 0
+    assert cache._should_compute("cfg0", _const_indicator()) is True
+    cache._pathways["cfg0"].history = [(0, _dummy_residual())]
+
+    for step in range(1, 4):
+        cache.state.step = step
+        assert cache._should_compute("cfg0", _const_indicator()) is False
+        cache._pathways["cfg0"].consecutive_cached += 1
+
+    cache.state.step = 4
+    assert cache._should_compute("cfg0", _const_indicator()) is True
+
+
+@pytest.mark.L0
+@pytest.mark.CPU
+def test_diffusion_cache_uses_tuned_defaults() -> None:
+    cache = DiffusionCache(num_steps=10)
+
+    assert cache.config.diffusion_cache_thresh == pytest.approx(0.25)
+    assert cache.config.max_consecutive_cached == 2
+
+
+@pytest.mark.L0
+@pytest.mark.CPU
+@pytest.mark.parametrize("value", [-1, 1.5, True])
+def test_invalid_max_consecutive_cached_rejected(value: object) -> None:
+    with pytest.raises(ValueError, match="max_consecutive_cached"):
+        DiffusionCache(num_steps=10, config={"max_consecutive_cached": value})