For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/get-started.md).
Mojo function
flash_attention_hw_supported
def flash_attention_hw_supported[qkv_type: DType]() -> Bool
Return True if the current GPU supports flash attention for qkv_type.
NVIDIA GPUs support all dtypes. AMD GPUs require bfloat16 or a
float8 type. Returns False on CPUs and unsupported GPU types so
callers can gracefully fall back to a reference implementation.
Parameters:
- qkv_type (
DType): The element data type of the Q/K/V tensors.
Returns:
Bool: True when flash attention is available for qkv_type on the
detected accelerator.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!