flashinfer.fused_moe.cake / cake_cute (MXFP4 SiTU routed MoE)¶

The Cake MXFP4 x MXFP8 SiTU routed MoE is the Kimi K3 W4A8 expert layer (MXFP4 weights, MXFP8 activations, the SiTU activation with per-expert beta / linear_beta, top-k finalize into a BF16 output) reproduced as one set of Cake programs and lowered to two independent backends. Both packages come from the same pinned Cake revision and carry a content-addressed manifest (schema cake.library_export.v5) whose producer_revision is that revision:

backend="cake" – flashinfer.fused_moe.cake: CUDA C++ translation units with tvm_ffi_utils bindings under csrc/fused_moe/cake_mxfp4_situ_moe/cuda/sm_103a/, JIT-built through FlashInfer’s TVM-FFI JIT. Keeps the pointer TMA ABI of the validated production build (a caller-owned descriptor workspace per launch, owned by the plan).
backend="cake_cute" – flashinfer.fused_moe.cake_cute: CuTe DSL device modules with cutedsl_tvm_ffi bindings under csrc/fused_moe/cake_mxfp4_situ_moe/cute/sm_103a/, compiled through the DSL’s cute.compile on first use (pip install nvidia-cutlass-dsl). By-value grid-constant TMA ABI.

Both target sm_103a (B300 / GB300). The kernels are the exact reproduction of the hand-written CuTe DSL swap-AB MoE chain: same warp roles, barrier and pipeline structure, tile forms, TMEM layout, epilogue instruction classes, programmatic-dependent-launch placement, persistent scheduler protocol and workspace layout; numerics are frozen (FP32 accumulation, SiTU FP32 sequence, UE8M0 round-up requantization, BF16 route-weight scaling before the reduce-add).

Host plan¶

Each package ships the same host plan, rendered per backend (cake_mxfp4_situ_moe_plan.py). Its decision table is the Cake plan’s own source, rendered verbatim from the pinned revision; it reproduces the hand-written CuteDslMxfp4MoEWrapper / Mxfp4MoESwapAbPlan decision for decision (swapab_max_tokens, tile policy, group rows, stage depths, weight grouping, fused-routing cap, PDL chain, workspace layout, launch sequence).

The plan is decided per device: CakeMxfp4MoEWrapper reads the CUDA device’s SM count once at construction (the current device; sm_count_source == "device"), exactly as the hand-written runPostTopKPipeline sizes its cooperative routing launch, and every SM-dependent choice follows from that one count – the cooperative moe_sort grid (SMs - 8 CTAs: 140 on the 148-SM B300, 144 on the 152-SM GB300), its bounded / generic per-thread state (T * top_k <= 4 * (SMs - 8) * 896) and the persistent GEMM grids (min(tiles, SMs)). plan() refuses a device whose SM count differs from the one the plan was decided for. An explicit sm_count= is the caller’s decision (a plan for another part; the runner then requires the cooperative grid to stay co-resident); without CUDA the host-only decision table uses the contract’s 148. The served rows’ selections are the same on 148 and 152 SMs; the contract evidence (gate oracle, bitwise / BF16-class agreement, timings) was measured on the 148-SM B300.

import torch
from flashinfer.fused_moe.cake import CakeMxfp4MoEWrapper, prepare_cake_mxfp4_weights
# or: from flashinfer.fused_moe.cake_cute import ...

# Once per rank (pure torch): canonical [up | gate] MXFP4 bank -> tile-major kernel operands.
weights = prepare_cake_mxfp4_weights(w1, w1_scale, w2, w2_scale)

runner = CakeMxfp4MoEWrapper(
    num_experts, top_k, hidden_size, intermediate_size,
    num_local_experts=num_local_experts, local_expert_offset=local_expert_offset,  # expert parallel
    # intermediate_shard=intermediate_size // tp_size,                            # or MoE tensor parallel
    backend="cake",                                                                # "cake_cute" in cake_cute
)
decision = runner.decide(num_tokens)          # host only; decision.supported / decision.reason
workspace = torch.empty(runner.get_workspace_size(num_tokens), dtype=torch.uint8, device="cuda")

# Outside CUDA-graph capture: binds the caller-owned buffers, loads the three kernels, runs one warmup.
plan = runner.plan(x, x_sf, topk_ids, topk_weights,
                   weights["w1"], weights["w1_sf"], weights["w2"], weights["w2_sf"],
                   beta=beta, linear_beta=linear_beta, workspace=workspace, output=out)
plan.run()                                    # graph-capturable; no allocation, no host sync

Inputs: x [T, H] float8_e4m3fn with x_sf [T, H/32] UE8M0 codes (uint8); topk_ids [T, top_k] int32 global expert ids; topk_weights [T, top_k] BF16 or FP32 (None selects the packed BF16-in-int32 route layout); beta / linear_beta FP32 [num_local_experts]; out [T, H] BF16. The weight bank is the canonical rank-local w1 [L, 2 * I_shard, H/2] (up rows then gate rows) with w1_scale [L, 2 * I_shard, H/32] and w2 [L, H, I_shard/2] with w2_scale [L, H, I_shard/32] (uint8).

plan refuses two classes of rows, both with NotImplementedError. decision.supported is False only when the hand-written selection needs an IR form this chain has not traced (the wide 192-row mixed192 form, MoE-TP rows of 8192 tokens and above); decision.reason names the missing form. Every other row reports decision.supported as True, but plan also refuses a row whose decision.path is not in EXECUTABLE_PATHS (the two_stage finalize on MoE-TP rows of 17 to 127 tokens is traced but not enqueued by this chain). Callers that need a hard answer before plan check decision.path in EXECUTABLE_PATHS as well as decision.supported. The dense grouped-GEMM path above swapab_max_tokens, the split two-stage path and the hybrid path run through their own executable chains (see executed_chains in the manifest).

Generated kernels¶

Kernels are selected through the trace-time constants recorded in each module’s route.form (kind, n_tile, kbps, m_group for the GEMM forms, the RoutingConfig fields for the routing forms); module file names are content fingerprints. gate names the validation each form went through in Cake: e2e_registry (compile, static analysis and GPU correctness on both backends in Cake’s kernel registry), plan_gate (selected by the plan on the contract rows and validated end to end through the plan), plan_variant (a routing-mode variant of a plan-gated form, reachable through the public API but not on a contract row).

Generated kernel forms (`route.stage`)¶
Stage	Kind	Backends	Gate
`gemm2_swapab_finalize_n8`	gemm2_swapab	cake, cake_cute	`e2e_registry`
`gemm2_swapab_finalize_n16`	gemm2_swapab	cake, cake_cute	`e2e_registry`
`gemm2_swapab_finalize_n32`	gemm2_swapab	cake, cake_cute	`e2e_registry`
`gemm2_swapab_finalize_n32_k12`	gemm2_swapab	cake, cake_cute	`e2e_registry`
`gemm2_swapab_finalize_n64`	gemm2_swapab	cake, cake_cute	`e2e_registry`
`gemm2_swapab_finalize_n128`	gemm2_swapab	cake, cake_cute	`e2e_registry`
`gemm2_swapab_finalize_n192_2cta`	gemm2_swapab	cake, cake_cute	`e2e_registry`
`gemm2_swapab_partial_n8_m2`	gemm2_swapab_partial	cake, cake_cute	`e2e_registry`
`gemm2_swapab_partial_n16_m2`	gemm2_swapab_partial	cake, cake_cute	`e2e_registry`
`gemm2_swapab_partial_n32_m2`	gemm2_swapab_partial	cake, cake_cute	`e2e_registry`
`gemm1_swapab_situ_n32`	gemm1_swapab	cake, cake_cute	`e2e_registry`
`gemm1_swapab_situ_n8_k8`	gemm1_swapab	cake, cake_cute	`e2e_registry`
`gemm1_swapab_situ_n16`	gemm1_swapab	cake, cake_cute	`e2e_registry`
`gemm1_swapab_situ_n64_rowgroup`	gemm1_swapab	cake, cake_cute	`e2e_registry`
`gemm1_swapab_situ_n192_2cta_rowgroup`	gemm1_swapab	cake, cake_cute	`e2e_registry`
`gemm2_dense_finalize_n192`	gemm2_dense	cake, cake_cute	`e2e_registry`
`gemm2_dense_finalize_n128`	gemm2_dense	cake, cake_cute	`e2e_registry`
`gemm2_dense_finalize_n256`	gemm2_dense	cake, cake_cute	`e2e_registry`
`gemm2_dense_finalize_n256_2cta`	gemm2_dense	cake, cake_cute	`e2e_registry`
`gemm2_dense_finalize_n256_c12`	gemm2_dense	cake, cake_cute	`e2e_registry`
`gemm2_dense_finalize_n192_c12`	gemm2_dense	cake, cake_cute	`e2e_registry`
`gemm2_dense_finalize_n256_rg`	gemm2_dense	cake, cake_cute	`e2e_registry`
`gemm2_dense_finalize_n256_rg_early`	gemm2_dense	cake, cake_cute	`e2e_registry`
`gemm1_dense_situ_m128_n256`	gemm1_dense	cake, cake_cute	`e2e_registry`
`gemm1_dense_situ_m128_n256_rowgroup`	gemm1_dense	cake, cake_cute	`e2e_registry`
`gemm1_dense_situ_m128_n128_rowgroup`	gemm1_dense	cake, cake_cute	`e2e_registry`
`gemm1_dense_situ_m128_n256_zero_fill`	gemm1_dense	cake, cake_cute	`e2e_registry`
`gemm1_dense_situ_m256_n256_2cta_zero_fill_secondary`	gemm1_dense	cake, cake_cute	`e2e_registry`
`gemm1_dense_situ_m128_n256_rowgroup_early`	gemm1_dense	cake, cake_cute	`e2e_registry`
`gemm1_dense_situ_m128_n128_rowgroup_early`	gemm1_dense	cake, cake_cute	`e2e_registry`
`routing_ep8_decode`	routing	cake, cake_cute	`e2e_registry`
`routing_cluster8`	routing	cake, cake_cute	`e2e_registry`
`routing_split`	routing	cake, cake_cute	`e2e_registry`
`token_index_top8`	token_index	cake, cake_cute	`e2e_registry`
`route_preprocess_packed`	route_preprocess	cake, cake_cute	`e2e_registry`
`dispatch_mixed_lists`	dispatch	cake, cake_cute	`e2e_registry`
`dispatch_mixed192_dual`	dispatch_mixed	cake, cake_cute	`e2e_registry`
`moe_sort_init_t384`	moe_sort_init	cake	`e2e_registry`
`moe_sort_coop_t384_dual_mixed`	moe_sort_coop	cake	`e2e_registry`
`moe_sort_init_t896`	moe_sort_init	cake	`e2e_registry`
`moe_sort_coop_t896_bounded`	moe_sort_coop	cake	`e2e_registry`
`moe_sort_coop_t896_bounded_dual`	moe_sort_coop	cake	`e2e_registry`
`moe_sort_coop_t896_dual`	moe_sort_coop	cake	`e2e_registry`
`moe_sort_coop_t896_bounded_dual_mixed`	moe_sort_coop	cake	`e2e_registry`
`moe_sort_coop_t896_dual_mixed`	moe_sort_coop	cake	`e2e_registry`
`finalize_top8`	finalize	cake, cake_cute	`e2e_registry`
`finalize_top8_split`	finalize	cake, cake_cute	`e2e_registry`
`gemm2_swapab_finalize_n8_m2`	gemm2_swapab	cake, cake_cute	`plan_gate`
`gemm2_swapab_finalize_n8_k8`	gemm2_swapab	cake, cake_cute	`plan_gate`
`gemm2_swapab_finalize_n8_k12`	gemm2_swapab	cake, cake_cute	`plan_gate`
`gemm2_swapab_finalize_n32_k8`	gemm2_swapab	cake, cake_cute	`plan_gate`
`gemm1_swapab_situ_n8_k8_nodp`	gemm1_swapab	cake, cake_cute	`plan_gate`
`gemm1_swapab_situ_n16_nodp`	gemm1_swapab	cake, cake_cute	`plan_gate`
`gemm1_swapab_situ_n32_nodp`	gemm1_swapab	cake, cake_cute	`plan_gate`
`gemm2_swapab_partial_n8_m2_nodp`	gemm2_swapab_partial	cake, cake_cute	`plan_gate`
`gemm2_swapab_partial_n16_m2_nodp`	gemm2_swapab_partial	cake, cake_cute	`plan_gate`
`gemm2_swapab_partial_n32_m2_nodp`	gemm2_swapab_partial	cake, cake_cute	`plan_gate`
`gemm1_swapab_situ_n16_nol2`	gemm1_swapab	cake, cake_cute	`plan_gate`
`gemm1_swapab_situ_n32_nol2`	gemm1_swapab	cake, cake_cute	`plan_gate`
`gemm2_swapab_finalize_n16_k8_nol2`	gemm2_swapab	cake, cake_cute	`plan_gate`
`gemm2_swapab_finalize_n32_k8_nol2`	gemm2_swapab	cake, cake_cute	`plan_gate`
`routing_separate_bf16_c1`	routing	cake, cake_cute	`plan_gate`
`routing_separate_bf16_c1_single`	routing	cake, cake_cute	`plan_gate`
`routing_separate_bf16_c2`	routing	cake, cake_cute	`plan_gate`
`routing_separate_bf16_c2_split`	routing	cake, cake_cute	`plan_gate`
`routing_separate_bf16_c3`	routing	cake, cake_cute	`plan_gate`
`routing_separate_bf16_c3_split`	routing	cake, cake_cute	`plan_gate`
`routing_separate_bf16_c4`	routing	cake, cake_cute	`plan_gate`
`routing_separate_bf16_c4_split`	routing	cake, cake_cute	`plan_gate`
`routing_separate_bf16_c5`	routing	cake, cake_cute	`plan_gate`
`routing_separate_bf16_c5_split`	routing	cake, cake_cute	`plan_gate`
`routing_separate_bf16_c6`	routing	cake, cake_cute	`plan_gate`
`routing_separate_bf16_c6_split`	routing	cake, cake_cute	`plan_gate`
`routing_separate_bf16_c7`	routing	cake, cake_cute	`plan_gate`
`routing_separate_bf16_c7_split`	routing	cake, cake_cute	`plan_gate`
`routing_separate_bf16_c8`	routing	cake, cake_cute	`plan_gate`
`routing_separate_bf16_c8_split`	routing	cake, cake_cute	`plan_gate`
`routing_separate_bf16_c8_l16384`	routing	cake, cake_cute	`plan_gate`
`routing_separate_bf16_c8_split_l16384`	routing	cake, cake_cute	`plan_gate`
`route_preprocess_separate_bf16`	route_preprocess	cake, cake_cute	`plan_gate`
`route_preprocess_separate_bf16_noclear`	route_preprocess	cake, cake_cute	`plan_gate`
`finalize_top16_split`	finalize	cake, cake_cute	`plan_gate`
`dispatch_wide_narrow_lists`	dispatch	cake, cake_cute	`plan_gate`
`gemm1_dense_situ_m128_n256_nopdl`	gemm1_dense	cake, cake_cute	`plan_gate`
`gemm1_dense_situ_m128_n256_zero_fill_nopdl`	gemm1_dense	cake, cake_cute	`plan_gate`
`gemm1_dense_situ_m256_n256_2cta_zero_fill_secondary_nopdl`	gemm1_dense	cake, cake_cute	`plan_gate`
`gemm2_dense_finalize_n192_c12_nopdl`	gemm2_dense	cake, cake_cute	`plan_gate`
`gemm2_dense_finalize_n192_nopdl`	gemm2_dense	cake, cake_cute	`plan_gate`
`gemm2_dense_finalize_n256_2cta_nopdl`	gemm2_dense	cake, cake_cute	`plan_gate`
`gemm2_dense_finalize_n256_c12_nopdl`	gemm2_dense	cake, cake_cute	`plan_gate`
`gemm2_dense_finalize_n256_nopdl`	gemm2_dense	cake, cake_cute	`plan_gate`
`moe_sort_coop_t896_bounded_dual_nopdl`	moe_sort_coop	cake	`plan_gate`
`moe_sort_coop_t896_bounded_nopdl`	moe_sort_coop	cake	`plan_gate`
`moe_sort_coop_t896_dual_nopdl`	moe_sort_coop	cake	`plan_gate`
`moe_sort_init_t896_nopdl`	moe_sort_init	cake	`plan_gate`

Forms this family has but a backend does not build are declared in that package’s manifest under contract.unsupported_forms with the reason (backend="cake_cute": the cooperative moe_sort pair, whose register state crosses grid.sync, which the CuTe DSL backend refuses; the CUDA package carries both kernels).

Manifest¶

contract of each manifest records the operator, the backend, the Cake revision, the architecture, the PDL launch state, the form table (with gate and backends per form), the served geometry, the contract rows the plan’s decision table serves (with the hand-written path of each, which of them each package carries with every launch form, and the EXECUTABLE_PATHS the shipped chain runs; plan() refuses decision-served rows outside them), the unsupported forms, and the SHA-256 of every integration file of the package (plan, kernel loader, __init__, tests, this page). The kernel loaders authenticate every generated source against the manifest before it is compiled.