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

depth_supported_by_gpu

def depth_supported_by_gpu[depth: Int, mask_t: MHAMask, config: MHAConfig[config.dtype], info: GPUInfo]() -> Bool

Return True if the given head depth is supported for flash attention on this GPU.

Checks the combination of depth, GPU architecture (info), and algorithm variant to decide whether the optimised kernel path is available. For example, depth 128 is universally supported, depth 64 requires SM80+, depth 512 requires SM100 or AMD.

Parameters:

depth (Int): Attention head depth (key/value dimension per head).
mask_t (MHAMask): Mask type; some depths require mask_safe_out_of_bounds.
config (MHAConfig[config.dtype]): MHA tile configuration, used to check the algorithm variant.
info (GPUInfo): GPU architecture descriptor.

Returns:

Bool: True when the optimised flash-attention kernel supports depth on the given GPU.