Analyzing DSA Operators

After the prerequisite study of DSA and FlashAttention, it is time to analyze existing DSA operators.

For now, the analysis is based entirely on implementations in the SGLang inference engine.

Because the GPUs available to me are H-series cards, the discussion is primarily based on the Hopper architecture.

SGLang nsa backend

The NSA prefill/decode attention operator options include:

• flashmla_sparse: the flash_mla_sparse_fwd operator from flash_mla. It runs on Hopper and Blackwell GPUs and requires bf16 Q and KV inputs.
• flashmla_kv: the flash_mla_with_kvcache operator from flash_mla. It runs on Hopper and Blackwell GPUs and requires bf16 Q and an fp8 K cache.
• fa3: the flash_attn_with_kvcache operator from flash_attn. It runs only on Hopper GPUs and requires bf16 Q and KV inputs.
• tilelang: a tilelang implementation that runs on GPUs, HPUs, and NPUs.
• aiter: the Aiter operator for AMD hardware. It can be used only for decoding.

Based on benchmarks, the default configurations for H200 and B200 are:

• H200: flashmla_sparse for prefill attention—short-sequence prefill uses MHA through FlashAttention varlen—and fa3 for decode attention.
• B200: flashmla_auto for prefill attention—short-sequence prefill uses MHA through TRT-LLM ragged—and flashmla_kv for decode attention.

flashmla-sparse

Quick QA

Q: Why not use FA3?

A: Although fa3 supports KV pages and paged addressing naturally permits random access, it cannot handle different queries using different KV addresses unless execution degenerates to one query per SM, which is extremely inefficient.

Q: Why not use MHA mode?

A：

• Dense prefill, MHA vs. MQA: MHA requires less computation and prefill is compute-bound, so choose MHA.
• Sparse prefill, MHA vs. MQA:
  • Top-k greatly reduces the number of KV entries, weakening the compute-bound assumption.
  • MQA does not need to expand K/V, avoiding a huge $L_q \times top\text{-}k \times n_{heads}$ intermediate tensor; MQA’s I/O advantage therefore becomes more significant.

op

fa3

Reference

• sglang doc