flashinfer.fused_moe.CudnnGroupedGemmNvfp4Config

class flashinfer.fused_moe.CudnnGroupedGemmNvfp4Config

cuDNN grouped-GEMM NVFP4 backend (grouped_mm_fp4 with 16-wide blocks).

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

  • Activations: the NVFP4 activation pack, packed E2M1 uint8 [M, H/2] values (M tokens of hidden size H) with token-major float8_e4m3fn [M, H/16] block scales and a global scale of one (see prepare_activations() for the magnitude range that encoding resolves).

  • 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 to NVFP4 with a global scale of one, straight into that layout, before GEMM2.

  • Weights: packed E2M1 with E4M3 block scales and one global scale per expert (see prepare_weights()). GEMM1’s output rows are multiplied by their expert’s global dequant; GEMM2’s is applied by the finalize, or to the unfinalized rows.

  • 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 packed E2M1 plus token-major E4M3 block scales.

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

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

supported(arch)