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`	Reusable THD metadata for local-query/global-key cuDNN DSA.
`_CudnnSparseAttention`	Pair FlashMLA forward with cuDNN backward for latent THD attention.

Functions

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

Data

_ATTENTION_HEAD_DIM

_FLASH_MLA_TOPK_ALIGNMENT

_INDEX_HEAD_DIM

_MAX_TOPK

_TOPK_ROW_ALIGNMENT

_TOPK_SCRATCH_INT32_FACTOR

_TOPK_SCRATCH_LIMIT_BYTES

_VALUE_HEAD_DIM

__all__

API

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

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

Bases: Function

Pair FlashMLA forward with cuDNN backward for latent THD attention.

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],

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.

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

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

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.

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.

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.

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.

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

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.

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.

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 | NoneDefaults to None

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

cu_seqlens_padded

torch.Tensor | NoneDefaults to None

Optional global padded packed-layout offsets.

packed_metadata

CudnnDsaPackedMetadata | NoneDefaults to 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.

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:

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 | NoneDefaults to 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

boolDefaults to 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 | NoneDefaults to 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.

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.

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 | NoneDefaults to 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 | NoneDefaults to 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 | NoneDefaults to None

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

padding_mask

torch.Tensor | NoneDefaults to 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,

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

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

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

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

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

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

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

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

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