flashinfer.experimental.cake_moe_grouped_gemm.grouped_gemm_wgrad

flashinfer.experimental.cake_moe_grouped_gemm.grouped_gemm_wgrad(g: Tensor, x: Tensor, offs: Tensor, out: Tensor | None = None, out_dtype: dtype | None = None, num_groups: int | None = None) → Tensor

Warning

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

dW[e] = G[offs[e-1]:offs[e]].T @ X[offs[e-1]:offs[e]] for every group e.

g is [sum_m, N] bfloat16, x is [sum_m, K] bfloat16, offs is [E] int32 device end offsets (E = offs.numel(); pass num_groups=E to hand in m_indptr [E + 1] instead) and out is [E, N, K] in out_dtype (bfloat16 or float32; defaults to the dtype of out, else bfloat16). fp32 accumulation, bitwise deterministic, exact zeros for empty groups. N % 256 == 0 and K a multiple of the selected k tile (256, or 512 when chosen) are required. See grouped_gemm_fwd for the prepared form.