flashinfer.experimental.cake_moe_grouped_gemm.grouped_gemm_dgrad

flashinfer.experimental.cake_moe_grouped_gemm.grouped_gemm_dgrad(g: Tensor, w: Tensor, offs: Tensor, out: Tensor | None = None) → Tensor

Warning

cake_moe_grouped_gemm.grouped_gemm_dgrad is experimental: it provides no compatibility guarantees and may change or be removed without deprecation.

dX[offs[e-1]:offs[e]] = G[offs[e-1]:offs[e]] @ W[e] for every group e.

g is [sum_m, N] bfloat16, w is [E, N, K] bfloat16 and is read in place (no transpose copy), offs is [E] int32 device end offsets (or m_indptr [E + 1]) and out is [sum_m, K] bfloat16. K % 256 == 0 and N % 64 == 0 are required. See grouped_gemm_fwd for the prepared form.