KDA 0.5 FP8 Block-Scale MoE

CuTe-DSL implementation of the MLSys 2026 FlashInfer contest definition moe_fp8_block_scale_ds_routing_topk8_ng8_kg4_e32_h7168_i2048.

The kernel fuses DeepSeek-V3-style routing with two persistent grouped FP8 GEMMs. It targets NVIDIA Blackwell SM100/SM103 GPUs and is tuned for the official contest workloads.

1.5723x geometric-mean speedup over the official FlashInfer wrapper baseline across all 19 official MoE workloads.

Requirements

• NVIDIA GPU with compute capability 10.0 or 10.3
• CUDA 13 compatible driver
• PyTorch 2.12.1 with CUDA 13.0 (tested)
• nvidia-cutlass-dsl==4.6.0 (tested)

Usage

from kernels import get_kernel

moe = get_kernel(
    "humanfia-lab/moe-fp8-block-scale",
    version=1,
    trust_remote_code=True,
)

output = moe.run(
    routing_logits,
    routing_bias,
    hidden_states,
    hidden_states_scale,
    gemm1_weights,
    gemm1_weights_scale,
    gemm2_weights,
    gemm2_weights_scale,
    local_expert_offset,
    routed_scaling_factor,
)

The output is a BF16 tensor of shape [seq_len, 7168].

Fixed geometry

Performance

Across the 19 official MoE workloads, KDA 0.5 achieves a 1.5723x geometric- mean speedup over the official FlashInfer wrapper baseline. Every result in both columns passes the benchmark correctness check.

Measurements use an NVIDIA B200 with CUPTI cold-L2 kernel-span timing, 3 warmup iterations, 50 measured iterations, and 3 trials. KDA 0.5 is measured with CUDA 13.0, PyTorch 2.12.1+cu130, and CuTe DSL 4.6.0. The FlashInfer wrapper baseline uses the contest-era stack: flashinfer-python==0.6.8.post1, PyTorch 2.12.0+cu132, and CuTe DSL 4.5.0.

Source

The implementation is published from commit 9cf7b6dedefb65704fe016ffe25653c582ee8003 of humanfia/mlsys2026-flashinfer-contest-kda0.5, file kernels/moe/kernel.py.

Available functions

• run