Check that torch has flash attention before taking the flash path (fixes CUDA on Windows)
On Windows, the official CUDA builds of torch ship without flash attention. _fused only checks the GPU generation, so any Ampere or newer card takes the flash path and the first system_one call fails:
RuntimeError: USE_FLASH_ATTENTION was not enabled for build.
attention.py, line 80, in causal_attention
out = torch.ops.aten._flash_attention_forward(...)
Setup: Windows 11, RTX laptop GPU (Ampere or newer), torch 2.6.0+cu124, transformers 5.22.2, bf16.
This PR adds torch.backends.cuda.is_flash_attention_available() to _fused. When flash is missing, every call site already falls back to its SDPA path (causal_attention, _windowed and Sees.attend), so nothing else changes. On builds that have flash, _fused returns the same value as before.
def _fused(q: torch.Tensor) -> bool:
"""Whether the flash kernels take these queries: a GPU in half precision, ROCm or CUDA from Ampere on."""
if not q.is_cuda or q.dtype not in (torch.float16, torch.bfloat16):
return False
if not torch.backends.cuda.is_flash_attention_available():
return False
return bool(torch.version.hip) or torch.cuda.get_device_capability(q.device)[0] >= 8
Tested with this change on the machine above: system_one runs, and our 20-sentence French intent set (one choice question, 9 options) gives the same answers as with _fused forced off, at a median of 91 ms per call. With is_flash_attention_available forced to True, the patched _fused takes the flash path again, so builds that have flash keep their current behaviour.
The same check is in d1-Omni-2B/attention.py (identical file) and d1-3B/hybrid.py (_fused, line 123). I can open the same PR there if that helps.
Context: we're testing d1-2B as the intent router of a local voice assistant (Mnemosyne OS), on French speech. Thanks for the open d1 release!