Flere frontmodeller bruker sparse attention for å presse ned inferenskostnaden, men ingen har lagt ut kode for å trene det effektivt. Det er utgangspunktet utvikleren bak Flash-MSA selv beskriver i bloggposten sin: kjerner skrevet i CuTeDSL for Hopper- og Blackwell-GPUer, utviklet på leide H100- og B200-kort.
«Dette er ikke en offisiell implementasjon, og jeg er ikke tilknyttet MiniMax.» — utvikleren bak Flash-MSA
MSA skiller seg fra DeepSeeks sparse attention på tre punkter. Utvelgelsen skjer i blokker på 128 nøkler i stedet for enkelttokens, noe som gir bedre caching. Hovedattention bruker GQA i stedet for MLA, og det er poenget som gjør kjernene brukbare utenfor de kinesiske labbene: formuleringen i GLM-5.2 og DeepSeek V4 er tilpasset MLA, som vestlige labber ikke trener med. Og hvert proxy-hode kan velge sitt eget utvalg nøkler, i stedet for å score hele laget under ett.
Det tekniske grepet som gjør treningen billig, ligger i bakoverpasseringen. Ved å utlede KL-leddet analytisk blir gradienten til proxy-scoren rett og slett proxy-sannsynligheten minus hovedsannsynligheten, så tapet aldri må materialiseres i delt minne. Blokkindeksene fra foroverpasseringen gjenbrukes, og bare proxy-foroverpasseringen er kvadratisk i kontekstlengde. Resten kjører lineært.
Korrektheten er sjekket mot en referanseimplementasjon i PyTorch, med cosinuslikhet mellom 0,9985 og 0,9996 for både utdata og gradienter i bf16, der vanlig toleranse ligger på 0,01. Begrensningen utvikleren selv peker på: den fusjonerte bakoverpasseringen bruker 138 registre per tråd og 105 KB delt minne, som låser den til én CTA per SM og gir en teoretisk okkupans på 12,5 prosent mot 18,75 prosent i Flash-Attention. Kontekstparallellisme mangler også, så du treffer minnetaket før du kommer veldig langt ut i kontekstlengde. For alle som har villet eksperimentere med sparse attention i pretrening uten å skrive CUDA-kjerner selv, er dette likevel første gang koden finnes.