flashinfer.fused_moe.trtllm_fp8_block_scale_moe

flashinfer.fused_moe.trtllm_fp8_block_scale_moe(routing_logits: Tensor, routing_bias: Tensor | None, hidden_states: Tensor, hidden_states_scale: Tensor, gemm1_weights: Tensor, gemm1_weights_scale: Tensor, gemm2_weights: Tensor, gemm2_weights_scale: Tensor, 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, use_shuffled_weight: bool = False, weight_layout: int = 0, do_finalize: bool = True, enable_pdl: bool | None = None, tune_max_num_tokens: int = 8192, fp8_quantization_type: Fp8QuantizationType = Fp8QuantizationType.DeepSeekFp8, num_fused_shared_experts: int | None = None, activation_type: int = 3, norm_topk_prob: bool = True, routing_replay_out: Tensor | None = None, gemm1_alpha: Tensor | None = None, gemm1_beta: Tensor | None = None, gemm1_clamp_limit: Tensor | None = None, output: Tensor | None = None, valid_hidden_size: int | None = None, valid_intermediate_size: int | None = None, *, gemm1_bias: Tensor | None = None, gemm2_bias: Tensor | None = None, workspace_buffer: Tensor | None = None) → List[Tensor] | Tensor

FP8 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. May be None.

  • hidden_states (torch.Tensor) – [seq_len, hidden_size] tensor of input hidden states. float16, bfloat16, or float8_e4m3fn (block scale must match: see hidden_states_scale).

  • hidden_states_scale (torch.Tensor) – [hidden_size // 128, seq_len] tensor of hidden-states block scales.

  • gemm1_weights (torch.Tensor) – First-layer weights. [num_experts, M, hidden_size] when weight_layout == WeightLayout.MajorK (0), or [num_experts, M // 128, hidden_size, 128] when weight_layout == WeightLayout.BlockMajorK (2). M is 2 * intermediate_size for gated activations and intermediate_size for non-gated activations. For gated activations each expert’s logical [M, hidden_size] matrix, before the use_shuffled_weight preprocessing (reorder_rows_for_gated_act_gemm() then shuffle_matrix_a) and any layout conversion, stacks the two halves as [X1; X2]: rows [0, intermediate_size) are X1 (linear/up) and rows [intermediate_size, 2 * intermediate_size) are X2 (gate). Weights stored with the halves interleaved in fixed-width row blocks are de-interleaved per expert at load time with deinterleave_gated_rows(); no kernel variant reads an interleaved FC1.

  • gemm1_weights_scale (torch.Tensor) – [num_experts, 2*intermediate_size // (32 if mxfp8 else 128), hidden_size // (32 if mxfp8 else 128)] first-layer block scales.

  • gemm2_weights (torch.Tensor) – Second-layer weights. [num_experts, hidden_size, intermediate_size] when weight_layout == WeightLayout.MajorK, or [num_experts, hidden_size // 128, intermediate_size, 128] when weight_layout == WeightLayout.BlockMajorK.

  • gemm2_weights_scale (torch.Tensor) – [num_experts, hidden_size // (32 if mxfp8 else 128), intermediate_size // (32 if mxfp8 else 128)] second-layer block scales.

  • 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). See trtllm_bf16_moe().

  • use_shuffled_weight (bool) – Whether to use the shuffled weight layout (default False).

  • weight_layout (int) –

    Weight layout for gemm1_weights / gemm2_weights; matches flashinfer.tllm_enums.WeightLayout. Allowed values for this function depend on fp8_quantization_type: DeepSeekFp8 accepts MajorK or BlockMajorK; MxFp8 requires MajorK. Default 0 (MajorK).

    • 0 MajorK — K-major, logical shape [Mn, K].

    • 1 MajorMn — M-major (A) / N-major (B), logical shape [K, Mn]. Not supported by this function.

    • 2 BlockMajorK — Blocked along K, logical shape [K / blockK, Mn, blockK] (blockK is fixed at 128 B). Only valid when ``fp8_quantization_type`` is ``DeepSeekFp8``.

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

  • enable_pdl (Optional[bool]) – Whether to enable Programmatic Dependent Launch. None (default) lets the runtime auto-select on SM90+.

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

  • fp8_quantization_type (Fp8QuantizationType) – FP8 quantization scheme (default Fp8QuantizationType.DeepSeekFp8).

  • num_fused_shared_experts (Optional[int]) – Number of shared experts to fuse into the MoE kernel (default None / 0). When > 0, the weight tensors must have num_experts + num_fused_shared_experts in the expert dimension. Expert parallelism (EP) is not yet supported together with fused shared experts: when this is > 0 you must pass local_expert_offset == 0 and local_num_experts == num_experts (all routed experts local), otherwise a ValueError is raised.

  • activation_type (int) – Activation type (default 3 — Swiglu). 3 Swiglu; 4 Geglu; 6 Relu2; 7 SwigluStep (MxFp8 only); 9 Identity.

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

  • gemm1_alpha (Optional[torch.Tensor]) – Optional [local_num_experts] float32 per-expert SwiGLU OA alpha parameter. Supported for Fp8QuantizationType.MxFp8 and Fp8QuantizationType.DeepSeekFp8 with ActivationType.Swiglu. Any subset of gemm1_alpha, gemm1_beta, gemm1_clamp_limit can be provided independently. When None (default), alpha=1.0 is used. Let GEMM1 output be split as X1 (linear/up half) and X2 (gate half). The activation output is X2 * sigmoid(alpha * X2) * (X1 + beta). Pass raw values; neither block-scale recipe carries a scalar dequant scale, so no host-side conversion is applied.

  • gemm1_beta (Optional[torch.Tensor]) – Optional [local_num_experts] float32 per-expert SwiGLU OA beta parameter. Supported for Fp8QuantizationType.MxFp8 and Fp8QuantizationType.DeepSeekFp8 with ActivationType.Swiglu. When None (default), beta=0.0 is used.

  • gemm1_clamp_limit (Optional[torch.Tensor]) – Optional [local_num_experts] float32 per-expert clamp limit. With Swiglu, supported for MxFp8 and DeepSeekFp8 and clamps X1 to [-limit, limit] and X2 from above before SiLU. With SwigluStep, supported only for fused MxFp8 and clamps X1 and silu(X2) at the physical limit. None uses physical 7 for SwigluStep. MXFP8 has no global FC1 dequant scale, so raw accumulator and physical units coincide. A trainer-side SwiGLU given as a clamp value v (linear half clamped to [-v, v], gate capped at v) and a linear offset o (act(gate) * (linear + o)), such as Megatron-Core’s activation_func_clamp_value and glu_linear_offset, maps to gemm1_clamp_limit=v, gemm1_beta=o and gemm1_alpha=None (1.0) with ActivationType.Swiglu.

  • output (Optional[torch.Tensor]) – Optional in-place output tensor of shape [seq_len, hidden_size]. Allocated internally when None (default).

  • 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 MxFP8 the padded hidden_size must be a multiple of 256 (128 is rejected). 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 MXFP8; 1.0 for the float32 DeepSeek scales), 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.

  • gemm1_bias (Optional[torch.Tensor]) – Optional FC1 bias, float32 (keyword-only). Shape is [local_num_experts, 2 * intermediate_size] for gated activations and [local_num_experts, intermediate_size] for non-gated activations. This tensor must be in the same row layout as gemm1_weights.

  • gemm2_bias (Optional[torch.Tensor]) – Optional [local_num_experts, hidden_size] FC2 bias, float32 (keyword-only). This tensor must be in the same row layout as gemm2_weights.

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

Returns:

Final MoE output when do_finalize is True, otherwise [gemm2_output, expert_weights, expanded_idx_to_permuted_idx].

Return type:

torch.Tensor or List[torch.Tensor]