Új Triton kernel gyorsítja a Vision Transformer modelleket az ArXiv szerint

A Vision Transformer (ViT) modellek token-metszési módszerei elméletileg jelentős, másodfokú csökkenést ígérnek a figyelmi FLOP-okban azáltal, hogy elhagyják az informatívnak nem ítélt képrészleteket. Azonban a gyakorlatban, a legmodernebb változó hosszúságú figyelmi API-k, mint a FlashAttention-2 varlen vagy a PyTorch NestedTensor SDPA használatakor, a valós idejű figyelmi késleltetés nem csökkent arányosan — írja az ArXiv ML.

A kutatók szerint a probléma egy úgynevezett dispatch-overhead szűk keresztmetszetben rejlik. A ViT-kre jellemző rövid, metszés utáni szekvenciahosszak (197 token vagy kevesebb) esetén a tényleges mátrixaritmetika mindössze néhány mikroszekundum alatt lezajlik, miközben a gazdagép oldali dispatch útvonal 60-90 mikroszekundumot emészt fel.

Erre a problémára válaszul fejlesztettek ki egy könnyűsúlyú, kétirányú Triton figyelmi kernelt. Ennek a kernelnek a dispatch alsó határa 40 mikroszekundum, ami körülbelül 1,5-szer alacsonyabb, mint a FlashAttention-2 varlen esetében. Ez a csökkentés lehetővé teszi, hogy a token-metszésből származó megtakarítások jobban megmutatkozzanak a valós idejű teljesítményben.

Az új rendszert egy teljes pack-attend-unpack folyamatba integrálták. A fejlesztők szerint a rendszer akár 2,24-szeres végpontok közötti átviteli sebességet is elérhet a párnázott PyTorch SDPA-hoz képest, ami jelentős előrelépést jelent a Vision Transformer modellek hatékonyságában.