flashinfer.fused_moe.CudnnGroupedGemmMxfp8Config

class flashinfer.fused_moe.CudnnGroupedGemmMxfp8Config

cuDNN grouped-GEMM MXFP8 backend (grouped_mm_mxfp8).

  • Pipeline: expert-sorted permutation, GEMM1, the typed activation, GEMM2 and the finalize, on flashinfer.grouped_mm.grouped_mm_mxfp8().

  • Activations: the MXFP8 activation pack, E4M3 [M, H] values (M tokens of hidden size H) with token-major uint8 [M, H/32] UE8M0 block scales (see prepare_activations()).

  • Rows and scale rows are permuted into expert order, each expert segment padded to a multiple of 128 rows; the permuted scale rows are swizzled into the GEMM’s block-scale layout, and the BF16 intermediate is quantized straight into that layout before GEMM2.

  • Weights: E4M3 with UE8M0 block scales per expert (see prepare_weights()).

  • hidden_size and intermediate_size must be multiples of 128.

  • 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].

  • 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.

__init__() → None

Methods

__init__()

prepare_activations(hidden_states_bf16)

Quantize BF16 activations to E4M3 plus token-major UE8M0 block scales.

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

Quantize canonical BF16 weights to per-expert MXFP8 with swizzled block scales.

supported(arch)