The v3 stable-ABI build runs the forward pass fine, but any backward with num_kv_heads != num_heads raises. MHA (num_kv_heads == num_heads) is unaffected, and the v2 cxx11 build is fine on the same torch, so this looks specific to the stable-ABI build rather than to a torch version.
import sys
import torch
from kernels import get_kernel
def check(version):
flash_attn = get_kernel("kernels-community/flash-attn2", version=version)
print(version, "->", sys.modules[flash_attn.__name__].__file__.split("/build/")[-1].split("/")[0])
for num_heads, num_kv_heads in ((8, 8), (8, 2)):
q = torch.randn(1, 8, num_heads, 64, dtype=torch.bfloat16, device="cuda", requires_grad=True)
k = torch.randn(1, 8, num_kv_heads, 64, dtype=torch.bfloat16, device="cuda", requires_grad=True)
v = torch.randn(1, 8, num_kv_heads, 64, dtype=torch.bfloat16, device="cuda", requires_grad=True)
flash_attn.flash_attn_func(q, k, v, causal=True).sum().backward()
print(f" q={num_heads} kv={num_kv_heads}: backward ok")
check(2) # torch212-cxx11-cu130-x86_64-linux: both fine
check(3) # torch-stable-abi210-cu130-x86_64-linux: MHA fine, GQA raises .../build/torch-stable-abi210-cu130-x86_64-linux/flash_attn_interface.py:291, in _flash_attn_backward
) = flash_attn.bwd(
.../torch/_ops.py:1279, in __call__
RuntimeError: torch_call_dispatcher( "aten::sum", "IntList_out", stack.data(), (...)) API call failed
at .../torch-2.13.0/.../torch/include/torch/csrc/stable/ops.h, line 1036
| torch |
v2 (cxx11) |
v3 (stable-abi210) |
| 2.12.1+cu130 |
MHA ok, GQA ok |
MHA ok, GQA raises |
| 2.13.0+cu130 |
no CUDA variant |
MHA ok, GQA raises |
flash_attn_varlen_func fails the same way, so it looks like the bwd op's aten::sum call through the stable ABI rather than anything varlen-specific. The overload exists on both torch versions (torch.ops.aten.sum.overloads() lists IntList_out). One possibly relevant detail: the error expands TORCH_ABI_VERSION to 2.13 and points at a nix-store torch 2.13.0 header, so the stable-abi210 build seems to have been compiled with a 2.13 ABI version rather than 2.10.
Same thing through transformers, any GQA model, here Qwen2.5-7B (28 query heads, 4 kv heads):
import torch
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained(
"Qwen/Qwen2.5-7B",
attn_implementation="kernels-community/flash-attn2",
dtype=torch.bfloat16,
).cuda()
input_ids = torch.randint(0, 100, (1, 32), device="cuda")
loss = model(input_ids=input_ids, labels=input_ids).loss # forward: fine
loss.backward() # RuntimeError, same as abovefrom_pretrained and the forward pass both succeed, the first .backward() raises. So inference is unaffected, but fine-tuning any GQA model (which is to say most models) is blocked wherever v3 is what gets resolved. That is already the case on torch 2.13, where v3 is the only build with a CUDA variant (#1082).
Env: kernels 0.16.0, transformers 5.16.0.dev0, H100 (cc 9.0), torch-stable-abi210-cu130-x86_64-linux, torch
2.13.0+cu130 and 2.12.1+cu130.
The v3 stable-ABI build runs the forward pass fine, but any backward with
num_kv_heads != num_headsraises. MHA (num_kv_heads == num_heads) is unaffected, and the v2cxx11build is fine on the same torch, so this looks specific to the stable-ABI build rather than to a torch version.cxx11)stable-abi210)flash_attn_varlen_funcfails the same way, so it looks like thebwdop'saten::sumcall through the stable ABI rather than anything varlen-specific. The overload exists on both torch versions (torch.ops.aten.sum.overloads()listsIntList_out). One possibly relevant detail: the error expandsTORCH_ABI_VERSIONto 2.13 and points at a nix-store torch 2.13.0 header, so thestable-abi210build seems to have been compiled with a 2.13 ABI version rather than 2.10.Same thing through transformers, any GQA model, here Qwen2.5-7B (28 query heads, 4 kv heads):
from_pretrainedand the forward pass both succeed, the first.backward()raises. So inference is unaffected, but fine-tuning any GQA model (which is to say most models) is blocked wherever v3 is what gets resolved. That is already the case on torch 2.13, where v3 is the only build with a CUDA variant (#1082).Env: kernels 0.16.0, transformers 5.16.0.dev0, H100 (cc 9.0),
torch-stable-abi210-cu130-x86_64-linux, torch2.13.0+cu130 and 2.12.1+cu130.