> This page is for version 0.5.0 · 26.06.
> For other versions, use one of these documentation indexes:
> - Nightly (default): https://docs.nvidia.com/nemo/automodel/nightly/llms.txt
> - Latest: https://docs.nvidia.com/nemo/automodel/latest/llms.txt
> - 0.5.0 · 26.06: https://docs.nvidia.com/nemo/automodel/v0.5/llms.txt
> - 0.4.0 · 26.04: https://docs.nvidia.com/nemo/automodel/v0.4/llms.txt

> 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)                 | Run FFPA CuTeDSL dense backward with the caller's globally merged output and LSE.  |
| [`_ffpa_dense_fwd`](#nemo_automodel-components-attention-ffpa_attention-_ffpa_dense_fwd)                 | Run FFPA CuTeDSL dense forward on 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)               | Run FFPA CuTeDSL varlen backward with the caller's globally merged output and LSE. |
| [`_ffpa_varlen_fwd`](#nemo_automodel-components-attention-ffpa_attention-_ffpa_varlen_fwd)               | Run FFPA CuTeDSL varlen forward on packed THD inputs.                              |
| [`_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)   | Route the Hugging Face attention interface through FFPA or an eligible fallback.   |
| [`ffpa_mask`](#nemo_automodel-components-attention-ffpa_attention-ffpa_mask)                             | Build the attention mask used by the registered FFPA backend.                      |
| [`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

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

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

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

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

[`_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)

[`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]
```

Run FFPA CuTeDSL dense backward with the caller's globally merged output and LSE.

**Parameters:**

**`grad_out`** `torch.Tensor`

Output-gradient tensor of shape \[batch, query\_heads, query\_sequence, value\_head\_dim] in SDPA
layout.

---

**`q`** `torch.Tensor`

Float16 or bfloat16 CUDA tensor of shape \[batch, query\_heads, query\_sequence, head\_dim] in SDPA
layout. QKV, `grad_out`, and `out` share dtype and device.

---

**`k`** `torch.Tensor`

Key tensor of shape \[batch, key\_value\_heads, key\_value\_sequence, head\_dim] in SDPA layout.

---

**`v`** `torch.Tensor`

Value tensor of shape \[batch, key\_value\_heads, key\_value\_sequence, value\_head\_dim] in SDPA layout.

---

**`out`** `torch.Tensor`

Caller's globally merged output tensor of shape
\[batch, query\_heads, query\_sequence, value\_head\_dim] in SDPA layout.

---

**`lse`** `torch.Tensor`

Caller's globally merged float32 LSE tensor of shape \[batch, query\_heads, query\_sequence] on the
same device as `q`.

---

**`scale`** `float`

Pre-softmax scaling factor applied to query-key scores.

---

**`causal`** `bool`

Whether to apply the causal mask.

---

**Returns:** `torch.Tensor`

A `(dq, dk, dv)` tuple in SDPA layout. Each gradient has the shape, dtype, and device of its

```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]
```

Run FFPA CuTeDSL dense forward on SDPA-layout tensors.

The returned per-chunk LSE supports the ring online-softmax merge.

**Parameters:**

**`q`** `torch.Tensor`

Float16 or bfloat16 CUDA tensor of shape \[batch, query\_heads, query\_sequence, head\_dim] in SDPA
layout. QKV share dtype, device, and head dimension.

---

**`k`** `torch.Tensor`

Key tensor of shape \[batch, key\_value\_heads, key\_value\_sequence, head\_dim] in SDPA layout.

---

**`v`** `torch.Tensor`

Value tensor of shape \[batch, key\_value\_heads, key\_value\_sequence, value\_head\_dim] in SDPA layout.

---

**`scale`** `float`

Pre-softmax scaling factor applied to query-key scores.

---

**`causal`** `bool`

Whether to apply the causal mask.

---

**Returns:** `torch.Tensor`

A tuple containing an output tensor of shape \[batch, query\_heads, query\_sequence, value\_head\_dim] with

```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]
```

Run FFPA CuTeDSL varlen backward with the caller's globally merged output and LSE.

**Parameters:**

**`grad_out_pack`** `torch.Tensor`

Output-gradient tensor of shape \[total\_query\_tokens, query\_heads, value\_head\_dim] in
packed THD layout.

---

**`q_pack`** `torch.Tensor`

Float16 or bfloat16 CUDA tensor of shape \[total\_query\_tokens, query\_heads, head\_dim] in packed
THD layout. QKV, `grad_out_pack`, and `out_pack` share dtype and device and are made contiguous
before dispatch.

---

**`k_pack`** `torch.Tensor`

Key tensor of shape \[total\_key\_value\_tokens, key\_value\_heads, head\_dim] in packed THD layout.

---

**`v_pack`** `torch.Tensor`

Value tensor of shape \[total\_key\_value\_tokens, key\_value\_heads, value\_head\_dim] in packed THD
layout.

---

**`out_pack`** `torch.Tensor`

Caller's globally merged output tensor of shape
\[total\_query\_tokens, query\_heads, value\_head\_dim] in packed THD layout.

---

**`lse_pack`** `torch.Tensor`

Caller's globally merged float32 LSE tensor of shape \[query\_heads, total\_query\_tokens] on the
same device as `q_pack`.

---

**`cu_q`** `torch.Tensor`

Int32 CUDA tensor of shape \[batch + 1] on the same device as `q_pack`, containing cumulative
query-token offsets.

---

**`cu_k`** `torch.Tensor`

Int32 CUDA tensor of shape \[batch + 1] on the same device as `q_pack`, containing cumulative
key/value-token offsets.

---

**`max_q`** `int`

Maximum query sequence length represented by `cu_q`.

---

**`max_k`** `int`

Maximum key/value sequence length represented by `cu_k`.

---

**`scale`** `float`

Pre-softmax scaling factor applied to query-key scores.

---

**`causal`** `bool`

Whether to apply the tail-aligned causal mask within each packed sequence.

---

**Returns:** `torch.Tensor`

A `(dq, dk, dv)` tuple in packed THD layout. Each gradient has the shape, dtype, and device of its

```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]
```

Run FFPA CuTeDSL varlen forward on packed THD inputs.

Windowing, softcap, and packed GQA are disabled at this seam.

**Parameters:**

**`q_pack`** `torch.Tensor`

Float16 or bfloat16 CUDA tensor of shape \[total\_query\_tokens, query\_heads, head\_dim] in packed
THD layout. QKV share dtype and device and are made contiguous before dispatch.

---

**`k_pack`** `torch.Tensor`

Key tensor of shape \[total\_key\_value\_tokens, key\_value\_heads, head\_dim] in packed THD layout.

---

**`v_pack`** `torch.Tensor`

Value tensor of shape \[total\_key\_value\_tokens, key\_value\_heads, value\_head\_dim] in packed THD
layout.

---

**`cu_q`** `torch.Tensor`

Int32 CUDA tensor of shape \[batch + 1] on the same device as `q_pack`, containing cumulative
query-token offsets from 0 through `total_query_tokens`.

---

**`cu_k`** `torch.Tensor`

Int32 CUDA tensor of shape \[batch + 1] on the same device as `q_pack`, containing cumulative
key/value-token offsets from 0 through `total_key_value_tokens`.

---

**`max_q`** `int`

Maximum query sequence length represented by `cu_q`.

---

**`max_k`** `int`

Maximum key/value sequence length represented by `cu_k`.

---

**`scale`** `float`

Pre-softmax scaling factor applied to query-key scores.

---

**`causal`** `bool`

Whether to apply the tail-aligned causal mask within each packed sequence.

---

**Returns:** `torch.Tensor`

A tuple containing an output tensor of shape \[total\_query\_tokens, query\_heads, value\_head\_dim] with the

```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, typing.Any] | None
```

```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]
```

Route the Hugging Face attention interface through FFPA or an eligible fallback.

**Parameters:**

**`module`** `nn.Module`

Attention module providing the training state and `head_dim` used for FFPA eligibility.

---

**`query`** `torch.Tensor`

Query tensor of shape \[batch, query\_heads, query\_sequence, head\_dim] in SDPA layout.

---

**`key`** `torch.Tensor`

Key tensor of shape \[batch, key\_value\_heads, key\_value\_sequence, head\_dim] in SDPA layout.

---

**`value`** `torch.Tensor`

Value tensor of shape \[batch, key\_value\_heads, key\_value\_sequence, value\_head\_dim] in SDPA layout.
QKV share dtype and device. FFPA requires float16 or bfloat16 CUDA tensors with head dimension 512.

---

**`attention_mask`** `torch.Tensor | None`

Optional boolean padding tensor of shape \[batch, key\_value\_sequence] or additive tensor
of shape \[batch, 1, query\_sequence, key\_value\_sequence]. A Flex `BlockMask` produced by
`ffpa_mask` may also be supplied for the FlexAttention route.

---

**`dropout`** `float | int` — default: 0.0

Dropout probability. Nonzero training dropout makes the FFPA route ineligible.

---

**`scaling`** `float | None` — default: None

Optional pre-softmax scaling factor applied to query-key scores.

---

**`softcap`** `float | None` — default: None

Optional score cap; when present, the eager fallback owns execution.

---

**`**kwargs`** `Any` — default: \{}

Additional Hugging Face attention arguments forwarded to the selected backend.

---

**Returns:** `torch.Tensor`

A tuple containing an output tensor of shape \[batch, query\_sequence, query\_heads, value\_head\_dim] with

```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'
```

Build the attention mask used by the registered FFPA backend.

**Parameters:**

**`batch_size`** `int`

Number of sequences in the batch.

---

**`q_length`** `int`

Query sequence length.

---

**`kv_length`** `int`

Key/value sequence length.

---

**`q_offset`** `int` — default: 0

Global offset of the first query token.

---

**`kv_offset`** `int` — default: 0

Global offset of the first key/value token.

---

**`mask_function`** `Callable | None` — default: None

Optional index-based mask predicate.

---

**`attention_mask`** `torch.Tensor | None` — default: None

Optional padding mask tensor of shape `[batch, kv_sequence]`
or a prebuilt mask tensor of shape `[batch, 1, query_sequence, kv_sequence]`.

---

**`dtype`** `torch.dtype` — default: torch.float32

Floating-point dtype used when a backend needs an additive mask.

---

**`**kwargs`** `Any` — default: \{}

Mask-factory metadata forwarded by Transformers.

---

**Returns:** `'torch.Tensor | BlockMask | None'`

A boolean tensor of shape `[batch, kv_sequence]` in padding layout or

```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._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_HEAD_DIM = 512
```

```python
nemo_automodel.components.attention.ffpa_attention._FFPA_HIGH_LEVEL: tuple[Callable, Any] | None = None
```

```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.logger = logging.getLogger(__name__)
```