flashinfer.fused_moe.trtllm_fp8_per_channel_scale_routed_moe

flashinfer.fused_moe.trtllm_fp8_per_channel_scale_routed_moe(topk_ids: Tensor, routing_bias: Tensor | None, hidden_states: Tensor, hidden_states_scale: Tensor, gemm1_weights: Tensor, gemm1_per_channel_weight_scale: Tensor, output1_scale_scalar: Tensor, output1_scale_gate_scalar: Tensor, gemm2_weights: Tensor, gemm2_per_channel_weight_scale: Tensor, output2_scale_scalar: 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, use_routing_scales_on_input: bool = False, routing_method_type: int = 0, do_finalize: bool = True, enable_pdl: bool | None = None, tune_max_num_tokens: int = 8192, activation_type: int = 3) → List[Tensor] | Tensor

FP8 per-token activation/per-channel weight MoE with pre-computed routing.

Parameters:

  • topk_ids (torch.Tensor) – [seq_len, top_k] int32 tensor of packed expert indices and weights with format (expert_id << 16) | (weight_bf16.view(int16)).

  • routing_bias (Optional[torch.Tensor]) – [num_experts] tensor of routing bias. May be None.

  • hidden_states (torch.Tensor) – [seq_len, hidden_size] tensor of input hidden states.

  • hidden_states_scale (torch.Tensor) – [seq_len, 1] FP32 per-token dequantization multipliers.

  • gemm1_weights (torch.Tensor) – [num_experts, M, hidden_size] FP8 first-layer weights, where M is 2 * intermediate_size for gated activations and intermediate_size otherwise.

  • gemm1_per_channel_weight_scale (torch.Tensor) – [local_num_experts, M] per-channel weight dequantization multipliers for GEMM1, in the same shuffled row order as gemm1_weights.

  • output1_scale_scalar (torch.Tensor) – [local_num_experts] per-expert output scales for GEMM1.

  • output1_scale_gate_scalar (torch.Tensor) – [local_num_experts] per-expert gate scales for GEMM1.

  • gemm2_weights (torch.Tensor) – [num_experts, hidden_size, intermediate_size] FP8 second-layer weights.

  • gemm2_per_channel_weight_scale (torch.Tensor) – [local_num_experts, hidden_size] per-channel dequantization multipliers for GEMM2, in the same shuffled row order as gemm2_weights.

  • output2_scale_scalar (torch.Tensor) – [local_num_experts] per-expert output scales for GEMM2.

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

  • use_routing_scales_on_input (bool) – Whether to apply routing scales to the input (default False).

  • routing_method_type (int) – Routing method (default 0). Matches flashinfer.tllm_enums.RoutingMethodType.

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

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

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

  • activation_type (int) – Activation type (default 3 — Swiglu).

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]