flashinfer.fused_moe.BGMVMoEBlackwellPlan

class flashinfer.fused_moe.BGMVMoEBlackwellPlan(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)

Pointer-stable SM100 BGMV MoE shrink+expand execution plan.

The plan owns caller-visible FP32 accumulation and shrink workspaces. Its first eager run captures the exact launch sequence into a CUDA Graph; later calls replay that graph on the same stream. If called while an outer CUDA Graph is being captured, the constituent kernels are enqueued directly.

__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) → 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.