flashinfer.fused_moe.cute_dsl_fused_moe_bf16

flashinfer.fused_moe.cute_dsl_fused_moe_bf16(x: Tensor, token_selected_experts: Tensor, token_final_scales: Tensor, w1_weight: Tensor, w2_weight: Tensor, num_experts: int, top_k: int, num_local_experts: int | None = None, local_expert_offset: int = 0, use_fused_finalize: bool = True, moe_output: Tensor | None = None, enable_pdl: bool = True, *, intermediate_buffer: Tensor | None = None, tile_size: int | None = None, gemm1_tile_n: int | None = None, gemm2_tile_n: int | None = None, gemm2_tile_k: int | None = None, gemm2_cluster_shape_mn: Tuple[int, int] | None = None, gemm2_raster_along_m: bool | None = None) → Tensor

SM90 CuTe-DSL fused MoE forward (BF16/FP16, unquantized).

out[t] = sum_k scale[t,k] * ffn_expert(x[t]; e[t,k]) with ffn(x; e) = (silu(x @ w1_gate[e].T) * (x @ w1_up[e].T)) @ w2_weight[e].T.

Supported configuration:

  • Arch: SM90 (Hopper) only.

  • Dtypes: bf16 or fp16 activations and weights (must match), fp32 accumulation; output dtype = input dtype. No quantized paths.

  • Activation: SwiGLU (SiLU-gated) only, fused into GEMM1.

  • Routing: pre-routed contract only — the caller runs the router and passes global expert ids plus normalized scales. top_k is a compile-time constant of the kernels.

  • Parallelism: TP by weight shapes; EP via num_local_experts + local_expert_offset (tokens routed entirely outside the local shard contribute zeros).

  • Shapes: hidden % 64 == 0 and (GEMM1 reduction) no tile-32 fallback, 2I % 64 == 0, I % 32 == 0 (weight interleave and GEMM2 tile-k 32 fallback for I % 64 != 0); num_tokens == 0 is supported.

  • Execution: CUDA-graph capturable; PDL on by default; fused finalize (default) is atomic and not bitwise-reproducible — use_fused_finalize=False selects the deterministic path.

Tile selection goes through the FlashInfer AutoTuner. Under the autotune() context every enumerated Sm90MoeTactic (capped at the top-2 legal N tiles per GEMM) is profiled and the per-bucket winner is cached; outside it the cached winner (or the heuristic auto-selection, as the default tactic) dispatches. Explicit tile / cluster / raster / buffer keyword overrides bypass the tuner:

with autotune(True):
    output = cute_dsl_fused_moe_bf16(...)
Parameters:

  • x – [num_tokens, hidden] bf16/fp16.

  • token_selected_experts – [num_tokens, top_k] int32.

  • token_final_scales – [num_tokens, top_k] float32, normalized by the caller.

  • w1_weight – [num_local_experts, 2I, hidden] — up/gate interleaved at 32 columns. Callers may cache this repack; the in-tree reference is interleave_up_gate_sm90().

  • w2_weight – [num_local_experts, hidden, I].

  • num_experts – Total (global) expert count.

  • top_k – Experts per token.

  • num_local_experts – Experts held by this rank (EP shard); defaults to num_experts.

  • local_expert_offset – Global id of this shard’s first expert.

  • use_fused_finalize – True (default) fuses the router-scaled scatter-reduce into GEMM2. The top-k combine then accumulates in the output dtype (cp.reduce.async.bulk.add): one output-dtype rounding per route on top of the bf16/fp16 intermediate hand-off, and not bitwise-reproducible across runs. False uses the deterministic two-stage path: GEMM2 scatters unscaled rows in expanded (token, slot) order, then moe_unpermute applies the scales and combines in float32 in a fixed order — one final rounding, at the cost of an extra kernel and the expanded intermediate.

  • moe_output – Optional pre-allocated [num_tokens, hidden] output (contents overwritten; zeroed internally for the fused finalize).

  • enable_pdl – True (default) launches both GEMMs (and the deterministic path’s moe_unpermute) with Programmatic Dependent Launch so each kernel’s prologue overlaps its predecessor’s tail. Numerics are unaffected. Part of the kernel compile cache key.

  • intermediate_buffer – Optional pre-allocated GEMM1 output buffer (advanced, keyword-only; bypasses the tuner like the tile overrides).

  • tile_size – Tile size shared by moe_sort and both GEMMs (64 or 128; keyword-only). Default None auto-selects: 64 below an average of 64 rows per local expert (num_tokens * top_k / num_local_experts) — small/decode batches pad each expert to a full M-tile, so the smaller tile roughly halves the wasted MMA work; 128 from one full tile per expert up (fatter tiles amortize per-tile fixed costs and B loads). Tiny reductions (per-rank I < 192) switch to 128 already at 16 rows per expert.

  • gemm1_tile_n – N tile for GEMM1 (None auto-selects; keyword-only).

  • gemm2_tile_n – N tile for GEMM2 (None auto-selects; keyword-only).

  • gemm2_tile_k – K tile for GEMM2, 64 or 32 (None auto-selects via _default_gemm2_tile_k(): 32 when 64 does not divide the per-rank I, and on prefill tiles with I >= 384 where its doubled pipeline depth wins; 64 otherwise. Keyword-only).

  • gemm2_cluster_shape_mn – GEMM2 CTA cluster shape, (1, 1) or (1, 2) (None applies the fallback heuristic; keyword-only). (1, 2) requires an even GEMM2 N-tile count.

  • gemm2_raster_along_m – GEMM2 tile raster order (None auto-selects; keyword-only). M-major confines the finalize scatter-RMW working set to one L2-resident output column band per CTA wave, at the cost of re-reading each A tile once per N tile — the default enables it only where that trade wins: prefill tiles with a large output working set (>= 32 MiB for per-rank I <= 192, >= 64 MiB above) and per-rank I <= 384.

Returns:

[num_tokens, hidden] in x’s dtype.