> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server.

# nemo_automodel.components.attention.ffpa_attention

FFPA attention bindings for HF `ALL_ATTENTION_FUNCTIONS` / `ALL_MASK_ATTENTION_FUNCTIONS`.

Routes `head_dim=512` bf16/fp16 layers through the CuTeDSL FFPA kernel,
sliding-window (flex `BlockMask`) layers through FlexAttention, and falls back
to SDPA (or eager for softcap) for other unsupported configurations.

## Module Contents

### Functions

| Name                                                                                                     | Description                                                                               |
| -------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------------------- |
| [`_ffpa_dense_bwd`](#nemo_automodel-components-attention-ffpa_attention-_ffpa_dense_bwd)                 | FFPA CuTeDSL *dense* backward using the *caller-supplied* out/lse.                        |
| [`_ffpa_dense_fwd`](#nemo_automodel-components-attention-ffpa_attention-_ffpa_dense_fwd)                 | FFPA CuTeDSL *dense* forward on `[B, H, N, D]` SDPA-layout tensors.                       |
| [`_ffpa_low_level_ready`](#nemo_automodel-components-attention-ffpa_attention-_ffpa_low_level_ready)     | -                                                                                         |
| [`_ffpa_varlen_bwd`](#nemo_automodel-components-attention-ffpa_attention-_ffpa_varlen_bwd)               | FFPA CuTeDSL varlen backward using the *caller-supplied* out/lse. Returns `(dq, dk, dv)`. |
| [`_ffpa_varlen_fwd`](#nemo_automodel-components-attention-ffpa_attention-_ffpa_varlen_fwd)               | FFPA CuTeDSL varlen forward on packed THD inputs. Returns `(out[T,Hq,D], lse[Hq,T])`.     |
| [`_ffpa_varlen_ready`](#nemo_automodel-components-attention-ffpa_attention-_ffpa_varlen_ready)           | Whether the FFPA CuTeDSL *varlen* ops are importable and registered.                      |
| [`_get_eager`](#nemo_automodel-components-attention-ffpa_attention-_get_eager)                           | -                                                                                         |
| [`_get_ffpa_high_level`](#nemo_automodel-components-attention-ffpa_attention-_get_ffpa_high_level)       | -                                                                                         |
| [`_get_flex`](#nemo_automodel-components-attention-ffpa_attention-_get_flex)                             | -                                                                                         |
| [`_get_sdpa`](#nemo_automodel-components-attention-ffpa_attention-_get_sdpa)                             | -                                                                                         |
| [`_is_block_mask`](#nemo_automodel-components-attention-ffpa_attention-_is_block_mask)                   | True for a flex `BlockMask` — the marker that a layer should run on FlexAttention.        |
| [`_warn_once`](#nemo_automodel-components-attention-ffpa_attention-_warn_once)                           | -                                                                                         |
| [`ffpa_attention_forward`](#nemo_automodel-components-attention-ffpa_attention-ffpa_attention_forward)   | HF attention-interface forward backed by `torch.ops.ffpa_attn._fwd_cute`.                 |
| [`ffpa_mask`](#nemo_automodel-components-attention-ffpa_attention-ffpa_mask)                             | Mask factory for `ALL_MASK_ATTENTION_FUNCTIONS["ffpa"]`.                                  |
| [`register_ffpa_attention`](#nemo_automodel-components-attention-ffpa_attention-register_ffpa_attention) | Register `"ffpa"` in HF attention/mask registries. Idempotent.                            |
| [`setup_ffpa_backend`](#nemo_automodel-components-attention-ffpa_attention-setup_ffpa_backend)           | Validate FFPA preconditions, then register the HF `"ffpa"` backend.                       |

### Data

[`_CUTEDSL_BACKEND`](#nemo_automodel-components-attention-ffpa_attention-_CUTEDSL_BACKEND)

[`_EAGER_FN`](#nemo_automodel-components-attention-ffpa_attention-_EAGER_FN)

[`_FALLBACK_WARNED`](#nemo_automodel-components-attention-ffpa_attention-_FALLBACK_WARNED)

[`_FFPA_FN`](#nemo_automodel-components-attention-ffpa_attention-_FFPA_FN)

[`_FFPA_HEAD_DIM`](#nemo_automodel-components-attention-ffpa_attention-_FFPA_HEAD_DIM)

[`_FFPA_LOW_LEVEL_READY`](#nemo_automodel-components-attention-ffpa_attention-_FFPA_LOW_LEVEL_READY)

[`_FLEX_FN`](#nemo_automodel-components-attention-ffpa_attention-_FLEX_FN)

[`_REGISTERED`](#nemo_automodel-components-attention-ffpa_attention-_REGISTERED)

[`_SDPA_FN`](#nemo_automodel-components-attention-ffpa_attention-_SDPA_FN)

[`_VARLEN_WIN_NONE`](#nemo_automodel-components-attention-ffpa_attention-_VARLEN_WIN_NONE)

[`__all__`](#nemo_automodel-components-attention-ffpa_attention-__all__)

[`logger`](#nemo_automodel-components-attention-ffpa_attention-logger)

### API

```python
nemo_automodel.components.attention.ffpa_attention._ffpa_dense_bwd(
    grad_out: torch.Tensor,
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    out: torch.Tensor,
    lse: torch.Tensor,
    scale: float,
    causal: bool
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]
```

FFPA CuTeDSL *dense* backward using the *caller-supplied* out/lse.

Returns `(dq, dk, dv)` all in `[B, H, N, D]` SDPA layout (dK/dV reduced to
`Hkv` for GQA). Exposed as an explicit call (not the package's autograd) so
the ring backward can feed the globally merged out/lse, not a chunk-local one.

```python
nemo_automodel.components.attention.ffpa_attention._ffpa_dense_fwd(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    scale: float,
    causal: bool
) -> tuple[torch.Tensor, torch.Tensor]
```

FFPA CuTeDSL *dense* forward on `[B, H, N, D]` SDPA-layout tensors.

Returns `(out[B, Hq, Nq, D], lse[B, Hq, Nq] fp32)`; handles GQA and causal/non-causal
internally. Exposed (not via the autograd `ffpa_attn_func`) so the ring gets the
per-chunk `lse` for its online-softmax merge.

```python
nemo_automodel.components.attention.ffpa_attention._ffpa_low_level_ready() -> bool
```

```python
nemo_automodel.components.attention.ffpa_attention._ffpa_varlen_bwd(
    grad_out_pack: torch.Tensor,
    q_pack: torch.Tensor,
    k_pack: torch.Tensor,
    v_pack: torch.Tensor,
    out_pack: torch.Tensor,
    lse_pack: torch.Tensor,
    cu_q: torch.Tensor,
    cu_k: torch.Tensor,
    max_q: int,
    max_k: int,
    scale: float,
    causal: bool
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]
```

FFPA CuTeDSL varlen backward using the *caller-supplied* out/lse. Returns `(dq, dk, dv)`.

Exposed as an explicit op call (not the package's varlen autograd) because the
ring backward feeds the globally merged out/lse, not a chunk-local one.

```python
nemo_automodel.components.attention.ffpa_attention._ffpa_varlen_fwd(
    q_pack: torch.Tensor,
    k_pack: torch.Tensor,
    v_pack: torch.Tensor,
    cu_q: torch.Tensor,
    cu_k: torch.Tensor,
    max_q: int,
    max_k: int,
    scale: float,
    causal: bool
) -> tuple[torch.Tensor, torch.Tensor]
```

FFPA CuTeDSL varlen forward on packed THD inputs. Returns `(out[T,Hq,D], lse[Hq,T])`.

Wraps `_varlen_fwd_cute` with the fixed no-window/no-softcap/no-pack-gqa sentinels.

```python
nemo_automodel.components.attention.ffpa_attention._ffpa_varlen_ready() -> bool
```

Whether the FFPA CuTeDSL *varlen* ops are importable and registered.

```python
nemo_automodel.components.attention.ffpa_attention._get_eager() -> typing.Callable | None
```

```python
nemo_automodel.components.attention.ffpa_attention._get_ffpa_high_level() -> tuple[typing.Callable | None, typing.Any]
```

```python
nemo_automodel.components.attention.ffpa_attention._get_flex() -> typing.Callable | None
```

```python
nemo_automodel.components.attention.ffpa_attention._get_sdpa() -> typing.Callable | None
```

```python
nemo_automodel.components.attention.ffpa_attention._is_block_mask(
    mask: typing.Any
) -> bool
```

True for a flex `BlockMask` — the marker that a layer should run on FlexAttention.

```python
nemo_automodel.components.attention.ffpa_attention._warn_once(
    reason: str,
    message: str,
    target: str = 'sdpa'
) -> None
```

```python
nemo_automodel.components.attention.ffpa_attention.ffpa_attention_forward(
    module: torch.nn.Module,
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    attention_mask: torch.Tensor | None,
    dropout: float | int = 0.0,
    scaling: float | None = None,
    softcap: float | None = None,
    kwargs: typing.Any = {}
) -> tuple[torch.Tensor, torch.Tensor | None]
```

HF attention-interface forward backed by `torch.ops.ffpa_attn._fwd_cute`.

```python
nemo_automodel.components.attention.ffpa_attention.ffpa_mask(
    batch_size: int,
    q_length: int,
    kv_length: int,
    q_offset: int = 0,
    kv_offset: int = 0,
    mask_function: typing.Callable | None = None,
    attention_mask: torch.Tensor | None = None,
    dtype: torch.dtype = torch.float32,
    kwargs: typing.Any = {}
) -> 'torch.Tensor | BlockMask | None'
```

Mask factory for `ALL_MASK_ATTENTION_FUNCTIONS["ffpa"]`.

```python
nemo_automodel.components.attention.ffpa_attention.register_ffpa_attention() -> bool
```

Register `"ffpa"` in HF attention/mask registries. Idempotent.

```python
nemo_automodel.components.attention.ffpa_attention.setup_ffpa_backend(
    cp_size: int,
    has_packed_sequence: bool
) -> None
```

Validate FFPA preconditions, then register the HF `"ffpa"` backend.

The `"ffpa"` backend calls the FFPA op directly, bypassing both the ring-CP
SDPA swap and packed-sequence masking; for CP or packing use
`attn_implementation="sdpa"` + `text_config.cp_full_attn_backend="ffpa"` instead.

```python
nemo_automodel.components.attention.ffpa_attention._CUTEDSL_BACKEND: Any = None
```

```python
nemo_automodel.components.attention.ffpa_attention._EAGER_FN: Callable | None = None
```

```python
nemo_automodel.components.attention.ffpa_attention._FALLBACK_WARNED: set[str] = set()
```

```python
nemo_automodel.components.attention.ffpa_attention._FFPA_FN: Callable | None = None
```

```python
nemo_automodel.components.attention.ffpa_attention._FFPA_HEAD_DIM = 512
```

```python
nemo_automodel.components.attention.ffpa_attention._FFPA_LOW_LEVEL_READY: bool | None = None
```

```python
nemo_automodel.components.attention.ffpa_attention._FLEX_FN: Callable | None = None
```

```python
nemo_automodel.components.attention.ffpa_attention._REGISTERED: bool = False
```

```python
nemo_automodel.components.attention.ffpa_attention._SDPA_FN: Callable | None = None
```

```python
nemo_automodel.components.attention.ffpa_attention._VARLEN_WIN_NONE = -2 ** 31
```

```python
nemo_automodel.components.attention.ffpa_attention.__all__ = ['ffpa_attention_forward', 'ffpa_mask', 'register_ffpa_attention', 'setup_ffpa_b...
```

```python
nemo_automodel.components.attention.ffpa_attention.logger = logging.getLogger(__name__)
```