Resumen Técnico: XGrammar-2

Declaración del Problema

Los agentes de Modelos de Lenguaje Grandes (LLM) modernos dependen cada vez más de la generación estructurada dinámica, particularmente para la llamada a herramientas y los protocolos de respuesta. A diferencia de las tareas tradicionales de generación estructurada donde el esquema de salida es estático y conocido de antemano, las cargas de trabajo agénticas exhiben una dinamismo estructural significativo:

1. Dinamismo Inter-solicitud: Cada solicitud puede exponer un conjunto diferente de herramientas y esquemas (a menudo con control de acceso por herramienta), lo que conduce a un espacio combinatorio enorme de gramáticas de salida posibles. Los motores existentes que almacenan en caché gramáticas completas a nivel de solicitud se vuelven ineficaces, forzando un preprocesamiento costoso por solicitud y aumentando el Tiempo hasta el Primer Token (TTFT).
2. Dinamismo Intra-solicitud: Dentro de una sola solicitud, el modelo debe cambiar entre texto libre y subgramáticas estructuradas basándose en desencadenantes (por ejemplo, un nombre de herramienta específico o una etiqueta de canal). Expresar estos cambios condicionales utilizando la Forma Normal de Backus-Naur (BNF) estándar o la BNF Extendida (EBNF) es engorroso, ineficiente y escala mal con el número de herramientas.

Los métodos existentes de decodificación restringida (por ejemplo, XGrammar, Outlines, llguidance) luchan con estas dinámicas. O bien requieren una compilación completa de la gramática de antemano (prohibitiva para conjuntos de herramientas dinámicos) o carecen de mecanismos eficientes para manejar la intercalación de texto libre y restricciones estructuradas dentro de un solo flujo de generación.

Metodología

XGrammar-2 es un motor de generación estructurada diseñado específicamente para cargas de trabajo agénticas dinámicas. Su arquitectura se basa en dos principios de diseño centrales: soporte de primera clase para el cambio de estructura activado por etiquetas y reutilización granular entre solicitudes con diferentes estructuras de salida.

1. TagDispatch: Semántica de Despacho Dinámico

Para abordar el dinamismo intra-solicitud, XGrammar-2 introduce TagDispatch, un intrínseco de gramática que permite un cambio eficiente entre texto libre y subgramáticas estructuradas.

• Mecanismo: TagDispatch se parametriza mediante una lista de pares etiqueta-gramática $(t_i, G_i)$ y un conjunto de cadenas de parada.
• Operación: El motor opera en dos modos:
  • Modo de Despacho: El motor acepta texto ordinario mientras utiliza un autómata Aho–Corasick (AC) para hacer coincidir incrementalmente las etiquetas registradas contra el texto generado.
  • Modo Despachado: Al coincidir una etiqueta, el motor cambia a la subgramática correspondiente $G_i$. Una vez que $G_i$ se completa, el control regresa al modo de despacho.
• Ventaja: Esto evita la necesidad de codificar lógica condicional compleja en reglas EBNF masivas, reduciendo significativamente el tamaño de la gramática y la complejidad de la compilación.

2. Caché de Gramática Cruzada

Para abordar el dinamismo inter-solicitud, XGrammar-2 implementa una Caché de Gramática Cruzada para reutilizar cachés de máscaras de tokens entre diferentes gramáticas que comparten subestructuras comunes.

• Desafío: Determinar si dos Máquinas de Estados Finitos (FSM) son equivalentes es difícil cuando contienen referencias a reglas y ciclos. Además, las cachés de máscaras de tokens dependen no solo de la estructura FSM, sino también de cómo se hace referencia a esa FSM (expansión de contexto).
• Solución: El sistema utiliza un algoritmo de hashing jerárquico:
  • Construye un grafo de referencia FSM.
  • Para porciones acíclicas, calcula hashes estructurales de abajo hacia arriba, incorporando los hashes de las FSM referenciadas.
  • Para referencias cíclicas, utiliza hashes provisionales seguidos de un refinamiento de hash de ciclo para garantizar identificadores estructurales únicos.
• Estrategia de Reutilización: Cuando se encuentra un nuevo estado de gramática, el sistema verifica el grupo de caché global. Si se encuentra una coincidencia perfecta (misma FSM y mirada hacia adelante), la caché se reutiliza. Si solo existe una coincidencia parcial (misma FSM, diferente mirada hacia adelante), el sistema vuelve a verificar los tokens inciertos y actualiza la caché.

3. Caché Adaptativa de Máscara de Tokens Basada en Earley

Lejos de los Autómatas de Pila (PDA) utilizados en versiones anteriores, XGrammar-2 utiliza un analizador Earley para su caché adaptativa de máscaras de tokens.

• Razonamiento: Los estados PDA pueden crecer exponencialmente en gramáticas no deterministas. El analizador Earley mantiene un conjunto de estados de análisis parciales (frontera) que es más eficiente para gramáticas libres de contexto complejas.
• Implementación: La caché se construye solo para estados escaneables (donde el siguiente símbolo es un terminal). Los estados no escaneables se expanden mediante operaciones de predicción y completado en tiempo de ejecución. Los tokens se categorizan como aceptados, rechazados o dependientes del contexto, siendo necesarios para estos últimos controles de contexto en tiempo de ejecución.

4. Compilación Justo a Tiempo (JIT)

Para mitigar el alto costo de compilación de gramáticas grandes y dinámicas:

• Estrategia: En lugar de materializar la caché completa de máscaras de tokens de antemano, XGrammar-2 emplea un sistema JIT configurable.
• Ejecución: Durante la fase de prellenado, el sistema estima el costo de generar cachés para estados específicos. Prioriza compilar los $K$ estados más costosos en tiempo durante el prellenado, superponiendo la compilación con el prellenado. Los estados restantes se compilan bajo demanda durante la fase de decodificación, superponiendo la generación de máscaras con la decodificación. Esto tiene como objetivo ocultar la latencia de compilación.

5. Compresión de Estado de Repetición

Para manejar estructuras de repetición comunes en esquemas JSON (por ejemplo, MinLength, MaxItems):

• Problema: La expansión ingenua de reglas de repetición conduce a un aumento lineal en los estados de la gramática en relación con la cantidad de repetición.
• Solución: El motor comprime los estados de repetición. Para rangos de repetición grandes, transforma la estructura en una secuencia de una regla base repetida $t$ veces, seguida de un estado comprimido que representa el rango restante. Esto limita el número de estados a una constante independientemente del límite de repetición, reduciendo significativamente el tamaño de la caché y la incertidumbre.

Resultados Clave

Se realizaron experimentos utilizando conjuntos de datos como BFCL-v3 y JSONSchemaBench, comparando XGrammar-2 con motores de última generación (XGrammar, llguidance, Outlines) en varios modelos (Qwen, serie Llama).

• Velocidad de Compilación: XGrammar-2 logra una compilación más de 6 veces más rápida en comparación con motores de generación estructurada anteriores en configuraciones dinámicas. Por ejemplo, mientras que XGrammar requería >1000 ms para compilar gramáticas dinámicas, XGrammar-2 logró ~10 ms.
• Sobrecarga de Extremo a Extremo: El sistema incurre en una sobrecarga de extremo a extremo cercana a cero en sistemas modernos de servicio LLM. En escenarios de llamada a funciones, la brecha de latencia entre la generación restringida (XGrammar-2) y la generación no restringida fue menor al 6%.
• Rendimiento: XGrammar-2 demostró un rendimiento de tokens superior en comparación con XGrammar a través de diferentes tamaños de lote y escalas de modelos (por ejemplo, una mejora de rendimiento de 2.5× en Llama-3.2-1B con tamaño de lote 16).
• Escalabilidad: La sobrecarga de memoria de la caché de gramática cruzada se mantiene baja y escala con el número total de herramientas únicas utilizadas en lugar del número de solicitudes, evitando la explosión de memoria en entornos dinámicos.
• Ablación: Eliminar JIT aumentó el tiempo de preprocesamiento en 8.1×; eliminar la Caché de Gramática Cruzada aumentó significativamente el tiempo de generación de máscaras; la Compresión de Estado de Repetición redujo el tiempo de preprocesamiento en 99.6× para casos de repetición de cola larga.

Significado y Afirmaciones

El documento afirma que XGrammar-2 proporciona una evolución necesaria para los sistemas de servicio LLM para soportar aplicaciones agénticas. Al desacoplar el manejo de conjuntos de herramientas dinámicos de la compilación de gramáticas estáticas, permite:

1. Fiabilidad: Aplicación estricta de la validez estructural (por ejemplo, Formato de Respuesta OpenAI Harmony, esquemas de llamada a herramientas) sin depender de correcciones a posteriori.
2. Eficiencia: La capacidad de servir cargas de trabajo dinámicas con costos de compilación y tiempo de ejecución comparables a las cargas de trabajo estáticas.
3. Adopción: El motor se ha integrado en marcos de código abierto como SGLang y vLLM, demostrando su viabilidad práctica en sistemas de inferencia de grado de producción.

Los autores posicionan a XGrammar-2 no como un reemplazo para las técnicas de entrenamiento de modelos, sino como un motor de tiempo de inferencia complementario que puede combinarse con estrategias existentes de ajuste fino o prompting para garantizar una salida estructurada robusta en escenarios de agentes complejos.