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['blackwell'] = 'blackwell', shrink_out: Tensor | None = None, y_accum: Tensor | None = None) → BGMVMoEBlackwellPlan

Prepare the generated SM100 BGMV MoE pipeline for graph replay.

This optimized path currently supports one LoRA slice, rank 32, hidden sizes 2688 or 3072, BF16/FP16 inputs, and exact SM100 devices. 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.

Parameters:

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

  • lora_a_weights – One LoRA-A tensor with shape [num_loras, num_experts, 32, hidden_size].

  • lora_b_weights – One LoRA-B tensor with shape [num_loras, num_experts, hidden_size, 32].

  • 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 both LoRA tensors.

  • backend – Backend selector. Only "blackwell" is supported.

  • shrink_out – Optional pointer-stable FP32 shrink workspace.

  • y_accum – Optional pointer-stable FP32 output accumulator.

Returns:

A reusable graph-backed execution plan whose run method returns the FP32 accumulated output.