flashinfer.fused_moe.trtllm_fp8_per_channel_scale_moe

flashinfer.fused_moe.trtllm_fp8_per_channel_scale_moe(routing_logits: 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, routing_method_type: int = 0, do_finalize: bool = True, enable_pdl: bool | None = None, tune_max_num_tokens: int = 8192, activation_type: int = 3, norm_topk_prob: bool = True) → List[Tensor] | Tensor

FP8 per-channel scale MoE operation.

Parameters:

  • routing_logits – [seq_len, num_experts] tensor of routing logits

  • routing_bias – [num_experts] tensor of routing bias

  • hidden_states – [seq_len, hidden_size] tensor of input hidden states

  • hidden_states_scale – [seq_len, 1] FP32 per-token dequantization multipliers

  • gemm1_weights – [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 – [local_num_experts, M] per-channel weight dequantization multipliers for gemm1, in the same shuffled row order as gemm1_weights

  • output1_scale_scalar – [local_num_experts] per-expert output scales for gemm1

  • output1_scale_gate_scalar – [local_num_experts] per-expert gate scales for gemm1

  • gemm2_weights – [num_experts, hidden_size, intermediate_size] FP8 second layer weights

  • gemm2_per_channel_weight_scale – [local_num_experts, hidden_size] per-channel dequantization multipliers for gemm2, in the same shuffled row order as gemm2_weights

  • output2_scale_scalar – [local_num_experts] per-expert output scales for gemm2

  • num_experts – Total number of experts

  • top_k – Number of experts to route to per token

  • n_group – Number of expert groups

  • topk_group – Number of groups to consider for top-k routing

  • intermediate_size – Size of intermediate layer

  • local_expert_offset – Offset of local experts in global expert space

  • local_num_experts – Number of experts handled by this device

  • routed_scaling_factor – Scaling factor for routing

  • use_routing_scales_on_input – Whether to use routing scales on input

  • routing_method_type – Type of routing method to use (default: 0)

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

  • enable_pdl – Whether to enable Programmatic Dependent Launch (PDL). Auto-enabled for >= sm90.

  • tune_max_num_tokens – Maximum number of tokens for tuning. (default: 8192)

  • activation_type – Type of activation function (default: 3 - Swiglu)

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

Returns:

when do_finalize=True, returns the final MoE output. otherwise, returns the intermediate results (gemm2_output, expert_weights, expanded_idx_to_permuted_idx).