NUMA Awareness & Thread Management
Relevant source files
The following files were used as context for generating this wiki page:
Colibri implements NUMA (Non-Uniform Memory Access) awareness and advanced thread management to optimize the inference of massive Mixture-of-Experts (MoE) models. The engine targets large-scale hardware configurations, including multi-socket servers, where memory latency and bandwidth vary across physical CPU nodes. By interleaving resident weights and managing thread affinity, Colibri minimizes cross-socket bottlenecks during high-throughput decode operations.
Thread management focuses on OpenMP optimization, utilizing active wait policies and specific spin counts to reduce re-wake latency in parallel regions. These optimizations are critical for the expert matmul operations, which consist of numerous small, back-to-back computational tasks.
NUMA Awareness Architecture
Colibri handles NUMA-aware expert placement to improve performance on multi-socket systems. When the COLI_NUMA=1 environment variable is active, the engine interleaves resident weights across different memory controllers. This distribution prevents a single memory controller from becoming a bottleneck, especially when streaming large volumes of parameters from RAM.
Sources: README.md:144-145, CHANGELOG.md:158-159, README.it.md:65-67
Memory Interleaving Logic
The engine distributes resident weights (attention, shared experts, and embeddings) across available NUMA nodes. This approach is particularly beneficial for models like GLM-5.2, which requires significant RAM residency even in its quantized int4 form.
The diagram shows the decision path for enabling NUMA interleaving during the model initialization phase.
Sources: README.md:144-145, CHANGELOG.md:158-159
Thread Management & OpenMP Tuning
Colibri uses specific OpenMP configurations to maximize CPU utilization and minimize task-switching overhead. The engine employs a "hot-thread" tuning strategy where the parallel team remains active between computational regions.
OpenMP Environment Configuration
The engine programmatically sets environment variables before execution to ensure optimal thread behavior. If these variables are not present, the engine re-executes itself to apply the tuned settings.
| Variable | Value | Description |
|---|---|---|
OMP_WAIT_POLICY | active | Keeps threads spinning to avoid sleep/wake latency. |
GOMP_SPINCOUNT | 200000 | Sets the number of iterations threads spin before sleeping. |
OMP_PROC_BIND | close | Binds threads to cores close to each other for cache locality. |
OMP_DYNAMIC | FALSE | Disables dynamic adjustment of the number of threads. |
Sources: c/inkling.c:836-848, CHANGELOG.md:95-97
Thread Re-exec Mechanism
To apply these tuned settings effectively, especially on Linux systems, the engine uses the following logic:
#ifndef COLI_CUDA
if (!getenv("COLI_OMP_TUNED") && !getenv("COLI_NO_OMP_TUNE")) {
setenv("OMP_WAIT_POLICY","active",0);
setenv("GOMP_SPINCOUNT","200000",0);
setenv("OMP_PROC_BIND","close",0);
setenv("OMP_DYNAMIC","FALSE",0);
setenv("COLI_OMP_TUNED","1",1);
#ifdef __linux__
execv("/proc/self/exe", argv);
perror("[OMP] execv self-reexec failed, running untuned");
#endif
}
#endif
Sources: c/inkling.c:836-848
Parallel Processing Regions
The engine utilizes #pragma omp parallel directives across several core functions to handle the high volume of parameters in 744B-class models.
Expert Matmul Parallelization
For MoE layers, matmul operations are parallelized at the row level. In the inkling.c implementation, the gate and up projections run as a single matmul over fused rows to reduce the number of parallel regions opened.
The sequence diagram illustrates how the engine reduces overhead by fusing expert matmul operations into fewer parallel regions.
Sources: c/inkling.c:609-620, c/inkling.c:836-840
Depthwise Causal Conv Parallelization
Thread management also extends to depthwise causal short convolutions (sconv_apply). The engine processes individual channels in parallel while maintaining causal history in a thread-local buffer.
Sources: c/inkling.c:237-254
Performance Considerations
- CUDA Interaction: When
COLI_CUDAis enabled, the active spinning policy for OpenMP is typically bypassed. A spinning CPU thread team can starve the CUDA driver during stream synchronization. - Affinity Mask Bug: Previous versions of Colibri had a bug where re-execution preserved the CPU affinity mask, jailing all threads onto a single core. This resulted in a ~20x slowdown. Version 1.1.0 and later ensure correct affinity management.
- Multi-Socket Speedup: NUMA-aware placement provides a measured +13–40% speedup on multi-socket systems compared to standard allocations.
Sources: c/inkling.c:841-843, CHANGELOG.md:95-97, CHANGELOG.md:158-159
Summary
NUMA Awareness and Thread Management are foundational to Colibri's ability to run 744B models on consumer and workstation hardware. By strategically managing memory placement across NUMA nodes and fine-tuning OpenMP thread behavior, the engine achieves high throughput while minimizing the latency inherent in massive Mixture-of-Experts architectures.