fix(trainer): rework FA4 custom op to the FA2/FA3 varlen signature, dispatched via _funcs

src/prime_rl/trainer/activation_checkpointing.py

@@ -42,6 +42,7 @@
         "_c10d_functional",
         "flash_attn",
         "flash_attn_3",
+        "prime_rl_attn",
         "prime_rl_collectives",
         "prime_rl_ring",
     }

src/prime_rl/trainer/models/layers/attn.py

@@ -22,10 +22,172 @@
 except ImportError:
     flash_attn_3_varlen_func = None  # type: ignore
 
+# flash-attention-4
+# flash_attn_4_varlen_func stays the raw kernel: gpt_oss and afmoe import it from here.
 try:
     from flash_attn.cute import flash_attn_varlen_func as flash_attn_4_varlen_func
+    from flash_attn.cute.interface import _flash_attn_bwd, _flash_attn_fwd
 except ImportError:
     flash_attn_4_varlen_func = None  # type: ignore
+    _flash_attn_fwd = None  # type: ignore
+    _flash_attn_bwd = None  # type: ignore
+
+
+# FA4's flash_attn_varlen_func is a Python autograd.Function that Dynamo cannot trace. Inside a
+# checkpointed block that graph break drops the whole block to eager, so FA4's kernels are exposed
+# as opaque custom ops with FA2/FA3's varlen signature and dispatched through FlashAttention._funcs.
+# That keeps _compute_attention version-free and lets selective AC save the op output.
+@torch.library.custom_op("prime_rl_attn::flash_attn_4_varlen", mutates_args=())
+def flash_attn_4_varlen(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    max_seqlen_q: int | None,
+    max_seqlen_k: int | None,
+    causal: bool,
+    window_size_left: int | None,
+    window_size_right: int | None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    out, lse, _, _ = _flash_attn_fwd(
+        q,
+        k,
+        v,
+        cu_seqlens_q=cu_seqlens_q,
+        cu_seqlens_k=cu_seqlens_k,
+        max_seqlen_q=max_seqlen_q,
+        max_seqlen_k=max_seqlen_k,
+        causal=causal,
+        window_size_left=window_size_left,
+        window_size_right=window_size_right,
+        return_lse=True,
+    )
+    return out, lse
+
+
+@flash_attn_4_varlen.register_fake
+def _(q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, causal, window_size_left, window_size_right):
+    out = q.new_empty((*q.shape[:-1], v.shape[-1]))
+    lse = q.new_empty((q.shape[1], q.shape[0]), dtype=torch.float32)
+    return out, lse
+
+
+@torch.library.custom_op("prime_rl_attn::flash_attn_4_varlen_backward", mutates_args=())
+def flash_attn_4_varlen_backward(
+    dout: torch.Tensor,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    out: torch.Tensor,
+    lse: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    max_seqlen_q: int | None,
+    max_seqlen_k: int | None,
+    causal: bool,
+    window_size_left: int | None,
+    window_size_right: int | None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    dq, dk, dv = _flash_attn_bwd(
+        q,
+        k,
+        v,
+        out,
+        dout,
+        lse,
+        causal=causal,
+        window_size_left=window_size_left,
+        window_size_right=window_size_right,
+        cu_seqlens_q=cu_seqlens_q,
+        cu_seqlens_k=cu_seqlens_k,
+        max_seqlen_q=max_seqlen_q,
+        max_seqlen_k=max_seqlen_k,
+    )
+    return dq, dk, dv
+
+
+@flash_attn_4_varlen_backward.register_fake
+def _(
+    dout,
+    q,
+    k,
+    v,
+    out,
+    lse,
+    cu_seqlens_q,
+    cu_seqlens_k,
+    max_seqlen_q,
+    max_seqlen_k,
+    causal,
+    window_size_left,
+    window_size_right,
+):
+    return torch.empty_like(q), torch.empty_like(k), torch.empty_like(v)
+
+
+def _flash_attn_4_varlen_setup_context(ctx, inputs, output) -> None:
+    q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, causal, window_size_left, window_size_right = (
+        inputs
+    )
+    out, lse = output
+    ctx.save_for_backward(q, k, v, out, lse, cu_seqlens_q, cu_seqlens_k)
+    ctx.max_seqlen_q = max_seqlen_q
+    ctx.max_seqlen_k = max_seqlen_k
+    ctx.causal = causal
+    ctx.window_size_left = window_size_left
+    ctx.window_size_right = window_size_right
+
+
+def _flash_attn_4_varlen_autograd(ctx, dout: torch.Tensor, _dlse: torch.Tensor | None):
+    q, k, v, out, lse, cu_seqlens_q, cu_seqlens_k = ctx.saved_tensors
+    dq, dk, dv = flash_attn_4_varlen_backward(
+        dout.contiguous(),
+        q,
+        k,
+        v,
+        out,
+        lse,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        ctx.max_seqlen_q,
+        ctx.max_seqlen_k,
+        ctx.causal,
+        ctx.window_size_left,
+        ctx.window_size_right,
+    )
+    return dq, dk, dv, None, None, None, None, None, None, None
+
+
+flash_attn_4_varlen.register_autograd(_flash_attn_4_varlen_autograd, setup_context=_flash_attn_4_varlen_setup_context)
+
+
+def _flash_attn_4_varlen_func(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    max_seqlen_q: int | None,
+    max_seqlen_k: int | None,
+    causal: bool = False,
+    window_size: tuple[int, int] = (-1, -1),
+) -> torch.Tensor:
+    """FA4 with the same call convention as FA2/FA3's flash_attn_varlen_func, for _funcs dispatch."""
+    window_size_left, window_size_right = window_size
+    out, _ = flash_attn_4_varlen(
+        q,
+        k,
+        v,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        max_seqlen_q,
+        max_seqlen_k,
+        causal,
+        None if window_size_left == -1 else window_size_left,
+        None if window_size_right == -1 else window_size_right,
+    )
+    return out
 
 
 @dataclass
@@ -55,7 +217,7 @@ class FlashAttention(nn.Module):
     _funcs = {
         2: flash_attn_varlen_func,
         3: flash_attn_3_varlen_func,
-        4: flash_attn_4_varlen_func,
+        4: _flash_attn_4_varlen_func if _flash_attn_fwd is not None else None,
     }
 
     def __init__(self, config: AttentionConfig, flash_attn_version: int = 2):
@@ -104,15 +266,7 @@ def _compute_attention(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
         sliding_window = getattr(self, "sliding_window", None)
         if sliding_window is not None:
             kwargs["window_size"] = (sliding_window - 1, 0)
-        if self._flash_attn_version == 4:
-            # FA4's flash_attn_varlen_func has qv as the 4th positional arg,
-            # so cu_seqlens must be passed as keyword args to avoid misalignment.
-            kwargs["cu_seqlens_q"] = cu_seqlens
-            kwargs["cu_seqlens_k"] = cu_seqlens
-            out, _ = self.func(q, k, v, **kwargs)
-        else:
-            out = self.func(q, k, v, cu_seqlens, cu_seqlens, max_seqlen, max_seqlen, **kwargs)
-        return out
+        return self.func(q, k, v, cu_seqlens, cu_seqlens, max_seqlen, max_seqlen, **kwargs)
 
     def forward(
         self,