fix(mot): gapped replay on GPUs FlashAttention-4 has no kernels for (sm120) by hakuturu583 · Pull Request #302 · NVIDIA/cosmos-framework

hakuturu583 · 2026-10-21T23:58:54Z

What

Maskless gapped replay (cosmos3::gapped_replay_forward) runs FlashAttention-4's varlen forward with seqused_q / seqused_k, and FA4 has kernels for sm90 / sm100 / sm103 only. On Blackwell workstation GPUs (RTX PRO 6000, RTX 50 series: sm120) every causal multiview checkpoint that uses it, e.g. Cosmos3-Nano-Sim-Auto, is refused at startup:

Maskless replay ... needs FlashAttention-4 on Hopper (sm90) or Blackwell (sm100/sm103), got sm120

This adds a PyTorch path for the forward, used where FA4 has no kernels (or with COSMOS_GAPPED_REPLAY_TORCH=1):

Each launch's segments are gathered into one packed batch: the used query rows of every segment, and its used keys / values. The batch runs through PyTorch's FlashAttention-2 varlen kernel (aten._flash_attention_forward, which has no seqused_q). The output and the log-sum-exp are then scattered back to the segment rows.
The rows FA4 would not write are left as they are:
- rows past a segment's used length;
- a segment with no keys, which gets a zero output and a -inf LSE.
The gather is built once per launch and fold, on the plan's runtime, so the segment lengths are read on the host once per plan rather than at every step.
Head sizes or dtypes FA2 does not take run as FP32 matmuls per segment.
Forward only (inference). The backward and the FA4 path are unchanged.

gapped_replay_unavailable_reason returns None where this path applies.

Tests

New maskless_attention_torch_replay_test.py checks the forward against a per-segment FP64 reference. The segments have gaps, a segment with no keys, and rows past each segment's used length.
- FP32 path: on CPU.
- FA2 path: bf16 on CUDA.
- It also checks that untouched rows keep their values, and how the gather packs the used rows.
On an RTX PRO 6000 (sm120, CUDA 13.0, torch 2.10):
- the new tests and multiview_maskless_exact_test.py pass (332 passed, 16 skipped);
- the gapped-replay tests of omni_mot_causal_model_test.py pass (-k "replay or gapped or FlashAttention": 15 passed).

Tried on an RTX PRO 6000

Cosmos3-Nano-Sim-Auto, joint camera (480p, 1 view) + LiDAR transfer, 81 frames, with cosmos_framework.scripts.inference:

replay forward	wall time of the rollout + decode
per-segment FP32 matmuls (a first version)	241 s
FlashAttention-2 gather (this PR)	48 s

The kernel alone: 2.0 ms against 37.7 ms per launch.
The two runs' first frame agrees to 40 dB PSNR. The difference then grows along the autoregressive rollout, as expected for BF16 against FP32 attention.

🤖 Generated with Claude Code

https://claude.ai/code/session_01RVDWXJs3WKZx7KruLfASrT

…sm120) Maskless replay launches FA4's varlen forward with seqused_q/seqused_k, and FA4 only builds for sm90/sm100/sm103, so RTX PRO 6000 / RTX 50 (sm120) refused Cosmos3-Nano-Sim-Auto. There (or with COSMOS_GAPPED_REPLAY_TORCH=1) each launch gathers its segments' used rows into one packed batch for PyTorch's FlashAttention-2 varlen kernel and scatters the output and LSE back; head sizes FA2 does not take run in FP32 matmuls. Forward only (inference). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01RVDWXJs3WKZx7KruLfASrT Signed-off-by: Masaya Kataoka <ms.kataoka@gmail.com>

pengcuo

Thanks for the detailed PR and the RTX PRO 6000 measurements. Two things block this for me; details inline.

1. Premise. FA4 does ship kernels for sm120. flash_attn/cute/interface.py dispatches arch // 10 in [8, 12] to the FlashAttentionForwardSm80 CuTe kernel (already present before the pinned 4.0.0b24 was cut, e.g. commit b54df16 of 2026-07-23), and _flash_attn_bwd asserts arch // 10 in [9, 10, 11, 12]. What refuses sm120 is our own gate in gapped_replay_unavailable_reason, not FA4. The one known defect on that path is Dao-AILab/flash-attention#2763 "[CuTe, SM80/SM120] Guard invalid varlen forward tiles" (merged 2026-09-08; the first release after it is 4.0.0b30), which is exactly the gapped varlen launch pattern used here. #295 takes that route: admit sm120 in the gate and bump FA4 to b30, with no second attention implementation. A 120-line PyTorch/FA2 fallback with its own numerics, host-side planning and no backward is a lot of surface for a problem the existing backend already covers.

What would change my mind: run the existing replay GPU tests (multiview_maskless_exact_test.py, omni_mot_causal_model_test.py -k "replay or gapped") on your RTX PRO 6000 with the #295 branch (gate + FA4 b30). If the FA4 sm120 path fails them, that evidence belongs in this PR and justifies a fallback. If they pass, I would rather land #295 and close this one.

2. The gate now admits training on sm120 with no backward. gapped_replay_unavailable_reason returns None for every CUDA device that is not sm90/sm10x (sm80/86/89/110/120/121, ...), and require_gapped_replay runs at model build for training as well as inference. Forward goes through the torch path, but _gapped_replay_backward is unchanged: it still calls _flash_attn_interface() and FA4's kernels. Without flash_attn.cute installed that is an ImportError inside autograd on the first backward; with it installed, a forward/backward kernel mismatch nobody has tested. If the fallback stays, the gate must stay closed for training, or the backward must raise a clear error when _torch_replay(device) is true.

Smaller points inline. CI has not run on this PR yet (fork workflow approval pending).

pengcuo · 2026-10-22T02:20:30Z

+    if os.environ.get("COSMOS_GAPPED_REPLAY_TORCH") == "1":
+        return True
+    major, minor = torch.cuda.get_device_capability(device)
+    return (major, minor) != (9, 0) and major != 10


This routes every CUDA device that is not sm90/sm10x to the torch path by elimination: sm80/86/89 (fine for FA2), sm110 (Thor, which FA4 supports for fwd+bwd), and anything below sm80, where aten._flash_attention_forward itself raises an opaque error. If the fallback is kept, select by what FA4 actually lacks and reject the rest with a message, rather than treating "not sm90/sm100" as "torch".

pengcuo · 2026-10-22T02:20:30Z

        return f"it needs a CUDA device, got {device}"
    if torch.version.hip is not None:
        return "it needs NVIDIA CUDA, not ROCm"
+    if _torch_replay(device):


This is where training on sm120 now passes the gate (review body, item 2). _gapped_replay_backward below is unchanged and still requires FA4, so a fine-tune on an RTX PRO 6000 builds, runs its forward, and fails inside autograd.

pengcuo · 2026-10-22T02:20:30Z

+    """Whether gapped replay runs in PyTorch (:func:`_torch_varlen_fwd`): on a GPU FA4 has no
+    kernels for (sm120: RTX PRO 6000 / RTX 50 Blackwell), or with ``COSMOS_GAPPED_REPLAY_TORCH=1``.
+    Forward only: inference."""
+    import os


Module-level import, please; the file already imports its stdlib dependencies at the top. Also COSMOS_GAPPED_REPLAY_TORCH is a new user-facing switch; if it stays, document it with the other COSMOS_* variables in docs/faq.md.

pengcuo · 2026-10-22T02:20:30Z

+
+    def __init__(self, cu_q: torch.Tensor, used_q: torch.Tensor, cu_k: torch.Tensor, used_k: torch.Tensor) -> None:
+        device = cu_q.device
+        qs, qn = cu_q[:-1].tolist(), used_q.tolist()


These tolist() calls are host syncs inside the opaque custom op, executed lazily on the first forward that sees each (launch, fold). The segment geometry already exists as Python ints in _Runtime.__init__ (q_values / k_values, which segment_tensors is derived from), so the gathers can be built there without a device round-trip and without a lazily populated cache on the runtime object.

pengcuo · 2026-10-22T02:20:30Z

+    values: torch.Tensor,  # [m,H_KV,Dv]
+    plan: _VarlenGather,
+    scale: float,
+    block_rows: int = 2048,


scores is materialized as [H_KV, block_rows, kn] fp32. With kn around 1e5 keys and 8 KV heads that is roughly 6.5 GB per block, which is an OOM on exactly the workstation GPUs this path targets. Derive the block size from a byte budget (max(1, budget // (H_KV * kn * 4))) instead of a fixed row count.

pengcuo · 2026-10-22T02:20:30Z

+        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(


The [H, total_q] LSE layout that lse[:, plan.q] = block_lse relies on is PyTorch's unpadded_lse varlen path (checked in v2.10.0 flash_api.cpp: softmax_lse = at::empty({num_heads, total_q})). Older torch returned [B, H, max_q], which would fail the assignment with a shape error; a one-line comment naming the dependency would help whoever moves the torch pin.

- Select the PyTorch path only on sm12x (or COSMOS_GAPPED_REPLAY_TORCH=1); other non-sm90/sm10x GPUs are rejected with a clear message as before. - The path is forward-only: _gapped_replay_backward now raises a clear error instead of failing inside FA4 during training. - Build the row gathers from the runtime's host offsets, removing the tolist() device syncs inside the custom op. - Size the FP32 matmul fallback's score blocks from a byte budget. - Comment the [H, total_q] LSE layout the scatter relies on, move the os import to module level, and document COSMOS_GAPPED_REPLAY_TORCH. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01KL4MFoym4MUrzAj3WXDFmj Signed-off-by: Masaya Kataoka <ms.kataoka@gmail.com>

pengcuo requested changes Oct 22, 2026

View reviewed changes

pengcuo mentioned this pull request Oct 22, 2026

fix: admit SM120 and use FA4 with invalid-tile fix #295

Draft

biofoolgreen mentioned this pull request Oct 26, 2026

fix(sim-auto): support SM120 replay with dense FA4 batches #312

Open

Navigation Menu

Repository navigation

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

fix(mot): gapped replay on GPUs FlashAttention-4 has no kernels for (sm120)#302
hakuturu583 wants to merge 2 commits into
NVIDIA:mainfrom
hakuturu583:fix/sm120-gapped-replay

hakuturu583 commented Oct 21, 2026

Uh oh!

pengcuo left a comment

Uh oh!

pengcuo Oct 22, 2026

Uh oh!

pengcuo Oct 22, 2026

Uh oh!

pengcuo Oct 22, 2026

Uh oh!

pengcuo Oct 22, 2026

Uh oh!

pengcuo Oct 22, 2026

Uh oh!

pengcuo Oct 22, 2026

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

hakuturu583 commented Oct 21, 2026

What

Tests

Tried on an RTX PRO 6000

Uh oh!

pengcuo left a comment

Choose a reason for hiding this comment

Uh oh!

pengcuo Oct 22, 2026

Choose a reason for hiding this comment

Uh oh!

pengcuo Oct 22, 2026

Choose a reason for hiding this comment

Uh oh!

pengcuo Oct 22, 2026

Choose a reason for hiding this comment

Uh oh!

pengcuo Oct 22, 2026

Choose a reason for hiding this comment

Uh oh!

pengcuo Oct 22, 2026

Choose a reason for hiding this comment

Uh oh!

pengcuo Oct 22, 2026

Choose a reason for hiding this comment

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants