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`	FFPA CuTeDSL dense backward using the caller-supplied out/lse.
`_ffpa_dense_fwd`	FFPA CuTeDSL dense forward on `[B, H, N, D]` SDPA-layout tensors.
`_ffpa_low_level_ready`	-
`_ffpa_varlen_bwd`	FFPA CuTeDSL varlen backward using the caller-supplied out/lse. Returns `(dq, dk, dv)`.
`_ffpa_varlen_fwd`	FFPA CuTeDSL varlen forward on packed THD inputs. Returns `(out[T,Hq,D], lse[Hq,T])`.
`_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`	HF attention-interface forward backed by `torch.ops.ffpa_attn._fwd_cute`.
`ffpa_mask`	Mask factory for `ALL_MASK_ATTENTION_FUNCTIONS["ffpa"]`.
`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]

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.

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.

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]

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.

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.

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

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]

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

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"].

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._CUTEDSL_BACKEND: Any = None

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_FN: Callable | None = None

nemo_automodel.components.attention.ffpa_attention._FFPA_HEAD_DIM = 512

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__)