flashinfer.fused_moe.CudnnGroupedGemmBf16Config

class flashinfer.fused_moe.CudnnGroupedGemmBf16Config

cuDNN grouped-GEMM BF16 backend (grouped_mm_bf16).

  • Pipeline: sort the pre-routed assignments by expert into a contiguous m_indptr layout, gather the activations, GEMM1, the typed activation, GEMM2 and the finalize, on flashinfer.grouped_mm.grouped_mm_bf16().

  • Activations: raw bfloat16 [M, H] (M tokens of hidden size H).

  • hidden_size and intermediate_size must be multiples of 8.

  • Weights: prepare_weights() takes the canonical [up, gate] BF16 weights as they are.

  • Permute and finalize: FlashInfer’s moe_sort, moe_permute and moe_unpermute kernels.

  • Activation functions: SwiGLU with default scalars, GeGLU, GELU, ReLU, SiLU and Identity, on FlashInfer’s moe_activation kernel.

  • MoEFinalizeConfig(do_finalize=False) returns the unfinalized [gemm2_out, expert_weights, token_to_row] for the caller to combine.

  • Expert parallelism: assignments to non-local experts are masked out.

  • Routing weights: packed, or unpacked FP32 / BF16.

  • Each token’s topk_ids must name distinct experts.

  • Not in the default backend list, since cuDNN builds one execution plan per token bucket: select it with BackendOptions((CudnnGroupedGemmBf16Config(),)).

__init__() → None

Methods

__init__()

prepare_weights(w1_bf16, w2_bf16, *, ...[, ...])

Build the cudnn_grouped_gemm_bf16 weight view from canonical BF16 weights.

supported(arch)