> 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.glm_moe_dsa.kernels.cudnn_dsa

cuDNN and FlashMLA kernels for the split GLM-5.2 DSA path.

## Module Contents

### Classes

| Name                                                                                                               | Description                                                         |
| ------------------------------------------------------------------------------------------------------------------ | ------------------------------------------------------------------- |
| [`CudnnDsaPackedMetadata`](#nemo_automodel-components-models-glm_moe_dsa-kernels-cudnn_dsa-CudnnDsaPackedMetadata) | Reusable THD metadata for local-query/global-key cuDNN DSA.         |
| [`_CudnnSparseAttention`](#nemo_automodel-components-models-glm_moe_dsa-kernels-cudnn_dsa-_CudnnSparseAttention)   | Pair FlashMLA forward with cuDNN backward for latent THD attention. |

### Functions

| Name                                                                                                                                     | Description                                                                |
| ---------------------------------------------------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------- |
| [`_compact_and_sort_indices`](#nemo_automodel-components-models-glm_moe_dsa-kernels-cudnn_dsa-_compact_and_sort_indices)                 | Canonicalize global indices `[T, K]` and return valid lengths `[T]`.       |
| [`_pad_attention_heads`](#nemo_automodel-components-models-glm_moe_dsa-kernels-cudnn_dsa-_pad_attention_heads)                           | Pad query `[T, H, 576]` and sink `[H]` to `padded_heads`.                  |
| [`_padded_head_count`](#nemo_automodel-components-models-glm_moe_dsa-kernels-cudnn_dsa-_padded_head_count)                               | Return the FlashMLA-supported head count for one SM generation.            |
| [`_require_available`](#nemo_automodel-components-models-glm_moe_dsa-kernels-cudnn_dsa-_require_available)                               | Raise when either optional runtime required by the split kernel is absent. |
| [`_require_cuda_tensors`](#nemo_automodel-components-models-glm_moe_dsa-kernels-cudnn_dsa-_require_cuda_tensors)                         | Validate that arbitrary-layout input tensors share one SM90+ CUDA device.  |
| [`_topk_wrapper_chunked`](#nemo_automodel-components-models-glm_moe_dsa-kernels-cudnn_dsa-_topk_wrapper_chunked)                         | Select top-k for FP32 scores `[T, S_max]` using causal lengths `[T]`.      |
| [`_unpack_packed_metadata`](#nemo_automodel-components-models-glm_moe_dsa-kernels-cudnn_dsa-_unpack_packed_metadata)                     | Validate reusable packed metadata without a CUDA-to-host synchronization.  |
| [`_validate_topk`](#nemo_automodel-components-models-glm_moe_dsa-kernels-cudnn_dsa-_validate_topk)                                       | Validate GLM-5.2's fixed sparse-selection width.                           |
| [`cudnn_indexer_topk`](#nemo_automodel-components-models-glm_moe_dsa-kernels-cudnn_dsa-cudnn_indexer_topk)                               | Compute GLM-5.2 packed-THD indexer top-k with cuDNN Frontend.              |
| [`cudnn_sparse_attention`](#nemo_automodel-components-models-glm_moe_dsa-kernels-cudnn_dsa-cudnn_sparse_attention)                       | Run split GLM-5.2 sparse MLA with FlashMLA forward and cuDNN backward.     |
| [`is_cudnn_dsa_available`](#nemo_automodel-components-models-glm_moe_dsa-kernels-cudnn_dsa-is_cudnn_dsa_available)                       | Return whether both optional libraries required by cuDNN DSA import.       |
| [`prepare_cudnn_dsa_packed_metadata`](#nemo_automodel-components-models-glm_moe_dsa-kernels-cudnn_dsa-prepare_cudnn_dsa_packed_metadata) | Build segmented THD metadata for local queries and gathered global keys.   |

### Data

[`_ATTENTION_HEAD_DIM`](#nemo_automodel-components-models-glm_moe_dsa-kernels-cudnn_dsa-_ATTENTION_HEAD_DIM)

[`_FLASH_MLA_TOPK_ALIGNMENT`](#nemo_automodel-components-models-glm_moe_dsa-kernels-cudnn_dsa-_FLASH_MLA_TOPK_ALIGNMENT)

[`_INDEX_HEAD_DIM`](#nemo_automodel-components-models-glm_moe_dsa-kernels-cudnn_dsa-_INDEX_HEAD_DIM)

[`_MAX_TOPK`](#nemo_automodel-components-models-glm_moe_dsa-kernels-cudnn_dsa-_MAX_TOPK)

[`_TOPK_ROW_ALIGNMENT`](#nemo_automodel-components-models-glm_moe_dsa-kernels-cudnn_dsa-_TOPK_ROW_ALIGNMENT)

[`_TOPK_SCRATCH_INT32_FACTOR`](#nemo_automodel-components-models-glm_moe_dsa-kernels-cudnn_dsa-_TOPK_SCRATCH_INT32_FACTOR)

[`_TOPK_SCRATCH_LIMIT_BYTES`](#nemo_automodel-components-models-glm_moe_dsa-kernels-cudnn_dsa-_TOPK_SCRATCH_LIMIT_BYTES)

[`_VALUE_HEAD_DIM`](#nemo_automodel-components-models-glm_moe_dsa-kernels-cudnn_dsa-_VALUE_HEAD_DIM)

[`__all__`](#nemo_automodel-components-models-glm_moe_dsa-kernels-cudnn_dsa-__all__)

### API

```python
class nemo_automodel.components.models.glm_moe_dsa.kernels.cudnn_dsa.CudnnDsaPackedMetadata(
    starts: torch.Tensor,
    causal_lengths: torch.Tensor,
    query_valid: torch.Tensor,
    valid_row_indices: torch.Tensor | None,
    segment_cu_q: torch.Tensor,
    segment_cu_k: torch.Tensor,
    q_causal_offsets: torch.Tensor,
    key_source_indices: torch.Tensor | None,
    max_seqlen_q: int,
    max_seqlen_k: int,
    total_key_tokens: int,
    all_rows_nonempty: bool
)
```

Dataclass

Reusable THD metadata for local-query/global-key cuDNN DSA.

**`all_rows_nonempty`** `bool`

---

**`causal_lengths`** `Tensor`

---

**`key_source_indices`** `Tensor | None`

---

**`max_seqlen_k`** `int`

---

**`max_seqlen_q`** `int`

---

**`q_causal_offsets`** `Tensor`

---

**`query_valid`** `Tensor`

---

**`segment_cu_k`** `Tensor`

---

**`segment_cu_q`** `Tensor`

---

**`starts`** `Tensor`

---

**`total_key_tokens`** `int`

---

**`valid_row_indices`** `Tensor | None`

---

```python
class nemo_automodel.components.models.glm_moe_dsa.kernels.cudnn_dsa._CudnnSparseAttention()
```

**Bases:** `Function`

Pair FlashMLA forward with cuDNN backward for latent THD attention.

```python
nemo_automodel.components.models.glm_moe_dsa.kernels.cudnn_dsa._CudnnSparseAttention.backward(
    ctx,
    grad_output: torch.Tensor
)
```

staticmethod

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

**Parameters:**

**`ctx`**

Autograd context populated by :meth:`forward`.

---

**`grad_output`** `torch.Tensor`

CUDA BF16 output gradient of shape `[T_q, H, 512]`.

---

**Returns:**

Gradients for the eight forward inputs: query `[T_q, H, 576]`,

```python
nemo_automodel.components.models.glm_moe_dsa.kernels.cudnn_dsa._CudnnSparseAttention.forward(
    ctx,
    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`**

Autograd context used to save forward tensors and scalar metadata.

---

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

CUDA BF16 query tensor of shape `[T_q, H, 576]`.

---

**`kv_latent`** `torch.Tensor`

CUDA BF16 gathered K/V tensor of shape `[T_k, 1, 576]`.

---

**`topk_indices`** `torch.Tensor`

CUDA int32 tensor of shape `[T_q, 1, K]` containing
global padded-storage K/V coordinates and a `-1` suffix.

---

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

---

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

---

**Returns:** `torch.Tensor`

CUDA BF16 latent values of shape `[T_q, H, 512]`.

```python
nemo_automodel.components.models.glm_moe_dsa.kernels.cudnn_dsa._compact_and_sort_indices(
    indices: torch.Tensor,
    key_count: int
) -> tuple[torch.Tensor, torch.Tensor]
```

Canonicalize global indices `[T, K]` and return valid lengths `[T]`.

```python
nemo_automodel.components.models.glm_moe_dsa.kernels.cudnn_dsa._pad_attention_heads(
    q: torch.Tensor,
    attn_sink: torch.Tensor,
    padded_heads: int
) -> tuple[torch.Tensor, torch.Tensor]
```

Pad query `[T, H, 576]` and sink `[H]` to `padded_heads`.

```python
nemo_automodel.components.models.glm_moe_dsa.kernels.cudnn_dsa._padded_head_count(
    num_heads: int,
    major: int
) -> int
```

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

```python
nemo_automodel.components.models.glm_moe_dsa.kernels.cudnn_dsa._require_available() -> None
```

Raise when either optional runtime required by the split kernel is absent.

```python
nemo_automodel.components.models.glm_moe_dsa.kernels.cudnn_dsa._require_cuda_tensors(
    operation: str,
    tensors: torch.Tensor = ()
) -> tuple[int, int]
```

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

```python
nemo_automodel.components.models.glm_moe_dsa.kernels.cudnn_dsa._topk_wrapper_chunked(
    scores: torch.Tensor,
    seq_lens: torch.Tensor,
    topk: int
) -> torch.Tensor
```

Select top-k for FP32 scores `[T, S_max]` using causal lengths `[T]`.

```python
nemo_automodel.components.models.glm_moe_dsa.kernels.cudnn_dsa._unpack_packed_metadata(
    packed_metadata: nemo_automodel.components.models.glm_moe_dsa.kernels.cudnn_dsa.CudnnDsaPackedMetadata,
    total_query_tokens: int,
    total_key_tokens: int,
    device: torch.device
) -> nemo_automodel.components.models.glm_moe_dsa.kernels.cudnn_dsa.CudnnDsaPackedMetadata
```

Validate reusable packed metadata without a CUDA-to-host synchronization.

**Parameters:**

**`packed_metadata`** `CudnnDsaPackedMetadata`

Metadata whose per-query fields have shape `[T_q]` and
whose key-source indices address the gathered padded-storage K/V tensor.

---

**`total_query_tokens`** `int`

Expected local query row count `T_q`.

---

**`total_key_tokens`** `int`

Expected gathered padded-storage K/V row count `T_k`.

---

**`device`** `torch.device`

CUDA device shared by the metadata and kernel inputs.

---

**Returns:** `CudnnDsaPackedMetadata`

The validated metadata object, unchanged.

```python
nemo_automodel.components.models.glm_moe_dsa.kernels.cudnn_dsa._validate_topk(
    index_topk: int
) -> None
```

Validate GLM-5.2's fixed sparse-selection width.

```python
nemo_automodel.components.models.glm_moe_dsa.kernels.cudnn_dsa.cudnn_indexer_topk(
    index_q: torch.Tensor,
    index_k: torch.Tensor,
    head_weights: torch.Tensor,
    cu_seqlens: torch.Tensor,
    index_topk: int,
    query_indices: torch.Tensor | None = None,
    cu_seqlens_padded: torch.Tensor | None = None,
    packed_metadata: nemo_automodel.components.models.glm_moe_dsa.kernels.cudnn_dsa.CudnnDsaPackedMetadata | None = None
) -> torch.Tensor
```

Compute GLM-5.2 packed-THD indexer top-k with cuDNN Frontend.

**Parameters:**

**`index_q`** `torch.Tensor`

Rank-local indexer query in THD layout, BF16 `[T_q, H_index, 128]`.

---

**`index_k`** `torch.Tensor`

Gathered global indexer key in TD layout, BF16 `[T_k, 128]`.

---

**`head_weights`** `torch.Tensor`

Already-scaled per-head weights in TH layout, FP32 or BF16
`[T, H_index]`. The caller owns both `H_index**-0.5` and
`128**-0.5` scaling; this function does not rescale them.

---

**`cu_seqlens`** `torch.Tensor`

Compact packed-sequence offsets, CUDA int32
`[num_sequences + 1]`.

---

**`index_topk`** `int`

Fixed output width `K` in `[1, 2048]`.

---

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

Optional contiguous global padded coordinates for local query
rows, `[T_q]`. Absent means CP=1 identity coordinates.

---

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

Optional global padded packed-layout offsets.

---

**`packed_metadata`** `CudnnDsaPackedMetadata | None` — default: None

Optional metadata object returned by
:func:`prepare_cudnn_dsa_packed_metadata`. Supplying it avoids rebuilding
and synchronizing the same metadata in every full-indexer layer.

---

**Returns:** `torch.Tensor`

CUDA int32 top-k indices in global padded-storage THD coordinates,

**Raises:**

* `RuntimeError`: If optional kernels, CUDA, or SM90+ are unavailable.
* `TypeError`: If tensor dtypes or `index_topk` are invalid.
* `ValueError`: If tensor shapes or compact packed metadata are invalid.

```python
nemo_automodel.components.models.glm_moe_dsa.kernels.cudnn_dsa.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 split GLM-5.2 sparse MLA with FlashMLA forward and cuDNN backward.

**Parameters:**

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

Absorbed MLA query, CUDA BF16 THD `[T_q, H, 576]`. The final
dimension is `kv_lora_rank + qk_rope_head_dim` (`512 + 64`).

---

**`kv_latent`** `torch.Tensor`

Shared latent key/value, CUDA BF16 THD `[T_kv, 1, 576]`.

---

**`topk_indices`** `torch.Tensor`

Global padded-storage K/V indices, CUDA int32
`[T_q, 1, K]` with `-1` for invalid slots.

---

**`softmax_scale`** `float`

Already-computed MLA attention scale. It is forwarded
unchanged to both FlashMLA and cuDNN backward.

---

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

Optional int32 valid-prefix lengths `[T_q]` prepared once
from packed causal metadata. When supplied, `topk_indices` must already
contain the indexer's canonical compact, ascending prefix.

---

**`all_rows_nonempty`** `bool` — default: False

Whether every query has a positive valid-prefix length.
The model supplies this cached metadata flag to keep unpadded inputs on
the allocation-free backward path.

---

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

Optional cached int64 row indices whose valid-prefix length
is positive. The model supplies these once per stage for padded inputs so
every attention layer can compact without rescanning CUDA metadata.

---

**Returns:** `torch.Tensor`

Latent sparse-attention output, CUDA BF16 `[T_q, H, 512]`. The caller

**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, top-k, or scale are invalid.

```python
nemo_automodel.components.models.glm_moe_dsa.kernels.cudnn_dsa.is_cudnn_dsa_available() -> bool
```

Return whether both optional libraries required by cuDNN DSA import.

```python
nemo_automodel.components.models.glm_moe_dsa.kernels.cudnn_dsa.prepare_cudnn_dsa_packed_metadata(
    cu_seqlens: torch.Tensor,
    total_key_tokens: int,
    max_seqlen: int | torch.Tensor | None = None,
    query_indices: torch.Tensor | None = None,
    cu_seqlens_padded: torch.Tensor | None = None,
    padding_mask: torch.Tensor | None = None
) -> nemo_automodel.components.models.glm_moe_dsa.kernels.cudnn_dsa.CudnnDsaPackedMetadata
```

Build segmented THD metadata for local queries and gathered global keys.

This validation intentionally performs one device-to-host synchronization. The model
prepares the object once per pipeline stage and reuses it across every indexer and
shared-attention layer.

**Parameters:**

**`cu_seqlens`** `torch.Tensor`

Cumulative compact real-token lengths, int32 tensor of shape
`[sequences + 1]`.

---

**`total_key_tokens`** `int`

Number of tokens in the gathered, padded THD key tensor.

---

**`max_seqlen`** `int | torch.Tensor | None` — default: None

Optional precomputed maximum sequence length as a Python integer or
scalar integer tensor. Its value is checked against `cu_seqlens`.

---

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

Optional contiguous integer tensor of shape `[T_q]` with
global padded-storage coordinates for the local query rows. When absent,
queries cover all global key tokens (CP=1).

---

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

Optional cumulative padded-storage boundaries, int32 tensor
of shape `[sequences + 1]`. When absent, `cu_seqlens` also defines
storage coordinates.

---

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

Optional local boolean padding mask of shape `[T_q]`. This
remains authoritative when THD preprocessing has absorbed trailing pack
padding into `cu_seqlens`.

---

**Returns:** `CudnnDsaPackedMetadata`

Metadata with per-query global padded-storage starts and real causal lengths,

```python
nemo_automodel.components.models.glm_moe_dsa.kernels.cudnn_dsa._ATTENTION_HEAD_DIM = 576
```

```python
nemo_automodel.components.models.glm_moe_dsa.kernels.cudnn_dsa._FLASH_MLA_TOPK_ALIGNMENT = 512
```

```python
nemo_automodel.components.models.glm_moe_dsa.kernels.cudnn_dsa._INDEX_HEAD_DIM = 128
```

```python
nemo_automodel.components.models.glm_moe_dsa.kernels.cudnn_dsa._MAX_TOPK = 2048
```

```python
nemo_automodel.components.models.glm_moe_dsa.kernels.cudnn_dsa._TOPK_ROW_ALIGNMENT = 512
```

```python
nemo_automodel.components.models.glm_moe_dsa.kernels.cudnn_dsa._TOPK_SCRATCH_INT32_FACTOR = 2
```

```python
nemo_automodel.components.models.glm_moe_dsa.kernels.cudnn_dsa._TOPK_SCRATCH_LIMIT_BYTES = 2 * 1024 * 1024 * 1024
```

```python
nemo_automodel.components.models.glm_moe_dsa.kernels.cudnn_dsa._VALUE_HEAD_DIM = 512
```

```python
nemo_automodel.components.models.glm_moe_dsa.kernels.cudnn_dsa.__all__ = ['CudnnDsaPackedMetadata', 'cudnn_indexer_topk', 'cudnn_sparse_attention', 'is_c...
```