nemo_automodel.components.models.glm_moe_dsa.kernels.tilelang_sparse_mla_fwd

nemo_automodel.components.models.glm_moe_dsa.kernels.tilelang_sparse_mla_fwd.sparse_mla_fwd(
    heads,
    dim,
    tail_dim,
    topk,
    kv_group = 1,
    sm_scale = None,
    is_causal = True,
    CP0 = True,
    block_I = 64,
    num_stages = 2,
    threads = 256
)

nemo_automodel.components.models.glm_moe_dsa.kernels.tilelang_sparse_mla_fwd.sparse_mla_fwd_interface(
    q,
    kv,
    indices,
    sm_scale = None,
    return_p_sum: bool = False,
    d_v = 512,
    block_I = 64,
    num_stages = 2,
    threads = 256
)