Demostración de TAKT: predicción de series temporales en Rust (DLinear y PatchTST sobre ETT)

Dos modelos estándar de predicción a largo plazo, DLinear (LTSF-Linear) y PatchTST, entrenados con los scripts oficiales sobre los conjuntos de datos ETT, traducidos de PyTorch a Rust puro (sin BLAS ni runtime de ML) y compilados dos veces: la compilación original de ese código Rust y la compilación de TAKT del mismo código. Las predicciones de ambas compilaciones son idénticas bit a bit. Aquí encontrará programas listos para ambas compilaciones, el código fuente de la versión original, un script que recrea las ventanas de entrada a partir de los datos oficiales, un script de benchmark y nuestras mediciones.

Qué calculan los programas

Todas las ventanas de la partición de prueba de un conjunto ETT (límites oficiales, features M, StandardScaler ajustado sobre la partición de entrenamiento): de entrada, 336 pasos anteriores x 7 canales; de salida, los 96 pasos siguientes x 7 canales; float32. ETTh1/ETTh2: 2.785 ventanas; ETTm1/ETTm2: 11.425 ventanas. Los pesos del checkpoint del conjunto seleccionado están compilados dentro del programa (--data).

El bucle de prueba oficial de PatchTST descarta el último lote incompleto de 128 ventanas; de ahí la columna «primeras 2.688». La traducción a Rust no es idéntica bit a bit a PyTorch (otro orden de suma; la mayor diferencia absoluta respecto a PyTorch eager es de 2e-6 a 2e-5); la compilación de TAKT es idéntica bit a bit al original en Rust.

Medición (banco de TAKT)

Medido el 3 de octubre de 2026. AMD Threadripper PRO 5975WX (Zen 3), Linux. Un CCX de 8 núcleos reservado, con sus hilos SMT gemelos inactivos; las ejecuciones de un hilo se fijan a un núcleo, y las de ocho hilos, a los 8 núcleos del CCX. Para cada conjunto: un par de calentamiento y luego 9 rondas en orden alterno, mediana. Tiempo: el bucle de predicción, cronometrado por el propio programa (se excluyen la lectura de las ventanas y la escritura de las predicciones, que cuestan lo mismo en ambas compilaciones). Ciclos: ciclos en modo usuario del mismo bucle (perf stat, sumados en todos los hilos). Ambas compilaciones comparten el mismo programa de control, que conserva en el proceso la memoria de heap liberada (mallopt de glibc), como haría un servicio de larga duración, para que los búferes temporales de cada ventana no se conviertan en fallos de página. Se midieron exactamente los archivos de bin/, con tools/run_bench.sh.

DLinear: 7,1–7,3× en tiempo y 7,2–7,4× en ciclos, con uno y con ocho hilos. PatchTST: 2,3× en tiempo y en ciclos. Con ocho hilos, ambas compilaciones gastan en total aproximadamente los mismos ciclos que con uno. El proceso completo, incluidas la lectura de 26–107 MB de ventanas y la escritura de las predicciones, es 5,7–6,0× más rápido para DLinear con un hilo, 3,0–3,1× con ocho y 2,3× para PatchTST. Los detalles, las muestras sin procesar y los tiempos del proceso completo están en measurements.json.

Ejecución

python3 tools/make_windows.py --out data --targets   # descarga los CSV oficiales de ETT (verificados con sha256),
                                                     # escribe data/windows_*_test.f32 (y los valores objetivo)
mkdir -p out
for d in ETTh1 ETTh2 ETTm1 ETTm2; do
  bin/ltsf-dlinear-takt --data $d --windows data/windows_${d}_test.f32 --threads 1 --out out/dlinear_$d.f32
done
for d in ETTh1 ETTh2; do
  bin/ltsf-patchtst-takt --data $d --windows data/windows_${d}_test.f32 --threads 1 --out out/patchtst_$d.f32
done
sha256sum -c expected.sha256                 # ventanas, valores objetivo y predicciones como en nuestra medición
python3 tools/score.py data/targets_ETTh1_test.f32 out/dlinear_ETTh1.f32

tools/run_bench.sh -m dlinear -n 9 -c 2          # ambas compilaciones, intercaladas, verificación byte a byte, núcleo 2
tools/run_bench.sh -m dlinear -n 9 -t 8 -c 2-9   # ocho hilos en los núcleos 2-9
tools/run_bench.sh -m patchtst -n 9 -c 2

--threads T admite cualquier T >= 1; las predicciones no dependen de T. Los archivos de entrada y salida son float32 little-endian sin formato (n x 336 x 7 y n x 96 x 7); tools/make_windows.py solo necesita numpy.

Linux x86-64 (glibc). Los programas están compilados para x86-64-v3 (AVX2, FMA, BMI2: Intel Haswell y posteriores, AMD Zen y posteriores); no se usan instrucciones FMA. PatchTST llama a erff de la biblioteca C del sistema: con una glibc cuya erff difiera de la de glibc 2.39 (Ubuntu 24.04), sus predicciones pueden diferir de expected.sha256, mientras que el original y la compilación de TAKT siguen coincidiendo byte a byte.

Equivalencia

• En cada ronda medida, las predicciones de las dos compilaciones fueron idénticas byte a byte (recuentos en la tabla).
• Ambas compilaciones reproducen, byte a byte, las predicciones guardadas en nuestra ejecución anterior: DLinear en cuatro conjuntos y PatchTST en dos, particiones de prueba y de validación, 1, 8 y 32 hilos: 72 de 72 archivos.
• tools/make_windows.py reproduce byte a byte las ventanas de los cargadores de datos oficiales de LTSF-Linear / PatchTST (los cuatro conjuntos, prueba y validación; para ello se reimplementó el analizador de números de coma flotante predeterminado de pandas).

Recompilar el original

cd source
RUSTFLAGS="-C target-cpu=x86-64-v3 --remap-path-prefix=$HOME/.rustup=. \
  --remap-path-prefix=$HOME/.cargo/registry/src/index.crates.io-1949cf8c6b5b557f=." \
  cargo +1.96.0 build --release --locked
strip --strip-all -o ltsf-dlinear-original target/release/ltsf-dlinear

Con Rust 1.96.0 de rustup (con el componente rust-src), esto reproduce bit a bit bin/ltsf-dlinear-original y bin/ltsf-patchtst-original en nuestra máquina; con otra configuración cambian las rutas incrustadas, pero no las predicciones.

Comparación anterior con PyTorch (no reproducible con este paquete)

El 2 de octubre de 2026 medimos los mismos modelos en la misma máquina en otras condiciones: sin reserva del banco (39 CPU lógicas compartidas con otras tareas); ambas versiones en Rust compiladas como bibliotecas compartidas para la CPU nativa y llamadas dentro del proceso; PyTorch 2.14 en la CPU, eager y torch.compile (Inductor), lotes de 64 o 512. Para cada implementación, el número de hilos (1, 8 o 32) y el tamaño de lote se eligieron por la menor latencia en la partición de validación, y después se cronometró la partición de prueba (mediana de 5).

Con un hilo, el orden fue distinto: PyTorch con Inductor y lotes fue más rápido que la compilación de TAKT (DLinear ETTh1: 37,9 ms frente a 111,4 ms; PatchTST ETTh1: 2,03 s frente a 5,78 s). En esa ejecución, la compilación de TAKT fue 5,6–5,9× (DLinear) y 2,2–2,5× (PatchTST) más rápida que el Rust original con los 32 hilos seleccionados, y 7,2–7,3× y 2,3× más rápida con un hilo. Todas las cifras de esa ejecución están en measurements.json (earlier_measurement_2026_10_02).

Contenido

• bin/: ltsf-dlinear-original, ltsf-dlinear-takt, ltsf-patchtst-original, ltsf-patchtst-takt (código Rust con enlace estático, pesos compilados dentro, sin símbolos);
• source/: la versión original: programa de control de línea de comandos, código del modelo generado y pesos;
• tools/: make_windows.py (ventanas de entrada), score.py (MSE/MAE), run_bench.sh (benchmark);
• expected.sha256, measurements.json, SHA256SUMS, LICENSE.

Los datos ETT no se incluyen (CC BY-ND 4.0; los descarga tools/make_windows.py). El código fuente de la compilación de TAKT no se publica: entregamos compilaciones.