> 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.optimized_kernels

Optional GLM-5.2 DSA sparse-kernel dispatch.

The dense torch path in `layers.py` (`GlmMoeDsaIndexer` / `GlmMoeDsaMLA`)
is the numerical reference and the default. This module provides optional
TileLang- and cuDNN-backed sparse paths, selected by `backend.attn`:

* the fused **lighting indexer** (logits + top-k), and
* the gather-top-k **sparse MLA** attention (reads only the selected KV; no
  `[T, T]` mask is materialized).

The TileLang kernels are vendored from THUDM's slime GLM-5.2 plugin under
`nemo_automodel.components.models.glm_moe_dsa.kernels` (see that package's
`__init__` for attribution); the cuDNN adapter lives in the same model-owned
kernel package. Optional entry points are imported with `safe_import_from` so
environments without their dependencies can still import the model and use
another attention backend.

Mirrors the structure of `deepseek_v4/optimized_kernels.py`.

## Module Contents

### Functions

| Name                                                                                                                                         | Description                                                                   |
| -------------------------------------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------- |
| [`_all_cuda`](#nemo_automodel-components-models-glm_moe_dsa-optimized_kernels-_all_cuda)                                                     | -                                                                             |
| [`_generate_padded_varlen_mask_params`](#nemo_automodel-components-models-glm_moe_dsa-optimized_kernels-_generate_padded_varlen_mask_params) | Build per-query key windows for padded THD layout.                            |
| [`cudnn_indexer_topk`](#nemo_automodel-components-models-glm_moe_dsa-optimized_kernels-cudnn_indexer_topk)                                   | Select fixed-width top-k indices with the cuDNN DSA indexer.                  |
| [`cudnn_sparse_attention`](#nemo_automodel-components-models-glm_moe_dsa-optimized_kernels-cudnn_sparse_attention)                           | Run cuDNN DSA attention on packed latent Q/KV tensors.                        |
| [`is_cudnn_dsa_available`](#nemo_automodel-components-models-glm_moe_dsa-optimized_kernels-is_cudnn_dsa_available)                           | Return whether all optional cuDNN DSA adapter entry points are available.     |
| [`is_dsa_kernel_available`](#nemo_automodel-components-models-glm_moe_dsa-optimized_kernels-is_dsa_kernel_available)                         | Return whether the optional TileLang kernel package for `name` is importable. |
| [`prepare_cudnn_dsa_packed_metadata`](#nemo_automodel-components-models-glm_moe_dsa-optimized_kernels-prepare_cudnn_dsa_packed_metadata)     | Prepare reusable local-query/global-key packed metadata once per stage.       |
| [`should_use_tilelang`](#nemo_automodel-components-models-glm_moe_dsa-optimized_kernels-should_use_tilelang)                                 | Decide whether to run the TileLang kernel; raise if forced but unavailable.   |
| [`tilelang_indexer_topk`](#nemo_automodel-components-models-glm_moe_dsa-optimized_kernels-tilelang_indexer_topk)                             | Fused lighting-indexer top-k selection (THD/varlen).                          |
| [`tilelang_sparse_attention`](#nemo_automodel-components-models-glm_moe_dsa-optimized_kernels-tilelang_sparse_attention)                     | Gather-top-k sparse MLA attention on the absorbed latent representation.      |

### Data

[`DsaIndexerBackend`](#nemo_automodel-components-models-glm_moe_dsa-optimized_kernels-DsaIndexerBackend)

[`DsaSparseAttentionBackend`](#nemo_automodel-components-models-glm_moe_dsa-optimized_kernels-DsaSparseAttentionBackend)

[`_CUDNN_DSA_MODULE`](#nemo_automodel-components-models-glm_moe_dsa-optimized_kernels-_CUDNN_DSA_MODULE)

### API

```python
nemo_automodel.components.models.glm_moe_dsa.optimized_kernels._all_cuda(
    tensors: torch.Tensor = ()
) -> bool
```

```python
nemo_automodel.components.models.glm_moe_dsa.optimized_kernels._generate_padded_varlen_mask_params(
    cu_seqlens: torch.Tensor,
    cu_seqlens_padded: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]
```

Build per-query key windows for padded THD layout.

`cu_seqlens` stores real document lengths compactly. `cu_seqlens_padded`
stores their offsets in the actual flattened token tensor, including CP
padding between documents. TileLang top-k indices refer to the flattened
token tensor, so starts/ends must be in the padded coordinate space while
still excluding CP padding keys for real query tokens.

```python
nemo_automodel.components.models.glm_moe_dsa.optimized_kernels.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
```

Select fixed-width top-k indices with the cuDNN DSA indexer.

**Parameters:**

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

Contiguous CUDA bfloat16 packed-THD query tensor of shape
`[tokens, index_heads, index_head_dim]`.

---

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

Contiguous CUDA bfloat16 packed-THD key tensor of shape
`[tokens, index_head_dim]`.

---

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

Contiguous CUDA FP32 tensor of shape `[tokens, index_heads]`,
scaled by `index_heads**-0.5 * index_head_dim**-0.5`. The adapter casts
it to the kernel dtype without applying any additional scaling.

---

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

CUDA int32 tensor of shape `[sequences + 1]` containing
cumulative lengths for the compact packed-token layout.

---

**`index_topk`** `int`

Fixed number of key indices emitted for every query token.

---

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

Optional CUDA integer tensor of shape `[tokens]` containing
global padded-storage query coordinates for context-parallel layouts.

---

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

Optional CUDA int32 tensor of shape `[sequences + 1]`
containing cumulative lengths for a padded packed-token layout.

---

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

Optional reusable segmented metadata prepared once for the
current packed stage.

---

**Returns:** `torch.Tensor`

Contiguous CUDA int32 tensor of shape `[tokens, 1, index_topk]`. Invalid

```python
nemo_automodel.components.models.glm_moe_dsa.optimized_kernels.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 cuDNN DSA attention on packed latent Q/KV tensors.

**Parameters:**

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

Contiguous CUDA bfloat16 packed-query tensor of shape
`[tokens, heads, kv_lora_rank + rope_head_dim]`.

---

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

Contiguous CUDA bfloat16 packed latent-KV tensor of shape
`[tokens, 1, kv_lora_rank + rope_head_dim]`.

---

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

Contiguous CUDA int32 tensor of shape `[tokens, 1, index_topk]`.

---

**`softmax_scale`** `float`

MLA attention scale applied to query-key scores.

---

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

Optional int32 valid-prefix lengths of shape `[tokens]`.

---

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

Whether every query has a positive valid-prefix length.

---

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

Optional cached int64 indices of nonempty query rows.

---

**Returns:** `torch.Tensor`

CUDA bfloat16 latent-attention tensor of shape `[tokens, heads, kv_lora_rank]`.

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

Return whether all optional cuDNN DSA adapter entry points are available.

```python
nemo_automodel.components.models.glm_moe_dsa.optimized_kernels.is_dsa_kernel_available(
    name: typing.Literal['indexer', 'sparse_attn']
) -> bool
```

Return whether the optional TileLang kernel package for `name` is importable.

```python
nemo_automodel.components.models.glm_moe_dsa.optimized_kernels.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
```

Prepare reusable local-query/global-key packed metadata once per stage.

**Parameters:**

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

CUDA int32 compact real-token offsets of shape
`[sequences + 1]`.

---

**`total_key_tokens`** `int`

Number of rows in the gathered padded-storage K/V tensor.

---

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

Optional maximum real sequence length as an integer or scalar
integer tensor.

---

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

Optional CUDA integer tensor of shape `[local_tokens]`
containing each local query's global padded-storage coordinate.

---

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

Optional CUDA int32 padded-storage boundaries of shape
`[sequences + 1]`.

---

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

Optional local boolean tensor of shape `[local_tokens]`;
`True` marks a padded query row.

---

**Returns:** `CudnnDsaPackedMetadata`

Segmented cuDNN metadata whose query and key-source fields use global

```python
nemo_automodel.components.models.glm_moe_dsa.optimized_kernels.should_use_tilelang(
    backend: str,
    available: bool,
    kernel_name: str,
    tensors: tuple[torch.Tensor, ...],
    require_bf16: bool = False
) -> bool
```

Decide whether to run the TileLang kernel; raise if forced but unavailable.

`backend="tilelang"` forces the kernel (and raises a clear error if it cannot run);
`backend="auto"` silently falls back to torch when the kernel is unavailable;
`backend="torch"` always uses the torch reference.

```python
nemo_automodel.components.models.glm_moe_dsa.optimized_kernels.tilelang_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
) -> torch.Tensor
```

Fused lighting-indexer top-k selection (THD/varlen).

**Parameters:**

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

`[T, index_n_heads, index_head_dim]` bf16 (rope already applied).

---

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

`[T, index_head_dim]` bf16 (k\_norm + rope already applied).

---

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

`[T, index_n_heads]` fp32; the caller must fold the index
`softmax_scale` into the weight (the kernel computes `relu(q·k) * w`
with no internal scale), i.e. `weights_proj(x) * index_n_heads**-0.5 *
index_head_dim**-0.5`.

---

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

`[num_seq + 1]` cumulative sequence lengths of the packed batch.

---

**`index_topk`** `int`

number of keys to keep (e.g. `2048`).

---

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

Optional global THD token indices for the local query
rows. Used by context parallelism when `index_q` is sharded but
`index_k` has been all-gathered in global token order.

---

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

Optional cumulative lengths in the padded THD token
layout. CP-packed datasets pad each document to a CP multiple, so
local query indices address this padded layout rather than the
compact `cu_seqlens` layout.

---

**Returns:** `torch.Tensor`

`topk_indices` `[T, 1, index_topk]` int32 (`-1` for invalid/causal-masked),

```python
nemo_automodel.components.models.glm_moe_dsa.optimized_kernels.tilelang_sparse_attention(
    q: torch.Tensor,
    kv_latent: torch.Tensor,
    topk_indices: torch.Tensor,
    w_vc: torch.Tensor,
    softmax_scale: float
) -> torch.Tensor
```

Gather-top-k sparse MLA attention on the absorbed latent representation.

**Parameters:**

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

`[T, n_heads, kv_lora_rank + qk_rope_head_dim]` bf16 — the absorbed query
`cat([q_nope @ w_kc, q_pe], -1)` (e.g. `512 + 64 = 576`).

---

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

`[T, 1, kv_lora_rank + qk_rope_head_dim]` bf16 — the latent KV
`cat([kv_compressed, k_pe], -1)`.

---

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

`[T, 1, index_topk]` int32 (`-1` sentinel).

---

**`w_vc`** `torch.Tensor`

`[n_heads, v_head_dim, kv_lora_rank]` — the value up-projection used to
map the latent attention output back to `v_head_dim`.

---

**`softmax_scale`** `float`

MLA attention scale `mscale**2 / sqrt(qk_head_dim)` (NOT the
kernel's `1/sqrt(dim+tail)` default).

---

**Returns:** `torch.Tensor`

`attn_out` `[T, n_heads, v_head_dim]` bf16.

```python
nemo_automodel.components.models.glm_moe_dsa.optimized_kernels.DsaIndexerBackend = Literal['torch', 'tilelang', 'cudnn', 'auto']
```

```python
nemo_automodel.components.models.glm_moe_dsa.optimized_kernels.DsaSparseAttentionBackend = Literal['torch', 'tilelang', 'cudnn', 'auto']
```

```python
nemo_automodel.components.models.glm_moe_dsa.optimized_kernels._CUDNN_DSA_MODULE = 'nemo_automodel.components.models.glm_moe_dsa.kernels.cudnn_dsa'
```