AMD ROCm: 4-valna interleave optimizacija FP8 GEMM kernela za Instinct MI355X (CDNA 4) udostručuje register budžet po valu

AMD ROCm tim objavio je detaljan tehnički opis nove optimizacijske strategije za FP8 GEMM (General Matrix Multiply) kernele na AMD Instinct MI355X grafičkim procesorima s CDNA 4 arhitekturom. Optimizacija nadilazi standardni 8-valni ping-pong pristup uvođenjem 4-valnog interleave dizajna koji dramatično povećava register budžet po valu i eliminira memorijske konflikte na razini hardware-a.

Zašto 4-valni pristup dominira nad 8-valnim za FP8 GEMM?

Ključna razlika leži u alokaciji VGPR (Vector General Purpose Registers) registara. U 8-valnom dizajnu, koji postavlja “dva vala po SIMD jedinici”, svaki val mora dijeliti registarski prostor — što ograničava veličinu izlaznog bloka koji može obraditi. 4-valni pristup dodjeljuje jedan val po SIMD jedinici, dajući mu pristup punom budžetu od 512 VGPR registara.

Ovo direktno utječe na veličinu obradivog tile-a: 4-valni pristup može koristiti izlazne tile od 128×128 elemenata nasuprot 64×128 u 8-valnom dizajnu. Globalni tile za cijeli kernel iznosi 256×256×128, a svaki od četiri vala obrađuje četiri bloka od po 64×64 elementa koristeći 16×16 MFMA (Matrix Fused Multiply-Add) operacije unutar svakog bloka.

Kako 4-wave interleave eliminira memorijsku latenciju na MI355X?

Središnji izazov u GPU kernelima za matrično množenje je skrivanje memorijske latencije — GPU mora biti zaposlen MFMA računanjem dok čeka na učitavanje podataka iz memorije. 4-wave pristup rješava ovo kroz precizno prepletanje instrukcija na razini sub-tile granularnosti, bez oslanjanja na hardverski rasporedjivač.

Implementacija koristi __builtin_amdgcn_sched_barrier(0) za prisilno poštivanje redosljeda instrukcija — kompajleru se zabranjuje grupiranje svih memorijskih operacija iza računskih. Svaki interleaved_block poziv prekriva 16 MFMA instrukcija s 8 LDS-to-register učitavanja (ds_read_b128), postižući konzistentno prekrivanje bez ping-pong sinkronizacijske složenosti.

Kako XOR-bazirani LDS swizzle eliminira konflikte memorijskih banaka?

MI355X LDS (Local Data Share) memorija organizirana je u 64 banke od po 4 bajta, s periodom od 256 bajta. Problem je strukturni: 16×16×128 MFMA instrukcije inherentno pristupaju stupcima matrice (column-major), što pri standardnom razmještaju uzrokuje 8-struke konflikte banaka — osam pristupa istovremeno čeka na istu banku.

AMD rješava ovo matematičkom transformacijom offseta:

row_bits = (offset % (16 * 128) >> 7) / 2
mask = row_bits << 4
swizzled_offset = offset ^ mask

XOR operacija remapira adrese tako da se sve četiri faze ds_read_b128 instrukcija odvijaju bez ijednog konflikta banaka. Rezultat je puna propusnost LDS memorije umjesto efektivnih 12,5% u slučaju 8-strukih konflikata.

Hardverska platforma i chiplet-aware optimizacija

Sve optimizacije ciljaju AMD Instinct MI355X s CDNA 4 arhitekturom. MI355X ima 8-chipletnu (XCD) topologiju — svaki XCD ima vlastiti L2 cache, a Last Level Cache se dijeli između XCD-ova. ROCm tim doda chiplet-aware grid swizzle koji remapira dodjelu thread blokova nitima koristeći morton curve-like mapping, maksimiziranjem iskorištenosti L2 unutar jednog XCD-a i LLJ između XCD-ova.

Benchmarking je korišten s 1000 warm-up i 1000 benchmark iteracija na normalno distribuiranim podacima s rotirajućim bufferima. ROCm tim navodi “konzistentne performanse kroz različite ROCm verzije” bez potrebe za ručnim pragma tuning-om — za razliku od 8-valnih varijanti koje su bile osjetljive na promjene kompajlera.