LLM-Augmented Bayesian Optimization with Adaptive Policies for Chiplet Design Space Exploration

Hang Liu; Zhengfei Chen; Zhengrui Chen; Hao Geng; Yu Li; Qi Sun

cs.LGOct 10, 2026

LLM-Augmented Bayesian Optimization with Adaptive Policies for Chiplet Design Space Exploration

Authors: Hang Liu, Zhengfei Chen, Zhengrui Chen, Hao Geng, Yu Li, Qi Sun

Organizations: Zhejiang University, Hangzhou, China · ChipFlux, Hangzhou, China · ShanghaiTech University, Shanghai, China

Abstract

Chiplet technology is a key approach for extending Moore's Law. However, its design process is constrained by two primary challenges: the scarcity of early-stage platforms for rapidly evaluating chiplet-based heterogeneous systems, and the complex relationships between design parameters and optimization metrics. To address these challenges, we propose a synergistic framework that integrates a Large Language Model (LLM) with Bayesian Optimization (BO), enabling architectural-level simulation and automated optimization for the 2.5D/3D CPU-GPU heterogeneous chiplets. We extend the open-source chiplet simulator to evaluate 3D designs by supporting graph topologies. Further, to enhance the exploration efficiency of BO, we introduce an LLM that generates high-quality policies to guide the optimization process, enabled by our proposed adaptive directional policy reasoning (ADPR) and an exploration gain-aware fine-tuning method. Finally, we use the policy-guided BO to perform an automated and efficient exploration of the complex chiplet design space. Experiments on 2.5D and 3D chiplet reveal that our framework improves design space exploration efficiency.

Figures & tables

Figure 1. Overview of our automated chiplet design framework, synergizing an extended simulator for architectural evaluation of 2.5D/3D chiplets with LLM-guided Bayesian Optimization for efficient design optimization. Overview of our automated chiplet design framework, synergizing an extended simulator for architectural evaluation of 2.5D/3D chiplets with LLM-guided Bayesian Optimization for efficient design optimization.

Figure 2. Workflow of the LLM-augmented BO framework. Adaptive directional policy reasoning and exploration gain-aware fine-tuning enable the LLM to generate effective directional policies that guide BO. Workflow of the causal policy-driven LLM–BO synergistic framework. Our causal policy synthesis chain-of-thought and exploration gain-aware fine-tuning method train the LLM to generate high-quality causal policies that guide BO.

Figure 3. Detailed optimization flow of our LLM-augmented Bayesian optimization framework. Detailed optimization flow of our causal policy-driven LLM-BO synergistic framework.

Module	Parameters	Candidates
GPU	CoresPerCluster	1, 2
GPU	CacheModel	0, 1
CPU	DispatchWidth	1, 2, 3, 4
	ROBSize	32, 64, 96, 128
	ICacheSize	32, 64
	ICacheAssociativity	4, 8

Table 1. Configurable parameters of the heterogeneous CPU-GPU chiplet architectures.

Dataset	Metric	XGBoost	BOOM Explorer	MCT Explorer	LLAMBO (w/Qwen)	LLAMBO (w/GPT)	LLM-DSE (w/Qwen)	LLM-DSE (w/GPT)	Ours
3D chiplet	HV $\uparrow$	0.7509	0.7466	0.7376	0.7881	0.7549	0.8210	0.8209	0.8356 ( ↑13.3% )
3D chiplet	ADRS $\downarrow$	0.1116	0.1028	0.1378	0.1129	0.0986	0.1405	0.1417	0.0983 ( ↓28.7% )
2.5D chiplet	HV $\uparrow$	0.7502	0.7630	0.7366	0.7732	0.7809	0.7564	0.7638	0.8051 ( ↑ 9.3% )
2.5D chiplet	ADRS $\downarrow$	0.1239	0.0859	0.1239	0.1175	0.1095	0.1315	0.1289	0.0835 ( ↓32.6% )

Table 2. Comparison of DSE performance against state-of-the-art methods. The “Qwen” and “GPT” headers refer to the Qwen-235B and GPT-5.4 models.

Figure 4. Comparison of hypervolume coverage for different methods over time on two datasets, left is 3D chiplet with MLP, and right is 2.5D chiplet with matrix multiplication. Comparison of hypervolume coverage for different methods over time on two datasets, left is 3D chiplet with MLP, and right is 2.5D chiplet with matrix multiplication.

Figure 5. Comparison of runtime distribution and hypervolume across different methods on the MLP benchmark. Comparison of runtime distribution and hypervolume across different methods on the MLP benchmark.

Dataset	$T_{\mathrm{lim}}$ ( ${}^{\circ}\mathrm{C}$ )	Sampled $T_{\mathrm{peak}}$ ( ${}^{\circ}\mathrm{C}$ )	Pareto $T_{\mathrm{peak}}^{\max}$ ( ${}^{\circ}\mathrm{C}$ )
2.5D chiplet	85	[42.3, 61.8]	57.9
3D chiplet	85	[45.6, 68.4]	63.7

Table 3. Thermal feasibility of designs from our method.

Method	Hypervolume $\uparrow$	ADRS $\downarrow$
Standard BO	0.7428	0.1272
Base LLM	0.7764	0.1525
LLM w/ ADPR	0.8117	0.1486
LLM w/ Sim-FT	0.7979	0.1610
LLM w/ Sim-FT + ADPR	0.8247	0.1247
Ours (Gain-Aware FT + ADPR)	0.8356	0.0983

Table 4. Ablation study of our framework’s key components.

Large Language Model-Guided Optimization

Large Language Model-Guided Optimization uses language models to propose, steer, and refine solutions to hard search and design problems.

See all 605 papers

LUMINA: LLM-Guided GPU Architecture Exploration via Bottleneck Analysis
Mar 6, 2026Tao Zhang, Rui Ma, Shuotao Xu +2Automated Heuristic Design Large Language Model-Guided Optimization
Microsoft Research China
HYDRA: A Heterogeneous Chiplet DSE Framework for Serving Dynamic Hybrid LLM Workloads
Aug 19, 2026Jiahao Lin, Alish Kanani, Sangwan Lee +2Hybrid Mamba-Transformer Architectures LLM Serving
University of Wisconsin–Madison · University of Ulsan
AgentDSE: Reasoning-Augmented Architectural Design Space Exploration
Jun 20, 2026Chenyu Wang, Jiahe Caroline Shi, David Kong +4Large Language Model-Guided Optimization Electronic Design Automation
Harvard University · Massachusetts Institute of Technology · Columbia University

Mar 6, 2026J/K move · Enter open · S save

LUMINA: LLM-Guided GPU Architecture Exploration via Bottleneck Analysis

Tao Zhang, Rui Ma, Shuotao Xu, Yongqiang Xiong+1

Microsoft Research China

Automated Heuristic Design Large Language Model-Guided Optimization

Open paper AskPDF arXiv

GPU design space exploration (DSE) for modern AI workloads, such as Large-Language Model (LLM) inference, is challenging because of GPUs' vast, multi-modal design spaces, high simulation costs, and complex design optimization objectives (e.g. performance, power and area trade-offs). Existing automated DSE methods are often prohibitively expensive, either requiring an excessive number of exploration samples or depending on intricate, manually crafted analyses of interdependent critical paths guided by human heuristics. We present LUMINA, an LLM-driven GPU architecture exploration framework that leverage AI to enhance the DSE efficiency and efficacy for GPUs. LUMINA extracts architectural knowledge from simulator code and performs sensitivity studies to automatically compose DSE rules,which are auto-corrected during exploration. A core component of LUMINA is a DSE Benchmark that comprehensively evaluates and enhances LLMs' capabilities across three fundamental skills required for architecture optimization, which provides a principled and reproducible basis for model selection and ensuring consistent architectural reasoning. In the design space with 4.7 million possible samples, LUMINA identifies 6 designs of better performance and area than an A100 GPU efficiently, using only 20 steps via LLM-assisted bottleneck analysis. In comparison, LUMINA achieves 17.5x higher than design space exploration efficiency, and 32.9% better designs (i.e. Pareto Hypervolume) than Machine-Learning baselines, showcasing its ability to deliver high-quality design guidance with minimal search cost.

Abstract

Large Language Model-Guided Optimization

Explore similar work

LUMINA: LLM-Guided GPU Architecture Exploration via Bottleneck Analysis

HYDRA: A Heterogeneous Chiplet DSE Framework for Serving Dynamic Hybrid LLM Workloads

AgentDSE: Reasoning-Augmented Architectural Design Space Exploration