nemo_automodel.components.models.common.cudnn_sparse_attention

Shared FlashMLA-forward/cuDNN-backward sparse latent attention.

Module Contents

Classes

Name	Description
`_CudnnSparseAttention`	Pair FlashMLA forward with cuDNN backward for latent THD attention.

Functions

Name	Description
`_compact_and_sort_indices`	Canonicalize sparse indices into an ascending valid prefix.
`_pad_attention_heads`	Pad query and attention-sink head axes for FlashMLA.
`_padded_head_count`	Return the FlashMLA-supported query-head count for one SM generation.
`_require_available`	Raise when either optional sparse-attention runtime is unavailable.
`_require_cuda_tensors`	Validate that arbitrary-layout input tensors share one SM90+ CUDA device.
`cudnn_sparse_attention`	Run sparse latent attention with FlashMLA forward and cuDNN backward.
`is_cudnn_sparse_attention_available`	Return whether the cuDNN backward and FlashMLA forward runtimes import.

Data

_FLASH_MLA_TOPK_ALIGNMENT

_SUPPORTED_ATTENTION_HEAD_DIMS

_VALUE_HEAD_DIM

__all__

API

class nemo_automodel.components.models.common.cudnn_sparse_attention._CudnnSparseAttention()

Bases: Function

Pair FlashMLA forward with cuDNN backward for latent THD attention.

nemo_automodel.components.models.common.cudnn_sparse_attention._CudnnSparseAttention.backward(
    ctx: typing.Any,
    grad_output: torch.Tensor
) -> tuple[torch.Tensor | None, ...]

staticmethod

Map output gradients to query and gathered latent-KV layouts.

Parameters:

ctx

Any

Autograd context populated by :meth:forward.

grad_output

torch.Tensor

CUDA BF16 output gradient of shape [query_tokens, heads, 512].

Returns: torch.Tensor | None

Gradients for the eight forward inputs: query tensor of shape

nemo_automodel.components.models.common.cudnn_sparse_attention._CudnnSparseAttention.forward(
    ctx: typing.Any,
    q: torch.Tensor,
    kv_latent: torch.Tensor,
    topk_indices: torch.Tensor,
    softmax_scale: float,
    padded_heads: int,
    topk_length: torch.Tensor | None,
    all_rows_nonempty: bool,
    valid_row_indices: torch.Tensor | None
) -> torch.Tensor

staticmethod

Run FlashMLA forward and save tensors required by cuDNN backward.

Parameters:

ctx

Any

Autograd context used to save forward tensors and scalar metadata.

torch.Tensor

CUDA BF16 query tensor of shape [query_tokens, heads, head_dim].

kv_latent

torch.Tensor

CUDA BF16 latent K/V tensor of shape [key_tokens, 1, head_dim].

topk_indices

torch.Tensor

CUDA int32 tensor of shape [query_tokens, 1, sparse_width] containing global K/V coordinates and invalid entries marked -1.

softmax_scale

float

Scale applied to query-key scores.

padded_heads

int

FlashMLA-compatible padded head count.

topk_length

torch.Tensor | None

Optional int32 valid-prefix lengths of shape [query_tokens].

all_rows_nonempty

bool

Whether every query has a positive valid prefix.

valid_row_indices

torch.Tensor | None

Optional int64 indices of nonempty queries with shape [valid_query_tokens].

Returns: torch.Tensor

CUDA BF16 latent values of shape [query_tokens, heads, 512].

nemo_automodel.components.models.common.cudnn_sparse_attention._compact_and_sort_indices(
    indices: torch.Tensor,
    key_count: int
) -> tuple[torch.Tensor, torch.Tensor]

Canonicalize sparse indices into an ascending valid prefix.

Parameters:

indices

torch.Tensor

Integer tensor of shape [query_tokens, sparse_width] with global K/V coordinates and negative invalid entries.

key_count

int

Number of rows in the flattened K/V tensor.

Returns: torch.Tensor

A contiguous int32 index tensor of shape [query_tokens, sparse_width]

nemo_automodel.components.models.common.cudnn_sparse_attention._pad_attention_heads(
    q: torch.Tensor,
    attn_sink: torch.Tensor,
    padded_heads: int
) -> tuple[torch.Tensor, torch.Tensor]

Pad query and attention-sink head axes for FlashMLA.

Parameters:

torch.Tensor

Query tensor of shape [query_tokens, heads, head_dim].

attn_sink

torch.Tensor

FP32 attention-sink tensor of shape [heads].

padded_heads

int

FlashMLA-compatible output head count.

Returns: torch.Tensor

Query tensor of shape [query_tokens, padded_heads, head_dim] and

nemo_automodel.components.models.common.cudnn_sparse_attention._padded_head_count(
    num_heads: int,
    major: int
) -> int

Return the FlashMLA-supported query-head count for one SM generation.

nemo_automodel.components.models.common.cudnn_sparse_attention._require_available() -> None

Raise when either optional sparse-attention runtime is unavailable.

nemo_automodel.components.models.common.cudnn_sparse_attention._require_cuda_tensors(
    operation: str,
    tensors: torch.Tensor = ()
) -> tuple[int, int]

Validate that arbitrary-layout input tensors share one SM90+ CUDA device.

Parameters:

operation

str

Name included in validation errors.

*tensors

torch.TensorDefaults to ()

Tensors with arbitrary shapes that must share one CUDA device.

Returns: tuple[int, int]

CUDA compute capability as (major, minor).

nemo_automodel.components.models.common.cudnn_sparse_attention.cudnn_sparse_attention(
    q: torch.Tensor,
    kv_latent: torch.Tensor,
    topk_indices: torch.Tensor,
    softmax_scale: float,
    topk_length: torch.Tensor | None = None,
    all_rows_nonempty: bool = False,
    valid_row_indices: torch.Tensor | None = None
) -> torch.Tensor

Run sparse latent attention with FlashMLA forward and cuDNN backward.

Parameters:

torch.Tensor

Contiguous CUDA BF16 query tensor of shape [query_tokens, heads, head_dim], where head_dim is 512 or 576.

kv_latent

torch.Tensor

Contiguous CUDA BF16 latent K/V tensor of shape [key_tokens, 1, head_dim].

topk_indices

torch.Tensor

Contiguous CUDA int32 tensor of shape [query_tokens, 1, sparse_width] with global K/V coordinates and invalid entries marked -1.

softmax_scale

float

Scale forwarded unchanged to FlashMLA and cuDNN backward.

topk_length

torch.Tensor | NoneDefaults to None

Optional contiguous CUDA int32 valid-prefix lengths of shape [query_tokens]. When supplied, topk_indices must already contain a compact, ascending valid prefix.

all_rows_nonempty

boolDefaults to False

Whether every query has a positive valid-prefix length.

valid_row_indices

torch.Tensor | NoneDefaults to None

Optional contiguous CUDA int64 indices of nonempty queries with shape [valid_query_tokens].

Returns: torch.Tensor

Contiguous CUDA BF16 latent output tensor of shape

Raises:

RuntimeError: If optional kernels, CUDA, or SM90+ are unavailable.
TypeError: If compute tensors are not BF16 or indices are not int32.
ValueError: If tensor layouts, dimensions, sparse width, or scale are invalid.

nemo_automodel.components.models.common.cudnn_sparse_attention.is_cudnn_sparse_attention_available() -> bool

Return whether the cuDNN backward and FlashMLA forward runtimes import.

nemo_automodel.components.models.common.cudnn_sparse_attention._FLASH_MLA_TOPK_ALIGNMENT = 512

nemo_automodel.components.models.common.cudnn_sparse_attention._SUPPORTED_ATTENTION_HEAD_DIMS = (512, 576)

nemo_automodel.components.models.common.cudnn_sparse_attention._VALUE_HEAD_DIM = 512

nemo_automodel.components.models.common.cudnn_sparse_attention.__all__ = ['cudnn_sparse_attention', 'is_cudnn_sparse_attention_available']