Backends rápidos

Portabilidad de TileLang

Los cinco kernels de laya/tl_kernels.py pueden ejecutarse en el objetivo CPU (c) de TileLang mediante compile_cpu. Es una especialización escalar de fp32 explícita, no un backend de aceleración por CPU para Agent.accelerate. Requiere TileLang y un compilador de C++ local. No introduce ninguna dependencia de servicio en tiempo de ejecución.

import torch
from laya import tl_kernels as K

A = torch.randn(17, 67)
W = torch.randn(70, 67)
b = torch.randn(70)
C = torch.empty(17, 70)
kernel = K.compile_cpu(K.gemm_kernel, 70, 67, bias=True, act="gelu")
kernel(A, W, b, C)

compile_cpu(factory, *args, **kwargs) acepta las mismas opciones de forma y de operación que las cinco factorías de GPU. Selecciona cpu=True, dtype="float32", el objetivo c y desactiva la vectorización. Todas las entradas y salidas de coma flotante deben ser tensores float32 de CPU contiguos; las longitudes de atención siguen siendo int32. Convierte las activaciones de 16 bits con .cpu().float().contiguous() antes de llamarlo. LayerNorm sigue actualizando su tensor residual in situ, y RoPE sigue actualizando las columnas Q/K in situ. Conserva el kernel compilado para reutilizar sus dimensiones dinámicas.

La especialización de CPU sustituye las asignaciones de fragmentos y de shared por búferes locales, usa bucles T.grid en serie y omite la anotación swizzle de GPU de la atención. Los GEMM usan la implementación escalar de CPU de TileLang y las reducciones usan búferes locales. Los algoritmos por bloques originales, los predicados de padding, las máscaras de atención y el softmax en línea siguen compartiéndose con la implementación de GPU. Los intermedios de CPU se quedan en fp32, incluidas las probabilidades de atención; los intermedios de GPU conservan su dtype original. No se afirma ninguna aceleración por CPU ni un backend de CPU para el modelo completo.

Nueva comprobación en TileLang 0.1.14

Observado en Linux, Python 3.12.13, torch 2.11.0+cu130, TileLang 0.1.14 y una RTX 4070 Ti SUPER. La preocupación anterior de portabilidad se refiere a compilar la especialización de GPU sin cambios, no a la posibilidad de generar código para la CPU. En el commit base fa9a2a7, esta página de documentación no estaba en el checkout.

La prueba compila factory.get_tir(...) con target="c" y la configuración de la pasada FAST de la GPU. Estos son los textos de diagnóstico exactos (se omiten ubicaciones de código y trazas de pila). Se comprobaron tanto bf16 como fp16 para cada kernel.

Los cuatro primeros fallos de fragmento/reducción son tvm.error.InternalError. El fallo de generación de código bf16 también es un InternalError. RoPE en FP16 lanza RuntimeError: Compilation Failed! seguido de la invocación del compilador y el código fuente; el diagnóstico anterior se emite en stderr. Sus conversiones vectoriales generadas también fallan al convertir vectores float de vuelta a vectores half.

Para aislar el soporte de dtype del soporte de fragmentos, las pruebas compilan cada kernel de nuevo con cpu=True (búferes locales y bucles en serie), conservan bf16 y desactivan la vectorización. Los cinco fallan entonces exactamente con:

Cannot convert type bfloat16 to C type

Por tanto, los fragmentos son el primer obstáculo para GEMM, GEGLU y la atención, las reducciones de fragmentos para LayerNorm, y bf16 es de forma independiente un obstáculo para los cinco. RoPE no tiene asignación de fragmentos ni reducción; GEMM y GEGLU no tienen ninguna operación T.reduce_* explícita. Las reducciones de la atención quedan inicialmente enmascaradas por su fallo de diseño. Sondas separadas de fragmentos solo fp32 reproducen el error de relleno y el error de reducción tanto para reduce_sum como para reduce_max, sin ningún GEMM ni bf16. Reemplazar solo los scopes también produjo este error semántico para el búfer de GEMM (los otros búferes afectados eran Ci, x y s):

[Tilelang Semantic Check] Local buffer `C_l` is indexed by T.Parallel loop variable `i`. Local buffers are thread-private and do not participate in parallel layout inference. Use T.serial/T.vectorized/T.unroll for per-thread local indexing, or T.alloc_fragment when the indexed dimension should be distributed across threads.

La especialización de fp32 en serie elimina estos obstáculos. Las aserciones de diagnóstico están fijadas a la versión 0.1.14 y se omiten en otra versión, donde los errores deberían volver a comprobarse. Las pruebas numéricas de CPU siguen ejecutándose en otras versiones.

Comprobaciones numéricas

Ejecuta python -m pytest tests/test_fast_cpu.py -q -s. En el entorno anterior: 58 superadas. Las comprobaciones solo de CPU no requieren CUDA; solo las comparaciones con GPU se omiten sin ella. La cobertura incluye teselas GEMM M/N/K desiguales, todos los epílogos de GEMM, GEGLU, todas las combinaciones de residual/bias, valores residuales grandes, wraparound de posición de RoPE y columnas V intactas, formas de atención estáticas/dinámicas, ventanas deslizantes, teselas parciales, longitudes desiguales, secuencias vacías y salidas de padding finitas.

La semilla es 1234. Las comparaciones con GPU usan valores de entrada idénticos redondeados a bf16 o fp16 y luego promovidos a fp32 para la ejecución en CPU. Son tolerancias absolutas para los fixtures acotados, no una garantía para magnitudes arbitrarias o profundidad de modelo. Las comparaciones de atención usan filas de query válidas, como en tests/test_fast.py.

La tolerancia CPU/referencia es 2e-5 (2e-6 para RoPE). Las actualizaciones del flujo residual son exactamente iguales, incluido el caso sin residual.

Evidencia de preservación de la GPU

La palabra clave cpu toma el valor false por defecto. Los valores por defecto existentes de bf16/fp16, los scopes de asignación de GPU, los bucles paralelos, los swizzles y las opciones FAST no cambian. Las llamadas a GPU en fp32 por defecto siguen lanzando ValueError.

Las ejecuciones antes/después usaron el módulo original extraído con git show fa9a2a7:laya/tl_kernels.py y el módulo modificado, respectivamente. El original se cargó como laya.tl_kernels mediante importlib para las ejecuciones de referencia; el resto del código de Laya y el entorno de Python permanecieron iguales.

• python -m pytest tests/test_fast.py -q, con las pruebas full-forward configuradas para cargar el checkpoint inglés en caché que se identifica más abajo: 13 superadas antes; 13 superadas después. Al principio, sin checkpoint, eran 11 superadas / 2 omitidas. Ambas ejecuciones completas emitieron la advertencia existente de limitación de temperatura del checkpoint (choice:11+=0.10058280825614929 -> 0.5).
• Captura con semilla de las pruebas de kernel existentes: 24 tensores de salida idénticos bit a bit, incluidas las actualizaciones residuales; diferencia máxima antes/después 0.
• Código fuente CUDA generado para las cinco formas de sonda anteriores, en ambos dtypes: 10/10 idénticos byte a byte. Esto también cubre RoPE, ausente de la suite fast original.
• python benchmarks/parity_fast.py --model "$MODEL" --dtype bf16 --json ... y el comando fp16 equivalente: 288 preguntas sobre 60 estados por dtype. Todos los registros JSON antes/después (probabilidades fp32, stock y fast) comparan exactamente iguales; diferencia máxima de probabilidad antes/después 0.

MODEL era el snapshot inglés en caché convaiinnovations/laya 55cf4c4ebb4ebe31b2550e8bdf3bd21b99753851; las ejecuciones usaron HF_HUB_OFFLINE=1. Aquí no se afirma ninguna comparación con un checkpoint multilingüe o de decisiones tipadas.

Comprobaciones del repositorio

Todos los comandos usaron /home/ckl/projects/S/laya/.venv/bin/python; ruff y zensical provenían del directorio bin de ese entorno virtual.

La nueva suite está registrada como necesitada del extra opcional de TileLang y un compilador de C++ en las exenciones existentes de la prueba de empaquetado. La prueba de contrato de API fija el selector de CPU aditivo solo por palabra clave, el dtype de GPU por defecto sin cambios y la firma de compile_cpu, sin importar TileLang en el entorno base de CI.

AOTInductor

DecisionModel se puede exportar, compilar en un paquete .pt2 y cargar con torch._inductor.aoti_load_package. El paquete devuelve tanto los logits de decisión como los logits de acción. La tokenización, el padding, la calibración de temperatura y el formato de respuestas siguen siendo responsabilidad de quien llama; esto no añade un backend de Agent ni cambia su ruta de ejecución por defecto.

El PR #472 cerrado documentó el obstáculo de dtype anterior. El estado agrupado y las características de confianza de la cabeza de acción se calculan en fp32, incluso cuando los pesos de un modelo son explícitamente bf16. Sin autocast, el primer linear de acción recibía por tanto una entrada fp32 y pesos bf16. Ahora la pasada hacia adelante convierte la entrada concatenada al dtype de los pesos de la cabeza solo fuera de autocast. Softmax, la entropía y los logits de decisión devueltos conservan sus cálculos en fp32. La inferencia eager existente con parámetros fp32, incluida la AMP fp16/bf16, mantiene su numérica; los dtypes mixtos de pesos/autocast también conservan la conversión original de autocast.

Exporta una copia de evaluación separada convertida a bf16, fuera de autocast. Declara el eje de tokens como 16 * Dim("tokens16", ...) para satisfacer las guardas de alineación de la atención. Las filas y los recuentos de marcadores pueden ser dinámicos de forma independiente. No asumas que una exportación capturada bajo autocast se puede empaquetar fuera de ese contexto: usa dtypes de pesos explícitos para esta receta.

Reproducir sin conexión

La comprobación basada en aserciones usa por defecto un ModernBERT diminuto inicializado aleatoriamente, sin acceso a la red. Pasa un directorio de checkpoint local para una medición con un modelo real. La falta de CUDA, de las API de AOTInductor o de un compilador de C++ produce un SKIP explícito; los fallos de compilación y de paridad en una instalación compatible hacen fallar la comprobación.

python scripts/check_aoti.py --output-dir /tmp/laya-aoti-smoke
HF_HUB_OFFLINE=1 TORCHINDUCTOR_CACHE_DIR=/tmp/laya-aoti/cache \
  python scripts/check_aoti.py --model /path/to/local/multilingual \
  --output-dir /tmp/laya-aoti

El directorio de salida contiene decision.pt2, results.json, inputs.pt y eager.pt. El JSON incluye instantáneas completas de nvidia-smi, tiempos de exportación/empaquetado/carga, bytes del paquete, latencia y deltas máximos absolutos de logit/probabilidad para ambas cabezas. Los artefactos binarios y las cachés del compilador se mantienen deliberadamente fuera del repositorio.

Medido en una RTX 4070 Ti SUPER

2026-10-03, Linux x86-64, Python 3.12.13, torch 2.11.0+cu130, transformers 5.17.0, controlador NVIDIA 615.71.09, 16.376 MiB de VRAM. Checkpoint: convaiinnovations/laya, subdirectorio multilingual en la revisión 1c5edc17a7acd8701df6fc341c0d179f1c62c982. La referencia es main fa9a2a7. Las instantáneas completas de la máquina y las mediciones sin redondear están en aoti_multilingual_rtx4070.json.

El fallo reproducido es en el empaquetado, tras una exportación correcta: mat1 and mat2 must have the same dtype, but got Float and BFloat16. Cada ejecución usó una caché de Inductor separada e inicialmente vacía; la compilación de AMP se ejecutó antes del empaquetado, así que el tiempo del paquete no es una medición de arranque en frío de un intérprete nuevo. La comprobación en un proceso nuevo cargó solo el paquete y las entradas guardadas (sin checkpoint), con torch.compile y los puntos de entrada del empaquetado bloqueados. Reprodujo las respuestas de ambas cabezas. PyTorch aun así compiló una pequeña sonda de capacidad AVX de CPU en la caché inicialmente vacía; no se compiló ningún grafo de modelo ni kernel CUDA. Por tanto, una afirmación rotunda de «ninguna actividad del compilador al cargar» sería inexacta en este entorno de ejecución.

Las mismas entradas guardadas contienen siete filas de decisión en tres lotes, con texto real tokenizado de facturación/reembolso, preguntas choice/score/noul, padding y recuentos de marcadores válidos desiguales. Cada latencia es la mediana de cinco grupos de 30 pasadas hacia adelante tras diez calentamientos, con cronometraje de reloj de pared sincronizado. torch.compile(dynamic=True) usa el modo por defecto. En estas cifras de latencia no se incluyen grafos CUDA explícitos, tokenización, exportación ni compilación.

Aquí AMP significa parámetros fp32 con autocast de bf16. bf16 explícito significa que los pesos y el flujo residual son bf16, sin autocast; usa esas columnas para la comparación de ejecución más cercana al paquete. bf16 explícito eager/compiled y AOTI no estaban disponibles antes de esta corrección.

Tanto el argmax de decisión como el de acción coinciden en 7/7 filas para ambas comparaciones de AOTI. Las probabilidades son salidas softmax sin procesar, sin calibración. La distribución de acción está saturada en estas entradas, así que su delta de probabilidad cero no implica logits subyacentes idénticos frente a AMP. La pasada hacia adelante compilada con bf16 explícito también difiere del bf16 explícito eager (delta máximo de logit de decisión 0.5, delta de probabilidad 0.00659859). La compilación en precisión reducida no es exacta bit a bit.

La GPU se compartió con aplicaciones de escritorio y otros trabajos de Python; la compilación en CPU también se compartió, y los relojes no estaban fijados. Estos son tiempos observados, no una afirmación de aceleración aislada. Instantáneas de nvidia-smi que enmarcan las ejecuciones:

Límites restantes

• El artefacto es específico de este checkpoint, esta precisión, esta pila de PyTorch/ejecución y este objetivo de GPU; no se probó la portabilidad a otro hardware o versiones de PyTorch.
• Esta comprobación exporta filas 1–8, tokens 32–512 en múltiplos de 16 y slots de marcador 2–8. Ejercita tres formas, incluidas formas distintas del ejemplo de exportación, en lugar de cada punto de esos rangos. Una opción válida puede usar slots de marcador con padding; un tensor real de un solo slot toma la rama separada de opción única y necesita una exportación separada. Las longitudes de token arbitrarias y una capa de producción de bucketing/dispatch quedan fuera de este cambio.
• Siete filas verifican la regresión del empaquetado, no la precisión amplia del checkpoint ni la paridad de confianza calibrada. Convertir una copia de exportación a bf16 difiere de conservar los pesos fp32 bajo autocast; la ruta eager existente en sí permanece sin cambios.
• El script requiere una compilación de PyTorch con capacidad CUDA y una cadena de herramientas de compilador local para crear el artefacto. El .pt2 incrusta el modelo y los kernels CUDA; no interviene ningún servicio alojado.