flashinfer.fused_moe¶

This module provides fused Mixture-of-Experts (MoE) operations optimized for different backends and data types.

Types and Enums¶

`RoutingMethodType`(value)	An enumeration.
`WeightLayout`(value)	An enumeration.

Shared activation helpers live in flashinfer.tllm_enums and are used by both the TRT-LLM and CuteDSL MoE paths.

is_gated_activation(activation_type)

Return whether the given activation type is a gated activation (e.g. SwiGLU family).

Unified MoE API¶

Backend-agnostic configuration and layer types. QuantConfig carries the MMA weight / activation formats and the layer output format as QuantFormat axes, plus the NVFP4 4over6 recipe (nvfp4_4over6); passing a recipe there is the supported alternative to the process-wide FLASHINFER_NVFP4_4OVER6* environment variables, described under flashinfer.quantization.

`MoEConfig`(routing, quant, experts, ...)	Top-level MoE configuration.
`RoutingConfig`(num_experts, top_k[, method, ...])	Expert routing parameters.
`QuantConfig`(weight, activation, output, *, ...)	Quantization scheme: MMA weight/activation formats plus the result format.
`QuantFormat`(value)	Numeric format of one MMA operand or of the MoE layer output.
`ExpertConfig`(intermediate_size[, ...])	Expert geometry.
`ExecutionConfig`([enable_pdl, ...])	Runtime execution parameters.
`MoEFinalizeConfig`([do_finalize, ...])	How the finalize (combine) step behaves.
`BackendOptions`([candidates])	Ordered list of backend candidates for dispatch and autotuning.
`MoEActivationPack`(hidden_states_q, ...[, ...])	Per-call backend-native activations plus routing inputs.
`MoEWeightPack`(native_views, ...)	Long-lived weight container with per-backend native materializations.
`CudnnFrostBf16Config`()	Frost BF16 MoE on SM107 and SM120, with packed precomputed routing.
`CudnnFrostMxfp8Config`()	Frost MXFP8 MoE on SM107, with packed precomputed routing.
`CudnnFrostNvfp4Config`()	Frost NVFP4 MoE on SM107, with packed precomputed routing.
`CudnnFrostMxfp8Mxfp4Config`()	Frost MXFP8-activation / MXFP4-weight MoE on SM107.

MoELayer is the official entry point of this API: both its constructor and its call operator are decorated with @flashinfer_api, so they participate in FLASHINFER_LOGLEVEL logging and FLASHINFER_DUMP_* capture. The lower-level per-backend functions above remain official in their own right — the two layers are designed to co-exist, and neither supersedes the other.

class flashinfer.fused_moe.MoELayer(config: MoEConfig, device: device | None = None)¶

Bases: object

Stateful MoE layer with cross-backend autotune.

TRTLLM runners bind their immutable launch metadata to each packed input list, so interleaved pack_inputs -> forward pairs cannot exchange weights or routing configuration. Other backend adapters may still retain per-call workspace or prepared-weight state; use one MoELayer per thread/stream until those adapters adopt the same convention.

Example

>>> layer = MoELayer(config)
>>> out = layer(act_pack, weight_pack)

__init__(config: MoEConfig, device: device | None = None)¶

Build the layer and the set of backend runners it will autotune over.

Every backend in config.backend whose hardware preconditions the target device satisfies gets a runner; the rest are skipped here rather than at call time, so an unsupported backend costs nothing per call.

Parameters:

config (MoEConfig) – Routing, quantization, expert geometry, activation, backend candidates and execution parameters. Cross-field constraints are validated by MoEConfig itself.
device (torch.device or None) – Device whose compute capability selects the usable backends. None → the current CUDA device.

Raises:

RuntimeError – If no configured backend is usable on this device’s architecture.

__call__(act_pack: MoEActivationPack, weight_pack: MoEWeightPack) → Tensor | List[Tensor]¶

Run the MoE layer, selecting and caching the fastest usable backend.

Only runners that support this pack’s routing_input_mode compete — not every backend has an in-kernel router — and the winner is cached per (token-count bucket, mode) so repeated calls skip reselection.

Parameters:

act_pack (MoEActivationPack) – Backend-native activations plus routing inputs for this call.
weight_pack (MoEWeightPack) – Expert weights, prepared for the selected backend’s native layout.

Returns:

The layer output. With config.finalize.do_finalize=False the unreduced TRTLLM intermediates are returned instead, as [gemm2_output, expert_weights, expanded_idx_to_permuted_idx], leaving the combine to the caller.

Return type:

torch.Tensor or list of torch.Tensor

Raises:

ValueError – If act_pack.num_tokens exceeds the execution.tune_max_num_tokens ceiling the layer was built with.
NotImplementedError – If no usable backend supports this pack’s routing_input_mode.

refresh_backends() → None¶

Refresh backend preparation and clear cached winners.

Call outside CUDA Graph capture after changing compiler or artifact configuration. Backends with fixed preparation snapshots adopt the current configuration; active packed calls and captured graphs keep their existing resources. The next layer call selects a winner again.

reset_winner() → None¶: Clear all cached per-bucket winners — next call re-tunes.

property winner_backend: str | None¶: Backend key selected on the most recent call, or None before first call.

cuDNN Frost backend¶

Frost is a stable MoELayer backend on SM107a, with additional BF16 support on SM120a. Its four configs are included in the default backend list and can also be selected through BackendOptions. Candidates with unavailable weight views or unsupported input shapes are skipped. It does not require FLASHINFER_ALLOW_EXPERIMENTAL_AUTO_BACKENDS.

Supported configurations¶
Config	Activation / weight format	Native weight-view key
`CudnnFrostBf16Config`	BF16 / BF16	`cudnn_frost_bf16`
`CudnnFrostMxfp8Config`	MXFP8 / MXFP8	`cudnn_frost_mxfp8`
`CudnnFrostNvfp4Config`	NVFP4 / NVFP4	`cudnn_frost_nvfp4`
`CudnnFrostMxfp8Mxfp4Config`	MXFP8 / MXFP4	`cudnn_frost_mxfp8_mxfp4`

Each config exposes prepare_weights and prepare_activations. Frost also accepts the corresponding canonical cutlass_* weight view without conversion; when both exist, the Frost-named view takes precedence. The public API is:

from flashinfer.fused_moe import (
    BackendOptions, CudnnFrostNvfp4Config, ExpertConfig, MoEActivationPack,
    MoEConfig, MoELayer, MoEWeightPack, QuantConfig, QuantFormat, RoutingConfig,
)

backend = CudnnFrostNvfp4Config
config = MoEConfig(
    routing=RoutingConfig(num_experts=num_experts, top_k=top_k),
    experts=ExpertConfig(intermediate_size=intermediate_size),
    quant=QuantConfig(weight=QuantFormat.NVFP4, activation=QuantFormat.NVFP4),
    backend=BackendOptions((backend(),)),
)
view = backend.prepare_weights(
    w1_bf16, w2_bf16, num_local_experts=num_experts,
    hidden_size=hidden_size, intermediate_size=intermediate_size,
    activation=config.activation,
)
weights = MoEWeightPack({"cudnn_frost_nvfp4": view})
xq, xsf = backend.prepare_activations(x_bf16, quant=config.quant)
activations = MoEActivationPack(xq, xsf, topk_ids, topk_weights)
layer = MoELayer(config)
output = layer(activations, weights)

The backend requires precomputed routing with contiguous int32 expert IDs and FP32 routing weights, and returns finalized BF16 output. It supports the default activation contracts; bias, expert parallelism, shared experts, custom activation parameters, and routing from logits are unsupported. Hidden/intermediate sizes must be positive multiples of 128. All four dtypes use measured model shortlists where available and a bounded, factorized tactic search for other supported shapes. Local expert counts are limited to 1024 and token counts to the configured tuning ceiling, at most 2**20. Kernel identities describe implementation content, independently of model dimensions.

Quantized FC1 weight scales accept expert-major or branch-major storage and may be updated in place. For CUDA Graph replay, keep tensor pointers, shapes and layouts fixed and update tensors before replay on the same stream.

After changing compiler or artifact configuration, call layer.refresh_backends() outside CUDA Graph capture, or reconstruct the layer.

The validated SM107 dependency baseline is PyTorch 2.14.1, CuTe DSL 4.8.0, TVM-FFI 0.1.14.post1, and external PTXAS 13.4.92. CUDA 13.5 PTXAS is also validated. Runtime checks additionally require the DSL primitives and PyTorch CUDA Graph resource-retention APIs used by the backend. Older dependencies remain usable by other FlashInfer backends; these requirements do not raise FlashInfer’s global dependency floor. Native JIT compilation also needs a complete CUDA toolkit for the target GPU, including cuBLAS headers used by the shared CUTLASS preparation helpers. A standalone PTXAS executable only supplies the assembler for Frost’s DSL kernels.

On SM107, set FLASHINFER_CUDNN_FROST_PTXAS=/path/to/ptxas when the DSL bundles an older assembler. CUDA_HOME alone does not replace that assembler. Unqualified assemblers are rejected when building or refreshing a runner. Automatic selection skips unavailable Frost candidates; explicit selection reports the missing requirement. Both external PTXAS and the bundled assembler must report CUDA 13.4 or newer. There is no patch-level minimum. Use the latest assembler patch release for SM107 tensor-map updates. SM120 BF16 kernels do not update tensor maps and retain capability-based compiler admission without this SM107-specific assembler minimum.

Prepare and optionally autotune before CUDA Graph capture. Frozen device sources ship in the wheel and are compiled for the requested geometry into the writable FlashInfer cache. The four fixed routing/finalize adapters are also registered for SM107 AOT builds, and the BF16 adapter is registered for SM120 AOT builds. Geometry-specialized DSL kernels retain source/JIT delivery. No runtime cuDNN Frontend installation is required.

Artifact catalogs use schema v4. Each architecture shares tile definitions; gated/non-gated activation groups and compatible store policies expand into kernel descriptors during preparation. Sources share exact activation fragments. Device occupancy, L2 budget, pipeline stages and workspace are derived once from the preparation device. Kernels support symbolic problem sizes with explicit alignment and bounds; measured shortlists accelerate selection for known shapes, and bounded heuristics supply candidates for other shapes.

Kernel IDs identify logical configurations. Source and compiler fingerprints qualify compiled kernels and persisted tuning choices. Artifact tooling should use cudnn_frost.manifest.read_manifest and write_manifest to read expanded records and preserve existing dimension catalogs. The reader also accepts schema v2/v3 exports; dimension_catalog.write_catalog factors complete exports into the dimension format.

Finalized precomputed MoELayer calls support FLASHINFER_TRACE_DUMP=1 and MoELayer.__call__.fi_trace(self=layer, act_pack=activations, weight_pack=weights). Traces preserve packed tensor shapes, dtypes, native weight views, and the layer configuration. FLASHINFER_LOGLEVEL continues to control API logging.

The old flashinfer.experimental.cudnn_frost_selected_kernels_moe_grouped_gemm imports are compatibility aliases for one release after graduation. The separate cudnn_frost_grouped_gemm1_swiglu API remains experimental.

Utility Functions¶

`convert_to_block_layout`(input_tensor, blockK)	Reshape a 2-D tensor into a 3-D block layout.
`reorder_rows_for_gated_act_gemm`(x)	Reorder rows of a weight tensor for the TensorRT-LLM gated-activation GEMM layout.
`interleave_moe_weights_for_sm90_mixed_gemm`(weight)	Interleave 4-bit packed MoE weights for the SM90 mixed-input GEMM.
`interleave_moe_scales_for_sm90_mixed_gemm`(scales)	Fold weight scales for the SM90 mixed-input MoE GEMM.
`preprocess_moe_weights_for_sm90_mixed_gemm_humming`(...)	Prepare MXFP4 weights for the SM90 Humming-style FP8 activation path.
`fused_topk_deepseek`(scores, bias, n_group, ...)	Fused expert routing with top-k selection for DeepSeek-V3.
`hash_topk`(router_logits, input_ids, tid2eid)	Hash-based MoE expert routing for DeepSeek-V4.

The E8M0 range-clamping, residual-scale factorization, and FP4 payload-rewrite scheme used by preprocess_moe_weights_for_sm90_mixed_gemm_humming is adapted from Humming.

AlphaMoE Router (SM100/SM103)¶

The standalone AlphaMoE frontend converts FP32 logits into top-k weights and an expert-grouped, block-aligned route plan. Selection and ordering among exactly equal routed logits are unspecified. The resulting plan can feed either the AlphaMoE W8A8 or NVFP4 compute path and can be reused across launches to avoid steady-state allocation.

`AlphaMoERoutePlan`(topk_weights, topk_ids, ...)	Device-resident routing outputs and reusable private workspace.
`allocate_alphamoe_route_plan`(logits, *, top_k)	Allocate a reusable `AlphaMoERoutePlan` for `logits`.
`alphamoe_fused_router`(logits, *, top_k[, ...])	Route FP32 logits and build an AlphaMoE aligned plan (SM100/SM103).

Multi-LoRA MoE (BGMV)¶

Batched Gather-Matrix-Vector kernels for serving multiple LoRA adapters on top of a Mixture-of-Experts layer (shrink + expand).

prepare_bgmv_moe() returns a graph-replayable plan. On SM90, SM100 and SM103 devices the generated Cake programs serve one LoRA slice with rank 8, 16, 32 or 64 and any hidden size that is a positive multiple of 8 (hidden 2688 and 3072 at rank 32 use the specialized measured bodies at up to 2048 tokens, plan.variant == "specialized"; everything else the runtime-hidden generic bundles, plan.variant == "generic"); the output has one owner per token, so replays are bitwise reproducible. Arbitrary pair order (for example expert-sorted dispatch) is served through a token->pair route index that the shrink kernels publish and the expand kernels read in O(1) per CTA; the generic shrink splits the hidden dimension over extra CTAs at small pair counts and reduces the partials in a fixed order. Other inputs fall back to the portable bgmv_moe_shrink / bgmv_moe_expand kernels (plan.backend_used == "portable") unless fallback=False is passed.

`bgmv_moe`(x, lora_a_weights, lora_b_weights, ...)	High-level multi-LoRA MoE BGMV: shrink + expand in one call.
`prepare_bgmv_moe`(x, lora_a_weights, ...[, ...])	Prepare a graph-replayable BGMV MoE shrink+expand pipeline.
`BGMVMoECakePlan`(module, *, y_accum, ...[, ...])	Pointer-stable SM90/SM100/SM103 Cake BGMV MoE shrink+expand execution plan.
`BGMVMoEPortablePlan`(*, y_accum, shrink_out, ...)	Prepared plan that runs the portable `bgmv_moe_shrink` + `bgmv_moe_expand` path.
`bgmv_moe_shrink`(y, x, w_ptr, ...[, ...])	MoE LoRA shrink operation: project input through LoRA-A matrices.
`bgmv_moe_expand`(y, x, w_ptr, ...[, finalize])	MoE LoRA expand operation: project through LoRA-B matrices.
`bgmv_moe_gemm1_lora_delta`(hidden_states, ...)	FC1 (gate_up_proj) LoRA delta for a routed MoE, in the layout consumed by `trtllm_*_moe`'s `gemm1_lora_delta`.
`bgmv_moe_gemm2_lora_delta`(...[, lora_dtype, ...])	FC2 (down_proj) LoRA delta for a routed MoE, to be ADDED to the MoE output.

CUTLASS Fused MoE¶

cutlass_fused_moe(input, ...[, ...])

Compute a Mixture of Experts (MoE) layer using CUTLASS backend.

cuTile Fused MoE¶

`CuTileBf16Config`()	cuTile BF16 backend.
`CuTileBf16Runner`(config, device)	Unified adapter for the cuTile BF16 MoE pipeline.
`CuTileFp8PerTensorBf16Config`()	cuTile per-tensor E4M3 weights with BF16 inputs to both GEMMs.
`CuTileFp8PerTensorBf16Runner`(config, device)
`CuTileFp8PerTensorConfig`()	cuTile per-tensor E4M3 weights and activations.
`CuTileFp8PerTensorRunner`(config, device)
`CuTileMxfp4Bf16Config`()	cuTile MXFP4-weight x BF16-activation backend.
`CuTileMxfp4Bf16Runner`(config, device)	cuTile MXFP4 weights and BF16 activations for both grouped GEMMs.
`CuTileMxfp4Config`()	cuTile MXFP4-weight x MXFP4-activation backend.
`CuTileMxfp4Runner`(config, device)	cuTile MXFP4 weights and activations for both grouped GEMMs.
`CuTileMxfp4Mxfp8Config`()	cuTile MXFP4 weights with MXFP8 inputs to both GEMMs.
`CuTileMxfp4Mxfp8Runner`(config, device)
`CuTileMxfp8Bf16Config`()	cuTile MXFP8 weights with BF16 inputs to both GEMMs.
`CuTileMxfp8Bf16Runner`(config, device)
`CuTileMxfp8Config`()	cuTile MXFP8 weights and activations, with E8M0 scales per 32 values.
`CuTileMxfp8Runner`(config, device)
`CuTileNvfp4Bf16Config`()	cuTile NVFP4-weight x BF16-activation backend.
`CuTileNvfp4Bf16Runner`(config, device)	cuTile NVFP4 weights and BF16 activations for both grouped GEMMs.
`CuTileNvfp4Config`()	cuTile NVFP4-weight x NVFP4-activation backend.
`CuTileNvfp4Runner`(config, device)	cuTile NVFP4 weights and inputs for both grouped GEMMs.

TensorRT-LLM Fused MoE¶

`trtllm_bf16_moe`(routing_logits, ...[, ...])	BF16 MoE operation with autotuning support.
`trtllm_bf16_routed_moe`(topk_ids, ...[, ...])	Pre-routed BF16 MoE operation with autotuning support.
`trtllm_fp4_block_scale_moe`(routing_logits, ...)	FP4 block-scaled MoE operation.
`trtllm_fp4_block_scale_routed_moe`(topk_ids, ...)	FP4 block scale MoE operation with pre-computed routing.
`trtllm_fp8_block_scale_moe`(routing_logits, ...)	FP8 block-scaled MoE operation.
`trtllm_fp8_block_scale_routed_moe`(topk_ids, ...)	Pre-routed FP8 block-scaled MoE operation.
`trtllm_fp8_per_channel_scale_moe`(...[, ...])	FP8 per-channel scale MoE operation.
`trtllm_fp8_per_channel_scale_routed_moe`(...)	FP8 per-token activation/per-channel weight MoE with pre-computed routing.
`trtllm_fp8_per_tensor_scale_moe`(...[, ...])	FP8 per-tensor-scale MoE operation.
`trtllm_fp8_per_tensor_scale_routed_moe`(...)	Pre-routed FP8 per-tensor-scale MoE operation.
`trtllm_mxint4_block_scale_moe`(...[, ...])	MXINT4 block-scaled MoE operation.
`trtllm_mxint4_block_scale_routed_moe`(...[, ...])	MxInt4 block-scale MoE with pre-computed routing.

AlphaMoE FP8 Block-Scaled MoE (SM100/SM103)¶

`alphamoe_interleave_gated_weights`(...)	Interleave gate/up weight rows into the alphamoe_sm100 device layout.
`alphamoe_fp8_block_scale_aligned_moe`(...[, ...])	Fused W8A8 block-scale MoE over a pre-aligned routing plan (SM100/SM103).

AlphaMoE NVFP4 (SM100/SM103)¶

The AlphaMoE path consumes packed E2M1 activations and weights with linear per-16 E4M3 scales. The aligned entry consumes an existing routing plan; the routed entry aligns the supplied expert IDs before compute. Shape-selected stages perform gate/up projection, SwiGLU, NVFP4 requantization and down projection. The aligned path uses FP32 scratch; selected routed paths use FP32 or BF16 route storage before finalizing into caller-owned BF16 output. Neither entry resets the caller’s initial output. Three contiguous FP32 [E] tensors provide the per-expert static ModelOpt scales: the gate scale is applied before SiLU, the up scale before SwiGLU multiplication, and the down scale before route weighting.

`alphamoe_nvfp4_aligned_moe`(hidden_states, ...)	Run AlphaMoE NVFP4 gate/up, SwiGLU, requantization and down compute.
`alphamoe_nvfp4_routed_moe`(hidden_states, ...)	Run route alignment, complete expert computation and weighted accumulation.
`prepare_nvfp4_w1_data`(w1)	Prepare immutable packed W1 panels once during model loading.
`prepare_nvfp4_w1_gate_up_data`(w1)	Return byte-preserving adjacent gate/up panels for the routed entry.
`prepare_nvfp4_w1_gate_up_scales`(...)	Pair existing prepared W1 scale bytes with adjacent gate/up data.
`prepare_nvfp4_w1_scales`(w1_scale)	Return native CP-layout U8 panels as `[E(N/128)(K/256),16,128]`.
`prepare_nvfp4_w2_data`(w2)	Prepare immutable packed W2 panels once during model loading.
`prepare_nvfp4_w2_data_k256`(w2)	Prepare immutable 128-byte-row W2 panels for the large-M token-tile route.
`prepare_nvfp4_w2_scales`(w2_scale)	Return native CP-layout U8 panels `[E(K/128)B,8,128]`.
`prepare_nvfp4_w2_scales_k256`(w2_scale)	Return 16-scale CP-layout U8 panels of the W2 scales as `[E(K/128)(N/512),16,128]`.

AlphaMoE NVFP4 prepared weight scales¶

The aligned entry consumes an existing routing plan. The routed entry also aligns raw expert IDs and initializes the accumulation buffer. Both preserve caller-owned BF16 output and retain the raw E4M3 scale inputs. Shape-selected compute stages perform gate/up projection, SwiGLU, NVFP4 requantization, down projection and weighted accumulation.

prepare_nvfp4_w1_scales and prepare_nvfp4_w2_scales permute existing scale bytes into immutable uint8 tensor-map panels. They perform no scale arithmetic or requantization. Call them after the final device-local weight layout is established and before warmup or graph capture. Keep the raw scale tensors and the resulting panels alive for the layer’s current weight load. A subsequent weight load must replace the panels; they are derived buffers and need not be saved in the model checkpoint.

For raw gate/up scales [E,N,K/16], the prepared shape is [E*(N/128)*(K/256),16,128]. For raw down scales [E,K,N/32], it is [E*(K/128)*(N/256),8,128]. At local dimensions E=256, N=1024, K=6144 these require an additional 96 MiB and 48 MiB, respectively, per weight pair. Preparation and its memory cost are separate from repeated-request kernel timing.

Supply the buffers through optional keywords on either entry:

from flashinfer.fused_moe import (
    prepare_nvfp4_w1_scales,
    prepare_nvfp4_w2_scales,
)

# Once, after device-local weight loading:
w1_scale_prepared = prepare_nvfp4_w1_scales(gemm1_weights_scale)
w2_scale_prepared = prepare_nvfp4_w2_scales(gemm2_weights_scale)
prepared_scales = {
    "w1_scale_prepared": w1_scale_prepared,
    "w2_scale_prepared": w2_scale_prepared,
}

# Existing aligned or routed call: retain every raw argument and append
# **prepared_scales. Reuse these same tensors for subsequent requests.

Omitting optional prepared scale tensors retains the existing scale paths. The routed prepared-data paths below use W1 scales; the adjacent gate/up path has its own paired scale carrier. Supplied prepared tensors must satisfy the API’s dtype, device, shape and alignment checks. No weight preparation occurs inside a routed request.

AlphaMoE NVFP4 prepared gate/up data¶

prepare_nvfp4_w1_data permutes contiguous uint8 packed W1 weights [E,N,K/2] into [E*(N/128)*(K/256),128,128] panels. It preserves every packed byte without dequantization or requantization and requires N % 128 == 0 and K % 256 == 0. The buffer occupies the same number of bytes as raw W1: an additional 768 MiB at E=256, N=1024, K=6144.

With compatible w1_scale_prepared and w1_data_prepared, the routed entry selects prepared-data paths for M=8, M=128 and M=512 at exactly N=1024, K=6144, E=256, top_k=8, block_m=8. Other dimensions retain their existing selection. Omitting w1_data_prepared retains the scale-only or raw path. These data keywords belong to the routed entry, not the aligned entry. The paths preserve caller output and use BF16 expert-route storage. The eight-token path also keeps a separate FP32 initial-output seed; the 128-token and 512-token prepared-data paths finalize directly into caller output in route order.

Prepare and pass ordinary data panels to the routed entry:

from flashinfer.fused_moe import prepare_nvfp4_w1_data

# Once, after the final device-local weight load:
w1_data_prepared = prepare_nvfp4_w1_data(gemm1_weights)
prepared_routed = dict(
    prepared_scales, w1_data_prepared=w1_data_prepared
)
# Append **prepared_routed to the existing routed call, retaining raw inputs.

AlphaMoE NVFP4 adjacent gate/up panels¶

For the exact 512-token shape above, the routed entry also accepts w1_gate_up_data_prepared and w1_gate_up_scale_prepared. These two carriers place each gate panel next to its matching up panel. Supply both or neither: supplying only one raises ValueError. Both must be contiguous uint8 tensors on the raw W1 device, with 16-byte-aligned addresses and these shapes, where R = E*(N/256)*(K/256):

data: [R,256,128] from prepare_nvfp4_w1_gate_up_data(gemm1_weights);
scales: [R,32,128] from prepare_nvfp4_w1_gate_up_scales( w1_scale_prepared, gemm1_weights.shape).

Here N counts the combined gate/up rows in raw W1 [E,N,K/2]. The preparation requires N % 256 == 0 and K % 256 == 0. The scale helper consumes the output of prepare_nvfp4_w1_scales, not raw E4M3 scales. Both helpers only permute bytes. At the supported local dimensions, the adjacent data and scale buffers add 768 MiB and 96 MiB, respectively. Retaining ordinary prepared-data panels as well incurs their separate memory cost.

A valid adjacent pair takes precedence at exactly M=512, N=1024, K=6144, E=256, top_k=8, block_m=8 with supported route metadata and 4-byte-aligned activation, W1 and W2 scale addresses. It does not require the separate w1_scale_prepared, w1_data_prepared or w2_scale_prepared arguments at call time. Raw weights and raw scales remain required arguments. For other shapes the adjacent pair does not select this specialization. Omitting both retains the existing prepared-data, scale-only or raw selection. Invalid supplied carriers are rejected rather than silently replaced.

Prepare the pair once and reuse it across routed requests:

from flashinfer.fused_moe import (
    prepare_nvfp4_w1_scales,
    prepare_nvfp4_w1_gate_up_data,
    prepare_nvfp4_w1_gate_up_scales,
)

w1_scales = prepare_nvfp4_w1_scales(gemm1_weights_scale)
adjacent_routed = {
    "w1_gate_up_data_prepared": prepare_nvfp4_w1_gate_up_data(gemm1_weights),
    "w1_gate_up_scale_prepared": prepare_nvfp4_w1_gate_up_scales(
        w1_scales, gemm1_weights.shape
    ),
}
# Append **adjacent_routed to the existing routed call; retain all raw inputs.

Prepare after the final device-local load or shard layout is established and before warmup or CUDA graph capture. Keep the raw tensors and selected panels alive and immutable; rebuild the panels whenever those weights or scales are replaced. Shape checks cannot establish that a panel belongs to the current weight values. The helpers do not maintain an automatic cache or fetch model files. Reuse existing read-only model files as loading inputs, and keep any optional derived-panel cache separate from the original checkpoint. Reuse layer-owned device panels across requests instead of rebuilding them per call. Prepared panels are derived buffers and need not be saved in the model checkpoint. Preparation time and memory remain outside routed kernel timing.

The aligned entry keeps its existing None return value. The routed entry returns the caller’s output tensor. Neither entry resets caller output; provide a zeroed output when a fresh result is wanted. The chosen implementation may use separate compute launches and scratch storage. It does not promise a single-kernel or entirely on-chip intermediate implementation.

Cake NVFP4 Warp Decode (SM100/SM103)¶

The Cake warp-decode runner is an explicit unified-MoE backend for exact SM100 and exact SM103. Select it with CakeWarpDecodeConfig(backend="cake"); it is not in the default backend list. Each device loads only its exact generated target. The current generated portfolio fails closed outside these contracts:

(activation, hidden_size, intermediate_size, num_experts, top_k) is exactly (SwiGLU(), 2048, 512, 512, 10), (SwiGLU(), 2048, 1536, 60, 4), (SwiGLU(), 2560, 768, 384, 4), (SiLU(), 6144, 1536, 192, 4), (SwiGLU(), 2048, 768, 128, 8), (SwiGLU(), 4096, 1536, 128, 8), (SwiGLU(), 2048, 512, 256, 8), (SwiGLU(), 4096, 1024, 512, 10), (SwiGLU(), 3072, 1536, 256, 8), (SwiGLU(alpha=1.702, beta=1.0, limit=7.0), 6144, 3072, 128, 4), (SiTU(gate_scale=4.0, linear_scale=25.0), 3584, 3072, 896, 16), or one of the sharded per-partition slices (SwiGLU(), 4096, 512, 512, 10), (SwiGLU(), 4096, 256, 512, 10), (SwiGLU(), 3072, 768, 256, 8), and (SwiGLU(), 3072, 384, 256, 8);
the token count is 1–32, routing is UnpackedPrecomputed with contiguous int32 expert IDs and BF16 routing weights;
quantization is NVFP4, finalization and PDL are enabled, and expert parallelism, fused shared experts, bias, and LoRA are disabled.

The backend reuses the physical weight and activation layouts prepared by TrtllmFp4Config. Logical GEMM1 weights have intermediate_size * (2 if activation.is_gated else 1) rows: SwiGLU and SiTU use 2 * intermediate_size gate/up rows while standalone SiLU uses intermediate_size rows. The packed E2M1 weights use the production MajorK 32-row MMA shuffle: physical row p is restored at logical row (p & ~31) + ((p & 7) << 2) + ((p & 31) >> 3). Their E4M3 block scales use the production R128c4 layout. Default SwiGLU and SiTU weight dictionaries can be registered for both backend keys without copying. Parameterized SwiGLU consumes the prepared per-expert gemm1_alpha, gemm1_beta, and gemm1_clamp_limit FP32 tensors. SiTU consumes gemm1_alpha as its gate scale and gemm1_beta as its linear scale, and requires no clamp tensor. The prepared values default to the configured activation; per-expert tensor changes are consumed on every launch and CUDA Graph replay.

Parameterized SwiGLU uses different beta/clamp units in the two backends. Cake consumes the logical parameters. The official TRT-LLM runner consumes beta / d and clamp_limit / d, where d = view["output1_scale_gate_scalar"]. For non-unit gate scales, prepare a separate official dictionary before launching; alpha and the physical weight and scale tensors remain shared:

official_view = dict(view)
d = view["output1_scale_gate_scalar"]
official_view["gemm1_beta"] = (view["gemm1_beta"] / d).contiguous()
official_view["gemm1_clamp_limit"] = (view["gemm1_clamp_limit"] / d).contiguous()

Keep the logical dictionary registered for "cake" and register official_view for "trtllm_fp4_routed". The derived buffers are FP32 per-expert tensors. After changing logical beta, clamp limit, or gate scale, refresh their contents with copy_ before the next official launch or replay. Perform conversion outside timing and CUDA Graph capture, preserving buffer addresses already used by a captured graph.

Standalone SiLU has no supported TRT-LLM routed-MoE peer and its dictionary must be registered only for "cake". The default-SwiGLU example below shares its dictionary directly:

cake = CakeWarpDecodeConfig(backend="cake")
activation = SwiGLU()  # Or SiLU() for H6144/I1536/E192/top-k 4.
view = cake.prepare_weights(
    w1_bf16,
    w2_bf16,
    num_local_experts=num_experts,
    hidden_size=hidden_size,
    intermediate_size=intermediate_size,
    activation=activation,
)

weights = MoEWeightPack()
weights.prepare_for("cake", view)
if activation == SwiGLU():
    weights.prepare_for("trtllm_fp4_routed", view)

x_q, x_scale = cake.prepare_activations(x_bf16)
activations = MoEActivationPack(
    x_q,
    x_scale,
    topk_ids,
    topk_weights_bf16,
    routing_input_mode=RoutingInputMode.UnpackedPrecomputed,
)
config = MoEConfig(
    routing=RoutingConfig(num_experts=num_experts, top_k=top_k),
    quant=QuantConfig(weight=QuantFormat.NVFP4, activation=QuantFormat.NVFP4),
    experts=ExpertConfig(intermediate_size=intermediate_size),
    activation=activation,
    backend=BackendOptions((cake,)),
    execution=ExecutionConfig(enable_pdl=True),
)
output = MoELayer(config)(activations, weights)

The runner prepares its route-map workspace before a timed launch or CUDA Graph capture and reuses it for the same token count and geometry. Warm up each shape and routing tensor before capturing it; an unseen workspace shape or an unvalidated routing-tensor generation during capture is rejected instead of initializing implicitly. Repeated calls reuse the routing validation receipt until a normal tensor is modified in place. Inference tensors lack a version counter, so their receipt is identity/storage based. Every later inference-mode mutation and graph replay must keep expert IDs in the configured range because neither path can be revalidated automatically. At most 64 live routing tensors are retained; validating another distinct tensor fails explicitly, so construct a new runner for another bounded lifetime.

The runner retains a bounded prepared-workspace cache keyed by execution stream and geometry. Preparation issues a generation receipt, and the binding records completion events so explicit re-preparation or release cannot overtake submitted work. Completion-event handles are retained in a bounded process-lifetime pool so a live CUDA graph cannot reference a destroyed handle; a generation whose accepted work cannot be recorded is quarantined instead of being reused. A recycled allocator address cannot inherit stale metadata. Ordinary MoELayer calls receive per-stream workspaces automatically. Keep the runner and its workspaces alive for the lifetime of any captured graph, and do not concurrently replay multiple low-level graph executables that share one receipt. Workspace receipts are positive, generation-specific, and single-use; an unknown, stale, or repeated release is rejected rather than treated as a successful retirement. The runner-owned receipt lease strongly retains a workspace until retirement; if retirement cannot prove completion, the storage remains quarantined until process exit rather than returning to PyTorch’s allocator. The 4096-address event pool is likewise process-lifetime and requires a process restart after exhaustion. The matching module is also registered in exact SM100 and SM103 AOT builds when MoE kernels are enabled. SM100 does not load or fall back to the SM103 module, and other compute capabilities are rejected.

`CakeWarpDecodeConfig`([backend])	Explicit Cake NVFP4 warp-decode backend for exact SM100 and SM103.
`CakeWarpDecodeRunner`(config, device)	Exact-SM100/SM103 Cake runner for calibrated NVFP4 decode geometries.

Prims-TS Unified MoE (SM100/SM103)¶

The Prims-TS runner is an explicit unified-MoE backend for SM100 and SM103. Select it with PrimsTsConfig(); it is not in the default backend list. MVP coverage is NVFP4×NVFP4 and BF16×BF16. Routing and finalize stay on the TRT-LLM Gen path; only the GEMM middle stage uses Prims-TS. The backend reuses the physical layouts prepared by TrtllmFp4Config (MajorK) and TrtllmBf16Config (BlockMajorK), so one view can be registered for both keys:

backend = PrimsTsConfig()
quant = QuantConfig(weight=QuantFormat.NVFP4, activation=QuantFormat.NVFP4)
view = backend.prepare_weights(
    w1_bf16,
    w2_bf16,
    quant=quant,
    num_local_experts=num_experts,
    hidden_size=hidden_size,
    intermediate_size=intermediate_size,
    activation=SwiGLU(),
)

weights = MoEWeightPack()
weights.prepare_for("prims_ts", view)
weights.prepare_for("trtllm_fp4_routed", view)

x_q, x_scale = backend.prepare_activations(x_bf16, quant=quant)
activations = MoEActivationPack(x_q, x_scale, topk_ids, topk_weights)
config = MoEConfig(
    routing=RoutingConfig(num_experts=num_experts, top_k=top_k),
    quant=quant,
    experts=ExpertConfig(intermediate_size=intermediate_size),
    backend=BackendOptions((backend,)),
)
output = MoELayer(config)(activations, weights)

intermediate_size must be a multiple of 128. Fused shared experts and LoRA are out of scope for this MVP.

`PrimsTsConfig`()	Explicit Prims-TS backend for SM100 and SM103.
`PrimsTsRunner`(config, device)	Unified adapter over the inner `PrimsTs*MoERunner` classes.

Prims-TS Fused MoE¶

Experimental SM100, SM103, and SM107 Prims-TS backends. Public entry points match the corresponding trtllm_* APIs.

`prims_ts_bf16_moe`(routing_logits, ...[, ...])	BF16 MoE using the Prims-TS batched GEMM backend on SM100/SM103/SM107.
`prims_ts_bf16_routed_moe`(topk_ids, ...[, ...])	BF16 Prims-TS MoE with precomputed packed or unpacked routing.
`prims_ts_fp4_block_scale_moe`(routing_logits, ...)	FP4 block-scaled MoE using the Prims-TS backend on SM100/SM103/SM107.
`prims_ts_fp4_block_scale_routed_moe`(...[, ...])	FP4 block-scaled MoE with precomputed routing (Prims-TS / SM100).
`prims_ts_fp8_block_scale_moe`(routing_logits, ...)	FP8 block-scaled MoE using Prims-TS on SM100/SM103/SM107.
`prims_ts_fp8_block_scale_routed_moe`(...[, ...])	Pre-routed FP8 block-scaled MoE using Prims-TS on SM100/SM103/SM107.
`prims_ts_fp8_per_tensor_scale_moe`(...[, ...])	FP8 per-tensor-scaled MoE using Prims-TS on SM100/SM103/SM107.

The *_partitioned variants split one MoE layer across two green-context streams with caller-owned weight shards and are supported on SM107 (Rubin) only (see flashinfer.prims_ts.moe.MoePartitionResources and flashinfer.prims_ts.moe.shard_prepared_moe_weights):

`prims_ts_bf16_moe_partitioned`(...[, ...])	Run two-way BF16 Prims-TS MoE with caller-owned shards and streams.
`prims_ts_fp4_block_scale_moe_partitioned`(...)	Run two-way FP4 Prims-TS MoE with caller-owned shards and streams.
`prims_ts_fp8_block_scale_moe_partitioned`(...)	Run partitioned MXFP8 block-scale MoE with caller-owned payloads.
`prims_ts_fp8_per_tensor_scale_moe_partitioned`(...)	Run partitioned FP8 per-tensor MoE with caller-owned payloads.

These symbols are defined in the backend modules and re-exported from flashinfer.fused_moe:

flashinfer.fused_moe.backends.prims_ts.bf16_op.prims_ts_bf16_moe()
flashinfer.fused_moe.backends.prims_ts.bf16_op.prims_ts_bf16_moe_partitioned()
flashinfer.fused_moe.backends.prims_ts.bf16_op.prims_ts_bf16_routed_moe()
flashinfer.fused_moe.backends.prims_ts.fp4_op.prims_ts_fp4_block_scale_moe()
flashinfer.fused_moe.backends.prims_ts.fp4_op.prims_ts_fp4_block_scale_moe_partitioned()
flashinfer.fused_moe.backends.prims_ts.fp4_op.prims_ts_fp4_block_scale_routed_moe()
flashinfer.fused_moe.backends.prims_ts.fp8_op.prims_ts_fp8_block_scale_moe()
flashinfer.fused_moe.backends.prims_ts.fp8_op.prims_ts_fp8_block_scale_moe_partitioned()
flashinfer.fused_moe.backends.prims_ts.fp8_op.prims_ts_fp8_block_scale_routed_moe()
flashinfer.fused_moe.backends.prims_ts.fp8_op.prims_ts_fp8_per_tensor_scale_moe()
flashinfer.fused_moe.backends.prims_ts.fp8_op.prims_ts_fp8_per_tensor_scale_moe_partitioned()

Standalone TRT-LLM Gen Routing¶

The routing stage the TRT-LLM Gen fused MoE launchers run before their GEMMs, exposed on its own so expert selection and the permutation/padding bookkeeping can be used (and tested) independently of quantization and GEMM configuration.

`trtllm_gen_routing`(routing_logits, ...[, ...])	Standalone trtllm-gen MoE routing (expert selection + permutation).
`TrtllmGenRoutingResult`(topk_ids, ...)	Outputs of the trtllm-gen MoE routing stage.

CuteDSL Fused MoE¶

The CuteDSL backends are conditionally available when the nvidia-cutlass-dsl package is installed.

`cute_dsl_fused_moe_bf16`(x, ...[, ...])	SM90 CuTe-DSL fused MoE forward (BF16/FP16, unquantized).
`cute_dsl_fused_moe`(x, x_sf, ...[, ...])	Run a fused MoE forward pass using CuTe-DSL block-scaled kernels.
`cute_dsl_fused_moe_nvfp4`(x, x_sf, ...[, ...])	Run a fused MoE forward pass using the CuTe-DSL NVFP4 kernels.
`cute_dsl_fused_moe_mxfp8_mxfp4`(x, x_sf, ...)	Run fused MoE with MXFP8 activations and packed MXFP4 weights.
`b12x_fused_moe`(x, w1_weight, w1_weight_sf, ...)	Run fused MoE on SM120/SM121 using b12x CuTe-DSL kernels.

class flashinfer.fused_moe.CuteDslMoEWrapper(num_experts: int, top_k: int, hidden_size: int, intermediate_size: int, use_cuda_graph: bool = False, max_num_tokens: int | None = None, num_local_experts: int | None = None, local_expert_offset: int = 0, tile_size: int = 128, sf_vec_size: int = 16, output_dtype: dtype = torch.bfloat16, device: str = 'cuda', 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, use_fused_finalize: bool = True, quant_mode: str = 'w4a4')¶

Bases: object

Wrapper class for CuteDSL MoE with CUDA graph and auto-tuning support.

With use_cuda_graph=True, the wrapper creates persistent CUDA stream and event resources outside graph capture, enabling async-memset / GEMM1 overlap during capture and replay. Auto-tuning is supported via the tactic parameter or autotune() context.

Supported architectures: SM100, SM103, and SM107. W4A8 is limited to SM100 and SM103.

num_experts¶: Total number of experts.

top_k¶: Number of experts per token.

hidden_size¶: Hidden dimension size.

intermediate_size¶: Intermediate dimension size.

use_cuda_graph¶: Whether the wrapper holds persistent stream/event resources for CUDA graph capture.

use_fused_finalize¶: Use atomic fused finalize; otherwise use the deterministic two-stage finalize.

quant_mode¶: Selected W4A4, W4A8, or W4A16 compute mode.

max_num_tokens¶: Deprecated; accepted for backwards compatibility but ignored.

Example (CUDA Graph):

>>> moe = CuteDslMoEWrapper(
...     num_experts=256, top_k=8,
...     hidden_size=7168, intermediate_size=2048,
...     use_cuda_graph=True,
... )
>>> # Warmup
>>> for _ in range(3):
...     output = moe.run(x, x_sf, topk_ids, topk_weights, w1, w1_sf, ...)
>>> # Capture
>>> g = torch.cuda.CUDAGraph()
>>> with torch.cuda.graph(g):
...     output = moe.run(x, x_sf, topk_ids, topk_weights, w1, w1_sf, ...)
>>> # Replay
>>> g.replay()

Example (Auto-tuning):

>>> moe = CuteDslMoEWrapper(num_experts=256, top_k=8, ...)
>>> # Run with auto-tuning
>>> with autotune(True):
...     output = moe.run(x, x_sf, topk_ids, topk_weights, w1, w1_sf, ...)

__init__(num_experts: int, top_k: int, hidden_size: int, intermediate_size: int, use_cuda_graph: bool = False, max_num_tokens: int | None = None, num_local_experts: int | None = None, local_expert_offset: int = 0, tile_size: int = 128, sf_vec_size: int = 16, output_dtype: dtype = torch.bfloat16, device: str = 'cuda', 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, use_fused_finalize: bool = True, quant_mode: str = 'w4a4')¶

Configure the CuTe-DSL block-scaled fused-MoE wrapper.

Parameters:

num_experts (int) – Total number of experts.
top_k (int) – Number of experts routed to per token.
hidden_size (int) – Hidden dimension size.
intermediate_size (int) – Intermediate dimension size after the fused activation.
use_cuda_graph (bool) – Create persistent CUDA stream/events for W4A4 async-memset overlap. W4A16 is CUDA-graph safe without those resources. Defaults to False.
max_num_tokens (Optional[int]) – Deprecated; accepted for backwards compatibility but ignored.
num_local_experts (Optional[int]) – Local experts for expert parallelism. Defaults to num_experts.
local_expert_offset (int) – Offset of local experts in the global expert space. Defaults to 0.
tile_size (int) – Tile size for moe_sort. Defaults to 128.
sf_vec_size (int) – Scale-factor vector size. Defaults to 16.
output_dtype (torch.dtype) – Output dtype. Defaults to torch.bfloat16.
device (str) – Device on which to allocate buffers. Defaults to "cuda".
enable_pdl (bool) – Enable Programmatic Dependent Launch. Defaults to True.
activation_type (int) – FC1 activation type. Use ActivationType.Swiglu for gated SwiGLU/SiTU, ActivationType.GegluTanh for tanh-approximate GeGLU, and ActivationType.Relu2 for non-gated ReLU^2. Setting situ_beta selects SiTU.
swiglu_alpha (float) – SwiGLU parameters. swiglu_oai is represented as ActivationType.Swiglu with non-default values.
swiglu_beta (float) – SwiGLU parameters. swiglu_oai is represented as ActivationType.Swiglu with non-default values.
swiglu_limit (float) – SwiGLU parameters. swiglu_oai is represented as ActivationType.Swiglu with non-default values.
situ_beta (Optional[float]) – When set with ActivationType.Swiglu, use the SiTU gate beta * tanh(gate / beta) * sigmoid(gate).
situ_linear_beta (Optional[float]) – Optional SiTU tanh clamp for the up branch.
use_fused_finalize (bool) – Use atomic fused finalize; otherwise use the deterministic two-stage finalize. Defaults to True.
quant_mode (str) – Compute mode: "w4a4", "w4a8", or "w4a16". Defaults to "w4a4". "nvfp4" is a deprecated alias for "w4a4".

get_valid_tactics() → list¶: Return list of valid tactics for this MoE configuration.

run(x: Tensor, x_sf: Tensor | None, token_selected_experts: Tensor, token_final_scales: Tensor, w1_weight: Tensor, w1_weight_sf: Tensor, w1_alpha: Tensor, fc2_input_scale: Tensor | None, w2_weight: Tensor, w2_weight_sf: Tensor, w2_alpha: Tensor, tactic: Tuple | None = None, *, per_token_scale: Tensor | None = None) → Tensor¶

Run the CuTe-DSL fused-MoE forward pass.

CUDA-graph safe when the wrapper was constructed with use_cuda_graph=True. Supports auto-tuning via the tactic argument or the surrounding autotune() context manager.

Parameters:

x (torch.Tensor) – Packed NVFP4 input for quant_mode="w4a4", MXFP8 input for quant_mode="w4a8", or BF16 input for quant_mode="w4a16".
x_sf (Optional[torch.Tensor]) – Scale factors for quant_mode="w4a4" or quant_mode="w4a8"; must be None for quant_mode="w4a16".
token_selected_experts (torch.Tensor) – Expert assignments of shape [num_tokens, top_k].
token_final_scales (torch.Tensor) – Routing weights of shape [num_tokens, top_k].
w1_weight (torch.Tensor) – GEMM1 weights (gate + up fused for gated activations, or a single projection for non-gated activations).
w1_weight_sf (torch.Tensor) – Scale factors for w1_weight.
w1_alpha (torch.Tensor) – Per-expert global scale for GEMM1.
fc2_input_scale (Optional[torch.Tensor]) – Global scale for W4A4 GEMM2 input quantization; must be None for W4A8 and W4A16.
w2_weight (torch.Tensor) – GEMM2 weights (down projection).
w2_weight_sf (torch.Tensor) – Scale factors for w2_weight.
w2_alpha (torch.Tensor) – Per-expert global scale for GEMM2.
tactic (Optional[Tuple]) – Tactic tuple, or None for auto-selection via the runtime tuner.
per_token_scale (Optional[torch.Tensor]) – Optional W4A4 per-token input row scale for GEMM1.

Returns:

Output tensor of shape [num_tokens, hidden_size].

Return type:

torch.Tensor

class flashinfer.fused_moe.CuteDslBf16MoEWrapper(num_experts: int, top_k: int, hidden_size: int, intermediate_size: int, num_local_experts: int | None = None, local_expert_offset: int = 0, output_dtype: dtype = torch.bfloat16, enable_pdl: bool = True, use_fused_finalize: 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)¶

Bases: object

SM90 CuTe-DSL fused-MoE wrapper (bf16/fp16, unquantized).

Holds the static MoE configuration so call sites only pass tensors. run is CUDA-graph capturable (the underlying pipeline is graph-safe) and delegates to cute_dsl_fused_moe_bf16(), which carries the API logging/trace decoration and documents the supported configuration. Auto-tuning is controlled by the autotune() context. Warm up the selected tactic once before CUDA-graph capture or serving so its process-local kernels are compiled.

Example (auto-tuning):

>>> moe = CuteDslBf16MoEWrapper(
...     num_experts=128, top_k=8, hidden_size=2048,
...     intermediate_size=768,
... )
>>> with autotune(True):
...     out = moe.run(x, topk_ids, topk_weights, w1_weight, w2_weight)

__init__(num_experts: int, top_k: int, hidden_size: int, intermediate_size: int, num_local_experts: int | None = None, local_expert_offset: int = 0, output_dtype: dtype = torch.bfloat16, enable_pdl: bool = True, use_fused_finalize: 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)¶

Configure the SM90 fused-MoE wrapper.

Parameters:

num_experts – Total (global) expert count.
top_k – Experts per token.
hidden_size – Model hidden dimension.
intermediate_size – Per-rank expert intermediate dimension (w1_weight is [E_local, 2*intermediate, hidden] interleaved for gated activations or [E_local, intermediate, hidden] for Relu2; w2_weight is [E_local, hidden, intermediate]).
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.
output_dtype – Output (= activation/weight) dtype, bf16 or fp16; used to allocate moe_output when the caller does not provide one.
enable_pdl – Launch the kernels with Programmatic Dependent Launch (default True; see cute_dsl_fused_moe_bf16()).
use_fused_finalize – True (default) fuses the router-scaled scatter-reduce into GEMM2; False selects the bitwise-reproducible two-stage finalize.
activation_type – GEMM1 activation (ActivationType.Swiglu, GegluTanh or Relu2); see cute_dsl_fused_moe_bf16().
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 with this scale.
situ_linear_beta – Optional SiTU tanh clamp of the up branch (requires situ_beta).

run(x: Tensor, token_selected_experts: Tensor, token_final_scales: Tensor, w1_weight: Tensor, w2_weight: Tensor, tactic: Tuple[Any, ...] | None = None, moe_output: Tensor | None = None, intermediate_buffer: Tensor | None = None) → Tensor¶

Run the fused MoE forward; see cute_dsl_fused_moe_bf16().

tactic dispatches that tactic tuple as given; None selects through the AutoTuner (the cached winner for these shapes when one exists, otherwise the fixed default).

class flashinfer.fused_moe.CuteDslMxfp8Mxfp4MoEWrapper(num_experts: int, top_k: int, hidden_size: int, intermediate_size: int, max_num_tokens: int | None = None, num_local_experts: int | None = None, local_expert_offset: int = 0, use_cuda_graph: bool = False, device: str = 'cuda', enable_pdl: bool = True, activation_type: int = 3, swiglu_alpha: float = 1.0, swiglu_beta: float = 0.0, swiglu_limit: float = 3.4028234663852886e+38)¶

Bases: CuteDslMoEWrapper

Production wrapper for the MXFP8 x MXFP4 fused-MoE pipeline.

Warning

This API will be deprecated in the future, please use CuteDslMoEWrapper with quant_mode="w4a8" instead.

Because the stream and event resources are reused, one wrapper instance is not reentrant or safe for concurrent calls. The first run binds the instance to that call’s CUDA stream; create one wrapper per stream.

__init__(num_experts: int, top_k: int, hidden_size: int, intermediate_size: int, max_num_tokens: int | None = None, num_local_experts: int | None = None, local_expert_offset: int = 0, use_cuda_graph: bool = False, device: str = 'cuda', enable_pdl: bool = True, activation_type: int = 3, swiglu_alpha: float = 1.0, swiglu_beta: float = 0.0, swiglu_limit: float = 3.4028234663852886e+38) → None¶: Initialize a reusable mixed-precision fused-MoE runner.

Warning

This API will be deprecated in the future, please use CuteDslMoEWrapper with quant_mode="w4a8" instead.

max_num_tokens is accepted for backwards compatibility but ignored. See CuteDslMoEWrapper for the full parameter documentation.

get_valid_tactics() → list¶: Return list of valid tactics for this MoE configuration.

run(x: Tensor, x_sf: Tensor, token_selected_experts: Tensor, token_final_scales: Tensor, w1_weight: Tensor, w1_weight_sf: Tensor, w1_alpha: Tensor, w2_weight: Tensor, w2_weight_sf: Tensor, w2_alpha: Tensor, tactic: Tuple[Any, ...] | None = None) → Tensor¶: Run the MXFP8 x MXFP4 fused-MoE forward pass.

Warning

This API will be deprecated in the future, please use CuteDslMoEWrapper.run() with quant_mode="w4a8" instead.

This entry point has no fc2_input_scale; it is forwarded as None. See CuteDslMoEWrapper.run() for the full parameter documentation.

class flashinfer.fused_moe.B12xMoEWrapper(num_experts: int, top_k: int, hidden_size: int, intermediate_size: int, *, use_cuda_graph: bool = False, max_num_tokens: int = 4096, num_local_experts: int | None = None, expert_map: Tensor | None = None, output_dtype: dtype = torch.bfloat16, device: str = 'cuda', activation: str = 'silu', swiglu_alpha: float = 1.702, swiglu_beta: float = 1.0, swiglu_limit: float | None = None, activation_precision: str = 'fp4', quant_mode: str | None = None, source_format: str = 'modelopt', shared_static_workspace: object | None = None, shared_dynamic_workspace: object | None = None, shared_output: Tensor | None = None)¶

Bases: object

B12x fused MoE wrapper for SM120/SM121 with CUDA graph support.

Pre-allocates workspace buffers for CUDA graph compatibility. Automatically selects micro/static/dynamic backend per call.

Parameters:

num_experts – Total number of experts.
top_k – Number of experts per token.
hidden_size – Hidden dimension size.
intermediate_size – Intermediate size.
use_cuda_graph – Pre-allocate buffers for CUDA graph compatibility.
max_num_tokens – Maximum tokens (only for use_cuda_graph=True).
num_local_experts – Experts whose weights this rank holds, for EP. Default: num_experts. Requires quant_mode=”w4a16” and expert_map when it differs from num_experts.
expert_map – Global-to-local expert map for EP, int32 [num_experts] with -1 for non-local experts. run() returns this rank’s zero-filled partial for the caller to sum across the EP group.
output_dtype – Output data type. Only torch.bfloat16 is currently supported. Default: torch.bfloat16.
device – Device for buffer allocation. Default: “cuda”.
activation – Activation — “silu”, “gelu_tanh”, “swigluoai_uninterleave”, or “relu2”. Default: “silu”. swiglu_alpha/beta/limit apply to swigluoai.
activation_precision – Backward-compatible alias for quant_mode. “fp4” selects quant_mode=”nvfp4”; “bf16” selects quant_mode=”w4a16”.
quant_mode – Quantization mode, “nvfp4”/”w4a4”, “mxfp4”, or “w4a16”. When set, this selects the backend and internal workspace family.
source_format – Source weight format for quant_mode=”w4a16”. Supports “modelopt” and “compressed_tensors”. Default: “modelopt”.
shared_static_workspace – Optional externally-allocated buffers reused instead of fresh allocations. Callers running many identically-shaped wrappers (e.g. one per MoE layer) can share a single set, since layers execute sequentially. Shapes must match this wrapper’s config. Only valid with use_cuda_graph=True.
shared_dynamic_workspace – Optional externally-allocated buffers reused instead of fresh allocations. Callers running many identically-shaped wrappers (e.g. one per MoE layer) can share a single set, since layers execute sequentially. Shapes must match this wrapper’s config. Only valid with use_cuda_graph=True.
shared_output – Optional externally-allocated buffers reused instead of fresh allocations. Callers running many identically-shaped wrappers (e.g. one per MoE layer) can share a single set, since layers execute sequentially. Shapes must match this wrapper’s config. Only valid with use_cuda_graph=True.

Example

>>> moe = B12xMoEWrapper(num_experts=256, top_k=8, ...)
>>> output = moe.run(x=hidden_states_bf16, ...)

__init__(num_experts: int, top_k: int, hidden_size: int, intermediate_size: int, *, use_cuda_graph: bool = False, max_num_tokens: int = 4096, num_local_experts: int | None = None, expert_map: Tensor | None = None, output_dtype: dtype = torch.bfloat16, device: str = 'cuda', activation: str = 'silu', swiglu_alpha: float = 1.702, swiglu_beta: float = 1.0, swiglu_limit: float | None = None, activation_precision: str = 'fp4', quant_mode: str | None = None, source_format: str = 'modelopt', shared_static_workspace: object | None = None, shared_dynamic_workspace: object | None = None, shared_output: Tensor | None = None)¶

Configure the b12x fused-MoE wrapper.

Parameters:

num_experts (int) – Total number of experts.
top_k (int) – Number of experts routed to per token.
hidden_size (int) – Hidden dimension size.
intermediate_size (int) – Intermediate dimension size.
use_cuda_graph (bool) – If True, pre-allocate workspace buffers sized for max_num_tokens so the wrapper can be captured into a CUDA graph. Defaults to False.
max_num_tokens (int) – Maximum batch size, only used when use_cuda_graph=True. Defaults to 4096.
num_local_experts (Optional[int]) – Number of experts whose weights this rank holds, for expert parallelism. Defaults to num_experts. Requires quant_mode="w4a16" and expert_map when it differs from num_experts.
expert_map (Optional[torch.Tensor]) – Global-to-local expert map for expert parallelism. See b12x_fused_moe(). Validated on the host once here, which may synchronize, and copied, so later caller mutations have no effect. Only supported with quant_mode="w4a16".
output_dtype (torch.dtype) – Output dtype. Only torch.bfloat16 is currently supported.
device (str) – Device on which to allocate workspace buffers. Defaults to "cuda".
activation (str) – Activation function — "silu" (gated SwiGLU), "gelu_tanh" (gated GeGLU, tanh-approx GELU), "swigluoai_uninterleave" (gated SwiGLU-OAI) or "relu2" (non-gated). Defaults to "silu".
swiglu_alpha (float) – SwiGLU-OAI parameters (only for "swigluoai_uninterleave"): gate*sigmoid(alpha*gate)*(up+beta) with optional clamp to swiglu_limit (None disables). Defaults 1.702 / 1.0 / None.
swiglu_beta (float) – SwiGLU-OAI parameters (only for "swigluoai_uninterleave"): gate*sigmoid(alpha*gate)*(up+beta) with optional clamp to swiglu_limit (None disables). Defaults 1.702 / 1.0 / None.
swiglu_limit (float) – SwiGLU-OAI parameters (only for "swigluoai_uninterleave"): gate*sigmoid(alpha*gate)*(up+beta) with optional clamp to swiglu_limit (None disables). Defaults 1.702 / 1.0 / None.
activation_precision (str) – Backward-compatible alias for quant_mode. "fp4" selects quant_mode="nvfp4"; "bf16" selects quant_mode="w4a16".
quant_mode (Optional[str]) – Quantization mode, "nvfp4" / "w4a4", "mxfp4", or "w4a16".
source_format (str) – Source weight format for quant_mode="w4a16" — "modelopt" (default) or "compressed_tensors".
shared_static_workspace (Optional[object]) – Externally allocated workspaces reused instead of fresh allocations. Callers running many identically-shaped wrappers (e.g. one per MoE layer) can share a single set, since layers execute sequentially. Shapes must match this wrapper’s config. Only valid with use_cuda_graph=True.
shared_dynamic_workspace (Optional[object]) – Externally allocated workspaces reused instead of fresh allocations. Callers running many identically-shaped wrappers (e.g. one per MoE layer) can share a single set, since layers execute sequentially. Shapes must match this wrapper’s config. Only valid with use_cuda_graph=True.
shared_output (Optional[torch.Tensor]) – Externally allocated output buffer, reused like the workspaces. Must be 2-D with shape (>= max_num_tokens, hidden_size), output_dtype, and the same device as this wrapper. Only valid with use_cuda_graph=True.

run(x: Tensor, w1_weight: Tensor, w1_weight_sf: Tensor, w2_weight: Tensor, w2_weight_sf: Tensor, token_selected_experts: Tensor, token_final_scales: Tensor, *, w1_alpha: Tensor, w2_alpha: Tensor, fc2_input_scale: Tensor | None = None, input_global_scale: Tensor | None = None) → Tensor¶

Run the b12x fused-MoE forward pass.

Parameters:

x (torch.Tensor) – Input activations of shape [num_tokens, hidden_size], bfloat16.
w1_weight (torch.Tensor) – FC1 weights, FP4-packed.
w1_weight_sf (torch.Tensor) – Scale factors for w1_weight.
w2_weight (torch.Tensor) – FC2 weights, FP4-packed.
w2_weight_sf (torch.Tensor) – Scale factors for w2_weight.
token_selected_experts (torch.Tensor) – Expert assignments of shape [num_tokens, top_k]. A negative id marks an unrouted slot (for example CUDA-graph padding): it contributes nothing, and a token whose slots are all negative produces a zero output row.
token_final_scales (torch.Tensor) – Routing weights of shape [num_tokens, top_k].
w1_alpha (torch.Tensor) – Per-expert global scale for FC1.
w2_alpha (torch.Tensor) – Per-expert global scale for FC2.
fc2_input_scale (Optional[torch.Tensor]) – Global scale for FC2 input quantization. Required for quant_mode="nvfp4"; accepted but ignored for "w4a16".
input_global_scale (Optional[torch.Tensor]) – Global scale for FC1 input quantization, scalar or [num_experts]. Defaults to w1_alpha; see b12x_fused_moe(). Ignored for "w4a16".

Returns:

Output tensor of shape [num_tokens, hidden_size].

Return type:

torch.Tensor

MonoMoE (Single-Kernel Block-FP8, SM90a)¶

Single-kernel top-K Mixture-of-Experts implementation specialized for the Qwen3.5-35B block-FP8 shape on Hopper (SM90a). The full pipeline — routing, up-projection, SiLU, down-projection and reduction — runs inside one kernel launch. Use has_monomoe() to check availability before calling.

`has_monomoe`()	Return True if the monomoe CUDA extension can be built and loaded.
`get_scratchpad_size_bytes`()	Return the global scratchpad size (bytes) required by the kernel.
`alloc_scratchpad`(device)	Allocate a zero-initialized scratchpad on `device` for the kernel.
`interleave_for_tma_wgmma_up`(w_fp8)	Repack fp8 up-projection weights for the Pair_Layout WGMMA A-tile.
`mono_moe`(activations_in, router_logits, ...)	Single-kernel block-FP8 top-K MoE (fixed E256/N512/K2048 shape, SM90a).