`core.inference.moe.flashinfer_nvfp4`#

FlashInfer routed-MoE support for per-token NVFP4 x NVFP4 squared-ReLU and SwiGLU experts.

Expert weights are quantized from MCore’s BF16 expert stacks with one FP32 global scale per expert. Activations are quantized online with one FP32 scale per token, so no calibration or static input scale is required. The experts run on FlashInfer’s CuTe DSL block-scaled grouped GEMMs.

Module Contents#

Classes#

FlashInferRoutedNVFP4Weights

FC1/FC2 expert stacks in the CuTe DSL block-scaled grouped-GEMM layout.

Functions#

`require_flashinfer_routed_nvfp4`	Raise an error when the per-token NVFP4 routed-MoE APIs are unavailable.
`_round_up`
`_global_encode_scale`	Return the one-element NVFP4 global encode scale for a weight matrix.
`_zero_pad`
`_pad_fc1`	Zero-pad one FC1 expert weight into the CuTe DSL row order.
`_quantize_expert`	Quantize one zero-padded `[rows, cols]` expert weight into CuTe DSL NVFP4 storage.
`_validate_stacks`	Return (experts, intermediate, hidden, gated); gated FC1 has twice FC2’s input width.
`_check_refresh_metadata`
`prepare_routed_nvfp4_weights`	Quantize BF16 squared-ReLU or SwiGLU expert stacks into the per-token NVFP4 kernel layout.
`quantize_routed_nvfp4_input`	Pad and quantize routed-MoE input with one FP32 scale per token.
`_get_cute_dsl_wrapper`
`flashinfer_routed_nvfp4_moe`	Run the FlashInfer per-token NVFP4 squared-ReLU or SwiGLU MoE.

Data#

`HAVE_FLASHINFER_ROUTED_NVFP4`
`NVFP4_BLOCK_SIZE`
`FLOAT4_E2M1_MAX`
`FLOAT8_E4M3_MAX`
`PER_TOKEN_ACTIVATION_GLOBAL_SCALE`
`_ALIGNMENT`
`_GLU_INTERLEAVE_ROWS`
`_CUTE_DSL_WRAPPERS`

API#

core.inference.moe.flashinfer_nvfp4.HAVE_FLASHINFER_ROUTED_NVFP4#: None

core.inference.moe.flashinfer_nvfp4.NVFP4_BLOCK_SIZE#: 16

core.inference.moe.flashinfer_nvfp4.FLOAT4_E2M1_MAX#: 6.0

core.inference.moe.flashinfer_nvfp4.FLOAT8_E4M3_MAX#: 448.0

core.inference.moe.flashinfer_nvfp4.PER_TOKEN_ACTIVATION_GLOBAL_SCALE#: None

core.inference.moe.flashinfer_nvfp4._ALIGNMENT#: 256

core.inference.moe.flashinfer_nvfp4._GLU_INTERLEAVE_ROWS#: 64

core.inference.moe.flashinfer_nvfp4._CUTE_DSL_WRAPPERS: dict#: None

core.inference.moe.flashinfer_nvfp4.require_flashinfer_routed_nvfp4() → None#: Raise an error when the per-token NVFP4 routed-MoE APIs are unavailable.

core.inference.moe.flashinfer_nvfp4._round_up(value: int, alignment: int) → int#

class core.inference.moe.flashinfer_nvfp4.FlashInferRoutedNVFP4Weights#

FC1/FC2 expert stacks in the CuTe DSL block-scaled grouped-GEMM layout.

Data is packed E2M1 [experts, N, K // 2]. *_scale is the 6D MMA view of the 128x4-swizzled E4M3 block scales held in *_scale_swizzled. *_alpha are the per-expert FP32 dequantization scales and fc2_input_scale is the per-token recipe constant for the FC2 input. Gated (SwiGLU) FC1 has 2 * padded_intermediate_size rows in FlashInfer’s interleaved up/gate order.

hidden_size: int#: None

intermediate_size: int#: None

padded_hidden_size: int#: None

padded_intermediate_size: int#: None

gated: bool#: None

fc1_data: torch.Tensor#: None

fc1_scale_swizzled: torch.Tensor#: None

fc1_scale: torch.Tensor#: None

fc1_alpha: torch.Tensor#: None

fc2_data: torch.Tensor#: None

fc2_scale_swizzled: torch.Tensor#: None

fc2_scale: torch.Tensor#: None

fc2_alpha: torch.Tensor#: None

fc2_input_scale: torch.Tensor#: None

property num_local_experts: int#: Number of experts in each stack.

core.inference.moe.flashinfer_nvfp4._global_encode_scale(weight: torch.Tensor) → torch.Tensor#: Return the one-element NVFP4 global encode scale for a weight matrix.

core.inference.moe.flashinfer_nvfp4._zero_pad( weight: torch.Tensor, padded_rows: int, padded_cols: int, ) → torch.Tensor#

core.inference.moe.flashinfer_nvfp4._pad_fc1( weight: torch.Tensor, gated: bool, padded_i: int, padded_h: int, ) → torch.Tensor#

Zero-pad one FC1 expert weight into the CuTe DSL row order.

MCore stores gated FC1 as [gate; up]. FlashInfer expects [up; gate] with each half padded separately and the halves interleaved in _GLU_INTERLEAVE_ROWS-row blocks.

core.inference.moe.flashinfer_nvfp4._quantize_expert( padded: torch.Tensor, data_out: torch.Tensor, swizzled_scale_out: torch.Tensor, alpha_out: torch.Tensor, ) → None#: Quantize one zero-padded [rows, cols] expert weight into CuTe DSL NVFP4 storage.

core.inference.moe.flashinfer_nvfp4._validate_stacks( fc1_weight: torch.Tensor, fc2_weight: torch.Tensor, ) → tuple[int, int, int, bool]#: Return (experts, intermediate, hidden, gated); gated FC1 has twice FC2’s input width.

core.inference.moe.flashinfer_nvfp4._check_refresh_metadata( out, experts, hidden_size, intermediate_size, padded_h, padded_i, gated, )#

core.inference.moe.flashinfer_nvfp4.prepare_routed_nvfp4_weights( fc1_weight: torch.Tensor, fc2_weight: torch.Tensor, out: core.inference.moe.flashinfer_nvfp4.FlashInferRoutedNVFP4Weights | None = None, ) → core.inference.moe.flashinfer_nvfp4.FlashInferRoutedNVFP4Weights#

Quantize BF16 squared-ReLU or SwiGLU expert stacks into the per-token NVFP4 kernel layout.

Parameters:

fc1_weight – [experts, intermediate, hidden] BF16 FC1 weights, or [experts, 2 * intermediate, hidden] as MCore’s [gate; up] for SwiGLU.
fc2_weight – [experts, hidden, intermediate] BF16 FC2 weights.
out – When provided, its tensors are refreshed in place. This preserves addresses captured by CUDA graphs across model refits.

core.inference.moe.flashinfer_nvfp4.quantize_routed_nvfp4_input( hidden_states: torch.Tensor, padded_hidden_size: int, ) → tuple[torch.Tensor, torch.Tensor, torch.Tensor]#

Pad and quantize routed-MoE input with one FP32 scale per token.

Returns packed E2M1 data [tokens, padded_hidden // 2], linear block scales [tokens, padded_hidden // 16] (as returned by FlashInfer), and FP32 per-token scales [tokens].

core.inference.moe.flashinfer_nvfp4._get_cute_dsl_wrapper( weights: core.inference.moe.flashinfer_nvfp4.FlashInferRoutedNVFP4Weights, num_experts: int, top_k: int, local_expert_offset: int, activation_type: int, device: torch.device, )#

core.inference.moe.flashinfer_nvfp4.flashinfer_routed_nvfp4_moe( hidden_states: torch.Tensor, routing_map: torch.Tensor, probabilities: torch.Tensor, weights: core.inference.moe.flashinfer_nvfp4.FlashInferRoutedNVFP4Weights, *, num_experts: int, local_expert_offset: int, activation_type: int, out: torch.Tensor | None = None, token_capacity: int | None = None, use_bounded_rows: bool = False, ) → torch.Tensor#

Run the FlashInfer per-token NVFP4 squared-ReLU or SwiGLU MoE.

Rows whose routing entries are -1 contribute zero output, so a graph-stable dispatcher buffer may be passed in full. When token_capacity is set and the controller marks the current graph safe, only that fixed prefix is quantized and processed, as for routed MXFP8; rows past it are left untouched in out.

core.inference.moe.flashinfer_nvfp4#