From 3c7fcd36b157836ae76e47232e91148960a9f594 Mon Sep 17 00:00:00 2001 From: Masaya Kataoka Date: Thu, 22 Oct 2026 08:30:02 +0900 Subject: [PATCH 1/2] fix(mot): gapped replay on GPUs FlashAttention-4 has no kernels for (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 Claude-Session: https://claude.ai/code/session_01RVDWXJs3WKZx7KruLfASrT Signed-off-by: Masaya Kataoka --- .../model/generator/mot/maskless_attention.py | 121 +++++++++++++++++- .../maskless_attention_torch_replay_test.py | 85 ++++++++++++ 2 files changed, 205 insertions(+), 1 deletion(-) create mode 100644 cosmos_framework/model/generator/mot/maskless_attention_torch_replay_test.py diff --git a/cosmos_framework/model/generator/mot/maskless_attention.py b/cosmos_framework/model/generator/mot/maskless_attention.py index 2e4d267b1..d6b509934 100644 --- a/cosmos_framework/model/generator/mot/maskless_attention.py +++ b/cosmos_framework/model/generator/mot/maskless_attention.py @@ -425,12 +425,121 @@ def _flash_attn_interface() -> ModuleType: return import_module("flash_attn.cute.interface") +def _torch_replay(device: torch.device) -> bool: + """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 + + 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 + + +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] + 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.""" + + 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() + ks, kn = cu_k[:-1].tolist(), used_k.tolist() + 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, + block_rows: int = 2048, +) -> 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: + 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" + if _torch_replay(device): + return None 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 @@ -599,6 +708,7 @@ def __init__( ) -> None: q_values: list[int] = [] k_values: list[int] = [] + self.torch_gathers: dict[tuple[int, int], _VarlenGather] = {} 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. @@ -844,7 +954,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 +977,14 @@ 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: + # Built once per launch and fold: the segment lengths are read on the host. + gather = runtime.torch_gathers.get((launch.offset, fold)) + if gather is None: + gather = runtime.torch_gathers[launch.offset, fold] = _VarlenGather(cu_q, used_q, cu_k, used_k) + _torch_varlen_fwd(q_part, k_span, v_span, gather, dim**-0.5, out, lse) + del k_span, v_span + continue flash_attn._flash_attn_fwd( q_part, k_span, diff --git a/cosmos_framework/model/generator/mot/maskless_attention_torch_replay_test.py b/cosmos_framework/model/generator/mot/maskless_attention_torch_replay_test.py new file mode 100644 index 000000000..36b14d4e1 --- /dev/null +++ b/cosmos_framework/model/generator/mot/maskless_attention_torch_replay_test.py @@ -0,0 +1,85 @@ +# 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 _launch(device: torch.device) -> tuple[torch.Tensor, ...]: + cu_q = torch.tensor([q[0] for q, _ in _SEGMENTS] + [_Q], device=device, dtype=torch.int32) + used_q = torch.tensor([q[1] for q, _ in _SEGMENTS], device=device, dtype=torch.int32) + cu_k = torch.tensor([k[0] for _, k in _SEGMENTS] + [_K], device=device, dtype=torch.int32) + used_k = torch.tensor([k[1] for _, k in _SEGMENTS], device=device, dtype=torch.int32) + return cu_q, used_q, cu_k, used_k + + +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 = m._VarlenGather(*_launch(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 = m._VarlenGather(*_launch(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.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) From e10fc0e1c4b2457c14f0bf2a5f83bd3cb25f5ebe Mon Sep 17 00:00:00 2001 From: Masaya Kataoka Date: Fri, 23 Oct 2026 05:26:35 +0000 Subject: [PATCH 2/2] fix(mot): address review on the sm12x gapped-replay PyTorch path - 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 Claude-Session: https://claude.ai/code/session_01KL4MFoym4MUrzAj3WXDFmj Signed-off-by: Masaya Kataoka --- .../model/generator/mot/maskless_attention.py | 83 ++++++++++++++----- .../maskless_attention_torch_replay_test.py | 53 ++++++++++-- docs/environment_variables.md | 1 + 3 files changed, 106 insertions(+), 31 deletions(-) diff --git a/cosmos_framework/model/generator/mot/maskless_attention.py b/cosmos_framework/model/generator/mot/maskless_attention.py index d6b509934..150fd8f4c 100644 --- a/cosmos_framework/model/generator/mot/maskless_attention.py +++ b/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,16 +428,23 @@ def _flash_attn_interface() -> ModuleType: return import_module("flash_attn.cute.interface") -def _torch_replay(device: torch.device) -> bool: - """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 +# 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 - major, minor = torch.cuda.get_device_capability(device) - return (major, minor) != (9, 0) and major != 10 + return torch.cuda.get_device_capability(device)[0] == 12 def _torch_varlen_fwd( @@ -462,6 +472,8 @@ def _torch_varlen_fwd( 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)) @@ -469,12 +481,11 @@ def _torch_varlen_fwd( class _VarlenGather: - """Which rows of a gapped varlen launch are used, packed for a dense varlen kernel.""" + """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, 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() - ks, kn = cu_k[:-1].tolist(), used_k.tolist() + 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 @@ -514,12 +525,13 @@ def _matmul_varlen_fwd( values: torch.Tensor, # [m,H_KV,Dv] plan: _VarlenGather, scale: float, - block_rows: int = 2048, ) -> 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): @@ -539,12 +551,16 @@ def gapped_replay_unavailable_reason(device: torch.device) -> str | None: if torch.version.hip is not None: return "it needs NVIDIA CUDA, not ROCm" if _torch_replay(device): + # Forward only: a backward through this path raises in _gapped_replay_backward. return None - major, minor = torch.cuda.get_device_capability(device) - if (major, minor) != (9, 0) and major != 10: + 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: @@ -708,7 +724,7 @@ def __init__( ) -> None: q_values: list[int] = [] k_values: list[int] = [] - self.torch_gathers: dict[tuple[int, int], _VarlenGather] = {} + 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. @@ -757,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. @@ -978,11 +1011,7 @@ def _gapped_replay_forward( 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: - # Built once per launch and fold: the segment lengths are read on the host. - gather = runtime.torch_gathers.get((launch.offset, fold)) - if gather is None: - gather = runtime.torch_gathers[launch.offset, fold] = _VarlenGather(cu_q, used_q, cu_k, used_k) - _torch_varlen_fwd(q_part, k_span, v_span, gather, dim**-0.5, out, lse) + _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( @@ -1031,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] diff --git a/cosmos_framework/model/generator/mot/maskless_attention_torch_replay_test.py b/cosmos_framework/model/generator/mot/maskless_attention_torch_replay_test.py index 36b14d4e1..f800e9ab5 100644 --- a/cosmos_framework/model/generator/mot/maskless_attention_torch_replay_test.py +++ b/cosmos_framework/model/generator/mot/maskless_attention_torch_replay_test.py @@ -14,12 +14,14 @@ _Q, _K = 950, 3600 -def _launch(device: torch.device) -> tuple[torch.Tensor, ...]: - cu_q = torch.tensor([q[0] for q, _ in _SEGMENTS] + [_Q], device=device, dtype=torch.int32) - used_q = torch.tensor([q[1] for q, _ in _SEGMENTS], device=device, dtype=torch.int32) - cu_k = torch.tensor([k[0] for _, k in _SEGMENTS] + [_K], device=device, dtype=torch.int32) - used_k = torch.tensor([k[1] for _, k in _SEGMENTS], device=device, dtype=torch.int32) - return cu_q, used_q, cu_k, used_k +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]) @@ -37,7 +39,7 @@ def _reference(q, k, v, scale): # [Q,H,D], [K,H,D], [K,H,Dv] -> rows -> (out [n def _check(q, k, v, *, out_tol: float, lse_tol: float) -> None: scale = q.shape[-1] ** -0.5 - gather = m._VarlenGather(*_launch(q.device)) + 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) @@ -58,7 +60,7 @@ def _check(q, k, v, *, out_tol: float, lse_tol: float) -> None: @pytest.mark.L0 @pytest.mark.CPU def test_gather_packs_used_rows() -> None: - gather = m._VarlenGather(*_launch(torch.device("cpu"))) + 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] @@ -76,6 +78,41 @@ def test_fp32_path_matches_reference() -> None: _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") diff --git a/docs/environment_variables.md b/docs/environment_variables.md index 6c5519f3f..f41cd87d8 100644 --- a/docs/environment_variables.md +++ b/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):