推理算子与内核——Attention/GEMM/采样对比、工程实践与真实项目

推理框架决定「谁进 batch、KV 放哪」；真正吃掉 Prefill 算力与 Decode 带宽的，往往是下一层——算子与内核（attention / GEMM / norm / sample / MoE）。工程师日常不一定手写 CUDA，但只要 TTFT/TPOT 抖一下，Nsight 里跳出来的名字就会变成排障主语：flash_fwd、flashinfer、cudnn_sdpa、gemmex、某个 TRT plugin……

本篇是「算子层」的独立深潜：多方面对比、工程实践、真实项目落点。CUDA 编程与 tiling 细节看 gpu-kernel 系列 与本系列 03；Prefill/Decode 心智看 11；分页与调度看 12；引擎产品看 13；框架解剖学看 29；DeepSeek 生产切片看 28。本文补的是：服务视角下的算子地图 + 对比表 + 怎么接到引擎上。

文中对比均为定性工程直觉，不是独立复现的 tok/s 排行；版本与 GPU 世代会改变相对强弱。调用链钉到具体 commit/路径，以便你在 checkout 里复核。

〇、30 秒版

1. 框架与算子要一起看：Scheduler/KV 决定请求形状与内存布局，kernel 决定对应路径的设备成本；只改一边常常像「优化了空气」。
2. 先分 Prefill / Decode：同一 Attention 实现，两端收益曲线不同；Decode 还要额外看 paged 布局与 CUDA Graph。
3. 对比尺子分层：算法/参考实现、Kernel 库、引擎集成不是同一货架；GQA 与 MLA 分列；能力用「支持/有条件/不适用 + 版本」读，不是吞吐圆点。
4. GEMM / 量化 / 采样同样是一等公民：小 M 先分清权重带宽 vs launch；Sampler 与 grammar bitmask 分环，勿统称 fused sampler。
5. 实践顺序：复现 → TTFT/TPOT → 引擎日志里的 backend 名 → profiler → 再改开关；正确性回归先于微基准。

一、为啥要单独讲算子

1.1 框架优化解决不了的那一段

特别篇 29 把 Scheduler、KV、Parallel、Sampling、Runner 钉成产品。但 Runner 内部仍要回答：

算子层的工作，是在正确的 metadata（positions、block tables、cu_seqlens、dtype）下，选中能打满 Tensor Core 或 HBM 的那条实现，并与 Graph / 融合策略兼容。

1.2 和 gpu-kernel 系列怎么分工

• gpu-kernel：教你怎么写、怎么 profile、怎么把 occupancy / bank conflict 讲清楚（本篇不替代 14 FlashAttention / 15 Triton / 16 融合 专篇）。
• 本篇：教你在 vLLM / SGLang / TRT-LLM / DeepInfer 里如何认路、对比、决策，以及真实项目里算子挂在哪。

两套都要时：先读本篇建立地图，再按需下潜到 FA / Triton / CUTLASS 专篇。

二、算子地图：一次 step 会碰到谁

图 2 按 pre-norm decode 画主链（与常见 Transformer 服务实现一致）：

1. embed 后进入 ×N 层：
   • Pre-attention RMSNorm → QKV 投影 → RoPE（作用在每层 Q/K，不是入口做一次）
   • paged attention + KV append → output 投影 + residual
   • RMSNorm → MLP / MoE（Act 在 MLP 内）→ residual
2. 层循环结束后：final norm → LM head / logits → bitmask → sample → stop

Prefill 更偏算力墙；常见 Decode 更偏带宽墙（小 batch、长 KV）。与 11、29 的 Roofline 叙事一致——算子选型必须带着阶段标签。Logits 约束与采样的职责拆分见第五节，勿全部叫成 fused sampler。

三、Attention 家族：多方面对比

3.1 对比维度（先定尺子）

先把对象分成三类，再读能力，避免把「算法论文实现 / 通用库 / 引擎插件」当成同一货架商品：

图 3 用 ● 支持 / ◐ 有条件 / ○ 不适用，并在「条件/边界」列写短注；不是吞吐强弱分，也不是可替代版本化 support 表的矩阵。上线前用你的 GPU arch + 模型结构核对各项目文档（例如 vLLM Attention Backend Feature Support、FlashAttention 接口里的 block_table / split、cuDNN Attention 的 page-table 输入）。

3.2 代表路径（定性 + 条件）

• FlashAttention-2 / 3（FA2/FA3）：IO-aware tiling + online softmax 的事实标准之一；FA3 更吃 Hopper 一代硬件特性。Prefill 通常强。官方接口已能表达部分 KV-cache / block-table / split；是否在你的引擎版本里默认启用，要以该引擎的 support 表为准——不能把「引擎未启用」画成「FA paged 能力弱」。
• FlashDecoding / FlashDecoding++：在 batch×heads 并行度不足且 KV 很长 时，用 KV 切分 + LSE 合并补 decode；有 combine 成本，不是「seq_q 小就一定切」。
• FlashInfer：服务向库——paged、GQA、采样周边更齐；SGLang 与部分 vLLM 路径依赖深。MLA 覆盖看版本矩阵。适合当「引擎默认旁边的第二公民」。
• cuDNN SDPA：PyTorch SDPA 的重要后端；文档侧有 paged / page-table 能力叙述，但驱动与 cuDNN 小版本会改变实际路由，必须钉版本。
• xFormers：历史记忆高效注意力入口，现更多作对照与遗留。
• TensorRT-LLM plugins：编译期强融合、峰值延迟友好；换模型/换架构工程税高。
• DeepInfer / DeepSeek 栈：MLA（FlashMLA）、MoE 等与结构绑得紧的自研执行路径；公开切片面向 DeepSeek-V4.1-Flash，硬件 NVIDIA SM100（CUDA 13）/ Huawei Ascend 950（CANN 9.2），形态 offline inference（许可 Apache-2.0；完整在线服务的 KV/PD/LB 编排未开源——那是适用边界，不是开源打折）。深潜见 28，不当通用 FA 替代品。

3.3 Tiling 直觉（不替换论文）

核心不是「又一个更快的 matmul」，而是：避免物化完整注意力矩阵，在 SRAM 里完成 softmax 归约。Decode 时若 batch×heads 等并行维度吃不饱、且 KV 很长，标准 FA 可能 SM 利用率不足，FlashDecoding / split-KV 一类二次切分才更有价值（还要计入 LSE combine 成本）；并非所有 decode 都自动需要。框架必须把 paged 布局 与 头映射 喂进 metadata；只替换 .so 却保留错误布局，结果是正确性事故，不是加速。

四、GEMM、量化与 MLP/MoE

4.1 三层入口与选用条件

1. cuBLASLt：通用 + epilogue（bias、激活、缩放）的默认好物。Prefill 大 M 时优先确认是否已走 Tensor Core + 合理 epilogue，而不是先手写 CUTLASS。
2. CUTLASS / CuTe：要定制 tiling、pipeline、混合精度，或库路径覆盖不到的 shape 时再上；维护成本归你的团队。
3. 专用库：DeepGEMM、各家 AWQ/GPTQ/FP8 推理 kernel——名字好听，一定要在 profiler 里看到对应 symbol，否则只是「配置写了量化」。

Decode 小 M GEMM：不要从「M 小」直接推出「主因是 launch」。至少拆成可观测分支：

grouped GEMM 适用于多组可合批问题（典型 MoE experts），不是任意单个小 M 的通用解；persistent 只覆盖其调度/launch 部分，不自动消除权重带宽墙。

4.2 和量化的边界

• 权重量化改的是 GEMM 输入格式与反量化融合；失败常见于回退到「先反量化再 FP16 GEMM」。
• KV 量化改的是 Attention 读带宽与数值范围；失败常见于长上下文质量塌或某些 backend 不接该 dtype。
• 选型条件至少记四条：GPU arch 指令支持、引擎 backend 是否认这一格式、是否融了 scale/zero、质量门禁是否覆盖长上下文——展开见 14。

4.3 MoE

专家并行下，瓶颈经常在 dispatch + grouped GEMM + combine，而不是「单个 expert 的 MLP/GEMM 微优化」。何时上 grouped：专家多、token 散、单专家 GEMM 太小导致 launch 占比高时。读 SGLang / vLLM / DeepInfer 时，把 EP 通信与 grouped 路径标进地图，避免只优化 Attention。

五、采样、结构化输出与「慢路径陷阱」

把链路拆开，避免「一个 fused sampler」包打天下：

选用判据（最小集，定性）：

Sampling 在 decode 步频下并不「永远便宜」。实践上把 structured output 当一等回归项；压测保留 tool call / JSON 比例，不要只用裸补全。失败回退到通用路径时，要能在日志或 profiler 里看到——否则 A/B 微基准会骗人。

六、引擎如何选 backend（真实交界）

决策是一棵树，不是一条全球统一的优先序：

GPU arch / driver
  └─ attention 类型（MHA/GQA/MLA/sliding window…）
       └─ prefill or decode（含 chunked）
            └─ paged? block size? KV dtype?
                 └─ 要不要 CUDA Graph / 捕获约束
                      └─ 引擎该版本的 support 表 → 具体 backend

vLLM 等引擎的官方材料是按 架构 × 注意力类型 × 特性 给支持矩阵，并可能单独开 MLA 路径；不同 GPU 上的优先次序不同，也没有「永远 FA3→FA2→FlashInfer→Triton→math」这种全球默认。图 6 应读成「输入条件 → registry」，不要读成固定排名。

排障顺序反而稳定：

1. 先读引擎日志 / metrics 里的 实际 backend 名；
2. 再用 Nsight 对上 kernel symbol；
3. 最后才改调度参数或换库。

错误示范：为追博客里的「最新 FA」，关掉 paged 或 Graph，微基准赢了，线上尾延迟与显存一起输。

七、真实项目切片（钉版本调用链）

带着四个问题读仓库（比背 API 有用）：

1. Attention 默认后端是谁？paged decode 是否同一条路径？MLA 是否独立 backend？
2. Sampler / grammar bitmask / stop 各在哪一层？失败如何回退？
3. MoE 是 grouped GEMM 还是多次 GEMM + 手写 dispatch？
4. 量化加载后，实际 kernel 名是什么？

7.1 DeepInfer（本地可核，commit 5537415）

对照本机只读检出（553741557a57df44a46fb3ab72b6aafeaf945337），与 28 一致。边界先写死：

• 模型：DeepSeek-V4.1-Flash（默认 HF id 见 README）
• 硬件：NVIDIA SM100 / Huawei Ascend 950
• 形态：offline（入口 examples/offline_inference.py）
• 许可：Apache-2.0

LLMEngine.step                         # deepinfer/engine/llm_engine.py:860
  └─ _run_workers("execute_model_step") # :966
       └─ Worker.execute_model_step     # deepinfer/worker/worker.py:89
            └─ model_runner.execute_main_model → 层循环
                 └─ DeepseekDecoderLayer.forward
                      # deepinfer/model_executor/models/decoder_layer.py:97
                      └─ V4Attention
                           ├─ prefill: prefill_attn_varlen_sparse
                           │   # deepinfer/model_executor/layers/attention_backend.py:26
                           │   （deepinfer/model_executor/models/v4_modules/attn.py:617；底层 flash_mla）
                           └─ decode: decode_mla_mqa_sparse
                               # deepinfer/model_executor/layers/attention_backend.py:77
                               （v4_modules/attn.py:660；FlashMLA + KV cache）
            └─ Sampler.forward
                 # deepinfer/model_executor/layers/sampler.py:114
                 ├─ structured mask：prefill / decode 分支不同（见 §8.3）
                 │   # deepinfer/structure_output/mask.py · XGrammarUtils
                 └─（可选）tilelang fused sample
                      # deepinfer/model_executor/layers/tilelang_ops/cuda/fused_sample.py:19
                      fused_sample_tl

要点：structured mask 与 fused sample 不在同一文件——排查时应分别看 structure_output/mask.py 与 Sampler.forward / fused_sample_tl，不要假设「一个 fused sampler 包打天下」。DeepInfer 替换的是整段 执行路径/runtime（引擎内核 + V4.1 特化），不是「只换一个 FA .so」。

7.2 vLLM（公开 main，对照 sha 9499c949）

以下路径来自 GitHub raw，对照 main @ 9499c949（vllm/v1/attention/selector.py 近期触及该路径的 commit；以你本地 checkout 为准，main 会继续前进）：

# 初始化期（模型加载时选型，不是每步重选）
vllm/v1/attention/selector.py::get_attn_backend
  （AttentionSelectorConfig：use_mla / sliding_window / kv_cache_dtype / …）
  → backends/registry.py::AttentionBackendEnum
       ├─ FLASH_ATTN  → backends/flash_attn.py
       ├─ FLASHINFER  → backends/flashinfer.py
       └─ MLA 族      → backends/mla/（flashattn_mla、flashinfer_mla、flashmla、…）

# 运行期每步
Scheduler → Worker 准备 inputs / block tables / positions
  → 已选定 Attention backend 的 forward（prefill 或 paged decode）
  → MLP / MoE GEMM
  → Sampler / logits processors（含 structured 约束）
  →（可选）CUDA Graph replay 同一 decode 形状

官方阅读锚点：Attention Backend Feature Support。配置侧常见入口是 attention backend / backend_per_kind 一类开关（名称随版本变，以你的文档为准）。

SGLang 侧把更多权重放在 FlashInfer 与前缀友好路径；TRT-LLM 把融合推进 plugin/编译图；DeepInfer 仅在上述模型/硬件/offline 边界匹配时考虑——不要写成「凡 DeepSeek 都上 DeepInfer」。

八、实践：从「慢」定位到算子

8.1 工作流

1. 复现：固定模型、长度、并发、精度；关掉无关实验开关。
2. 指标：分开 TTFT 与 TPOT（或 ITL）；不要只看平均 tok/s。
3. 引擎日志：确认 backend、是否 Graph、是否量化路径。
4. Profiler：PyTorch profiler / Nsight Systems / Nsight Compute——先看时间轴，再下钻单个 kernel。

8.2 判读口诀

• Prefill 慢且 FA/GEMM 占比高 → 算力路径、融合、并行度。
• Decode 慢且 attn/memcpy 占比高 → KV 布局、paged kernel、batch 过小。
• GPU 利用率低、CPU 高 → 调度/Python、Graph 未捕获、grammar/bitmask 偏 CPU、采样回退。
• 换 backend 几乎无收益 → 瓶颈可能在 KV 管理或通信，不在 Attention 微内核。

8.3 worked diagnosis（源码证据链；无 Nsight 数字）

证据类型：本地可读源码路径 + 符号（DeepInfer 5537415）。不是脱敏 Nsight 截图——本环境无 SM100/Ascend，无法把 DeepInfer 当线上 backend 跑出真实时间轴；若你的环境也不是该硬件，应用 vLLM selector 日志确认实际 backend，而不是硬套 DeepInfer 结论。

场景：structured output 开启后，decode 变「粘」，有人直觉怪「fused sampler 不够快」。

1. 固定：引擎 commit、模型 id、是否启用 SO、硬件是否在 DeepInfer 支持列表内。
2. 读链：LLMEngine.step → Worker.execute_model_step → Sampler.forward（sampler.py:114）。
3. 观察（源码级）：
   • prefill 约束：structure_output/mask.py::XGrammarUtils.apply_vocab_mask_by_step_inplace（:25），由 Sampler.forward 在 :188 调用；
   • decode 约束：Sampler.forward → _verification（:411）→ apply_token_bitmask_by_step_inplace（sampler.py:437），再 process_sampling；不要把 :188 当成 decode 证据；
   • fused sample 内核在 deepinfer/model_executor/layers/tilelang_ops/cuda/fused_sample.py::fused_sample_tl（:19），与 mask 分文件；
   • Attention 热点若在 MLA，符号是 prefill_attn_varlen_sparse / decode_mla_mqa_sparse（attention_backend.py，底层 flash_mla），与 sampler 无关。
4. 判读：应分别看「grammar/bitmask 是否回退或每步重编译」与「sample 是否仍走通用排序」；把二者糊成「一个 fused sampler」会指错文件。
5. 下一步验证：无 DeepInfer 硬件时，改在 vLLM 上读日志里的实际 attention backend 名（get_attn_backend / registry），再用 Nsight 对 symbol；有硬件时再对 examples/offline_inference.py 采时间轴。
6. 结论写法：「在 DeepInfer 5537415 上，SO mask（prefill :188 / decode _verification:437）与 fused sample 分属不同文件；因此先核对 mask 路径而非只升 sampler。」——别人能按路径复核，无需相信虚构 ms。

8.4 正确性

换算子后，最低验收应覆盖：

• greedy 与参考实现对齐（greedy 本身不消耗采样 RNG）；
• 固定种子的随机采样统计（分布 / RNG 语义）；
• eager vs CUDA Graph 一致性；
• 边界 shape：短/长序列、paged block 边界；
• GQA 与 MLA 各自（若启用）路径；
• 量化 KV / 权重量化若开启，单独质量门禁（含 logits/中间量容差口径，而不只看最终 token）；
• 结构化输出场景；
• 自定义 kernel 路径跑 compute-sanitizer（至少 memcheck）。

不要用「看起来差不多」代替比对。

8.5 CUDA Graph 与融合边界

适合进 Graph 的，多是 shape 稳定的 decode step；动态 draft 长度、CPU 侧 grammar/stop、频繁切换 backend，都会让捕获变脆。

融合请分层说，别写成「一律塞进 Attention epilogue」：

grammar 编译 / bitmask 与 sampler 分环：融了可以很快，但要留可调试回退。跨请求 paged 更新不要为了微基准静态化而牺牲正确性。

九、选型路径

1. 通用 CUDA 服务 → 跟随 vLLM 默认 FA/FlashInfer，先测再换。
2. Agent / 前缀密集 → SGLang + FlashInfer 路径更同向。
3. 极致延迟且养得起编译栈 → TRT-LLM plugins。
4. DeepSeek 特化且匹配 → 仅当模型为 DeepSeek-V4.1-Flash、硬件为 SM100 或 Ascend 950、且接受 offline 形态时，才考虑 DeepInfer；先读 28。匹配不上就留在社区 CUDA 引擎。

二级：只缺 decode → 在并行度不足且 KV 长时查 FlashDecoding/FlashInfer；只缺 prefill → Hopper+ 可看 FA3，其他 arch 看 FA2/cuDNN/引擎默认；MoE 热点 → grouped GEMM/EP。引擎产品横评仍以 13 为准。具体 backend 以第六节的版本化选择树为准，本图不是跨硬件默认排名。

十、开放问题

1. 通用 FA 路径 vs 引擎内特化 kernel：谁该拥有 paged/MLA 第一公民地位？
2. 编译栈强融合 vs 动态图可调试：尾延迟与工程税如何权衡？
3. FP4/FP8 与注意力数值：服务默认精度如何选，才不在长上下文翻车？
4. PD 分离后，KV 传递与本地 Attention kernel 的责任边界画在哪？
5. 国产加速器上「同名算子」是否同语义？迁移时正确性回归成本谁付？

十一、阅读路径与交叉引用

参考（公开资料；版本会变）

• FlashAttention 系列：Tri Dao et al.；接口参数见上游 flash_attn_interface（含 block table / split 等，随版本变）
• FlashDecoding / FlashDecoding++：公开技术报告与实现
• FlashInfer：https://github.com/flashinfer-ai/flashinfer · sampling 说明见项目站点技术博文
• cuDNN Attention / SDPA：https://docs.nvidia.com/deeplearning/cudnn/latest/operations/Attention.html
• vLLM Attention Backend Feature Support：https://docs.vllm.ai/en/latest/design/attention_backends/；源码锚点 main @ 9499c949（vllm/v1/attention/selector.py、backends/registry.py、backends/mla/）
• TensorRT-LLM：https://github.com/NVIDIA/TensorRT-LLM
• DeepInfer：https://github.com/deepseek-ai/DeepInfer @ 5537415（深潜见 28；Apache-2.0；V4.1-Flash / SM100·Ascend / offline）
• CUTLASS / CuTe：https://github.com/NVIDIA/cutlass
• XGrammar（结构化约束与 CPU/GPU 分工）：arXiv:2411.15100

本篇定位回顾：在框架篇与 CUDA/算子编程系列之间，补上「推理服务里的算子层」——对比尺子、引擎交界、钉版本调用链、实践路径和项目落点；深度放在怎么选、怎么证伪，而不是再做一张无法复现的吞吐榜。