有关写代码

python

torch

1. efficient_attention 没有 forward AD:

from torch.nn.attention import SDPBackend, sdpa_kernel
with sdpa_kernel(SDPBackend.MATH):
    # do sth

跳过 efficient_attention, 使用 C++ 纯 MATH 实现

C++

SIMD