KDA 0.5 DSA Sparse Attention

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

The kernel implements sparse MLA decode attention using paged compressed-KV and RoPE caches, split-K flash-decoding, and specialized small-token paths for NVIDIA Blackwell SM100/SM103 GPUs.

38.3270x geometric-mean speedup over the official FlashInfer wrapper baseline across all 23 official DSA attention 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

dsa = get_kernel(
    "humanfia-lab/dsa-sparse-attention",
    version=1,
    trust_remote_code=True,
)

output, lse = dsa.run(
    q_nope,
    q_pe,
    ckv_cache,
    kpe_cache,
    sparse_indices,
    sm_scale,
)

Fixed geometry

The output has shape [num_tokens, 16, 512]; LSE has shape [num_tokens, 16].

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

Published from commit 9cf7b6dedefb65704fe016ffe25653c582ee8003 of humanfia/mlsys2026-flashinfer-contest-kda0.5, file kernels/dsa_attention/kernel.py.

Available functions

• run