flashinfer.fused_moe.cute_dsl.cute_dsl_fused_moe_bf16

flashinfer.fused_moe.cute_dsl.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, activation_type: int = 3, swiglu_alpha: float = 1.0, swiglu_beta: float = 0.0, swiglu_limit: float = 3.4028234663852886e+38, situ_beta: float | None = None, situ_linear_beta: float | None = None, *, intermediate_buffer: Tensor | None = None, tactic: Tuple[Any, ...] | 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, for the default SwiGLU, 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, fused into GEMM1: ActivationType.Swiglu (default; the OAI variant via swiglu_alpha/swiglu_beta/swiglu_limit, SiTU via situ_beta), ActivationType.GegluTanh, or the non-gated ActivationType.Relu2.

  • 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 % 8 == 0; I % 32 == 0 for gated activations (32-column up/gate interleave), I % 8 == 0 for Relu2; partial last K and N tiles are handled; 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 tactic the runner offers for the token bucket (get_valid_tactics()) is profiled and the per-bucket winner is cached; outside it the cached winner (or the fixed default, DEFAULT_SM90_MOE_TACTIC: tile 128, 64-wide N tiles, cluster (1, 1), N-major walks) dispatches. An explicit tactic bypasses 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 – Gated activations: [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(). Relu2: [num_local_experts, I, hidden], no interleave.

  • 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.

  • activation_type – GEMM1 activation, see the supported configuration above. Together with the constants below it is part of the kernel compile key and of the autotuner cache key.

  • swiglu_alpha – SwiGLU sigmoid multiplier (Swiglu only).

  • swiglu_beta – SwiGLU up-projection bias (Swiglu only).

  • swiglu_limit – SwiGLU clamp limit (Swiglu only).

  • situ_beta – With Swiglu, selects the SiTU gate beta * tanh(gate / beta) * sigmoid(gate).

  • situ_linear_beta – Optional SiTU tanh clamp of the up branch (requires situ_beta).

  • intermediate_buffer – Optional pre-allocated GEMM1 output buffer (advanced, keyword-only).

  • tactic – Optional (tile_size, gemm1_tactic, gemm2_tactic) tuple (see sm90_tuner) dispatched as given, validated by the kernel wrappers; None selects through the AutoTuner.

Returns:

[num_tokens, hidden] in x’s dtype.