MoonMath AI har udgivet en open source bf16 forward-attention-kernel skrevet i HIP til AMDs MI300X og hævder bedre ydeevne end AMDs AITER v3 på alle testede shapes og afrundingsmodi. De rapporterede geomean-forbedringer er cirka 1.18×, 1.15× og 1.08× (per rounding-mode) med peak omkring 1.26×. For inference-workloads, hvor attention er tung, er den størrelsesorden værd at jagte – hvis tallene holder i jeres miljø.
En hurtigere forward-kernel kan sænke latenstid eller øge throughput uden at røre applikationslaget. Det er ren motoroptimering. Ikke mere pynt.
Hvad er udgivet
Der er tale om en fused forward-attention-kernel i bf16 til AMD MI300X (CDNA3) med gfx942 som ISA-target. Implementeringen er i HIP, ikke håndskrevet assembly, og udgivet under MIT-licens. Kernen understøtter inputlayouts BSHD og BHSD uden transpose, head_dim på 128, vilkårlige sekvenslængder og cross-attention.
Begrænsningerne er tydelige: ingen causal mask, ingen GQA og ingen varlen batching. Velegnet til udvalgte workloads – men ikke en direkte erstatning i klassisk autoregressiv serving, hvor causal masking og uens længder er normalen.
Hvordan kernen arbejder
Kernen beregner softmax(QKᵀ/√d)·V i én pipeline og leverer bf16-output. Pipeline-designet forsøger at holde matrixkernen travl, mens softmax og loads overlapper, hvor det kan lade sig gøre. Stram orkestrering, uden at hælde hele vejen over i specialiseret assembly.
Praktisk betyder det mindre friktion ved review og portering: HIP-kode med et par velvalgte lav-niveau-greb.
Hvorfor MoonMath melder speedup
Geomean-speedups opgives som cirka 1.18×, 1.15× og 1.08× (per rounding-mode) med topgevinster op til cirka 1.26×. Geomean udjævner variationer på tværs af shapes og kan skjule både tabere og vindere. Ligger produktionen i bestemte kombinationer af batch, heads og sekvenslængde, kan jeres reelle gevinst afvige en del fra geomean.
Miljøet bag tallene er ikke specificeret i detaljer (batchstørrelser, seqlens, warmup, driver-versioner). Læs derfor resultaterne nøgternt – og mål selv, før I beslutter jer.
Kernen i tricket
Den tekniske finesse er små inline-asm-wrappere: én instruktion i en __device__ __forceinline__-funktion med udvidede constraints. Så bindes input og output korrekt, mens compileren styrer registertildeling. Det undgår overflødige v_mov på akkumulatorer og holder koden tæt på ren HIP.
Det er en mellemvej. Intrinsics kan blive for abstrakte, fuld inline-asm er skrøbelig. Her styres få kritiske instruktioner – resten overlades til compileren.
Memory-placement og mikroarkitektur
MoonMath tilskriver store dele af gevinsten memory-placement: K i LDS med dobbeltbuffer, V holdt varm i L1, og Q samt akkumulatorer i registre. På CDNA3 afhænger MFMA-gennemløb af, at data lander rigtigt til tiden. Ellers forsvinder headroom hurtigt.
Der peges også på MFMA-tiling, fx 16×16×16 bf16 frem for 32×32×8, for at sænke registertryk på akkumulatorer og skabe plads til dybere prefetch. Tankegangen ligger i familie med nyere FlashAttention-ideer, men justeret til CDNA3s asynkrone memory-moves.
Numerik samlet ét sted
Outputs er bf16. Rounding følger AITER per mode. Ifølge udmeldingen ligger alle endelige værdier inden for 1 bf16-ULP af AITER. NaN/Inf håndteres bit-identisk. Resultater beskrives som deterministiske. For drift er pointen enkel: mindre støj ved kernel-skift og en skarp kontrakt for CI, hvor afvigelser kan fanges tidligt.
Det gør validering håndterbar og reproducerbar.
Begrænsninger og åbne spørgsmål
Kernen er designet til og valideret på MI300X/gfx942. HIP hjælper portabilitet, men asm-wrappere kan i praksis binde til specifikke opcodes; adfærd på andre AMD-acceleratorer eller efter en ISA-revision er ikke afklaret i materialet.
Funktionelt mangler causal mask, GQA og varlen batching – features mange produktionsmodeller forventer. Udviklingstest omtales som udført på bare-metal via HotAisle i den primære dækning; uden flere detaljer er det kontekst, ikke en reproduktionsvejledning.
Community-signal
En community-PR i SGLang omtales med cirka 1.23× hurtigere Wan2.1 video-diffusion uden rapporteret kvalitetsregression. Interessant som casenote, men ikke en uafhængig bekræftelse.
Hvad det betyder i praksis
Kører I MI300X i produktion, kan det her være en drop-in gevinst i udvalgte inference-trin – for de former, der passer. Mål i jeres egen stack, på jeres batchmønstre og seqlens, og på de driver-/firmware-versioner I faktisk anvender. Små forskelle i ROCm, compiler og BIOS kan flytte nålen.
Udnyt determinismen i CI: byg per-mode rounding-tests, ULP-kontrol pr. tensorblok og edge-cases som lange sekvenser og ekstreme logits. Mål både microbenchmarks og end-to-end throughput/latenstid, så gevinsten ikke forsvinder i scheduler-overhead.
Autoregressiv serving, de praktiske hager
Manglen på causal mask blokerer for standard LLM-serving. Ingen GQA og ingen varlen batching gør udnyttelse i produktionsservere sværere, hvor tokens pr. sekund ofte afhænger af at kunne mixe længder og grupper. En realistisk integrationsvej er at beholde AITER til de autoregressive trin og teste MoonMath-kernen i ikke-causale sekvenser og cross-attention, hvor formen passer.
Vil man bygge bro, kræver det ingeniørarbejde: effektiv causal masking med bibeholdt memory-placement, eller udvidet inputhåndtering til varlen – uden at sprænge registerbudgettet.
Fra mikrobenchmark til driftstal
Map et repræsentativt mikrobenchmark-shape til en konkret servicevej. Vælg jeres hyppigste B, H og S, og mål kernel-tid for attention alene. Oversæt det til p50/p95 for et helt request ved at indregne de øvrige lag og framework-overheads i en end-to-end-profil.
Så bliver en 1.18× kernel-speedup til et tal, I faktisk kan handle på – fx om den ender som 1.08× på p95, afhænger af resten af stakken. Tjek for nye flaskehalse: KV-cache, allocatorer, dispatch.
Portering, drift og versionstyring
HIP sænker adgangsbarrieren i forhold til ren assembly, men performance er følsom over for hele ROCm-stakken, inkl. compiler, driver og BIOS/firmware. Frys versioner, dokumentér build-flags og genmål efter hver opdatering. Små ændringer i scheduling eller registertildeling kan spise gevinsterne, hvis wrappere og compiler ikke længere spiller sammen.
Hav en ren backout-vej, så I hurtigt kan skifte tilbage til AITER, hvis noget skrider i produktion.
Risici og tradeoffs
Open source giver indblik og hurtige iterationer, men kan også skubbe vedligehold over på brugerne, hvis upstream står stille. Kompatibilitet med ROCm eller ROCm-next er central. Skifter compilerens registerallokering adfærd, kan wrappere kræve justering for at undgå ekstra kopier og lavere occupancy.
Planlæg test, monitorering og fallback som en del af rollout – ikke bagefter.
Hvad man bør gøre før adoption
Et kort checkpoint før udrulning:
- Benchmark-protokol: Fastlæg shapes (B, S, H, D=128), tre rounding-modes, 3–5 sekvenslængder inkl. en meget lang. Brug warmup og derefter mindst 100 iterationer pr. shape.
- Numerisk validering: Mål ULP-afvigelse mod reference (AITER), og tjek bit-identisk NaN/Inf. Flag al ikke-determinisme.
- Stack-versioner: Frys ROCm, compiler, kernel-driver og firmware. Log build-flags og GPU-tilstande.
- Driftsprofil: Mål både latency p50/p95 og tokens pr. sekund under realistisk batching og load.
- Fallback-plan: Hav en konfig-switch tilbage til AITER ved regressioner.
Hvad konkurrenter og AMD sandsynligvis gør
For AMD ligger oplagte svar i at optimere AITER yderligere på MI300X eller indarbejde de bedste ideer i officielle biblioteker. Leverandører og cloud-udbydere vil formentlig teste integration i inference-rammer, især til diffusion og cross-attention, hvor fravær af causal mask ikke spænder ben for brug.
Afslutning
En relativt ren HIP-tilgang med få, målrettede asm-wrappere kan – i de rapporterede scenarier – slå AMDs egen kernel. Det peger på, at der stadig er tocifrede procenter gemt i softwaren, når memory-placement og instruktionsvalg rammer rigtigt. Matcher jeres workloads begrænsningerne, er her potentiel ydeevne at hente. Mål i jeres miljø, før I ruller bredt ud. Man opdager først forskellen, når man sidder med det i hænderne.
