flashinfer.experimental.cake_stepfun_moe.stepfun_moe

flashinfer.experimental.cake_stepfun_moe.stepfun_moe(*, family: str, hidden_states: Tensor, hidden_states_scale: Tensor | None = None, gemm1_weights: Tensor, gemm1_weights_scale: Tensor | None = None, output1_scale_scalar: Tensor | None = None, output1_scale_gate_scalar: Tensor | None = None, gemm1_clamp_limit: Tensor, per_token_scale: Tensor | None = None, gemm2_weights: Tensor, gemm2_weights_scale: Tensor | None = None, output2_scale_scalar: Tensor | None = None, routing_logits: Tensor | None = None, topk_ids: Tensor | None = None, topk_weights: Tensor | None = None, num_experts: int, top_k: int, local_expert_offset: int, local_num_experts: int, plan: Any = None, output: Tensor | None = None, do_finalize: bool = True, enable_pdl: bool = True, reserved_sms: int = 8, global_scale_inv: float = 0.0003720238095238095, poison_routing_tail: bool = False) → list[Tensor]

Warning

cake_stepfun_moe.stepfun_moe is experimental: it provides no compatibility guarantees and may change or be removed without deprecation.

One StepFun fused-MoE forward on the exported Cake kernels.

Prepares and launches in one call: routing -> FC1 -> (requantization) -> FC2 -> finalize on the current torch stream. family is one of nvfp4, nvfp4_bf16tok, bf16, fp8, mxfp8; the tensors are the trtllm-gen weight / activation views of that family (see the package README), gemm1_clamp_limit the per-expert fp32 StepFun limits in the units the kernels clamp, and exactly one of routing_logits (bf16 or fp32 [num_tokens, num_experts]) or topk_ids + topk_weights (int32 / bf16-or-fp32 [num_tokens, top_k]) selects the routing input. plan is None (the planner’s centre tile), a tile or a Plan.

Returns [output] (bf16 [num_tokens, hidden_size]) when do_finalize, else [gemm2_output, expert_weights, expanded_idx_to_permuted_idx] in the trtllm-gen layout. For repeated launches over the same tensors (or CUDA Graph capture) use prepare_stepfun_moe once and call launch() on the returned object.