Cake SiTU backend

flashinfer.fused_moe.cutlass_fused_moe(..., backend="cake") executes the complete BF16-input routed SiTU operation: group-16 NVFP4 input quantization, routing, FC1, SiTU, intermediate NVFP4 requantization, FC2, and BF16 weighted finalization. It supports SM100a/SM103a with TP8-local H=3584, I=384, E=896, top-k=16 and 1 through 16384 tokens. Expert IDs are int32 and routing weights are BF16. The caller supplies valid expert IDs in [0,896); routing values never participate in host dispatch. Output is local to the TP rank; collectives belong to the caller. The existing CUTLASS default and Cake warp-decode API are unchanged.

Prepared weights and scales

The Cake selection uses the TRTLLM shuffled group-16 NVFP4 layout, including the fused FC1 gate/up row permutation and block-scale interleave. Ordinary CUTLASS packed weight bytes are a different layout and must not be reused. Prepare the TRTLLM FP4 weights once with TrtllmFp4Config.prepare_weights using activation=SiTU() or supply tensors already in that exact layout. Storage is uint8 [896,768,1792] for FC1 and [896,3584,192] for FC2; block-scale storage is uint8 or float8_e4m3fn [896,768,224] and [896,3584,24] respectively. All tensors are contiguous.

The existing six-entry NVFP4 quant_scales list is passed directly:

  1. qX: one float32 activation encoding scale.

  2. FC1 block scales in the prepared layout.

  3. dX*dW13: float32 [896] FC1 dequantization scales.

  4. qA: float32 [896] intermediate encoding scales.

  5. FC2 block scales in the prepared layout.

  6. dA*dW2: float32 [896] FC2 dequantization scales.

qX=1/dX and qA=1/dA. Expand a shared intermediate scale to a contiguous 896-element tensor during model preparation; submission performs no scale arithmetic or broadcast allocation. Weight preparation’s scales must be carried with its packed bytes. Default SiTU is (4*tanh(g/4)*sigmoid(g))*(25*tanh(u/25)). Existing situ_beta and situ_linear_beta arguments can supply caller-owned float32 [896] gate and linear smoothing parameters directly. Omitting them uses the workspace’s prepared 4/25 constants.

Caller-owned execution

from flashinfer.fused_moe import (
    cutlass_fused_moe, cutlass_fused_moe_workspace_size,
    cake_fused_moe_prepare_workspace,
)
import torch
from flashinfer.tllm_enums import ActivationType

# x, expert_ids, route_weights, w13, w2 and the six scales are prepared.
nbytes = cutlass_fused_moe_workspace_size(
    max_tokens, 3584, 384, 896, 16,
    x_dtype=torch.bfloat16, weight_dtype=torch.uint8,
    activation_type=ActivationType.Situ, tp_size=8, backend="cake",
)
workspace = torch.empty(nbytes, dtype=torch.uint8, device=x.device)
output = torch.empty_like(x)
cake_fused_moe_prepare_workspace(workspace, x.shape[0], backend="cake")
cutlass_fused_moe(
    x, expert_ids, route_weights, w13, w2, torch.bfloat16, quant_scales,
    activation_type=ActivationType.Situ, tp_size=8, backend="cake",
    workspace_buffer=workspace, output=output,
)

The size query uses host metadata only. Explicit preparation initializes default SiTU constants and loads the route before capture; it allocates no tensor storage. For the pooled FC2 routes (8, 16, 32, 64, 128, 256, 512 and 1024 tokens), the workspace includes a four-byte work counter appended after the output scratch. Preparation resolves the FC2 pool (from the current GPU’s SM count on the 8- and 16-token routes, a fixed row count otherwise); the routing launch seeds this counter on every call. Prepare every token count before its first submission. A maximum-size buffer can hold every smaller prepared shape. Preparation and use must have normal stream ordering; each concurrently executing call needs separate output and workspace storage. Submission owns no graph and makes no CUDA allocation. The caller may capture or replay the complete call. The generated per-stage PDL and tensor-map ABI are preserved; enable_pdl=False is unsupported.