Skip to navigation
ReferenceFull Library ReferenceNemo AutomodelNemo AutomodelComponentsModelsGlm Moe DsaKernelsnemo_automodel.components.models.glm_moe_dsa.kernels.indexer

nemo_automodel.components.models.glm_moe_dsa.kernels.indexer

View as Markdown

Module Contents

Classes

NameDescription
IndexerFunction-

Functions

NameDescription
_topk_in_row_chunksBound temporary memory while preserving independent per-row top-k selection.
generate_varlen_mask_params-
lighting_indexer-
pytorch_extract_topk_scores-

Data

_TOPK_MAX_ELEMENTS_PER_CALL

API

class nemo_automodel.components.models.glm_moe_dsa.kernels.indexer.IndexerFunction()

Bases: Function

nemo_automodel.components.models.glm_moe_dsa.kernels.indexer.IndexerFunction.backward(
    ctx,
    grad_scores,
    grad_indices
)
staticmethod
nemo_automodel.components.models.glm_moe_dsa.kernels.indexer.IndexerFunction.forward(
    ctx,
    index_q: torch.Tensor,
    index_k: torch.Tensor,
    weights: torch.Tensor,
    cu_seqlen_ks: torch.Tensor,
    cu_seqlen_ke: torch.Tensor,
    topk: int,
    topk_indices: torch.Tensor | None = None
)
staticmethod
nemo_automodel.components.models.glm_moe_dsa.kernels.indexer._topk_in_row_chunks(
    logits: torch.Tensor,
    topk: int
) -> tuple[torch.Tensor, torch.Tensor]

Bound temporary memory while preserving independent per-row top-k selection.

nemo_automodel.components.models.glm_moe_dsa.kernels.indexer.generate_varlen_mask_params(
    cu_seqlens
)
 
nemo_automodel.components.models.glm_moe_dsa.kernels.indexer.lighting_indexer(
    index_q: torch.Tensor,
    index_k: torch.Tensor,
    weights: torch.Tensor,
    cu_seqlen_ks: torch.Tensor,
    cu_seqlen_ke: torch.Tensor,
    topk: int,
    topk_indices: torch.Tensor | None = None
)
 
nemo_automodel.components.models.glm_moe_dsa.kernels.indexer.pytorch_extract_topk_scores(
    logits,
    topk_indices,
    dim = -1
)
 
nemo_automodel.components.models.glm_moe_dsa.kernels.indexer._TOPK_MAX_ELEMENTS_PER_CALL = 32 * 1024 * 1024