nemo_automodel.components.models.glm_moe_dsa.kernels.tilelang_indexer_fwd

Module Contents

Functions

Name	Description
`clean_logits_`	-
`indexer_fwd_interface`	-
`tl_indexer_fwd_impl`	-

API

nemo_automodel.components.models.glm_moe_dsa.kernels.tilelang_indexer_fwd.clean_logits_(
    threads: int = 512,
    block_K: int = 4096
)

nemo_automodel.components.models.glm_moe_dsa.kernels.tilelang_indexer_fwd.indexer_fwd_interface(
    q,
    kv,
    weights,
    cu_seqlen_ks,
    cu_seqlen_ke,
    clean_logits = True
)

nemo_automodel.components.models.glm_moe_dsa.kernels.tilelang_indexer_fwd.tl_indexer_fwd_impl(
    heads,
    index_dim,
    block_N = 256,
    num_stages = 3,
    threads = 512,
    block_Q = None
)