fix(mot): gapped replay on GPUs FlashAttention-4 has no kernels for (sm120)

cosmos_framework/model/generator/mot/maskless_attention.py

@@ -38,7 +38,9 @@
 FlashAttention-4 is not a torch.library operator: traced by torch.compile it returns garbage.
 The replay therefore runs inside two opaque custom operators whose only plan input is a CPU
 handle, so a new pack is a new handle value and never a recompile.
-Execution requires NVIDIA CUDA, FlashAttention-4 and Triton; there is no alternate backend.
+Execution requires NVIDIA CUDA, FlashAttention-4 and Triton. On workstation Blackwell (sm12x), or
+with ``COSMOS_GAPPED_REPLAY_TORCH=1``, the forward instead packs each launch's used rows for PyTorch's
+FlashAttention-2 varlen kernel; that path is inference-only and its backward raises.
 
 The kernels read gap rows (partial tiles, backward preprocessing) but never write them. Gap
 rows are real tokens of the same pack and must be finite, as every token must be.
@@ -47,6 +49,7 @@
 import bisect
 import heapq
 import itertools
+import os
 import weakref
 from collections.abc import Sequence
 from dataclasses import dataclass, fields
@@ -425,17 +428,139 @@ def _flash_attn_interface() -> ModuleType:
     return import_module("flash_attn.cute.interface")
 
 
+# Byte budget of one FP32 score block in the matmul fallback (scores and their exponent).
+_MATMUL_SCORE_BYTES = 1 << 29
+
+
+def _fa4_replay_device(device: torch.device) -> bool:
+    """Hopper (sm90) or datacenter Blackwell (sm100/sm103): the GPUs the FA4 replay tests cover."""
+    major, minor = torch.cuda.get_device_capability(device)
+    return (major, minor) == (9, 0) or major == 10
+
+
+def _torch_replay(device: torch.device) -> bool:
+    """Whether gapped replay runs in PyTorch (:func:`_torch_varlen_fwd`) instead of FA4: on
+    workstation Blackwell (sm12x: RTX PRO 6000 / RTX 50), or on any GPU with
+    ``COSMOS_GAPPED_REPLAY_TORCH=1``. Forward only: :func:`_gapped_replay_backward` refuses it."""
+    if os.environ.get("COSMOS_GAPPED_REPLAY_TORCH") == "1":
+        return True
+    return torch.cuda.get_device_capability(device)[0] == 12
+
+
+def _torch_varlen_fwd(
+    q: torch.Tensor,  # [N*G,H_KV,D]
+    k: torch.Tensor,  # [K_span,H_KV,D]
+    v: torch.Tensor,  # [K_span,H_KV,Dv]
+    plan: "_VarlenGather",
+    scale: float,
+    out: torch.Tensor,  # [N*G,H_KV,Dv], written on segment rows only
+    lse: torch.Tensor,  # [H_KV,N*G] FP32, written on segment rows only
+) -> None:
+    """FA4's non-causal varlen forward with ``seqused`` (what gapped replay launches), in PyTorch:
+    each segment's ``used_q`` rows attend to its ``used_k`` keys; other rows are left as they are.
+
+    The segments' used rows are gathered into one packed batch for PyTorch's FlashAttention-2
+    varlen kernel (``aten._flash_attention_forward``, which has no ``seqused_q``), and scattered
+    back. Head sizes FlashAttention-2 does not take run in FP32 matmuls."""
+    if plan.empty_q is not None:
+        out[plan.empty_q] = 0
+        lse[:, plan.empty_q] = -torch.inf
+    if plan.q is None:
+        return
+    queries, keys, values = plan.take(q, plan.q), plan.take(k, plan.k), plan.take(v, plan.k)
+    if q.shape[-1] == v.shape[-1] and q.shape[-1] <= 256 and q.dtype in (torch.float16, torch.bfloat16):
+        result, block_lse, *_ = torch.ops.aten._flash_attention_forward(
+            queries, keys, values, plan.cu_q, plan.cu_k, plan.max_q, plan.max_k, 0.0, False, False, scale=scale
+        )  # [n,H_KV,Dv], [H_KV,n]
+        # Unpadded varlen input makes the LSE [H, total_q] (padded input gives [B, H, max_q]);
+        # the row scatter below relies on that layout.
+    else:
+        result, block_lse = _matmul_varlen_fwd(queries, keys, values, plan, scale)
+    plan.put(out, plan.q, result.to(out.dtype))
+    lse[:, plan.q] = block_lse
+
+
+class _VarlenGather:
+    """Which rows of a gapped varlen launch are used, packed for a dense varlen kernel.
+
+    Built from the host segment geometry (starts and used lengths), so no device sync."""
+
+    def __init__(self, qs: list[int], qn: list[int], ks: list[int], kn: list[int], device: torch.device) -> None:
+        live = [(a, n, b, m) for a, n, b, m in zip(qs, qn, ks, kn, strict=True) if n > 0 and m > 0]
+        empty = [(a, n) for a, n, m in zip(qs, qn, kn, strict=True) if n > 0 and m <= 0]
+        self.empty_q = self._rows([(a, n) for a, n in empty], device) if empty else None
+        if not live:
+            self.q = self.k = None
+            return
+        self.q = self._rows([(a, n) for a, n, _, _ in live], device)
+        self.k = self._rows([(b, m) for _, _, b, m in live], device)
+        self.cu_q = torch.tensor([0, *itertools.accumulate(n for _, n, _, _ in live)], device=device, dtype=torch.int32)
+        self.cu_k = torch.tensor([0, *itertools.accumulate(m for _, _, _, m in live)], device=device, dtype=torch.int32)
+        self.max_q = max(n for _, n, _, _ in live)
+        self.max_k = max(m for _, _, _, m in live)
+        self.lengths = [(n, m) for _, n, _, m in live]
+
+    @staticmethod
+    def _rows(ranges: list[tuple[int, int]], device: torch.device) -> slice | torch.Tensor:
+        """A slice when the ranges are back to back, else the row indexes."""
+        if all(a + n == b for (a, n), (b, _) in zip(ranges, ranges[1:])):
+            return slice(ranges[0][0], ranges[-1][0] + ranges[-1][1])
+        return torch.cat([torch.arange(a, a + n, device=device) for a, n in ranges])
+
+    @staticmethod
+    def take(tensor: torch.Tensor, rows: slice | torch.Tensor) -> torch.Tensor:
+        return tensor[rows] if isinstance(rows, slice) else tensor.index_select(0, rows)
+
+    @staticmethod
+    def put(tensor: torch.Tensor, rows: slice | torch.Tensor, values: torch.Tensor) -> None:
+        if isinstance(rows, slice):
+            tensor[rows] = values
+        else:
+            tensor.index_copy_(0, rows, values)
+
+
+def _matmul_varlen_fwd(
+    queries: torch.Tensor,  # [n,H_KV,D]
+    keys: torch.Tensor,  # [m,H_KV,D]
+    values: torch.Tensor,  # [m,H_KV,Dv]
+    plan: _VarlenGather,
+    scale: float,
+) -> tuple[torch.Tensor, torch.Tensor]:  # [n,H_KV,Dv], [H_KV,n] FP32
+    out = queries.new_empty((queries.shape[0], queries.shape[1], values.shape[-1]), dtype=torch.float32)
+    lse = queries.new_empty((queries.shape[1], queries.shape[0]), dtype=torch.float32)
+    q0 = k0 = 0
+    for qn, kn in plan.lengths:
+        # Two FP32 [H_KV,rows,kn] temporaries (scores and their exponent) per block.
+        block_rows = max(1, _MATMUL_SCORE_BYTES // (2 * 4 * queries.shape[1] * kn))
+        key = keys[k0 : k0 + kn].float().transpose(0, 1)  # [H_KV,kn,D]
+        value = values[k0 : k0 + kn].float().transpose(0, 1)  # [H_KV,kn,Dv]
+        for lo in range(q0, q0 + qn, block_rows):
+            hi = min(lo + block_rows, q0 + qn)
+            scores = torch.matmul(queries[lo:hi].float().transpose(0, 1), key.transpose(1, 2)) * scale
+            block_lse = torch.logsumexp(scores, dim=-1)  # [H_KV,r]
+            out[lo:hi] = torch.matmul(torch.exp(scores - block_lse[..., None]), value).transpose(0, 1)
+            lse[:, lo:hi] = block_lse
+        q0, k0 = q0 + qn, k0 + kn
+    return out, lse
+
+
 def gapped_replay_unavailable_reason(device: torch.device) -> str | None:
     """Why gapped replay cannot run on ``device``, or ``None`` if it can."""
     if device.type != "cuda":
         return f"it needs a CUDA device, got {device}"
     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:
+    if _torch_replay(device):
+        # Forward only: a backward through this path raises in _gapped_replay_backward.
+        return None
+    if not _fa4_replay_device(device):
         # 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}"
+        major, minor = torch.cuda.get_device_capability(device)
+        return (
+            "it needs FlashAttention-4 on Hopper (sm90) or Blackwell (sm100/sm103), or the inference-only "
+            f"PyTorch path on sm12x or with COSMOS_GAPPED_REPLAY_TORCH=1, got sm{major}{minor}"
+        )
     try:
         _flash_attn_interface()
     except ImportError as error:
@@ -599,6 +724,7 @@ def __init__(
     ) -> None:
         q_values: list[int] = []
         k_values: list[int] = []
+        self.device = device
 
         def launch(q_segments: list[tuple[int, int]], q_total: int, k_segments: list[tuple[int, int]]) -> _Launch:
             # FA4 reads segment i at start cu_seqlens[i] with length seqused[i]; cu_seqlens[-1] is the span.
@@ -647,7 +773,24 @@ def launch(q_segments: list[tuple[int, int]], q_total: int, k_segments: list[tup
         self.parts = tuple(parts)
         self.q_offsets = torch.tensor(q_values, dtype=torch.int32, device=device)  # [N_offsets]
         self.k_offsets = torch.tensor(k_values, dtype=torch.int32, device=device)  # [N_offsets]
+        self._q_values, self._k_values = q_values, k_values  # host copies for the PyTorch path
         self._segments: dict[int, dict[int, _SegmentTensors]] = {}
+        self._gathers: dict[tuple[int, int], _VarlenGather] = {}
+
+    def varlen_gather(self, launch: _Launch, fold: int) -> _VarlenGather:
+        """The PyTorch path's packing of one launch, from the host offsets: no device sync."""
+        key = (launch.offset, fold)
+        if key not in self._gathers:
+            start, count = launch.offset, launch.count
+            q, k = self._q_values, self._k_values
+            self._gathers[key] = _VarlenGather(
+                [row * fold for row in q[start : start + count]],
+                [length * fold for length in q[start + count + 1 : start + 2 * count + 1]],
+                k[start : start + count],
+                k[start + count + 1 : start + 2 * count + 1],
+                self.device,
+            )
+        return self._gathers[key]
 
     def segment_tensors(self, launch: _Launch, fold: int) -> _SegmentTensors:  # [B+1], [B], [B+1], [B]
         """``cu_seqlens_q``, ``seqused_q``, ``cu_seqlens_k``, ``seqused_k`` views for one launch.
@@ -844,7 +987,8 @@ def _gapped_replay_forward(
     handle: torch.Tensor,  # [1] int64 on CPU
 ) -> tuple[torch.Tensor, torch.Tensor]:  # [1,Q,H,Dv], [1,Q,H] FP32
     """Run each part's native forward and merge by FP32 log-sum-exp. Rows nothing reaches are zero."""
-    flash_attn = _flash_attn_interface()
+    torch_replay = _torch_replay(q.device)
+    flash_attn = None if torch_replay else _flash_attn_interface()
     if q.shape[2] % k.shape[2] != 0:
         raise ValueError(f"Query heads ({q.shape[2]}) must be a multiple of KV heads ({k.shape[2]}).")
     runtime = _runtime(handle)
@@ -866,6 +1010,10 @@ def _gapped_replay_forward(
             cu_q, used_q, cu_k, used_k = runtime.segment_tensors(launch, fold)
             k_span = k_part[launch.k_lo : launch.k_hi]  # [K_span,H_KV,D]
             v_span = v_part[launch.k_lo : launch.k_hi]  # [K_span,H_KV,Dv]
+            if torch_replay:
+                _torch_varlen_fwd(q_part, k_span, v_span, runtime.varlen_gather(launch, fold), dim**-0.5, out, lse)
+                del k_span, v_span
+                continue
             flash_attn._flash_attn_fwd(
                 q_part,
                 k_span,
@@ -912,6 +1060,14 @@ def _gapped_replay_backward(
     Every launch uses the same folded query layout and segment arguments as forward.
     Rows nothing reaches are zero.
     """
+    if _torch_replay(q.device):
+        major, minor = torch.cuda.get_device_capability(q.device)
+        raise RuntimeError(
+            "Gapped replay's PyTorch path is forward-only (inference); training needs FlashAttention-4 "
+            f"on Hopper (sm90) or Blackwell (sm100/sm103), got sm{major}{minor}"
+            + (" with COSMOS_GAPPED_REPLAY_TORCH=1" if os.environ.get("COSMOS_GAPPED_REPLAY_TORCH") == "1" else "")
+            + "."
+        )
     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_torch_replay_test.py

@@ -0,0 +1,122 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: OpenMDW-1.1
+
+"""The PyTorch gapped-replay forward (GPUs without FlashAttention-4 kernels) against a per-segment reference."""
+
+import pytest
+import torch
+
+from cosmos_framework.model.generator.mot import maskless_attention as m
+
+# (q start, q used), (k start, k used): gaps between segments, a segment with no keys,
+# and rows past each segment's used length that the launch must leave untouched.
+_SEGMENTS = [((0, 290), (0, 900)), ((300, 200), (900, 1200)), ((500, 40), (2100, 0)), ((550, 350), (2100, 1500))]
+_Q, _K = 950, 3600
+
+
+def _gather(device: torch.device) -> m._VarlenGather:
+    return m._VarlenGather(
+        [q[0] for q, _ in _SEGMENTS],
+        [q[1] for q, _ in _SEGMENTS],
+        [k[0] for _, k in _SEGMENTS],
+        [k[1] for _, k in _SEGMENTS],
+        device,
+    )
+
+
+def _reference(q, k, v, scale):  # [Q,H,D], [K,H,D], [K,H,Dv] -> rows -> (out [n,H,Dv], lse [H,n])
+    result = {}
+    for (qs, qn), (ks, kn) in _SEGMENTS:
+        if kn == 0:
+            result[(qs, qn)] = None
+            continue
+        scores = torch.einsum("qhd,khd->hqk", q[qs : qs + qn].double(), k[ks : ks + kn].double()) * scale
+        lse = torch.logsumexp(scores, dim=-1)
+        out = torch.einsum("hqk,khd->qhd", torch.exp(scores - lse[..., None]), v[ks : ks + kn].double())
+        result[(qs, qn)] = (out, lse)
+    return result
+
+
+def _check(q, k, v, *, out_tol: float, lse_tol: float) -> None:
+    scale = q.shape[-1] ** -0.5
+    gather = _gather(q.device)
+    out = q.new_full((_Q, q.shape[1], v.shape[-1]), 7.0)
+    lse = q.new_full((q.shape[1], _Q), 7.0, dtype=torch.float32)
+    m._torch_varlen_fwd(q, k, v, gather, scale, out, lse)
+    reference = _reference(q, k, v, scale)
+    written = torch.zeros(_Q, dtype=torch.bool)
+    for (qs, qn), expected in reference.items():
+        written[qs : qs + qn] = True
+        if expected is None:  # a segment with no keys: zero output, -inf LSE
+            assert (out[qs : qs + qn] == 0).all()
+            assert torch.isneginf(lse[:, qs : qs + qn]).all()
+            continue
+        torch.testing.assert_close(out[qs : qs + qn].double(), expected[0], atol=out_tol, rtol=out_tol)
+        torch.testing.assert_close(lse[:, qs : qs + qn].double(), expected[1], atol=lse_tol, rtol=lse_tol)
+    assert (out[~written.to(out.device)] == 7.0).all()  # rows no segment uses are left as they are
+    assert (lse[:, ~written.to(lse.device)] == 7.0).all()
+
+
+@pytest.mark.L0
+@pytest.mark.CPU
+def test_gather_packs_used_rows() -> None:
+    gather = _gather(torch.device("cpu"))
+    assert gather.lengths == [(290, 900), (200, 1200), (350, 1500)]
+    assert gather.cu_q.tolist() == [0, 290, 490, 840]
+    assert gather.cu_k.tolist() == [0, 900, 2100, 3600]
+    assert (gather.max_q, gather.max_k) == (350, 1500)
+    assert not isinstance(gather.q, slice)  # q has gaps
+    assert gather.k == slice(0, 3600)  # k does not
+    assert gather.empty_q == slice(500, 540)  # the segment with no keys
+
+
+@pytest.mark.L0
+@pytest.mark.CPU
+def test_fp32_path_matches_reference() -> None:
+    torch.manual_seed(0)
+    q, k, v = torch.randn(_Q, 4, 64), torch.randn(_K, 4, 64), torch.randn(_K, 4, 64)
+    _check(q, k, v, out_tol=1e-4, lse_tol=1e-4)
+
+
+@pytest.mark.L0
+@pytest.mark.CPU
+def test_fp32_path_small_score_budget_matches_reference(monkeypatch: pytest.MonkeyPatch) -> None:
+    # One query row per block: the blocks a byte budget picks must tile every segment exactly.
+    monkeypatch.setattr(m, "_MATMUL_SCORE_BYTES", 1)
+    torch.manual_seed(0)
+    q, k, v = torch.randn(_Q, 2, 32), torch.randn(_K, 2, 32), torch.randn(_K, 2, 32)
+    _check(q, k, v, out_tol=1e-4, lse_tol=1e-4)
+
+
+@pytest.mark.L0
+@pytest.mark.CPU
+@pytest.mark.parametrize(
+    ("capability", "torch_replay", "admitted"),
+    [
+        ((9, 0), False, True),
+        ((10, 0), False, True),
+        ((12, 0), True, True),
+        ((8, 6), False, False),
+        ((11, 0), False, False),
+    ],
+)
+def test_device_gate(monkeypatch: pytest.MonkeyPatch, capability, torch_replay: bool, admitted: bool) -> None:
+    monkeypatch.delenv("COSMOS_GAPPED_REPLAY_TORCH", raising=False)
+    monkeypatch.setattr(torch.version, "hip", None)
+    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda device=None: capability)
+    monkeypatch.setattr(m, "_flash_attn_interface", lambda: None)
+    device = torch.device("cuda")
+    assert m._torch_replay(device) is torch_replay
+    assert (m.gapped_replay_unavailable_reason(device) is None) is admitted
+    monkeypatch.setenv("COSMOS_GAPPED_REPLAY_TORCH", "1")
+    assert m._torch_replay(device)
+    assert m.gapped_replay_unavailable_reason(device) is None
+
+
+@pytest.mark.L0
+@pytest.mark.GPU
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="FlashAttention-2 needs CUDA")
+def test_flash_attention_2_path_matches_reference() -> None:
+    torch.manual_seed(0)
+    q, k, v = (torch.randn(n, 8, 128, device="cuda", dtype=torch.bfloat16) for n in (_Q, _K, _K))
+    _check(q, k, v, out_tol=2e-2, lse_tol=1e-3)

docs/environment_variables.md

@@ -3,6 +3,7 @@
 - `IMAGINAIRE_OUTPUT_ROOT=outputs/train`: Output directory for training DCP checkpoints (default is repo-relative). We recommend at least 1 TB of free disk space.
 - `COSMOS_SMOKE=1`: Enable smoke test.
 - `COSMOS_VERBOSE=1`: Enable verbose console output.
+- `COSMOS_GAPPED_REPLAY_TORCH=1`: Run maskless replay's gapped same-view attention forward in PyTorch (FlashAttention-2 varlen) instead of FlashAttention-4 on any NVIDIA GPU. Workstation Blackwell (sm12x: RTX PRO 6000 / RTX 50) uses this path without the variable. Inference only: a backward through it raises.
 - `WANDB_API_KEY`: Required when the recipe TOML's `job.wandb_mode = "online"`. Recipes default to `wandb_mode = "disabled"`; flip the TOML and export this key to log to [Weights & Biases](https://wandb.ai/).
 
 [Hugging Face](https://huggingface.co/docs/huggingface_hub/en/package_reference/environment_variables):