> 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.models.common.cudnn_sparse_attention

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

## Module Contents

### Classes

| Name                                                                                                             | Description                                                         |
| ---------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------- |
| [`_CudnnSparseAttention`](#nemo_automodel-components-models-common-cudnn_sparse_attention-_CudnnSparseAttention) | Pair FlashMLA forward with cuDNN backward for latent THD attention. |

### Functions

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

### Data

[`_FLASH_MLA_TOPK_ALIGNMENT`](#nemo_automodel-components-models-common-cudnn_sparse_attention-_FLASH_MLA_TOPK_ALIGNMENT)

[`_SUPPORTED_ATTENTION_HEAD_DIMS`](#nemo_automodel-components-models-common-cudnn_sparse_attention-_SUPPORTED_ATTENTION_HEAD_DIMS)

[`_VALUE_HEAD_DIM`](#nemo_automodel-components-models-common-cudnn_sparse_attention-_VALUE_HEAD_DIM)

[`__all__`](#nemo_automodel-components-models-common-cudnn_sparse_attention-__all__)

### API

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

**Bases:** `Function`

Pair FlashMLA forward with cuDNN backward for latent THD attention.

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

```python
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.

---

**`q`** `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]`.

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

```python
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:**

**`q`** `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

```python
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.

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

Raise when either optional sparse-attention runtime is unavailable.

```python
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.Tensor` — default: ()

Tensors with arbitrary shapes that must share one CUDA device.

---

**Returns:** `tuple[int, int]`

CUDA compute capability as `(major, minor)`.

```python
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:**

**`q`** `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 | None` — default: 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`** `bool` — default: False

Whether every query has a positive valid-prefix length.

---

**`valid_row_indices`** `torch.Tensor | None` — default: 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.

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

Return whether the cuDNN backward and FlashMLA forward runtimes import.

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

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

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

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