flashinfer.fused_moe.trtllm_fp4_block_scale_moe

flashinfer.fused_moe.trtllm_fp4_block_scale_moe(routing_logits: ~torch.Tensor, routing_bias: ~torch.Tensor | None, hidden_states: ~torch.Tensor, hidden_states_scale: ~torch.Tensor | None, gemm1_weights: ~torch.Tensor, gemm1_weights_scale: ~torch.Tensor, gemm1_bias: ~torch.Tensor | None, gemm1_alpha: ~torch.Tensor | None, gemm1_beta: ~torch.Tensor | None, gemm1_clamp_limit: ~torch.Tensor | None, gemm2_weights: ~torch.Tensor, gemm2_weights_scale: ~torch.Tensor, gemm2_bias: ~torch.Tensor | None, output1_scale_scalar: ~torch.Tensor | None, output1_scale_gate_scalar: ~torch.Tensor | None, output2_scale_scalar: ~torch.Tensor | None, num_experts: int, top_k: int, n_group: int | None, topk_group: int | None, intermediate_size: int, local_expert_offset: int, local_num_experts: int, routed_scaling_factor: float | None, routing_method_type: int = 0, do_finalize: bool = True, enable_pdl: bool | None = None, activation_type: int = 3, per_token_scale: ~torch.Tensor | None = None, output: ~torch.Tensor | None = None, tune_max_num_tokens: int = 8192, norm_topk_prob: bool = True, routing_replay_out: ~torch.Tensor | None = None, num_fused_shared_experts: int | None = None, valid_hidden_size: int | None = None, valid_intermediate_size: int | None = None, hidden_states_scale_layout: ~flashinfer.tllm_enums.SfLayout | None = None, workspace_buffer: ~torch.Tensor | None = None, nvfp4_4over6: ~flashinfer.quantization.nvfp4_quantization_utils.NVFP44Over6Config | None = <nvfp4_4over6 unset: read FLASHINFER_NVFP4_4OVER6*>) → List[Tensor]

FP4 block-scaled MoE operation.

Parameters:

  • routing_logits (torch.Tensor) – [seq_len, num_experts] tensor of routing logits. float32 or bfloat16.

  • routing_bias (Optional[torch.Tensor]) – [num_experts] tensor of routing bias, bfloat16 or float32 (independent of routing_logits’s dtype). May be None.

  • hidden_states (torch.Tensor) – Hidden states of shape [seq_len, hidden_size // 2] (NVFP4) or [seq_len, hidden_size] (MXFP8 / bfloat16). Supports bfloat16, MXFP8, and NVFP4 (packed into uint8).

  • hidden_states_scale (Optional[torch.Tensor]) – Block scales for MXFP8 / NVFP4 hidden states of shape [seq_len, hidden_size // (32 if mxfp8 else 16)]. Dtype is float8. The equivalent flat [seq_len * hidden_size // (32 if mxfp8 else 16)] buffer returned by mxfp8_quantize() with is_sf_swizzled_layout=False (the linear layout) is also accepted. Declare which layout the buffer is in with hidden_states_scale_layout; only the linear layout is supported.

  • gemm1_weights (torch.Tensor) – [num_experts, M, hidden_size // 2] packed FP4 FC1 weights, dtype uint8. M is 2 * intermediate_size for gated activations and intermediate_size for non-gated ones (6 Relu2, 9 Identity).

  • gemm1_weights_scale (torch.Tensor) – [num_experts, M, hidden_size // (32 if mxfp4 else 16)] FC1 weight block scales, dtype float8, with the same M as gemm1_weights.

  • gemm1_bias (Optional[torch.Tensor]) – [num_experts, M] FC1 bias, float32, with the same M as gemm1_weights.

  • gemm1_alpha (Optional[torch.Tensor]) – [num_experts] swiglu alpha, float32. For SiTU this is [local_num_experts], finite and positive; None materializes per-expert alpha=1.

  • gemm1_beta (Optional[torch.Tensor]) – [num_experts] swiglu beta, float32. For SiTU this is [local_num_experts], finite and positive; None materializes per-expert beta=1.

  • gemm1_clamp_limit (Optional[torch.Tensor]) – [num_experts] swiglu clamp limit, float32. For SwigluStep, this is the raw accumulator limit (physical limit divided by output1_scale_gate_scalar); None uses physical 7. The fused epilogue caps silu(gate) and the up branch. Alpha and beta are not accepted for SwigluStep. For SiTU a provided limit is per-local-expert, finite, and positive; it clamps x0 to [-limit, limit] and x1 from above.

  • gemm2_weights (torch.Tensor) – [num_experts, hidden_size, intermediate_size] packed FP4 FC2 weights, dtype uint8.

  • gemm2_weights_scale (torch.Tensor) – [num_experts, hidden_size, intermediate_size // (32 if mxfp4 else 16)] FC2 weight block scales, dtype float8.

  • gemm2_bias (Optional[torch.Tensor]) – [num_experts, hidden_size] FC2 bias, float32.

  • output1_scale_scalar (Optional[torch.Tensor]) – [local_num_experts] scaling factors for the first-layer activation output.

  • output1_scale_gate_scalar (Optional[torch.Tensor]) – [local_num_experts] scaling factors for the first-layer gate output.

  • output2_scale_scalar (Optional[torch.Tensor]) – [local_num_experts] scaling factors for the second-layer output.

  • num_experts (int) – Total number of experts.

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

  • n_group (Optional[int]) – Number of expert groups.

  • topk_group (Optional[int]) – Number of groups to consider for top-k routing.

  • intermediate_size (int) – Size of the intermediate layer.

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

  • local_num_experts (int) – Number of experts handled by this device.

  • routed_scaling_factor (Optional[float]) – Scaling factor for routing.

  • routing_method_type (int) –

    Routing method (default 0). Selects the routing-kernel pipeline; matches flashinfer.tllm_enums.RoutingMethodType.

    • 0 Default — Softmax → TopK.

    • 1 Renormalize — TopK → Softmax.

    • 2 DeepSeekV3 — Sigmoid → RoutingBiasAdd → Top-2 in group → Top-topk_group groups → Top-top_k experts from the selected groups.

    • 3 Llama4 — Top-1 → Sigmoid.

    • 4 RenormalizeNaive — Softmax → TopK → Renormalize (Qwen3 style).

    • 5 TopK — TopK only (no softmax/sigmoid).

    • 6 SigmoidRenorm — Sigmoid → TopK → Renormalize (divide by the sum of the top-K weights).

    • 7 MiniMax2 — Sigmoid + Bias → TopK → ScaledSumNormalize (routeScale = 1.0, epsilon = 1e-20).

    • 8 Sigmoid — Sigmoid → TopK (no renormalization).

    • 9 TopKSigmoid — TopK → Sigmoid (no renormalization).

    • 10 Unspecified — reserved.

  • do_finalize (bool) – Whether to finalize the output (default True).

  • enable_pdl (Optional[bool]) – Whether to enable Programmatic Dependent Launch.

  • activation_type (int) – Activation type (default 3 — Swiglu). 3 Swiglu; 4 Geglu; 6 Relu2; 7 SwigluStep; 9 Identity. 10 SiTU computes beta*tanh(x0/beta) * alpha*tanh(x1/alpha)*sigmoid(x1).

  • per_token_scale (Optional[torch.Tensor]) – [seq_len] per-token scaling factors, float32.

  • output (Optional[torch.Tensor]) – Optional in-place [seq_len, hidden_size] output tensor.

  • tune_max_num_tokens (int) – Maximum number of tokens for autotuning (default 8192).

  • norm_topk_prob (bool) – Whether to normalize the top-k probabilities (default True).

  • routing_replay_out (Optional[torch.Tensor]) – Optional int16 tensor of shape (num_tokens_or_larger, top_k) used to capture the selected expert IDs during routing. Column order matches topk_indices. When None (default) the kernel skips the write entirely. The buffer may be larger than num_tokens for CUDA-graph pre-allocation; only rows [0, num_tokens) are written. With fused shared experts the layout stays routed-only: dim1 is top_k, not top_k + num_fused_shared_experts, and shared slots are not recorded.

  • num_fused_shared_experts (Optional[int]) – Number of shared experts to fuse into the MoE kernel (default None / 0). When > 0, every per-expert tensor (gemm1_weights, gemm1_weights_scale, gemm2_weights, gemm2_weights_scale, output*_scale_scalar, biases) must have num_experts + num_fused_shared_experts rows in the expert dimension — the shared-expert weights are appended after the routed ones. Every token is unconditionally routed to the shared experts with weight 1.0. With do_finalize=False, the returned expert_weights and expanded_idx_to_permuted_idx cover top_k + num_fused_shared_experts slots per token.

  • hidden_states_scale_layout (Optional[flashinfer.tllm_enums.SfLayout]) –

    Layout of hidden_states_scale. This is validation only: the trtllm-gen routed GEMM can consume only layout_linear, so declaring any swizzled layout raises NotImplementedError instead of reading the buffer as linear bytes and silently producing wrong results.

    Deprecated since version Omitting: this argument is deprecated. None (the default) preserves today’s behavior — the linear layout is inferred — and emits a DeprecationWarning. The layout genuinely cannot be recovered from the tensor: a 128x4-swizzled buffer has exactly the linear numel whenever num_tokens % 128 == 0. Both the inference and this argument’s optionality are deprecated together, so a future release can make it required and drop the inference in one step.

  • valid_hidden_size (Optional[int]) – Valid (unpadded) hidden dimension. When provided, the hidden_size implied by the tensor shapes is treated as padded and only the valid region is contracted. Matching TensorRT-LLM, the returned tensor is exactly valid_hidden_size wide: GEMM2 computes roundUp(valid_hidden_size, 128) columns (that is what the FC2 weights, weight scales and bias are sized for) and finalize writes only the leading valid_hidden_size of them. A caller-supplied output must therefore have valid_hidden_size columns. For NVFP4 the padded hidden_size must be a multiple of 256 (128 gives wrong results). Default None (use the full hidden_size).

  • valid_intermediate_size (Optional[int]) –

    Valid (unpadded) intermediate dimension. When provided, intermediate_size is treated as padded and only the valid region is computed. Default None (use the full intermediate_size).

    The kernels tile expert weights in 128-row blocks, so intermediate_size (and with it the FC1 row count M) must be a multiple of 128; a model dimension that is not (1856, 928) is supported by padding and declaring, not by passing it directly. Zero-pad the weights to a multiple of 128 (the gate and up halves of FC1 separately, so the rows stay [gate; up] at the padded size), fill the padded block scales with a neutral finite value (E8M0 127, i.e. 2**0, for MXFP4; E4M3 1.0 for NVFP4), pass the padded size as intermediate_size and the real one here. The kernel reads the [valid, roundUp(valid, 128)) stretch of the padding and never reads the rest, so it must hold zeros and finite scales. hidden_size follows the same rule through valid_hidden_size.

  • workspace_buffer (Optional[torch.Tensor]) – Caller-owned 1-D uint8 CUDA buffer of at least trtllm_fused_moe_workspace_size() bytes. When given, the launch carves all of its internal scratch (routing maps, intermediate GEMM outputs and batched-GEMM workspaces) from it instead of allocating per call, so MoE layers that run sequentially on one stream can share a single buffer. Requires do_finalize=True and no gemm1_lora_delta. Autotuning profiles from the same buffer; scratch that does not fit (a larger token count) is allocated per call. Default None (allocate per call).

  • nvfp4_4over6 (NVFP44Over6Config or None) –

    NVFP4 4over6 recipe for the per-token requantization of the GEMM2 input (the GEMM1 input is quantized by the caller).

    • omitted (the default): read the legacy FLASHINFER_NVFP4_4OVER6* environment variables.

    • None: 4over6 off, environment ignored.

    • NVFP44Over6Config: that recipe. Requires per_token_scale; no other path has a 4over6 requantization, so a pinned recipe raises ValueError.

Returns:

[output] when do_finalize is True, otherwise [gemm2_output, expert_weights, expanded_idx_to_permuted_idx]. The expert_weights tensor is always bfloat16 (the routing kernel emits bf16 weights for every routing method), regardless of the routing_logits dtype — including the do_finalize=False path and fp32 DeepSeekV3 logits.

Return type:

List[torch.Tensor]