Backend Dispatch 与 Fallback

四层名字必须分开

“调用了 FlashAttention”“matmul 走了 cuBLAS”这类结论，不能只从 Python 源码得出。一条简化执行链是：

API semantics
  torch.nn.functional.scaled_dot_product_attention / x @ W
        ↓
framework operator / lowering
  aten::scaled_dot_product_attention / aten::matmul → aten::mm
        ↓
backend or library policy
  math / Flash SDPA / cuDNN；BLAS、generated or custom GEMM path
        ↓
concrete GPU kernel for this shape and version
  profiler 中看到的 kernel name

• API 定义调用参数和预期数学语义。
• framework operator / graph lowering 是框架内部表示，可能被 decomposition、fusion 或 compilation 改写。
• backend/library 提供一组实现及选择策略；并不等同某一个 kernel。
• kernel 是本次 device、dtype、shape、layout 和配置实际启动的代码。

缩放点积注意力定义 SDPA 的 reference 结果；FlashAttention解释其中一种 IO-aware 算法。Backend dispatch 页回答的是“这一次调用实际走了哪条实现路径，以及不满足条件时发生什么”。

Kernel Backend、Triton Blocking 与 Persistent Kernel继续向下一层展开：dispatch 选定实现资源后，具体 kernel 如何把 tensor 映射为 tile、program、warp，以及为何一个可运行的 custom kernel 不自动胜过成熟 backend。

Dispatch key 不只是 dtype

SDPA 常见条件

对 $Q,K,V:[B,H,T,D]$，选择可能依赖：

• CPU/CUDA device 与 GPU compute capability；
• FP16、BF16、FP32、FP64 等 dtype 和 accumulator 路径；
• $D$、序列长度、batch/head 数及 alignment；
• contiguous/stride/layout；
• causal 或 explicit mask、mask shape；
• dropout、GQA、nested/ragged tensor 与软件版本支持矩阵；
• deterministic、backend enable/disable 与编译/捕获上下文。

同一个 scaled_dot_product_attention API 因此可能成为 fused Flash kernel，也可能分解成 matmul、mask、softmax、matmul。两者要保持同一数学契约，但低精度累加顺序不同，不要求 bitwise identical。

GEMM 常见条件

对 $X:[M,K]$、$W:[K,N]$，x @ W 的选择还会考虑 transpose、leading dimensions、dtype、TF32、Tensor Core alignment、epilogue、workspace、batching、确定性和可用 library/kernels。Profiler 中出现 aten::matmul / aten::mm 只证明 framework operator；CUDA event 才进一步给出本次 kernel 证据，而 kernel 名也不一定足以可靠反推完整 library 调用栈。

推理量化会增加更多 dispatch keys：bit width、packing、group size、scale layout、activation dtype 和 GPU capability。若 fast path 不支持其中任一条件，框架可能反量化后走浮点 GEMM、换用通用 kernel，或拒绝执行。

模型加载决定权重最终 dtype、stride、packing 与 device/rank placement；这些不是静态文件名的装饰，而是本页 dispatch 的实际输入条件。加载成功不代表首次前向不会 repack 或 fallback。

同一输入的三条 SDPA 路径

当前实验使用 CUDA FP16 Q=K=V:[1,8,256,64]：

默认 FP16 与强制 math 输出最大绝对差为 1.953e-03，在本例 FP16 容差内通过 assert_close。这说明 fallback 的目标是保存 API 语义，而不是保存逐 bit 的累加轨迹。

可运行的 profiler 证据

实验问题：同一个 SDPA API 能否在默认与强制 math 配置下产生不同 operator/kernel events？不支持的 dtype 在默认策略和 Flash-only 策略下分别怎样处理？x @ W 的 API、ATen operator 与 CUDA kernel 名是否处于不同层？

import platform,warnings
import torch
import torch.nn.functional as F
from torch.nn.attention import SDPBackend,sdpa_kernel
from torch.profiler import ProfilerActivity,profile

torch.manual_seed(157)
DEVICE=torch.device("cuda" if torch.cuda.is_available() else "cpu")

def profile_names(fn):
    activities=[ProfilerActivity.CPU]
    if DEVICE.type=="cuda": activities.append(ProfilerActivity.CUDA)
    with profile(activities=activities) as prof: output=fn()
    if DEVICE.type=="cuda": torch.cuda.synchronize()
    operators=sorted({e.key for e in prof.key_averages()
        if "scaled_dot_product" in e.key or e.key in {"aten::matmul","aten::mm"}})
    kernels=[] if DEVICE.type!="cuda" else sorted({e.name for e in prof.events()
        if str(e.device_type)=="DeviceType.CUDA"})
    return output,operators,kernels

def main():
    if DEVICE.type=="cuda":
        q=torch.randn(1,8,256,64,device=DEVICE,dtype=torch.float16)
    else:
        q=torch.randn(1,2,32,16,device=DEVICE,dtype=torch.float32)
    default,default_ops,default_kernels=profile_names(
        lambda:F.scaled_dot_product_attention(q,q,q,is_causal=True))
    with sdpa_kernel(SDPBackend.MATH):
        math,math_ops,math_kernels=profile_names(
            lambda:F.scaled_dot_product_attention(q,q,q,is_causal=True))
    tol=(2e-3,2e-3) if q.dtype==torch.float16 else (1e-5,1e-6)
    torch.testing.assert_close(default,math,rtol=tol[0],atol=tol[1])

    fallback_ops=[]; forced_flash="not tested on CPU"
    if DEVICE.type=="cuda":
        q64=torch.randn(1,2,32,16,device=DEVICE,dtype=torch.float64)
        _,fallback_ops,_=profile_names(
            lambda:F.scaled_dot_product_attention(q64,q64,q64,is_causal=True))
        assert any("_math" in name for name in fallback_ops)
        try:
            with warnings.catch_warnings():
                warnings.simplefilter("ignore",UserWarning)
                with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
                    F.scaled_dot_product_attention(q64,q64,q64,is_causal=True)
        except RuntimeError as error: forced_flash=type(error).__name__
        else: raise AssertionError("float64 unexpectedly ran with flash-only SDPA")
        x=torch.randn(128,256,device=DEVICE,dtype=torch.float16)
        weight=torch.randn(256,192,device=DEVICE,dtype=torch.float16)
    else:
        x=torch.randn(32,64); weight=torch.randn(64,48)
    _,gemm_ops,gemm_kernels=profile_names(lambda:x@weight)
    assert "aten::scaled_dot_product_attention" in default_ops
    assert any("_math" in name for name in math_ops)
    assert "aten::matmul" in gemm_ops and "aten::mm" in gemm_ops

    print(f"python={platform.python_version()} torch={torch.__version__}")
    print(f"device={DEVICE} SDPA_QKV={tuple(q.shape)} dtype={q.dtype}")
    print(f"default_sdpa_ops={default_ops}")
    print(f"forced_math_ops={math_ops}")
    print(f"default_vs_math_max_abs_diff={(default-math).abs().max():.3e}")
    print(f"float64_default_fallback_ops={fallback_ops}")
    print(f"float64_forced_flash_result={forced_flash}")
    print(f"matmul operator_events={gemm_ops}")
    print(f"default_sdpa_cuda_kernels={default_kernels[:2]}")
    print(f"forced_math_cuda_kernels={math_kernels[:3]}")
    print(f"matmul_cuda_kernel={gemm_kernels[:1]}")

if __name__=="__main__":
    with torch.inference_mode(): main()

完整文件位于 examples/backend_dispatch.py。2026-08-13 在 PyTorch 2.13.0+cu130、RTX 3080 上，默认 SDPA profiler event 为 aten::_scaled_dot_product_flash_attention，强制 math 为 aten::_scaled_dot_product_attention_math；默认 FP64 走 math fallback，Flash-only FP64 抛出 RuntimeError。x @ W 同时显示 aten::matmul、aten::mm 与一个具体 CUDA GEMM kernel 名。

Profiler 本身会引入开销，本实验不比较性能。kernel 名、backend 支持和 fallback 行为都可能随 PyTorch/CUDA、GPU 与输入变化；生产证据必须记录完整环境和代表性 shape，不能把本次 event 名写成永久架构事实。

如何诊断“为什么没有命中 fast path”

建议按从语义到实现的顺序收集证据：

1. 固定可复现输入并先用 reference/另一 backend 做输出 assert_close。
2. 记录 device、capability、dtype、shape、stride、mask/dropout、GQA 与相关 flags。
3. 使用 backend controls 分别允许/强制候选路径；区分自动 fallback 与强制失败。
4. 用 profiler 查看 framework events、CUDA kernels、layout conversions 和同步。
5. 再做 warmup、同步正确的多轮 benchmark；不能从 kernel 名直接推出端到端收益。

常见误诊包括：

• 只看高层 API 名称，宣称已经命中 FlashAttention 或量化 kernel；
• 只看一次 warning，不核对 profiler 和输出正确性；
• fast path 报错后把 mask/dtype 改掉，却没有确认模型语义是否也被改变；
• benchmark 默认路径与强制路径时使用不同 shape、不同精度容差或不同同步边界；
• torch.compile 后只找 eager ATen 名称，忽略 graph lowering 和 fusion 已改变边界；
• CUDA Graph capture 了 fallback 路径，此后 replay 很稳定，却稳定地重放了不期望的 kernels。

GPU Runtime进一步说明 torch.compile graph lowering/fusion 和 stream/event 依赖会改变 profiler 边界；dispatch 证据要在实际 eager/compiled/captured 模式下分别收集。

Dispatch 也不是请求 scheduler。Continuous batching 决定本轮有哪些序列和 shape；dispatch 再针对形成的算子输入选 backend。动态 batch 改变 shape 后可能跨过某个 kernel 支持/效率边界，因此二者共同影响 latency，但处于不同层。

参考资料

• PyTorch — scaled_dot_product_attention
• PyTorch — sdpa_kernel
• PyTorch Profiler
• NVIDIA cuBLAS documentation