flashinfer.fused_moe.cute_dsl_fused_moe_nvfp4

flashinfer.fused_moe.cute_dsl_fused_moe_nvfp4(x: Tensor, x_sf: Tensor | None, token_selected_experts: Tensor, token_final_scales: Tensor, w1_weight: Tensor, w1_weight_sf: Tensor, w1_alpha: Tensor, fc2_input_scale: Tensor | None, w2_weight: Tensor, w2_weight_sf: Tensor, w2_alpha: Tensor, num_experts: int, top_k: int, num_local_experts: int | None = None, local_expert_offset: int = 0, output_dtype: dtype = torch.bfloat16, use_fused_finalize: bool = True, moe_output: Tensor | None = None, aux_stream: Stream | None = None, enable_pdl: bool = True, activation_type: int = 3, swiglu_alpha: float = 1.0, swiglu_beta: float = 0.0, swiglu_limit: float = 3.4028234663852886e+38, situ_beta: float | None = None, situ_linear_beta: float | None = None, *, quant_mode: str = 'w4a4', per_token_scale: Tensor | None = None) → Tensor

Run a fused MoE forward pass using the CuTe-DSL NVFP4 kernels.

Supported architectures: SM100, SM103. This is the simple functional API; for CUDA-graph support use CuteDslMoEWrapper instead.

Auto-tuning is controlled by the autotune() context manager:

with autotune(True):
    output = cute_dsl_fused_moe_nvfp4(...)
Parameters:

  • x (torch.Tensor) – Packed NVFP4 input for quant_mode="w4a4" or BF16 input for quant_mode="w4a16".

  • x_sf (Optional[torch.Tensor]) – Scale factors for quant_mode="w4a4"; must be None for quant_mode="w4a16".

  • token_selected_experts (torch.Tensor) – Expert assignments of shape [num_tokens, top_k].

  • token_final_scales (torch.Tensor) – Routing weights of shape [num_tokens, top_k].

  • w1_weight (torch.Tensor) – GEMM1 weights (gate + up fused for gated activations, or a single projection for non-gated activations).

  • w1_weight_sf (torch.Tensor) – Scale factors for w1_weight.

  • w1_alpha (torch.Tensor) – Per-expert global scale for GEMM1.

  • fc2_input_scale (Optional[torch.Tensor]) – Global scale for W4A4 GEMM2 input quantization; must be None for W4A16 because GEMM1 output stays in BF16.

  • w2_weight (torch.Tensor) – GEMM2 weights (down projection).

  • w2_weight_sf (torch.Tensor) – Scale factors for w2_weight.

  • w2_alpha (torch.Tensor) – Per-expert global scale for GEMM2.

  • num_experts (int) – Total number of experts.

  • top_k (int) – Number of experts routed to per token.

  • num_local_experts (Optional[int]) – Local experts for expert parallelism. Defaults to num_experts.

  • local_expert_offset (int) – Offset of local experts in the global expert space. Defaults to 0.

  • output_dtype (torch.dtype) – Output dtype. Defaults to torch.bfloat16.

  • use_fused_finalize (bool) – Use atomic fused finalize; otherwise use the deterministic two-stage finalize. Defaults to True.

  • moe_output (Optional[torch.Tensor]) – Pre-allocated output buffer. Allocated internally if None.

  • aux_stream (Optional[torch.cuda.Stream]) – Optional auxiliary CUDA stream used to overlap setup work with the main computation.

  • enable_pdl (bool) – Enable Programmatic Dependent Launch. Defaults to True.

  • activation_type (int) – FC1 activation type. Use ActivationType.Swiglu for gated SwiGLU/SiTU, ActivationType.GegluTanh for tanh-approximate GeGLU, and ActivationType.Relu2 for non-gated ReLU^2. Setting situ_beta selects SiTU; swiglu_oai is represented as ActivationType.Swiglu with non-default swiglu_alpha/beta/limit.

  • swiglu_alpha (float) – SwiGLU parameters.

  • swiglu_beta (float) – SwiGLU parameters.

  • swiglu_limit (float) – SwiGLU parameters.

  • quant_mode (str) – Compute mode: "w4a4" / "nvfp4" or "w4a16". Defaults to "w4a4".

  • situ_beta (Optional[float]) – When set with ActivationType.Swiglu, use the SiTU gate beta * tanh(gate / beta) * sigmoid(gate).

  • situ_linear_beta (Optional[float]) – Optional SiTU tanh clamp for the up branch.

  • per_token_scale (Optional[torch.Tensor]) – Optional W4A4 per-token input row scale for GEMM1.

Returns:

Output tensor of shape [num_tokens, hidden_size].

Return type:

torch.Tensor