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.