Chunked Prefill

El prellenado por bloques es una técnica de serving de modelos que divide un mensaje de entrada largo en grupos más pequeños de tokens y los procesa en varios pasos, intercalándolos con la generación de tokens de otras solicitudes. Sin esta técnica, un motor de inferencia que admite un mensaje largo dedica un paso entero a ese mensaje, y todas las respuestas que ya se están transmitiendo a otros usuarios se pausan hasta que termina. El artículo de 2023 sobre SARATHI introdujo el prellenado por bloques, y su continuación, Sarathi-Serve, creó una planificación «sin pausas» basada en él. Dynamic SplitFuse, de DeepSpeed-FastGen, aplica la misma idea. El prellenado por bloques viene ahora activado de forma predeterminada en vLLM y está disponible en SGLang y TensorRT-LLM, los motores que sirven muchos modelos de lenguaje grandes (LLM) de pesos abiertos.

Prellenado, decodificación y pausas en la generación#

La inferencia autorregresiva de Transformer tiene dos fases:

• Prellenado: El modelo procesa en paralelo todos los tokens del mensaje, calcula sus claves y valores y los almacena en la caché KV. El prellenado realiza muchas operaciones por byte de pesos leído, así que normalmente está limitado por la capacidad de cálculo.
• Decodificación: El modelo genera un token por paso para cada solicitud y lee toda la caché cada vez. La decodificación realiza pocas operaciones aritméticas por byte leído, así que normalmente está limitada por el ancho de banda de memoria.

Un paso de GPU que incluye el prellenado de 20 000 tokens dura mucho más que un paso de decodificación, así que todas las solicitudes del lote tienen que esperar. Sarathi-Serve llama a estas pausas interrupciones de la generación, y los usuarios cuyos propios mensajes son breves las perciben como picos de latencia entre tokens.

Cómo funciona el prellenado por bloques#

El planificador asigna a cada paso un presupuesto de tokens, que en vLLM se establece con max_num_batched_tokens. Primero añade un token por cada decodificación en curso y, después, completa el presupuesto restante con bloques de los mensajes pendientes. Si un mensaje no cabe, continúa en el paso siguiente, ampliando su caché KV bloque a bloque hasta procesar todo el mensaje y empezar a decodificarlo. Mezclar las dos fases en un mismo paso también permite aprovechar mejor la GPU, ya que la decodificación deja capacidad de cálculo sin utilizar que pueden ocupar los bloques de prellenado:

La fragmentación por bloques solo cambia cuándo se procesan los tokens del mensaje. Cada bloque presta atención a las claves y los valores almacenados en caché de todos los bloques anteriores y a sus propios tokens, mediante una máscara causal; por eso, el resultado es el mismo que con un único prellenado, como muestra una capa de atención en PyTorch:

import torch
import torch.nn.functional as F

torch.manual_seed(0)
prompt_len, chunk = 1000, 256
q, k, v = torch.randn(3, 1, 8, prompt_len, 64).unbind(0)  # one layer: 8 heads, 64 dims each

# One-shot prefill: all prompt tokens attend causally in a single pass
full = F.scaled_dot_product_attention(q, k, v, is_causal=True)

# Chunked prefill: each chunk attends to the cached keys and values of earlier chunks plus itself
outputs = []
for start in range(0, prompt_len, chunk):
    end = min(start + chunk, prompt_len)
    mask = torch.ones(end - start, end, dtype=torch.bool).tril(diagonal=start)  # causal over absolute positions
    q_chunk, k_cached, v_cached = q[..., start:end, :], k[..., :end, :], v[..., :end, :]
    outputs.append(F.scaled_dot_product_attention(q_chunk, k_cached, v_cached, attn_mask=mask))
chunked = torch.cat(outputs, dim=2)

print(f"{len(outputs)} chunks, same result as one pass: {torch.allclose(full, chunked, atol=1e-5)}")

Muestra 4 chunks, same result as one pass: True: el mensaje de 1000 tokens se procesa en bloques de 256, 256, 256 y 232 tokens, y las salidas coinciden con las de un único paso.

Técnicas de serving relacionadas#

• Agrupación continua: Añade y elimina solicitudes en cada paso en lugar de esperar a que termine un lote entero. El prellenado por bloques permite planificar mensajes largos dentro de ese ciclo.
• PagedAttention: Almacena la caché KV en páginas de tamaño fijo, de modo que un mensaje cuyo prellenado se realiza en varios pasos puede ampliar la caché sin reservar memoria de antemano para toda su longitud.
• Almacenamiento en caché de mensajes: Reutiliza las claves y los valores almacenados en caché de un prefijo de mensaje repetido, lo que elimina por completo el prellenado de esa parte.
• Prellenado y decodificación desagregados: Ejecuta las dos fases en grupos de GPU separados y transfiere entre ellos la caché KV, de modo que cada grupo pueda escalarse y ajustarse a su propio cuello de botella.
• Decodificación especulativa: Acelera la fase de decodificación proponiendo varios tokens y verificándolos en un solo paso.

Aplicaciones#

El prellenado por bloques es especialmente útil cuando un servidor atiende mensajes largos y cortos. Un asistente de generación aumentada por recuperación (RAG) que llena su ventana de contexto con un informe de 50 páginas puede prellenarlo por bloques mientras siguen transmitiéndose los chats de otros usuarios. Los modelos de visión y lenguaje (VLM) añaden cientos o miles de tokens de imagen a cada mensaje, así que las canalizaciones de cámaras que envían fotogramas para describirlos generan precisamente esa combinación. La interfaz de LLM de Ultralytics envía texto e imágenes a cualquier punto de conexión compatible con OpenAI, incluido un servidor vLLM autoalojado, donde las solicitudes de una canalización YOLO26 para describir los fotogramas en los que detecta personas se prellenan por bloques junto con el tráfico de otros usuarios.

Ajuste del presupuesto de tokens#

El presupuesto contrapone dos métricas de latencia. Un presupuesto pequeño deja poco espacio para el prellenado en cada paso, lo que mantiene baja la latencia entre tokens para los usuarios que reciben resultados en streaming y retrasa el tiempo hasta el primer token de los mensajes nuevos; un presupuesto grande produce el efecto contrario. Cada bloque posterior de un mensaje largo también vuelve a leer la caché KV, cada vez mayor, de los bloques anteriores, así que los bloques muy pequeños generan más tráfico de memoria. Los operadores ajustan el presupuesto teniendo en cuenta la latencia de inferencia mediana y de cola, el tiempo hasta el primer token y el rendimiento con longitudes de mensaje realistas. La guía de optimización de vLLM cita 2048 tokens como un presupuesto que favorece la latencia entre tokens y recomienda más de 8192 tokens para lograr el máximo rendimiento con modelos más pequeños en GPU grandes.