fix(sim-auto): enable native FA4 replay on SM120

cosmos_framework/model/generator/mot/maskless_attention.py

@@ -60,8 +60,8 @@
 import torch
 import triton
 
-from cosmos_framework.model.attention.utils.environment import is_torch_compiling
 from cosmos_framework.configs.base.defaults.multiview_attention import temporal_window_bounds
+from cosmos_framework.model.attention.utils.environment import is_torch_compiling
 from cosmos_framework.model.generator.mot import maskless_kernels as kernels
 from cosmos_framework.model.generator.mot.causal_flex_attention import (
     _ROLE_PADDING,
@@ -432,10 +432,10 @@ def gapped_replay_unavailable_reason(device: torch.device) -> str | None:
     if torch.version.hip is not None:
         return "it needs NVIDIA CUDA, not ROCm"
     major, minor = torch.cuda.get_device_capability(device)
-    if (major, minor) != (9, 0) and major != 10:
-        # Hopper and Blackwell use different FA4 kernels; the replay GPU tests cover
-        # the segment isolation and untouched gap rows required by both paths.
-        return f"it needs FlashAttention-4 on Hopper (sm90) or Blackwell (sm100/sm103), got sm{major}{minor}"
+    if (major, minor) not in ((9, 0), (12, 0)) and major != 10:
+        # FA4 >= 4.0.0b30 guards invalid varlen tiles on SM120; keep the
+        # existing gapped replay path rather than introducing a second backend.
+        return f"it needs FlashAttention-4 on Hopper (sm90) or Blackwell (sm100/sm103/sm120), got sm{major}{minor}"
     try:
         _flash_attn_interface()
     except ImportError as error:
@@ -912,6 +912,8 @@ def _gapped_replay_backward(
     Every launch uses the same folded query layout and segment arguments as forward.
     Rows nothing reaches are zero.
     """
+    if torch.cuda.get_device_capability(q.device) == (12, 0):
+        raise NotImplementedError("SM120 replay is inference-only; backward is not qualified.")
     flash_attn = _flash_attn_interface()
     runtime = _runtime(handle)
     q, k, v = (tensor.contiguous() for tensor in (q, k, v))  # [1,Q,H,D], [1,KV,H_KV,D], [1,KV,H_KV,Dv]

cosmos_framework/model/generator/mot/maskless_attention_test.py

@@ -0,0 +1,131 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: OpenMDW-1.1
+
+"""SM120 replay compatibility against independent dense attention references."""
+
+import math
+from types import ModuleType
+
+import pytest
+import torch
+
+from cosmos_framework.model.generator.mot import maskless_attention as replay
+
+
+@pytest.mark.L0
+@pytest.mark.gpus(0)
+@pytest.mark.parametrize("capability", [(9, 0), (10, 0), (10, 3), (12, 0)])
+def test_replay_supported_architectures(monkeypatch: pytest.MonkeyPatch, capability: tuple[int, int]) -> None:
+    monkeypatch.setattr(torch.version, "hip", None)
+    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda device: capability)
+    monkeypatch.setattr(replay, "_flash_attn_interface", lambda: ModuleType("flash_attn"))
+    assert replay.gapped_replay_unavailable_reason(torch.device("cuda")) is None
+
+
+@pytest.mark.L0
+@pytest.mark.gpus(0)
+@pytest.mark.parametrize("capability", [(8, 0), (8, 9), (9, 1), (11, 0), (12, 1)])
+def test_replay_rejects_unqualified_architectures(monkeypatch: pytest.MonkeyPatch, capability: tuple[int, int]) -> None:
+    monkeypatch.setattr(torch.version, "hip", None)
+    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda device: capability)
+    reason = replay.gapped_replay_unavailable_reason(torch.device("cuda"))
+    assert reason is not None
+    assert f"sm{capability[0]}{capability[1]}" in reason
+
+
+@pytest.mark.L0
+@pytest.mark.gpus(0)
+@pytest.mark.parametrize(("device", "hip", "message"), [("cpu", None, "CUDA device"), ("cuda", "6.0", "not ROCm")])
+def test_replay_rejects_non_nvidia_devices(
+    monkeypatch: pytest.MonkeyPatch, device: str, hip: str | None, message: str
+) -> None:
+    monkeypatch.setattr(torch.version, "hip", hip)
+
+    def unexpected_capability_query(device):
+        pytest.fail("Non-NVIDIA devices must be rejected before querying CUDA capability")
+
+    monkeypatch.setattr(torch.cuda, "get_device_capability", unexpected_capability_query)
+    reason = replay.gapped_replay_unavailable_reason(torch.device(device))
+    assert reason is not None
+    assert message in reason
+
+
+@pytest.mark.L0
+@pytest.mark.gpus(0)
+def test_sm120_replay_requires_flash_attention(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(torch.version, "hip", None)
+    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda device: (12, 0))
+
+    def missing_backend():
+        raise ImportError("FA4 is not installed")
+
+    monkeypatch.setattr(replay, "_flash_attn_interface", missing_backend)
+    with pytest.raises(RuntimeError, match="could not be imported: FA4 is not installed"):
+        replay.require_gapped_replay(torch.device("cuda"))
+
+
+def _require_sm120() -> None:
+    if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (12, 0):
+        pytest.skip("Requires an SM120 GPU")
+    pytest.importorskip("flash_attn.cute.interface")
+
+
+@pytest.mark.L0
+@pytest.mark.gpus(1)
+@pytest.mark.parametrize("dim", [64, 128])
+@pytest.mark.parametrize(("heads", "kv_heads"), [(8, 2), (8, 8), (16, 1)])
+def test_sm120_replay_matches_dense_partition_reference(dim: int, heads: int, kv_heads: int) -> None:
+    _require_sm120()
+    torch.manual_seed(20261026)
+    q = torch.randn(1, 257, heads, dim, device="cuda", dtype=torch.bfloat16)
+    k = torch.randn(1, 389, kv_heads, dim, device="cuda", dtype=torch.bfloat16)
+    v = torch.randn_like(k)
+    segments = [
+        replay.GappedSegment(0, 31, 10, 71),
+        replay.GappedSegment(55, 74, 100, 101),
+        replay.GappedSegment(55, 74, 240, 100),
+        replay.GappedSegment(199, 41, 3, 93),
+        replay.GappedSegment(140, 31, 310, 71),
+    ]
+    caption = replay.ReplayMasklessPass(
+        "caption",
+        torch.tensor([0, 5, 55, 100, 199, 225], device="cuda"),
+        torch.tensor([1, 8, 12, 190, 388], device="cuda"),
+        (0, 6),
+        (0, 5),
+    )
+    plan = replay.build_gapped_replay_plan(segments, [caption], torch.device("cuda"))
+    out, lse = replay._gapped_replay_forward(q, k, v, plan.handle)
+
+    # Count each partition's edges independently, including duplicated keys.
+    count = torch.zeros(257, 389, device="cuda")
+    for segment in segments:
+        count[
+            segment.q_row : segment.q_row + segment.q_len,
+            segment.k_row : segment.k_row + segment.k_len,
+        ] += 1
+    count[caption.q_gather[:, None], caption.kv_gather[None, :]] += 1
+    keys = k.float().repeat_interleave(heads // kv_heads, dim=2)
+    values = v.float().repeat_interleave(heads // kv_heads, dim=2)
+    scores = torch.einsum("bqhd,bkhd->bhqk", q.float(), keys) / math.sqrt(dim) + count.log()[None, None]
+    valid = count.sum(-1) > 0
+    weights = torch.nan_to_num(scores.softmax(-1))
+    expected = torch.einsum("bhqk,bkhd->bqhd", weights, values)
+    expected_lse = scores.logsumexp(-1).transpose(1, 2)
+    assert torch.isfinite(out).all()
+    torch.testing.assert_close(out.float(), expected, atol=0.025, rtol=0)
+    torch.testing.assert_close(lse[:, valid], expected_lse[:, valid], atol=0.01, rtol=0)
+    assert torch.count_nonzero(out[:, ~valid]) == 0
+    assert torch.isneginf(lse[:, ~valid]).all()
+
+
+@pytest.mark.L0
+@pytest.mark.gpus(1)
+def test_sm120_replay_backward_is_explicitly_unsupported() -> None:
+    _require_sm120()
+    q = torch.zeros(1, 1, 1, 64, device="cuda", dtype=torch.bfloat16)
+    lse = torch.zeros(1, 1, 1, device="cuda")
+    # An unregistered handle proves the guard runs before plan lookup or a kernel launch.
+    handle = torch.tensor([-1], dtype=torch.int64)
+    with pytest.raises(NotImplementedError, match="inference-only"):
+        replay._gapped_replay_backward(q, q, q, q, lse, q, handle)

docs/inference_sim_auto.md

@@ -40,6 +40,14 @@ Rolling replay runs its same-camera attention through the FlashAttention-4
 kernels. The `cu130-train` group installed there includes FlashAttention-4, so
 no separate installation is needed.
 
+On SM120 GPUs such as RTX 6000D, replay inference uses the same native
+FlashAttention-4 gapped-varlen path as the other supported architectures.
+The pinned FlashAttention-4 4.0.0b30 release includes the SM80/SM120 invalid-tile
+fix required by this workload. Upgrade the environment through the setup command
+above; admitting SM120 with the older 4.0.0b24 pin can cause illegal memory access.
+SM120 replay backward is not qualified and is explicitly rejected; SM120 support
+is inference only.
+
 ## Checkpoint
 
 Download the published snapshot and pass its local directory. A bare Hugging

pyproject.toml

@@ -195,7 +195,7 @@ vllm = [
 # Match nvcr.io/nvidia/pytorch:25.11-py3: https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes/rel-26-01.html
 cu128 = [
   # FA4 CuTeDSL kernels; the cu13 extra is only needed by CUDA 13 environments.
-  "flash-attn-4==4.0.0b24",
+  "flash-attn-4==4.0.0b30",
   "flash-attn-3-nv==1.0.3+cu128.torch210; platform_machine == 'x86_64'",
   "flash-attn==2.7.4.post1+cu128.torch210; platform_machine == 'x86_64'",
   "natten==0.21.6.dev6+cu128.torch210; platform_machine == 'x86_64'",
@@ -230,7 +230,7 @@ cu128-train = [
 # Match nvcr.io/nvidia/pytorch:25.11-py3: https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes/rel-25-11.html
 cu130 = [
   # Match the upstream FA4 installation on both x86_64 and aarch64.
-  "flash-attn-4[cu13]==4.0.0b24",
+  "flash-attn-4[cu13]==4.0.0b30",
   "flash-attn-3-nv==1.0.3+cu130.torch210; platform_machine == 'x86_64'",
   "flash-attn==2.7.4.post1+cu130.torch210; platform_machine == 'x86_64'",
   # GB300/B300 (sm_103a): use the custom build that includes sm_103a kernels (both arches).
@@ -266,7 +266,7 @@ cu130-train = [
 ]
 # torch 2.13 on CUDA 13.0. Same layout as cu130 but pinned to torch 2.13.
 cu130-torch213 = [
-  "flash-attn-4[cu13]==4.0.0b24",
+  "flash-attn-4[cu13]==4.0.0b30",
   #"flash-attn-3-nv==1.0.3+cu130.torch210; platform_machine == 'x86_64'", # TODO: torch213 build (cu130.torch210 wheels are cp313-only + torch210 ABI)
   #"flash-attn==2.7.4.post1+cu130.torch210; platform_machine == 'x86_64'", # TODO: torch213 build (cu130.torch210 wheels are cp313-only + torch210 ABI)
   # torch2.13 build (natten 0.21.6, sm 80/86/89/90/100/103/110/120/121) — matches this group's torch/ABI.