KDA 0.5 GDN Prefill

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

The implementation combines a latency-oriented small-batch kernel, a persistent TCGen05 path, and a split-sequence path for chain-bound long sequences on NVIDIA B200.

10.3567x geometric-mean speedup over the official FlashInfer wrapper baseline across all 100 official GDN prefill workloads.

Requirements

• NVIDIA B200 or another GPU with compute capability 10.0
• 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

gdn = get_kernel(
    "humanfia-lab/gdn-prefill",
    version=1,
    trust_remote_code=True,
)

output, new_state = gdn.run(
    q,
    k,
    v,
    state,
    A_log,
    a,
    dt_bias,
    b,
    cu_seqlens,
    scale,
)

Fixed geometry

The output has shape [total_seq_len, 8, 128]; the returned state has shape [num_seqs, 8, 128, 128].

Performance

The complete per-workload FlashInfer comparison is available in flashinfer-baseline.csv.

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 and licensing

Published from commit 9cf7b6dedefb65704fe016ffe25653c582ee8003 of humanfia/mlsys2026-flashinfer-contest-kda0.5, directory kernels/gdn_prefill/.

Most files are MIT-licensed. gdn_core.py and gdn_tile_scheduler.py retain their upstream BSD-3-Clause notices, while gdn_runtime.py retains the FlashInfer Apache-2.0 notice. The complete notices are included in THIRD_PARTY_NOTICES.md in the kernel artifact.

Available functions

• run