Az ügynöki (agentic) és hosszú kontextusú munkaterhelések terjedésével a kontextushosszok növekednek, és az attention aránya a teljes inferenciaidőn belül jelentősen megnő. Mivel ma már az attention jelenti a költségdomináns részt, nemcsak a megvalósítás, hanem maga az architektúratervezés is egyre inkább meghatározza a modell futtatási teljesítményét. E gondolat mentén NVIDIA vizsgálta meg, hogy a group size (egy KV fejhez tartozó query-fejek száma, G), a fejdimen zió (Hsz) és a szekvenciahossz hogyan befolyásolja a sűrű attention (dense attention) teljesítményét, valamint hogyan párhuzosítják az attentiont GPU-kon.
Elemzésük két alapra támaszkodik: analitikus képletekre a GEMM-alakú számításokból és mérésekre prefill és decode kernel-ekről FP8 használatával mind attention számításhoz, mind a KV cache olvasásához.
Fontos rövidítések és értemezés
- PB: Prefill batch size
- DB: Decode batch size
- QH: query head-ek száma
- KH: KV head-ek száma (MHA: KH = QH, GQA: KH = QH/G, MQA: KH = 1)
- G: group size = QH/KH (annyi query head osztozik egy KV head-en)
- Hsz: head dimenzió (tipikusan 64, 128 vagy 256)
- ISL: input sequence length (prefillben a lekérdezett tokenek száma)
- KVSL: átlagos KV cache hossz egy decode iteráción
Prefill és decode: két külön feladat
- Prefill: a teljes promptot párhuzamosan dolgozza fel, nagy GEMM-M (= ISL × G) matmulokra, amelyek jellemzően compute-bound (számítás-korlátozottak).
- Decode: spekulatív dekódolás nélkül tokenenként dolgozik, kis GEMM-M (= G) matmulok keletkeznek, és a KV cache-ből történő HBM-olvasások miatt memory-bound (memória-korlátozott). Spekulatív dekódolás növelheti a GEMM-M értékét, eltolva a decode-ot compute-bound irányba.
A két fázis fő különbségeit táblázatos formában is érdemes megjegyezni: a lekérdezés hossza (ISL vs 1 token), a KV kontextus (ISL vs teljes KV cache), az attention GEMM-M és a domináns szűk keresztmetszet más.
Arithmetic intensity és a roofline viselkedés
A GPU teljesítményét a roofline modell a számítási és sávszélességi korlátok mentén határolja. Az arithmetic intensity (FLOP-ok / memória-bájtok) dönti el, hogy egy kernel compute- vagy memory-bound. Prefill tipikusan a ridge pont felett van (compute-bound), míg decode a rampen van (memory-bound). Spekulatív dekódolás növeli a decode arithmetic intensity-jét.
Hogyan számítja a FlashAttention az attentiont GPU-n?
A FlashAttention nem materializálja a teljes attention mátrixot; HBM-ból SRAM-be streameli a Q, K és V darabjait, és egyetlen átfutásban fűzi össze a következő lépéseket:
- BMM1: lekérdezések és kulcsok pontszámait kiszámolja (Q·K^T)
- Online softmax: a pontszámokat futó max és összeg segítségével normalizálja
- BMM2: a normalizált súlyokkal aggregálja a V-ket A BMM-ek Tensor Core-okon futnak, a softmax exponenciális része pedig a special-function egységeken.
GEMM-alakok prefill és decode esetén
A BMM1 és BMM2 dimenziói (Batch, M, N, K) fázisonként a következők:
- BMM1 prefill: Batch = PB × KH, M = ISL × G, N = ISL, K = Hsz
- BMM1 decode: Batch = DB × KH, M = 1 × G, N = KVSL, K = Hsz
- BMM2 prefill: Batch = PB × KH, M = ISL × G, N = Hsz, K = ISL
- BMM2 decode: Batch = DB × KH, M = 1 × G, N = Hsz, K = KVSL
Fontos megjegyzés: decode-nál a GEMM-M általában G = 8–16 körüli, ami jóval kisebb egy GPU tile-M (>64) értékénél, ezért kevés párhuzamos munkát kínál egy tile számára. Nagyobb G csökkenti az egy tokenre eső KV betöltést és jobban amortizálja a memóriaműveleteket.
Group size (G)
A group size azt mutatja, hány query head osztozik egy KV head-en. MHA-nál G=1, GQA-nál tipikusan G=4,8,16…, MQA-nál G=QH.
- Prefill: az arithmetic intensity közelít 2×ISL felé nagy G esetén. Gyakorlatilag prefillt az ISL dominálja, ezért G változtatása kismértékben befolyásolja a prefill futási időt (G=1…64 között <1% változás egy mért esetben).
- Decode: az arithmetic intensity közelítően ≈ 2×G, tehát G duplázása nagyjából duplázza a decode arithmetic intensity-jét és ekvivalens mértékben csökkentheti a runtime-ot. Például G növelése 1-ről 8-ra akár ≈8× hatékonyságnövekedést eredményezhet memóriaterhelés csökkenése miatt. A mérések szerint decode runtime körülbelül 2× javul minden G duplázásnál, amíg más fix overhead-ek és a flash-decoding párhuzamosítási korlátai nem dominálnak.
Irányelv 1: decode-hatékonyság miatt érdemes G-t magasra választani; prefill érzéketlen rá. Spekulatív dekódolás további eszköz a hatékonyság növelésére adott G mellett.
Head dimension (Hsz)
A Hsz nem változtatja meg az arithmetic intensity-t, mert egyszerre növeli a FLOP-okat és a mozgatott bájtokat. Ennek ellenére a mért futási idők növekednek Hsz szerint, mivel az attention-kernelt háromféle munkaforma befolyásolja eltérő módon:
- Matmul: Hsz növelésével nő, de a GPU tile-okhoz igazodóan lépcsőzetesen (a javasolt dimenziók többszörösei hatékonyabbak).
- Memória (KV állapot): nő Hsz-szal, és a GPU 128 bájtos mozgatási egységei miatt a 128 többszörösei hatékonyak.
- Softmax: független Hsz-tól, mert a softmax a pontszám mátrixon operál (ISL × ISL vagy 1 × KVSL szerkezet).
Gyakorlati következmény: Hsz=128 vagy 256 általában a legjobb kompromisszum: sorba áll a GPU tile- és cache-szervezésével, miközben nem közelíti meg a tensor memory (TMEM) kapacitási határait; Hsz=64 gyakran „kifizeti” a 128-as tile költségét, Hsz≥512 viszont a TMEM korlátjához nyomhat.
Irányelv 2: használjon Hsz=128 vagy 256-ot, hogy illeszkedjen a GPU hardverhez és 128-bájtos átvitelekhez.
Szekvenciahossz (ISL / KVSL)
A szekvenciahossz aszimmetrikusan hat prefillre és decode-ra:
- Prefill: O(ISL²) műveletigénnyel skálázódik, mert minden token minden tokenre figyel. Az arithmetic intensity ISL-függőként növekszik, így prefill compute-bound marad és óriási ISL-nél a futásidő négyszereződik ISL duplázásakor (amíg a fix overhead-ek nem dominálnak).
- Decode: minden lépésnél a teljes KV cache-t olvassák, így a per-step költség KB-ben lineárisan nő KVSL-lel. Decode kvar(arányosan) O(KVSL) szerint skálázódik és marad memory-bound. Duplázás kb. duplázza a decode runtime-ot, rövid KVSL-nél azonban a fix overhead-ek visszafogják az ideális skálázást.
Irányelv 3: csökkentse a hatékony KV állapotot, ahol lehetséges — például KV-cache tömörítéssel, ritkább vagy csúszó ablakos attentionnel, vagy hibrid architektúrákkal (például Nemotron 3), ahol csak egyes rétegek viselik a globális KV állapot növekedését.
Párhuzamosítás: Tensor parallelism és a KH szerepe
A Tensor Parallelism (TP) fejenként osztja meg az attentiont a GPU-k között: minden GPU QH/TP query head-et és KH/TP KV head-et kap. TP gyakorlatilag a head-ek mentén shard-ol, nem a tokeneken.
Gyakorlati korlát: a KH-nak egyenletesen kell oszlania a GPU-k között. Ha TP > KH, akkor egy KV head több rank-re terjed, és minden rank-re másolat kerül a KV cache-ből, ami memória- és sávszélesség-túlterhelést okoz duplikáció nélkül nem javítva a párhuzamosítást. Ezért érdemes TP ≤ KH értéket tartani, hogy minden GPU legalább egy teljes csoportot (egy KV head és hozzá tartozó G query head) birtokolhasson.
Kevés KV head-et használó modellek (például Nemotron 3 két KV fejjel vagy MQA egy KV fejjel) gyorsan kimerítik a TP lehetőségeit. Ilyen esetekben jobb stratégiák: Attention Data Parallelism (ADP) a kérés-szintű shardoláshoz, vagy KV Parallelism (KVP) a hosszú KV cache-ek shardingjához, míg az FFN-eket Expert Parallelism (EP) skálázhatja. TensorRT-LLM például a Wide EP-t (ADP + EP) és a Helix Parallelism-t (KVP + EP) használja.
Irányelv 4: hagyja, hogy a KH határozza meg a párhuzamosítási stratégiát. Tartsa TP ≤ KH, és kevés KH-t használó modelleknél válasszon ADP- vagy KVP-alapú skálázást plusz EP-t az FFN-re.
Összefoglaló ellenőrzőlista (co-design ajánlások)
- Válassza meg a group size-ot (G) a decode-hatékonyság érdekében, és törekedjen nagy G-re; prefill nem érzékeny rá.
- Állítsa a head dimenziót Hsz=128 vagy 256-ra a hardverrel való illeszkedés és a TMEM-költségek optimalizálása miatt.
- Minimalizálja a hatékony KV állapotot (KV-cache tömörítés, sparse/sliding attention, vagy hibrid modellek), mert prefill O(ISL²)-t, míg decode O(KVSL)-t skáláz.
- A párhuzamosítás stratégia tervezésénél KH adja az irányt: tartsa TP ≤ KH; kevés KH esetén használjon ADP/KVP + EP megoldásokat (például Wide EP vagy Helix Parallelism).
Köszönetnyilvánítás
A poszt az NVIDIA több csapatának együttműködésében készült. Köszönet Timmy Liu, Jatin Mitra, Tiyasa Mitra, Bhargava Gopireddy, Brian Pharris, Julien Demouth és Eduardo Alvarez közreműködéséért.
