flashinfer.fused_moe.prims_ts_fp8_per_tensor_scale_moe_partitioned

flashinfer.fused_moe.prims_ts_fp8_per_tensor_scale_moe_partitioned(routing_logits: Tensor | None, routing_bias: Tensor | None, topk_ids: Tensor | None, expert_weights: Tensor | None, hidden_states: Tensor, gemm1_weights: Sequence[Tensor], output1_scales_scalar: Tensor, output1_scales_gate_scalar: Tensor, gemm2_weights: Sequence[Tensor], output2_scales_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, partition_resources: MoePartitionResources, partition_strategy: Literal['tp_n_tp_n', 'tp_n_tp_k'], routed_scaling_factor: float | None = 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, norm_topk_prob: bool = True, routing_replay_out: Tensor | None = None, output: Tensor | None = None, *, weight_layout: int = 0) → List[Tensor] | Tensor

Run partitioned FP8 per-tensor MoE with caller-owned payloads.

Same computation as prims_ts_fp8_per_tensor_scale_moe(), split across two green-context partitions on SM107 (Rubin) only. Per-channel weight-scale vectors and gemm1_clamp_limit are not supported.

Parameters:

  • routing_logits (Optional[torch.Tensor]) – [seq_len, num_experts] routing logits, or None when using precomputed routing.

  • routing_bias (Optional[torch.Tensor]) – Optional [num_experts] routing bias.

  • topk_ids (Optional[torch.Tensor]) – int32 precomputed routing when routing_logits is None: packed (expert_id, weight) values when expert_weights is None, otherwise unpacked expert indices.

  • expert_weights (Optional[torch.Tensor]) – Precomputed routing weights for unpacked topk_ids.

  • hidden_states (torch.Tensor) – float8_e4m3fn activations.

  • gemm1_weights (Sequence[torch.Tensor]) – Two float8_e4m3fn FC1 shards, one per partition, sharded with flashinfer.prims_ts.moe.shard_prepared_moe_weights.

  • output1_scales_scalar (torch.Tensor) – Per-expert FC1 output scales.

  • output1_scales_gate_scalar (torch.Tensor) – Per-expert FC1 gate scales.

  • gemm2_weights (Sequence[torch.Tensor]) – Two float8_e4m3fn FC2 shards.

  • output2_scales_scalar (torch.Tensor) – Per-expert FC2 output scales.

  • num_experts (int) – Total number of experts.

  • top_k (int) – Experts selected per token.

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

  • topk_group (Optional[int]) – Groups considered for top-k routing.

  • intermediate_size (int) – Unsharded intermediate (FFN) width.

  • local_expert_offset (int) – Global offset of the first local expert.

  • local_num_experts (int) – Number of experts resident on this device.

  • partition_resources (MoePartitionResources) – Caller-owned green-context streams and reusable fork/join events. Construct before CUDA graph capture and keep alive for every launch or replay that uses it.

  • partition_strategy (PartitionStrategy) – "tp_n_tp_n" splits FC2 output rows; "tp_n_tp_k" splits FC2 K. FC1 rows are split in both. Must match the weight sharding.

  • routed_scaling_factor (Optional[float]) – Optional routing scale.

  • use_routing_scales_on_input (bool) – Apply routing scales on the input path when True; requires routing_logits.

  • routing_method_type (int) – Routing method selector (default 0).

  • do_finalize (bool) – If True, return the finalized MoE output.

  • enable_pdl (Optional[bool]) – Enable Programmatic Dependent Launch when supported.

  • tune_max_num_tokens (int) – Autotune token-bucket upper bound (default 8192).

  • activation_type (int) – Activation enum value (default Swiglu).

  • norm_topk_prob (bool) – Normalize top-k routing probabilities.

  • routing_replay_out (Optional[torch.Tensor]) – Optional buffer that captures selected expert IDs.

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

  • weight_layout (int) – Prims-TS weight layout enum value (default MajorK). Keyword-only.

Returns:

Same return contract as prims_ts_fp8_per_tensor_scale_moe().

Return type:

torch.Tensor or List[torch.Tensor]