Pipelined Decoding para Optimización de Inferencia de LLM en GPU

La inferencia de modelos de lenguaje autoregresivos, donde cada token generado depende del anterior, presenta un desafío fundamental en la utilización eficiente de las unidades de procesamiento gráfico (GPU). A pesar de que las GPU realizan la mayor parte del cómputo intensivo, la coordinación secuencial entre la CPU (para tareas de "housekeeping" como planificación, selección de tokens y gestión de metadatos) y la GPU (para la ejecución del modelo) introduce latencias significativas. Estas latencias, conocidas como 'GPU bubbles', resultan en periodos de inactividad de la GPU, reduciendo el throughput efectivo.

El problema se agrava porque el trabajo de GPU por token es relativamente pequeño, mientras que el overhead de la CPU es un costo fijo por cada ciclo de decodificación. La tesis central de Photon es que al superponer estas tareas de CPU y GPU a través de una decodificación pipelined, se puede eliminar o reducir drásticamente estos 'bubbles', permitiendo que la GPU opere de manera casi continua. Esto es particularmente relevante en un contexto donde la velocidad de las GPU y la complejidad de los modelos continúan aumentando, haciendo que el cuello de botella se desplace hacia la eficiencia de la orquestación entre CPU y GPU.

Arquitectura del Sistema

Photon implementa la decodificación pipelined mediante tres mecanismos clave. Primero, utiliza 'ping-pong slots', que son conjuntos duplicados de buffers (input, output, KV cache) en memoria GPU y host (page-locked) para evitar colisiones. Mientras un slot está siendo procesado por la GPU, el otro puede ser preparado o sus resultados leídos por la CPU. Esto permite que la GPU inicie el cómputo del siguiente token mientras la CPU aún procesa el token actual. Los kernels se encolan en un CUDA stream compartido para el cómputo, pero las copias device-to-host de los tokens sampled se realizan en copy streams separados, ancladas a CUDA events para no bloquear el cómputo principal.

Segundo, se desacopla la fase de forward (cómputo del modelo) de la fase de sampling (selección del token). El forward para el token t+1 puede ejecutarse antes de que el token t haya sido completamente committed por la CPU. La sampling para t+1, que puede depender del token t (especialmente en constrained decoding donde una máscara de tokens permitidos se genera dinálogos al contexto), se retrasa hasta que el commit de t se completa. Este orden commit-before-finalize asegura la corrección sin introducir un bloqueo en la ruta crítica del forward.

Tercero, se gestiona la finalización de solicitudes (requests) con un mecanismo de 'zombies' y refcounting. Cuando una secuencia alcanza un token de parada (EOS) o su límite de longitud, se marca como finalized, pero no se elimina inmediatamente si aún está referenciada por un forward en vuelo (indicado por inflight_refs). El commit de un token zombie se salta, evitando mutaciones de estado innecesarias. La liberación de recursos (páginas de KV cache, slots LoRA) ocurre solo cuando inflight_refs llega a cero, reemplazando lógica compleja de cancelación con un esquema simple de reference counting. Este pipeline unificado también maneja prefill (procesamiento inicial del prompt) y decode (generación token a token) de manera indistinta, maximizando la superposición de trabajo CPU/GPU en ambos escenarios.

Fundamentos Teóricos

El problema de la coordinación entre unidades de procesamiento heterogéneas y la optimización del throughput mediante pipelining es un tema recurrente en la arquitectura de computadoras y sistemas distribuidos. Conceptos como el Instruction Level Parallelism (ILP) y el pipelining de CPUs, popularizados en la década de 1980 con arquitecturas RISC, buscan mantener las unidades de ejecución ocupadas solapando fases de instrucción. En el contexto de sistemas distribuidos, el principio de batching y pipelining para amortizar costos fijos de comunicación o coordinación es fundamental, similar a cómo se aplica en message queues o RPC frameworks para reducir la latencia efectiva. El uso de CUDA streams y events para la gestión asíncrona de tareas en la GPU se alinea con los principios de asynchronous I/O y non-blocking operations para evitar el bloqueo de hilos de ejecución, un patrón bien establecido en la programación concurrente. La gestión de recursos con reference counting para evitar condiciones de carrera y asegurar la liberación segura de memoria es un patrón clásico en la gestión de memoria de sistemas operativos y lenguajes de programación, como se describe en trabajos fundamentales sobre garbage collection o smart pointers.