flashinfer.fused_moe.mono_moe

flashinfer.fused_moe.mono_moe(activations_in: Tensor, router_logits: Tensor, expert_weights_up: Tensor, expert_scales_up: Tensor, expert_weights_down: Tensor, expert_scales_down: Tensor, top_k: int, scoring_func: str = 'softmax', renormalize: bool = True, out: Tensor | None = None, scratchpad: Tensor | None = None, interleave_up: bool = True) → Tensor

Single-kernel block-FP8 top-K MoE (Qwen3.5-35B shape, SM90a only).

Fixed shape: E=256 experts, N=512, K=2048 hidden, up to BS=8 tokens.

Parameters:

  • activations_in – bf16 input activations [M, K] (M <= 8).

  • router_logits – bf16 router logits [M, E].

  • expert_weights_up – fp8_e4m3 up/gate weights [E, 2*N, K]. By default this function applies interleave_for_tma_wgmma_up(); pass interleave_up=False if the tensor is already interleaved.

  • expert_scales_up – fp32 block-wise scales [E, 2N/128, K/128].

  • expert_weights_down – fp8_e4m3 down weights [E, K, N] (raw row-major).

  • expert_scales_down – fp32 block-wise scales [E, K/128, N/128].

  • top_k – experts selected per token (1..8).

  • scoring_func – "sigmoid" or "softmax".

  • renormalize – renormalize the top-K weights to sum to 1.

  • out – optional bf16 output buffer [M, K]; allocated if omitted.

  • scratchpad – optional reusable uint8 scratchpad from alloc_scratchpad(); allocated per-call if omitted.

  • interleave_up – apply the gate/up TMA repack to expert_weights_up (default True).

Returns:

bf16 MoE output [M, K].