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`	Run FFPA CuTeDSL dense backward with the caller’s globally merged output and LSE.
`_ffpa_dense_fwd`	Run FFPA CuTeDSL dense forward on SDPA-layout tensors.
`_ffpa_low_level_ready`	-
`_ffpa_varlen_bwd`	Run FFPA CuTeDSL varlen backward with the caller’s globally merged output and LSE.
`_ffpa_varlen_fwd`	Run FFPA CuTeDSL varlen forward on packed THD inputs.
`_ffpa_varlen_ready`	Whether the FFPA CuTeDSL varlen ops are importable and registered.
`_get_eager`	-
`_get_ffpa_high_level`	-
`_get_flex`	-
`_get_sdpa`	-
`_is_block_mask`	True for a flex `BlockMask` — the marker that a layer should run on FlexAttention.
`_warn_once`	-
`ffpa_attention_forward`	Route the Hugging Face attention interface through FFPA or an eligible fallback.
`ffpa_mask`	Build the attention mask used by the registered FFPA backend.
`register_ffpa_attention`	Register `"ffpa"` in HF attention/mask registries. Idempotent.
`setup_ffpa_backend`	Validate FFPA preconditions, then register the HF `"ffpa"` backend.

Data

_FFPA_LOW_LEVEL_READY

API

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.

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.

torch.Tensor

Key tensor of shape [batch, key_value_heads, key_value_sequence, head_dim] in SDPA layout.

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

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:

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.

torch.Tensor

Key tensor of shape [batch, key_value_heads, key_value_sequence, head_dim] in SDPA layout.

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

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

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

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

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

Whether the FFPA CuTeDSL varlen ops are importable and registered.

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

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

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

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

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.

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

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 | intDefaults to 0.0

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

scaling

float | NoneDefaults to None

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

softcap

float | NoneDefaults to None

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

**kwargs

AnyDefaults to {}

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

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

intDefaults to 0

Global offset of the first query token.

kv_offset

intDefaults to 0

Global offset of the first key/value token.

mask_function

Callable | NoneDefaults to None

Optional index-based mask predicate.

attention_mask

torch.Tensor | NoneDefaults to 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.dtypeDefaults to torch.float32

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

**kwargs

AnyDefaults to {}

Mask-factory metadata forwarded by Transformers.

Returns: 'torch.Tensor | BlockMask | None'

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

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

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.

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

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

nemo_automodel.components.attention.ffpa_attention._FFPA_HEAD_DIM = 512

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

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

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

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

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

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

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

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