NumKong: Optimización de Operaciones Numéricas de Precisión Mixta en CPUs Modernas

El problema fundamental que NumKong aborda es la brecha de rendimiento y la falta de soporte para tipos de datos numéricos de precisión mixta (como Float16, BFloat16, Float8, Int4) en las bibliotecas de álgebra lineal tradicionales (BLAS) y frameworks de ML en CPUs. A medida que la industria de la IA se inclina hacia modelos más grandes y la inferencia en el borde, la necesidad de reducir el footprint de memoria y aumentar el throughput computacional sin sacrificar la estabilidad numérica se vuelve crítica. Las CPUs modernas han comenzado a incorporar "tensor cores" o extensiones de matriz (AMX en Intel, SME en Arm, RVV en RISC-V) que permiten operaciones de multiplicación-acumulación (MAC) de alta densidad con estos tipos de baja precisión. Sin embargo, la programación de estas extensiones es compleja y específica de cada arquitectura, lo que dificulta la portabilidad y la optimización.

NumKong se posiciona como una solución que abstrae esta complejidad, proporcionando una colección de kernels altamente optimizados que explotan las capacidades SIMD y de "tiled tensor operations" de las CPUs contemporáneas. Esto permite a los ingenieros Staff+ y Arquitectos aprovechar el hardware subyacente para acelerar cargas de trabajo intensivas en números, como la búsqueda vectorial, el procesamiento de imágenes y la inferencia de LLM, en un rango más amplio de plataformas y lenguajes de programación, más allá de los ecosistemas dominantes de GPU y Python. La tesis es que, con una ingeniería de bajo nivel cuidadosa, las CPUs pueden desempeñar un papel mucho más significativo en la computación de precisión mixta, complementando a las GPUs y extendiendo las capacidades de IA al edge y al navegador.

Arquitectura del Sistema

NumKong se estructura como una colección de kernels SIMD de bajo nivel, escritos en C99, C++23, Rust, Swift, JavaScript, GoLang y Python, que se dirigen a arquitecturas de CPU específicas: Intel (AVX-512, AMX, VNNI), Arm (NEON, SVE, SME) y RISC-V (RVV). La arquitectura se basa en un diseño modular donde cada kernel está altamente especializado para una operación numérica y un tipo de dato particular, aprovechando las instrucciones SIMD nativas y las extensiones de matriz de cada ISA.

Las decisiones de diseño clave incluyen: (1) Soporte explícito para precisión mixta y sub-byte: Implementación de tipos como Float16, BFloat16, Float8, Float6, Int4/UInt4, con conversiones y operaciones optimizadas. (2) Optimización de "tiled tensor operations": Aprovechamiento de AMX (Intel) y SME (Arm) para operaciones de multiplicación-acumulación de matrices (GEMM) de alta densidad. (3) Manejo explícito de la memoria: NumKong no realiza asignaciones ocultas ni gestiona hilos internamente; el usuario es responsable de los buffers y la paralelización, lo que permite un control preciso y evita problemas de concurrencia. (4) Fusión de epílogos: Las operaciones comunes post-GEMM (como normalización, cálculo de distancias L2, argmax) se fusionan directamente en el kernel de GEMM para reducir el movimiento de datos y las pasadas adicionales. (5) Algoritmos de alta precisión: Implementación de algoritmos como Dot2 (Ogita, Rump, Oishi) para sumas compensadas en Float64, garantizando una precisión sub-ULP incluso en vectores largos. (6) Patrones de diseño de kernels: Uso de técnicas como transformaciones algebraicas para dot products simétricos en Int8/Int4, lookups tabulares para conversiones de tipos de datos, y máscaras de convergencia para algoritmos iterativos (ej. Vincenty). (7) Interoperabilidad: Integración con ecosistemas existentes a través del buffer protocol de Python y punteros crudos en C/C++, permitiendo el uso con NumPy, PyTorch y tensores de GPU (via CUDA unified memory).

```c
// Simplified example for signed Int8 dot product
// a: signed Int8 vector, b: signed Int8 vector
// sum_b: sum of elements in b

__m512i a_transformed = _mm512_xor_si512(a_vec, _mm512_set1_epi8(0x80));
__m512i dot_product_partial = _mm512_dpbusd_epi32(acc, a_transformed, b_vec);
__m512i sum_b_partial = _mm512_sad_epu8(b_vec, _mm512_setzero_si512());

// After loop, final correction:
result -= (sum_b << 7); // Equivalent to sum_b * 128
```

```c
// Simplified example for E2M3 to BFloat16 conversion on x86 AVX-512
// e2m3_val: 8-bit E2M3 value
// lut_table: precomputed 32-entry BFloat16 lookup table in a ZMM register

uint8_t magnitude = e2m3_val & 0x1F; // Mask for 5-bit magnitude
__m512i sign_bit = _mm512_slli_epi16(_mm512_set1_epi8(e2m3_val & 0x80), 7); // Shift sign to BF16 position
__m512i bf16_magnitude = _mm512_permutexvar_epi16(_mm512_set1_epi8(magnitude), lut_table);
__m512i result = _mm512_or_si512(bf16_magnitude, sign_bit);
```

```c
// Simplified example for Vincenty iteration with AVX-512
// lambda_vec: vector of current lambda values
// prev_lambda_vec: vector of previous lambda values
// converged_mask: __mmask8 tracking converged lanes

while (converged_mask != 0xFF) { // Loop until all 8 lanes converge
    __m512d delta_lambda = _mm512_sub_pd(lambda_vec, prev_lambda_vec);
    // ... compute new_lambda_vec ...

    __mmask8 new_converged_mask = _mm512_cmp_pd_mask(new_lambda_vec, lambda_vec, _CMP_EQ_OQ);
    converged_mask |= new_converged_mask;

    // Guard divisions with masks
    __m512d div_result = _mm512_maskz_div_pd(converged_mask, numerator, denominator_for_coincident_points);
    // ... other masked operations ...

    prev_lambda_vec = lambda_vec;
    lambda_vec = new_lambda_vec;
}
```

Fundamentos Teóricos

El concepto de operaciones de multiplicación-acumulación (MAC) de alta densidad y precisión mixta tiene profundas raíces en la computación de alto rendimiento y la álgebra lineal numérica. Los algoritmos BLAS (Basic Linear Algebra Subprograms), estandarizados desde 1979, son el pilar de la computación científica, pero su diseño original se centró en Float32 y Float64. La evolución hacia tipos de datos de menor precisión en hardware moderno, como BFloat16 y Float16, y más recientemente Float8 y Float6, se alinea con la investigación en redes neuronales profundas, donde la tolerancia a errores numéricos permite ganancias masivas en rendimiento y eficiencia energética. Esto se predijo en trabajos como "Mixed-Precision Training for Deep Neural Networks" (Micikevicius et al., 2017), que demostró los beneficios de usar Float16 para acelerar el entrenamiento.

La implementación de sumas compensadas en NumKong, como el algoritmo Dot2, se basa directamente en la teoría de la aritmética de punto flotante de alta precisión, con trabajos seminales de Ogita, Rump y Oishi ("Accurate Summation and Dot Product with Applications to Interval Arithmetic", 2005) que proporcionan los fundamentos para garantizar una precisión O(1) independientemente de la longitud del vector. De manera similar, el esquema de Ozaki para la multiplicación de matrices de doble precisión utilizando hardware de precisión simple se basa en la investigación de Ozaki, Ogita, Oishi y Rump sobre la multiplicación de matrices precisas. La optimización de la alineación de mallas (Kabsch, Umeyama) se basa en algoritmos clásicos de álgebra lineal para la descomposición de valores singulares (SVD), como el método de Jacobi, adaptados para la vectorización SIMD. La gestión de la coherencia numérica y la eficiencia en el uso de la memoria caché, incluyendo la mitigación de "cache aliasing" con padding, se conecta con principios fundamentales de la arquitectura de computadoras y el diseño de algoritmos para sistemas jerárquicos de memoria, un tema recurrente en la literatura de optimización de rendimiento desde los primeros días de la computación.