flashinfer.fused_moe.cute_dsl_fused_moe¶

flashinfer.fused_moe.cute_dsl_fused_moe(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, num_experts: int, top_k: int, num_local_experts: int | None = None, local_expert_offset: int = 0, output_dtype: dtype = torch.bfloat16, use_fused_finalize: bool = True, moe_output: Tensor | None = None, aux_stream: Stream | 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, *, quant_mode: str = 'w4a4', per_token_scale: Tensor | None = None, tactic: Tuple | None = None, localized_weights: list | None = None, localized_streams: list | None = None, localized_sm_count: int | None = None, localized_memset_stream: Stream | None = None, localized_allow_nonlocalized: bool = False, autotune_use_cuda_graph: bool = False) → Tensor¶

Run a fused MoE forward pass using CuTe-DSL block-scaled kernels.

Supported architectures: SM100, SM103, and SM107. W4A8 is limited to SM100 and SM103. This is the simple functional API; for CUDA-graph support use CuteDslMoEWrapper instead.

Auto-tuning is controlled by the autotune() context manager:

with autotune(True):
    output = cute_dsl_fused_moe(...)

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.
num_experts (int) – Total number of experts.
top_k (int) – Number of experts routed to per token.
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.
output_dtype (torch.dtype) – Output dtype. Defaults to torch.bfloat16.
use_fused_finalize (bool) – Use atomic fused finalize; otherwise use the deterministic two-stage finalize. Defaults to True.
moe_output (Optional[torch.Tensor]) – Pre-allocated output buffer. Allocated internally if None.
aux_stream (Optional[torch.cuda.Stream]) – Optional auxiliary CUDA stream used to overlap setup work with the main computation.
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_oai is represented as ActivationType.Swiglu with non-default swiglu_alpha/beta/limit.
swiglu_alpha (float) – SwiGLU parameters.
swiglu_beta (float) – SwiGLU parameters.
swiglu_limit (float) – SwiGLU parameters.
quant_mode (str) – Compute mode: "w4a4", "w4a8", or "w4a16". Defaults to "w4a4". "nvfp4" is a deprecated alias for "w4a4".
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.
per_token_scale (Optional[torch.Tensor]) – Optional W4A4 per-token input row scale for GEMM1.
tactic (Optional[Tuple]) – Tactic tuple, or None for auto-selection via the runtime tuner.
localized_weights (Optional[list]) –
Per-locality-domain W4A4 weight-shard dictionaries. Supplying these enables localized execution; exactly two equal-width shards are required. Entry i contains w1_weight, w1_weight_sf, w2_weight, and w2_weight_sf for localized_streams[i]. Allocate all four tensors in the corresponding locality-domain memory pool, for example inside torch.cuda.use_mem_pool(domain_pool); slicing an allocation from another pool does not relocate its storage.

For gated W4A4 with E local experts, hidden size H, and intermediate size I, packed uint8 full weights have shapes [E, 2*I, H/2] (FC1) and [E, H, I/2] (FC2). Split dimension 1 into two contiguous halves: each domain receives [E, I, H/2] and [E, H/2, I/2]. Keep each interleaved 64-row up/gate pair together in FC1. Split the corresponding unswizzled scale rows identically, then swizzle each shard independently into the CuTe MMA scale layout. FC2 retains the full reduction dimension; its shards produce disjoint hidden-output columns.

Continue passing the original full-width w1_weight, w2_weight, and their *_sf tensors to the regular arguments; do not pass None or shard tensors there. With localized_allow_nonlocalized=False, neither the GEMMs nor the autotuner read them, so a placeholder that keeps the full-width metadata but releases its storage (for example w.untyped_storage().resize_(0)) is valid; only FLASHINFER_LOGLEVEL=5 statistics would touch its data. With True, retain their valid contents because autotuning may run and select the full-width path. w1_alpha, w2_alpha, and fc2_input_scale remain shared across domains.

Example with shards already allocated in their domain pools and scale factors already converted to the MMA layout:
```
shards = [
    {
        "w1_weight": w1_domain0,
        "w1_weight_sf": w1_sf_domain0,
        "w2_weight": w2_domain0,
        "w2_weight_sf": w2_sf_domain0,
    },
    {
        "w1_weight": w1_domain1,
        "w1_weight_sf": w1_sf_domain1,
        "w2_weight": w2_domain1,
        "w2_weight_sf": w2_sf_domain1,
    },
]
```
localized_streams (Optional[list]) – One long-lived green-context CUDA stream for each locality domain (for example GreenContext.Stream()). Work is forked from and joined to the caller’s current stream with torch.cuda.execute_on_streams when PyTorch provides it, otherwise with an equivalent in-repo helper.
localized_sm_count (Optional[int]) – Number of SMs available to each locality-domain stream, used to size the persistent kernel grid.
localized_memset_stream (Optional[torch.cuda.Stream]) – Optional CUDA stream used to overlap fused-finalize output zeroing with FC1.
localized_allow_nonlocalized (bool) – Include the full-width path in autotuning when localized weights are supplied. The caller must retain usable full-width weights.
autotune_use_cuda_graph (bool) – Profile tactics through CUDA Graph replay.

Returns:

Output tensor of shape [num_tokens, hidden_size].

Return type:

torch.Tensor