Skip to content

Repository files navigation

MLSys26 CuTe-DSL kernels via KDA-v0.5

Three pure CuTe-DSL kernels for the MLSys 2026 FlashInfer AI Kernel Generation Contest tracks, targeting NVIDIA B200 (SM100). The kernels were produced by the KDA Team using kernel-design-agents; this repository compares them against the contest's human winners on the official workloads. Note, we currently only share the generated kernels in the repo. More details will come later. Stay tuned!

Each kernels/<track>/kernel.py exposes the run(...) entry point of its flashinfer-bench definition:

Track Definition
GDN prefill gdn_prefill_qk4_v8_d128_k_last
MoE moe_fp8_block_scale_ds_routing_topk8_ng8_kg4_e32_h7168_i2048
DSA attention dsa_sparse_attention_h16_ckv512_kpe64_topk2048_ps64
kda-05

Results

For every official workload, all kernels run on the same inputs; each entry is the per-track geometric mean of per-workload latency ratios. The first three columns are speedups over the official FlashInfer wrapper baseline; the last column compares KDA 0.5 directly against the human SOTA kernels. KDA 0.1 is the KDA team's original contest submission; KDA 0.5 is this repository's release. Measured on an NVIDIA B200 with CUPTI cold-L2 kernel-span timing (warmup 3, 50 iterations, 3 trials).

Track KDA 0.1 Human SOTA KDA 0.5 KDA-0.5 Speedup over Human
GDN prefill 1.42x 6.06x 10.36x 1.688x
MoE 0.67x 1.33x 1.57x 1.173x
DSA attention 3.51x 27.55x 38.33x 1.408x

Per-workload latencies: results/b200_vs_human_best.csv and results/b200_vs_flashinfer_baseline.csv.

KDA 0.5 and the human kernels are measured in this repository's pinned environment (CUDA 13.0, torch 2.12.1+cu130, CuTe-DSL 4.6.0); the FlashInfer baseline and KDA 0.1 in the contest-era environment (flashinfer-python 0.6.8.post1, torch 2.12.0+cu132, CuTe-DSL 4.5.0) with the same protocol, workloads, and seeds. The contest-era environment is the baseline's fastest (and, for GDN, its only runnable) stack: the contest-era GDN baseline does not run on CuTe-DSL 4.6 (cute.arch.ProxyKind was removed), and where the baseline runs on the current stack it is 1.2-1.4x slower, so these ratios are the conservative reference (against the current-stack baseline instead, MoE is 2.18x and DSA is 44.92x). The human winners' Triton kernels run unmodified in both environments with latencies agreeing to ~2%, anchoring the cross-environment ratios.

The human SOTA references are the winners' own public repositories (GDN prefill: Kachua; MoE: Team Wombat; DSA attention: Dogacel), pinned as git submodules under human_best/; the benchmark refuses to run if a submodule is missing, off its pinned commit, or locally modified.

A few caveats, stated up front:

  • The ratios are direct latency comparisons under the current flashinfer-bench protocol (CUPTI kernel spans); contest scores were computed by the organizers' wall-clock-era harness and are not directly comparable.
  • Performance tuning targets the official workload shapes; correctness is general, but the same speedups should not be expected on other shapes.
  • The pinned human snapshots are public repositories that postdate the contest deadline by a few commits; they are not organizer-signed artifacts.

Reproducing

Requires an NVIDIA B200 and a CUDA 13 driver.

# 1. Clone with the human-winner submodules.
git clone --recurse-submodules https://github.com/humanfia/mlsys2026-flashinfer-contest-kda0.5.git
cd mlsys2026-flashinfer-contest-kda0.5
# (already cloned? run: git submodule update --init --recursive)

# 2. Install the pinned environment. CUPTI timing is mandatory;
#    the runner aborts rather than fall back to another timer.
uv sync

# 3. Download the official workloads into benchmark_data/ (gitignored).
uv run --no-sync python tools/download_workloads.py

# 4. Run the paired comparison.
uv run --no-sync python benchmarks/compare_human_best.py \
  --track all --warmup 3 --iterations 50 --trials 3 \
  --output benchmark_results/b200_human_comparison.json

For a quick smoke test, pass --track dsa_attention --max-workloads 1 to both the download and benchmark commands (GDN's value tensors dominate the full download).

The runner builds both sides from source with FlashInfer-Bench, checks each against the definition's reference implementation under the MLSys contest's official acceptance criteria (the FlashInfer-Bench evaluation settings, per track), then times both with FlashInfer's CUPTI cold-L2 kernel-span timer. Random inputs use a deterministic per-workload seed (--seed); safetensor inputs come straight from the dataset. The JSON output contains the same per-workload rows as results/b200_vs_human_best.csv. (The baseline CSV additionally needs the contest-era environment described above, which this repository does not pin.)

Acknowledgements

Thanks to the FlashInfer team and the MLSys organizers for the benchmark framework and workload data, and to Kachua, Team Wombat, and Dogacel for publishing their winning kernels. The GDN prefill implementation builds on FlashInfer's open-source SM100 Gated Delta Net kernel.

License

MIT (see LICENSE). Three GDN prefill files derive from FlashInfer's SM100 GDN implementation and keep their upstream BSD-3-Clause/Apache-2.0 notices; see THIRD_PARTY_NOTICES.md. The kernels under human_best/ are submodules of their authors' repositories under their own licenses.

About

No description, website, or topics provided.

Resources

Stars

23 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages