Disaggregated prefill e decode per l’inference LLM su SageMaker HyperPod

L’inference di un LLM ha due fasi molto diverse. Il prefill è compute-bound: elabora l’intero prompt in parallelo per generare la KV cache iniziale. Il decode è memory-bound: genera un token alla volta e richiede molta banda di memoria. Quando le due fasi condividono la stessa GPU, un prompt lungo blocca la generazione di token per ogni richiesta concorrente. Il Disaggregated Prefill and Decode (DPD), ora funzionalità dell’HyperPod Inference Operator, rimuove questa interferenza eseguendo ciascuna fase su pool di GPU separati, collegati tramite Elastic Fabric Adapter (EFA) con RDMA.

Quando conviene

Il DPD dà i vantaggi maggiori su workload streaming long-context e ad alta concorrenza: assistenti chat, pipeline agentic, endpoint di analisi documentale e RAG con contesti recuperati ampi. Vale la pena considerarlo quando i prompt superano regolarmente i 4.096 token, ci sono più utenti concorrenti e conta una consegna costante dei token. Sotto la soglia di routing (default 4.096 token) il router invia le richieste direttamente al decoder, evitando il costo di trasferire la KV cache quando non serve; così un singolo endpoint gestisce automaticamente traffico misto.

Architettura

L’implementazione poggia sul router del vLLM Production Stack, con LMCache che fornisce il layer di trasferimento della KV cache su NIXL ed EFA. Il router tokenizza il prompt e decide se seguire il percorso disaggregato; il prefiller calcola la KV cache e la spinge al decoder layer-by-layer, sfruttando anche una L1 CPU cache per i prefissi ricorrenti (system prompt, storia multi-turn, contesti di retrieval); il decoder riserva memoria GPU per i trasferimenti in arrivo e avvia la generazione a trasferimento completato. Lo stack di trasporto (LMCache PD, NIXL, libfabric, EFA) usa GPU-Direct RDMA: su ml.p5.48xlarge con 3.200 Gbps di EFA, un trasferimento di 8.000 token per Llama 3.3 70B richiede pochi millisecondi.

Prestazioni e scaling

Nei benchmark (4.096 token in input, 256 in output, concorrenza 8-32) il DPD mantiene piatta la latenza per token sotto carico: il miglioramento del TPOT va dal 22% al 66% su H100 e dal 28% al 48% su H200, con throughput in output fino a +35% su H100 e +64% su H200 e latenza end-to-end P50 migliore del 14-32% su H100. In cambio, il DPD introduce un modesto aumento del time to first token per il trasferimento della KV cache. Al momento il supporto è per un singolo decoder con più prefiller: si parte da un rapporto prefill-decode 1:1, salendo a 2:1 o 3:1 per workload prefill-heavy come summarization e RAG.

• Richiede almeno un nodo prefill e uno decode con EFA RDMA (famiglie P5 e P6); serve HyperPod Inference Operator 3.2 o superiore.
• Si attiva aggiungendo pochi campi (pdSpec) allo stesso InferenceEndpointConfig degli endpoint non disaggregati.
• Le strategie di routing (kvaware, session) massimizzano la cache locality tra i prefiller.

Fonte: Disaggregated prefill and decode for LLM inference on SageMaker HyperPod — https://aws.amazon.com/blogs/machine-learning/disaggregated-prefill-and-decode-for-llm-inference-on-sagemaker-hyperpod/