flashinfer.fused_moe.BGMVMoECakePlan

class flashinfer.fused_moe.BGMVMoECakePlan(module, *, y_accum: Tensor, shrink_out: Tensor, x: Tensor, lora_a: Tensor, lora_b: Tensor, sorted_token_ids: Tensor, expert_ids: Tensor, lora_indices: Tensor, topk_weights: Tensor, schedule_id: int, variant: Literal['specialized', 'generic'] = 'specialized', shrink_launch: Tuple[int, int] | None = None, grouped: bool = False, order_remap: bool = False, pdl_mode: int = 0)

Pointer-stable SM90/SM100/SM103 Cake BGMV MoE shrink+expand execution plan.

Runs the generated Cake programs (one owner per output token, no output atomics, bitwise-reproducible replays). backend_used is "cake"; variant is "specialized" for the hidden 2688/3072 x rank 32 bodies (up to 2048 tokens) and "generic" for the runtime-hidden rank 8/16/32/64 bundles.

__init__(module, *, y_accum: Tensor, shrink_out: Tensor, x: Tensor, lora_a: Tensor, lora_b: Tensor, sorted_token_ids: Tensor, expert_ids: Tensor, lora_indices: Tensor, topk_weights: Tensor, schedule_id: int, variant: Literal['specialized', 'generic'] = 'specialized', shrink_launch: Tuple[int, int] | None = None, grouped: bool = False, order_remap: bool = False, pdl_mode: int = 0) → None

Methods

__init__(module, *, y_accum, shrink_out, x, ...)
close()	Release graph resources after pending replay work completes.
run()	Run or replay the prepared zero+shrink+expand pipeline.

Attributes

backend_used