Optimizaciones de GLM-5.3, Parte 2: PCP, DCP e Hybrid HiSparse en 16 H200

TL;DR: Los contextos agénticos largos llevan a GLM-5.3 contra dos límites en los H200: la memoria de la caché KV y la velocidad de prefill. Abordamos ambos en 16 H200, divididos en un nodo de prefill y un nodo de decodificación:

• El paralelismo de contexto de prefill (PCP) divide cada prompt entre las GPU por token en lugar de por cabeza, lo cual se adapta mucho mejor a la atención dispersa de GLM-5.3 que el paralelismo de tensores.
• El paralelismo de contexto de decodificación (DCP) fragmenta la caché KV del motor de prefill por token, de modo que un contexto de 1M cabe en el nodo de prefill.
• Hybrid HiSparse permite que la decodificación mantenga la KV en memoria del host, de modo que una solicitud de 1M de tokens cabe.
• En AgentX, el prefill PCP8 × DCP8 con decodificación Hybrid HiSparse ofrece entre 1.7 y 3.2 veces la velocidad por usuario de un despliegue en paralelo de datos al mismo throughput, y es la única configuración que probamos que sirve contextos de 1M de tokens.

Servir contextos largos en despliegues con memoria limitada

Las cargas de trabajo agénticas tienen contextos largos que crecen con cada turno, lo que presiona tanto la memoria de la GPU como el prefill. En los H200, los pesos de GLM-5.3 ocupan la mayor parte de la memoria de cada GPU, dejando poco espacio para la caché KV de esos contextos largos. Y como cada turno añade más salidas de herramientas y contenidos de archivos al prefill, bajo carga los motores de prefill pueden convertirse en el cuello de botella de todo el despliegue.

La Parte 1 mostró cómo Hybrid HiSparse permite que contextos largos quepan en un despliegue agregado en un solo nodo de 8× H200. La desagregación de prefill/decodificación requiere más que eso. En el nodo de prefill, usamos PCP para acelerar los prompts largos paralelizando a través de los tokens del prompt. Con MLA, sin embargo, esto deja cada rango PCP con una copia completa de la caché KV, así que lo combinamos con DCP para eliminar esa replicación. En el nodo de decodificación, Hybrid HiSparse permite mayor concurrencia en contextos largos. El nuevo núcleo modular de vLLM (Model Runner V2) hace que componer todo esto sea mucho más fácil.

¿Por qué no TP o DP?

Dos formas estándar de ejecutar el prefill de GLM-5.3 en un nodo de 8× H200 son DP8 y TP8:

• DP8 asigna cada prompt a una sola GPU, de modo que su atención no puede usar las otras siete. Esto limita el tiempo hasta el primer token (TTFT) para prompts largos, y la caché KV del prompt debe caber en la memoria de esa GPU.
• TP8 particiona las cabezas de atención entre las GPU. Debido a que la KV latente de MLA es compartida por las 64 cabezas, cada GPU almacena la caché completa y la lee para servir solo 8 cabezas. El indexador, que puntúa cada token anterior para seleccionar los 2,048 a los que la MLA dispersa atiende, no se particiona en absoluto, por lo que su coste cuadrático se repite en cada GPU.

Ninguno de los dos esquemas distribuye el trabajo de atención de un prompt largo entre las ocho GPU sin mantener también toda su KV en cada GPU que lo sirve. Las dos secciones siguientes abordan cada mitad. PCP divide las consultas de un prompt entre las GPU, y DCP divide su KV.

PCP: dividir los tokens, no las cabezas

En nuestro motor de prefill de 8× H200, PCP divide cada prompt entre las ocho GPU por token: cada GPU calcula las consultas para un octavo del prompt, con las 64 cabezas. Una división ingenua estaría desequilibrada. Bajo atención causal, un token atiende a cada token anterior a él, así que la GPU que contiene el final del prompt haría mucho más trabajo que la que contiene el inicio. En cambio, PCP corta el prompt en 16 fragmentos y los empareja: la GPU toma un fragmento del frente y un fragmento del final, así que cada GPU recibe la misma mezcla de tokens baratos y costosos (la Figura 3 muestra esto como un patrón en V a través de las ocho GPU). Este es el sharding con balanceo de carga de Context Parallelism for Scalable Million-Token Inference. Finalmente, la nueva KV de cada GPU se recopila con all-gather, de modo que cada consulta aún puede atender al contexto completo.

La Figura 1 compara las dos divisiones. PCP corrige ambos problemas. Cada GPU indexa solo sus propios tokens, así que el indexador ya no se replica en cada GPU y cada GPU hace un octavo de él. Y cada fila de KV leída de la memoria ahora sirve a las 64 cabezas en lugar de 8, elevando la MLA dispersa de unos 26 a unos 191 FLOPs por byte movido,¹ cerca del punto donde el H200 se vuelve limitado por cómputo en BF16.²

No replicar el indexador importa sobre todo en los prefills largos. El costo del indexador crece cuadráticamente con la longitud del prompt, mientras que el del MLA disperso crece solo linealmente, porque cada consulta atiende a un conjunto fijo de 2,048 tokens. Contando los FLOPs de un prefill desde cero, el indexador supera al MLA disperso bajo TP8 una vez que el prompt alcanza unos 32K tokens; bajo PCP8, donde ninguna GPU lo replica, el punto de cruce se desplaza hasta unos 256K.

En conjunto, esto hace que PCP sea ideal para el prefill de prompts largos. Pero PCP todavía tiene dos costos. Cada GPU guarda todos los pesos de atención, mientras que los expertos MoE permanecen divididos entre las ocho GPUs con EP8. Y como el nuevo KV se reúne por completo (all-gather) en cada capa, cada GPU termina almacenando la caché KV completa. Esa caché replicada, no el cómputo, es lo que limita la longitud del contexto, y es lo que DCP aborda.

DCP: fragmentar también la caché KV

Una solicitud GLM-5.3 de 1M de tokens necesita 53.4 GiB de KV en FP8,³ pero cada GPU solo tiene espacio para unos 15 GiB. Paralelismo de Contexto en Decodificación (DCP, por sus siglas en inglés) lo resuelve dividiendo la propia caché: los tokens se reparten en round-robin entre las ocho GPUs, de modo que cada GPU almacena uno de cada ocho tokens, y la misma solicitud solo necesita 6.7 GiB (=53.4/8) por GPU. La Figura 2 lo compara con DP8 y TP8.

Nota: PCP divide las consultas; DCP divide el KV. Los nombres son históricos. DCP se creó para la decodificación, donde cada solicitud tiene un único token de consulta, de modo que la caché KV es lo único que vale la pena dividir. PCP se creó para el prefill, donde los muchos tokens de consulta del prompt son el trabajo. Nuestra configuración ejecuta ambos en el motor de prefill.

Como resultado de este sharding, la mayor parte del KV que las consultas de cada GPU necesitan vive ahora en otras GPUs. En la decodificación, DCP envía las consultas al KV: cada GPU atiende sobre su propio shard, y los resultados parciales se fusionan con su log-sum-exp. En este motor de prefill hace lo contrario: en cada capa, cada GPU hace all-gather de los shards en una copia temporal de contexto completo, atiende para sus propias consultas y libera la copia. Solo la caché persistente se divide entre ocho.

La Figura 3 lo reúne todo. Avanza paso a paso por una capa de PCP8 × DCP8 para ver qué KV almacena cada GPU, qué consultas calcula y qué cruza la interconexión. Cambia a DP8, TP8 o PCP8 para comparar. Ten en cuenta que los dos diseños no se alinean: la GPU almacena uno de cada ocho tokens pero calcula las consultas de los chunks y , por eso cada capa necesita el all-gather.

Ese all-gather no es gratuito: en cada capa reenvía todo el contexto en caché, así que su costo crece con la longitud del contexto. La Figura 4 lo aísla ejecutando el mismo chunk de prefill de 2K tokens en un nodo PCP8 con y sin DCP8, con igual contexto en caché. DCP8 añade alrededor de un 6% por paso hasta 64K tokens de contexto, subiendo al 15% a 248K. Ese es el precio de un pool de KV 8 veces mayor. Esta sobrecarga podría reducirse solapando el all-gather de cada capa con el cómputo de la capa anterior, una optimización en la que estamos trabajando actualmente.

Más allá de aproximadamente 250K tokens, PCP8 por sí solo no puede contener la caché KV; ahí DCP no es una sobrecarga sino lo que hace posibles los prompts largos. Juntos, PCP8 × DCP8 le dan al motor de prefill tanto velocidad como espacio.

El despliegue

Lo servimos como un despliegue desagregado de prefill/decodificación (P/D) en dos nodos de 8× H200. Ambos motores usan una caché KV en FP8 y decodificación especulativa MTP con tres draft tokens (MTP3):

• Prefill: un motor PCP8 × DCP8 × EP8, con 512 GiB de offload de KV en CPU para que los turnos posteriores reutilicen su prefijo.
• Decodificación: ocho rangos DP con EP8 y Hybrid HiSparse de Parte 1 (48 GiB de pool en host + 16 GiB de offload por rango), de modo que la concurrencia de decodificación está limitada por la memoria del host y no por la HBM.

Un router con conciencia de prefijos de llm-d envía cada solicitud al motor de prefill, el conector NIXL mueve su caché KV a un rango de decodificación, y el rango de decodificación transmite los tokens de vuelta.

Comparamos el prefill PCP8 × DCP8 contra el prefill DP8 (ocho rangos DP con EP8 y 64 GiB de offload de KV en CPU cada uno, los mismos 512 GiB en total), y Hybrid HiSparse contra el offload de KV simple en CPU en el nodo de decodificación, con los mismos 512 GiB de memoria de host en ambos casos (8 × 48 GiB de pool HiSparse + 8 × 16 GiB de offload para HiSparse). Cada configuración usa las mismas 16 GPUs, la misma memoria de host y la misma versión de vLLM.

Resultados

Ejecutamos el benchmark AgentX de SemiAnalysis: trazas de Claude Code reproducidas con reutilización de prefijos, un servidor nuevo por concurrencia, ejecuciones de una hora, y MTP3 con una longitud de aceptación sintética fijada en 2.99, de modo que cada configuración obtiene la misma aceleración especulativa. El throughput es tok/s lógico por GPU (incluyendo los cacheados): tokens de entrada más salida, contando la entrada con prefijo cacheado, sobre las 16 GPUs. La interactividad es el P90 de tokens de salida E2E-normalizados por usuario por segundo de AgentX, el eje x de la Figura 5.

Hasta c32, el despliegue está mayormente limitado por el prefill, así que el prefiller marca la mayor diferencia. Con decodificación HiSparse en ambas configuraciones, el prefill PCP8 × DCP8 da 1.7–3.2 veces la interactividad de DP8 al mismo throughput y reduce la TTFT mediana de 1.5–5.2 s a 0.8–0.9 s. El decodificador marca poca diferencia con esta carga: en c8 y c16 las dos configuraciones PCP8 × DCP8 están a pocos puntos porcentuales una de la otra, y la decodificación solo con offload va ligeramente por delante en c8.

A partir de c32 el despliegue se vuelve más limitado por la decodificación, y la decodificación HiSparse toma la delantera. En c64 casi duplica el throughput, a 8,340 tok/s/GPU lógicos a 29.6 tok/s/usuario, mientras que todas las demás configuraciones caen a 1–2 tok/s/usuario.

Con un contexto de 1M, el despliegue necesita ambas cosas. Un rango de prefill DP8 no puede contener una sola solicitud de 1M tokens, y tampoco puede un rango de decode solo con offload.

Hacer que todo se componga

Gran parte del trabajo difícil en vLLM consiste en hacer que características como estas se compongan entre sí. Cada una cambia cómo se distribuye, planifica o mueve la caché KV, y cada combinación debe mantenerse correcta y rápida.

Para el P/D, el KV tiene que cambiar de disposición entre los dos motores. El motor de prefill distribuye el KV de cada solicitud por token entre ocho GPUs. Cada rank de decodificación necesita el KV completo, incluido el de la capa MTP, en su propia disposición paginada. El conector NIXL hace este reensamblado entre los dos nodos, y la reutilización de prefijos desde el offload en CPU tiene que seguir funcionando junto con ello.

Cada configuración aquí se ejecuta en Model Runner V2. Esto mantiene el estado específico del modelo fuera de la ruta de ejecución común y ejecuta la decodificación especulativa sin sincronizaciones CPU–GPU. Hybrid HiSparse está construido sobre esto.

Pruébalo

Todo se ejecuta en vLLM v0.32.0 o posterior (la vllm/vllm-openai:v0.32.0 imagen). Los resultados anteriores se midieron en vLLM main en 9e38b4111a. El gist de reproducción fija 2537629eb2. Las banderas clave del motor de prefill:

--prefill-context-parallel-size 8 \
--decode-context-parallel-size 8 \
--dcp-comm-backend ag_rs \
--enable-expert-parallel \
--attention-config '{"sparse_mla_force_mqa":true}'

Sin sparse_mla_force_mqa, los lotes de prefill cuyos prompts caben dentro del top-k de 2,048 tokens toman la ruta densa MHA. Esta bandera los mantiene también en el kernel sparse-MLA (MQA), de modo que cada fila de KV sirve a todas las cabezas.

Agradecimientos

Este trabajo fue realizado por Roberto L. Castro (Red Hat), Simon Veitner (Red Hat), Matthew Bonanni (Red Hat), Lucas Wilkinson (Red Hat) y Nicolò Lucchesi (Mistral). Agradecemos a Maroon Ayoub (Red Hat) y Robert Shaw (Red Hat) por las pruebas de estrés del despliegue, la notificación de errores y la corrección de bugs.

Notas

Notas al pie

1. Por consulta, el sparse MLA lee 2,048 filas de KV seleccionadas de 660 bytes cada una (656 de KV FP8 en el formato fp8_ds_mla de FlashMLA, más un índice de 4 bytes), unos 1.35 MB. Cada cabeza añade 2,176 bytes para su consulta y salida BF16, y hace 2,176 FLOPs por fila (QK sobre 576 dimensiones, PV sobre 512). Así, 8 cabezas hacen 35.7 MFLOPs sobre 1.37 MB, unos 26 FLOPs por byte, y 64 cabezas hacen 285 MFLOPs sobre 1.49 MB, unos 191. ↩

2. Para el H200 SXM: 989 TFLOP/s densos en BF16 (los 1,979 de la hoja de datos asumen una dispersión 2:4) 4.8 TB/s 206 FLOPs por byte (hoja de datos de NVIDIA H200). ↩

3. Por token: 78 capas × 656 B de MLA KV, 21 capas de indexador × 132 B de claves del indexador (128 valores FP8 más una escala de 4 bytes), y 788 B para la capa MTP, o 54,728 B en total. Para = 1,048,576 tokens, eso es 54,728 MiB 53.4 GiB. ↩