flashinfer.experimental.cake_moe_grouped_gemm.cake_grouped_mm

flashinfer.experimental.cake_moe_grouped_gemm.cake_grouped_mm(x: Tensor, w: Tensor, offs: Tensor, deterministic: bool = True) → Tensor

Warning

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

Differentiable ragged grouped GEMM Y[offs[e-1]:offs[e]] = X[offs[e-1]:offs[e]] @ W[e].T.

xtorch.Tensor

[sum_m, K] bfloat16 rows of all groups, group after group.

wtorch.Tensor

[E, N, K] bfloat16 weights ([N, K] per group).

offstorch.Tensor

int32 device tensor: [E] cumulative end offsets or m_indptr [E + 1]. Read by the kernels, never on the host.

deterministicbool

Reduction mode of the gradients. The Cake programs always reduce in a fixed order (no atomics); False is accepted for API compatibility and currently runs the same deterministic programs.

torch.Tensor

[sum_m, N] bfloat16. backward produces x.grad [sum_m, K] bfloat16 and w.grad [E, N, K] in w’s dtype (bfloat16); use grouped_gemm_wgrad() directly for an fp32 weight gradient.