flashinfer.experimental.cake_moe_grouped_gemm.grouped_gemm_fwd

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

Warning

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

Y[offs[e-1]:offs[e]] = X[offs[e-1]:offs[e]] @ W[e].T for every group e.

xtorch.Tensor

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

wtorch.Tensor

[E, N, K] bfloat16 weights, one [N, K] matrix per group.

offstorch.Tensor

[E] int32 cumulative end offsets on the device (offs[-1] == sum_m), or m_indptr [E + 1] with a leading 0; read by the kernel, never on the host.

outOptional[torch.Tensor]

[sum_m, N] bfloat16 output (allocated when omitted); may be row-padded.

torch.Tensor

out. N % 256 == 0 and K % 64 == 0 are required.

Prepares and launches in one call. For repeated launches over the same tensors (or CUDA Graph capture) use prepare_grouped_gemm_fwd once and call launch() on the returned object.