Large Foundation Model (LFM) inference is both memory- and compute-intensive, traditionally relying on GPUs. However, the limited availability and high cost have motivated the adoption of high-performance general-purpose CPUs, especially emerging 3D-stacked Static Non-Uniform Cache Architecture (3D S-NUCA) systems. These architectures offer enhanced bandwidth and locality but suffer from severe thermal challenges and uneven cache latencies due to 3D Networks-on-Chip (NoC). Optimal management of thread migration and V/f scaling is non-trivial due to LFM kernel diversity and system heterogeneity. Existing thermal management approaches often rely on oversimplified analytical models and lack adaptability. We propose AILFM, an Active Imitation Learning (AIL)-based scheduling framework that learns near-optimal thermal-aware scheduling policies from Oracle demonstrations with minimal run-time overhead. AILFM accounts for both core-level performance heterogeneity and kernel-specific behavior in LFMs to maintain thermal safety while maximizing performance. Extensive experiments show that AILFM outperforms state-of-the-art baselines and generalizes well across diverse LFM workloads.
Figures & tables
Figure 1. An abstract representation of 3D-stacked S-NUCA systems.
Figure 2. Processing core AMDs in 3D-stacked S-NUCA processors
Kernel
IPS ↑
LLC MPKI ↓
3.0
3.5
4.0
4.5
3.0
3.5
4.0
4.5
Embedding
6.23
6.01
5.82
5.51
10
13
17
22
Self-Attention
7.92
6.83
6.03
5.10
21
25
32
46
FFN
6.59
6.42
6.21
6.14
7
7
9
11
LM Head
5.77
5.29
4.94
4.36
11
15
19
37
Table 1. IPS (×10 9 ) and LLC MPKI of ViT kernels at varying AMDs.
Figure 3. Overview of the Oracle and Learning Policies in AI L FM for thermal management of 3D S-NUCA systems. The Oracle Policy π∗ uses MoGPR to model MPKI–IPS–thermal dynamics and label kernel-migration utilities. The Learning Policy πθ applies MC Dropout for uncertainty estimation, enabling selective Oracle queries and confident autonomous actions.
Figure 4. Illustration of the MoGPR-based Oracle in AI L FM . The selected GP predicts migration utility from source and target core.
Figure 5. Normalized performance of AI L FM across LFMs under varying input sequence lengths, with values referenced to AI L FM at L=128 .
Figure 6. Normalized performance of AI L FM versus baselines.
Figure 7. Tpeak comparison between baselines of 3D S-NUCA.
Figure 8. Run-time overhead with increasing active cores.
Architecture
Norm. Time ↓
Tpeak [ ∘ C] ↓
speedup ↑
2D planar S-NUCA (8 × 8, 64 cores)
1.00
73.5
-
3D cores w/o S-NUCA (4 × 4 × 4, 64 cores)
0.88
77.1
1.14 ×
3D S-NUCA (4 × 4 × 4, 64 cores)
0.67
75.2
1.49 ×
3D cores w/o S-NUCA (6 × 6 × 6, 216 cores)
0.79
81.6
1.27 ×
3D S-NUCA (6 × 6 × 6, 216 cores)
0.61
83.4
1.64 ×
Table 2. Scalability across cache organizations and 3D topologies.
Figure 9. Comparison of AI L FM and DLFM shows AI L FM delivers higher IPS and lower, more stable temperatures via its MoGPR.
Method
Accuracy [%] ↑
Norm. Time ↓
Tpeak [ ∘ C] ↓
Overhead [%] ↓
Oracle (MoGPR, upper bound)
97.3
1
74.3
24.6
AI L FM (ours)
96.7
0.13
75.1
3.2
Table 3. Oracle upper bound versus AI L FM .
Model
Utility MAE ↓
Exec. Time (s) ↓
Queries/Epoch ↓
Random Forest
0.42
103.7
0.21
MLP
0.39
150.9
0.18
MoGPR (ours)
0.31
79.1
0.11
Table 4. Utility prediction accuracy with different models.
τ
Safe overrides (%) ↓
Violations (%) ↓
Overhead (%) ↓
0.05
7.9
0.07
2.5
0.10
6.1
0.09
1.0
0.15
5.4
0.07
0.2
0.20
5.1
0.23
1.6
0.30
4.8
0.92
3.2
Table 5. Sensitivity of the uncertainty threshold τ .
Method
Acc. [%] ↑
Norm. Time ↓
Tpeak [ ∘ C] ↓
AI L FM (ours)
96.7
0.77
75.1
✗ Kernel-aware
90.8
0.92
79.9
✗ Active querying
91.1
0.95
81.4
DLFM (direct learning)
90.3
1.00
82.7
Table 6. Ablation of kernel-aware modeling and active querying.
Although existing frameworks for large language model (LLM) inference on CPUs are mature, they fail to fully exploit the computation potential of many-core CPU platforms. Many-core CPUs are widely deployed in web servers and high-end networking devices, and are typically organized into multiple NUMA nodes that group cores and memory. Current frameworks largely overlook the substantial overhead of cross-NUMA memory access, limiting inference scalability and intelligence enabling on such platforms. To address this limitation, we build ArcLight, a lightweight LLM inference architecture designed from the ground up for many-core CPUs. ArcLight integrates efficient memory management and thread scheduling, and introduces finely controlled tensor parallelism to mitigate the cross-node memory access wall. Experimental results show that ArcLight significantly surpasses the performance ceiling of mainstream frameworks, achieving up to 46% higher inference throughput. Moreover, ArcLight maintains compatibility with arbitrary CPU devices. ArcLight is publicly available at https://github.com/OpenBMB/ArcLight.
Yuzhuang Xu, Xu Han, Yuxuan Li +1
Harbin Institute of Technology, Harbin, China · Tsinghua University, Beijing, China
Processing-In-Memory (PIM) has emerged as a promising technology for accelerating machine learning (ML) workloads. Specifically, non-volatile memory-based PIM architectures have enabled effective ML acceleration due to their ability to perform energy-efficient matrix-vector multiplication operations. However, these devices suffer from non-idealities such as thermal noise. This noise alters the stored values in the memory cells which correspond to actual model weights, compromising the inference accuracy. In this work, we introduce ThRIve, a noise-aware training methodology that leverages low-rank adaptation to enable thermally robust inference on heterogeneous PIM architectures. ThRIve selectively stores these low-rank noise-aware parameters on a hardware that is less susceptible to thermal noise, enabling robustness against temperature-induced noise variations. ThRIve mitigates the effects of thermal-noise and prevent the drop in inference accuracy across the entire operating temperature range. Experimental results demonstrate that ThRIve-enabled architectures maintain consistent inference accuracy, with the mean accuracy staying within 2% of the ideal (i.e., noise-free) accuracy, and the variation in accuracy across the entire operating temperature range remaining within 2% of the mean. The proposed methodology achieves accuracy and robustness comparable to thermally-resilient Static Random-Access Memory (SRAM)-based PIM systems, while delivering up to 5.4x reduction in energy-delay product (EDP) during CNN model inferencing.
Modern CPUs increasingly integrate matrix extensions, such as Arm Scalable Matrix Extension (SME), that provide high-throughput matrix execution within the CPU. For LLM inference, however, these units are not a universal replacement for conventional CPU cores: prefill, decode, attention, and KV-cache operations expose different arithmetic intensities, vector behavior, and layout requirements, while SME units and CPU cores still compete for shared memory bandwidth. This paper studies this mismatch through a roofline-based characterization of SME-enabled CPUs and uses the resulting model to guide operator-level execution choices. We present SMEPilot, an LLM inference engine that selects CPU-only, SME-only, or cooperative SME+CPU execution for each operator shape. SMEPilot partitions matrix work across SME and CPU cores at tile granularity, overlaps SME-suitable matrix stages with CPU-suitable vector stages in attention, and maintains layout state so packed tensor representations are reused rather than repeatedly rebuilt on critical paths. Across Llama-3.2-3B, Qwen3-4B, and Qwen3-30BA3B on phone, PC, and server platforms, SMEPilot improves end-to-end inference performance by up to 3.94×.