flashinfer.fused_moe.prepare_bgmv_moe

flashinfer.fused_moe.prepare_bgmv_moe(x: Tensor, lora_a_weights: List[Tensor], lora_b_weights: List[Tensor], sorted_token_ids: Tensor, expert_ids: Tensor, lora_indices: Tensor, topk_weights: Tensor, num_experts: int, *, backend: Literal['cake'] = 'cake', fallback: bool = True, shrink_out: Tensor | None = None, y_accum: Tensor | None = None, grouped: bool | None = None, order_remap: bool | None = None) → BGMVMoECakePlan | BGMVMoEPortablePlan

Prepare a graph-replayable BGMV MoE shrink+expand pipeline.

The generated Cake path supports one LoRA slice, LoRA rank 8, 16, 32 or 64, any hidden size that is a positive multiple of 8 (LoRA-B feat_out equal to it), BF16/FP16 inputs, and exact SM90 (H100/H200), SM100 (B200/GB200) or SM103 (B300/GB300) devices; each target runs its own cubin. Hidden sizes 2688 and 3072 at rank 32 use the specialized measured bodies at up to 2048 tokens (plan.variant == "specialized"); everything else, including those shapes at larger token counts, uses the runtime-hidden generic bundles (plan.variant == "generic"). Routing may be arbitrary; each output has one owner that accumulates routes in fixed input order, so identical prepared replays are bitwise reproducible. The contiguous top-k=2 layout takes the optimized fast path; any other pair order (for example expert-sorted dispatch) is served through a token->pair route index that the shrink kernels publish and the expand kernels read in O(1) per CTA (plan.route_index, 16 slots per token; tokens routed to more pairs take an exact serial scan). Pairs already in their contiguous position are implicit, so a contiguous launch publishes nothing. At small pair counts the generic shrink also splits the hidden dimension over extra CTAs (plan.shrink_launch) and the last CTA of each output tile reduces the FP32 partials in split order, so the result stays deterministic and the workspace is never reset.

Inputs outside that support set (other device capabilities, ranks, hidden sizes that are not multiples of 8, multiple slices) are served by a BGMVMoEPortablePlan running the portable bgmv_moe_shrink + bgmv_moe_expand kernels when fallback=True (default); one RuntimeWarning is emitted per distinct reason per process. With fallback=False such inputs raise ValueError as before. Invalid inputs (shape mismatches, wrong index dtypes, out-of-range routing indices, non-contiguous or CPU tensors) always raise.

Parameters:

  • x – Input activations with shape [num_tokens, hidden_size].

  • lora_a_weights – LoRA-A tensors, one per slice, each with shape [num_loras, num_experts, rank, hidden_size].

  • lora_b_weights – LoRA-B tensors, one per slice, each with shape [num_loras, num_experts, feat_out, rank].

  • sorted_token_ids – Routed token indices with shape [num_pairs].

  • expert_ids – Expert index for each routed pair.

  • lora_indices – LoRA index for each input token.

  • topk_weights – FP32 routing weight for each routed pair.

  • num_experts – Number of experts in the LoRA tensors.

  • backend – Backend selector; only "cake" (the generated Cake programs) is supported.

  • fallback – Serve unsupported inputs with the portable path instead of raising.

  • grouped – Generic Cake variant only. None (default) lets the plan choose the pair-grouped pipeline when the routes clearly outnumber the num_loras * num_experts (LoRA, expert) pairs (each pair’s weights are then streamed once per tile of routes and a deterministic per-token combine sums the FP32 route partials in ascending pair order; the plan owns an int32 grouping workspace and num_pairs * hidden FP32 partials). True / False force or disable it.

  • order_remap – Generic per-route Cake pipeline only. None (default) lets the plan dispatch the shrink in (LoRA, expert)-bin order on SM90 where the saved LoRA-A traffic outweighs the single-CTA ordering prologue (select_cake_bgmv_moe_order_remap); every route is still computed by the same code on the same operands, so the result is bitwise identical to the identity dispatch. True / False force or disable it; True is rejected together with the grouped pipeline or the specialized variant.

  • shrink_out – Optional pointer-stable shrink workspace with shape [num_slices, num_pairs, rank] and the weight dtype.

  • y_accum – Optional pointer-stable FP32 output accumulator with shape [num_tokens, sum(feat_out)]. The Cake path writes it through its row stride, so it may be a column slice of a wider row-major buffer; the portable fallback needs a contiguous accumulator.

Returns:

A reusable graph-backed execution plan whose run method returns the FP32 accumulated output; plan.backend_used is "cake" or "portable".