dsa-op

DSA算子分析

经过了前置的学习DSA以及FA的学习之后，是时候分析一下现有的dsa算子了

分析的来源目前均为推理引擎SGLang中现有的实现

由于我现在接触到的是H卡，所以下面的分析基本都以 Hopper 架构为基础

SGLang nsa backend

nsa prefill/decode 注意力算子的选项包括：

• flashmla_sparse：来自 flash_mla 库的 flash_mla_sparse_fwd 算子。可在 Hopper 和 Blackwell GPU 上运行。需要 bf16 格式的 q, kv 输入。
• flashmla_kv：来自 flash_mla 库的 flash_mla_with_kvcache 算子。可在 Hopper 和 Blackwell GPU 上运行。需要 bf16 格式的 q 和 fp8 格式的 k_cache 输入。
• fa3：来自 flash_attn 库的 flash_attn_with_kvcache 算子。仅能在 Hopper GPU 上运行。需要 bf16 格式的 q, kv 输入。
• tilelang：可在 GPU、HPU 和 NPU 上运行的 tilelang 实现。
• aiter：AMD HPU 上的 Aiter 算子。仅能用作 decode 算子。

基于benchmark，H200 和 B200 的默认配置设置如下：

• H200：flashmla_sparse prefill 注意力（短序列 prefill 通过 FlashAttention varlen 使用 MHA），fa3 decode 注意力
• B200：flashmla_auto prefill 注意力（短序列 prefill 通过 TRT-LLM ragged 使用 MHA），flashmla_kv decode 注意力

flashmla-sparse

Quick QA

Q：为什么不用fa3：

A：因为fa3的实现里虽然支持kv page且paged寻址天然就支持随机访问，但是它不能处理不同Q使用了不同的kv地址的情况，除非退化为每个SM只执行一个Q，效率极低

Q：为什么不用MHA mode

A：

• Dense prefill MHA vs MQA：MHA计算量小，prefill compute-bound → 选MHA
• Sparse prefill MHA vs MQA：
  • top-k大幅减少了KV数量 → compute-bound假设减弱
  • MQA不需要展开K/V → 避免了L_q × top-k × n_heads的巨型中间tensor→ MQA的IO优势反而更突出

op

fa3

Reference

• sglang doc