nemo_automodel.components.models.glm_moe_dsa.kernels.tilelang_sparse_mla_bwd

nemo_automodel.components.models.glm_moe_dsa.kernels.tilelang_sparse_mla_bwd.bwd(
    B,
    S,
    S_kv,
    H,
    D,
    D_tail,
    topk,
    kv_group = 1,
    sm_scale = None,
    is_causal = True,
    block_size = 32,
    num_stages = 0,
    threads = 128,
    indices_dtype = T.int32,
    dtype = T.bfloat16,
    accum_dtype = T.float32
)

nemo_automodel.components.models.glm_moe_dsa.kernels.tilelang_sparse_mla_bwd.postprocess(
    B,
    S_kv,
    D,
    D_tail,
    kv_group = 1,
    block_N = 64,
    threads = 128,
    dtype = T.bfloat16,
    accum_dtype = T.float32
)

nemo_automodel.components.models.glm_moe_dsa.kernels.tilelang_sparse_mla_bwd.preprocess(
    B,
    S,
    H,
    D,
    block_ND = 32,
    num_stages = 5,
    dtype = T.bfloat16,
    accum_dtype = T.float32
)

nemo_automodel.components.models.glm_moe_dsa.kernels.tilelang_sparse_mla_bwd.sparse_mla_bwd(
    q,
    kv,
    o,
    do,
    indices,
    lse,
    sm_scale = None,
    is_casual = True,
    return_kernel = False,
    delta = None
)