Luca Vaudano's Blog

Inference Enginering - Prefill vs decode

Questo è il mio primo blogpost in italiano. Questo è il primo appunto del learning path Inference engineering di Deep-ML. L'obiettivo è spiegare la differenza tra il prefill e il decode e perché vengono rispettivamente considerati compute-bandwith-bound e memory-bound.


1. Le due fasi dell'inferenza

Nel contesto dell'inferenza di un LLM, prefill e decode sono i due step che ne caratterizzano il funzionamento.

Il primo step, il prefill, consiste nel primo forward pass del modello: legge il system prompt e il prompt inserito dall'utente, costruisce la KV cache e produce anche il primo token della risposta. Tutti gli N token di ingresso vengono processati insieme, in un solo passaggio. Questa operazione è compute-bound, cioè limitata dalla quantità di calcoli che l'hardware riesce a fare.

La seconda operazione, il decode, è quella che in maniera auto-regressiva produce un token dopo l'altro. Il nome non ha niente a che vedere con la detokenizzazione, cioè la conversione degli ID numerici in stringhe di testo: deriva dalla fase di generazione del decoder. Questa operazione è memory-bandwidth-bound, cioè limitata dalla banda di memoria.


2. Arithmetic intensity e ridge point

Lo strumento che permette di dimostrare la differenza è l'arithmetic intensity (AI). La definizione è semplice: quanti FLOP eseguo per ogni byte che leggo dalla memoria.

AI = FLOP eseguiti / byte trasferiti da HBM        [FLOP/byte]

Il ridge point è invece una proprietà dell'hardware, non del carico di lavoro. È l'intersezione tra il limite di banda della memoria e il limite di throughput di calcolo di picco:

ridge point = peak FLOP/s / peak bandwidth         [FLOP/byte]

La distinzione è che l'AI descrive il lavoro che sto facendo, il ridge point descrive la macchina su cui lo faccio. Il confronto fra i due dice in quale regime mi trovo.

Prendendo una H100 SXM di NVIDIA, che ha circa 990 TFLOP/s in BF16 dense e circa 3.35 TB/s di HBM3, il ridge point è di circa 295 FLOP/byte. Sotto quella soglia il carico è memory-bound, cioè le unità di calcolo restano ferme ad aspettare i dati. Sopra è compute-bound, cioè la memoria riesce a stare al passo.

Roofline model: prefill contro decode su H100 throughput (FLOP/s) decode, B=1 decode, B=64 prefill 1 10 100 295 1k 10k memory-bound compute-bound arithmetic intensity (FLOP/byte)

3. Dove vivono i dati

Per capire perché il decode resti bloccato sulla banda serve sapere com'è fatta la memoria di una GPU.

Gerarchia di memoria di una GPU H100 HBM3 · 80 GB pesi e KV cache stanno qui 3.35 TB/s L2 cache · 50 MB condivisa da tutti gli SM SM 0 L1 + shared 4 tensor core SM 1 L1 + shared 4 tensor core SM 131 L1 + shared 4 tensor core 132 SM in totale · 256 KB di L1/shared ciascuno

HBM sta per High Bandwidth Memory: è la memoria grande della GPU, dell'ordine delle decine di GB, e contiene i pesi del modello e la KV cache. Sotto c'è la L2, da 50 MB, condivisa da tutti gli SM. Sotto ancora ci sono i 132 Streaming Multiprocessor, ognuno con la propria L1/shared memory da 256 KB e i propri quattro tensor core.

Un modello da 70B in BF16 occupa 140 GB di pesi, mentre la L2 ne tiene 50 MB. Sono tre ordini di grandezza di differenza. Non esiste nessuna strategia di caching che salvi la situazione: i pesi sono nel HBM e da lì vanno riletti a ogni passaggio.

Il prefill se lo può permettere, perché carica un blocco di pesi nella shared memory di un SM e lo riusa su tutte le N righe dell'input: il traffico da HBM viene ammortizzato su N token. Il decode no: carica lo stesso blocco, lo usa per un token, lo scarta.


4. I conti delle due fasi

Chiamo P il numero di parametri del modello e considero che ogni parametro occupi 2 byte (BF16).

Prefill su N token. Leggo i pesi una volta sola e li riuso per tutti gli N token. I FLOP eseguiti sono 2·N·P, i byte letti sono 2P, quindi l'AI vale circa N:

AI_prefill = 2·N·P / 2P = N

Con un prompt di 1000 token ottengo circa 1000 FLOP/byte, ben sopra il ridge point di 295. Il prefill è compute-bound.

Decode a batch 1. Leggo esattamente gli stessi 2P byte, ma li uso per produrre un solo token:

AI_decode = 2P / 2P = 1

Circa 1 FLOP/byte, cioè trecento volte sotto la soglia. Il decode è memory-bandwidth-bound.

Una conseguenza secondaria dello stesso conto: dato che AI_prefill ≈ N, esiste una lunghezza di prompt sotto la quale anche il prefill scivola in regime memory-bound. Su H100 quella soglia sta intorno ai 300 token. È una stima grossolana, perché trascura i FLOP dell'attention e gli overhead di lancio dei kernel, ma l'ordine di grandezza regge: un prompt di trenta token non satura una H100.


5. Il batching e il tetto asintotico

La reazione naturale davanti a un AI di 1 è processare più sequenze insieme. Se elaboro B sequenze nello stesso step, i pesi li leggo comunque una volta sola e li riuso per B token diversi: i FLOP diventano 2·B·P mentre i byte restano 2P, quindi l'AI vale circa B. Sembrerebbe risolto, basta portare B a 295.

Non è risolto, perché quel denominatore è incompleto. A ogni step di decode non leggo solo i pesi: leggo anche la KV cache, dato che l'attention deve confrontare il token nuovo con tutto il contesto precedente. E qui c'è un'asimmetria che cambia tutto. I pesi sono condivisi da tutte le sequenze del batch, la KV cache no: ogni sequenza ha la sua. Se raddoppio il batch, raddoppio anche i byte di KV da leggere.

La formula corretta è:

AI = 2·B·P / (2P + B·KV_seq)

I parametri da sostituire sono due e indipendenti: 2P, che dipende solo da quanto pesa il modello, e KV_seq, che dipende solo da quanta KV cache consuma una singola sequenza al contesto scelto. Il primo resta fisso qualunque cosa faccia il contesto, il secondo cresce linearmente con la lunghezza del contesto.

Numeratore e denominatore crescono entrambi con B. Finché il termine 2P domina il denominatore, la crescita del numeratore non trova contrappeso e l'AI sale quasi linearmente. Ma quando B·KV_seq diventa il termine grande, i due B si semplificano e resta 2P / KV_seq, un valore in cui B non compare più. Da quel punto in poi aggiungere sequenze non alza più l'intensità aritmetica: compra solo throughput, finché non si esaurisce la memoria.

Con numeri veri

Prendo Llama-3 8B con contesto da 8k token: 8 KV head (GQA), head_dim 128, 32 layer, 2 byte per valore, K e V.

KV per token = 2 × 8 × 128 × 2 × 32 = 128 KB
KV_seq (8k)  = 8192 × 128 KB ≈ 1 GB per sequenza
2P           = 16 GB

Sostituendo nella formula, con le quantità espresse in GB:

AI = 16·B / (16 + 1·B) = 16·B / (16 + B)

Quell'1 davanti a B è KV_seq, e sparisce dalla scrittura solo perché vale esattamente 1 GB in questo caso. È una coincidenza dell'esempio, non una proprietà della formula.

Saturazione dell'arithmetic intensity al crescere del batch AI (FLOP/byte) tetto = 2P / KV_seq = 16 0 5 10 15 20 0 64 128 192 256 batch size B ridge point H100 = 295 FLOP/byte, fuori scala verso l'alto

B=1 dà 0.94, B=16 dà 8, B=64 dà 12.8, B=256 dà 15.1. Quadruplicando il batch da 64 a 256, e quindi quadruplicando la memoria occupata dalle KV cache, l'AI cresce del 18%. Il tetto è 16 / 1 = 16 FLOP/byte contro un ridge point di 295: con questo contesto, su questo hardware, il decode non diventa compute-bound per nessun valore di B. Non è un problema di batch troppo piccolo, è un limite strutturale.

Rifaccio lo stesso conto abbassando il contesto a 512 token. Il modello non cambia, quindi 2P resta 16 GB; cambia solo KV_seq:

KV_seq (512) = 512 × 128 KB = 64 MB = 0.064 GB per sequenza

AI = 16·B / (16 + 0.064·B)

Ora B=16 dà 15, B=64 dà 51, B=256 dà 126, e il tetto sale a 16 / 0.064 = 250, a un passo dal ridge point.

contesto KV_seq tetto = 2P / KV_seq
512 0.064 GB 250
8k 1 GB 16
32k 4.2 GB 3.8

Il tetto è semplicemente il rapporto tra i byte di pesi e i byte di KV cache di una singola sequenza. Quadruplicando il contesto, il tetto si divide per quattro.

Da qui si capisce perché GQA e la quantizzazione della KV cache siano leve così importanti: agiscono su KV_seq, cioè alzano il tetto. Quantizzare i pesi invece agisce su 2P, il termine che a contesto lungo è già trascurabile. Sono due interventi che sembrano simili ma colpiscono termini opposti della stessa frazione.


6. Perché il decode lascia i tensor core in gran parte idle

L'ultima domanda richiede di scendere sotto il livello dell'aritmetica e guardare la forma delle operazioni.

Il prefill esegue un GEMM, cioè una moltiplicazione matrice per matrice. L'input è una matrice di forma [N × d], dove N sono i token del prompt. Il decode a batch 1 esegue un GEMV, cioè matrice per vettore: l'input è una singola riga di forma [1 × d]. La matrice dei pesi è identica nei due casi.

GEMM nel prefill contro GEMV nel decode prefill: GEMM × = N × d d × d N × d pesi riusati su N token decode: GEMV × = 1 × d d × d 1 × d pesi usati per 1 token stessa matrice di pesi, stessi byte letti da HBM, N volte meno lavoro

Da questa differenza di forma discendono due cause di inattività, distinte e indipendenti.

Prima causa: la forma non entra nel tile. Un tensor core non esegue moltiplicazioni scalari, esegue istruzioni MMA su blocchi di dimensione fissa, per esempio m16n8k16. L'istruzione pretende un frammento di input alto 16 righe. Nel decode a batch 1 le righe disponibili sono una: le altre quindici andrebbero riempite di zeri, pagando il ciclo intero per un sedicesimo di lavoro utile. A quel punto le librerie non usano affatto il tensor core per un GEMV e instradano il calcolo sui CUDA core. I tensor core sono idle in senso letterale: il kernel non li chiama.

Seconda causa: lo stallo su memoria. Anche ipotizzando efficienza piena, il calcolo di un GEMV finisce in una frazione del tempo necessario a portare i pesi dalla HBM all'SM. L'SM resta in attesa del blocco successivo. Nei profili di esecuzione questo appare come SM inattivi con il bus HBM saturo al 85-95%, l'esatto opposto del prefill, dove gli SM lavorano in continuazione e il bus sta intorno al 30%.

La differenza tra le due cause è operativa. La prima dice che lo strumento è sbagliato per quella forma, e il batching la risolve: già con B=16 il blocco si riempie. La seconda dice che comunque non ci sarebbe abbastanza lavoro da fare, e il batching non la risolve, perché come si è visto al punto 5 il tetto asintotico resta.


7. Cosa se ne fa un sistema di inferenza

Le due fasi hanno profili di risorse opposti sullo stesso hardware. Un H100 che esegue prefill usa i suoi 132 SM in continuazione e lascia la banda per metà libera; lo stesso H100 che esegue decode satura la banda e tiene le unità di calcolo quasi ferme. Il costo orario però è identico.

Il problema è aggravato dalle proporzioni tipiche di una richiesta. Una generazione da 300 token passa qualche centinaio di millisecondi in prefill e diversi secondi in decode: la GPU passa oltre il 90% del tempo nella fase in cui è sottoutilizzata.

Da qui nascono tre famiglie di rimedi.

Continuous batching. Raggruppare più sequenze nello stesso step di decode è la leva diretta sull'AI, con il limite asintotico visto sopra.

Chunked prefill. Spezzare un prefill lungo in blocchi e interlacciarli con gli step di decode. Migliora la ITL (inter-token latency) delle richieste già in generazione, che altrimenti si bloccherebbero per centinaia di millisecondi ogni volta che arriva un prompt lungo; il prezzo è un TTFT (time to first token) leggermente peggiore per la richiesta spezzettata. Mescolando blocchi ad alta AI con step a bassa AI alza anche l'utilizzo medio, ma non risolve il fatto che la stessa GPU debba coprire due profili opposti, e vada quindi dimensionata sul picco.

Disaggregazione prefill/decode. Separare fisicamente le due fasi su pool distinti, scegliendo hardware diverso per ciascuna. Per il prefill conviene privilegiare i FLOP di picco. Per il decode conviene privilegiare la banda HBM, che determina la latenza di ogni step e quindi la ITL, e la capacità HBM, che determina quante sequenze concorrenti stanno in KV cache e quindi il throughput. Quello che sul decode non serve è il picco di calcolo: pagarlo significa comprare tensor core che resteranno fermi.


8. Gli esercizi

Due problemi di Deep-ML mettono in codice esattamente il modello descritto sopra.

8.1 Compute Arithmetic Intensity and Classify Bottleneck

Il problema chiede di classificare un singolo kernel generico. È la formalizzazione minima del roofline: quattro numeri in ingresso, cinque in uscita.

def compute_arithmetic_intensity(flops: float, bytes_accessed: float,
                                 peak_performance: float, peak_bandwidth: float) -> dict:
    arithmetic_intensity = flops / bytes_accessed
    ridge_point = peak_performance / peak_bandwidth
    achieved_performance = min(arithmetic_intensity * peak_bandwidth, peak_performance)
    utilization_percent = achieved_performance / peak_performance * 100

    return {
        'arithmetic_intensity': round(arithmetic_intensity, 4),
        'ridge_point': round(ridge_point, 4),
        'bottleneck': 'compute-bound' if arithmetic_intensity >= ridge_point else 'memory-bound',
        'achieved_performance': round(achieved_performance, 4),
        'utilization_percent': round(utilization_percent, 4),
    }

Tre punti di attenzione.

Il primo è la distinzione fra grandezze del carico e grandezze della macchina. flops e bytes_accessed descrivono il kernel, peak_performance e peak_bandwidth descrivono l'hardware. L'arithmetic intensity mette in rapporto due grandezze del carico, il ridge point due grandezze della macchina. Non si mescolano mai.

Il secondo è la riga con min. La performance ottenibile non è sempre AI × peak_bandwidth: quella formula descrive il tratto diagonale del roofline, valido solo sotto il ridge point. Sopra il ridge il tetto è piatto e il valore è peak_performance. Prendere il minore dei due riproduce entrambi i tratti in una riga sola. Un controllo utile: i due argomenti di min devono avere la stessa unità di misura, cioè FLOP/s. Se uno dei due è una quantità e non una velocità, è sbagliato.

Il terzo è il confronto >= invece di >, perché il problema specifica che la parità esatta va classificata come compute-bound.

8.2 Classify LLM Phases

Il (secondo problema)[https://www.deep-ml.com/problems/417?ip=inference-engineer-general__inference-engineer] applica lo stesso schema alle due fasi dell'inferenza, calcolando FLOP e byte a partire dai parametri del modello.

def classify_llm_phases(num_params: int, sequence_length: int, batch_size: int,
                        bytes_per_param: int, peak_flops: float,
                        peak_bandwidth: float) -> dict:

    ridge_point = peak_flops / peak_bandwidth

    # --- PREFILL ---
    prefill_flops = 2 * num_params * sequence_length
    prefill_memory = num_params * bytes_per_param
    prefill_ai = prefill_flops / prefill_memory
    prefill_bottleneck = 'compute-bound' if prefill_ai >= ridge_point else 'memory-bound'
    prefill_achieved = min(prefill_ai * peak_bandwidth, peak_flops)
    prefill_util = prefill_achieved / peak_flops * 100

    # --- DECODE ---
    decode_flops = 2 * num_params * batch_size
    decode_memory = num_params * bytes_per_param
    decode_ai = decode_flops / decode_memory
    decode_bottleneck = 'compute-bound' if decode_ai >= ridge_point else 'memory-bound'
    decode_achieved = min(decode_ai * peak_bandwidth, peak_flops)
    decode_util = decode_achieved / peak_flops * 100

    return {
        'ridge_point': round(ridge_point, 4),
        'prefill': {
            'total_flops': prefill_flops,
            'memory_bytes': prefill_memory,
            'arithmetic_intensity': round(prefill_ai, 4),
            'bottleneck': prefill_bottleneck,
            'achieved_flops': round(prefill_achieved, 4),
            'utilization_percent': round(prefill_util, 4),
        },
        'decode': {
            'total_flops': decode_flops,
            'memory_bytes': decode_memory,
            'arithmetic_intensity': round(decode_ai, 4),
            'bottleneck': decode_bottleneck,
            'achieved_flops': round(decode_achieved, 4),
            'utilization_percent': round(decode_util, 4),
        },
    }

La formula dei FLOP è la stessa nelle due fasi, 2 × parametri × token processati: cambia solo cosa conta come token processati. Nel prefill è sequence_length, cioè la lunghezza del prompt; nel decode è batch_size, cioè quante sequenze avanzano di un token in quello step. I byte invece coincidono, perché in entrambe le fasi i pesi si leggono una volta.

Vale la pena notare cosa questo esercizio non modella: la KV cache non compare nel calcolo della memoria. Il denominatore è solo num_params × bytes_per_param. Di conseguenza l'AI del decode qui vale esattamente batch_size, che è la versione ingenua discussa all'inizio del punto 5, prima della correzione. In un sistema reale il termine B·KV_seq va aggiunto, ed è quello che impone il tetto asintotico.

Casi di test che coprono i rami significativi:


9. In sintesi

prefill decode
operazione matrice per matrice matrice per vettore
token per step N (tutto il prompt) B (uno per sequenza)
byte letti da HBM 2P 2P + B·KV_seq
arithmetic intensity ≈ N ≈ B, con tetto 2P/KV_seq
regime su H100 compute-bound sopra ~300 token memory-bandwidth-bound
tensor core saturi non invocati
risorsa da comprare FLOP di picco banda e capacità HBM

La domanda di partenza chiedeva perché il prefill sia compute-bound e il decode memory-bandwidth-bound. La risposta breve è che entrambe le fasi leggono gli stessi identici byte di pesi da HBM, ma il prefill li ammortizza su N token mentre il decode li spende per uno solo. Tutto il resto, dal ridge point ai tensor core inattivi fino alla disaggregazione dell'hardware, è conseguenza di quel rapporto.