Skip to content

b12285

  • b12285
  • 9d7b2d7
  • Verified

    This commit was created on GitHub.com and signed with GitHub’s verified signature.
  • Choose a tag to compare

  • b12285
  • 9d7b2d7
  • Choose a tag to compare

  • Verified

    This commit was created on GitHub.com and signed with GitHub’s verified signature.
@sihanyu03 sihanyu03 tagged this 11 Nov 08:45
* cuda : add MMA FlashAttention support for head sizes 40 and 72

Head sizes 40 and 72 (used e.g. in vision towers of Gemma 3/4 and Qwen3-VL) were excluded from every tensor-core dispatch branch and fell back to the tile kernel.

The fp16 MMA fragment covers 16 halfs in the head direction, so the in-kernel geometry is padded to the next multiple of 16 (40 -> 48, 72 -> 80). Register arrays, shared memory tile strides and MMA loop bounds use the padded size, while global memory loads and stores the real one.

Enabled on NVIDIA Turing and newer while Volta and AMD keep the tile kernel. The RDNA/CDNA config rows and the mma_tile_sizes specialisations are needed because the explicit instantiations are compiled for every target of a translation unit.

* tests : add the head size = 40 encoder perf shape

* cuda : address review comments on the head size 40 and 72 kernels
Assets 2
Loading