NVIDIA y Princeton presentan ReplaySSM para acelerar la decodificación de modelos híbridos

Investigadores de NVIDIA y del Dao AI Lab de la Universidad de Princeton presentaron ReplaySSM, un método para modelos de lenguaje híbridos que almacena temporalmente entradas recientes de modelos de espacio de estados (SSM) en vez de escribir todo el estado en memoria en cada paso de token. Representa el estado mediante un punto de control y un búfer corto de entradas, separando el cálculo de salida de las actualizaciones de estado para simplificar la reversión de la decodificación especulativa y limitar la acumulación de errores en lecturas de baja precisión.

En pruebas con cuatro modelos Nemotron 3 y Qwen3.5 con lote 256, el equipo comunicó aceleraciones de 1,46 a 1,80 veces en los kernels de decodificación y de 1,20 a 1,47 veces de extremo a extremo; las mismas pruebas no mostraron mejoras con lote 1. En dos modelos de clase 120B, el rendimiento de generación especulativa aumentó 1,85 y 2,10 veces con lote 512. Son resultados del equipo bajo sus propias condiciones de hardware y software; el método se ha integrado en TensorRT-LLM, vLLM y FlashInfer.