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

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.

Multi-LoRA MoE (BGMV)

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

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 the generated SM100 BGMV MoE pipeline for graph replay.

BGMVMoEBlackwellPlan(module, *, y_accum, ...)

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

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.

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.

Cake NVFP4 Warp Decode (SM103)

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

  • (hidden_size, intermediate_size, num_experts, top_k) is exactly (2048, 512, 512, 10) or (2048, 1536, 60, 4);

  • the token count is 1–32, routing is UnpackedPrecomputed with contiguous int32 expert IDs and BF16 routing weights, and the activation is the default SwiGLU();

  • 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. One prepared weight dictionary can therefore be registered for both backend keys without copying:

cake = CakeWarpDecodeConfig(backend="cake")
view = cake.prepare_weights(
    w1_bf16,
    w2_bf16,
    num_local_experts=num_experts,
    hidden_size=2048,
    intermediate_size=intermediate_size,
)

weights = MoEWeightPack()
weights.prepare_for("cake", view)
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(variant=QuantVariant.NVFP4),
    experts=ExpertConfig(intermediate_size=intermediate_size),
    activation=SwiGLU(),
    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 module is also registered in SM103 AOT builds when MoE kernels are enabled.

CakeWarpDecodeConfig([backend])

Explicit Cake NVFP4 warp-decode backend for exact SM103.

CakeWarpDecodeRunner(config, device)

Exact-SM103 Cake runner for two calibrated NVFP4 decode geometries.

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.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, 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 – Local experts for EP. Default: num_experts.

  • 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, 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 local experts for expert parallelism. Defaults to num_experts.

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

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